diff --git a/ingestion/src/metadata/ingestion/source/database/bigquery/metadata.py b/ingestion/src/metadata/ingestion/source/database/bigquery/metadata.py index 2a5f055f071..eed889e84bd 100644 --- a/ingestion/src/metadata/ingestion/source/database/bigquery/metadata.py +++ b/ingestion/src/metadata/ingestion/source/database/bigquery/metadata.py @@ -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 diff --git a/ingestion/tests/unit/test_handle_partitions.py b/ingestion/tests/unit/test_handle_partitions.py index c0ff3afa363..038454c2d92 100644 --- a/ingestion/tests/unit/test_handle_partitions.py +++ b/ingestion/tests/unit/test_handle_partitions.py @@ -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() diff --git a/ingestion/tests/unit/topology/database/test_bigquery.py b/ingestion/tests/unit/topology/database/test_bigquery.py index f89358a84c1..089b40adbb6 100644 --- a/ingestion/tests/unit/topology/database/test_bigquery.py +++ b/ingestion/tests/unit/topology/database/test_bigquery.py @@ -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"