acryl-datahub 0.15.0.5rc10__py3-none-any.whl → 0.15.0.6rc2__py3-none-any.whl

datahub/ingestion/source/snowflake/snowflake_report.py

@@ -104,6 +104,7 @@ class SnowflakeV2Report(
     schemas_scanned: int = 0
     databases_scanned: int = 0
     tags_scanned: int = 0
+    streams_scanned: int = 0
 
     include_usage_stats: bool = False
     include_operational_stats: bool = False
@@ -113,6 +114,7 @@ class SnowflakeV2Report(
     table_lineage_query_secs: float = -1
     external_lineage_queries_secs: float = -1
     num_tables_with_known_upstreams: int = 0
+    num_streams_with_known_upstreams: int = 0
     num_upstream_lineage_edge_parsing_failed: int = 0
     num_secure_views_missing_definition: int = 0
     num_structured_property_templates_created: int = 0
@@ -131,6 +133,8 @@ class SnowflakeV2Report(
     num_get_tags_for_object_queries: int = 0
     num_get_tags_on_columns_for_table_queries: int = 0
 
+    num_get_streams_for_schema_queries: int = 0
+
     rows_zero_objects_modified: int = 0
 
     _processed_tags: MutableSet[str] = field(default_factory=set)
@@ -157,6 +161,8 @@ class SnowflakeV2Report(
                 return
             self._scanned_tags.add(name)
             self.tags_scanned += 1
+        elif ent_type == "stream":
+            self.streams_scanned += 1
         else:
             raise KeyError(f"Unknown entity {ent_type}.")
 

datahub/ingestion/source/snowflake/snowflake_schema.py

@@ -6,6 +6,7 @@ from datetime import datetime
 from typing import Callable, Dict, Iterable, List, MutableMapping, Optional
 
 from datahub.ingestion.api.report import SupportsAsObj
+from datahub.ingestion.source.common.subtypes import DatasetSubTypes
 from datahub.ingestion.source.snowflake.constants import SnowflakeObjectDomain
 from datahub.ingestion.source.snowflake.snowflake_connection import SnowflakeConnection
 from datahub.ingestion.source.snowflake.snowflake_query import (
@@ -14,7 +15,7 @@ from datahub.ingestion.source.snowflake.snowflake_query import (
 )
 from datahub.ingestion.source.sql.sql_generic import BaseColumn, BaseTable, BaseView
 from datahub.utilities.file_backed_collections import FileBackedDict
-from datahub.utilities.prefix_batch_builder import build_prefix_batches
+from datahub.utilities.prefix_batch_builder import PrefixGroup, build_prefix_batches
 from datahub.utilities.serialized_lru_cache import serialized_lru_cache
 
 logger: logging.Logger = logging.getLogger(__name__)
@@ -100,6 +101,9 @@ class SnowflakeTable(BaseTable):
     def is_hybrid(self) -> bool:
         return self.type is not None and self.type == "HYBRID TABLE"
 
+    def get_subtype(self) -> DatasetSubTypes:
+        return DatasetSubTypes.TABLE
+
 
 @dataclass
 class SnowflakeView(BaseView):
@@ -109,6 +113,9 @@ class SnowflakeView(BaseView):
     column_tags: Dict[str, List[SnowflakeTag]] = field(default_factory=dict)
     is_secure: bool = False
 
+    def get_subtype(self) -> DatasetSubTypes:
+        return DatasetSubTypes.VIEW
+
 
 @dataclass
 class SnowflakeSchema:
@@ -118,6 +125,7 @@ class SnowflakeSchema:
     comment: Optional[str]
     tables: List[str] = field(default_factory=list)
     views: List[str] = field(default_factory=list)
+    streams: List[str] = field(default_factory=list)
     tags: Optional[List[SnowflakeTag]] = None
 
 
@@ -131,6 +139,32 @@ class SnowflakeDatabase:
     tags: Optional[List[SnowflakeTag]] = None
 
 
+@dataclass
+class SnowflakeStream:
+    name: str
+    created: datetime
+    owner: str
+    source_type: str
+    type: str
+    stale: str
+    mode: str
+    invalid_reason: str
+    owner_role_type: str
+    database_name: str
+    schema_name: str
+    table_name: str
+    comment: Optional[str]
+    columns: List[SnowflakeColumn] = field(default_factory=list)
+    stale_after: Optional[datetime] = None
+    base_tables: Optional[str] = None
+    tags: Optional[List[SnowflakeTag]] = None
+    column_tags: Dict[str, List[SnowflakeTag]] = field(default_factory=dict)
+    last_altered: Optional[datetime] = None
+
+    def get_subtype(self) -> DatasetSubTypes:
+        return DatasetSubTypes.SNOWFLAKE_STREAM
+
+
 class _SnowflakeTagCache:
     def __init__(self) -> None:
         # self._database_tags[<database_name>] = list of tags applied to database
@@ -208,6 +242,7 @@ class SnowflakeDataDictionary(SupportsAsObj):
             self.get_tables_for_database,
             self.get_views_for_database,
             self.get_columns_for_schema,
+            self.get_streams_for_database,
             self.get_pk_constraints_for_schema,
             self.get_fk_constraints_for_schema,
         ]
@@ -431,9 +466,18 @@ class SnowflakeDataDictionary(SupportsAsObj):
             # For massive schemas, use a FileBackedDict to avoid memory issues.
             columns = FileBackedDict()
 
-        object_batches = build_prefix_batches(
-            all_objects, max_batch_size=10000, max_groups_in_batch=5
-        )
+        # Single prefix table case (for streams)
+        if len(all_objects) == 1:
+            object_batches = [
+                [PrefixGroup(prefix=all_objects[0], names=[], exact_match=True)]
+            ]
+        else:
+            # Build batches for full schema scan
+            object_batches = build_prefix_batches(
+                all_objects, max_batch_size=10000, max_groups_in_batch=5
+            )
+
+        # Process batches
         for batch_index, object_batch in enumerate(object_batches):
             if batch_index > 0:
                 logger.info(
@@ -611,3 +655,63 @@ class SnowflakeDataDictionary(SupportsAsObj):
             tags[column_name].append(snowflake_tag)
 
         return tags
+
+    @serialized_lru_cache(maxsize=1)
+    def get_streams_for_database(
+        self, db_name: str
+    ) -> Dict[str, List[SnowflakeStream]]:
+        page_limit = SHOW_VIEWS_MAX_PAGE_SIZE
+
+        streams: Dict[str, List[SnowflakeStream]] = {}
+
+        first_iteration = True
+        stream_pagination_marker: Optional[str] = None
+        while first_iteration or stream_pagination_marker is not None:
+            cur = self.connection.query(
+                SnowflakeQuery.streams_for_database(
+                    db_name,
+                    limit=page_limit,
+                    stream_pagination_marker=stream_pagination_marker,
+                )
+            )
+
+            first_iteration = False
+            stream_pagination_marker = None
+
+            result_set_size = 0
+            for stream in cur:
+                result_set_size += 1
+
+                stream_name = stream["name"]
+                schema_name = stream["schema_name"]
+                if schema_name not in streams:
+                    streams[schema_name] = []
+                streams[stream["schema_name"]].append(
+                    SnowflakeStream(
+                        name=stream["name"],
+                        created=stream["created_on"],
+                        owner=stream["owner"],
+                        comment=stream["comment"],
+                        source_type=stream["source_type"],
+                        type=stream["type"],
+                        stale=stream["stale"],
+                        mode=stream["mode"],
+                        database_name=stream["database_name"],
+                        schema_name=stream["schema_name"],
+                        invalid_reason=stream["invalid_reason"],
+                        owner_role_type=stream["owner_role_type"],
+                        stale_after=stream["stale_after"],
+                        table_name=stream["table_name"],
+                        base_tables=stream["base_tables"],
+                        last_altered=stream["created_on"],
+                    )
+                )
+
+            if result_set_size >= page_limit:
+                # If we hit the limit, we need to send another request to get the next page.
+                logger.info(
+                    f"Fetching next page of streams for {db_name} - after {stream_name}"
+                )
+                stream_pagination_marker = stream_name
+
+        return streams

datahub/ingestion/source/snowflake/snowflake_schema_gen.py

@@ -21,7 +21,6 @@ from datahub.ingestion.glossary.classification_mixin import (
 from datahub.ingestion.source.aws.s3_util import make_s3_urn_for_lineage
 from datahub.ingestion.source.common.subtypes import (
     DatasetContainerSubTypes,
-    DatasetSubTypes,
 )
 from datahub.ingestion.source.snowflake.constants import (
     GENERIC_PERMISSION_ERROR_KEY,
@@ -48,6 +47,7 @@ from datahub.ingestion.source.snowflake.snowflake_schema import (
     SnowflakeFK,
     SnowflakePK,
     SnowflakeSchema,
+    SnowflakeStream,
     SnowflakeTable,
     SnowflakeTag,
     SnowflakeView,
@@ -58,6 +58,7 @@ from datahub.ingestion.source.snowflake.snowflake_utils import (
     SnowflakeIdentifierBuilder,
     SnowflakeStructuredReportMixin,
     SnowsightUrlBuilder,
+    split_qualified_name,
 )
 from datahub.ingestion.source.sql.sql_utils import (
     add_table_to_schema_container,
@@ -70,6 +71,7 @@ from datahub.ingestion.source.sql.sql_utils import (
 )
 from datahub.ingestion.source_report.ingestion_stage import (
     EXTERNAL_TABLE_DDL_LINEAGE,
+    LINEAGE_EXTRACTION,
     METADATA_EXTRACTION,
     PROFILING,
 )
@@ -81,6 +83,7 @@ from datahub.metadata.com.linkedin.pegasus2avro.common import (
     TimeStamp,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
+    DatasetLineageTypeClass,
     DatasetProperties,
     ViewProperties,
 )
@@ -420,73 +423,126 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         schema_name = snowflake_schema.name
 
         if self.config.extract_tags != TagOption.skip:
-            snowflake_schema.tags = self.tag_extractor.get_tags_on_object(
-                schema_name=schema_name, db_name=db_name, domain="schema"
-            )
+            self._process_tags(snowflake_schema, schema_name, db_name, domain="schema")
 
         if self.config.include_technical_schema:
             yield from self.gen_schema_containers(snowflake_schema, db_name)
 
-        # We need to do this first so that we can use it when fetching columns.
+        tables, views, streams = [], [], []
+
         if self.config.include_tables:
             tables = self.fetch_tables_for_schema(
                 snowflake_schema, db_name, schema_name
             )
+            db_tables[schema_name] = tables
+            yield from self._process_tables(
+                tables, snowflake_schema, db_name, schema_name
+            )
+
         if self.config.include_views:
             views = self.fetch_views_for_schema(snowflake_schema, db_name, schema_name)
+            yield from self._process_views(
+                views, snowflake_schema, db_name, schema_name
+            )
 
-        if self.config.include_tables:
-            db_tables[schema_name] = tables
+        if self.config.include_streams:
+            self.report.num_get_streams_for_schema_queries += 1
+            streams = self.fetch_streams_for_schema(
+                snowflake_schema, db_name, schema_name
+            )
+            yield from self._process_streams(streams, snowflake_schema, db_name)
 
-            if self.config.include_technical_schema:
-                data_reader = self.make_data_reader()
-                for table in tables:
-                    table_wu_generator = self._process_table(
-                        table, snowflake_schema, db_name
-                    )
+        if self.config.include_technical_schema and snowflake_schema.tags:
+            yield from self._process_tags_in_schema(snowflake_schema)
 
-                    yield from classification_workunit_processor(
-                        table_wu_generator,
-                        self.classification_handler,
-                        data_reader,
-                        [db_name, schema_name, table.name],
-                    )
+        if (
+            not snowflake_schema.views
+            and not snowflake_schema.tables
+            and not snowflake_schema.streams
+        ):
+            self.structured_reporter.info(
+                title="No tables/views/streams found in schema",
+                message="If objects exist, please grant REFERENCES or SELECT permissions on them.",
+                context=f"{db_name}.{schema_name}",
+            )
 
-        if self.config.include_views:
-            if self.aggregator:
-                for view in views:
-                    view_identifier = self.identifiers.get_dataset_identifier(
+    def _process_tags(
+        self,
+        snowflake_schema: SnowflakeSchema,
+        schema_name: str,
+        db_name: str,
+        domain: str,
+    ) -> None:
+        snowflake_schema.tags = self.tag_extractor.get_tags_on_object(
+            schema_name=schema_name, db_name=db_name, domain=domain
+        )
+
+    def _process_tables(
+        self,
+        tables: List[SnowflakeTable],
+        snowflake_schema: SnowflakeSchema,
+        db_name: str,
+        schema_name: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        if self.config.include_technical_schema:
+            data_reader = self.make_data_reader()
+            for table in tables:
+                table_wu_generator = self._process_table(
+                    table, snowflake_schema, db_name
+                )
+                yield from classification_workunit_processor(
+                    table_wu_generator,
+                    self.classification_handler,
+                    data_reader,
+                    [db_name, schema_name, table.name],
+                )
+
+    def _process_views(
+        self,
+        views: List[SnowflakeView],
+        snowflake_schema: SnowflakeSchema,
+        db_name: str,
+        schema_name: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        if self.aggregator:
+            for view in views:
+                view_identifier = self.identifiers.get_dataset_identifier(
+                    view.name, schema_name, db_name
+                )
+                if view.is_secure and not view.view_definition:
+                    view.view_definition = self.fetch_secure_view_definition(
                         view.name, schema_name, db_name
                     )
-                    if view.is_secure and not view.view_definition:
-                        view.view_definition = self.fetch_secure_view_definition(
-                            view.name, schema_name, db_name
-                        )
-                    if view.view_definition:
-                        self.aggregator.add_view_definition(
-                            view_urn=self.identifiers.gen_dataset_urn(view_identifier),
-                            view_definition=view.view_definition,
-                            default_db=db_name,
-                            default_schema=schema_name,
-                        )
-                    elif view.is_secure:
-                        self.report.num_secure_views_missing_definition += 1
+                if view.view_definition:
+                    self.aggregator.add_view_definition(
+                        view_urn=self.identifiers.gen_dataset_urn(view_identifier),
+                        view_definition=view.view_definition,
+                        default_db=db_name,
+                        default_schema=schema_name,
+                    )
+                elif view.is_secure:
+                    self.report.num_secure_views_missing_definition += 1
 
-            if self.config.include_technical_schema:
-                for view in views:
-                    yield from self._process_view(view, snowflake_schema, db_name)
+        if self.config.include_technical_schema:
+            for view in views:
+                yield from self._process_view(view, snowflake_schema, db_name)
 
-        if self.config.include_technical_schema and snowflake_schema.tags:
+    def _process_streams(
+        self,
+        streams: List[SnowflakeStream],
+        snowflake_schema: SnowflakeSchema,
+        db_name: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        for stream in streams:
+            yield from self._process_stream(stream, snowflake_schema, db_name)
+
+    def _process_tags_in_schema(
+        self, snowflake_schema: SnowflakeSchema
+    ) -> Iterable[MetadataWorkUnit]:
+        if snowflake_schema.tags:
             for tag in snowflake_schema.tags:
                 yield from self._process_tag(tag)
 
-        if not snowflake_schema.views and not snowflake_schema.tables:
-            self.structured_reporter.info(
-                title="No tables/views found in schema",
-                message="If tables exist, please grant REFERENCES or SELECT permissions on them.",
-                context=f"{db_name}.{schema_name}",
-            )
-
     def fetch_secure_view_definition(
         self, table_name: str, schema_name: str, db_name: str
     ) -> Optional[str]:
@@ -729,7 +785,7 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
 
     def gen_dataset_workunits(
         self,
-        table: Union[SnowflakeTable, SnowflakeView],
+        table: Union[SnowflakeTable, SnowflakeView, SnowflakeStream],
         schema_name: str,
         db_name: str,
     ) -> Iterable[MetadataWorkUnit]:
@@ -786,13 +842,7 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         if dpi_aspect:
             yield dpi_aspect
 
-        subTypes = SubTypes(
-            typeNames=(
-                [DatasetSubTypes.VIEW]
-                if isinstance(table, SnowflakeView)
-                else [DatasetSubTypes.TABLE]
-            )
-        )
+        subTypes = SubTypes(typeNames=[table.get_subtype()])
 
         yield MetadataChangeProposalWrapper(
             entityUrn=dataset_urn, aspect=subTypes
@@ -843,28 +893,50 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
 
     def get_dataset_properties(
         self,
-        table: Union[SnowflakeTable, SnowflakeView],
+        table: Union[SnowflakeTable, SnowflakeView, SnowflakeStream],
         schema_name: str,
         db_name: str,
     ) -> DatasetProperties:
         custom_properties = {}
 
         if isinstance(table, SnowflakeTable):
-            if table.clustering_key:
-                custom_properties["CLUSTERING_KEY"] = table.clustering_key
-
-            if table.is_hybrid:
-                custom_properties["IS_HYBRID"] = "true"
-
-            if table.is_dynamic:
-                custom_properties["IS_DYNAMIC"] = "true"
-
-            if table.is_iceberg:
-                custom_properties["IS_ICEBERG"] = "true"
+            custom_properties.update(
+                {
+                    k: v
+                    for k, v in {
+                        "CLUSTERING_KEY": table.clustering_key,
+                        "IS_HYBRID": "true" if table.is_hybrid else None,
+                        "IS_DYNAMIC": "true" if table.is_dynamic else None,
+                        "IS_ICEBERG": "true" if table.is_iceberg else None,
+                    }.items()
+                    if v
+                }
+            )
 
         if isinstance(table, SnowflakeView) and table.is_secure:
             custom_properties["IS_SECURE"] = "true"
 
+        elif isinstance(table, SnowflakeStream):
+            custom_properties.update(
+                {
+                    k: v
+                    for k, v in {
+                        "SOURCE_TYPE": table.source_type,
+                        "TYPE": table.type,
+                        "STALE": table.stale,
+                        "MODE": table.mode,
+                        "INVALID_REASON": table.invalid_reason,
+                        "OWNER_ROLE_TYPE": table.owner_role_type,
+                        "TABLE_NAME": table.table_name,
+                        "BASE_TABLES": table.base_tables,
+                        "STALE_AFTER": (
+                            table.stale_after.isoformat() if table.stale_after else None
+                        ),
+                    }.items()
+                    if v
+                }
+            )
+
         return DatasetProperties(
             name=table.name,
             created=(
@@ -909,7 +981,9 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         ).as_workunit()
 
     def gen_column_tags_as_structured_properties(
-        self, dataset_urn: str, table: Union[SnowflakeTable, SnowflakeView]
+        self,
+        dataset_urn: str,
+        table: Union[SnowflakeTable, SnowflakeView, SnowflakeStream],
     ) -> Iterable[MetadataWorkUnit]:
         for column_name in table.column_tags:
             schema_field_urn = SchemaFieldUrn(dataset_urn, column_name).urn()
@@ -922,7 +996,7 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
 
     def gen_schema_metadata(
         self,
-        table: Union[SnowflakeTable, SnowflakeView],
+        table: Union[SnowflakeTable, SnowflakeView, SnowflakeStream],
         schema_name: str,
         db_name: str,
     ) -> SchemaMetadata:
@@ -1214,3 +1288,158 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
                 "External table ddl lineage extraction failed",
                 exc=e,
             )
+
+    def fetch_streams_for_schema(
+        self, snowflake_schema: SnowflakeSchema, db_name: str, schema_name: str
+    ) -> List[SnowflakeStream]:
+        try:
+            streams: List[SnowflakeStream] = []
+            for stream in self.get_streams_for_schema(schema_name, db_name):
+                stream_identifier = self.identifiers.get_dataset_identifier(
+                    stream.name, schema_name, db_name
+                )
+
+                self.report.report_entity_scanned(stream_identifier, "stream")
+
+                if not self.filters.is_dataset_pattern_allowed(
+                    stream_identifier, SnowflakeObjectDomain.STREAM
+                ):
+                    self.report.report_dropped(stream_identifier)
+                else:
+                    streams.append(stream)
+            snowflake_schema.streams = [stream.name for stream in streams]
+            return streams
+        except Exception as e:
+            if isinstance(e, SnowflakePermissionError):
+                error_msg = f"Failed to get streams for schema {db_name}.{schema_name}. Please check permissions."
+                raise SnowflakePermissionError(error_msg) from e.__cause__
+            else:
+                self.structured_reporter.warning(
+                    "Failed to get streams for schema",
+                    f"{db_name}.{schema_name}",
+                    exc=e,
+                )
+                return []
+
+    def get_streams_for_schema(
+        self, schema_name: str, db_name: str
+    ) -> List[SnowflakeStream]:
+        streams = self.data_dictionary.get_streams_for_database(db_name)
+
+        return streams.get(schema_name, [])
+
+    def _process_stream(
+        self,
+        stream: SnowflakeStream,
+        snowflake_schema: SnowflakeSchema,
+        db_name: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        schema_name = snowflake_schema.name
+
+        try:
+            # Retrieve and register the schema without metadata to prevent columns from mapping upstream
+            stream.columns = self.get_columns_for_stream(stream.table_name)
+            yield from self.gen_dataset_workunits(stream, schema_name, db_name)
+
+            if self.config.include_column_lineage:
+                with self.report.new_stage(f"*: {LINEAGE_EXTRACTION}"):
+                    self.populate_stream_upstreams(stream, db_name, schema_name)
+
+        except Exception as e:
+            self.structured_reporter.warning(
+                "Failed to get columns for stream:", stream.name, exc=e
+            )
+
+    def get_columns_for_stream(
+        self,
+        source_object: str,  # Qualified name of source table/view
+    ) -> List[SnowflakeColumn]:
+        """
+        Get column information for a stream by getting source object columns and adding metadata columns.
+        Stream includes all columns from source object plus metadata columns like:
+        - METADATA$ACTION
+        - METADATA$ISUPDATE
+        - METADATA$ROW_ID
+        """
+        columns: List[SnowflakeColumn] = []
+
+        source_parts = split_qualified_name(source_object)
+
+        source_db, source_schema, source_name = source_parts
+
+        # Get columns from source object
+        source_columns = self.data_dictionary.get_columns_for_schema(
+            source_schema, source_db, itertools.chain([source_name])
+        ).get(source_name, [])
+
+        # Add all source columns
+        columns.extend(source_columns)
+
+        # Add standard stream metadata columns
+        metadata_columns = [
+            SnowflakeColumn(
+                name="METADATA$ACTION",
+                ordinal_position=len(columns) + 1,
+                is_nullable=False,
+                data_type="VARCHAR",
+                comment="Type of DML operation (INSERT/DELETE)",
+                character_maximum_length=10,
+                numeric_precision=None,
+                numeric_scale=None,
+            ),
+            SnowflakeColumn(
+                name="METADATA$ISUPDATE",
+                ordinal_position=len(columns) + 2,
+                is_nullable=False,
+                data_type="BOOLEAN",
+                comment="Whether row is from UPDATE operation",
+                character_maximum_length=None,
+                numeric_precision=None,
+                numeric_scale=None,
+            ),
+            SnowflakeColumn(
+                name="METADATA$ROW_ID",
+                ordinal_position=len(columns) + 3,
+                is_nullable=False,
+                data_type="NUMBER",
+                comment="Unique row identifier",
+                character_maximum_length=None,
+                numeric_precision=38,
+                numeric_scale=0,
+            ),
+        ]
+
+        columns.extend(metadata_columns)
+
+        return columns
+
+    def populate_stream_upstreams(
+        self, stream: SnowflakeStream, db_name: str, schema_name: str
+    ) -> None:
+        """
+        Populate Streams upstream tables
+        """
+        self.report.num_streams_with_known_upstreams += 1
+        if self.aggregator:
+            source_parts = split_qualified_name(stream.table_name)
+            source_db, source_schema, source_name = source_parts
+
+            dataset_identifier = self.identifiers.get_dataset_identifier(
+                stream.name, schema_name, db_name
+            )
+            dataset_urn = self.identifiers.gen_dataset_urn(dataset_identifier)
+
+            upstream_identifier = self.identifiers.get_dataset_identifier(
+                source_name, source_schema, source_db
+            )
+            upstream_urn = self.identifiers.gen_dataset_urn(upstream_identifier)
+
+            logger.debug(
+                f"""upstream_urn: {upstream_urn}, downstream_urn: {dataset_urn}"""
+            )
+
+            self.aggregator.add_known_lineage_mapping(
+                upstream_urn=upstream_urn,
+                downstream_urn=dataset_urn,
+                lineage_type=DatasetLineageTypeClass.COPY,
+            )