acryl-datahub 1.0.0rc10__py3-none-any.whl → 1.0.0rc11__py3-none-any.whl

datahub/ingestion/source/redshift/query.py

@@ -31,40 +31,62 @@ class RedshiftCommonQuery:
         AND (datname <> ('template1')::name)
         """
 
-    list_schemas: str = """SELECT distinct n.nspname AS "schema_name",
-        'local' as schema_type,
-        null as schema_owner_name,
-        '' as schema_option,
-        null as external_database
-        FROM pg_catalog.pg_class c
-        LEFT JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
-        JOIN pg_catalog.pg_user u ON u.usesysid = c.relowner
-        WHERE c.relkind IN ('r','v','m','S','f')
-        AND   n.nspname !~ '^pg_'
-        AND   n.nspname != 'information_schema'
-UNION ALL
-SELECT  schemaname as schema_name,
-        CASE s.eskind
-            WHEN '1' THEN 'GLUE'
-            WHEN '2' THEN 'HIVE'
-            WHEN '3' THEN 'POSTGRES'
-            WHEN '4' THEN 'REDSHIFT'
-            ELSE 'OTHER'
-        END as schema_type,
-        -- setting user_name to null as we don't use it now now and it breaks backward compatibility due to additional permission need
-        -- usename as schema_owner_name,
-        null as schema_owner_name,
-        esoptions as schema_option,
-        databasename as external_database
+    # NOTE: although schema owner id is available in tables, we do not use it
+    # as getting username from id requires access to pg_catalog.pg_user_info
+    # which is available only to superusers.
+    # NOTE: Need union here instead of using svv_all_schemas, in order to get
+    # external platform related lineage
+    # NOTE: Using database_name filter for svv_redshift_schemas, as  otherwise
+    # schemas from other shared databases also show up.
+    @staticmethod
+    def list_schemas(database: str) -> str:
+        return f"""
+        SELECT 
+            schema_name,
+            schema_type,
+            schema_option,
+            cast(null as varchar(256)) as external_platform,
+            cast(null as varchar(256)) as external_database
+        FROM svv_redshift_schemas
+        WHERE database_name = '{database}'
+          AND schema_name != 'pg_catalog' and schema_name != 'information_schema'
+    UNION ALL
+        SELECT 
+            schemaname as schema_name,
+            'external' as schema_type,
+            esoptions as schema_option,
+            CASE s.eskind
+                WHEN '1' THEN 'GLUE'
+                WHEN '2' THEN 'HIVE'
+                WHEN '3' THEN 'POSTGRES'
+                WHEN '4' THEN 'REDSHIFT'
+                ELSE 'OTHER'
+            END as external_platform,
+            databasename as external_database
         FROM SVV_EXTERNAL_SCHEMAS as s
-        -- inner join pg_catalog.pg_user_info as i on i.usesysid = s.esowner
         ORDER BY SCHEMA_NAME;
         """
 
+    @staticmethod
+    def get_database_details(database):
+        return f"""\
+            select 
+                database_name,
+                database_type, 
+                database_options 
+            from svv_redshift_databases 
+            where database_name='{database}';"""
+
+    # NOTE: although table owner id is available in tables, we do not use it
+    # as getting username from id requires access to pg_catalog.pg_user_info
+    # which is available only to superusers.
+    # NOTE: Tables from shared database are not available in pg_catalog.pg_class
     @staticmethod
     def list_tables(
-        skip_external_tables: bool = False,
+        skip_external_tables: bool = False, is_shared_database: bool = False
     ) -> str:
+        # NOTE: it looks like description is available only in pg_description
+        # So this remains preferrred way
         tables_query = """
  SELECT  CASE c.relkind
                 WHEN 'r' THEN 'TABLE'
@@ -83,8 +105,6 @@ SELECT  schemaname as schema_name,
                 WHEN 8 THEN 'ALL'
             END AS "diststyle",
             c.relowner AS "owner_id",
-            -- setting user_name to null as we don't use it now now and it breaks backward compatibility due to additional permission need
-            -- u.usename AS "owner_name",
             null as "owner_name",
             TRIM(TRAILING ';' FROM pg_catalog.pg_get_viewdef (c.oid,TRUE)) AS "view_definition",
             pg_catalog.array_to_string(c.relacl,'\n') AS "privileges",
@@ -98,11 +118,11 @@ SELECT  schemaname as schema_name,
         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 pg_catalog.pg_user u ON u.usesysid = c.relowner
         WHERE c.relkind IN ('r','v','m','S','f')
         AND   n.nspname !~ '^pg_'
         AND   n.nspname != 'information_schema'
 """
+
         external_tables_query = """
         SELECT 'EXTERNAL_TABLE' as tabletype,
             NULL AS "schema_oid",
@@ -125,13 +145,62 @@ SELECT  schemaname as schema_name,
         ORDER BY "schema",
                 "relname"
 """
-        if skip_external_tables:
+        shared_database_tables_query = """
+        SELECT table_type as tabletype,
+            NULL AS "schema_oid",
+            schema_name AS "schema",
+            NULL AS "rel_oid",
+            table_name AS "relname",
+            NULL as "creation_time",
+            NULL AS "diststyle",
+            table_owner AS "owner_id",
+            NULL AS "owner_name",
+            NULL AS "view_definition",
+            table_acl AS "privileges",
+            NULL as "location",
+            NULL as parameters,
+            NULL as input_format,
+            NULL As output_format,
+            NULL as serde_parameters,
+            NULL as table_description
+        FROM svv_redshift_tables
+        ORDER BY "schema",
+                "relname"
+"""
+        if is_shared_database:
+            return shared_database_tables_query
+        elif skip_external_tables:
             return tables_query
         else:
             return f"{tables_query} UNION {external_tables_query}"
 
-    # Why is this unused. Is this a bug?
-    list_columns: str = """
+    @staticmethod
+    def list_columns(is_shared_database: bool = False) -> str:
+        if is_shared_database:
+            return """
+            SELECT
+              schema_name as "schema",
+              table_name as "table_name",
+              column_name as "name",
+              encoding as "encode",
+              -- Spectrum represents data types differently.
+              -- Standardize, so we can infer types.
+              data_type AS "type",
+              distkey as "distkey",
+              sortkey as "sortkey",
+              (case when is_nullable = 'no' then TRUE else FALSE end)  as "notnull",
+              null as "comment",
+              null as "adsrc",
+              ordinal_position as "attnum",
+              data_type AS "format_type",
+              column_default as "default",
+              null as "schema_oid",
+              null as "table_oid"
+            FROM SVV_REDSHIFT_COLUMNS
+            WHERE 1 and schema = '{schema_name}'
+            ORDER BY "schema", "table_name", "attnum"
+"""
+        return """
             SELECT
               n.nspname as "schema",
               c.relname as "table_name",
@@ -362,6 +431,29 @@ ORDER BY target_schema, target_table, filename
     ) -> str:
         raise NotImplementedError
 
+    @staticmethod
+    def list_outbound_datashares() -> str:
+        return """SELECT \
+            share_type, \
+            share_name, \
+            trim(producer_namespace) as producer_namespace, \
+            source_database \
+        FROM svv_datashares
+        WHERE share_type='OUTBOUND'\
+        """
+
+    @staticmethod
+    def get_inbound_datashare(database: str) -> str:
+        return f"""SELECT \
+            share_type, \
+            share_name, \
+            trim(producer_namespace) as producer_namespace, \
+            consumer_database \
+        FROM svv_datashares
+        WHERE share_type='INBOUND'
+        AND consumer_database= '{database}'\
+        """
+
 
 class RedshiftProvisionedQuery(RedshiftCommonQuery):
     @staticmethod

datahub/ingestion/source/redshift/redshift.py

@@ -33,7 +33,10 @@ from datahub.ingestion.api.source import (
     TestableSource,
     TestConnectionReport,
 )
-from datahub.ingestion.api.source_helpers import create_dataset_props_patch_builder
+from datahub.ingestion.api.source_helpers import (
+    auto_workunit,
+    create_dataset_props_patch_builder,
+)
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.glossary.classification_mixin import (
     ClassificationHandler,
@@ -45,6 +48,7 @@ from datahub.ingestion.source.common.subtypes import (
     DatasetSubTypes,
 )
 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
@@ -52,6 +56,7 @@ 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 (
     RedshiftColumn,
+    RedshiftDatabase,
     RedshiftDataDictionary,
     RedshiftSchema,
     RedshiftTable,
@@ -150,76 +155,6 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
     - Table, row, and column statistics via optional SQL profiling
     - Table lineage
     - Usage statistics
-
-    ### Prerequisites
-
-    This source needs to access system tables that require extra permissions.
-    To grant these permissions, please alter your datahub Redshift user the following way:
-    ```sql
-    ALTER USER datahub_user WITH SYSLOG ACCESS UNRESTRICTED;
-    GRANT SELECT ON pg_catalog.svv_table_info to datahub_user;
-    GRANT SELECT ON pg_catalog.svl_user_info to datahub_user;
-    ```
-
-    :::note
-
-    Giving a user unrestricted access to system tables gives the user visibility to data generated by other users. For example, STL_QUERY and STL_QUERYTEXT contain the full text of INSERT, UPDATE, and DELETE statements.
-
-    :::
-
-    ### Lineage
-
-    There are multiple lineage collector implementations as Redshift does not support table lineage out of the box.
-
-    #### stl_scan_based
-    The stl_scan based collector uses Redshift's [stl_insert](https://docs.aws.amazon.com/redshift/latest/dg/r_STL_INSERT.html) and [stl_scan](https://docs.aws.amazon.com/redshift/latest/dg/r_STL_SCAN.html) system tables to
-    discover lineage between tables.
-    Pros:
-    - Fast
-    - Reliable
-
-    Cons:
-    - Does not work with Spectrum/external tables because those scans do not show up in stl_scan table.
-    - If a table is depending on a view then the view won't be listed as dependency. Instead the table will be connected with the view's dependencies.
-
-    #### sql_based
-    The sql_based based collector uses Redshift's [stl_insert](https://docs.aws.amazon.com/redshift/latest/dg/r_STL_INSERT.html) to discover all the insert queries
-    and uses sql parsing to discover the dependencies.
-
-    Pros:
-    - Works with Spectrum tables
-    - Views are connected properly if a table depends on it
-
-    Cons:
-    - Slow.
-    - Less reliable as the query parser can fail on certain queries
-
-    #### mixed
-    Using both collector above and first applying the sql based and then the stl_scan based one.
-
-    Pros:
-    - Works with Spectrum tables
-    - Views are connected properly if a table depends on it
-    - A bit more reliable than the sql_based one only
-
-    Cons:
-    - Slow
-    - May be incorrect at times as the query parser can fail on certain queries
-
-    :::note
-
-    The redshift stl redshift tables which are used for getting data lineage retain at most seven days of log history, and sometimes closer to 2-5 days. This means you cannot extract lineage from queries issued outside that window.
-
-    :::
-
-    ### Profiling
-    Profiling runs sql queries on the redshift cluster to get statistics about the tables. To be able to do that, the user needs to have read access to the tables that should be profiled.
-
-    If you don't want to grant read access to the tables you can enable table level profiling which will get table statistics without reading the data.
-    ```yaml
-    profiling:
-      profile_table_level_only: true
-    ```
     """
 
     # TODO: Replace with standardized types in sql_types.py
@@ -330,6 +265,9 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
         self.config: RedshiftConfig = config
         self.report: RedshiftReport = RedshiftReport()
         self.classification_handler = ClassificationHandler(self.config, self.report)
+        self.datashares_helper = RedshiftDatasharesHelper(
+            self.config, self.report, self.ctx.graph
+        )
         self.platform = "redshift"
         self.domain_registry = None
         if self.config.domain:
@@ -361,6 +299,7 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
             is_serverless=self.config.is_serverless
         )
 
+        self.db: Optional[RedshiftDatabase] = None
         self.db_tables: Dict[str, Dict[str, List[RedshiftTable]]] = {}
         self.db_views: Dict[str, Dict[str, List[RedshiftView]]] = {}
         self.db_schemas: Dict[str, Dict[str, RedshiftSchema]] = {}
@@ -424,6 +363,11 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
 
         database = self.config.database
         logger.info(f"Processing db {database}")
+
+        self.db = self.data_dictionary.get_database_details(connection, database)
+        self.report.is_shared_database = (
+            self.db is not None and self.db.is_shared_database
+        )
         with self.report.new_stage(METADATA_EXTRACTION):
             self.db_tables[database] = defaultdict()
             self.db_views[database] = defaultdict()
@@ -563,7 +507,9 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
 
             schema_columns: Dict[str, Dict[str, List[RedshiftColumn]]] = {}
             schema_columns[schema.name] = self.data_dictionary.get_columns_for_schema(
-                conn=connection, schema=schema
+                conn=connection,
+                schema=schema,
+                is_shared_database=self.report.is_shared_database,
             )
 
             if self.config.include_tables:
@@ -887,6 +833,7 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
         tables, views = self.data_dictionary.get_tables_and_views(
             conn=connection,
             skip_external_tables=self.config.skip_external_tables,
+            is_shared_database=self.report.is_shared_database,
         )
         for schema in tables:
             if not is_schema_allowed(
@@ -1029,6 +976,28 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
         database: str,
         lineage_extractor: RedshiftSqlLineageV2,
     ) -> Iterable[MetadataWorkUnit]:
+        if self.config.include_share_lineage:
+            outbound_shares = self.data_dictionary.get_outbound_datashares(connection)
+            yield from auto_workunit(
+                self.datashares_helper.to_platform_resource(list(outbound_shares))
+            )
+
+            if self.db and self.db.is_shared_database:
+                inbound_share = self.db.get_inbound_share()
+                if inbound_share is None:
+                    self.report.warning(
+                        title="Upstream lineage of inbound datashare will be missing",
+                        message="Database options do not contain sufficient information",
+                        context=f"Database: {database}, Options {self.db.options}",
+                    )
+                else:
+                    for known_lineage in self.datashares_helper.generate_lineage(
+                        inbound_share, self.get_all_tables()[database]
+                    ):
+                        lineage_extractor.aggregator.add(known_lineage)
+
+        # TODO: distinguish between definition level lineage and audit log based lineage
+        # definition level lineage should never be skipped
         if not self._should_ingest_lineage():
             return
 

datahub/ingestion/source/redshift/redshift_schema.py

@@ -1,7 +1,8 @@
 import logging
+import re
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
-from typing import Dict, Iterable, List, Optional, Tuple
+from typing import Dict, Iterable, List, Optional, Tuple, Union
 
 import redshift_connector
 
@@ -41,6 +42,10 @@ class RedshiftTable(BaseTable):
     serde_parameters: Optional[str] = None
     last_altered: Optional[datetime] = None
 
+    @property
+    def is_external_table(self) -> bool:
+        return self.type == "EXTERNAL_TABLE"
+
 
 @dataclass
 class RedshiftView(BaseTable):
@@ -51,6 +56,10 @@ class RedshiftView(BaseTable):
     size_in_bytes: Optional[int] = None
     rows_count: Optional[int] = None
 
+    @property
+    def is_external_table(self) -> bool:
+        return self.type == "EXTERNAL_TABLE"
+
 
 @dataclass
 class RedshiftSchema:
@@ -59,8 +68,102 @@ class RedshiftSchema:
     type: str
     owner: Optional[str] = None
     option: Optional[str] = None
+    external_platform: Optional[str] = None
     external_database: Optional[str] = None
 
+    @property
+    def is_external_schema(self) -> bool:
+        return self.type == "external"
+
+
+@dataclass
+class PartialInboundDatashare:
+    share_name: str
+    producer_namespace_prefix: str
+    consumer_database: str
+
+    def get_description(self) -> str:
+        return (
+            f"Namespace Prefix {self.producer_namespace_prefix} Share {self.share_name}"
+        )
+
+
+@dataclass
+class OutboundDatashare:
+    share_name: str
+    producer_namespace: str
+    source_database: str
+
+    def get_key(self) -> str:
+        return f"{self.producer_namespace}.{self.share_name}"
+
+
+@dataclass
+class InboundDatashare:
+    share_name: str
+    producer_namespace: str
+    consumer_database: str
+
+    def get_key(self) -> str:
+        return f"{self.producer_namespace}.{self.share_name}"
+
+    def get_description(self) -> str:
+        return f"Namespace {self.producer_namespace} Share {self.share_name}"
+
+
+@dataclass
+class RedshiftDatabase:
+    name: str
+    type: str
+    options: Optional[str] = None
+
+    @property
+    def is_shared_database(self) -> bool:
+        return self.type == "shared"
+
+    # NOTE: ideally options are in form
+    # {"datashare_name":"xxx","datashare_producer_account":"1234","datashare_producer_namespace":"yyy"}
+    # however due to varchar(128) type of database table that captures options
+    # we may receive only partial information about inbound share
+    def get_inbound_share(
+        self,
+    ) -> Optional[Union[InboundDatashare, PartialInboundDatashare]]:
+        if not self.is_shared_database or not self.options:
+            return None
+
+        # Convert into single regex ??
+        share_name_match = re.search(r'"datashare_name"\s*:\s*"([^"]*)"', self.options)
+        namespace_match = re.search(
+            r'"datashare_producer_namespace"\s*:\s*"([^"]*)"', self.options
+        )
+        partial_namespace_match = re.search(
+            r'"datashare_producer_namespace"\s*:\s*"([^"]*)$', self.options
+        )
+
+        if not share_name_match:
+            # We will always at least get share name
+            return None
+
+        share_name = share_name_match.group(1)
+        if namespace_match:
+            return InboundDatashare(
+                share_name=share_name,
+                producer_namespace=namespace_match.group(1),
+                consumer_database=self.name,
+            )
+        elif partial_namespace_match:
+            return PartialInboundDatashare(
+                share_name=share_name,
+                producer_namespace_prefix=partial_namespace_match.group(1),
+                consumer_database=self.name,
+            )
+        else:
+            return PartialInboundDatashare(
+                share_name=share_name,
+                producer_namespace_prefix="",
+                consumer_database=self.name,
+            )
+
 
 @dataclass
 class RedshiftExtraTableMeta:
@@ -141,13 +244,31 @@ class RedshiftDataDictionary:
 
         return [db[0] for db in dbs]
 
+    @staticmethod
+    def get_database_details(
+        conn: redshift_connector.Connection, database: str
+    ) -> Optional[RedshiftDatabase]:
+        cursor = RedshiftDataDictionary.get_query_result(
+            conn,
+            RedshiftCommonQuery.get_database_details(database),
+        )
+
+        row = cursor.fetchone()
+        if row is None:
+            return None
+        return RedshiftDatabase(
+            name=database,
+            type=row[1],
+            options=row[2],
+        )
+
     @staticmethod
     def get_schemas(
         conn: redshift_connector.Connection, database: str
     ) -> List[RedshiftSchema]:
         cursor = RedshiftDataDictionary.get_query_result(
             conn,
-            RedshiftCommonQuery.list_schemas.format(database_name=database),
+            RedshiftCommonQuery.list_schemas(database),
         )
 
         schemas = cursor.fetchall()
@@ -158,8 +279,8 @@ class RedshiftDataDictionary:
                 database=database,
                 name=schema[field_names.index("schema_name")],
                 type=schema[field_names.index("schema_type")],
-                owner=schema[field_names.index("schema_owner_name")],
                 option=schema[field_names.index("schema_option")],
+                external_platform=schema[field_names.index("external_platform")],
                 external_database=schema[field_names.index("external_database")],
             )
             for schema in schemas
@@ -203,6 +324,7 @@ class RedshiftDataDictionary:
         self,
         conn: redshift_connector.Connection,
         skip_external_tables: bool = False,
+        is_shared_database: bool = False,
     ) -> Tuple[Dict[str, List[RedshiftTable]], Dict[str, List[RedshiftView]]]:
         tables: Dict[str, List[RedshiftTable]] = {}
         views: Dict[str, List[RedshiftView]] = {}
@@ -213,7 +335,10 @@ class RedshiftDataDictionary:
 
         cur = RedshiftDataDictionary.get_query_result(
             conn,
-            RedshiftCommonQuery.list_tables(skip_external_tables=skip_external_tables),
+            RedshiftCommonQuery.list_tables(
+                skip_external_tables=skip_external_tables,
+                is_shared_database=is_shared_database,
+            ),
         )
         field_names = [i[0] for i in cur.description]
         db_tables = cur.fetchall()
@@ -358,11 +483,15 @@ class RedshiftDataDictionary:
 
     @staticmethod
     def get_columns_for_schema(
-        conn: redshift_connector.Connection, schema: RedshiftSchema
+        conn: redshift_connector.Connection,
+        schema: RedshiftSchema,
+        is_shared_database: bool = False,
     ) -> Dict[str, List[RedshiftColumn]]:
         cursor = RedshiftDataDictionary.get_query_result(
             conn,
-            RedshiftCommonQuery.list_columns.format(schema_name=schema.name),
+            RedshiftCommonQuery.list_columns(
+                is_shared_database=is_shared_database
+            ).format(schema_name=schema.name),
         )
 
         table_columns: Dict[str, List[RedshiftColumn]] = {}
@@ -508,3 +637,34 @@ class RedshiftDataDictionary:
                     start_time=row[field_names.index("start_time")],
                 )
             rows = cursor.fetchmany()
+
+    @staticmethod
+    def get_outbound_datashares(
+        conn: redshift_connector.Connection,
+    ) -> Iterable[OutboundDatashare]:
+        cursor = conn.cursor()
+        cursor.execute(RedshiftCommonQuery.list_outbound_datashares())
+        for item in cursor.fetchall():
+            yield OutboundDatashare(
+                share_name=item[1],
+                producer_namespace=item[2],
+                source_database=item[3],
+            )
+
+    # NOTE: this is not used right now as it requires superuser privilege
+    # We can use this in future if the permissions are lowered.
+    @staticmethod
+    def get_inbound_datashare(
+        conn: redshift_connector.Connection,
+        database: str,
+    ) -> Optional[InboundDatashare]:
+        cursor = conn.cursor()
+        cursor.execute(RedshiftCommonQuery.get_inbound_datashare(database))
+        item = cursor.fetchone()
+        if item:
+            return InboundDatashare(
+                share_name=item[1],
+                producer_namespace=item[2],
+                consumer_database=item[3],
+            )
+        return None

datahub/ingestion/source/redshift/report.py

@@ -60,5 +60,8 @@ class RedshiftReport(
     sql_aggregator: Optional[SqlAggregatorReport] = None
     lineage_phases_timer: Dict[str, PerfTimer] = field(default_factory=dict)
 
+    is_shared_database: bool = False
+    outbound_shares_count: Optional[int] = None
+
     def report_dropped(self, key: str) -> None:
         self.filtered.append(key)