acryl-datahub 0.15.0.6rc2__py3-none-any.whl → 1.0.0__py3-none-any.whl

datahub/ingestion/source/redshift/datashares.py

@@ -0,0 +1,236 @@
+from typing import Dict, Iterable, List, Optional, Union
+
+from pydantic import BaseModel
+
+from datahub.api.entities.platformresource.platform_resource import (
+    ElasticPlatformResourceQuery,
+    PlatformResource,
+    PlatformResourceKey,
+    PlatformResourceSearchFields,
+)
+from datahub.emitter.mce_builder import make_dataset_urn_with_platform_instance
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.ingestion.source.redshift.config import RedshiftConfig
+from datahub.ingestion.source.redshift.redshift_schema import (
+    InboundDatashare,
+    OutboundDatashare,
+    PartialInboundDatashare,
+    RedshiftTable,
+    RedshiftView,
+)
+from datahub.ingestion.source.redshift.report import RedshiftReport
+from datahub.sql_parsing.sql_parsing_aggregator import KnownLineageMapping
+from datahub.utilities.search_utils import LogicalOperator
+
+
+class OutboundSharePlatformResource(BaseModel):
+    namespace: str
+    platform_instance: Optional[str]
+    env: str
+    source_database: str
+    share_name: str
+
+    def get_key(self) -> str:
+        return f"{self.namespace}.{self.share_name}"
+
+
+PLATFORM_RESOURCE_TYPE = "OUTBOUND_DATASHARE"
+
+
+class RedshiftDatasharesHelper:
+    """
+    Redshift datashares lineage generation relies on PlatformResource entity
+    to identify the producer namespace and its platform_instance and env
+
+    Ingestion of any database in namespace will
+    A. generate PlatformResource entity for all outbound shares in namespace.
+    B. generate lineage with upstream tables from another namespace, if the database
+        is created from an inbound share
+
+    """
+
+    def __init__(
+        self,
+        config: RedshiftConfig,
+        report: RedshiftReport,
+        graph: Optional[DataHubGraph],
+    ):
+        self.platform = "redshift"
+        self.config = config
+        self.report = report
+        self.graph = graph
+
+    def to_platform_resource(
+        self, shares: List[OutboundDatashare]
+    ) -> Iterable[MetadataChangeProposalWrapper]:
+        if not shares:
+            self.report.outbound_shares_count = 0
+            return
+
+        self.report.outbound_shares_count = len(shares)
+        # Producer namespace will be current namespace for all
+        # outbound data shares
+
+        for share in shares:
+            producer_namespace = share.producer_namespace
+            try:
+                platform_resource_key = PlatformResourceKey(
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    resource_type=PLATFORM_RESOURCE_TYPE,
+                    primary_key=share.get_key(),
+                )
+
+                value = OutboundSharePlatformResource(
+                    namespace=producer_namespace,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                    source_database=share.source_database,
+                    share_name=share.share_name,
+                )
+
+                platform_resource = PlatformResource.create(
+                    key=platform_resource_key,
+                    value=value,
+                    secondary_keys=[share.share_name, share.producer_namespace],
+                )
+
+                yield from platform_resource.to_mcps()
+
+            except Exception as exc:
+                self.report.warning(
+                    title="Downstream lineage to outbound datashare may not work",
+                    message="Failed to generate platform resource for outbound datashares",
+                    context=f"Namespace {share.producer_namespace} Share {share.share_name}",
+                    exc=exc,
+                )
+
+    def generate_lineage(
+        self,
+        share: Union[InboundDatashare, PartialInboundDatashare],
+        tables: Dict[str, List[Union[RedshiftTable, RedshiftView]]],
+    ) -> Iterable[KnownLineageMapping]:
+        upstream_share = self.find_upstream_share(share)
+
+        if not upstream_share:
+            return
+
+        for schema in tables:
+            for table in tables[schema]:
+                dataset_urn = self.gen_dataset_urn(
+                    f"{share.consumer_database}.{schema}.{table.name}",
+                    self.config.platform_instance,
+                    self.config.env,
+                )
+
+                upstream_dataset_urn = self.gen_dataset_urn(
+                    f"{upstream_share.source_database}.{schema}.{table.name}",
+                    upstream_share.platform_instance,
+                    upstream_share.env,
+                )
+
+                yield KnownLineageMapping(
+                    upstream_urn=upstream_dataset_urn, downstream_urn=dataset_urn
+                )
+
+    def find_upstream_share(
+        self, share: Union[InboundDatashare, PartialInboundDatashare]
+    ) -> Optional[OutboundSharePlatformResource]:
+        if not self.graph:
+            self.report.warning(
+                title="Upstream lineage of inbound datashare will be missing",
+                message="Missing datahub graph. Either use the datahub-rest sink or "
+                "set the top-level datahub_api config in the recipe",
+            )
+        else:
+            resources = self.get_platform_resources(self.graph, share)
+
+            if len(resources) == 0 or (
+                not any(
+                    [
+                        resource.resource_info is not None
+                        and resource.resource_info.resource_type
+                        == PLATFORM_RESOURCE_TYPE
+                        for resource in resources
+                    ]
+                )
+            ):
+                self.report.info(
+                    title="Upstream lineage of inbound datashare will be missing",
+                    message="Missing platform resource for share. "
+                    "Setup redshift ingestion for namespace if not already done. If ingestion is setup, "
+                    "check whether ingestion user has ALTER/SHARE permission to share.",
+                    context=share.get_description(),
+                )
+            else:
+                # Ideally we should get only one resource as primary key is namespace+share
+                # and type is "OUTBOUND_DATASHARE"
+                for resource in resources:
+                    try:
+                        assert (
+                            resource.resource_info is not None
+                            and resource.resource_info.value is not None
+                        )
+                        return resource.resource_info.value.as_pydantic_object(
+                            OutboundSharePlatformResource, True
+                        )
+                    except Exception as e:
+                        self.report.warning(
+                            title="Upstream lineage of inbound datashare will be missing",
+                            message="Failed to parse platform resource for outbound datashare",
+                            context=share.get_description(),
+                            exc=e,
+                        )
+
+        return None
+
+    def get_platform_resources(
+        self,
+        graph: DataHubGraph,
+        share: Union[InboundDatashare, PartialInboundDatashare],
+    ) -> List[PlatformResource]:
+        # NOTE: ideally we receive InboundDatashare and not PartialInboundDatashare.
+        # however due to varchar(128) type of database table that captures datashare options
+        # we may receive only partial information about inbound share
+        # Alternate option to get InboundDatashare using svv_datashares requires superuser
+        if isinstance(share, PartialInboundDatashare):
+            return list(
+                PlatformResource.search_by_filters(
+                    graph,
+                    ElasticPlatformResourceQuery.create_from()
+                    .group(LogicalOperator.AND)
+                    .add_field_match(
+                        PlatformResourceSearchFields.RESOURCE_TYPE,
+                        PLATFORM_RESOURCE_TYPE,
+                    )
+                    .add_field_match(
+                        PlatformResourceSearchFields.PLATFORM, self.platform
+                    )
+                    .add_field_match(
+                        PlatformResourceSearchFields.SECONDARY_KEYS,
+                        share.share_name,
+                    )
+                    .add_wildcard(
+                        PlatformResourceSearchFields.SECONDARY_KEYS.field_name,
+                        f"{share.producer_namespace_prefix}*",
+                    )
+                    .end(),
+                )
+            )
+        return list(
+            PlatformResource.search_by_key(
+                graph, key=share.get_key(), primary=True, is_exact=True
+            )
+        )
+
+    # TODO: Refactor and move to new RedshiftIdentifierBuilder class
+    def gen_dataset_urn(
+        self, datahub_dataset_name: str, platform_instance: Optional[str], env: str
+    ) -> str:
+        return make_dataset_urn_with_platform_instance(
+            platform=self.platform,
+            name=datahub_dataset_name,
+            platform_instance=platform_instance,
+            env=env,
+        )

datahub/ingestion/source/redshift/lineage.py

@@ -813,9 +813,13 @@ class RedshiftLineageExtractor:
             )
 
         tablename = table.name
-        if table.type == "EXTERNAL_TABLE":
+        if (
+            table.is_external_table()
+            and schema.is_external_schema()
+            and schema.external_platform
+        ):
             # external_db_params = schema.option
-            upstream_platform = schema.type.lower()
+            upstream_platform = schema.external_platform.lower()
             catalog_upstream = UpstreamClass(
                 mce_builder.make_dataset_urn_with_platform_instance(
                     upstream_platform,

datahub/ingestion/source/redshift/lineage_v2.py

@@ -401,11 +401,14 @@ class RedshiftSqlLineageV2(Closeable):
     ) -> None:
         for schema_name, tables in all_tables[self.database].items():
             for table in tables:
-                if table.type == "EXTERNAL_TABLE":
-                    schema = db_schemas[self.database][schema_name]
-
+                schema = db_schemas[self.database][schema_name]
+                if (
+                    table.is_external_table()
+                    and schema.is_external_schema()
+                    and schema.external_platform
+                ):
                     # external_db_params = schema.option
-                    upstream_platform = schema.type.lower()
+                    upstream_platform = schema.external_platform.lower()
 
                     table_urn = mce_builder.make_dataset_urn_with_platform_instance(
                         self.platform,
@@ -413,14 +416,26 @@ class RedshiftSqlLineageV2(Closeable):
                         platform_instance=self.config.platform_instance,
                         env=self.config.env,
                     )
-                    upstream_urn = mce_builder.make_dataset_urn_with_platform_instance(
-                        upstream_platform,
-                        f"{schema.external_database}.{table.name}",
-                        platform_instance=(
+                    if upstream_platform == self.platform:
+                        upstream_schema = schema.get_upstream_schema_name() or "public"
+                        upstream_dataset_name = (
+                            f"{schema.external_database}.{upstream_schema}.{table.name}"
+                        )
+                        upstream_platform_instance = self.config.platform_instance
+                    else:
+                        upstream_dataset_name = (
+                            f"{schema.external_database}.{table.name}"
+                        )
+                        upstream_platform_instance = (
                             self.config.platform_instance_map.get(upstream_platform)
                             if self.config.platform_instance_map
                             else None
-                        ),
+                        )
+
+                    upstream_urn = mce_builder.make_dataset_urn_with_platform_instance(
+                        upstream_platform,
+                        upstream_dataset_name,
+                        platform_instance=upstream_platform_instance,
                         env=self.config.env,
                     )
 

datahub/ingestion/source/redshift/profile.py

@@ -48,7 +48,7 @@ class RedshiftProfiler(GenericProfiler):
                 if not self.config.schema_pattern.allowed(schema):
                     continue
                 for table in tables[db].get(schema, {}):
-                    if table.type == "EXTERNAL_TABLE":
+                    if table.is_external_table() or self.report.is_shared_database:
                         if not self.config.profiling.profile_external_tables:
                             # Case 1: If user did not tell us to profile external tables, simply log this.
                             self.report.profiling_skipped_other[schema] += 1

datahub/ingestion/source/redshift/query.py

@@ -31,40 +31,64 @@ 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(
+        database: str,
         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 +107,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,12 +120,12 @@ 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 = """
+
+        external_tables_query = f"""
         SELECT 'EXTERNAL_TABLE' as tabletype,
             NULL AS "schema_oid",
             schemaname AS "schema",
@@ -122,16 +144,70 @@ SELECT  schemaname as schema_name,
             serde_parameters,
             NULL as table_description
         FROM pg_catalog.svv_external_tables
+        WHERE redshift_database_name='{database}'
+        ORDER BY "schema",
+                "relname"
+"""
+        shared_database_tables_query = f"""
+        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
+        WHERE database_name='{database}'
         ORDER BY "schema",
                 "relname"
 """
-        if skip_external_tables:
+        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(
+        database_name: str, schema_name: str, is_shared_database: bool = False
+    ) -> str:
+        if is_shared_database:
+            return f"""
+            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}'
+            AND database_name = '{database_name}'
+            ORDER BY "schema", "table_name", "attnum"
+"""
+        return f"""
             SELECT
               n.nspname as "schema",
               c.relname as "table_name",
@@ -206,6 +282,7 @@ SELECT  schemaname as schema_name,
               null as "table_oid"
             FROM SVV_EXTERNAL_COLUMNS
             WHERE 1 and schema = '{schema_name}'
+            AND redshift_database_name = '{database_name}'
             ORDER BY "schema", "table_name", "attnum"
 """
 
@@ -362,6 +439,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