yujunjun/OpenMetadata
1
0
Fork 0
mirror of https://github.com/open-metadata/OpenMetadata.git synced 2025-12-04 19:44:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix Bigquery Test connection for multiproject (#13380)

Browse Source 
Co-authored-by: Ayush Shah <ayush@getcollate.io>
This commit is contained in:
Mayur Singal 2023-10-05 14:50:42 +05:30 committed by GitHub
parent f879656f0a
commit 0090286924
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 22 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
ingestion/src/metadata/ingestion/source/database/bigquery/metadata.py
View File

@ -197,10 +197,19 @@ class BigquerySource(StoredProcedureMixin, CommonDbSourceService):
"""
def __init__(self, config, metadata):
# Check if the engine is established before setting project IDs
# This ensures that we don't try to set project IDs when there is no engine
# as per service connection config, which would result in an error.
self.test_connection = lambda: None
super().__init__(config, metadata)
self.temp_credentials = None
self.client = None
# Upon invoking the set_project_id method, we retrieve a comprehensive
# list of all project IDs. Subsequently, after the invokation,
# we proceed to test the connections for each of these project IDs
self.project_ids = self.set_project_id()
self.test_connection = self._test_connection
self.test_connection()
@classmethod
def create(cls, config_dict, metadata: OpenMetadata):
@ -217,11 +226,15 @@ class BigquerySource(StoredProcedureMixin, CommonDbSourceService):
_, project_ids = auth.default()
return project_ids if isinstance(project_ids, list) else [project_ids]
def test_connection(self) -> None:
for project_id in self.set_project_id():
self.set_inspector(project_id)
def _test_connection(self) -> None:
for project_id in self.project_ids:
inspector_details = get_inspector_details(
database_name=project_id, service_connection=self.service_connection
)
test_connection_fn = get_test_connection_fn(self.service_connection)
test_connection_fn(self.metadata, self.engine, self.service_connection)
test_connection_fn(
self.metadata, inspector_details.engine, self.service_connection
)
def query_table_names_and_types(
self, schema_name: str

5
ingestion/tests/unit/test_handle_partitions.py
View File

@ -97,17 +97,20 @@ class BigqueryUnitTest(TestCase):
"metadata.ingestion.source.database.bigquery.metadata.BigquerySource.set_project_id"
)
@patch(
"metadata.ingestion.source.database.bigquery.metadata.BigquerySource.test_connection"
"metadata.ingestion.source.database.bigquery.metadata.BigquerySource._test_connection"
)
@patch("metadata.ingestion.source.database.common_db_source.get_connection")
def __init__(
self,
methodName,
get_connection_common,
test_connection,
set_project_id,
create_generic_connection,
client,
) -> None:
super().__init__(methodName)
get_connection_common.return_value = Mock()
client.return_value = Mock()
create_generic_connection.return_value = Mock()
set_project_id.return_value = Mock()

2
ingestion/tests/unit/topology/database/test_bigquery.py
View File

@ -53,7 +53,7 @@ EXPECTED_URL = "https://console.cloud.google.com/bigquery?project=random-project
class BigqueryUnitTest(TestCase):
@patch(
"metadata.ingestion.source.database.bigquery.metadata.BigquerySource.test_connection"
"metadata.ingestion.source.database.bigquery.metadata.BigquerySource._test_connection"
)
@patch(
"metadata.ingestion.source.database.bigquery.metadata.BigquerySource.set_project_id"