openmetadata-ingestion 1.7.5.3__py3-none-any.whl → 1.8.1.0__py3-none-any.whl

metadata/ingestion/source/dashboard/tableau/client.py

@@ -110,7 +110,7 @@ class TableauClient:
             owner = self.tableau_server.users.get_by_id(owner_id) if owner_id else None
             if owner and owner.email:
                 owner_obj = TableauOwner(
-                    id=str(owner.id), name=owner.name, email=owner.email
+                    id=owner.id, name=owner.name, email=owner.email
                 )
                 self.owner_cache[owner_id] = owner_obj
                 return owner_obj
@@ -130,7 +130,7 @@ class TableauClient:
             try:
                 charts.append(
                     TableauChart(
-                        id=str(view.id),
+                        id=view.id,
                         name=view.name,
                         tags=view.tags,
                         owner=self.get_tableau_owner(view.owner_id),
@@ -212,10 +212,10 @@ class TableauClient:
                     workbook.views
                 )
                 workbook = TableauDashboard(
-                    id=str(workbook.id),
+                    id=workbook.id,
                     name=workbook.name,
                     project=TableauBaseModel(
-                        id=str(workbook.project_id), name=workbook.project_name
+                        id=workbook.project_id, name=workbook.project_name
                     ),
                     owner=self.get_tableau_owner(workbook.owner_id),
                     description=workbook.description,

metadata/ingestion/source/dashboard/tableau/metadata.py

@@ -257,7 +257,6 @@ class TableauSource(DashboardServiceSource):
                 ),
                 sql=self._get_datamodel_sql_query(data_model=data_model),
                 owners=self.get_owner_ref(dashboard_details=dashboard_details),
-                project=self.get_project_name(dashboard_details=dashboard_details),
             )
             yield Either(right=data_model_request)
             self.register_record_datamodel(datamodel_request=data_model_request)

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

@@ -68,21 +68,25 @@ def get_connection_url(connection: BigQueryConnection) -> str:
             connection.credentials.gcpConfig.projectId, SingleProjectId
         ):
             if not connection.credentials.gcpConfig.projectId.root:
-                return f"{connection.scheme.value}://{connection.credentials.gcpConfig.projectId or ''}"
+                return f"{connection.scheme.value}://{connection.billingProjectId or connection.credentials.gcpConfig.projectId.root or ''}"
             if (
                 not connection.credentials.gcpConfig.privateKey
                 and connection.credentials.gcpConfig.projectId.root
             ):
                 project_id = connection.credentials.gcpConfig.projectId.root
-                os.environ["GOOGLE_CLOUD_PROJECT"] = project_id
-            return f"{connection.scheme.value}://{connection.credentials.gcpConfig.projectId.root}"
+                os.environ["GOOGLE_CLOUD_PROJECT"] = (
+                    connection.billingProjectId or project_id
+                )
+            return f"{connection.scheme.value}://{connection.billingProjectId or connection.credentials.gcpConfig.projectId.root}"
         elif isinstance(connection.credentials.gcpConfig.projectId, MultipleProjectId):
             for project_id in connection.credentials.gcpConfig.projectId.root:
                 if not connection.credentials.gcpConfig.privateKey and project_id:
                     # Setting environment variable based on project id given by user / set in ADC
-                    os.environ["GOOGLE_CLOUD_PROJECT"] = project_id
-                return f"{connection.scheme.value}://{project_id}"
-            return f"{connection.scheme.value}://"
+                    os.environ["GOOGLE_CLOUD_PROJECT"] = (
+                        connection.billingProjectId or project_id
+                    )
+                return f"{connection.scheme.value}://{connection.billingProjectId or project_id}"
+            return f"{connection.scheme.value}://{connection.billingProjectId or ''}"
 
     # If gcpConfig is the JSON key path and projectId is defined, we use it by default
     elif (
@@ -92,13 +96,13 @@ def get_connection_url(connection: BigQueryConnection) -> str:
         if isinstance(  # pylint: disable=no-else-return
             connection.credentials.gcpConfig.projectId, SingleProjectId
         ):
-            return f"{connection.scheme.value}://{connection.credentials.gcpConfig.projectId.root}"
+            return f"{connection.scheme.value}://{connection.billingProjectId or connection.credentials.gcpConfig.projectId.root}"
 
         elif isinstance(connection.credentials.gcpConfig.projectId, MultipleProjectId):
             for project_id in connection.credentials.gcpConfig.projectId.root:
-                return f"{connection.scheme.value}://{project_id}"
+                return f"{connection.scheme.value}://{connection.billingProjectId or project_id}"
 
-    return f"{connection.scheme.value}://"
+    return f"{connection.scheme.value}://{connection.billingProjectId or ''}"
 
 
 def get_connection(connection: BigQueryConnection) -> Engine:

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

@@ -72,7 +72,9 @@ def get_inspector_details(
                 "lifetime"
             ] = new_service_connection.credentials.gcpImpersonateServiceAccount.lifetime
 
-    client = get_bigquery_client(project_id=database_name, **kwargs)
+    client = get_bigquery_client(
+        project_id=new_service_connection.billingProjectId or database_name, **kwargs
+    )
     engine = get_connection(new_service_connection)
     inspector = inspect(engine)
 
@@ -86,15 +88,15 @@ def get_pk_constraint(
     This function overrides to get primary key constraint
     """
     try:
-        constraints = PK_CACHE.get(f"{connection.engine.url.host}.{schema}")
+        project, schema = schema.split(".")
+        constraints = PK_CACHE.get(f"{project}.{schema}")
         if constraints is None:
             constraints = connection.engine.execute(
                 BIGQUERY_TABLE_CONSTRAINTS.format(
-                    project_id=connection.engine.url.host,
-                    schema_name=schema,
+                    project_id=project, schema_name=schema
                 )
             )
-            PK_CACHE[f"{connection.engine.url.host}.{schema}"] = constraints.fetchall()
+            PK_CACHE[f"{project}.{schema}"] = constraints.fetchall()
 
         col_name = []
         table_constraints = [row for row in constraints if row.table_name == table_name]
@@ -116,15 +118,15 @@ def get_foreign_keys(
     This function overrides to get foreign key constraint
     """
     try:
-        constraints = FK_CACHE.get(f"{connection.engine.url.host}.{schema}")
+        project, schema = schema.split(".")
+        constraints = FK_CACHE.get(f"{project}.{schema}")
         if constraints is None:
             constraints = connection.engine.execute(
                 BIGQUERY_FOREIGN_CONSTRAINTS.format(
-                    project_id=connection.engine.url.host,
-                    schema_name=schema,
+                    project_id=project, schema_name=schema
                 )
             )
-            FK_CACHE[f"{connection.engine.url.host}.{schema}"] = constraints.fetchall()
+            FK_CACHE[f"{project}.{schema}"] = constraints.fetchall()
 
         col_name = []
         table_constraints = [row for row in constraints if row.table_name == table_name]

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

@@ -83,7 +83,10 @@ from metadata.ingestion.source.database.bigquery.models import (
     BigQueryStoredProcedure,
 )
 from metadata.ingestion.source.database.bigquery.queries import (
+    BIGQUERY_GET_MATERIALIZED_VIEW_NAMES,
+    BIGQUERY_GET_SCHEMA_NAMES,
     BIGQUERY_GET_STORED_PROCEDURES,
+    BIGQUERY_GET_VIEW_NAMES,
     BIGQUERY_LIFE_CYCLE_QUERY,
     BIGQUERY_SCHEMA_DESCRIPTION,
     BIGQUERY_TABLE_AND_TYPE,
@@ -238,6 +241,7 @@ Inspector.get_all_table_ddls = get_all_table_ddls
 Inspector.get_table_ddl = get_table_ddl
 
 
+# pylint: disable=too-many-public-methods
 class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
     """
     Implements the necessary methods to extract
@@ -347,6 +351,33 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
             logger.debug(traceback.format_exc())
             raise InvalidSourceException(f"Error setting BigQuery project IDs: {exc}")
 
+    # pylint: disable=arguments-differ
+    def _get_columns_with_constraints(
+        self, schema_name: str, table_name: str, inspector: Inspector
+    ) -> Tuple[List, List, List]:
+        database_name = self.context.get().database
+        schema_name = f"{database_name}.{schema_name}"
+        return super()._get_columns_with_constraints(schema_name, table_name, inspector)
+
+    def _get_columns_internal(
+        self,
+        schema_name: str,
+        table_name: str,
+        db_name: str,
+        inspector: Inspector,
+        table_type: TableType = None,
+    ):
+        """
+        Get columns list
+        """
+
+        return inspector.get_columns(
+            table_name,
+            f"{db_name}.{schema_name}",
+            table_type=table_type,
+            db_name=db_name,
+        )
+
     def _test_connection(self) -> None:
         for project_id in self.project_ids:
             inspector_details = get_inspector_details(
@@ -375,7 +406,7 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
         table_names_and_types = (
             self.engine.execute(
                 BIGQUERY_TABLE_AND_TYPE.format(
-                    project_id=self.client.project, schema_name=schema_name
+                    project_id=self.context.get().database, schema_name=schema_name
                 )
             )
             or []
@@ -409,8 +440,6 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
         logic on how to handle table types, e.g., material views,...
         """
 
-        view_names = self.inspector.get_view_names(schema_name) or []
-
         if self.incremental.enabled:
             view_names = [
                 view_name
@@ -419,10 +448,50 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
                 in self.incremental_table_processor.get_not_deleted(schema_name)
             ]
 
-        return [
-            TableNameAndType(name=view_name, type_=TableType.View)
-            for view_name in view_names
-        ]
+        table_name_and_types = []
+        for table_type, query in {
+            TableType.View: BIGQUERY_GET_VIEW_NAMES,
+            TableType.MaterializedView: BIGQUERY_GET_MATERIALIZED_VIEW_NAMES,
+        }.items():
+            view_names = list(
+                map(
+                    lambda x: x[0],
+                    (
+                        self.engine.execute(
+                            query.format(
+                                project=self.context.get().database, dataset=schema_name
+                            )
+                        )
+                        or []
+                    ),
+                )
+            )
+            if self.incremental.enabled:
+                view_names = [
+                    view_name
+                    for view_name in view_names
+                    if view_name
+                    in self.incremental_table_processor.get_not_deleted(schema_name)
+                ]
+
+            table_name_and_types.extend(
+                [
+                    TableNameAndType(name=view_name, type_=table_type)
+                    for view_name in view_names
+                ]
+            )
+
+        return table_name_and_types
+
+    # pylint: disable=arguments-differ
+    @calculate_execution_time()
+    def get_table_description(
+        self, schema_name: str, table_name: str, inspector: Inspector
+    ) -> str:
+        schema_name = f"{self.context.get().database}.{schema_name}"
+        return super().get_table_description(
+            schema_name=schema_name, table_name=table_name, inspector=inspector
+        )
 
     def yield_tag(
         self, schema_name: str
@@ -430,7 +499,9 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
         """Build tag context"""
         try:
             # Fetching labels on the databaseSchema ( dataset ) level
-            dataset_obj = self.client.get_dataset(schema_name)
+            dataset_obj = self.client.get_dataset(
+                f"{self.context.get().database}.{schema_name}"
+            )
             if dataset_obj.labels:
                 for key, value in dataset_obj.labels.items():
                     yield from get_ometa_tag_and_classification(
@@ -477,7 +548,7 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
         try:
             query_resp = self.client.query(
                 BIGQUERY_SCHEMA_DESCRIPTION.format(
-                    project_id=self.client.project,
+                    project_id=self.context.get().database,
                     region=self.service_connection.usageLocation,
                     schema_name=schema_name,
                 )
@@ -505,7 +576,7 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
         3. Adds the Deleted Tables to the context
         """
         self.incremental_table_processor.set_changed_tables_map(
-            project=self.client.project,
+            project=self.context.get().database,
             dataset=schema_name,
             start_date=self.incremental.start_datetime_utc,
         )
@@ -526,6 +597,18 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
             ]
         )
 
+    def get_raw_database_schema_names(self) -> Iterable[str]:
+        if self.service_connection.__dict__.get("databaseSchema"):
+            yield self.service_connection.databaseSchema
+        else:
+            for schema in self.engine.execute(
+                BIGQUERY_GET_SCHEMA_NAMES.format(
+                    project=self.context.get().database,
+                    region=self.service_connection.usageLocation,
+                )
+            ):
+                yield schema[0]
+
     def _get_filtered_schema_names(
         self, return_fqn: bool = False, add_to_status: bool = True
     ) -> Iterable[str]:
@@ -575,7 +658,9 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
             ),
         )
         if self.source_config.includeTags:
-            dataset_obj = self.client.get_dataset(schema_name)
+            dataset_obj = self.client.get_dataset(
+                f"{self.context.get().database}.{schema_name}"
+            )
             if dataset_obj.labels:
                 database_schema_request_obj.tags = []
                 for label_classification, label_tag_name in dataset_obj.labels.items():
@@ -707,7 +792,9 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
 
             if self.source_config.includeDDL:
                 schema_definition = inspector.get_table_ddl(
-                    self.connection, table_name, schema_name
+                    self.connection,
+                    table_name,
+                    f"{self.context.get().database}.{schema_name}",
                 )
                 schema_definition = (
                     str(schema_definition).strip()
@@ -784,7 +871,9 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
         try:
             database = self.context.get().database
             table = self.client.get_table(fqn._build(database, schema_name, table_name))
-            columns = inspector.get_columns(table_name, schema_name, db_name=database)
+            columns = inspector.get_columns(
+                table_name, f"{database}.{schema_name}", db_name=database
+            )
             if (
                 hasattr(table, "external_data_configuration")
                 and hasattr(table.external_data_configuration, "hive_partitioning")

metadata/ingestion/source/database/bigquery/profiler/system.py

@@ -154,7 +154,7 @@ class BigQuerySystemMetricsSource(
                     "rowsAffected": getattr(q, rows_affected_field),
                 }
                 for q in query_results
-                if getattr(q, rows_affected_field) > 0
+                if getattr(q, rows_affected_field, -1) > 0
                 and q.project_id == project_id
                 and q.dataset_id == dataset_id
                 and q.table_name == table

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

@@ -63,7 +63,7 @@ BIGQUERY_SCHEMA_DESCRIPTION = textwrap.dedent(
 
 BIGQUERY_TABLE_AND_TYPE = textwrap.dedent(
     """
-    select table_name, table_type from `{project_id}`.{schema_name}.INFORMATION_SCHEMA.TABLES where table_type != 'VIEW'
+    select table_name, table_type from `{project_id}`.{schema_name}.INFORMATION_SCHEMA.TABLES where table_type NOT IN  ('VIEW', 'MATERIALIZED VIEW')
     """
 )
 
@@ -96,7 +96,7 @@ SELECT
   routine_name as name,
   routine_definition as definition,
   external_language as language
-FROM `{schema_name}`.INFORMATION_SCHEMA.ROUTINES
+FROM `{database_name}`.`{schema_name}`.INFORMATION_SCHEMA.ROUTINES
 WHERE routine_type in ('PROCEDURE', 'TABLE FUNCTION')
   AND routine_catalog = '{database_name}'
   AND routine_schema = '{schema_name}'
@@ -164,7 +164,7 @@ BIGQUERY_LIFE_CYCLE_QUERY = textwrap.dedent(
 select
 table_name as table_name,
 creation_time as created_at
-from `{schema_name}`.INFORMATION_SCHEMA.TABLES
+from `{database_name}`.`{schema_name}`.INFORMATION_SCHEMA.TABLES
 where table_schema = '{schema_name}'
 and table_catalog = '{database_name}'
 """
@@ -182,6 +182,18 @@ AND resource.labels.dataset_id = "{dataset}"
 AND timestamp >= "{start_date}"
 """
 
+BIGQUERY_GET_SCHEMA_NAMES = """
+SELECT schema_name FROM `{project}`.`region-{region}`.INFORMATION_SCHEMA.SCHEMATA;
+"""
+
+BIGQUERY_GET_VIEW_NAMES = """
+SELECT table_name FROM `{project}`.`{dataset}`.INFORMATION_SCHEMA.VIEWS;
+"""
+
+BIGQUERY_GET_MATERIALIZED_VIEW_NAMES = """
+SELECT table_name FROM `{project}`.`{dataset}`.INFORMATION_SCHEMA.MATERIALIZED_VIEWS;
+"""
+
 
 class BigQueryQueryResult(BaseModel):
     project_id: str

metadata/ingestion/source/database/common_db_source.py

@@ -61,10 +61,8 @@ from metadata.ingestion.connections.session import create_and_bind_thread_safe_s
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.models.patch_request import PatchedEntity, PatchRequest
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import (
-    get_connection,
-    kill_active_connections,
-)
+from metadata.ingestion.source.connections import get_connection
+from metadata.ingestion.source.connections_utils import kill_active_connections
 from metadata.ingestion.source.database.database_service import DatabaseServiceSource
 from metadata.ingestion.source.database.sql_column_handler import SqlColumnHandlerMixin
 from metadata.ingestion.source.database.sqlalchemy_source import SqlAlchemySource

metadata/ingestion/source/database/database_service.py

@@ -57,9 +57,6 @@ from metadata.ingestion.api.delete import delete_entity_from_source
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.models.life_cycle import OMetaLifeCycleData
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.models.topology import (
@@ -68,7 +65,7 @@ from metadata.ingestion.models.topology import (
     TopologyContextManager,
     TopologyNode,
 )
-from metadata.ingestion.source.connections import get_test_connection_fn
+from metadata.ingestion.source.connections import test_connection_common
 from metadata.utils import fqn
 from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.filters import filter_by_schema
@@ -604,8 +601,6 @@ class DatabaseServiceSource(
         """
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)