fix(ingest/powerbi): Snowflake column lineage

Author: mminichino

metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/pattern_handler.py

@@ -1,7 +1,7 @@
 import logging
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import Dict, List, Optional, Tuple, Type, cast
+from typing import Dict, List, Optional, Tuple, Type, cast, Any
 
 from lark import Tree
 
@@ -30,7 +30,13 @@
     ReferencedTable,
 )
 from datahub.ingestion.source.powerbi.rest_api_wrapper.data_classes import Table
-from datahub.sql_parsing.sqlglot_lineage import SqlParsingResult
+from datahub.sql_parsing.sqlglot_lineage import (
+    SqlParsingResult,
+    ColumnLineageInfo,
+    DownstreamColumnRef,
+    ColumnRef
+)
+from datahub.metadata.schema_classes import SchemaFieldDataTypeClass
 
 logger = logging.getLogger(__name__)
 
@@ -262,6 +268,35 @@ def parse_custom_sql(
             ),
         )
 
+    def create_table_column_lineage(
+            self, urn: str
+    ) -> List[ColumnLineageInfo]:
+        column_lineage = []
+
+        for column in self.table.columns:
+            downstream = DownstreamColumnRef(
+                table=self.table.name,
+                column=column.name,
+                column_type=SchemaFieldDataTypeClass(type=column.datahubDataType),
+                native_column_type=column.dataType or "UNKNOWN"
+            )
+
+            upstreams = [
+                ColumnRef(
+                    table=urn,
+                    column=column.name.lower(),
+                )
+            ]
+
+            column_lineage_info = ColumnLineageInfo(
+                downstream=downstream,
+                upstreams=upstreams
+            )
+
+            column_lineage.append(column_lineage_info)
+
+        return column_lineage
+
 
 class AmazonRedshiftLineage(AbstractLineage):
     def get_platform_pair(self) -> DataPlatformPair:
@@ -671,14 +706,16 @@ def create_lineage(
             qualified_table_name=qualified_table_name,
         )
 
+        column_lineage = self.create_table_column_lineage(urn)
+
         return Lineage(
             upstreams=[
                 DataPlatformTable(
                     data_platform_pair=self.get_platform_pair(),
                     urn=urn,
                 )
             ],
-            column_lineage=[],
+            column_lineage=column_lineage,
         )
 
 

metadata-ingestion/src/datahub/ingestion/source/powerbi/m_query/resolver.py

@@ -361,6 +361,9 @@ def resolve_to_lineage(
         )
 
         if output_variable is None:
+            logger.debug(
+                f"Table: {self.table.full_name}: output-variable not found in tree"
+            )
             self.reporter.report_warning(
                 f"{self.table.full_name}-output-variable",
                 "output-variable not found in table expression",
@@ -374,6 +377,9 @@ def resolve_to_lineage(
 
         # Each item is data-access function
         for f_detail in table_links:
+            logger.debug(
+                f"Processing data-access-function {f_detail.data_access_function_name}"
+            )
             # Get & Check if we support data-access-function available in M-Query
             supported_resolver = SupportedPattern.get_pattern_handler(
                 f_detail.data_access_function_name
@@ -390,6 +396,10 @@ def resolve_to_lineage(
 
             # From supported_resolver enum get respective handler like AmazonRedshift or Snowflake or Oracle or NativeQuery and create instance of it
             # & also pass additional information that will be need to generate lineage
+            logger.debug(
+                f"Creating instance of {supported_resolver.handler().__name__} "
+                f"for data-access-function {f_detail.data_access_function_name}"
+            )
             pattern_handler: AbstractLineage = supported_resolver.handler()(
                 ctx=ctx,
                 table=self.table,

metadata-ingestion/src/datahub/ingestion/source/powerbi/powerbi.py

@@ -238,6 +238,10 @@ def extract_lineage(
         upstream: List[UpstreamClass] = []
         cll_lineage: List[FineGrainedLineage] = []
 
+        logger.debug(
+            f"Extracting lineage for table {table.full_name} in dataset {table.dataset.name}"
+        )
+
         upstream_lineage: List[
             datahub.ingestion.source.powerbi.m_query.data_classes.Lineage
         ] = parser.get_upstream_tables(