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

metadata/ingestion/source/database/snowflake/metadata.py

@@ -13,6 +13,7 @@ Snowflake source module
 """
 import json
 import traceback
+from datetime import datetime
 from typing import Dict, Iterable, List, Optional, Tuple
 
 import sqlparse
@@ -28,7 +29,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,
 )
@@ -42,6 +45,7 @@ from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
 from metadata.generated.schema.type.basic import EntityName, SourceUrl
+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
@@ -51,6 +55,12 @@ from metadata.ingestion.source.database.common_db_source import (
     CommonDbSourceService,
     TableNameAndType,
 )
+from metadata.ingestion.source.database.external_table_lineage_mixin import (
+    ExternalTableLineageMixin,
+)
+from metadata.ingestion.source.database.incremental_metadata_extraction import (
+    IncrementalConfig,
+)
 from metadata.ingestion.source.database.life_cycle_query_mixin import (
     LifeCycleQueryMixin,
 )
@@ -66,6 +76,7 @@ from metadata.ingestion.source.database.snowflake.queries import (
     SNOWFLAKE_GET_CURRENT_ACCOUNT,
     SNOWFLAKE_GET_DATABASE_COMMENTS,
     SNOWFLAKE_GET_DATABASES,
+    SNOWFLAKE_GET_EXTERNAL_LOCATIONS,
     SNOWFLAKE_GET_ORGANIZATION_NAME,
     SNOWFLAKE_GET_SCHEMA_COMMENTS,
     SNOWFLAKE_GET_STORED_PROCEDURE_QUERIES,
@@ -79,7 +90,9 @@ from metadata.ingestion.source.database.snowflake.utils import (
     get_foreign_keys,
     get_pk_constraint,
     get_schema_columns,
+    get_schema_foreign_keys,
     get_table_comment,
+    get_table_ddl,
     get_table_names,
     get_table_names_reflection,
     get_unique_constraints,
@@ -96,7 +109,7 @@ from metadata.utils import fqn
 from metadata.utils.filters import filter_by_database
 from metadata.utils.helpers import get_start_and_end
 from metadata.utils.logger import ingestion_logger
-from metadata.utils.sqlalchemy_utils import get_all_table_comments
+from metadata.utils.sqlalchemy_utils import get_all_table_comments, get_all_table_ddls
 from metadata.utils.tag_utils import get_ometa_tag_and_classification
 
 ischema_names["VARIANT"] = VARIANT
@@ -125,35 +138,65 @@ SnowflakeDialect._current_database_schema = (  # pylint: disable=protected-acces
 SnowflakeDialect.get_pk_constraint = get_pk_constraint
 SnowflakeDialect.get_foreign_keys = get_foreign_keys
 SnowflakeDialect.get_columns = get_columns
+Inspector.get_all_table_ddls = get_all_table_ddls
+Inspector.get_table_ddl = get_table_ddl
+SnowflakeDialect._get_schema_foreign_keys = get_schema_foreign_keys
 
 
 class SnowflakeSource(
-    LifeCycleQueryMixin, StoredProcedureMixin, CommonDbSourceService, MultiDBSource
+    LifeCycleQueryMixin,
+    StoredProcedureMixin,
+    ExternalTableLineageMixin,
+    CommonDbSourceService,
+    MultiDBSource,
 ):
     """
     Implements the necessary methods to extract
     Database metadata from Snowflake Source
     """
 
-    def __init__(self, config, metadata):
+    def __init__(
+        self,
+        config,
+        metadata,
+        pipeline_name,
+        incremental_configuration: IncrementalConfig,
+    ):
         super().__init__(config, metadata)
         self.partition_details = {}
         self.schema_desc_map = {}
         self.database_desc_map = {}
+        self.external_location_map = {}
 
         self._account: Optional[str] = None
         self._org_name: Optional[str] = None
         self.life_cycle_query = SNOWFLAKE_LIFE_CYCLE_QUERY
+        self.context.get_global().deleted_tables = []
+        self.pipeline_name = pipeline_name
+        self.incremental = incremental_configuration
+
+        if self.incremental.enabled:
+            date = datetime.fromtimestamp(self.incremental.start_timestamp / 1000)
+            logger.info(
+                "Starting Incremental Metadata Extraction.\n\t Considering Table changes from %s",
+                date,
+            )
 
     @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: SnowflakeConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, SnowflakeConnection):
             raise InvalidSourceException(
                 f"Expected SnowflakeConnection, but got {connection}"
             )
-        return cls(config, metadata)
+
+        incremental_config = IncrementalConfig.create(
+            config.sourceConfig.config.incremental, pipeline_name, metadata
+        )
+        return cls(config, metadata, pipeline_name, incremental_config)
 
     @property
     def account(self) -> Optional[str]:
@@ -198,21 +241,33 @@ class SnowflakeSource(
                 ] = row.CLUSTERING_KEY
 
     def set_schema_description_map(self) -> None:
+        self.schema_desc_map.clear()
         results = self.engine.execute(SNOWFLAKE_GET_SCHEMA_COMMENTS).all()
         for row in results:
             self.schema_desc_map[(row.DATABASE_NAME, row.SCHEMA_NAME)] = row.COMMENT
 
     def set_database_description_map(self) -> None:
+        self.database_desc_map.clear()
         if not self.database_desc_map:
             results = self.engine.execute(SNOWFLAKE_GET_DATABASE_COMMENTS).all()
             for row in results:
                 self.database_desc_map[row.DATABASE_NAME] = row.COMMENT
 
+    def set_external_location_map(self, database_name: str) -> None:
+        self.external_location_map.clear()
+        results = self.engine.execute(
+            SNOWFLAKE_GET_EXTERNAL_LOCATIONS.format(database_name=database_name)
+        ).all()
+        self.external_location_map = {
+            (row.database_name, row.schema_name, row.name): row.location
+            for row in results
+        }
+
     def get_schema_description(self, schema_name: str) -> Optional[str]:
         """
         Method to fetch the schema description
         """
-        return self.schema_desc_map.get((self.context.database, schema_name))
+        return self.schema_desc_map.get((self.context.get().database, schema_name))
 
     def get_database_description(self, database_name: str) -> Optional[str]:
         """
@@ -237,13 +292,14 @@ class SnowflakeSource(
             self.set_partition_details()
             self.set_schema_description_map()
             self.set_database_description_map()
+            self.set_external_location_map(configured_db)
             yield configured_db
         else:
             for new_database in self.get_database_names_raw():
                 database_fqn = fqn.build(
                     self.metadata,
                     entity_type=Database,
-                    service_name=self.context.database_service,
+                    service_name=self.context.get().database_service,
                     database_name=new_database,
                 )
 
@@ -262,6 +318,7 @@ class SnowflakeSource(
                     self.set_partition_details()
                     self.set_schema_description_map()
                     self.set_database_description_map()
+                    self.set_external_location_map(new_database)
                     yield new_database
                 except Exception as exc:
                     logger.debug(traceback.format_exc())
@@ -319,15 +376,21 @@ class SnowflakeSource(
 
     def get_table_partition_details(
         self, table_name: str, schema_name: str, inspector: Inspector
-    ) -> Tuple[bool, TablePartition]:
+    ) -> Tuple[bool, Optional[TablePartition]]:
         cluster_key = self.partition_details.get(f"{schema_name}.{table_name}")
         if cluster_key:
             partition_columns = self.parse_column_name_from_expr(cluster_key)
             partition_details = TablePartition(
-                columns=self.__fix_partition_column_case(
-                    table_name, schema_name, inspector, partition_columns
-                ),
-                intervalType=IntervalType.COLUMN_VALUE,
+                columns=[
+                    PartitionColumnDetails(
+                        columnName=column,
+                        intervalType=PartitionIntervalTypes.COLUMN_VALUE,
+                        interval=None,
+                    )
+                    for column in self.__fix_partition_column_case(
+                        table_name, schema_name, inspector, partition_columns
+                    )
+                ]
             )
             return True, partition_details
         return False, None
@@ -340,7 +403,7 @@ class SnowflakeSource(
             try:
                 result = self.connection.execute(
                     SNOWFLAKE_FETCH_ALL_TAGS.format(
-                        database_name=self.context.database,
+                        database_name=self.context.get().database,
                         schema_name=schema_name,
                     )
                 )
@@ -353,8 +416,8 @@ class SnowflakeSource(
                     )
                     result = self.connection.execute(
                         SNOWFLAKE_FETCH_ALL_TAGS.format(
-                            database_name=f'"{self.context.database}"',
-                            schema_name=f'"{self.context.database_schema}"',
+                            database_name=f'"{self.context.get().database}"',
+                            schema_name=f'"{self.context.get().database_schema}"',
                         )
                     )
                 except Exception as inner_exc:
@@ -371,7 +434,7 @@ class SnowflakeSource(
                 fqn_elements = [name for name in row[2:] if name]
                 yield from get_ometa_tag_and_classification(
                     tag_fqn=fqn._build(  # pylint: disable=protected-access
-                        self.context.database_service, *fqn_elements
+                        self.context.get().database_service, *fqn_elements
                     ),
                     tags=[row[1]],
                     classification_name=row[0],
@@ -379,6 +442,41 @@ class SnowflakeSource(
                     classification_description="SNOWFLAKE TAG NAME",
                 )
 
+    def _get_table_names_and_types(
+        self, schema_name: str, table_type: TableType = TableType.Regular
+    ) -> List[TableNameAndType]:
+        table_type_to_params_map = {
+            TableType.Regular: {},
+            TableType.External: {"external_tables": True},
+            TableType.Transient: {"include_transient_tables": True},
+            TableType.Dynamic: {"dynamic_tables": True},
+        }
+
+        snowflake_tables = self.inspector.get_table_names(
+            schema=schema_name,
+            incremental=self.incremental,
+            **table_type_to_params_map[table_type],
+        )
+
+        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 in snowflake_tables.get_deleted()
+            ]
+        )
+
+        return [
+            TableNameAndType(name=table.name, type_=table_type)
+            for table in snowflake_tables.get_not_deleted()
+        ]
+
     def query_table_names_and_types(
         self, schema_name: str
     ) -> Iterable[TableNameAndType]:
@@ -390,31 +488,21 @@ class SnowflakeSource(
         This is useful for sources where we need fine-grained
         logic on how to handle table types, e.g., external, foreign,...
         """
-        table_list = [
-            TableNameAndType(name=table_name)
-            for table_name in self.inspector.get_table_names(
-                schema=schema_name,
-            )
-        ]
+        table_list = self._get_table_names_and_types(schema_name)
 
         table_list.extend(
-            [
-                TableNameAndType(name=table_name, type_=TableType.External)
-                for table_name in self.inspector.get_table_names(
-                    schema=schema_name, external_tables=True
-                )
-            ]
+            self._get_table_names_and_types(schema_name, table_type=TableType.External)
+        )
+
+        table_list.extend(
+            self._get_table_names_and_types(schema_name, table_type=TableType.Dynamic)
         )
 
         if self.service_connection.includeTransientTables:
             table_list.extend(
-                [
-                    TableNameAndType(name=table_name, type_=TableType.Transient)
-                    for table_name in self.inspector.get_table_names(
-                        schema=schema_name,
-                        include_transient_tables=True,
-                    )
-                ]
+                self._get_table_names_and_types(
+                    schema_name, table_type=TableType.Transient
+                )
             )
 
         return table_list
@@ -475,6 +563,38 @@ class SnowflakeSource(
             logger.error(f"Unable to get source url: {exc}")
         return None
 
+    def _get_view_names_and_types(
+        self, schema_name: str, materialized_views: bool = False
+    ) -> List[TableNameAndType]:
+        table_type = (
+            TableType.MaterializedView if materialized_views else TableType.View
+        )
+
+        snowflake_views = self.inspector.get_view_names(
+            schema=schema_name,
+            incremental=self.incremental,
+            materialized_views=materialized_views,
+        )
+
+        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=view.name,
+                )
+                for view in snowflake_views.get_deleted()
+            ]
+        )
+
+        return [
+            TableNameAndType(name=view.name, type_=table_type)
+            for view in snowflake_views.get_not_deleted()
+        ]
+
     def query_view_names_and_types(
         self, schema_name: str
     ) -> Iterable[TableNameAndType]:
@@ -486,29 +606,20 @@ class SnowflakeSource(
         This is useful for sources where we need fine-grained
         logic on how to handle table types, e.g., material views,...
         """
+        views = self._get_view_names_and_types(schema_name)
+        views.extend(
+            self._get_view_names_and_types(schema_name, materialized_views=True)
+        )
 
-        regular_views = [
-            TableNameAndType(name=view_name, type_=TableType.View)
-            for view_name in self.inspector.get_view_names(schema_name) or []
-        ]
-
-        materialized_views = [
-            TableNameAndType(name=view_name, type_=TableType.MaterializedView)
-            for view_name in self.inspector.get_view_names(
-                schema_name, materialized_views=True
-            )
-            or []
-        ]
-
-        return regular_views + materialized_views
+        return views
 
     def get_stored_procedures(self) -> Iterable[SnowflakeStoredProcedure]:
         """List Snowflake stored procedures"""
         if self.source_config.includeStoredProcedures:
             results = self.engine.execute(
                 SNOWFLAKE_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:
@@ -535,8 +646,8 @@ class SnowflakeSource(
         """
         res = self.engine.execute(
             SNOWFLAKE_DESC_STORED_PROCEDURE.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,
                 procedure_name=stored_procedure.name,
                 procedure_signature=stored_procedure.unquote_signature(),
             )
@@ -559,14 +670,14 @@ class SnowflakeSource(
                 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_source_url_root(
-                        database_name=self.context.database,
-                        schema_name=self.context.database_schema,
+                        database_name=self.context.get().database,
+                        schema_name=self.context.get().database_schema,
                     )
                     + f"/procedure/{stored_procedure.name}"
                     + f"{stored_procedure.signature if stored_procedure.signature else ''}"
@@ -599,3 +710,26 @@ class SnowflakeSource(
         )
 
         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/snowflake/models.py

@@ -12,7 +12,8 @@
 Snowflake models
 """
 import urllib
-from typing import Optional
+from datetime import datetime
+from typing import List, Optional
 
 from pydantic import BaseModel, Field, validator
 from requests.utils import quote
@@ -72,3 +73,25 @@ class SnowflakeStoredProcedure(BaseModel):
 
     def unquote_signature(self) -> Optional[str]:
         return urllib.parse.unquote(self.signature) if self.signature else "()"
+
+
+class SnowflakeTable(BaseModel):
+    """Models the items returned from the Table and View Queries used to get the entities to process.
+    :name: Holds the table/view name.
+    :deleted: Holds either a datetime if the table was deleted or None.
+    """
+
+    name: str
+    deleted: Optional[datetime]
+
+
+class SnowflakeTableList(BaseModel):
+    """Understands how to return the deleted and not deleted tables/views from a given list."""
+
+    tables: List[SnowflakeTable]
+
+    def get_deleted(self) -> List[SnowflakeTable]:
+        return [table for table in self.tables if table.deleted]
+
+    def get_not_deleted(self) -> List[SnowflakeTable]:
+        return [table for table in self.tables if not table.deleted]