feat(ingestion/trino): Add sibling support in ingestion

metadata-ingestion/docs/sources/trino/trino_recipe.yml

@@ -13,6 +13,14 @@ source:
     # options:
     #   connect_args:
     #     http_scheme: http
+
+    # Optional -- A mapping of trino catalog to its connector details like connector database, env and platform instance.
+    # This configuration is used to ingest lineage of datasets to connectors. Use catalog name as key.
+    # catalog_to_connector_details:
+    #   catalog_name:
+    #     connector_database: db_name
+    #     platform_instance: cloud_instance
+    #     env: DEV    
 
 sink:
   # sink configs

metadata-ingestion/src/datahub/ingestion/source/sql/trino.py

@@ -1,7 +1,8 @@
+import functools
 import json
 import uuid
 from textwrap import dedent
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, Iterable, List, Optional, Union
 
 import sqlalchemy
 import trino
@@ -16,6 +17,12 @@
 from trino.sqlalchemy import datatype
 from trino.sqlalchemy.dialect import TrinoDialect
 
+from datahub.configuration.source_common import (
+    EnvConfigMixin,
+    PlatformInstanceConfigMixin,
+)
+from datahub.emitter.mce_builder import make_dataset_urn_with_platform_instance
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SourceCapability,
@@ -25,12 +32,23 @@
     platform_name,
     support_status,
 )
+from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.extractor import schema_util
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
+    SqlWorkUnit,
     register_custom_type,
 )
-from datahub.ingestion.source.sql.sql_config import BasicSQLAlchemyConfig
+from datahub.ingestion.source.sql.sql_config import (
+    BasicSQLAlchemyConfig,
+    SQLCommonConfig,
+)
+from datahub.metadata.com.linkedin.pegasus2avro.common import Siblings
+from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
+    DatasetLineageType,
+    Upstream,
+    UpstreamLineage,
+)
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     MapTypeClass,
     NumberTypeClass,
@@ -42,6 +60,19 @@
 register_custom_type(datatype.MAP, MapTypeClass)
 register_custom_type(datatype.DOUBLE, NumberTypeClass)
 
+
+KNOWN_CONNECTOR_PLATFORM_MAPPING = {
+    "clickhouse": "clickhouse",
+    "hive": "hive",
+    "postgresql": "postgres",
+    "mysql": "mysql",
+    "iceberg": "iceberg",
+    "redshift": "redshift",
+    "bigquery": "bigquery",
+}
+
+TWO_TIER_CONNECTORS = ["clickhouse", "hive", "mysql", "iceberg"]
+
 # Type JSON was introduced in trino sqlalchemy dialect in version 0.317.0
 if version.parse(trino.__version__) >= version.parse("0.317.0"):
     register_custom_type(datatype.JSON, RecordTypeClass)
@@ -131,19 +162,47 @@ def _get_columns(self, connection, table_name, schema: str = None, **kw):  # typ
 TrinoDialect._get_columns = _get_columns
 
 
+@functools.lru_cache()
+def get_catalog_connector_name(
+    catalog_name: str, inspector: Inspector
+) -> Optional[str]:
+    if inspector.engine:
+        query = dedent(
+            """
+            SELECT *
+            FROM "system"."metadata"."catalogs"
+        """
+        ).strip()
+        res = inspector.engine.execute(sql.text(query))
+        catalog_connector_dict = {row.catalog_name: row.connector_name for row in res}
+        return catalog_connector_dict.get(catalog_name)
+    return None
+
+
+class ConnectorDetail(PlatformInstanceConfigMixin, EnvConfigMixin):
+    connector_database: Optional[str] = Field(default=None, description="")
+
+
 class TrinoConfig(BasicSQLAlchemyConfig):
     # defaults
     scheme: str = Field(default="trino", description="", hidden_from_docs=True)
 
+    catalog_to_connector_details: Dict[str, ConnectorDetail] = Field(
+        default={},
+        description="A mapping of trino catalog to its connector details like connector database, env and platform instance."
+        "This configuration is used to ingest lineage of datasets to connectors. Use catalog name as key.",
+    )
+
+    ingest_lineage_to_connectors: bool = Field(
+        default=True,
+        description="Whether lineage of datasets to connectors should be ingested",
+    )
+
     def get_identifier(self: BasicSQLAlchemyConfig, schema: str, table: str) -> str:
         identifier = f"{schema}.{table}"
         if self.database:  # TODO: this should be required field
             identifier = f"{self.database}.{identifier}"
-        return (
-            f"{self.platform_instance}.{identifier}"
-            if self.platform_instance
-            else identifier
-        )
+        return identifier
 
 
 @platform_name("Trino", doc_order=1)
@@ -175,6 +234,127 @@ def get_db_name(self, inspector: Inspector) -> str:
         else:
             return super().get_db_name(inspector)
 
+    def _get_source_dataset_urn(
+        self,
+        dataset_name: str,
+        inspector: Inspector,
+        schema: str,
+        table: str,
+    ) -> Optional[str]:
+        catalog_name = dataset_name.split(".")[0]
+        connector_name = get_catalog_connector_name(catalog_name, inspector)
+        if connector_name:
+            connector_platform_name = KNOWN_CONNECTOR_PLATFORM_MAPPING.get(
+                connector_name
+            )
+            if connector_platform_name:
+                connector_details = self.config.catalog_to_connector_details.get(
+                    catalog_name, ConnectorDetail()
+                )
+                if (
+                    connector_platform_name in TWO_TIER_CONNECTORS
+                ):  # connector is two tier
+                    return make_dataset_urn_with_platform_instance(
+                        platform=connector_platform_name,
+                        name=f"{schema}.{table}",
+                        platform_instance=connector_details.platform_instance,
+                        env=connector_details.env,
+                    )
+                elif (
+                    connector_details.connector_database
+                ):  # else connector is three tier
+                    return make_dataset_urn_with_platform_instance(
+                        platform=connector_platform_name,
+                        name=f"{connector_details.connector_database}.{schema}.{table}",
+                        platform_instance=connector_details.platform_instance,
+                        env=connector_details.env,
+                    )
+        return None
+
+    def gen_siblings_workunit(
+        self,
+        dataset_urn: str,
+        source_dataset_urn: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        """
+        Generate sibling workunit for both trino dataset and its connector source dataset
+        """
+        yield MetadataChangeProposalWrapper(
+            entityUrn=dataset_urn,
+            aspect=Siblings(primary=False, siblings=[source_dataset_urn]),
+        ).as_workunit()
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=source_dataset_urn,
+            aspect=Siblings(primary=True, siblings=[dataset_urn]),
+        ).as_workunit()
+
+    def gen_lineage_workunit(
+        self,
+        dataset_urn: str,
+        source_dataset_urn: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        """
+        Generate dataset to source connector lineage workunit
+        """
+        yield MetadataChangeProposalWrapper(
+            entityUrn=dataset_urn,
+            aspect=UpstreamLineage(
+                upstreams=[
+                    Upstream(dataset=source_dataset_urn, type=DatasetLineageType.VIEW)
+                ]
+            ),
+        ).as_workunit()
+
+    def _process_table(
+        self,
+        dataset_name: str,
+        inspector: Inspector,
+        schema: str,
+        table: str,
+        sql_config: SQLCommonConfig,
+    ) -> Iterable[Union[SqlWorkUnit, MetadataWorkUnit]]:
+        yield from super()._process_table(
+            dataset_name, inspector, schema, table, sql_config
+        )
+        if self.config.ingest_lineage_to_connectors:
+            dataset_urn = make_dataset_urn_with_platform_instance(
+                self.platform,
+                dataset_name,
+                self.config.platform_instance,
+                self.config.env,
+            )
+            source_dataset_urn = self._get_source_dataset_urn(
+                dataset_name, inspector, schema, table
+            )
+            if source_dataset_urn:
+                yield from self.gen_siblings_workunit(dataset_urn, source_dataset_urn)
+                yield from self.gen_lineage_workunit(dataset_urn, source_dataset_urn)
+
+    def _process_view(
+        self,
+        dataset_name: str,
+        inspector: Inspector,
+        schema: str,
+        view: str,
+        sql_config: SQLCommonConfig,
+    ) -> Iterable[Union[SqlWorkUnit, MetadataWorkUnit]]:
+        yield from super()._process_view(
+            dataset_name, inspector, schema, view, sql_config
+        )
+        if self.config.ingest_lineage_to_connectors:
+            dataset_urn = make_dataset_urn_with_platform_instance(
+                self.platform,
+                dataset_name,
+                self.config.platform_instance,
+                self.config.env,
+            )
+            source_dataset_urn = self._get_source_dataset_urn(
+                dataset_name, inspector, schema, view
+            )
+            if source_dataset_urn:
+                yield from self.gen_siblings_workunit(dataset_urn, source_dataset_urn)
+
     @classmethod
     def create(cls, config_dict, ctx):
         config = TrinoConfig.parse_obj(config_dict)

metadata-ingestion/tests/integration/trino/test_trino.py

@@ -1,4 +1,3 @@
-import re
 import subprocess
 
 import pytest
@@ -9,7 +8,7 @@
 from datahub.ingestion.run.pipeline import Pipeline
 from datahub.ingestion.sink.file import FileSinkConfig
 from datahub.ingestion.source.ge_profiling_config import GEProfilingConfig
-from datahub.ingestion.source.sql.trino import TrinoConfig
+from datahub.ingestion.source.sql.trino import ConnectorDetail, TrinoConfig
 from tests.test_helpers import fs_helpers, mce_helpers
 from tests.test_helpers.docker_helpers import wait_for_port
 
@@ -53,7 +52,6 @@ def loaded_trino(trino_runner):
 
 
 @freeze_time(FROZEN_TIME)
-@pytest.mark.xfail
 def test_trino_ingest(
     loaded_trino, test_resources_dir, pytestconfig, tmp_path, mock_time
 ):
@@ -73,7 +71,7 @@ def test_trino_ingest(
                     username="foo",
                     schema_pattern=AllowDenyPattern(allow=["^librarydb"]),
                     profile_pattern=AllowDenyPattern(
-                        allow=["library_catalog.librarydb.*"]
+                        allow=["postgresqldb.librarydb.*"]
                     ),
                     profiling=GEProfilingConfig(
                         enabled=True,
@@ -89,6 +87,12 @@ def test_trino_ingest(
                         include_field_histogram=True,
                         include_field_sample_values=True,
                     ),
+                    catalog_to_connector_details={
+                        "postgresqldb": ConnectorDetail(
+                            connector_database="postgres",
+                            platform_instance="local_server",
+                        )
+                    },
                 ).dict(),
             },
             "sink": {
@@ -169,8 +173,6 @@ def test_trino_hive_ingest(
 def test_trino_instance_ingest(
     loaded_trino, test_resources_dir, pytestconfig, tmp_path, mock_time
 ):
-    instance = "production_warehouse"
-    platform = "trino"
     mce_out_file = "trino_instance_mces.json"
     events_file = tmp_path / mce_out_file
     pipeline_config = {
@@ -183,6 +185,11 @@ def test_trino_instance_ingest(
                 username="foo",
                 platform_instance="production_warehouse",
                 schema_pattern=AllowDenyPattern(allow=["^db1"]),
+                catalog_to_connector_details={
+                    "hivedb": ConnectorDetail(
+                        platform_instance="local_server",
+                    )
+                },
             ).dict(),
         },
         "sink": {
@@ -197,40 +204,9 @@ def test_trino_instance_ingest(
     pipeline.pretty_print_summary()
     pipeline.raise_from_status(raise_warnings=True)
 
-    # Assert that all events generated have instance specific urns
-    urn_pattern = "^" + re.escape(
-        f"urn:li:dataset:(urn:li:dataPlatform:{platform},{instance}."
-    )
-    assert (
-        mce_helpers.assert_mce_entity_urn(
-            "ALL",
-            entity_type="dataset",
-            regex_pattern=urn_pattern,
-            file=events_file,
-        )
-        >= 0
-    ), "There should be at least one match"
-
-    assert (
-        mce_helpers.assert_mcp_entity_urn(
-            "ALL",
-            entity_type="dataset",
-            regex_pattern=urn_pattern,
-            file=events_file,
-        )
-        >= 0
-    ), "There should be at least one MCP"
-
-    # all dataset entities emitted must have a dataPlatformInstance aspect emitted
-    # there must be at least one entity emitted
-    assert (
-        mce_helpers.assert_for_each_entity(
-            entity_type="dataset",
-            aspect_name="dataPlatformInstance",
-            aspect_field_matcher={
-                "instance": f"urn:li:dataPlatformInstance:(urn:li:dataPlatform:{platform},{instance})"
-            },
-            file=events_file,
-        )
-        >= 1
+    # Verify the output.
+    mce_helpers.check_golden_file(
+        pytestconfig,
+        output_path=events_file,
+        golden_path=test_resources_dir / "trino_hive_instance_mces_golden.json",
     )