acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/vertexai/vertexai_config.py

@@ -0,0 +1,29 @@
+from typing import Dict, Optional
+
+from pydantic import Field
+
+from datahub.configuration.source_common import EnvConfigMixin
+from datahub.ingestion.source.common.gcp_credentials_config import GCPCredential
+
+
+class VertexAIConfig(EnvConfigMixin):
+    credential: Optional[GCPCredential] = Field(
+        default=None, description="GCP credential information"
+    )
+    project_id: str = Field(description=("Project ID in Google Cloud Platform"))
+    region: str = Field(
+        description=("Region of your project in Google Cloud Platform"),
+    )
+    bucket_uri: Optional[str] = Field(
+        default=None,
+        description=("Bucket URI used in your project"),
+    )
+    vertexai_url: Optional[str] = Field(
+        default="https://console.cloud.google.com/vertex-ai",
+        description=("VertexUI URI"),
+    )
+
+    def get_credentials(self) -> Optional[Dict[str, str]]:
+        if self.credential:
+            return self.credential.to_dict(self.project_id)
+        return None

datahub/ingestion/source/vertexai/vertexai_result_type_utils.py

@@ -0,0 +1,89 @@
+from typing import Optional, Union
+
+from google.cloud.aiplatform.base import VertexAiResourceNoun
+from google.cloud.aiplatform.jobs import _RunnableJob
+from google.cloud.aiplatform.training_jobs import _TrainingJob
+from google.cloud.aiplatform_v1.types import JobState, PipelineState, PipelineTaskDetail
+
+from datahub.metadata.schema_classes import RunResultTypeClass
+
+
+def get_automl_job_result_type(state: PipelineState) -> Union[str, RunResultTypeClass]:
+    state_mapping = {
+        PipelineState.PIPELINE_STATE_SUCCEEDED: RunResultTypeClass.SUCCESS,
+        PipelineState.PIPELINE_STATE_FAILED: RunResultTypeClass.FAILURE,
+        PipelineState.PIPELINE_STATE_CANCELLED: "Cancelled",
+        PipelineState.PIPELINE_STATE_PAUSED: "Paused",
+        PipelineState.PIPELINE_STATE_QUEUED: "Queued",
+        PipelineState.PIPELINE_STATE_RUNNING: "Running",
+        PipelineState.PIPELINE_STATE_UNSPECIFIED: "Unspecific",
+    }
+
+    return state_mapping.get(state, "UNKNOWN")
+
+
+def get_custom_job_result_type(state: JobState) -> Union[str, RunResultTypeClass]:
+    state_mapping = {
+        JobState.JOB_STATE_SUCCEEDED: RunResultTypeClass.SUCCESS,
+        JobState.JOB_STATE_FAILED: RunResultTypeClass.FAILURE,
+        JobState.JOB_STATE_CANCELLED: "Cancelled",
+        JobState.JOB_STATE_PAUSED: "Paused",
+        JobState.JOB_STATE_QUEUED: "Queued",
+        JobState.JOB_STATE_RUNNING: "Running",
+        JobState.JOB_STATE_CANCELLING: "Cancelling",
+        JobState.JOB_STATE_EXPIRED: "Expired",
+        JobState.JOB_STATE_UPDATING: "Updating",
+    }
+    return state_mapping.get(state, "UNKNOWN")
+
+
+def get_job_result_status(job: VertexAiResourceNoun) -> Union[str, RunResultTypeClass]:
+    if isinstance(job, _TrainingJob) and isinstance(job.state, PipelineState):
+        return get_automl_job_result_type(job.state)
+    elif isinstance(job, _RunnableJob) and isinstance(job.state, JobState):
+        return get_custom_job_result_type(job.state)
+    return "UNKNOWN"
+
+
+def get_execution_result_status(status: int) -> Union[str, RunResultTypeClass]:
+    """
+    State of the execution.
+    STATE_UNSPECIFIED = 0
+    PENDING = 1
+    RUNNING = 2
+    SUCCEEDED = 3
+    FAILED = 4
+    """
+    status_mapping = {
+        0: "STATE_UNSPECIFIED",
+        1: "PENDING",
+        2: "RUNNING",
+        3: RunResultTypeClass.SUCCESS,
+        4: RunResultTypeClass.FAILURE,
+    }
+    return status_mapping.get(status, "UNKNOWN")
+
+
+def get_pipeline_task_result_status(
+    status: Optional[PipelineTaskDetail.State],
+) -> Union[str, RunResultTypeClass]:
+    # TODO: DataProcessInstanceRunResultClass fails with status string except for SUCCESS, FAILURE, SKIPPED,
+    #  which will be fixed in the future
+    status_mapping = {
+        # PipelineTaskDetail.State.STATE_UNSPECIFIED: "STATE_UNSPECIFIED",
+        # PipelineTaskDetail.State.PENDING: "PENDING",
+        # PipelineTaskDetail.State.RUNNING: "RUNNING",
+        # PipelineTaskDetail.State.CANCEL_PENDING: "CANCEL_PENDING",
+        # PipelineTaskDetail.State.CANCELLING: "CANCELLING",
+        # PipelineTaskDetail.State.NOT_TRIGGERED: "NOT_TRIGGERED",
+        PipelineTaskDetail.State.SUCCEEDED: RunResultTypeClass.SUCCESS,
+        PipelineTaskDetail.State.FAILED: RunResultTypeClass.FAILURE,
+        PipelineTaskDetail.State.SKIPPED: RunResultTypeClass.SKIPPED,
+    }
+    if status is None:
+        return "UNKNOWN"
+    return status_mapping.get(status, "UNKNOWN")
+
+
+def is_status_for_run_event_class(status: Union[str, RunResultTypeClass]) -> bool:
+    return status in [RunResultTypeClass.SUCCESS, RunResultTypeClass.FAILURE]

datahub/ingestion/source_config/pulsar.py

@@ -2,6 +2,7 @@ import re
 from typing import Dict, List, Optional, Union
 from urllib.parse import urlparse
 
+import pydantic
 from pydantic import Field, validator
 
 from datahub.configuration.common import AllowDenyPattern
@@ -121,7 +122,8 @@ class PulsarSourceConfig(
             )
         return client_secret
 
-    @validator("web_service_url")
+    @pydantic.field_validator("web_service_url", mode="after")
+    @classmethod
     def web_service_url_scheme_host_port(cls, val: str) -> str:
         # Tokenize the web url
         url = urlparse(val)

datahub/ingestion/source_report/ingestion_stage.py

@@ -1,7 +1,9 @@
 import logging
+from collections import defaultdict
 from contextlib import AbstractContextManager
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
+from enum import Enum
 
 from datahub.utilities.perf_timer import PerfTimer
 from datahub.utilities.stats_collections import TopKDict
@@ -20,31 +22,68 @@ QUERIES_EXTRACTION = "Queries Extraction"
 PROFILING = "Profiling"
 
 
+class IngestionHighStage(Enum):
+    """
+    The high-level stages at the framework level
+    Team to add more stages as needed
+    """
+
+    PROFILING = "Profiling"
+    _UNDEFINED = "Ingestion"
+
+
 @dataclass
 class IngestionStageReport:
+    ingestion_high_stage_seconds: dict[IngestionHighStage, float] = field(
+        default_factory=lambda: defaultdict(float)
+    )
     ingestion_stage_durations: TopKDict[str, float] = field(default_factory=TopKDict)
 
-    def new_stage(self, stage: str) -> "IngestionStageContext":
-        return IngestionStageContext(stage, self)
+    def new_stage(
+        self, stage: str, high_stage: IngestionHighStage = IngestionHighStage._UNDEFINED
+    ) -> "IngestionStageContext":
+        return IngestionStageContext(stage, self, high_stage)
+
+    def new_high_stage(self, stage: IngestionHighStage) -> "IngestionStageContext":
+        return IngestionStageContext("", self, stage)
 
 
 @dataclass
 class IngestionStageContext(AbstractContextManager):
-    def __init__(self, stage: str, report: IngestionStageReport):
-        self._ingestion_stage = f"{stage} at {datetime.now(timezone.utc)}"
+    def __init__(
+        self,
+        stage: str,
+        report: IngestionStageReport,
+        high_stage: IngestionHighStage = IngestionHighStage._UNDEFINED,
+    ):
+        self._high_stage = high_stage
+        self._ingestion_stage = (
+            f"{stage} at {datetime.now(timezone.utc)}" if stage else ""
+        )
         self._timer: PerfTimer = PerfTimer()
         self._report = report
 
     def __enter__(self) -> "IngestionStageContext":
-        logger.info(f"Stage started: {self._ingestion_stage}")
+        if self._ingestion_stage:
+            logger.info(f"Stage started: {self._ingestion_stage}")
+        else:
+            logger.info(f"High stage started: {self._high_stage.value}")
         self._timer.start()
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         elapsed = self._timer.elapsed_seconds(digits=2)
-        logger.info(
-            f"Time spent in stage <{self._ingestion_stage}>: {elapsed} seconds",
-            stacklevel=2,
-        )
-        self._report.ingestion_stage_durations[self._ingestion_stage] = elapsed
-        return None
+        if self._ingestion_stage:
+            logger.info(
+                f"Time spent in stage <{self._ingestion_stage}>: {elapsed} seconds",
+                stacklevel=2,
+            )
+            # Store tuple as string to avoid serialization errors
+            key = f"({self._high_stage.value}, {self._ingestion_stage})"
+            self._report.ingestion_stage_durations[key] = elapsed
+        else:
+            logger.info(
+                f"Time spent in stage <{self._high_stage.value}>: {elapsed} seconds",
+                stacklevel=2,
+            )
+        self._report.ingestion_high_stage_seconds[self._high_stage] += elapsed

datahub/ingestion/transformer/add_dataset_dataproduct.py

@@ -54,7 +54,7 @@ class AddDatasetDataProduct(DatasetDataproductTransformer):
         data_products_container: Dict[str, DataProductPatchBuilder] = {}
         logger.debug("Generating dataproducts")
         is_container = self.config.is_container
-        for entity_urn in self.entity_map.keys():
+        for entity_urn in self.entity_map:
             data_product_urn = self.config.get_data_product_to_add(entity_urn)
             if data_product_urn:
                 if data_product_urn not in data_products:

datahub/ingestion/transformer/add_dataset_ownership.py

@@ -71,8 +71,24 @@ class AddDatasetOwnership(OwnershipTransformer):
 
         server_ownership = graph.get_ownership(entity_urn=urn)
         if server_ownership:
-            owners = {owner.owner: owner for owner in server_ownership.owners}
-            owners.update({owner.owner: owner for owner in mce_ownership.owners})
+            owners = {
+                (
+                    owner.owner,
+                    owner.type,
+                    owner.typeUrn,
+                ): owner
+                for owner in server_ownership.owners
+            }
+            owners.update(
+                {
+                    (
+                        owner.owner,
+                        owner.type,
+                        owner.typeUrn,
+                    ): owner
+                    for owner in mce_ownership.owners
+                }
+            )
             mce_ownership.owners = list(owners.values())
 
         return mce_ownership
@@ -86,7 +102,7 @@ class AddDatasetOwnership(OwnershipTransformer):
         logger.debug("Generating Ownership for containers")
         ownership_container_mapping: Dict[str, List[OwnerClass]] = {}
         for entity_urn, data_ownerships in (
-            (urn, self.config.get_owners_to_add(urn)) for urn in self.entity_map.keys()
+            (urn, self.config.get_owners_to_add(urn)) for urn in self.entity_map
         ):
             if not data_ownerships:
                 continue

datahub/ingestion/transformer/base_transformer.py

@@ -281,11 +281,14 @@ class BaseTransformer(Transformer, metaclass=ABCMeta):
                                 )
                             )
 
-                            record_metadata = _update_work_unit_id(
-                                envelope=envelope,
-                                aspect_name=mcp.aspect.get_aspect_name(),  # type: ignore
-                                urn=mcp.entityUrn,
-                            )
+                            if mcp.entityUrn:
+                                record_metadata = _update_work_unit_id(
+                                    envelope=envelope,
+                                    aspect_name=mcp.aspect.get_aspect_name(),  # type: ignore
+                                    urn=mcp.entityUrn,
+                                )
+                            else:
+                                record_metadata = envelope.metadata.copy()
 
                             yield RecordEnvelope(
                                 record=mcp,

datahub/ingestion/transformer/dataset_domain.py

@@ -125,7 +125,7 @@ class AddDatasetDomain(DatasetDomainTransformer):
             return domain_mcps
 
         for entity_urn, domain_to_add in (
-            (urn, self.config.get_domains_to_add(urn)) for urn in self.entity_map.keys()
+            (urn, self.config.get_domains_to_add(urn)) for urn in self.entity_map
         ):
             if not domain_to_add or not domain_to_add.domains:
                 continue

datahub/ingestion/transformer/set_browse_path.py

@@ -0,0 +1,112 @@
+import re
+from collections import defaultdict
+from typing import Dict, List, Optional, cast
+
+from datahub.configuration.common import (
+    TransformerSemanticsConfigModel,
+)
+from datahub.emitter.mce_builder import Aspect
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.transformer.base_transformer import (
+    BaseTransformer,
+    SingleAspectTransformer,
+)
+from datahub.metadata.schema_classes import (
+    BrowsePathEntryClass,
+    BrowsePathsV2Class,
+)
+from datahub.utilities.urns.urn import guess_entity_type
+
+
+class SetBrowsePathTransformerConfig(TransformerSemanticsConfigModel):
+    path: List[str]
+
+
+class SetBrowsePathTransformer(BaseTransformer, SingleAspectTransformer):
+    ctx: PipelineContext
+    config: SetBrowsePathTransformerConfig
+
+    def __init__(self, config: SetBrowsePathTransformerConfig, ctx: PipelineContext):
+        super().__init__()
+        self.ctx = ctx
+        self.config = config
+
+    def aspect_name(self) -> str:
+        return "browsePathsV2"
+
+    def entity_types(self) -> List[str]:
+        # This is an arbitrary list, might be adjusted if it makes sense. It might be reasonable to make it configurable
+        return ["dataset", "dataJob", "dataFlow", "chart", "dashboard", "container"]
+
+    @classmethod
+    def create(
+        cls, config_dict: dict, ctx: PipelineContext
+    ) -> "SetBrowsePathTransformer":
+        config = SetBrowsePathTransformerConfig.parse_obj(config_dict)
+        return cls(config, ctx)
+
+    @staticmethod
+    def _build_model(existing_browse_paths: BrowsePathsV2Class) -> Dict[str, List[str]]:
+        template_vars: Dict[str, List[str]] = {}
+        model: Dict[str, List[str]] = defaultdict(list)
+        for entry in existing_browse_paths.path or []:
+            if entry.urn:
+                entity_type = guess_entity_type(entry.urn)
+                model[entity_type].append(entry.urn)
+
+        for entity_type, urns in model.items():
+            template_vars[f"{entity_type}[*]"] = urns
+            for i, urn in enumerate(urns):
+                template_vars[f"{entity_type}[{i}]"] = [urn]
+
+        return template_vars
+
+    @classmethod
+    def _expand_nodes(
+        cls, templates: List[str], template_vars: Dict[str, List[str]]
+    ) -> BrowsePathsV2Class:
+        expanded_nodes: List[str] = []
+        for node in templates:
+            resolved_nodes = cls._resolve_template_to_nodes(node, template_vars)
+            expanded_nodes.extend(resolved_nodes)
+
+        processed_entries: List[BrowsePathEntryClass] = []
+        for node in expanded_nodes:
+            if not node or node.isspace():
+                continue
+            processed_entries.append(
+                BrowsePathEntryClass(
+                    id=node, urn=node if node.startswith("urn:") else None
+                )
+            )
+        return BrowsePathsV2Class(path=processed_entries)
+
+    def transform_aspect(
+        self, entity_urn: str, aspect_name: str, aspect: Optional[Aspect]
+    ) -> Optional[Aspect]:
+        template_vars: Dict[str, List[str]] = {}
+        if aspect is not None:
+            assert isinstance(aspect, BrowsePathsV2Class)
+            template_vars = self._build_model(aspect)
+        new_browse_paths: BrowsePathsV2Class = self._expand_nodes(
+            self.config.path, template_vars
+        )
+        if aspect is not None and not self.config.replace_existing:
+            for node in aspect.path:
+                new_browse_paths.path.append(node)
+
+        return cast(Aspect, new_browse_paths)
+
+    @staticmethod
+    def _resolve_template_to_nodes(
+        template_str: str, template_vars: Dict[str, List[str]]
+    ) -> List[str]:
+        # This mechanism can be made simpler (match against known variables only) or more complex (e.g. by using a
+        # proper templating engine, like jinja).
+        template_str = template_str.strip()
+        var_pattern = re.findall(r"^\$([a-zA-Z]+\[[0-9*]+]$)", template_str)
+
+        if not var_pattern:
+            return [template_str]
+
+        return template_vars.get(var_pattern[0], [])

datahub/integrations/assertion/common.py

@@ -3,6 +3,7 @@ from typing import List, Optional, Tuple, TypedDict
 
 from datahub.api.entities.assertion.assertion import BaseEntityAssertion
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.com.linkedin.pegasus2avro.dataset import DatasetProperties
 from datahub.metadata.com.linkedin.pegasus2avro.schema import SchemaMetadata
 from datahub.utilities.urns.urn import Urn
@@ -15,7 +16,7 @@ class ColumnDict(TypedDict):
 
 @lru_cache
 def get_qualified_name_from_datahub(urn: str) -> Optional[str]:
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         props: Optional[DatasetProperties] = graph.get_aspect(urn, DatasetProperties)
         if props is not None:
             return props.qualifiedName
@@ -24,7 +25,7 @@ def get_qualified_name_from_datahub(urn: str) -> Optional[str]:
 
 @lru_cache
 def get_schema_from_datahub(urn: str) -> Optional[List[ColumnDict]]:
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.INGESTION) as graph:
         schema: Optional[SchemaMetadata] = graph.get_aspect(urn, SchemaMetadata)
         if schema is not None:
             return [

datahub/integrations/assertion/snowflake/compiler.py

@@ -84,9 +84,10 @@ class SnowflakeAssertionCompiler(AssertionCompiler):
 
         dmf_definitions_path = self.output_dir / DMF_DEFINITIONS_FILE_NAME
         dmf_associations_path = self.output_dir / DMF_ASSOCIATIONS_FILE_NAME
-        with (dmf_definitions_path).open("w") as definitions, (
-            dmf_associations_path
-        ).open("w") as associations:
+        with (
+            (dmf_definitions_path).open("w") as definitions,
+            (dmf_associations_path).open("w") as associations,
+        ):
             for assertion_spec in assertion_config_spec.assertions:
                 result.report.num_processed += 1
                 try:

datahub/lite/lite_util.py

@@ -99,7 +99,7 @@ def get_datahub_lite(config_dict: dict, read_only: bool = False) -> "DataHubLite
         lite_class = lite_registry.get(lite_type)
     except KeyError as e:
         raise Exception(
-            f"Failed to find a registered lite implementation for {lite_type}. Valid values are {[k for k in lite_registry.mapping.keys()]}"
+            f"Failed to find a registered lite implementation for {lite_type}. Valid values are {[k for k in lite_registry.mapping]}"
         ) from e
 
     lite_specific_config = lite_class.get_config_class().parse_obj(
@@ -127,7 +127,7 @@ def get_datahub_lite(config_dict: dict, read_only: bool = False) -> "DataHubLite
                 return lite
         else:
             raise Exception(
-                f"Failed to find a registered forwarding sink for type {lite_local_config.forward_to.type}. Valid values are {[k for k in sink_registry.mapping.keys()]}"
+                f"Failed to find a registered forwarding sink for type {lite_local_config.forward_to.type}. Valid values are {[k for k in sink_registry.mapping]}"
             )
     else:
         return lite