openmetadata-ingestion 1.3.3.0rc1__py3-none-any.whl → 1.4.0.0rc1__py3-none-any.whl

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,7 +96,7 @@ 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
@@ -210,7 +219,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 +236,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 +293,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 +366,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 +415,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 +481,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 +505,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 +565,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 +579,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 +590,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())
@@ -485,7 +607,7 @@ class BigquerySource(
         if table_type == TableType.View:
             try:
                 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)
@@ -498,39 +620,47 @@ class BigquerySource(
 
     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 +741,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 +766,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 +805,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
+        ]

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

@@ -38,7 +38,7 @@ WHERE creation_time BETWEEN "{start_time}" AND "{end_time}"
 )
 
 BIGQUERY_TEST_STATEMENT = textwrap.dedent(
-    """SELECT query FROM `region-{region}`.INFORMATION_SCHEMA.JOBS_BY_PROJECT 
+    """SELECT query FROM `region-{region}`.INFORMATION_SCHEMA.JOBS_BY_PROJECT
     where creation_time > '{creation_date}' limit 1"""
 )
 
@@ -46,8 +46,8 @@ BIGQUERY_TEST_STATEMENT = textwrap.dedent(
 BIGQUERY_SCHEMA_DESCRIPTION = textwrap.dedent(
     """
     SELECT option_value as schema_description FROM
-    `{project_id}`.`region-{region}`.INFORMATION_SCHEMA.SCHEMATA_OPTIONS 
-    where schema_name = '{schema_name}' and option_name = 'description' 
+    `{project_id}`.`region-{region}`.INFORMATION_SCHEMA.SCHEMATA_OPTIONS
+    where schema_name = '{schema_name}' and option_name = 'description'
     and option_value is not null
     """
 )
@@ -62,27 +62,28 @@ BIGQUERY_TABLE_CONSTRAINTS = textwrap.dedent(
     """
     SELECT * 
     FROM `{project_id}`.{schema_name}.INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE 
-    WHERE table_name = '{table_name}' AND constraint_name LIKE '%pk$';
+    WHERE constraint_name LIKE '%pk$';
     """
 )
 
 BIGQUERY_FOREIGN_CONSTRAINTS = textwrap.dedent(
     """
-    SELECT 
-      c.table_name AS referred_table, 
-      r.table_schema as referred_schema, 
-      r.constraint_name as name, 
+    SELECT
+      c.table_name AS referred_table,
+      r.table_schema as referred_schema,
+      r.constraint_name as name,
       c.column_name as referred_columns,
-      c.column_name as constrained_columns 
+      c.column_name as constrained_columns,
+      r.table_name as table_name
     FROM `{project_id}`.{schema_name}.INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE c 
     JOIN `{project_id}`.{schema_name}.INFORMATION_SCHEMA.TABLE_CONSTRAINTS r ON c.constraint_name = r.constraint_name 
-    WHERE r.constraint_type = 'FOREIGN KEY' AND r.table_name='{table_name}';
+    WHERE r.constraint_type = 'FOREIGN KEY';
     """
 )
 
 BIGQUERY_GET_STORED_PROCEDURES = textwrap.dedent(
     """
-SELECT 
+SELECT
   routine_name as name,
   routine_definition as definition,
   external_language as language
@@ -151,7 +152,7 @@ ORDER BY procedure_start_time DESC
 
 BIGQUERY_LIFE_CYCLE_QUERY = textwrap.dedent(
     """
-select 
+select
 table_name as table_name,
 creation_time as created_at
 from `{schema_name}`.INFORMATION_SCHEMA.TABLES
@@ -159,3 +160,15 @@ where table_schema = '{schema_name}'
 and table_catalog = '{database_name}'
 """
 )
+
+BIGQUERY_GET_CHANGED_TABLES_FROM_CLOUD_LOGGING = """
+protoPayload.metadata.@type="type.googleapis.com/google.cloud.audit.BigQueryAuditMetadata"
+AND (
+    protoPayload.methodName = ("google.cloud.bigquery.v2.TableService.UpdateTable" OR "google.cloud.bigquery.v2.TableService.InsertTable" OR "google.cloud.bigquery.v2.TableService.PatchTable" OR "google.cloud.bigquery.v2.TableService.DeleteTable")
+    OR
+    (protoPayload.methodName = "google.cloud.bigquery.v2.JobService.InsertJob" AND (protoPayload.metadata.tableCreation:* OR protoPayload.metadata.tableChange:* OR protoPayload.metadata.tableDeletion:*))
+)
+AND resource.labels.project_id = "{project}"
+AND resource.labels.dataset_id = "{dataset}"
+AND timestamp >= "{start_date}"
+"""

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

@@ -45,7 +45,9 @@ class BigqueryQueryParserSource(QueryParserSource, ABC):
         self.database = self.project_id
 
     @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):

metadata/ingestion/source/database/bigtable/metadata.py

@@ -71,7 +71,9 @@ class BigtableSource(CommonNoSQLSource, MultiDBSource):
         self.tables: Dict[ProjectId, Dict[InstanceId, Dict[TableId, Table]]] = {}
 
     @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: BigTableConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, BigTableConnection):
@@ -96,7 +98,7 @@ class BigtableSource(CommonNoSQLSource, MultiDBSource):
         yield from self.client.project_ids()
 
     def get_schema_name_list(self) -> List[str]:
-        project_id = self.context.database
+        project_id = self.context.get().database
         try:
             # the first element is a list of instances
             # the second element is another collection (seems empty) and I do not know what is its purpose
@@ -113,7 +115,7 @@ class BigtableSource(CommonNoSQLSource, MultiDBSource):
             raise
 
     def get_table_name_list(self, schema_name: str) -> List[str]:
-        project_id = self.context.database
+        project_id = self.context.get().database
         try:
             instance = self._get_instance(project_id, schema_name)
             if instance is None:
@@ -146,7 +148,7 @@ class BigtableSource(CommonNoSQLSource, MultiDBSource):
     def get_table_columns_dict(
         self, schema_name: str, table_name: str
     ) -> Union[List[Dict], Dict]:
-        project_id = self.context.database
+        project_id = self.context.get().database
         try:
             table = self._get_table(project_id, schema_name, table_name)
             if table is None:

metadata/ingestion/source/database/clickhouse/metadata.py

@@ -104,7 +104,9 @@ class ClickhouseSource(CommonDbSourceService):
     """
 
     @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: ClickhouseConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, ClickhouseConnection):

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

@@ -16,7 +16,7 @@ import ast
 import traceback
 from abc import ABC
 from datetime import datetime
-from typing import List
+from typing import List, Optional
 
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.services.connections.database.clickhouseConnection import (
@@ -39,7 +39,9 @@ class ClickhouseQueryParserSource(QueryParserSource, ABC):
     """
 
     @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: ClickhouseConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, ClickhouseConnection):