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

metadata/ingestion/source/database/redshift/metadata.py

@@ -16,7 +16,7 @@ import re
 import traceback
 from typing import Dict, Iterable, List, Optional, Tuple
 
-from sqlalchemy import inspect, sql
+from sqlalchemy import sql
 from sqlalchemy.dialects.postgresql.base import PGDialect
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy_redshift.dialect import RedshiftDialect, RedshiftDialectMixin
@@ -32,7 +32,9 @@ from metadata.generated.schema.entity.data.storedProcedure import (
 )
 from metadata.generated.schema.entity.data.table import (
     ConstraintType,
-    IntervalType,
+    PartitionColumnDetails,
+    PartitionIntervalTypes,
+    Table,
     TableConstraint,
     TablePartition,
     TableType,
@@ -47,6 +49,7 @@ from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
 from metadata.generated.schema.type.basic import EntityName
+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.ometa.ometa_api import OpenMetadata
@@ -54,10 +57,16 @@ 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,
 )
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
+from metadata.ingestion.source.database.redshift.incremental_table_processor import (
+    RedshiftIncrementalTableProcessor,
+)
 from metadata.ingestion.source.database.redshift.models import RedshiftStoredProcedure
 from metadata.ingestion.source.database.redshift.queries import (
     REDSHIFT_GET_ALL_RELATION_INFO,
@@ -80,6 +89,10 @@ from metadata.ingestion.source.database.stored_procedures_mixin import (
     StoredProcedureMixin,
 )
 from metadata.utils import fqn
+from metadata.utils.execution_time_tracker import (
+    calculate_execution_time,
+    calculate_execution_time_generator,
+)
 from metadata.utils.filters import filter_by_database
 from metadata.utils.helpers import get_start_and_end
 from metadata.utils.logger import ingestion_logger
@@ -118,20 +131,41 @@ class RedshiftSource(
     Database metadata from Redshift Source
     """
 
-    def __init__(self, config, metadata):
+    def __init__(
+        self,
+        config: WorkflowSource,
+        metadata,
+        incremental_configuration: IncrementalConfig,
+    ):
         super().__init__(config, metadata)
         self.partition_details = {}
         self.life_cycle_query = REDSHIFT_LIFE_CYCLE_QUERY
+        self.context.get_global().deleted_tables = []
+        self.incremental = incremental_configuration
+        self.incremental_table_processor: Optional[
+            RedshiftIncrementalTableProcessor
+        ] = 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: RedshiftConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, RedshiftConnection):
             raise InvalidSourceException(
                 f"Expected RedshiftConnection, but got {connection}"
             )
-        return cls(config, metadata)
+        incremental_config = IncrementalConfig.create(
+            config.sourceConfig.config.incremental, pipeline_name, metadata
+        )
+        return cls(config, metadata, incremental_config)
 
     def get_partition_details(self) -> None:
         """
@@ -139,7 +173,7 @@ class RedshiftSource(
         """
         try:
             self.partition_details.clear()
-            results = self.engine.execute(REDSHIFT_PARTITION_DETAILS).fetchall()
+            results = self.connection.execute(REDSHIFT_PARTITION_DETAILS).fetchall()
             for row in results:
                 self.partition_details[f"{row.schema}.{row.table}"] = row.diststyle
         except Exception as exe:
@@ -158,6 +192,16 @@ class RedshiftSource(
             {"schema": schema_name},
         )
 
+        if self.incremental.enabled:
+            result = [
+                (name, relkind)
+                for name, relkind in result
+                if name
+                in self.incremental_table_processor.get_not_deleted(
+                    schema_name=schema_name
+                )
+            ]
+
         return [
             TableNameAndType(
                 name=name, type_=STANDARD_TABLE_TYPES.get(relkind, TableType.Regular)
@@ -165,6 +209,35 @@ class RedshiftSource(
             for name, relkind in result
         ]
 
+    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,...
+        """
+
+        result = self.inspector.get_view_names(schema_name) or []
+
+        if self.incremental.enabled:
+            result = [
+                name
+                for name in result
+                if name
+                in self.incremental_table_processor.get_not_deleted(
+                    schema_name=schema_name
+                )
+            ]
+
+        return [
+            TableNameAndType(name=table_name, type_=TableType.View)
+            for table_name in result
+        ]
+
     def get_configured_database(self) -> Optional[str]:
         if not self.service_connection.ingestAllDatabases:
             return self.service_connection.database
@@ -173,25 +246,59 @@ class RedshiftSource(
     def get_database_names_raw(self) -> Iterable[str]:
         yield from self._execute_database_query(REDSHIFT_GET_DATABASE_NAMES)
 
+    def _set_incremental_table_processor(self, database: str):
+        """Prepares the needed data for doing incremental metadata extration for a given database.
+
+        1. Queries Redshift to get the changes done after the `self.incremental.start_datetime_utc`
+        2. Sets the table map with the changes within the RedshiftIncrementalTableProcessor
+        3. Sets the deleted tables in the context
+        """
+        if self.incremental.enabled:
+            self.incremental_table_processor = RedshiftIncrementalTableProcessor.create(
+                self.connection, self.inspector.default_schema_name
+            )
+
+            self.incremental_table_processor.set_table_map(
+                database=database, 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=database,
+                    schema_name=schema_name,
+                    table_name=table_name,
+                )
+                for schema_name, table_name in self.incremental_table_processor.get_deleted()
+            )
+
     def get_database_names(self) -> Iterable[str]:
         if not self.config.serviceConnection.__root__.config.ingestAllDatabases:
-            self.inspector = inspect(self.engine)
             self.get_partition_details()
+
+            self._set_incremental_table_processor(
+                self.config.serviceConnection.__root__.config.database
+            )
+
             yield self.config.serviceConnection.__root__.config.database
         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,
                 )
 
                 if filter_by_database(
                     self.source_config.databaseFilterPattern,
-                    database_fqn
-                    if self.source_config.useFqnForFiltering
-                    else new_database,
+                    (
+                        database_fqn
+                        if self.source_config.useFqnForFiltering
+                        else new_database
+                    ),
                 ):
                     self.status.filter(database_fqn, "Database Filtered Out")
                     continue
@@ -199,6 +306,9 @@ class RedshiftSource(
                 try:
                     self.set_inspector(database_name=new_database)
                     self.get_partition_details()
+
+                    self._set_incremental_table_processor(new_database)
+
                     yield new_database
                 except Exception as exc:
                     logger.debug(traceback.format_exc())
@@ -206,26 +316,34 @@ class RedshiftSource(
                         f"Error trying to connect to database {new_database}: {exc}"
                     )
 
-    def _get_partition_key(self, diststyle: str) -> Optional[List[str]]:
+    def _get_partition_key(self, diststyle: str) -> Optional[str]:
         try:
             regex = re.match(r"KEY\((\w+)\)", diststyle)
             if regex:
-                return [regex.group(1)]
+                return regex.group(1)
         except Exception as err:
             logger.debug(traceback.format_exc())
             logger.warning(err)
         return None
 
+    @calculate_execution_time()
     def get_table_partition_details(
         self, table_name: str, schema_name: str, inspector: Inspector
-    ) -> Tuple[bool, TablePartition]:
+    ) -> Tuple[bool, Optional[TablePartition]]:
         diststyle = self.partition_details.get(f"{schema_name}.{table_name}")
         if diststyle:
-            partition_details = TablePartition(
-                columns=self._get_partition_key(diststyle),
-                intervalType=IntervalType.COLUMN_VALUE,
-            )
-            return True, partition_details
+            distkey = self._get_partition_key(diststyle)
+            if distkey is not None:
+                partition_details = TablePartition(
+                    columns=[
+                        PartitionColumnDetails(
+                            columnName=distkey,
+                            intervalType=PartitionIntervalTypes.COLUMN_VALUE,
+                            interval=None,
+                        )
+                    ]
+                )
+                return True, partition_details
         return False, None
 
     def process_additional_table_constraints(
@@ -254,15 +372,16 @@ class RedshiftSource(
     def get_stored_procedures(self) -> Iterable[RedshiftStoredProcedure]:
         """List Snowflake stored procedures"""
         if self.source_config.includeStoredProcedures:
-            results = self.engine.execute(
+            results = self.connection.execute(
                 REDSHIFT_GET_STORED_PROCEDURES.format(
-                    schema_name=self.context.database_schema,
+                    schema_name=self.context.get().database_schema,
                 )
             ).all()
             for row in results:
                 stored_procedure = RedshiftStoredProcedure.parse_obj(dict(row))
                 yield stored_procedure
 
+    @calculate_execution_time_generator()
     def yield_stored_procedure(
         self, stored_procedure: RedshiftStoredProcedure
     ) -> Iterable[Either[CreateStoredProcedureRequest]]:
@@ -278,9 +397,9 @@ class RedshiftSource(
                 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,
                 ),
             )
             yield Either(right=stored_procedure_request)
@@ -304,7 +423,7 @@ class RedshiftSource(
         start, _ = get_start_and_end(self.source_config.queryLogDuration)
         query = REDSHIFT_GET_STORED_PROCEDURE_QUERIES.format(
             start_date=start,
-            database_name=self.context.database,
+            database_name=self.context.get().database,
         )
 
         queries_dict = self.procedure_queries_dict(
@@ -312,3 +431,26 @@ class RedshiftSource(
         )
 
         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/redshift/models.py

@@ -11,10 +11,14 @@
 """
 Redshift models
 """
-from typing import Optional
+import re
+from typing import Dict, List, Optional, Tuple
 
 from pydantic import BaseModel
 
+TableName = str
+SchemaName = str
+
 
 class RedshiftStoredProcedure(BaseModel):
     """Redshift stored procedure list query results"""
@@ -22,3 +26,72 @@ class RedshiftStoredProcedure(BaseModel):
     name: str
     owner: Optional[str]
     definition: str
+
+
+class RedshiftTableChangeQueryRegex(BaseModel):
+    """Redshift Table Change Query Regex Model. Used for Incremental Extraction."""
+
+    regex: re.Pattern
+    deleted: bool
+
+
+class RedshiftTable(BaseModel):
+    """Redshift Simplified Table Model. Used for Incremental Extraction."""
+
+    name: str
+    deleted: bool
+
+
+class RedshiftTableMap(BaseModel):
+    """Redshift TableMap Model. Used for Incremental Extraction"""
+
+    table_map: Dict[SchemaName, Dict[TableName, RedshiftTable]]
+
+    @classmethod
+    def default(cls) -> "RedshiftTableMap":
+        """Creates a new RedshiftTableMap with an Empty dict."""
+        return cls(table_map={})
+
+    def update(self, schema: SchemaName, table: RedshiftTable):
+        """Updates the Redshift Table Map.
+
+        If the table is already registered in the map, it does nothing.
+        This relies on the processing of the table being ordered by latest changes first.
+        """
+        if schema not in self.table_map:
+            self.table_map[schema] = {table.name: table}
+        else:
+            if table.name not in self.table_map[schema]:
+                self.table_map[schema][table.name] = table
+
+    def get_deleted(
+        self, schema_name: Optional[SchemaName] = None
+    ) -> List[Tuple[SchemaName, TableName]]:
+        """Returns all deleted table names for a given schema."""
+        if schema_name:
+            return [
+                (schema_name, table.name)
+                for table in self.table_map.get(schema_name, {}).values()
+                if table.deleted
+            ]
+
+        deleted_tables = []
+
+        for schema_name in self.table_map:
+            deleted_tables.extend(
+                [
+                    (schema_name, table.name)
+                    for table in self.table_map.get(schema_name, {}).values()
+                    if table.deleted
+                ]
+            )
+
+        return deleted_tables
+
+    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/redshift/queries.py

@@ -248,7 +248,7 @@ REDSHIFT_GET_ALL_RELATIONS = """
         AS "diststyle",
         c.relowner AS "owner_id",
         u.usename AS "owner_name",
-        TRIM(TRAILING ';' FROM 
+        TRIM(TRAILING ';' FROM
         'create view ' || n.nspname || '.' || c.relname || ' as ' ||pg_catalog.pg_get_viewdef(c.oid, true))
         AS "view_definition",
         pg_catalog.array_to_string(c.relacl, '\n') AS "privileges"
@@ -345,7 +345,7 @@ select
 from SP_HISTORY sp
   join Q_HISTORY q
     on sp.procedure_session_id = q.query_session_id
-   and q.query_start_time between sp.procedure_start_time and sp.procedure_end_time 
+   and q.query_start_time between sp.procedure_start_time and sp.procedure_end_time
    and q.query_end_time between sp.procedure_start_time and sp.procedure_end_time
 order by procedure_start_time DESC
     """
@@ -360,3 +360,18 @@ where o.schema = '{schema_name}'
 and o.database = '{database_name}'
 """
 )
+
+REDSHIFT_TABLE_CHANGES_QUERY = """
+SELECT
+    query_text
+FROM SYS_QUERY_HISTORY
+WHERE status = 'success'
+  AND (
+    query_type = 'DDL' OR
+    (query_type = 'UTILITY' AND query_text ilike '%%COMMENT ON%%') OR
+    (query_type = 'CTAS' AND query_text ilike '%%CREATE TABLE%%')
+  )
+  and database_name = '{database}'
+  and end_time >= '{start_date}'
+ORDER BY end_time DESC
+"""

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

@@ -38,7 +38,9 @@ class RedshiftQueryParserSource(QueryParserSource, ABC):
     filters: str
 
     @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: RedshiftConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, RedshiftConnection):

metadata/ingestion/source/database/redshift/utils.py

@@ -33,6 +33,7 @@ from metadata.ingestion.source.database.redshift.queries import (
     REDSHIFT_GET_SCHEMA_COLUMN_INFO,
     REDSHIFT_TABLE_COMMENTS,
 )
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.sqlalchemy_utils import get_table_comment_wrapper
 
 sa_version = Version(sa.__version__)
@@ -45,6 +46,7 @@ ischema_names.update(REDSHIFT_ISCHEMA_NAMES)
 
 
 # pylint: disable=protected-access
+@calculate_execution_time()
 @reflection.cache
 def get_columns(self, connection, table_name, schema=None, **kw):
     """
@@ -80,6 +82,7 @@ def get_columns(self, connection, table_name, schema=None, **kw):
     return columns
 
 
+@calculate_execution_time()
 def _get_column_info(self, *args, **kwargs):
     """
     Get column info
@@ -112,6 +115,7 @@ def _get_column_info(self, *args, **kwargs):
     return column_info
 
 
+@calculate_execution_time()
 @reflection.cache
 def _get_schema_column_info(
     self, connection, schema=None, **kw
@@ -130,13 +134,12 @@ def _get_schema_column_info(
     """
     schema_clause = f"AND schema = '{schema if schema else ''}'"
     all_columns = defaultdict(list)
-    with connection.connect() as cnct:
-        result = cnct.execute(
-            REDSHIFT_GET_SCHEMA_COLUMN_INFO.format(schema_clause=schema_clause)
-        )
-        for col in result:
-            key = RelationKey(col.table_name, col.schema, connection)
-            all_columns[key].append(col)
+    result = connection.execute(
+        REDSHIFT_GET_SCHEMA_COLUMN_INFO.format(schema_clause=schema_clause)
+    )
+    for col in result:
+        key = RelationKey(col.table_name, col.schema, connection)
+        all_columns[key].append(col)
     return dict(all_columns)
 
 
@@ -270,6 +273,7 @@ def _get_charlen(format_type):
     return charlen
 
 
+@calculate_execution_time()
 @reflection.cache
 def _get_pg_column_info(  # pylint: disable=too-many-locals,too-many-arguments, unused-argument
     self,
@@ -353,6 +357,7 @@ def _get_pg_column_info(  # pylint: disable=too-many-locals,too-many-arguments,
     return column_info
 
 
+@calculate_execution_time()
 @reflection.cache
 def get_table_comment(
     self, connection, table_name, schema=None, **kw  # pylint: disable=unused-argument
@@ -366,6 +371,7 @@ def get_table_comment(
     )
 
 
+@calculate_execution_time()
 @reflection.cache
 def _get_all_relation_info(self, connection, **kw):  # pylint: disable=unused-argument
     # pylint: disable=consider-using-f-string

metadata/ingestion/source/database/salesforce/metadata.py

@@ -81,7 +81,9 @@ class SalesforceSource(DatabaseServiceSource):
         self.database_source_state = set()
 
     @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: SalesforceConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, SalesforceConnection):
@@ -112,7 +114,7 @@ class SalesforceSource(DatabaseServiceSource):
         yield Either(
             right=CreateDatabaseRequest(
                 name=database_name,
-                service=self.context.database_service,
+                service=self.context.get().database_service,
             )
         )
 
@@ -135,8 +137,8 @@ class SalesforceSource(DatabaseServiceSource):
                 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,
                 ),
             )
         )
@@ -150,7 +152,7 @@ class SalesforceSource(DatabaseServiceSource):
 
         :return: tables or views, depending on config
         """
-        schema_name = self.context.database_schema
+        schema_name = self.context.get().database_schema
 
         try:
             if self.service_connection.sobjectName:
@@ -165,9 +167,9 @@ class SalesforceSource(DatabaseServiceSource):
                     table_fqn = fqn.build(
                         self.metadata,
                         entity_type=Table,
-                        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,
                         table_name=table_name,
                     )
                     if filter_by_table(
@@ -240,9 +242,9 @@ class SalesforceSource(DatabaseServiceSource):
                 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=self.get_source_url(
                     table_name=table_name,