openmetadata-ingestion 1.3.4.0__py3-none-any.whl → 1.4.0.0rc2__py3-none-any.whl

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

@@ -35,6 +35,9 @@ from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
+FK_CACHE = {}
+PK_CACHE = {}
+
 
 class InspectorWrapper(BaseModel):
     client: Any
@@ -80,14 +83,18 @@ def get_pk_constraint(
     This function overrides to get primary key constraint
     """
     try:
-        table_constraints = connection.engine.execute(
-            BIGQUERY_TABLE_CONSTRAINTS.format(
-                project_id=connection.engine.url.host,
-                schema_name=schema,
-                table_name=table_name,
+        constraints = PK_CACHE.get(f"{connection.engine.url.host}.{schema}")
+        if constraints is None:
+            constraints = connection.engine.execute(
+                BIGQUERY_TABLE_CONSTRAINTS.format(
+                    project_id=connection.engine.url.host,
+                    schema_name=schema,
+                )
             )
-        )
+            PK_CACHE[f"{connection.engine.url.host}.{schema}"] = constraints.fetchall()
+
         col_name = []
+        table_constraints = [row for row in constraints if row.table_name == table_name]
         for table_constraint in table_constraints:
             col_name.append(table_constraint.column_name)
         return {"constrained_columns": tuple(col_name)}
@@ -106,14 +113,18 @@ def get_foreign_keys(
     This function overrides to get foreign key constraint
     """
     try:
-        table_constraints = connection.engine.execute(
-            BIGQUERY_FOREIGN_CONSTRAINTS.format(
-                project_id=connection.engine.url.host,
-                schema_name=schema,
-                table_name=table_name,
+        constraints = FK_CACHE.get(f"{connection.engine.url.host}.{schema}")
+        if constraints is None:
+            constraints = connection.engine.execute(
+                BIGQUERY_FOREIGN_CONSTRAINTS.format(
+                    project_id=connection.engine.url.host,
+                    schema_name=schema,
+                )
             )
-        )
+            FK_CACHE[f"{connection.engine.url.host}.{schema}"] = constraints.fetchall()
+
         col_name = []
+        table_constraints = [row for row in constraints if row.table_name == table_name]
         for table_constraint in table_constraints:
             col_name.append(
                 {

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

@@ -0,0 +1,86 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Bigquery Incremental Table processing logic
+"""
+from datetime import datetime
+from typing import List
+
+import google.cloud.logging
+from google.cloud.logging_v2.entries import LogEntry
+
+from metadata.ingestion.source.database.bigquery.models import (
+    BigQueryTable,
+    BigQueryTableMap,
+    SchemaName,
+    TableName,
+)
+from metadata.ingestion.source.database.bigquery.queries import (
+    BIGQUERY_GET_CHANGED_TABLES_FROM_CLOUD_LOGGING,
+)
+
+
+class BigQueryIncrementalTableProcessor:
+    def __init__(self, client: google.cloud.logging.Client):
+        self._client = client
+        self._changed_tables_map = BigQueryTableMap(table_map={})
+
+    @classmethod
+    def from_project(cls, project: str) -> "BigQueryIncrementalTableProcessor":
+        client = google.cloud.logging.Client(project=project)
+        return cls(client)
+
+    def _is_table_deleted(self, entry: LogEntry) -> bool:
+        if "tableDeletion" in entry.payload.get("metadata").keys():
+            return True
+        return False
+
+    def set_changed_tables_map(
+        self,
+        project: str,
+        dataset: str,
+        start_date: datetime,
+    ):
+        table_map = {}
+
+        resource_names = [f"projects/{project}"]
+        filters = BIGQUERY_GET_CHANGED_TABLES_FROM_CLOUD_LOGGING.format(
+            project=project,
+            dataset=dataset,
+            start_date=start_date.strftime("%Y-%m-%dT%H:%M:%SZ"),
+        )
+
+        entries = self._client.list_entries(
+            resource_names=resource_names,
+            filter_=filters,
+            order_by=google.cloud.logging.DESCENDING,
+        )
+
+        for entry in entries:
+            table_name = entry.payload.get("resourceName", "").split("/")[-1]
+            timestamp = entry.timestamp
+            deleted = self._is_table_deleted(entry)
+
+            if table_name not in table_map:
+                table_map[table_name] = BigQueryTable(
+                    name=table_name, timestamp=timestamp, deleted=deleted
+                )
+        self._changed_tables_map.add(dataset, table_map)
+
+    def get_deleted(self, schema_name: SchemaName) -> List[TableName]:
+        if self._changed_tables_map:
+            return self._changed_tables_map.get_deleted(schema_name)
+        return []
+
+    def get_not_deleted(self, schema_name: SchemaName) -> List[TableName]:
+        if self._changed_tables_map:
+            return self._changed_tables_map.get_not_deleted(schema_name)
+        return []

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

@@ -33,7 +33,9 @@ from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.storedProcedure import StoredProcedureCode
 from metadata.generated.schema.entity.data.table import (
-    IntervalType,
+    PartitionColumnDetails,
+    PartitionIntervalTypes,
+    Table,
     TablePartition,
     TableType,
 )
@@ -51,6 +53,7 @@ from metadata.generated.schema.security.credentials.gcpValues import (
 )
 from metadata.generated.schema.type.basic import EntityName, SourceUrl
 from metadata.generated.schema.type.tagLabel import TagLabel
+from metadata.ingestion.api.delete import delete_entity_by_name
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
@@ -61,6 +64,9 @@ from metadata.ingestion.source.database.bigquery.helper import (
     get_inspector_details,
     get_pk_constraint,
 )
+from metadata.ingestion.source.database.bigquery.incremental_table_processor import (
+    BigQueryIncrementalTableProcessor,
+)
 from metadata.ingestion.source.database.bigquery.models import (
     STORED_PROC_LANGUAGE_MAP,
     BigQueryStoredProcedure,
@@ -77,6 +83,9 @@ from metadata.ingestion.source.database.common_db_source import (
     CommonDbSourceService,
     TableNameAndType,
 )
+from metadata.ingestion.source.database.incremental_metadata_extraction import (
+    IncrementalConfig,
+)
 from metadata.ingestion.source.database.life_cycle_query_mixin import (
     LifeCycleQueryMixin,
 )
@@ -87,10 +96,14 @@ from metadata.ingestion.source.database.stored_procedures_mixin import (
 )
 from metadata.utils import fqn
 from metadata.utils.credentials import GOOGLE_CREDENTIALS
-from metadata.utils.filters import filter_by_database
+from metadata.utils.filters import filter_by_database, filter_by_schema
 from metadata.utils.helpers import get_start_and_end
 from metadata.utils.logger import ingestion_logger
-from metadata.utils.sqlalchemy_utils import is_complex_type
+from metadata.utils.sqlalchemy_utils import (
+    get_all_table_ddls,
+    get_table_ddl,
+    is_complex_type,
+)
 from metadata.utils.tag_utils import get_ometa_tag_and_classification, get_tag_label
 from metadata.utils.tag_utils import get_tag_labels as fetch_tag_labels_om
 
@@ -201,6 +214,9 @@ BigQueryDialect._build_formatted_table_id = (  # pylint: disable=protected-acces
 BigQueryDialect.get_pk_constraint = get_pk_constraint
 BigQueryDialect.get_foreign_keys = get_foreign_keys
 
+Inspector.get_all_table_ddls = get_all_table_ddls
+Inspector.get_table_ddl = get_table_ddl
+
 
 class BigquerySource(
     LifeCycleQueryMixin, StoredProcedureMixin, CommonDbSourceService, MultiDBSource
@@ -210,7 +226,7 @@ class BigquerySource(
     Database metadata from Bigquery Source
     """
 
-    def __init__(self, config, metadata):
+    def __init__(self, config, metadata, incremental_configuration: IncrementalConfig):
         # 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.
@@ -227,15 +243,32 @@ class BigquerySource(
         self.test_connection = self._test_connection
         self.test_connection()
 
+        self.context.get_global().deleted_tables = []
+        self.incremental = incremental_configuration
+        self.incremental_table_processor: Optional[
+            BigQueryIncrementalTableProcessor
+        ] = None
+
+        if self.incremental.enabled:
+            logger.info(
+                "Starting Incremental Metadata Extraction.\n\t Considering Table changes from %s",
+                self.incremental.start_datetime_utc,
+            )
+
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: BigQueryConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, BigQueryConnection):
             raise InvalidSourceException(
                 f"Expected BigQueryConnection, but got {connection}"
             )
-        return cls(config, metadata)
+        incremental_config = IncrementalConfig.create(
+            config.sourceConfig.config.incremental, pipeline_name, metadata
+        )
+        return cls(config, metadata, incremental_config)
 
     @staticmethod
     def set_project_id() -> List[str]:
@@ -267,18 +300,56 @@ class BigquerySource(
         This is useful for sources where we need fine-grained
         logic on how to handle table types, e.g., external, foreign,...
         """
+        table_names_and_types = (
+            self.engine.execute(
+                BIGQUERY_TABLE_AND_TYPE.format(
+                    project_id=self.client.project, schema_name=schema_name
+                )
+            )
+            or []
+        )
+
+        if self.incremental.enabled:
+            table_names_and_types = [
+                (table_name, table_type)
+                for table_name, table_type in table_names_and_types
+                if table_name
+                in self.incremental_table_processor.get_not_deleted(schema_name)
+            ]
 
         return [
             TableNameAndType(
                 name=table_name,
                 type_=_bigquery_table_types.get(table_type, TableType.Regular),
             )
-            for table_name, table_type in self.engine.execute(
-                BIGQUERY_TABLE_AND_TYPE.format(
-                    project_id=self.client.project, schema_name=schema_name
-                )
-            )
-            or []
+            for table_name, table_type in table_names_and_types
+        ]
+
+    def query_view_names_and_types(
+        self, schema_name: str
+    ) -> Iterable[TableNameAndType]:
+        """
+        Connect to the source database to get the view
+        name and type. By default, use the inspector method
+        to get the names and pass the View type.
+
+        This is useful for sources where we need fine-grained
+        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
+                for view_name in view_names
+                if view_name
+                in self.incremental_table_processor.get_not_deleted(schema_name)
+            ]
+
+        return [
+            TableNameAndType(name=view_name, type_=TableType.View)
+            for view_name in view_names
         ]
 
     def yield_tag(
@@ -302,7 +373,7 @@ class BigquerySource(
                         include_tags=self.source_config.includeTags,
                     )
             # Fetching policy tags on the column level
-            list_project_ids = [self.context.database]
+            list_project_ids = [self.context.get().database]
             if not self.service_connection.taxonomyProjectID:
                 self.service_connection.taxonomyProjectID = []
             list_project_ids.extend(self.service_connection.taxonomyProjectID)
@@ -351,6 +422,59 @@ class BigquerySource(
             )
         return ""
 
+    def _prepare_schema_incremental_data(self, schema_name: str):
+        """Prepares the data for Incremental Extraction.
+
+        1. Queries Cloud Logging for the changes
+        2. Sets the table map with the changes within the BigQueryIncrementalTableProcessor
+        3. Adds the Deleted Tables to the context
+        """
+        self.incremental_table_processor.set_changed_tables_map(
+            project=self.client.project,
+            dataset=schema_name,
+            start_date=self.incremental.start_datetime_utc,
+        )
+
+        self.context.get_global().deleted_tables.extend(
+            [
+                fqn.build(
+                    metadata=self.metadata,
+                    entity_type=Table,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
+                    schema_name=schema_name,
+                    table_name=table_name,
+                )
+                for table_name in self.incremental_table_processor.get_deleted(
+                    schema_name
+                )
+            ]
+        )
+
+    def _get_filtered_schema_names(
+        self, return_fqn: bool = False, add_to_status: bool = True
+    ) -> Iterable[str]:
+        for schema_name in self.get_raw_database_schema_names():
+            schema_fqn = fqn.build(
+                self.metadata,
+                entity_type=DatabaseSchema,
+                service_name=self.context.get().database_service,
+                database_name=self.context.get().database,
+                schema_name=schema_name,
+            )
+            if filter_by_schema(
+                self.source_config.schemaFilterPattern,
+                schema_fqn if self.source_config.useFqnForFiltering else schema_name,
+            ):
+                if add_to_status:
+                    self.status.filter(schema_fqn, "Schema Filtered Out")
+                continue
+
+            if self.incremental.enabled:
+                self._prepare_schema_incremental_data(schema_name)
+
+            yield schema_fqn if return_fqn else schema_name
+
     def yield_database_schema(
         self, schema_name: str
     ) -> Iterable[CreateDatabaseSchemaRequest]:
@@ -364,12 +488,12 @@ class BigquerySource(
             database=fqn.build(
                 metadata=self.metadata,
                 entity_type=Database,
-                service_name=self.context.database_service,
-                database_name=self.context.database,
+                service_name=self.context.get().database_service,
+                database_name=self.context.get().database,
             ),
             description=self.get_schema_description(schema_name),
             sourceUrl=self.get_source_url(
-                database_name=self.context.database,
+                database_name=self.context.get().database,
                 schema_name=schema_name,
             ),
         )
@@ -388,8 +512,8 @@ class BigquerySource(
         yield Either(right=database_schema_request_obj)
 
     def get_table_obj(self, table_name: str):
-        schema_name = self.context.database_schema
-        database = self.context.database
+        schema_name = self.context.get().database_schema
+        database = self.context.get().database
         bq_table_fqn = fqn._build(database, schema_name, table_name)
         return self.client.get_table(bq_table_fqn)
 
@@ -448,7 +572,8 @@ class BigquerySource(
             self.temp_credentials_file_path.append(os.environ[GOOGLE_CREDENTIALS])
         self.client = inspector_details.client
         self.engine = inspector_details.engine
-        self.inspector = inspector_details.inspector
+        thread_id = self.context.get_current_thread_id()
+        self._inspector_map[thread_id] = inspector_details.inspector
 
     def get_configured_database(self) -> Optional[str]:
         return None
@@ -461,7 +586,7 @@ class BigquerySource(
             database_fqn = fqn.build(
                 self.metadata,
                 entity_type=Database,
-                service_name=self.context.database_service,
+                service_name=self.context.get().database_service,
                 database_name=project_id,
             )
             if filter_by_database(
@@ -472,6 +597,10 @@ class BigquerySource(
             else:
                 try:
                     self.set_inspector(database_name=project_id)
+                    if self.incremental.enabled:
+                        self.incremental_table_processor = (
+                            BigQueryIncrementalTableProcessor.from_project(project_id)
+                        )
                     yield project_id
                 except Exception as exc:
                     logger.debug(traceback.format_exc())
@@ -479,58 +608,78 @@ class BigquerySource(
                         f"Error trying to connect to database {project_id}: {exc}"
                     )
 
-    def get_view_definition(
+    def get_schema_definition(
         self, table_type: str, table_name: str, schema_name: str, inspector: Inspector
     ) -> Optional[str]:
-        if table_type == TableType.View:
-            try:
+        """
+        Get the DDL statement or View Definition for a table
+        """
+        try:
+            if table_type == TableType.View:
                 view_definition = inspector.get_view_definition(
-                    fqn._build(self.context.database, schema_name, table_name)
+                    fqn._build(self.context.get().database, schema_name, table_name)
                 )
                 view_definition = (
-                    "" if view_definition is None else str(view_definition)
+                    f"CREATE VIEW {schema_name}.{table_name} AS {str(view_definition)}"
+                    if view_definition is not None
+                    else None
                 )
-            except NotImplementedError:
-                logger.warning("View definition not implemented")
-                view_definition = ""
-            return f"CREATE VIEW {schema_name}.{table_name} AS {view_definition}"
+                return view_definition
+
+            schema_definition = inspector.get_table_ddl(
+                self.connection, table_name, schema_name
+            )
+            schema_definition = (
+                str(schema_definition) if schema_definition is not None else None
+            )
+            return schema_definition
+        except NotImplementedError:
+            logger.warning("Schema definition not implemented")
         return None
 
     def get_table_partition_details(
         self, table_name: str, schema_name: str, inspector: Inspector
-    ) -> Tuple[bool, TablePartition]:
+    ) -> Tuple[bool, Optional[TablePartition]]:
         """
         check if the table is partitioned table and return the partition details
         """
-        database = self.context.database
+        database = self.context.get().database
         table = self.client.get_table(fqn._build(database, schema_name, table_name))
         if table.time_partitioning is not None:
             if table.time_partitioning.field:
                 table_partition = TablePartition(
-                    interval=str(table.time_partitioning.type_),
-                    intervalType=IntervalType.TIME_UNIT.value,
+                    columns=[
+                        PartitionColumnDetails(
+                            columnName=table.time_partitioning.field,
+                            interval=str(table.time_partitioning.type_),
+                            intervalType=PartitionIntervalTypes.TIME_UNIT,
+                        )
+                    ]
                 )
-                table_partition.columns = [table.time_partitioning.field]
                 return True, table_partition
-
             return True, TablePartition(
-                interval=str(table.time_partitioning.type_),
-                intervalType=IntervalType.INGESTION_TIME.value,
+                columns=[
+                    PartitionColumnDetails(
+                        columnName="_PARTITIONTIME"
+                        if table.time_partitioning.type_ == "HOUR"
+                        else "_PARTITIONDATE",
+                        interval=str(table.time_partitioning.type_),
+                        intervalType=PartitionIntervalTypes.INGESTION_TIME,
+                    )
+                ]
             )
         if table.range_partitioning:
-            table_partition = TablePartition(
-                intervalType=IntervalType.INTEGER_RANGE.value,
+            table_partition = PartitionColumnDetails(
+                columnName=table.range_partitioning.field,
+                intervalType=PartitionIntervalTypes.INTEGER_RANGE,
+                interval=None,
             )
             if hasattr(table.range_partitioning, "range_") and hasattr(
                 table.range_partitioning.range_, "interval"
             ):
                 table_partition.interval = table.range_partitioning.range_.interval
-            if (
-                hasattr(table.range_partitioning, "field")
-                and table.range_partitioning.field
-            ):
-                table_partition.columns = [table.range_partitioning.field]
-            return True, table_partition
+            table_partition.columnName = table.range_partitioning.field
+            return True, TablePartition(columns=[table_partition])
         return False, None
 
     def clean_raw_data_type(self, raw_data_type):
@@ -611,8 +760,8 @@ class BigquerySource(
         if self.source_config.includeStoredProcedures:
             results = self.engine.execute(
                 BIGQUERY_GET_STORED_PROCEDURES.format(
-                    database_name=self.context.database,
-                    schema_name=self.context.database_schema,
+                    database_name=self.context.get().database,
+                    schema_name=self.context.get().database_schema,
                 )
             ).all()
             for row in results:
@@ -636,14 +785,14 @@ class BigquerySource(
                 databaseSchema=fqn.build(
                     metadata=self.metadata,
                     entity_type=DatabaseSchema,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
-                    schema_name=self.context.database_schema,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
+                    schema_name=self.context.get().database_schema,
                 ),
                 sourceUrl=SourceUrl(
                     __root__=self.get_stored_procedure_url(
-                        database_name=self.context.database,
-                        schema_name=self.context.database_schema,
+                        database_name=self.context.get().database,
+                        schema_name=self.context.get().database_schema,
                         # Follow the same building strategy as tables
                         table_name=stored_procedure.name,
                     )
@@ -675,3 +824,26 @@ class BigquerySource(
         )
 
         return queries_dict
+
+    def mark_tables_as_deleted(self):
+        """
+        Use the current inspector to mark tables as deleted
+        """
+        if self.incremental.enabled:
+            if not self.context.get().__dict__.get("database"):
+                raise ValueError(
+                    "No Database found in the context. We cannot run the table deletion."
+                )
+
+            if self.source_config.markDeletedTables:
+                logger.info(
+                    f"Mark Deleted Tables set to True. Processing database [{self.context.get().database}]"
+                )
+                yield from delete_entity_by_name(
+                    self.metadata,
+                    entity_type=Table,
+                    entity_names=self.context.get_global().deleted_tables,
+                    mark_deleted_entity=self.source_config.markDeletedTables,
+                )
+        else:
+            yield from super().mark_tables_as_deleted()

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

@@ -11,7 +11,8 @@
 """
 BigQuery models
 """
-from typing import Optional
+from datetime import datetime
+from typing import Dict, List, Optional
 
 from pydantic import BaseModel, Field
 
@@ -20,6 +21,9 @@ from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
+TableName = str
+SchemaName = str
+
 STORED_PROC_LANGUAGE_MAP = {
     "SQL": Language.SQL,
     "JAVASCRIPT": Language.JavaScript,
@@ -34,3 +38,33 @@ class BigQueryStoredProcedure(BaseModel):
     language: Optional[str] = Field(
         None, description="Will only be informed for non-SQL routines."
     )
+
+
+class BigQueryTable(BaseModel):
+    name: TableName
+    timestamp: datetime
+    deleted: bool
+
+
+class BigQueryTableMap(BaseModel):
+    table_map: Dict[SchemaName, Dict[TableName, BigQueryTable]]
+
+    def add(self, schema_name: SchemaName, table_map: Dict[TableName, BigQueryTable]):
+        """Adds a new schema table map."""
+        self.table_map[schema_name] = table_map
+
+    def get_deleted(self, schema_name: SchemaName) -> List[TableName]:
+        """Returns all deleted table names for a given schema."""
+        return [
+            table.name
+            for table in self.table_map.get(schema_name, {}).values()
+            if table.deleted
+        ]
+
+    def get_not_deleted(self, schema_name: SchemaName) -> List[TableName]:
+        """Returns all not deleted table names for a given schema."""
+        return [
+            table.name
+            for table in self.table_map.get(schema_name, {}).values()
+            if not table.deleted
+        ]