openmetadata-ingestion 1.1.5.0__py3-none-any.whl → 1.1.6.0__py3-none-any.whl

metadata/ingestion/source/database/bigquery/query_parser.py

@@ -12,6 +12,7 @@
 Handle big query usage extraction
 """
 from abc import ABC
+from copy import deepcopy
 from datetime import datetime
 
 from google import auth
@@ -25,7 +26,9 @@ from metadata.generated.schema.entity.services.connections.metadata.openMetadata
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.ingestion.api.source import InvalidSourceException
+from metadata.generated.schema.security.credentials.gcpValues import MultipleProjectId
+from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.source.database.bigquery.helper import get_inspector_details
 from metadata.ingestion.source.database.query_parser_source import QueryParserSource
 
 
@@ -65,3 +68,18 @@ class BigqueryQueryParserSource(QueryParserSource, ABC):
     def set_project_id():
         _, project_id = auth.default()
         return project_id
+
+    def get_engine(self):
+        if isinstance(
+            self.service_connection.credentials.gcpConfig.projectId, MultipleProjectId
+        ):
+            project_ids = deepcopy(
+                self.service_connection.credentials.gcpConfig.projectId
+            )
+            for project_id in project_ids.__root__:
+                inspector_details = get_inspector_details(
+                    project_id, self.service_connection
+                )
+                yield inspector_details.engine
+        else:
+            yield self.engine

metadata/ingestion/source/database/lineage_source.py

@@ -78,25 +78,28 @@ class LineageSource(QueryParserSource, ABC):
         Given an engine, iterate over the query results to
         yield a TableQuery with query parsing info
         """
-        with self.engine.connect() as conn:
-            rows = conn.execute(
-                self.get_sql_statement(
-                    start_time=self.start,
-                    end_time=self.end,
-                )
-            )
-            for row in rows:
-                query_dict = dict(row)
-                try:
-                    yield TableQuery(
-                        query=query_dict["query_text"],
-                        databaseName=self.get_database_name(query_dict),
-                        serviceName=self.config.serviceName,
-                        databaseSchema=self.get_schema_name(query_dict),
+        for engine in self.get_engine():
+            with engine.connect() as conn:
+                rows = conn.execute(
+                    self.get_sql_statement(
+                        start_time=self.start,
+                        end_time=self.end,
                     )
-                except Exception as exc:
-                    logger.debug(traceback.format_exc())
-                    logger.warning(f"Error processing query_dict {query_dict}: {exc}")
+                )
+                for row in rows:
+                    query_dict = dict(row)
+                    try:
+                        yield TableQuery(
+                            query=query_dict["query_text"],
+                            databaseName=self.get_database_name(query_dict),
+                            serviceName=self.config.serviceName,
+                            databaseSchema=self.get_schema_name(query_dict),
+                        )
+                    except Exception as exc:
+                        logger.debug(traceback.format_exc())
+                        logger.warning(
+                            f"Error processing query_dict {query_dict}: {exc}"
+                        )
 
     def next_record(self) -> Iterable[AddLineageRequest]:
         """

metadata/ingestion/source/database/query_parser_source.py

@@ -112,6 +112,9 @@ class QueryParserSource(Source[Union[TableQuery, AddLineageRequest]], ABC):
             return f"{self.filters} AND {self.source_config.filterCondition}"
         return self.filters
 
+    def get_engine(self):
+        yield self.engine
+
     def close(self):
         """
         By default, there is nothing to close

metadata/ingestion/source/database/usage_source.py

@@ -93,37 +93,38 @@ class UsageSource(QueryParserSource, ABC):
                 f"{(self.start + timedelta(days=days + 1)).date()}"
             )
             try:
-                with self.engine.connect() as conn:
-                    rows = conn.execute(
-                        self.get_sql_statement(
-                            start_time=self.start + timedelta(days=days),
-                            end_time=self.start + timedelta(days=days + 1),
+                for engine in self.get_engine():
+                    with engine.connect() as conn:
+                        rows = conn.execute(
+                            self.get_sql_statement(
+                                start_time=self.start + timedelta(days=days),
+                                end_time=self.start + timedelta(days=days + 1),
+                            )
                         )
-                    )
-                    queries = []
-                    for row in rows:
-                        row = dict(row)
-                        try:
-                            queries.append(
-                                TableQuery(
-                                    query=self.format_query(row["query_text"]),
-                                    userName=row["user_name"],
-                                    startTime=str(row["start_time"]),
-                                    endTime=str(row["end_time"]),
-                                    analysisDate=row["start_time"],
-                                    aborted=self.get_aborted_status(row),
-                                    databaseName=self.get_database_name(row),
-                                    duration=row.get("duration"),
-                                    serviceName=self.config.serviceName,
-                                    databaseSchema=self.get_schema_name(row),
+                        queries = []
+                        for row in rows:
+                            row = dict(row)
+                            try:
+                                queries.append(
+                                    TableQuery(
+                                        query=self.format_query(row["query_text"]),
+                                        userName=row["user_name"],
+                                        startTime=str(row["start_time"]),
+                                        endTime=str(row["end_time"]),
+                                        analysisDate=row["start_time"],
+                                        aborted=self.get_aborted_status(row),
+                                        databaseName=self.get_database_name(row),
+                                        duration=row.get("duration"),
+                                        serviceName=self.config.serviceName,
+                                        databaseSchema=self.get_schema_name(row),
+                                    )
                                 )
-                            )
-                        except Exception as exc:
-                            logger.debug(traceback.format_exc())
-                            logger.warning(
-                                f"Unexpected exception processing row [{row}]: {exc}"
-                            )
-                yield TableQueries(queries=queries)
+                            except Exception as exc:
+                                logger.debug(traceback.format_exc())
+                                logger.warning(
+                                    f"Unexpected exception processing row [{row}]: {exc}"
+                                )
+                    yield TableQueries(queries=queries)
             except Exception as exc:
                 logger.debug(traceback.format_exc())
                 logger.error(f"Source usage processing error: {exc}")