acryl-datahub 0.15.0rc25__py3-none-any.whl → 0.15.0.1rc1__py3-none-any.whl

datahub/ingestion/source/looker/looker_source.py

@@ -145,7 +145,9 @@ class LookerDashboardSource(TestableSource, StatefulIngestionSourceBase):
         self.source_config: LookerDashboardSourceConfig = config
         self.reporter: LookerDashboardSourceReport = LookerDashboardSourceReport()
         self.looker_api: LookerAPI = LookerAPI(self.source_config)
-        self.user_registry: LookerUserRegistry = LookerUserRegistry(self.looker_api)
+        self.user_registry: LookerUserRegistry = LookerUserRegistry(
+            self.looker_api, self.reporter
+        )
         self.explore_registry: LookerExploreRegistry = LookerExploreRegistry(
             self.looker_api, self.reporter, self.source_config
         )
@@ -1673,5 +1675,14 @@ class LookerDashboardSource(TestableSource, StatefulIngestionSourceBase):
                 yield usage_mcp.as_workunit()
             self.reporter.report_stage_end("usage_extraction")
 
+        # Dump looker user resource mappings.
+        logger.info("Ingesting looker user resource mapping workunits")
+        self.reporter.report_stage_start("user_resource_extraction")
+        yield from auto_workunit(
+            self.user_registry.to_platform_resource(
+                self.source_config.platform_instance
+            )
+        )
+
     def get_report(self) -> SourceReport:
         return self.reporter

datahub/ingestion/source/mlflow.py

@@ -38,16 +38,30 @@ T = TypeVar("T")
 class MLflowConfig(EnvConfigMixin):
     tracking_uri: Optional[str] = Field(
         default=None,
-        description="Tracking server URI. If not set, an MLflow default tracking_uri is used (local `mlruns/` directory or `MLFLOW_TRACKING_URI` environment variable)",
+        description=(
+            "Tracking server URI. If not set, an MLflow default tracking_uri is used"
+            " (local `mlruns/` directory or `MLFLOW_TRACKING_URI` environment variable)"
+        ),
     )
     registry_uri: Optional[str] = Field(
         default=None,
-        description="Registry server URI. If not set, an MLflow default registry_uri is used (value of tracking_uri or `MLFLOW_REGISTRY_URI` environment variable)",
+        description=(
+            "Registry server URI. If not set, an MLflow default registry_uri is used"
+            " (value of tracking_uri or `MLFLOW_REGISTRY_URI` environment variable)"
+        ),
     )
     model_name_separator: str = Field(
         default="_",
         description="A string which separates model name from its version (e.g. model_1 or model-1)",
     )
+    base_external_url: Optional[str] = Field(
+        default=None,
+        description=(
+            "Base URL to use when constructing external URLs to MLflow."
+            " If not set, tracking_uri is used if it's an HTTP URL."
+            " If neither is set, external URLs are not generated."
+        ),
+    )
 
 
 @dataclass
@@ -279,12 +293,23 @@ class MLflowSource(Source):
         )
         return urn
 
-    def _make_external_url(self, model_version: ModelVersion) -> Union[None, str]:
+    def _get_base_external_url_from_tracking_uri(self) -> Optional[str]:
+        if isinstance(
+            self.client.tracking_uri, str
+        ) and self.client.tracking_uri.startswith("http"):
+            return self.client.tracking_uri
+        else:
+            return None
+
+    def _make_external_url(self, model_version: ModelVersion) -> Optional[str]:
         """
         Generate URL for a Model Version to MLflow UI.
         """
-        base_uri = self.client.tracking_uri
-        if base_uri.startswith("http"):
+        base_uri = (
+            self.config.base_external_url
+            or self._get_base_external_url_from_tracking_uri()
+        )
+        if base_uri:
             return f"{base_uri.rstrip('/')}/#/models/{model_version.name}/versions/{model_version.version}"
         else:
             return None

datahub/ingestion/source/powerbi/config.py

@@ -9,7 +9,7 @@ from pydantic.class_validators import root_validator
 
 import datahub.emitter.mce_builder as builder
 from datahub.configuration.common import AllowDenyPattern, ConfigModel
-from datahub.configuration.source_common import DatasetSourceConfigMixin
+from datahub.configuration.source_common import DatasetSourceConfigMixin, PlatformDetail
 from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated
 from datahub.ingestion.source.common.subtypes import BIAssetSubTypes
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
@@ -232,19 +232,6 @@ def default_for_dataset_type_mapping() -> Dict[str, str]:
     return dict_
 
 
-class PlatformDetail(ConfigModel):
-    platform_instance: Optional[str] = pydantic.Field(
-        default=None,
-        description="DataHub platform instance name. To generate correct urn for upstream dataset, this should match "
-        "with platform instance name used in ingestion "
-        "recipe of other datahub sources.",
-    )
-    env: str = pydantic.Field(
-        default=builder.DEFAULT_ENV,
-        description="The environment that all assets produced by DataHub platform ingestion source belong to",
-    )
-
-
 class DataBricksPlatformDetail(PlatformDetail):
     """
     metastore is an additional field used in Databricks connector to generate the dataset urn

datahub/ingestion/source/powerbi/dataplatform_instance_resolver.py

@@ -2,8 +2,8 @@ import logging
 from abc import ABC, abstractmethod
 from typing import Union
 
+from datahub.configuration.source_common import PlatformDetail
 from datahub.ingestion.source.powerbi.config import (
-    PlatformDetail,
     PowerBiDashboardSourceConfig,
     PowerBIPlatformDetail,
 )

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

@@ -5,13 +5,13 @@ from typing import Dict, List, Optional, Tuple, Type, cast
 
 from lark import Tree
 
+from datahub.configuration.source_common import PlatformDetail
 from datahub.emitter import mce_builder as builder
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.source.powerbi.config import (
     Constant,
     DataBricksPlatformDetail,
     DataPlatformPair,
-    PlatformDetail,
     PowerBiDashboardSourceConfig,
     PowerBiDashboardSourceReport,
     PowerBIPlatformDetail,

datahub/ingestion/source/snowflake/snowflake_v2.py

@@ -540,6 +540,7 @@ class SnowflakeV2Source(
                 identifiers=self.identifiers,
                 schema_resolver=schema_resolver,
                 discovered_tables=discovered_datasets,
+                graph=self.ctx.graph,
             )
 
             # TODO: This is slightly suboptimal because we create two SqlParsingAggregator instances with different configs

datahub/ingestion/source/sql/mssql/job_models.py

@@ -1,11 +1,17 @@
 from dataclasses import dataclass, field
 from typing import Dict, List, Optional, Union
 
-from datahub.emitter.mce_builder import make_data_flow_urn, make_data_job_urn
+from datahub.emitter.mce_builder import (
+    make_data_flow_urn,
+    make_data_job_urn,
+    make_data_platform_urn,
+    make_dataplatform_instance_urn,
+)
 from datahub.metadata.schema_classes import (
     DataFlowInfoClass,
     DataJobInfoClass,
     DataJobInputOutputClass,
+    DataPlatformInstanceClass,
 )
 
 
@@ -204,6 +210,18 @@ class MSSQLDataJob:
             status=self.status,
         )
 
+    @property
+    def as_maybe_platform_instance_aspect(self) -> Optional[DataPlatformInstanceClass]:
+        if self.entity.flow.platform_instance:
+            return DataPlatformInstanceClass(
+                platform=make_data_platform_urn(self.entity.flow.orchestrator),
+                instance=make_dataplatform_instance_urn(
+                    platform=self.entity.flow.orchestrator,
+                    instance=self.entity.flow.platform_instance,
+                ),
+            )
+        return None
+
 
 @dataclass
 class MSSQLDataFlow:
@@ -238,3 +256,14 @@ class MSSQLDataFlow:
             customProperties=self.flow_properties,
             externalUrl=self.external_url,
         )
+
+    @property
+    def as_maybe_platform_instance_aspect(self) -> Optional[DataPlatformInstanceClass]:
+        if self.entity.platform_instance:
+            return DataPlatformInstanceClass(
+                platform=make_data_platform_urn(self.entity.orchestrator),
+                instance=make_dataplatform_instance_urn(
+                    self.entity.orchestrator, self.entity.platform_instance
+                ),
+            )
+        return None

datahub/ingestion/source/sql/mssql/source.py

@@ -639,6 +639,13 @@ class SQLServerSource(SQLAlchemySource):
             aspect=data_job.as_datajob_info_aspect,
         ).as_workunit()
 
+        data_platform_instance_aspect = data_job.as_maybe_platform_instance_aspect
+        if data_platform_instance_aspect:
+            yield MetadataChangeProposalWrapper(
+                entityUrn=data_job.urn,
+                aspect=data_platform_instance_aspect,
+            ).as_workunit()
+
         if include_lineage:
             yield MetadataChangeProposalWrapper(
                 entityUrn=data_job.urn,
@@ -654,6 +661,13 @@ class SQLServerSource(SQLAlchemySource):
             entityUrn=data_flow.urn,
             aspect=data_flow.as_dataflow_info_aspect,
         ).as_workunit()
+
+        data_platform_instance_aspect = data_flow.as_maybe_platform_instance_aspect
+        if data_platform_instance_aspect:
+            yield MetadataChangeProposalWrapper(
+                entityUrn=data_flow.urn,
+                aspect=data_platform_instance_aspect,
+            ).as_workunit()
         # TODO: Add SubType when it appear
 
     def get_inspectors(self) -> Iterable[Inspector]:

datahub/ingestion/source/tableau/tableau.py

@@ -645,7 +645,7 @@ def report_user_role(report: TableauSourceReport, server: Server) -> None:
         # the site-role might be different on another site
         logged_in_user: UserInfo = UserInfo.from_server(server=server)
 
-        if not logged_in_user.is_site_administrator_explorer():
+        if not logged_in_user.has_site_administrator_explorer_privileges():
             report.warning(
                 title=title,
                 message=message,
@@ -896,10 +896,9 @@ class TableauSiteSource:
         return f"/{self.config.env.lower()}{self.no_env_browse_prefix}"
 
     def _re_authenticate(self):
-        tableau_auth: Union[
-            TableauAuth, PersonalAccessTokenAuth
-        ] = self.config.get_tableau_auth(self.site_id)
-        self.server.auth.sign_in(tableau_auth)
+        # Sign-in again may not be enough because Tableau sometimes caches invalid sessions
+        # so we need to recreate the Tableau Server object
+        self.server = self.config.make_tableau_client(self.site_id)
 
     @property
     def site_content_url(self) -> Optional[str]:

datahub/ingestion/source/tableau/tableau_constant.py

@@ -82,4 +82,6 @@ PROJECT = "Project"
 SITE = "Site"
 IS_UNSUPPORTED_CUSTOM_SQL = "isUnsupportedCustomSql"
 SITE_PERMISSION = "sitePermission"
-SITE_ROLE = "SiteAdministratorExplorer"
+ROLE_SITE_ADMIN_EXPLORER = "SiteAdministratorExplorer"
+ROLE_SITE_ADMIN_CREATOR = "SiteAdministratorCreator"
+ROLE_SERVER_ADMIN = "ServerAdministrator"

datahub/ingestion/source/tableau/tableau_server_wrapper.py

@@ -11,8 +11,12 @@ class UserInfo:
     site_role: str
     site_id: str
 
-    def is_site_administrator_explorer(self):
-        return self.site_role == c.SITE_ROLE
+    def has_site_administrator_explorer_privileges(self):
+        return self.site_role in [
+            c.ROLE_SITE_ADMIN_EXPLORER,
+            c.ROLE_SITE_ADMIN_CREATOR,
+            c.ROLE_SERVER_ADMIN,
+        ]
 
     @staticmethod
     def from_server(server: Server) -> "UserInfo":

datahub/ingestion/source/tableau/tableau_validation.py

@@ -28,7 +28,7 @@ def check_user_role(
 
     try:
         # TODO: Add check for `Enable Derived Permissions`
-        if not logged_in_user.is_site_administrator_explorer():
+        if not logged_in_user.has_site_administrator_explorer_privileges():
             capability_dict[c.SITE_PERMISSION] = CapabilityReport(
                 capable=False,
                 failure_reason=f"{failure_reason} Their current role is {logged_in_user.site_role}.",

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -490,7 +490,7 @@ class SqlParsingAggregator(Closeable):
             self._exit_stack.push(self._query_usage_counts)
 
         # Tool Extractor
-        self._tool_meta_extractor = ToolMetaExtractor()
+        self._tool_meta_extractor = ToolMetaExtractor.create(graph)
         self.report.tool_meta_report = self._tool_meta_extractor.report
 
     def close(self) -> None:

datahub/sql_parsing/tool_meta_extractor.py

@@ -1,3 +1,4 @@
+import contextlib
 import json
 import logging
 from dataclasses import dataclass, field
@@ -5,8 +6,15 @@ from typing import Callable, Dict, List, Optional, Tuple, Union
 
 from typing_extensions import Protocol
 
+from datahub.api.entities.platformresource.platform_resource import (
+    ElasticPlatformResourceQuery,
+    PlatformResource,
+    PlatformResourceSearchFields,
+)
 from datahub.ingestion.api.report import Report
+from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.urns import CorpGroupUrn, CorpUserUrn
+from datahub.utilities.search_utils import LogicalOperator
 from datahub.utilities.stats_collections import int_top_k_dict
 
 UrnStr = str
@@ -31,6 +39,7 @@ def _get_last_line(query: str) -> str:
 @dataclass
 class ToolMetaExtractorReport(Report):
     num_queries_meta_extracted: Dict[str, int] = field(default_factory=int_top_k_dict)
+    failures: List[str] = field(default_factory=list)
 
 
 class ToolMetaExtractor:
@@ -42,14 +51,81 @@ class ToolMetaExtractor:
     by warehouse query logs.
     """
 
-    def __init__(self) -> None:
-        self.report = ToolMetaExtractorReport()
+    def __init__(
+        self,
+        report: ToolMetaExtractorReport,
+        looker_user_mapping: Optional[Dict[str, str]] = None,
+    ) -> None:
+        self.report = report
         self.known_tool_extractors: List[Tuple[str, Callable[[QueryLog], bool]]] = [
             (
                 "mode",
                 self._extract_mode_query,
-            )
+            ),
+            (
+                "looker",
+                self._extract_looker_query,
+            ),
         ]
+        # maps user id (as string) to email address
+        self.looker_user_mapping = looker_user_mapping
+
+    @classmethod
+    def create(
+        cls,
+        graph: Optional[DataHubGraph] = None,
+    ) -> "ToolMetaExtractor":
+        report = ToolMetaExtractorReport()
+        looker_user_mapping = None
+        if graph:
+            try:
+                looker_user_mapping = cls.extract_looker_user_mapping_from_graph(
+                    graph, report
+                )
+            except Exception as e:
+                report.failures.append(
+                    f"Unexpected error during Looker user metadata extraction: {str(e)}"
+                )
+
+        return cls(report, looker_user_mapping)
+
+    @classmethod
+    def extract_looker_user_mapping_from_graph(
+        cls, graph: DataHubGraph, report: ToolMetaExtractorReport
+    ) -> Optional[Dict[str, str]]:
+        looker_user_mapping = None
+        query = (
+            ElasticPlatformResourceQuery.create_from()
+            .group(LogicalOperator.AND)
+            .add_field_match(PlatformResourceSearchFields.PLATFORM, "looker")
+            .add_field_match(
+                PlatformResourceSearchFields.RESOURCE_TYPE,
+                "USER_ID_MAPPING",
+            )
+            .end()
+        )
+        platform_resources = list(
+            PlatformResource.search_by_filters(query=query, graph_client=graph)
+        )
+
+        if len(platform_resources) > 1:
+            report.failures.append(
+                "Looker user metadata extraction failed. Found more than one looker user id mappings."
+            )
+        else:
+            platform_resource = platform_resources[0]
+
+            if (
+                platform_resource
+                and platform_resource.resource_info
+                and platform_resource.resource_info.value
+            ):
+                with contextlib.suppress(ValueError, AssertionError):
+                    value = platform_resource.resource_info.value.as_raw_json()
+                    if value:
+                        looker_user_mapping = value
+
+        return looker_user_mapping
 
     def _extract_mode_query(self, entry: QueryLog) -> bool:
         """
@@ -78,14 +154,49 @@ class ToolMetaExtractor:
 
         return True
 
+    def _extract_looker_query(self, entry: QueryLog) -> bool:
+        """
+        Returns:
+            bool: whether QueryLog entry is that of looker and looker user info
+            is extracted into entry.
+        """
+        if not self.looker_user_mapping:
+            return False
+
+        last_line = _get_last_line(entry.query_text)
+
+        if not (last_line.startswith("--") and "Looker Query Context" in last_line):
+            return False
+
+        start_quote_idx = last_line.index("'")
+        end_quote_idx = last_line.rindex("'")
+        if start_quote_idx == -1 or end_quote_idx == -1:
+            return False
+
+        looker_json_raw = last_line[start_quote_idx + 1 : end_quote_idx]
+        looker_json = json.loads(looker_json_raw)
+
+        user_id = str(looker_json["user_id"])
+        email = self.looker_user_mapping.get(user_id)
+        if not email:
+            return False
+
+        original_user = entry.user
+
+        entry.user = email_to_user_urn(email)
+        entry.extra_info = entry.extra_info or {}
+        entry.extra_info["user_via"] = original_user
+
+        return True
+
     def extract_bi_metadata(self, entry: QueryLog) -> bool:
         for tool, meta_extractor in self.known_tool_extractors:
             try:
                 if meta_extractor(entry):
                     self.report.num_queries_meta_extracted[tool] += 1
                     return True
-            except Exception:
-                logger.debug("Tool metadata extraction failed with error : {e}")
+            except Exception as e:
+                logger.debug(f"Tool metadata extraction failed with error : {e}")
         return False