acryl-datahub 0.15.0rc3__py3-none-any.whl → 0.15.0rc5__py3-none-any.whl

{acryl_datahub-0.15.0rc3.dist-info → acryl_datahub-0.15.0rc5.dist-info}/RECORD

@@ -1,4 +1,4 @@
-datahub/__init__.py,sha256=bjNaEX_th93Zf5oFOCSAaZbNy-DUb3y4CxlHrAvZOkI,574
+datahub/__init__.py,sha256=c5YiGS9ajJPufFiwc_4_Bv9DF1Ha6s0H9dd-rtKRF3Y,574
 datahub/__main__.py,sha256=pegIvQ9hzK7IhqVeUi1MeADSZ2QlP-D3K0OQdEg55RU,106
 datahub/entrypoints.py,sha256=3-qSfXAx3Z0FEkBV5tlO8fQr4xk4ySeDRMVTpS5Xd6A,7793
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -187,7 +187,7 @@ datahub/ingestion/source/confluent_schema_registry.py,sha256=_h9D8bUXoaGcwgwB94d
 datahub/ingestion/source/csv_enricher.py,sha256=xjCbcsSMM8l_ASCRAnNsUGKuYMrD1lec19Waixub1EM,29498
 datahub/ingestion/source/demo_data.py,sha256=yzA_R-wfSX2WPz0i5ukYlscpmpb0Pt8D7EkhtKfftvo,1286
 datahub/ingestion/source/elastic_search.py,sha256=qFUVNzynTVJTabASTjGMu8Qhf9UpNbEtSBFjaPQjBJE,22641
-datahub/ingestion/source/feast.py,sha256=NYaAjzLVRhmMKDawBwN0OL8AMyKDLsxOwEj3YFX0wIA,14244
+datahub/ingestion/source/feast.py,sha256=uZpeUkJsiNlvZcUkARiEuZT_3n6sbGc0yFzwqhtnefA,18103
 datahub/ingestion/source/file.py,sha256=pH-Qkjh5FQ2XvyYPE7Z8XEY4vUk_SUHxm8p8IxG12tU,15879
 datahub/ingestion/source/ge_data_profiler.py,sha256=JqTonv8y7Re4Rfn2YKOEaLufiiAOWKfK1XQvJfV5dvs,64126
 datahub/ingestion/source/ge_profiling_config.py,sha256=P-9pd20koFvpxeEL_pqFvKWWz-qnpZ6XkELUyBKr7is,10807
@@ -312,14 +312,14 @@ datahub/ingestion/source/git/git_import.py,sha256=5CT6vMDb0MDctCtShnxb3JVihULtvk
 datahub/ingestion/source/grafana/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/grafana/grafana_source.py,sha256=3pU3xodPgS5lmnjuQ_u7F0XPzD_Y8MnPlMxRJ86qz4g,4960
 datahub/ingestion/source/iceberg/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/ingestion/source/iceberg/iceberg.py,sha256=lfCettf4F_BjWt7Wh65AhXHerXPP5KLFEo_LCbEZZnI,25537
+datahub/ingestion/source/iceberg/iceberg.py,sha256=fjqp3VBW5W5-54X_-ubkRZiAmdHvuMbxRbC4UYzEr4U,25900
 datahub/ingestion/source/iceberg/iceberg_common.py,sha256=TS3_ZYZ47Fe02CmzEo1z0pvy7yjXuG1VlwqNxa0U6pc,8506
 datahub/ingestion/source/iceberg/iceberg_profiler.py,sha256=hLT1Le_TEUoFXvsJSlrRB1qbTiTe-YVGCof5TFHMyd8,9908
 datahub/ingestion/source/identity/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/identity/azure_ad.py,sha256=GdmJFD4UMsb5353Z7phXRf-YsXR2woGLRJwBXUkgXq0,28809
 datahub/ingestion/source/identity/okta.py,sha256=PnRokWLG8wSoNZlXJiRZiW6APTEHO09q4n2j_l6m3V0,30756
 datahub/ingestion/source/kafka/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/ingestion/source/kafka/kafka.py,sha256=D54QXYoEQJYVb7oWyvLVCGbJunS8ZTGYlMMdaBqpmMs,25475
+datahub/ingestion/source/kafka/kafka.py,sha256=QUw8VCmqIhZJvUiFJmFmekFmy4nXCLD4EKJNC6jk6Y4,26092
 datahub/ingestion/source/kafka/kafka_connect.py,sha256=5KUlhn3876c41Z3kx5l4oJhbu0ekXZQRdxmu52vb_v8,55167
 datahub/ingestion/source/kafka/kafka_schema_registry_base.py,sha256=13XjSwqyVhH1CJUFHAbWdmmv_Rw0Ju_9HQdBmIzPNNA,566
 datahub/ingestion/source/looker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -347,18 +347,19 @@ datahub/ingestion/source/looker/view_upstream.py,sha256=k278-uwh8uspdREpjE_uqks4
 datahub/ingestion/source/metadata/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/metadata/business_glossary.py,sha256=eRVRpQI0ZX5OofS1BUhNihFOfWih70TIAkJM7zaMH80,17577
 datahub/ingestion/source/metadata/lineage.py,sha256=XiZGuY6k3O9qBmgo7AzosIndJHwrvEhapVLdRlDxCuc,9507
-datahub/ingestion/source/powerbi/__init__.py,sha256=jYWYtaz95cyQAgEpZK1kxu1aKjOBAyDUlE77UdkCd6g,76
-datahub/ingestion/source/powerbi/config.py,sha256=T7E3YpfkLJlOJLr_xKfAVXhQRu_fkMqJbPtWp_v4k-8,23237
+datahub/ingestion/source/powerbi/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+datahub/ingestion/source/powerbi/config.py,sha256=LV8BOm2zzF9t0RMwQVVUNB0bStzBPo8A6JkaW0xlgsQ,23241
 datahub/ingestion/source/powerbi/dataplatform_instance_resolver.py,sha256=AIU89lVPoCWlzc_RfUjDJwRQ11akPtnGpBTluBMCKio,2242
 datahub/ingestion/source/powerbi/powerbi-lexical-grammar.rule,sha256=PRWzuZMMhKdOVoAaE8csHvUFbZHxYe5meJHgrqlgiuw,19795
-datahub/ingestion/source/powerbi/powerbi.py,sha256=AuJPYakZ0Uko8lQUkJkeKPj2mtNmOHndUSwcwPOteMU,54294
+datahub/ingestion/source/powerbi/powerbi.py,sha256=7UsAEqaFlkWONcXJdQ2hotUYYn46ks6Fe71KXEMh7lI,54495
 datahub/ingestion/source/powerbi/m_query/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/ingestion/source/powerbi/m_query/data_classes.py,sha256=NLEF8uNQNhtOlMf42F7jHAd8C0TIwDCDOmMfLeRpyk8,1278
+datahub/ingestion/source/powerbi/m_query/data_classes.py,sha256=s2Nckmr50hxae5gPFcIfpyLzYpaMH56Q9nsDsEgi_-k,2243
 datahub/ingestion/source/powerbi/m_query/native_sql_parser.py,sha256=zzKVDGeUM3Yv3-zNah4D6mSnr6jXsstNuLmzczcPQEE,3683
-datahub/ingestion/source/powerbi/m_query/parser.py,sha256=YZXE6L1SYsfoT_sXolrVOfmWQXZOaoN9_KQiA1cdGAQ,5626
-datahub/ingestion/source/powerbi/m_query/resolver.py,sha256=7lLdz8u3Xc2RIoNCwfwCJMj48naEQdNcbxPt-Vmt45w,50429
+datahub/ingestion/source/powerbi/m_query/parser.py,sha256=pB1LGdb02Ryf8Pr8JPSgiOmLE6mEAgDbKodtKOOY6LU,5782
+datahub/ingestion/source/powerbi/m_query/pattern_handler.py,sha256=K9kyX-pMFjOOCAvedKXXce7xG-cHwTGQTrBT5GFXEIg,32865
+datahub/ingestion/source/powerbi/m_query/resolver.py,sha256=v2DJYT70Vtw9NyIwFecP7tHGcVxBWH9UqV7vbpYPhws,16985
 datahub/ingestion/source/powerbi/m_query/tree_function.py,sha256=i7HM9Oyj9XdJpNfG2lE8puxeuc47aSJ-5dPdBEcw2WU,6165
-datahub/ingestion/source/powerbi/m_query/validator.py,sha256=FFAwBl_WmV1SSrYhSbyq8FUsTD0vd0ncdysmkgdKxj8,1008
+datahub/ingestion/source/powerbi/m_query/validator.py,sha256=crG-VZy2XPieiDliP9yVMgiFcc8b2xbZyDFEATXqEAQ,1155
 datahub/ingestion/source/powerbi/rest_api_wrapper/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/powerbi/rest_api_wrapper/data_classes.py,sha256=xqAsnNUCP44Wd1rE1m_phbKtNCMJTFJfOX4_2varadg,8298
 datahub/ingestion/source/powerbi/rest_api_wrapper/data_resolver.py,sha256=O2XTVBdXteIgQF8Lss_t2RhRSsRMmMyWrAoNonDMQFI,39604
@@ -971,8 +972,8 @@ datahub_provider/operators/datahub_assertion_operator.py,sha256=uvTQ-jk2F0sbqqxp
 datahub_provider/operators/datahub_assertion_sensor.py,sha256=lCBj_3x1cf5GMNpHdfkpHuyHfVxsm6ff5x2Z5iizcAo,140
 datahub_provider/operators/datahub_operation_operator.py,sha256=aevDp2FzX7FxGlXrR0khoHNbxbhKR2qPEX5e8O2Jyzw,174
 datahub_provider/operators/datahub_operation_sensor.py,sha256=8fcdVBCEPgqy1etTXgLoiHoJrRt_nzFZQMdSzHqSG7M,168
-acryl_datahub-0.15.0rc3.dist-info/METADATA,sha256=l71CWZc9fD8L_4p0NFRk9ruU6w4QEZgGJq3RObW_ljA,171129
-acryl_datahub-0.15.0rc3.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-acryl_datahub-0.15.0rc3.dist-info/entry_points.txt,sha256=VcQx0dnqaYLyeY_L5OaX7bLmmE-Il7TAXkxCKvEn2bA,9432
-acryl_datahub-0.15.0rc3.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-0.15.0rc3.dist-info/RECORD,,
+acryl_datahub-0.15.0rc5.dist-info/METADATA,sha256=SiiSRUUBz-MJZHnnGpn6342w8hkW9COEktKjlJDQQuw,171117
+acryl_datahub-0.15.0rc5.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+acryl_datahub-0.15.0rc5.dist-info/entry_points.txt,sha256=3jOfMXB66r8zRDaqzRYpNc0tK-oUO-3tXlnGYDdVAmg,9440
+acryl_datahub-0.15.0rc5.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-0.15.0rc5.dist-info/RECORD,,

{acryl_datahub-0.15.0rc3.dist-info → acryl_datahub-0.15.0rc5.dist-info}/entry_points.txt

@@ -73,7 +73,7 @@ okta = datahub.ingestion.source.identity.okta:OktaSource
 openapi = datahub.ingestion.source.openapi:OpenApiSource
 oracle = datahub.ingestion.source.sql.oracle:OracleSource
 postgres = datahub.ingestion.source.sql.postgres:PostgresSource
-powerbi = datahub.ingestion.source.powerbi:PowerBiDashboardSource
+powerbi = datahub.ingestion.source.powerbi.powerbi:PowerBiDashboardSource
 powerbi-report-server = datahub.ingestion.source.powerbi_report_server:PowerBiReportServerDashboardSource
 preset = datahub.ingestion.source.preset:PresetSource
 presto = datahub.ingestion.source.sql.presto:PrestoSource

datahub/__init__.py

@@ -3,7 +3,7 @@ import warnings
 
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "0.15.0rc3"
+__version__ = "0.15.0rc5"
 
 
 def is_dev_mode() -> bool:

datahub/ingestion/source/feast.py

@@ -42,10 +42,14 @@ from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import (
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
 from datahub.metadata.schema_classes import (
     BrowsePathsClass,
+    GlobalTagsClass,
     MLFeaturePropertiesClass,
     MLFeatureTablePropertiesClass,
     MLPrimaryKeyPropertiesClass,
+    OwnerClass,
+    OwnershipClass,
     StatusClass,
+    TagAssociationClass,
 )
 
 # FIXME: ValueType module cannot be used as a type
@@ -91,6 +95,24 @@ class FeastRepositorySourceConfig(ConfigModel):
     environment: str = Field(
         default=DEFAULT_ENV, description="Environment to use when constructing URNs"
     )
+    # owner_mappings example:
+    # This must be added to the recipe in order to extract owners, otherwise NO owners will be extracted
+    # owner_mappings:
+    #   - feast_owner_name: "<owner>"
+    #     datahub_owner_urn: "urn:li:corpGroup:<owner>"
+    #     datahub_ownership_type: "BUSINESS_OWNER"
+    owner_mappings: Optional[List[Dict[str, str]]] = Field(
+        default=None, description="Mapping of owner names to owner types"
+    )
+    enable_owner_extraction: bool = Field(
+        default=False,
+        description="If this is disabled, then we NEVER try to map owners. "
+        "If this is enabled, then owner_mappings is REQUIRED to extract ownership.",
+    )
+    enable_tag_extraction: bool = Field(
+        default=False,
+        description="If this is disabled, then we NEVER try to extract tags.",
+    )
 
 
 @platform_name("Feast")
@@ -215,10 +237,15 @@ class FeastRepositorySource(Source):
         """
 
         feature_view_name = f"{self.feature_store.project}.{feature_view.name}"
+        aspects = (
+            [StatusClass(removed=False)]
+            + self._get_tags(entity)
+            + self._get_owners(entity)
+        )
 
         entity_snapshot = MLPrimaryKeySnapshot(
             urn=builder.make_ml_primary_key_urn(feature_view_name, entity.name),
-            aspects=[StatusClass(removed=False)],
+            aspects=aspects,
         )
 
         entity_snapshot.aspects.append(
@@ -243,10 +270,11 @@ class FeastRepositorySource(Source):
         Generate an MLFeature work unit for a Feast feature.
         """
         feature_view_name = f"{self.feature_store.project}.{feature_view.name}"
+        aspects = [StatusClass(removed=False)] + self._get_tags(field)
 
         feature_snapshot = MLFeatureSnapshot(
             urn=builder.make_ml_feature_urn(feature_view_name, field.name),
-            aspects=[StatusClass(removed=False)],
+            aspects=aspects,
         )
 
         feature_sources = []
@@ -295,13 +323,18 @@ class FeastRepositorySource(Source):
         """
 
         feature_view_name = f"{self.feature_store.project}.{feature_view.name}"
+        aspects = (
+            [
+                BrowsePathsClass(paths=[f"/feast/{self.feature_store.project}"]),
+                StatusClass(removed=False),
+            ]
+            + self._get_tags(feature_view)
+            + self._get_owners(feature_view)
+        )
 
         feature_view_snapshot = MLFeatureTableSnapshot(
             urn=builder.make_ml_feature_table_urn("feast", feature_view_name),
-            aspects=[
-                BrowsePathsClass(paths=[f"/feast/{self.feature_store.project}"]),
-                StatusClass(removed=False),
-            ],
+            aspects=aspects,
         )
 
         feature_view_snapshot.aspects.append(
@@ -360,6 +393,64 @@ class FeastRepositorySource(Source):
 
         return MetadataWorkUnit(id=on_demand_feature_view_name, mce=mce)
 
+    # If a tag is specified in a Feast object, then the tag will be ingested into Datahub if enable_tag_extraction is
+    # True, otherwise NO tags will be ingested
+    def _get_tags(self, obj: Union[Entity, FeatureView, FeastField]) -> list:
+        """
+        Extracts tags from the given object and returns a list of aspects.
+        """
+        aspects: List[Union[GlobalTagsClass]] = []
+
+        # Extract tags
+        if self.source_config.enable_tag_extraction:
+            if obj.tags.get("name"):
+                tag_name: str = obj.tags["name"]
+                tag_association = TagAssociationClass(
+                    tag=builder.make_tag_urn(tag_name)
+                )
+                global_tags_aspect = GlobalTagsClass(tags=[tag_association])
+                aspects.append(global_tags_aspect)
+
+        return aspects
+
+    # If an owner is specified in a Feast object, it will only be ingested into Datahub if owner_mappings is specified
+    # and enable_owner_extraction is True in FeastRepositorySourceConfig, otherwise NO owners will be ingested
+    def _get_owners(self, obj: Union[Entity, FeatureView, FeastField]) -> list:
+        """
+        Extracts owners from the given object and returns a list of aspects.
+        """
+        aspects: List[Union[OwnershipClass]] = []
+
+        # Extract owner
+        if self.source_config.enable_owner_extraction:
+            owner = getattr(obj, "owner", None)
+            if owner:
+                # Create owner association, skipping if None
+                owner_association = self._create_owner_association(owner)
+                if owner_association:  # Only add valid owner associations
+                    owners_aspect = OwnershipClass(owners=[owner_association])
+                    aspects.append(owners_aspect)
+
+        return aspects
+
+    def _create_owner_association(self, owner: str) -> Optional[OwnerClass]:
+        """
+        Create an OwnerClass instance for the given owner using the owner mappings.
+        """
+        if self.source_config.owner_mappings is not None:
+            for mapping in self.source_config.owner_mappings:
+                if mapping["feast_owner_name"] == owner:
+                    ownership_type_class: str = mapping.get(
+                        "datahub_ownership_type", "TECHNICAL_OWNER"
+                    )
+                    datahub_owner_urn = mapping.get("datahub_owner_urn")
+                    if datahub_owner_urn:
+                        return OwnerClass(
+                            owner=datahub_owner_urn,
+                            type=ownership_type_class,
+                        )
+        return None
+
     @classmethod
     def create(cls, config_dict, ctx):
         config = FeastRepositorySourceConfig.parse_obj(config_dict)

datahub/ingestion/source/iceberg/iceberg.py

@@ -9,6 +9,7 @@ from pyiceberg.exceptions import (
     NoSuchIcebergTableError,
     NoSuchNamespaceError,
     NoSuchPropertyException,
+    NoSuchTableError,
 )
 from pyiceberg.schema import Schema, SchemaVisitorPerPrimitiveType, visit
 from pyiceberg.table import Table
@@ -104,7 +105,7 @@ logging.getLogger("azure.core.pipeline.policies.http_logging_policy").setLevel(
 @capability(SourceCapability.DESCRIPTIONS, "Enabled by default.")
 @capability(
     SourceCapability.OWNERSHIP,
-    "Optionally enabled via configuration by specifying which Iceberg table property holds user or group ownership.",
+    "Automatically ingests ownership information from table properties based on `user_ownership_property` and `group_ownership_property`",
 )
 @capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
 class IcebergSource(StatefulIngestionSourceBase):
@@ -192,9 +193,7 @@ class IcebergSource(StatefulIngestionSourceBase):
                     table = thread_local.local_catalog.load_table(dataset_path)
                     time_taken = timer.elapsed_seconds()
                     self.report.report_table_load_time(time_taken)
-                LOGGER.debug(
-                    f"Loaded table: {table.identifier}, time taken: {time_taken}"
-                )
+                LOGGER.debug(f"Loaded table: {table.name()}, time taken: {time_taken}")
                 yield from self._create_iceberg_workunit(dataset_name, table)
             except NoSuchPropertyException as e:
                 self.report.report_warning(
@@ -206,12 +205,20 @@ class IcebergSource(StatefulIngestionSourceBase):
                 )
             except NoSuchIcebergTableError as e:
                 self.report.report_warning(
-                    "no-iceberg-table",
+                    "not-an-iceberg-table",
                     f"Failed to create workunit for {dataset_name}. {e}",
                 )
                 LOGGER.warning(
                     f"NoSuchIcebergTableError while processing table {dataset_path}, skipping it.",
                 )
+            except NoSuchTableError as e:
+                self.report.report_warning(
+                    "no-such-table",
+                    f"Failed to create workunit for {dataset_name}. {e}",
+                )
+                LOGGER.warning(
+                    f"NoSuchTableError while processing table {dataset_path}, skipping it.",
+                )
             except Exception as e:
                 self.report.report_failure("general", f"Failed to create workunit: {e}")
                 LOGGER.exception(

datahub/ingestion/source/kafka/kafka.py

@@ -148,7 +148,7 @@ def get_kafka_consumer(
 ) -> confluent_kafka.Consumer:
     consumer = confluent_kafka.Consumer(
         {
-            "group.id": "test",
+            "group.id": "datahub-kafka-ingestion",
             "bootstrap.servers": connection.bootstrap,
             **connection.consumer_config,
         }
@@ -164,6 +164,25 @@ def get_kafka_consumer(
     return consumer
 
 
+def get_kafka_admin_client(
+    connection: KafkaConsumerConnectionConfig,
+) -> AdminClient:
+    client = AdminClient(
+        {
+            "group.id": "datahub-kafka-ingestion",
+            "bootstrap.servers": connection.bootstrap,
+            **connection.consumer_config,
+        }
+    )
+    if CallableConsumerConfig.is_callable_config(connection.consumer_config):
+        # As per documentation, we need to explicitly call the poll method to make sure OAuth callback gets executed
+        # https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/index.html#kafka-client-configuration
+        logger.debug("Initiating polling for kafka admin client")
+        client.poll(timeout=30)
+        logger.debug("Initiated polling for kafka admin client")
+    return client
+
+
 @dataclass
 class KafkaSourceReport(StaleEntityRemovalSourceReport):
     topics_scanned: int = 0
@@ -278,13 +297,7 @@ class KafkaSource(StatefulIngestionSourceBase, TestableSource):
     def init_kafka_admin_client(self) -> None:
         try:
             # TODO: Do we require separate config than existing consumer_config ?
-            self.admin_client = AdminClient(
-                {
-                    "group.id": "test",
-                    "bootstrap.servers": self.source_config.connection.bootstrap,
-                    **self.source_config.connection.consumer_config,
-                }
-            )
+            self.admin_client = get_kafka_admin_client(self.source_config.connection)
         except Exception as e:
             logger.debug(e, exc_info=e)
             self.report.report_warning(

datahub/ingestion/source/powerbi/__init__.py

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

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]

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"

datahub/ingestion/source/powerbi/m_query/parser.py

@@ -7,6 +7,7 @@ from typing import Dict, List
 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,