fix(ingestion/redshift): Fix for Redshift COPY-based lineage

datahub-project · Oct 7, 2024 · d618a69 · d618a69
1 parent 7bdb7bb
commit d618a69
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 5 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/redshift/lineage_v2.py b/metadata-ingestion/src/datahub/ingestion/source/redshift/lineage_v2.py
@@ -334,19 +334,26 @@ def _process_view_lineage(self, lineage_row: LineageRow) -> None:
         )
 
     def _process_copy_command(self, lineage_row: LineageRow) -> None:
-        source = self._lineage_v1._get_sources(
+        logger.debug(f"Processing COPY command for lineage row: {lineage_row}")
+        sources = self._lineage_v1._get_sources(
             lineage_type=LineageCollectorType.COPY,
             db_name=self.database,
             source_schema=None,
             source_table=None,
             ddl=None,
             filename=lineage_row.filename,
-        )[0]
+        )
+        logger.debug(f"Recognized sources: {sources}")
+        source = sources[0]
         if not source:
+            logger.debug("Ignoring command since couldn't recognize proper source")
             return
         s3_urn = source[0].urn
-
+        logger.debug(f"Recognized s3 dataset urn: {s3_urn}")
         if not lineage_row.target_schema or not lineage_row.target_table:
+            logger.debug(
+                f"Didn't find target schema (found: {lineage_row.target_schema}) or target table (found: {lineage_row.target_table})"
+            )
             return
         target = self._make_filtered_target(lineage_row)
         if not target:

diff --git a/metadata-ingestion/src/datahub/ingestion/source/redshift/query.py b/metadata-ingestion/src/datahub/ingestion/source/redshift/query.py
@@ -950,7 +950,7 @@ def list_copy_commands_sql(
                     distinct
                         "schema" as target_schema,
                         "table" as target_table,
-                        c.file_name
+                        c.file_name as filename
                 from
                     SYS_QUERY_DETAIL as si
                 join SYS_LOAD_DETAIL as c on

diff --git a/metadata-ingestion/src/datahub/sql_parsing/sql_parsing_aggregator.py b/metadata-ingestion/src/datahub/sql_parsing/sql_parsing_aggregator.py
@@ -613,7 +613,9 @@ def add_known_lineage_mapping(
             upstream_urn: The upstream dataset URN.
             downstream_urn: The downstream dataset URN.
         """
-
+        logger.debug(
+            f"Adding lineage to the map, downstream: {downstream_urn}, upstream: {upstream_urn}"
+        )
         self.report.num_known_mapping_lineage += 1
 
         # We generate a fake "query" object to hold the lineage.