acryl-datahub 1.2.0.10rc4__py3-none-any.whl → 1.2.0.10rc5__py3-none-any.whl

datahub/ingestion/source/snaplogic/snaplogic_lineage_extractor.py

@@ -0,0 +1,107 @@
+from datetime import datetime
+from typing import Iterable, Optional, Tuple
+
+import requests
+
+from datahub.ingestion.api.source import (
+    SourceReport,
+)
+from datahub.ingestion.source.snaplogic.snaplogic_config import SnaplogicConfig
+from datahub.ingestion.source.state.redundant_run_skip_handler import (
+    RedundantLineageRunSkipHandler,
+)
+
+
+class SnaplogicLineageExtractor:
+    """
+    A class to interact with the SnapLogic API.
+    """
+
+    def __init__(
+        self,
+        config: SnaplogicConfig,
+        redundant_run_skip_handler: Optional[RedundantLineageRunSkipHandler],
+        report: SourceReport,
+    ):
+        self.config = config
+        self.report = report
+        self.redundant_run_skip_handler = redundant_run_skip_handler
+        self.start_time, self.end_time = self._get_time_window()
+
+    def get_lineages(self) -> Iterable[dict]:
+        """Generator function that yields lineage records one at a time as they are fetched."""
+        page = 0
+        has_more = True
+        records_processed = 0
+
+        try:
+            while has_more:
+                params = {
+                    "format": "OPENLINEAGE",
+                    "start_ts": str(int(self.start_time.timestamp() * 1000)),
+                    "end_ts": str(int(self.end_time.timestamp() * 1000)),
+                    "page": str(page),
+                }
+
+                self.report.info(
+                    message=f"Fetching lineage data - page: {page}, start_ts: {self.start_time}, end_ts: {self.end_time}",
+                    title="Lineage Fetch",
+                )
+                headers = {"User-Agent": "datahub-connector/1.0"}
+                response = requests.get(
+                    url=f"{self.config.base_url}/api/1/rest/public/catalog/{self.config.org_name}/lineage",
+                    params=params,
+                    headers=headers,
+                    auth=(
+                        self.config.username,
+                        self.config.password.get_secret_value(),
+                    ),
+                )
+                response.raise_for_status()
+
+                data = response.json()
+                content = data["content"]
+
+                # Yield records one at a time
+                for record in content:
+                    records_processed += 1
+                    yield record
+
+                # Check if we need to fetch more pages
+                has_more = (
+                    len(content) >= 20
+                )  # If we got full page size, there might be more
+                page += 1
+
+            self.report.info(
+                message=f"Completed fetching lineage data. Total records processed: {records_processed}",
+                title="Lineage Fetch Complete",
+            )
+
+        except Exception as e:
+            self.report.report_failure(
+                message="Error fetching lineage data",
+                exc=e,
+                title="Lineage Fetch Error",
+            )
+            raise
+
+    def _get_time_window(self) -> Tuple[datetime, datetime]:
+        if self.redundant_run_skip_handler:
+            return self.redundant_run_skip_handler.suggest_run_time_window(
+                self.config.start_time, self.config.end_time
+            )
+        else:
+            return self.config.start_time, self.config.end_time
+
+    def update_stats(self):
+        if self.redundant_run_skip_handler:
+            # Update the checkpoint state for this run.
+            self.redundant_run_skip_handler.update_state(
+                self.config.start_time,
+                self.config.end_time,
+            )
+
+    def report_status(self, step: str, status: bool) -> None:
+        if self.redundant_run_skip_handler:
+            self.redundant_run_skip_handler.report_current_run_status(step, status)

datahub/ingestion/source/snaplogic/snaplogic_parser.py

@@ -0,0 +1,168 @@
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+
+@dataclass
+class Dataset:
+    name: str
+    display_name: str
+    fields: List[Dict] = field(default_factory=list)
+    platform: str = "snaplogic"
+    platform_instance: Optional[str] = None
+    type: Optional[str] = None  # INPUT or OUTPUT
+    env: str = "PROD"
+
+
+@dataclass
+class Pipeline:
+    name: str
+    id: str
+    namespace: str
+
+
+@dataclass
+class Task:
+    name: str
+    id: str
+    namespace: str
+
+
+@dataclass
+class ColumnMapping:
+    input_dataset: Dataset
+    output_dataset: Dataset
+    input_field: str
+    output_field: str
+
+
+class SnapLogicParser:
+    def __init__(self, case_insensitive_namespaces: list[str], namespace_mapping: dict):
+        self.case_insensitive_namespaces = case_insensitive_namespaces
+        self.namespace_mapping = namespace_mapping
+        self.platform_mapping = {
+            "sqlserver": "mssql",
+        }
+
+    def _parse_platform(self, namespace: str) -> str:
+        type_part = namespace.split("://")[0] if "://" in namespace else namespace
+
+        return self.platform_mapping.get(type_part.lower(), type_part.lower())
+
+    def extract_task_from_lineage(self, lineage: dict) -> Task:
+        job = lineage.get("job")
+        if not job:
+            raise ValueError("Job information is missing in the lineage data.")
+        name = job.get("name")
+        namespace = job.get("namespace")
+
+        return Task(
+            id=name,
+            name=name.rsplit(":", 1)[0],
+            namespace=self._parse_platform(namespace),
+        )
+
+    def extract_pipeline_from_lineage(self, lineage: dict) -> Pipeline:
+        parent_run = lineage.get("run", {}).get("facets", {}).get("parent", {})
+        job = parent_run.get("job", {})
+        name = job.get("name")
+        namespace = job.get("namespace")
+        pipeline_snode_id = parent_run.get("_producer").split("#pipe_snode=")[1]
+        return Pipeline(
+            id=pipeline_snode_id, name=name, namespace=self._parse_platform(namespace)
+        )
+
+    def _get_case_sensitive_value(self, value: str, namespace: str) -> str:
+        """Transform value to lowercase if namespace is case-insensitive."""
+        return value.lower() if namespace in self.case_insensitive_namespaces else value
+
+    def _create_dataset_info(
+        self,
+        namespace: str,
+        name: str,
+        display_name: str,
+        type: str,
+        fields: Optional[List[Dict]] = None,
+    ) -> Dataset:
+        """Create a Dataset instance with proper case sensitivity."""
+        return Dataset(
+            platform=self._parse_platform(namespace),
+            name=self._get_case_sensitive_value(name, namespace),
+            display_name=display_name or name,
+            fields=fields or [],
+            env="PROD",
+            platform_instance=self.namespace_mapping.get(namespace, None),
+            type=type,
+        )
+
+    def extract_columns_mapping_from_lineage(
+        self, lineage: dict
+    ) -> List[ColumnMapping]:
+        outputs = lineage.get("outputs", [])
+        lineages = []
+
+        for output in outputs:
+            output_namespace = output.get("namespace")
+            output_name = output.get("name", "")
+            column_lineage = (
+                output.get("facets", {}).get("columnLineage", {}).get("fields", {})
+            )
+
+            for field_name, field_dict in column_lineage.items():
+                output_field = self._get_case_sensitive_value(
+                    field_name, output_namespace
+                )
+
+                for input_field in field_dict.get("inputFields", []):
+                    input_namespace = input_field.get("namespace")
+                    input_name = input_field.get("name", "")
+                    input_field_name = input_field.get("field", "")
+
+                    lineages.append(
+                        ColumnMapping(
+                            input_dataset=self._create_dataset_info(
+                                input_namespace, input_name, input_name, "INPUT"
+                            ),
+                            output_dataset=self._create_dataset_info(
+                                output_namespace, output_name, output_name, "OUTPUT"
+                            ),
+                            input_field=self._get_case_sensitive_value(
+                                input_field_name, input_namespace
+                            ),
+                            output_field=output_field,
+                        )
+                    )
+
+        return lineages
+
+    def extract_datasets_from_lineage(self, lineage: dict) -> List[Dataset]:
+        inputs = lineage.get("inputs", {})
+        outputs = lineage.get("outputs", {})
+
+        datasets = []
+        for dataset, dataset_type in [
+            *[(input_dataset, "INPUT") for input_dataset in inputs],
+            *[(output_dataset, "OUTPUT") for output_dataset in outputs],
+        ]:
+            namespace = dataset.get("namespace")
+            name = dataset.get("name", "")
+            fields = dataset.get("facets", {}).get("schema", {}).get("fields", [])
+            display_name = name
+
+            # Transform names to lowercase if namespace is in case_insensitive_namespaces
+            if namespace in self.case_insensitive_namespaces:
+                name = name.lower()
+                fields = [
+                    {**field, "name": field.get("name", "").lower()} for field in fields
+                ]
+
+            datasets.append(
+                self._create_dataset_info(
+                    namespace=namespace,
+                    name=name,
+                    fields=fields,
+                    display_name=display_name,
+                    type=dataset_type,
+                )
+            )
+
+        return datasets

datahub/ingestion/source/snaplogic/snaplogic_utils.py

@@ -0,0 +1,31 @@
+from datahub.metadata.schema_classes import (
+    BooleanTypeClass,
+    NumberTypeClass,
+    SchemaFieldDataTypeClass,
+    StringTypeClass,
+)
+
+
+class SnaplogicUtils:
+    @staticmethod
+    def get_datahub_type(type_str: str) -> SchemaFieldDataTypeClass:
+        """
+        Maps a string-based type to a DataHub SchemaFieldDataTypeClass.
+
+        Args:
+            type_str (str): The input type (e.g., "string", "int", "boolean").
+
+        Returns:
+            SchemaFieldDataTypeClass: The mapped DataHub type.
+        """
+        normalized_type = type_str.lower()
+
+        if normalized_type in ["string", "varchar"]:
+            return SchemaFieldDataTypeClass(type=StringTypeClass())
+        elif normalized_type in ["number", "long", "float", "double", "int"]:
+            return SchemaFieldDataTypeClass(type=NumberTypeClass())
+        elif normalized_type == "boolean":
+            return SchemaFieldDataTypeClass(type=BooleanTypeClass())
+        else:
+            # Default fallback: String
+            return SchemaFieldDataTypeClass(type=StringTypeClass())

datahub/ingestion/source/tableau/tableau.py

@@ -594,13 +594,13 @@ class TableauConfig(
     )
 
     extract_lineage_from_unsupported_custom_sql_queries: bool = Field(
-        default=False,
-        description="[Experimental] Whether to extract lineage from unsupported custom sql queries using SQL parsing",
+        default=True,
+        description="[Experimental] Extract lineage from Custom SQL queries using DataHub's SQL parser in cases where the Tableau Catalog API fails to return lineage for the query.",
     )
 
     force_extraction_of_lineage_from_custom_sql_queries: bool = Field(
         default=False,
-        description="[Experimental] Force extraction of lineage from custom sql queries using SQL parsing, ignoring Tableau metadata",
+        description="[Experimental] Force extraction of lineage from Custom SQL queries using DataHub's SQL parser, even when the Tableau Catalog API returns lineage already.",
     )
 
     sql_parsing_disable_schema_awareness: bool = Field(

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/sdk/_shared.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import warnings
+from abc import ABC, abstractmethod
 from datetime import datetime
 from typing import (
     TYPE_CHECKING,
@@ -61,6 +62,7 @@ DataPlatformInstanceUrnOrStr: TypeAlias = Union[str, DataPlatformInstanceUrn]
 DataPlatformUrnOrStr: TypeAlias = Union[str, DataPlatformUrn]
 
 ActorUrn: TypeAlias = Union[CorpUserUrn, CorpGroupUrn]
+ActorUrnOrStr: TypeAlias = Union[str, ActorUrn]
 StructuredPropertyUrnOrStr: TypeAlias = Union[str, StructuredPropertyUrn]
 StructuredPropertyValueType: TypeAlias = Union[str, float, int]
 StructuredPropertyInputType: TypeAlias = Dict[
@@ -110,6 +112,130 @@ def parse_time_stamp(ts: Optional[models.TimeStampClass]) -> Optional[datetime]:
     return parse_ts_millis(ts.time)
 
 
+class ChangeAuditStampsMixin(ABC):
+    """Mixin class for managing audit stamps on entities."""
+
+    __slots__ = ()
+
+    @abstractmethod
+    def _get_audit_stamps(self) -> models.ChangeAuditStampsClass:
+        """Get the audit stamps from the entity properties."""
+        pass
+
+    @abstractmethod
+    def _set_audit_stamps(self, audit_stamps: models.ChangeAuditStampsClass) -> None:
+        """Set the audit stamps on the entity properties."""
+        pass
+
+    @property
+    def last_modified(self) -> Optional[datetime]:
+        """Get the last modification timestamp from audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        if audit_stamps.lastModified.time == 0:
+            return None
+        return datetime.fromtimestamp(
+            audit_stamps.lastModified.time / 1000
+        )  # supports only seconds precision
+
+    def set_last_modified(self, last_modified: datetime) -> None:
+        """Set the last modification timestamp in audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        audit_stamps.lastModified.time = make_ts_millis(last_modified)
+        self._set_audit_stamps(audit_stamps)
+
+    @property
+    def last_modified_by(self) -> Optional[str]:
+        """Get the last modification actor from audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        if audit_stamps.lastModified.actor == builder.UNKNOWN_USER:
+            return None
+        return audit_stamps.lastModified.actor
+
+    def set_last_modified_by(self, last_modified_by: ActorUrnOrStr) -> None:
+        """Set the last modification actor in audit stamps."""
+        if isinstance(last_modified_by, str):
+            last_modified_by = make_user_urn(last_modified_by)
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        audit_stamps.lastModified.actor = str(last_modified_by)
+        self._set_audit_stamps(audit_stamps)
+
+    @property
+    def created_at(self) -> Optional[datetime]:
+        """Get the creation timestamp from audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        if audit_stamps.created.time == 0:
+            return None
+        return datetime.fromtimestamp(
+            audit_stamps.created.time / 1000
+        )  # supports only seconds precision
+
+    def set_created_at(self, created_at: datetime) -> None:
+        """Set the creation timestamp in audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        audit_stamps.created.time = make_ts_millis(created_at)
+        self._set_audit_stamps(audit_stamps)
+
+    @property
+    def created_by(self) -> Optional[ActorUrnOrStr]:
+        """Get the creation actor from audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        if audit_stamps.created.actor == builder.UNKNOWN_USER:
+            return None
+        return audit_stamps.created.actor
+
+    def set_created_by(self, created_by: ActorUrnOrStr) -> None:
+        """Set the creation actor in audit stamps."""
+        if isinstance(created_by, str):
+            created_by = make_user_urn(created_by)
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        audit_stamps.created.actor = str(created_by)
+        self._set_audit_stamps(audit_stamps)
+
+    @property
+    def deleted_on(self) -> Optional[datetime]:
+        """Get the deletion timestamp from audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        if audit_stamps.deleted is None or audit_stamps.deleted.time == 0:
+            return None
+        return datetime.fromtimestamp(
+            audit_stamps.deleted.time / 1000
+        )  # supports only seconds precision
+
+    def set_deleted_on(self, deleted_on: datetime) -> None:
+        """Set the deletion timestamp in audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        # Default constructor sets deleted to None
+        if audit_stamps.deleted is None:
+            audit_stamps.deleted = models.AuditStampClass(
+                time=0, actor=builder.UNKNOWN_USER
+            )
+        audit_stamps.deleted.time = make_ts_millis(deleted_on)
+        self._set_audit_stamps(audit_stamps)
+
+    @property
+    def deleted_by(self) -> Optional[ActorUrnOrStr]:
+        """Get the deletion actor from audit stamps."""
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        if (
+            audit_stamps.deleted is None
+            or audit_stamps.deleted.actor == builder.UNKNOWN_USER
+        ):
+            return None
+        return audit_stamps.deleted.actor
+
+    def set_deleted_by(self, deleted_by: ActorUrnOrStr) -> None:
+        """Set the deletion actor in audit stamps."""
+        if isinstance(deleted_by, str):
+            deleted_by = make_user_urn(deleted_by)
+        audit_stamps: models.ChangeAuditStampsClass = self._get_audit_stamps()
+        if audit_stamps.deleted is None:
+            audit_stamps.deleted = models.AuditStampClass(
+                time=0, actor=builder.UNKNOWN_USER
+            )
+        audit_stamps.deleted.actor = str(deleted_by)
+        self._set_audit_stamps(audit_stamps)
+
+
 class HasPlatformInstance(Entity):
     __slots__ = ()