refactor(ingest/powerbi): organize code within the module based on responsibilities

metadata-ingestion/setup.py

@@ -773,7 +773,7 @@
         "trino = datahub.ingestion.source.sql.trino:TrinoSource",
         "starburst-trino-usage = datahub.ingestion.source.usage.starburst_trino_usage:TrinoUsageSource",
         "nifi = datahub.ingestion.source.nifi:NifiSource",
-        "powerbi = datahub.ingestion.source.powerbi:PowerBiDashboardSource",
+        "powerbi = datahub.ingestion.source.powerbi.powerbi:PowerBiDashboardSource",
         "powerbi-report-server = datahub.ingestion.source.powerbi_report_server:PowerBiReportServerDashboardSource",
         "iceberg = datahub.ingestion.source.iceberg.iceberg:IcebergSource",
         "vertica = datahub.ingestion.source.sql.vertica:VerticaSource",

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

@@ -1 +0,0 @@
-from datahub.ingestion.source.powerbi.powerbi import PowerBiDashboardSource

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

@@ -173,7 +173,7 @@ class SupportedDataPlatform(Enum):
         datahub_data_platform_name="redshift",
     )
 
-    DATABRICK_SQL = DataPlatformPair(
+    DATABRICKS_SQL = DataPlatformPair(
         powerbi_data_platform_name="Databricks", datahub_data_platform_name="databricks"
     )
 
@@ -313,8 +313,8 @@ class PowerBiDashboardSourceConfig(
         " Note: This field works in conjunction with 'workspace_type_filter' and both must be considered when filtering workspaces.",
     )
 
-    # Dataset type mapping PowerBI support many type of data-sources. Here user need to define what type of PowerBI
-    # DataSource need to be mapped to corresponding DataHub Platform DataSource. For example PowerBI `Snowflake` is
+    # Dataset type mapping PowerBI support many type of data-sources. Here user needs to define what type of PowerBI
+    # DataSource needs to be mapped to corresponding DataHub Platform DataSource. For example, PowerBI `Snowflake` is
     # mapped to DataHub `snowflake` PowerBI `PostgreSQL` is mapped to DataHub `postgres` and so on.
     dataset_type_mapping: Union[
         Dict[str, str], Dict[str, PlatformDetail]

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

@@ -1,10 +1,14 @@
 import os
 from abc import ABC
 from dataclasses import dataclass
-from typing import Any, Dict, Optional
+from enum import Enum
+from typing import Any, Dict, List, Optional
 
 from lark import Tree
 
+from datahub.ingestion.source.powerbi.config import DataPlatformPair
+from datahub.sql_parsing.sqlglot_lineage import ColumnLineageInfo
+
 TRACE_POWERBI_MQUERY_PARSER = os.getenv("DATAHUB_TRACE_POWERBI_MQUERY_PARSER", False)
 
 
@@ -30,7 +34,7 @@ class IdentifierAccessor(AbstractIdentifierAccessor):
 
         "[Schema="public",Item="order_date"]" is "items" in ItemSelector. Data of items varies as per DataSource
 
-        "public_order_date" is in "next" of ItemSelector. The "next" will be None if this identifier is leaf i.e. table
+        "public_order_date" is in "next" of ItemSelector. The "next" will be None if this identifier is leaf i.e., table
 
     """
 
@@ -53,3 +57,31 @@ class ReferencedTable:
     database: str
     schema: str
     table: str
+
+
+@dataclass
+class DataPlatformTable:
+    data_platform_pair: DataPlatformPair
+    urn: str
+
+
+@dataclass
+class Lineage:
+    upstreams: List[DataPlatformTable]
+    column_lineage: List[ColumnLineageInfo]
+
+    @staticmethod
+    def empty() -> "Lineage":
+        return Lineage(upstreams=[], column_lineage=[])
+
+
+class FunctionName(Enum):
+    NATIVE_QUERY = "Value.NativeQuery"
+    POSTGRESQL_DATA_ACCESS = "PostgreSQL.Database"
+    ORACLE_DATA_ACCESS = "Oracle.Database"
+    SNOWFLAKE_DATA_ACCESS = "Snowflake.Databases"
+    MSSQL_DATA_ACCESS = "Sql.Database"
+    DATABRICK_DATA_ACCESS = "Databricks.Catalogs"
+    GOOGLE_BIGQUERY_DATA_ACCESS = "GoogleBigQuery.Database"
+    AMAZON_REDSHIFT_DATA_ACCESS = "AmazonRedshift.Database"
+    DATABRICK_MULTI_CLOUD_DATA_ACCESS = "DatabricksMultiCloud.Catalogs"

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

@@ -7,6 +7,7 @@
 import lark
 from lark import Lark, Tree
 
+import datahub.ingestion.source.powerbi.m_query.data_classes
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.source.powerbi.config import (
     PowerBiDashboardSourceConfig,
@@ -65,7 +66,7 @@ def get_upstream_tables(
     ctx: PipelineContext,
     config: PowerBiDashboardSourceConfig,
     parameters: Dict[str, str] = {},
-) -> List[resolver.Lineage]:
+) -> List[datahub.ingestion.source.powerbi.m_query.data_classes.Lineage]:
     if table.expression is None:
         logger.debug(f"There is no M-Query expression in table {table.full_name}")
         return []
@@ -127,12 +128,14 @@ def get_upstream_tables(
     reporter.m_query_parse_successes += 1
 
     try:
-        lineage: List[resolver.Lineage] = resolver.MQueryResolver(
+        lineage: List[
+            datahub.ingestion.source.powerbi.m_query.data_classes.Lineage
+        ] = resolver.MQueryResolver(
             table=table,
             parse_tree=parse_tree,
             reporter=reporter,
             parameters=parameters,
-        ).resolve_to_data_platform_table_list(
+        ).resolve_to_lineage(
             ctx=ctx,
             config=config,
             platform_instance_resolver=platform_instance_resolver,