acryl-datahub 0.15.0rc2__py3-none-any.whl → 0.15.0rc4__py3-none-any.whl

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

@@ -1,7 +1,7 @@
 import logging
 from typing import Optional, Tuple
 
-from datahub.ingestion.source.powerbi.m_query import resolver
+import datahub.ingestion.source.powerbi.m_query.data_classes
 
 logger = logging.getLogger(__name__)
 
@@ -14,12 +14,18 @@ def validate_parse_tree(
     :param native_query_enabled: Whether user want to extract lineage from native query
     :return: True or False.
     """
-    function_names = [fun.value for fun in resolver.FunctionName]
+    function_names = [
+        fun.value
+        for fun in datahub.ingestion.source.powerbi.m_query.data_classes.FunctionName
+    ]
     if not any(fun in expression for fun in function_names):
         return False, "DataAccess function is not present in M-Query expression."
 
     if native_query_enabled is False:
-        if resolver.FunctionName.NATIVE_QUERY.value in function_names:
+        if (
+            datahub.ingestion.source.powerbi.m_query.data_classes.FunctionName.NATIVE_QUERY.value
+            in function_names
+        ):
             return (
                 False,
                 "Lineage extraction from native query is disabled. Enable native_query_parsing in recipe",

datahub/ingestion/source/powerbi/powerbi.py

@@ -10,6 +10,7 @@ from typing import Iterable, List, Optional, Tuple, Union
 import more_itertools
 
 import datahub.emitter.mce_builder as builder
+import datahub.ingestion.source.powerbi.m_query.data_classes
 import datahub.ingestion.source.powerbi.rest_api_wrapper.data_classes as powerbi_data_classes
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import ContainerKey, gen_containers
@@ -42,12 +43,13 @@ from datahub.ingestion.source.powerbi.config import (
     Constant,
     PowerBiDashboardSourceConfig,
     PowerBiDashboardSourceReport,
+    SupportedDataPlatform,
 )
 from datahub.ingestion.source.powerbi.dataplatform_instance_resolver import (
     AbstractDataPlatformInstanceResolver,
     create_dataplatform_instance_resolver,
 )
-from datahub.ingestion.source.powerbi.m_query import parser, resolver
+from datahub.ingestion.source.powerbi.m_query import parser
 from datahub.ingestion.source.powerbi.rest_api_wrapper.powerbi_api import PowerBiAPI
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
@@ -182,7 +184,9 @@ class Mapper:
         return [schema_mcp]
 
     def make_fine_grained_lineage_class(
-        self, lineage: resolver.Lineage, dataset_urn: str
+        self,
+        lineage: datahub.ingestion.source.powerbi.m_query.data_classes.Lineage,
+        dataset_urn: str,
     ) -> List[FineGrainedLineage]:
         fine_grained_lineages: List[FineGrainedLineage] = []
 
@@ -234,7 +238,9 @@ class Mapper:
         upstream: List[UpstreamClass] = []
         cll_lineage: List[FineGrainedLineage] = []
 
-        upstream_lineage: List[resolver.Lineage] = parser.get_upstream_tables(
+        upstream_lineage: List[
+            datahub.ingestion.source.powerbi.m_query.data_classes.Lineage
+        ] = parser.get_upstream_tables(
             table=table,
             reporter=self.__reporter,
             platform_instance_resolver=self.__dataplatform_instance_resolver,
@@ -1294,7 +1300,7 @@ class PowerBiDashboardSource(StatefulIngestionSourceBase, TestableSource):
     def validate_dataset_type_mapping(self):
         powerbi_data_platforms: List[str] = [
             data_platform.value.powerbi_data_platform_name
-            for data_platform in resolver.SupportedDataPlatform
+            for data_platform in SupportedDataPlatform
         ]
 
         for key in self.source_config.dataset_type_mapping.keys():
@@ -1481,7 +1487,7 @@ class PowerBiDashboardSource(StatefulIngestionSourceBase, TestableSource):
 
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
         # As modified_workspaces is not idempotent, hence workunit processors are run later for each workspace_id
-        # This will result in creating checkpoint for each workspace_id
+        # This will result in creating a checkpoint for each workspace_id
         if self.source_config.modified_since:
             return []  # Handle these in get_workunits_internal
         else:
@@ -1492,7 +1498,7 @@ class PowerBiDashboardSource(StatefulIngestionSourceBase, TestableSource):
 
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         """
-        Datahub Ingestion framework invoke this method
+        Datahub Ingestion framework invokes this method
         """
         logger.info("PowerBi plugin execution is started")
         # Validate dataset type mapping

datahub/telemetry/telemetry.py

@@ -7,7 +7,7 @@ import sys
 import uuid
 from functools import wraps
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, TypeVar
 
 from mixpanel import Consumer, Mixpanel
 from typing_extensions import ParamSpec
@@ -16,10 +16,12 @@ import datahub as datahub_package
 from datahub.cli.config_utils import DATAHUB_ROOT_FOLDER
 from datahub.cli.env_utils import get_boolean_env_variable
 from datahub.configuration.common import ExceptionWithProps
-from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.schema_classes import _custom_package_path
 from datahub.utilities.perf_timer import PerfTimer
 
+if TYPE_CHECKING:
+    from datahub.ingestion.graph.client import DataHubGraph
+
 logger = logging.getLogger(__name__)
 
 DATAHUB_FOLDER = Path(DATAHUB_ROOT_FOLDER)
@@ -117,7 +119,11 @@ class Telemetry:
     tracking_init: bool = False
     sentry_enabled: bool = False
 
+    context_properties: Dict[str, Any] = {}
+
     def __init__(self):
+        self.context_properties = {}
+
         if SENTRY_DSN:
             self.sentry_enabled = True
             try:
@@ -157,6 +163,9 @@ class Telemetry:
             except Exception as e:
                 logger.debug(f"Error connecting to mixpanel: {e}")
 
+        # Initialize the default properties for all events.
+        self.set_context()
+
     def update_config(self) -> bool:
         """
         Update the config file with the current client ID and enabled status.
@@ -238,18 +247,22 @@ class Telemetry:
 
         return False
 
-    def update_capture_exception_context(
+    def set_context(
         self,
-        server: Optional[DataHubGraph] = None,
+        server: Optional["DataHubGraph"] = None,
         properties: Optional[Dict[str, Any]] = None,
     ) -> None:
+        self.context_properties = {
+            **self._server_props(server),
+            **(properties or {}),
+        }
+
         if self.sentry_enabled:
             from sentry_sdk import set_tag
 
             properties = {
                 **_default_telemetry_properties(),
-                **self._server_props(server),
-                **(properties or {}),
+                **self.context_properties,
             }
 
             for key in properties:
@@ -297,7 +310,6 @@ class Telemetry:
         self,
         event_name: str,
         properties: Optional[Dict[str, Any]] = None,
-        server: Optional[DataHubGraph] = None,
     ) -> None:
         """
         Send a single telemetry event.
@@ -323,14 +335,15 @@ class Telemetry:
 
             properties = {
                 **_default_telemetry_properties(),
-                **self._server_props(server),
+                **self.context_properties,
                 **properties,
             }
             self.mp.track(self.client_id, event_name, properties)
         except Exception as e:
             logger.debug(f"Error reporting telemetry: {e}")
 
-    def _server_props(self, server: Optional[DataHubGraph]) -> Dict[str, str]:
+    @classmethod
+    def _server_props(cls, server: Optional["DataHubGraph"]) -> Dict[str, str]:
         if not server:
             return {
                 "server_type": "n/a",
@@ -435,6 +448,7 @@ def with_telemetry(
                             **call_props,
                             "status": "error",
                             **_error_props(e),
+                            "code": e.code,
                         },
                     )
                 telemetry_instance.capture_exception(e)