acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/redshift/profile.py

@@ -42,9 +42,9 @@ class RedshiftProfiler(GenericProfiler):
                 "max_overflow", self.config.profiling.max_workers
             )
 
-        for db in tables.keys():
+        for db in tables:
             profile_requests = []
-            for schema in tables.get(db, {}).keys():
+            for schema in tables.get(db, {}):
                 if not self.config.schema_pattern.allowed(schema):
                     continue
                 for table in tables[db].get(schema, {}):

datahub/ingestion/source/redshift/query.py

@@ -44,7 +44,7 @@ class RedshiftCommonQuery:
         SELECT 
             schema_name,
             schema_type,
-            schema_option,
+            cast(null as varchar(1024)) as schema_option,
             cast(null as varchar(256)) as external_platform,
             cast(null as varchar(256)) as external_database
         FROM svv_redshift_schemas
@@ -89,7 +89,7 @@ class RedshiftCommonQuery:
     ) -> str:
         # NOTE: it looks like description is available only in pg_description
         # So this remains preferrred way
-        tables_query = """
+        tables_query = f"""
  SELECT  CASE c.relkind
                 WHEN 'r' THEN 'TABLE'
                 WHEN 'v' THEN 'VIEW'
@@ -120,6 +120,7 @@ class RedshiftCommonQuery:
         LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
         LEFT JOIN pg_class_info as ci on c.oid = ci.reloid
         LEFT JOIN pg_catalog.pg_description pgd ON pgd.objsubid = 0 AND pgd.objoid = c.oid
+        JOIN svv_redshift_schemas rs ON rs.schema_name = n.nspname AND rs.database_name = '{database}'
         WHERE c.relkind IN ('r','v','m','S','f')
         AND   n.nspname !~ '^pg_'
         AND   n.nspname != 'information_schema'
@@ -128,23 +129,24 @@ class RedshiftCommonQuery:
         external_tables_query = f"""
         SELECT 'EXTERNAL_TABLE' as tabletype,
             NULL AS "schema_oid",
-            schemaname AS "schema",
+            t.schemaname AS "schema",
             NULL AS "rel_oid",
-            tablename AS "relname",
+            t.tablename AS "relname",
             NULL as "creation_time",
             NULL AS "diststyle",
             NULL AS "owner_id",
             NULL AS "owner_name",
             NULL AS "view_definition",
             NULL AS "privileges",
-            "location",
-            parameters,
-            input_format,
-            output_format,
-            serde_parameters,
+            t."location",
+            t.parameters,
+            t.input_format,
+            t.output_format,
+            t.serde_parameters,
             NULL as table_description
-        FROM pg_catalog.svv_external_tables
-        WHERE redshift_database_name='{database}'
+        FROM pg_catalog.svv_external_tables t
+        JOIN SVV_EXTERNAL_SCHEMAS s ON t.schemaname = s.schemaname
+        WHERE t.redshift_database_name='{database}'
         ORDER BY "schema",
                 "relname"
 """
@@ -232,11 +234,12 @@ class RedshiftCommonQuery:
               ON att.attrelid = c.oid
             LEFT JOIN pg_catalog.pg_attrdef ad
               ON (att.attrelid, att.attnum) = (ad.adrelid, ad.adnum)
+            JOIN svv_redshift_schemas rs ON rs.schema_name = n.nspname AND rs.database_name = '{database_name}'
             WHERE n.nspname !~ '^pg_'
               AND   n.nspname != 'information_schema'
               AND att.attnum > 0
               AND NOT att.attisdropped
-              and schema = '{schema_name}'
+              and n.nspname = '{schema_name}'
             UNION
             SELECT
               view_schema as "schema",
@@ -263,26 +266,27 @@ class RedshiftCommonQuery:
             WHERE 1 and schema = '{schema_name}'
             UNION
             SELECT
-              schemaname as "schema",
-              tablename as "table_name",
-              columnname as "name",
+              c.schemaname as "schema",
+              c.tablename as "table_name",
+              c.columnname as "name",
               null as "encode",
               -- Spectrum represents data types differently.
               -- Standardize, so we can infer types.
-              external_type AS "type",
+              c.external_type AS "type",
               null as "distkey",
               0 as "sortkey",
               null as "notnull",
               null as "comment",
               null as "adsrc",
               null as "attnum",
-              external_type AS "format_type",
+              c.external_type AS "format_type",
               null as "default",
               null as "schema_oid",
               null as "table_oid"
-            FROM SVV_EXTERNAL_COLUMNS
-            WHERE 1 and schema = '{schema_name}'
-            AND redshift_database_name = '{database_name}'
+            FROM SVV_EXTERNAL_COLUMNS c
+            JOIN SVV_EXTERNAL_SCHEMAS s ON c.schemaname = s.schemaname
+            WHERE c.schemaname = '{schema_name}'
+            AND c.redshift_database_name = '{database_name}'
             ORDER BY "schema", "table_name", "attnum"
 """
 

datahub/ingestion/source/redshift/redshift.py

@@ -1,5 +1,4 @@
 import functools
-import itertools
 import logging
 from collections import defaultdict
 from typing import Dict, Iterable, List, Optional, Type, Union
@@ -10,6 +9,7 @@ import humanfriendly
 import pydantic
 import redshift_connector
 
+from datahub.configuration.common import AllowDenyPattern
 from datahub.configuration.pattern_utils import is_schema_allowed
 from datahub.emitter.mce_builder import (
     make_data_platform_urn,
@@ -46,12 +46,12 @@ from datahub.ingestion.source.common.data_reader import DataReader
 from datahub.ingestion.source.common.subtypes import (
     DatasetContainerSubTypes,
     DatasetSubTypes,
+    SourceCapabilityModifier,
 )
 from datahub.ingestion.source.redshift.config import RedshiftConfig
 from datahub.ingestion.source.redshift.datashares import RedshiftDatasharesHelper
 from datahub.ingestion.source.redshift.exception import handle_redshift_exceptions_yield
-from datahub.ingestion.source.redshift.lineage import RedshiftLineageExtractor
-from datahub.ingestion.source.redshift.lineage_v2 import RedshiftSqlLineageV2
+from datahub.ingestion.source.redshift.lineage import RedshiftSqlLineage
 from datahub.ingestion.source.redshift.profile import RedshiftProfiler
 from datahub.ingestion.source.redshift.redshift_data_reader import RedshiftDataReader
 from datahub.ingestion.source.redshift.redshift_schema import (
@@ -70,7 +70,6 @@ from datahub.ingestion.source.sql.sql_utils import (
     add_table_to_schema_container,
     gen_database_container,
     gen_database_key,
-    gen_lineage,
     gen_schema_container,
     gen_schema_key,
     get_dataplatform_instance_aspect,
@@ -90,8 +89,8 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
 from datahub.ingestion.source_report.ingestion_stage import (
     LINEAGE_EXTRACTION,
     METADATA_EXTRACTION,
-    PROFILING,
     USAGE_EXTRACTION_INGESTION,
+    IngestionHighStage,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.common import SubTypes, TimeStamp
 from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
@@ -114,7 +113,6 @@ from datahub.metadata.com.linkedin.pegasus2avro.schema import (
 )
 from datahub.metadata.schema_classes import GlobalTagsClass, TagAssociationClass
 from datahub.utilities import memory_footprint
-from datahub.utilities.dedup_list import deduplicate_list
 from datahub.utilities.mapping import Constants
 from datahub.utilities.perf_timer import PerfTimer
 from datahub.utilities.registries.domain_registry import DomainRegistry
@@ -125,7 +123,14 @@ logger: logging.Logger = logging.getLogger(__name__)
 @platform_name("Redshift")
 @config_class(RedshiftConfig)
 @support_status(SupportStatus.CERTIFIED)
-@capability(SourceCapability.CONTAINERS, "Enabled by default")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.DATABASE,
+        SourceCapabilityModifier.SCHEMA,
+    ],
+)
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
 @capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
@@ -138,14 +143,17 @@ logger: logging.Logger = logging.getLogger(__name__)
 @capability(SourceCapability.SCHEMA_METADATA, "Enabled by default")
 @capability(
     SourceCapability.USAGE_STATS,
-    "Enabled by default, can be disabled via configuration `include_usage_statistics`",
+    "Optionally enabled via `include_usage_statistics`",
+)
+@capability(
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
 )
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
 @capability(
     SourceCapability.CLASSIFICATION,
     "Optionally enabled via `classification.enabled`",
     supported=True,
 )
+@capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
 class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
     """
     This plugin extracts the following:
@@ -354,7 +362,23 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
             ).workunit_processor,
         ]
 
+    def _warn_deprecated_configs(self):
+        if (
+            self.config.match_fully_qualified_names is not None
+            and not self.config.match_fully_qualified_names
+            and self.config.schema_pattern is not None
+            and self.config.schema_pattern != AllowDenyPattern.allow_all()
+        ):
+            self.report.report_warning(
+                message="Please update `schema_pattern` to match against fully qualified schema name `<database_name>.<schema_name>` and set config `match_fully_qualified_names : True`."
+                "Current default `match_fully_qualified_names: False` is only to maintain backward compatibility. "
+                "The config option `match_fully_qualified_names` will be removed in future and the default behavior will be like `match_fully_qualified_names: True`.",
+                context="Config option deprecation warning",
+                title="Config option deprecation warning",
+            )
+
     def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]:
+        self._warn_deprecated_configs()
         connection = self._try_get_redshift_connection(self.config)
 
         if connection is None:
@@ -395,40 +419,25 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
             memory_footprint.total_size(self.db_views)
         )
 
-        if self.config.use_lineage_v2:
-            with RedshiftSqlLineageV2(
-                config=self.config,
-                report=self.report,
-                context=self.ctx,
-                database=database,
-                redundant_run_skip_handler=self.redundant_lineage_run_skip_handler,
-            ) as lineage_extractor:
-                yield from lineage_extractor.aggregator.register_schemas_from_stream(
-                    self.process_schemas(connection, database)
-                )
-
-                with self.report.new_stage(LINEAGE_EXTRACTION):
-                    yield from self.extract_lineage_v2(
-                        connection=connection,
-                        database=database,
-                        lineage_extractor=lineage_extractor,
-                    )
-
-            all_tables = self.get_all_tables()
-        else:
-            yield from self.process_schemas(connection, database)
+        with RedshiftSqlLineage(
+            config=self.config,
+            report=self.report,
+            context=self.ctx,
+            database=database,
+            redundant_run_skip_handler=self.redundant_lineage_run_skip_handler,
+        ) as lineage_extractor:
+            yield from lineage_extractor.aggregator.register_schemas_from_stream(
+                self.process_schemas(connection, database)
+            )
 
-            all_tables = self.get_all_tables()
+            with self.report.new_stage(LINEAGE_EXTRACTION):
+                yield from self.extract_lineage_v2(
+                    connection=connection,
+                    database=database,
+                    lineage_extractor=lineage_extractor,
+                )
 
-            if (
-                self.config.include_table_lineage
-                or self.config.include_view_lineage
-                or self.config.include_copy_lineage
-            ):
-                with self.report.new_stage(LINEAGE_EXTRACTION):
-                    yield from self.extract_lineage(
-                        connection=connection, all_tables=all_tables, database=database
-                    )
+        all_tables = self.get_all_tables()
 
         if self.config.include_usage_statistics:
             with self.report.new_stage(USAGE_EXTRACTION_INGESTION):
@@ -437,7 +446,7 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
                 )
 
         if self.config.is_profiling_enabled():
-            with self.report.new_stage(PROFILING):
+            with self.report.new_high_stage(IngestionHighStage.PROFILING):
                 profiler = RedshiftProfiler(
                     config=self.config,
                     report=self.report,
@@ -940,45 +949,11 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
 
             self.report.usage_extraction_sec[database] = timer.elapsed_seconds(digits=2)
 
-    def extract_lineage(
-        self,
-        connection: redshift_connector.Connection,
-        database: str,
-        all_tables: Dict[str, Dict[str, List[Union[RedshiftView, RedshiftTable]]]],
-    ) -> Iterable[MetadataWorkUnit]:
-        if not self._should_ingest_lineage():
-            return
-
-        lineage_extractor = RedshiftLineageExtractor(
-            config=self.config,
-            report=self.report,
-            context=self.ctx,
-            redundant_run_skip_handler=self.redundant_lineage_run_skip_handler,
-        )
-
-        with PerfTimer() as timer:
-            lineage_extractor.populate_lineage(
-                database=database, connection=connection, all_tables=all_tables
-            )
-
-            self.report.lineage_extraction_sec[f"{database}"] = timer.elapsed_seconds(
-                digits=2
-            )
-            yield from self.generate_lineage(
-                database, lineage_extractor=lineage_extractor
-            )
-
-            if self.redundant_lineage_run_skip_handler:
-                # Update the checkpoint state for this run.
-                self.redundant_lineage_run_skip_handler.update_state(
-                    self.config.start_time, self.config.end_time
-                )
-
     def extract_lineage_v2(
         self,
         connection: redshift_connector.Connection,
         database: str,
-        lineage_extractor: RedshiftSqlLineageV2,
+        lineage_extractor: RedshiftSqlLineage,
     ) -> Iterable[MetadataWorkUnit]:
         if self.config.include_share_lineage:
             outbound_shares = self.data_dictionary.get_outbound_datashares(connection)
@@ -1041,40 +1016,6 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
 
         return True
 
-    def generate_lineage(
-        self, database: str, lineage_extractor: RedshiftLineageExtractor
-    ) -> Iterable[MetadataWorkUnit]:
-        logger.info(f"Generate lineage for {database}")
-        for schema in deduplicate_list(
-            itertools.chain(self.db_tables[database], self.db_views[database])
-        ):
-            if (
-                database not in self.db_schemas
-                or schema not in self.db_schemas[database]
-            ):
-                logger.warning(
-                    f"Either database {database} or {schema} exists in the lineage but was not discovered earlier. Something went wrong."
-                )
-                continue
-
-            table_or_view: Union[RedshiftTable, RedshiftView]
-            for table_or_view in (
-                []
-                + self.db_tables[database].get(schema, [])
-                + self.db_views[database].get(schema, [])
-            ):
-                datahub_dataset_name = f"{database}.{schema}.{table_or_view.name}"
-                dataset_urn = self.gen_dataset_urn(datahub_dataset_name)
-
-                lineage_info = lineage_extractor.get_lineage(
-                    table_or_view,
-                    dataset_urn,
-                    self.db_schemas[database][schema],
-                )
-                if lineage_info:
-                    # incremental lineage generation is taken care by auto_incremental_lineage
-                    yield from gen_lineage(dataset_urn, lineage_info)
-
     def add_config_to_report(self):
         self.report.stateful_lineage_ingestion_enabled = (
             self.config.enable_stateful_lineage_ingestion

datahub/ingestion/source/redshift/redshift_schema.py

@@ -15,6 +15,7 @@ from datahub.ingestion.source.sql.sql_generic import BaseColumn, BaseTable
 from datahub.metadata.com.linkedin.pegasus2avro.schema import SchemaField
 from datahub.sql_parsing.sqlglot_lineage import SqlParsingResult
 from datahub.utilities.hive_schema_to_avro import get_schema_fields_for_hive_column
+from datahub.utilities.perf_timer import PerfTimer
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -243,9 +244,13 @@ class RedshiftDataDictionary:
         conn: redshift_connector.Connection, query: str
     ) -> redshift_connector.Cursor:
         cursor: redshift_connector.Cursor = conn.cursor()
-
-        logger.debug(f"Query : {query}")
-        cursor.execute(query)
+        with PerfTimer() as timer:
+            query_hash_id = hash(query)
+            logger.info(f"Executing query [{query_hash_id}]\n{query}")
+            cursor.execute(query)
+            logger.info(
+                f"Time taken query [{query_hash_id}: {timer.elapsed_seconds():.3f} seconds"
+            )
         return cursor
 
     @staticmethod
@@ -545,8 +550,7 @@ class RedshiftDataDictionary:
         conn: redshift_connector.Connection,
         query: str,
     ) -> Iterable[LineageRow]:
-        cursor = conn.cursor()
-        cursor.execute(query)
+        cursor = RedshiftDataDictionary.get_query_result(conn=conn, query=query)
         field_names = [i[0] for i in cursor.description]
 
         rows = cursor.fetchmany()
@@ -603,9 +607,7 @@ class RedshiftDataDictionary:
         conn: redshift_connector.Connection,
         query: str,
     ) -> Iterable[TempTableRow]:
-        cursor = conn.cursor()
-
-        cursor.execute(query)
+        cursor = RedshiftDataDictionary.get_query_result(conn=conn, query=query)
 
         field_names = [i[0] for i in cursor.description]
 
@@ -662,8 +664,9 @@ class RedshiftDataDictionary:
     def get_outbound_datashares(
         conn: redshift_connector.Connection,
     ) -> Iterable[OutboundDatashare]:
-        cursor = conn.cursor()
-        cursor.execute(RedshiftCommonQuery.list_outbound_datashares())
+        cursor = RedshiftDataDictionary.get_query_result(
+            conn=conn, query=RedshiftCommonQuery.list_outbound_datashares()
+        )
         for item in cursor.fetchall():
             yield OutboundDatashare(
                 share_name=item[1],
@@ -678,8 +681,10 @@ class RedshiftDataDictionary:
         conn: redshift_connector.Connection,
         database: str,
     ) -> Optional[InboundDatashare]:
-        cursor = conn.cursor()
-        cursor.execute(RedshiftCommonQuery.get_inbound_datashare(database))
+        cursor = RedshiftDataDictionary.get_query_result(
+            conn=conn,
+            query=RedshiftCommonQuery.get_inbound_datashare(database),
+        )
         item = cursor.fetchone()
         if item:
             return InboundDatashare(

datahub/ingestion/source/redshift/report.py

@@ -4,7 +4,6 @@ from typing import Dict, Optional
 
 from datahub.ingestion.glossary.classification_mixin import ClassificationReportMixin
 from datahub.ingestion.source.sql.sql_report import SQLSourceReport
-from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
 from datahub.ingestion.source_report.time_window import BaseTimeWindowReport
 from datahub.sql_parsing.sql_parsing_aggregator import SqlAggregatorReport
 from datahub.utilities.lossy_collections import LossyDict
@@ -15,7 +14,6 @@ from datahub.utilities.stats_collections import TopKDict
 @dataclass
 class RedshiftReport(
     SQLSourceReport,
-    IngestionStageReport,
     BaseTimeWindowReport,
     ClassificationReportMixin,
 ):

datahub/ingestion/source/redshift/usage.py

@@ -25,6 +25,7 @@ from datahub.ingestion.source.redshift.query import (
     RedshiftServerlessQuery,
 )
 from datahub.ingestion.source.redshift.redshift_schema import (
+    RedshiftDataDictionary,
     RedshiftTable,
     RedshiftView,
 )
@@ -182,15 +183,17 @@ class RedshiftUsageExtractor:
         self.report.num_operational_stats_filtered = 0
 
         if self.config.include_operational_stats:
-            with self.report.new_stage(USAGE_EXTRACTION_OPERATIONAL_STATS):
-                with PerfTimer() as timer:
-                    # Generate operation aspect workunits
-                    yield from self._gen_operation_aspect_workunits(
-                        self.connection, all_tables
-                    )
-                    self.report.operational_metadata_extraction_sec[
-                        self.config.database
-                    ] = timer.elapsed_seconds(digits=2)
+            with (
+                self.report.new_stage(USAGE_EXTRACTION_OPERATIONAL_STATS),
+                PerfTimer() as timer,
+            ):
+                # Generate operation aspect workunits
+                yield from self._gen_operation_aspect_workunits(
+                    self.connection, all_tables
+                )
+                self.report.operational_metadata_extraction_sec[
+                    self.config.database
+                ] = timer.elapsed_seconds(digits=2)
 
         # Generate aggregate events
         with self.report.new_stage(USAGE_EXTRACTION_USAGE_AGGREGATION):
@@ -261,8 +264,7 @@ class RedshiftUsageExtractor:
         connection: redshift_connector.Connection,
         all_tables: Dict[str, Dict[str, List[Union[RedshiftView, RedshiftTable]]]],
     ) -> Iterable[RedshiftAccessEvent]:
-        cursor = connection.cursor()
-        cursor.execute(query)
+        cursor = RedshiftDataDictionary.get_query_result(conn=connection, query=query)
         results = cursor.fetchmany()
         field_names = [i[0] for i in cursor.description]
         while results:

datahub/ingestion/source/s3/report.py

@@ -1,19 +1,21 @@
 import dataclasses
 from dataclasses import field as dataclass_field
-from typing import List
 
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalSourceReport,
 )
+from datahub.utilities.lossy_collections import LossyList
 
 
 @dataclasses.dataclass
 class DataLakeSourceReport(StaleEntityRemovalSourceReport):
     files_scanned = 0
-    filtered: List[str] = dataclass_field(default_factory=list)
+    filtered: LossyList[str] = dataclass_field(default_factory=LossyList)
+    number_of_files_filtered: int = 0
 
     def report_file_scanned(self) -> None:
         self.files_scanned += 1
 
     def report_file_dropped(self, file: str) -> None:
         self.filtered.append(file)
+        self.number_of_files_filtered += 1