feat(ingestion/powerbi): support multiple tables as upstream in native SQL parsing

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

@@ -121,6 +121,12 @@ class DataPlatformPair:
     powerbi_data_platform_name: str
 
 
+@dataclass
+class PowerBIPlatformDetail:
+    data_platform_pair: DataPlatformPair
+    data_platform_server: str
+
+
 class SupportedDataPlatform(Enum):
     POSTGRES_SQL = DataPlatformPair(
         powerbi_data_platform_name="PostgreSQL", datahub_data_platform_name="postgres"
@@ -382,6 +388,15 @@ class PowerBiDashboardSourceConfig(
         description="The instance of the platform that all assets produced by this recipe belong to",
     )
 
+    # Enable advance sql construct
+    enable_advance_lineage_sql_construct: bool = pydantic.Field(
+        default=False,
+        description="Whether to enable advance native sql construct for parsing like join, sub-queries. "
+        "along this flag , the native_query_parsing should be enabled. "
+        "By default convert_lineage_urns_to_lowercase is enabled, in-case if you have disabled it in previous ingestion execution then it may break lineage "
+        "as this option generates the upstream datasets URN in lowercase.",
+    )
+
     @validator("dataset_type_mapping")
     @classmethod
     def map_data_platform(cls, value):

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

@@ -5,16 +5,16 @@
 from datahub.ingestion.source.powerbi.config import (
     PlatformDetail,
     PowerBiDashboardSourceConfig,
+    PowerBIPlatformDetail,
 )
-from datahub.ingestion.source.powerbi.m_query.resolver import DataPlatformTable
 
 logger = logging.getLogger(__name__)
 
 
 class AbstractDataPlatformInstanceResolver(ABC):
     @abstractmethod
     def get_platform_instance(
-        self, dataplatform_table: DataPlatformTable
+        self, data_platform_detail: PowerBIPlatformDetail
     ) -> PlatformDetail:
         pass
 
@@ -32,10 +32,10 @@ class ResolvePlatformInstanceFromDatasetTypeMapping(
     BaseAbstractDataPlatformInstanceResolver
 ):
     def get_platform_instance(
-        self, dataplatform_table: DataPlatformTable
+        self, data_platform_detail: PowerBIPlatformDetail
     ) -> PlatformDetail:
         platform: Union[str, PlatformDetail] = self.config.dataset_type_mapping[
-            dataplatform_table.data_platform_pair.powerbi_data_platform_name
+            data_platform_detail.data_platform_pair.powerbi_data_platform_name
         ]
 
         if isinstance(platform, PlatformDetail):
@@ -48,13 +48,13 @@ class ResolvePlatformInstanceFromServerToPlatformInstance(
     BaseAbstractDataPlatformInstanceResolver
 ):
     def get_platform_instance(
-        self, dataplatform_table: DataPlatformTable
+        self, data_platform_detail: PowerBIPlatformDetail
     ) -> PlatformDetail:
         return (
             self.config.server_to_platform_instance[
-                dataplatform_table.datasource_server
+                data_platform_detail.data_platform_server
             ]
-            if dataplatform_table.datasource_server
+            if data_platform_detail.data_platform_server
             in self.config.server_to_platform_instance
             else PlatformDetail.parse_obj({})
         )

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

@@ -1,8 +1,15 @@
 import logging
-from typing import List
+from typing import List, Optional
 
 import sqlparse
 
+from datahub.ingestion.api.common import PipelineContext
+from datahub.utilities.sqlglot_lineage import (
+    SchemaResolver,
+    SqlParsingResult,
+    sqlglot_lineage,
+)
+
 SPECIAL_CHARACTERS = ["#(lf)", "(lf)"]
 
 logger = logging.getLogger()
@@ -45,3 +52,48 @@ def get_tables(native_query: str) -> List[str]:
         from_index = from_index + 1
 
     return tables
+
+
+def parse_custom_sql(
+    ctx: PipelineContext,
+    query: str,
+    schema: Optional[str],
+    database: Optional[str],
+    platform: str,
+    env: str,
+    platform_instance: Optional[str],
+) -> Optional["SqlParsingResult"]:
+
+    logger.debug("Using sqlglot_lineage to parse custom sql")
+
+    sql_query = remove_special_characters(query)
+    logger.debug(f"Parsing sql={sql_query}")
+
+    parsed_result: Optional["SqlParsingResult"] = None
+    try:
+        schema_resolver = (
+            ctx.graph._make_schema_resolver(
+                platform=platform,
+                platform_instance=platform_instance,
+                env=env,
+            )
+            if ctx.graph is not None
+            else SchemaResolver(
+                platform=platform,
+                platform_instance=platform_instance,
+                env=env,
+                graph=None,
+            )
+        )
+
+        parsed_result = sqlglot_lineage(
+            sql_query,
+            schema_resolver=schema_resolver,
+            default_db=database,
+            default_schema=schema,
+        )
+    except Exception as e:
+        logger.debug(f"Fail to prase query {query}", exc_info=e)
+        logger.warning("Fail to parse custom SQL")
+
+    return parsed_result

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

@@ -6,7 +6,14 @@
 import lark
 from lark import Lark, Tree
 
-from datahub.ingestion.source.powerbi.config import PowerBiDashboardSourceReport
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.source.powerbi.config import (
+    PowerBiDashboardSourceConfig,
+    PowerBiDashboardSourceReport,
+)
+from datahub.ingestion.source.powerbi.dataplatform_instance_resolver import (
+    AbstractDataPlatformInstanceResolver,
+)
 from datahub.ingestion.source.powerbi.m_query import resolver, validator
 from datahub.ingestion.source.powerbi.m_query.data_classes import (
     TRACE_POWERBI_MQUERY_PARSER,
@@ -45,7 +52,9 @@ def _parse_expression(expression: str) -> Tree:
 def get_upstream_tables(
     table: Table,
     reporter: PowerBiDashboardSourceReport,
-    native_query_enabled: bool = True,
+    platform_instance_resolver: AbstractDataPlatformInstanceResolver,
+    ctx: PipelineContext,
+    config: PowerBiDashboardSourceConfig,
     parameters: Dict[str, str] = {},
 ) -> List[resolver.DataPlatformTable]:
     if table.expression is None:
@@ -58,7 +67,7 @@ def get_upstream_tables(
         parse_tree: Tree = _parse_expression(table.expression)
 
         valid, message = validator.validate_parse_tree(
-            parse_tree, native_query_enabled=native_query_enabled
+            parse_tree, native_query_enabled=config.native_query_parsing
         )
         if valid is False:
             assert message is not None
@@ -84,7 +93,11 @@ def get_upstream_tables(
             parse_tree=parse_tree,
             reporter=reporter,
             parameters=parameters,
-        ).resolve_to_data_platform_table_list()
+        ).resolve_to_data_platform_table_list(
+            ctx=ctx,
+            config=config,
+            platform_instance_resolver=platform_instance_resolver,
+        )
 
     except BaseException as e:
         reporter.report_warning(table.full_name, "Failed to process m-query expression")