acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/snaplogic/snaplogic.py

@@ -0,0 +1,355 @@
+from typing import Iterable, List, Optional
+
+from datahub.emitter.mce_builder import (
+    make_data_flow_urn,
+    make_data_job_urn,
+    make_data_platform_urn,
+    make_dataset_urn_with_platform_instance,
+    make_schema_field_urn,
+)
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.api.decorators import (
+    SupportStatus,
+    capability,
+    config_class,
+    platform_name,
+    support_status,
+)
+from datahub.ingestion.api.source import (
+    MetadataWorkUnitProcessor,
+    SourceCapability,
+    SourceReport,
+)
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.ingestion.source.snaplogic.snaplogic_config import SnaplogicConfig
+from datahub.ingestion.source.snaplogic.snaplogic_lineage_extractor import (
+    SnaplogicLineageExtractor,
+)
+from datahub.ingestion.source.snaplogic.snaplogic_parser import (
+    ColumnMapping,
+    Dataset,
+    SnapLogicParser,
+)
+from datahub.ingestion.source.snaplogic.snaplogic_utils import SnaplogicUtils
+from datahub.ingestion.source.state.redundant_run_skip_handler import (
+    RedundantLineageRunSkipHandler,
+)
+from datahub.ingestion.source.state.stale_entity_removal_handler import (
+    StaleEntityRemovalHandler,
+    StaleEntityRemovalSourceReport,
+)
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+    StatefulIngestionSourceBase,
+)
+from datahub.metadata.schema_classes import (
+    DataFlowInfoClass,
+    DataJobInfoClass,
+    DataJobInputOutputClass,
+    DatasetPropertiesClass,
+    FineGrainedLineageClass,
+    FineGrainedLineageDownstreamTypeClass,
+    OtherSchemaClass,
+    SchemaFieldClass,
+    SchemaMetadataClass,
+)
+
+
+@platform_name("SnapLogic")
+@config_class(SnaplogicConfig)
+@support_status(SupportStatus.TESTING)
+@capability(
+    SourceCapability.PLATFORM_INSTANCE,
+    "SnapLogic does not support platform instances",
+    supported=False,
+)
+@capability(SourceCapability.LINEAGE_COARSE, "Enabled by default")
+@capability(SourceCapability.LINEAGE_FINE, "Enabled by default")
+@capability(SourceCapability.DELETION_DETECTION, "Not supported yet", supported=False)
+class SnaplogicSource(StatefulIngestionSourceBase):
+    """
+    A source plugin for ingesting lineage and metadata from SnapLogic.
+    """
+
+    def __init__(self, config: SnaplogicConfig, ctx: PipelineContext):
+        super().__init__(config, ctx)
+        self.config = config
+        self.report = StaleEntityRemovalSourceReport()
+        self.graph: Optional[DataHubGraph] = ctx.graph
+        self.snaplogic_parser = SnapLogicParser(
+            config.case_insensitive_namespaces, self.config.namespace_mapping
+        )
+        self.redundant_lineage_run_skip_handler: Optional[
+            RedundantLineageRunSkipHandler
+        ] = None
+        if self.config.enable_stateful_lineage_ingestion:
+            self.redundant_lineage_run_skip_handler = RedundantLineageRunSkipHandler(
+                source=self,
+                config=self.config,
+                pipeline_name=ctx.pipeline_name,
+                run_id=ctx.run_id,
+            )
+        self.snaplogic_lineage_extractor = SnaplogicLineageExtractor(
+            config=config,
+            redundant_run_skip_handler=self.redundant_lineage_run_skip_handler,
+            report=self.report,
+        )
+
+    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
+        try:
+            self.report.info(
+                message="Starting lineage ingestion from SnapLogic",
+                title="Lineage Ingestion",
+            )
+
+            records_processed = 0
+            for lineage in self.snaplogic_lineage_extractor.get_lineages():
+                try:
+                    for workunit in self._process_lineage_record(lineage):
+                        yield workunit
+                    records_processed += 1
+
+                    if records_processed % 20 == 0:
+                        self.report.info(
+                            message=f"Processed {records_processed} lineage records",
+                            title="Lineage Ingestion Progress",
+                        )
+                except Exception as e:
+                    self.report.report_failure(
+                        message="Failed to process lineage record",
+                        context=str(lineage),
+                        exc=e,
+                    )
+            self.report.info(
+                message=f"Completed processing {records_processed} lineage records",
+                title="Lineage Ingestion Complete",
+            )
+            self.snaplogic_lineage_extractor.report_status("lineage_ingestion", True)
+            self.snaplogic_lineage_extractor.update_stats()
+        except Exception as e:
+            self.report.report_failure(message="Failed to fetch lineages", exc=e)
+            self.snaplogic_lineage_extractor.report_status("lineage_ingestion", False)
+
+    def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
+        return [
+            *super().get_workunit_processors(),
+            StaleEntityRemovalHandler.create(
+                self, self.config, self.ctx
+            ).workunit_processor,
+        ]
+
+    def _process_lineage_record(self, lineage: dict) -> Iterable[MetadataWorkUnit]:
+        """Process a lineage record to create pipeline and task workunits with relationships."""
+        producer = lineage.get("producer")
+        if not producer:
+            return
+        pipeline_snode_id = producer.split("#pipe_snode=")[1]
+        if not pipeline_snode_id:
+            return
+        datasets = self.snaplogic_parser.extract_datasets_from_lineage(lineage)
+        pipeline = self.snaplogic_parser.extract_pipeline_from_lineage(lineage)
+        task = self.snaplogic_parser.extract_task_from_lineage(lineage)
+        columns_mapping = self.snaplogic_parser.extract_columns_mapping_from_lineage(
+            lineage
+        )
+
+        # Create pipeline MCP
+        for pipeline_workunit in self.create_pipeline_mcp(
+            name=pipeline.name,
+            pipeline_snode_id=pipeline.id,
+            namespace=pipeline.namespace,
+        ):
+            self.report.report_workunit(pipeline_workunit)
+            yield pipeline_workunit
+
+        # Create dataset MCP
+        for dataset in datasets:
+            for dataset_workunit in self.create_dataset_mcp(
+                dataset_name=dataset.name,
+                dataset_display_name=dataset.display_name,
+                fields=dataset.fields,
+                platform=dataset.platform,
+                platform_instance=dataset.platform_instance,
+            ):
+                self.report.report_workunit(dataset_workunit)
+                yield dataset_workunit
+
+        # Create task MCP
+        for task_workunit in self.create_task_mcp(
+            name=task.name,
+            task_id=task.id,
+            namespace=task.namespace,
+            pipeline_snode_id=pipeline_snode_id,
+            input_datasets=[dataset for dataset in datasets if dataset.type == "INPUT"],
+            output_datasets=[
+                dataset for dataset in datasets if dataset.type == "OUTPUT"
+            ],
+            columns_mapping=columns_mapping,
+        ):
+            self.report.report_workunit(task_workunit)
+            yield task_workunit
+
+    def create_task_mcp(
+        self,
+        task_id: str,
+        name: str,
+        namespace: str,
+        pipeline_snode_id: str,
+        input_datasets: list[Dataset],
+        output_datasets: list[Dataset],
+        columns_mapping: list[ColumnMapping],
+    ) -> Iterable[MetadataWorkUnit]:
+        """Create MCPs for a task (snap) including metadata and lineage."""
+        job_urn = make_data_job_urn(
+            orchestrator=namespace,
+            flow_id=pipeline_snode_id,
+            job_id=task_id,
+            cluster="PROD",
+        )
+        yield MetadataChangeProposalWrapper(
+            entityUrn=job_urn,
+            aspect=DataJobInfoClass(
+                name=name,
+                description="",
+                externalUrl=f"{self.config.base_url}/sl/designer.html?v=21818#pipe_snode={pipeline_snode_id}",
+                type="SNAPLOGIC_SNAP",
+            ),
+        ).as_workunit()
+
+        # Helper functions
+        def dataset_urn(d: Dataset) -> str:
+            return make_dataset_urn_with_platform_instance(
+                d.platform, d.name, d.platform_instance
+            )
+
+        def field_urn(d, f):
+            return make_schema_field_urn(dataset_urn(d), f["name"])
+
+        # Emit lineage
+        yield MetadataChangeProposalWrapper(
+            entityUrn=job_urn,
+            aspect=DataJobInputOutputClass(
+                inputDatasets=[dataset_urn(d) for d in input_datasets],
+                outputDatasets=[dataset_urn(d) for d in output_datasets],
+                inputDatasetFields=[
+                    field_urn(d, f) for d in input_datasets for f in d.fields
+                ],
+                outputDatasetFields=[
+                    field_urn(d, f) for d in output_datasets for f in d.fields
+                ],
+                fineGrainedLineages=[
+                    FineGrainedLineageClass(
+                        upstreamType=FineGrainedLineageDownstreamTypeClass.FIELD_SET,
+                        upstreams=[
+                            make_schema_field_urn(
+                                make_dataset_urn_with_platform_instance(
+                                    cl.input_dataset.platform,
+                                    cl.input_dataset.name,
+                                    cl.input_dataset.platform_instance,
+                                    cl.input_dataset.env,
+                                ),
+                                cl.input_field,
+                            )
+                        ],
+                        downstreamType=FineGrainedLineageDownstreamTypeClass.FIELD_SET,
+                        downstreams=[
+                            make_schema_field_urn(
+                                make_dataset_urn_with_platform_instance(
+                                    cl.output_dataset.platform,
+                                    cl.output_dataset.name,
+                                    cl.output_dataset.platform_instance,
+                                    cl.output_dataset.env,
+                                ),
+                                cl.output_field,
+                            )
+                        ],
+                    )
+                    for cl in columns_mapping
+                ],
+            ),
+        ).as_workunit()
+
+    def create_dataset_mcp(
+        self,
+        dataset_name: str,
+        dataset_display_name: str,
+        fields: list[dict],
+        platform: str = "snaplogic",
+        env: str = "PROD",
+        platform_instance: Optional[str] = None,
+    ) -> Iterable[MetadataWorkUnit]:
+        dataset_urn = make_dataset_urn_with_platform_instance(
+            platform=platform,
+            name=dataset_name,
+            env=env,
+            platform_instance=platform_instance,
+        )
+
+        # Skip dataset creation if:
+        # 1. The platform is not "snaplogic" AND
+        # 2. Either:
+        #    a) The config `create_non_snaplogic_datasets` is disabled (False), meaning
+        #       we do not create datasets for non-snaplogic platforms, OR
+        #    b) The dataset already exists in DataHub (`self.graph.exists(dataset_urn)`).
+        if platform != "snaplogic" and (
+            not self.config.create_non_snaplogic_datasets
+            or (self.graph and self.graph.exists(dataset_urn))
+        ):
+            return
+
+        dataset_properties = DatasetPropertiesClass(
+            name=dataset_display_name,
+            qualifiedName=dataset_name,
+        )
+        schema_fields = [
+            SchemaFieldClass(
+                fieldPath=field["name"],
+                type=SnaplogicUtils.get_datahub_type(field.get("type", "Varchar")),
+                nativeDataType=field.get("type", "Varchar"),
+            )
+            for field in fields
+        ]
+        schema_metadata = SchemaMetadataClass(
+            schemaName=dataset_name,
+            platform=make_data_platform_urn(platform),
+            version=0,
+            hash="",
+            platformSchema=OtherSchemaClass(rawSchema=""),
+            fields=schema_fields,
+        )
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=dataset_urn, aspect=dataset_properties
+        ).as_workunit()
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=dataset_urn, aspect=schema_metadata
+        ).as_workunit()
+
+    def create_pipeline_mcp(
+        self, name: str, namespace: str, pipeline_snode_id: str
+    ) -> Iterable[MetadataWorkUnit]:
+        flow_urn = make_data_flow_urn(
+            orchestrator=namespace, flow_id=pipeline_snode_id, cluster="PROD"
+        )
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=flow_urn,
+            aspect=DataFlowInfoClass(
+                name=name,
+                description="",
+                externalUrl=f"{self.config.base_url}/sl/designer.html?v=21818#pipe_snode={pipeline_snode_id}",
+            ),
+        ).as_workunit()
+
+    def get_report(self) -> SourceReport:
+        return self.report
+
+    def close(self) -> None:
+        super().close()
+
+    @classmethod
+    def create(cls, config_dict: dict, ctx: PipelineContext) -> "SnaplogicSource":
+        config = SnaplogicConfig.parse_obj(config_dict)
+        return cls(config, ctx)

datahub/ingestion/source/snaplogic/snaplogic_config.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+from pydantic import Field, SecretStr
+
+from datahub.ingestion.source.state.stale_entity_removal_handler import (
+    StatefulIngestionConfigBase,
+    StatefulStaleMetadataRemovalConfig,
+)
+from datahub.ingestion.source.state.stateful_ingestion_base import (
+    StatefulLineageConfigMixin,
+    StatefulUsageConfigMixin,
+)
+
+
+class SnaplogicConfig(
+    StatefulIngestionConfigBase, StatefulLineageConfigMixin, StatefulUsageConfigMixin
+):
+    platform: str = "SnapLogic"
+    username: str = Field(description="Username")
+    password: SecretStr = Field(description="Password")
+    base_url: str = Field(
+        default="https://elastic.snaplogic.com",
+        description="Url to your SnapLogic instance: `https://elastic.snaplogic.com`, or similar. Used for making API calls to SnapLogic.",
+    )
+    org_name: str = Field(description="Organization name from SnapLogic instance")
+    namespace_mapping: dict = Field(
+        default={}, description="Mapping of namespaces to platform instances"
+    )
+    case_insensitive_namespaces: list = Field(
+        default=[],
+        description="List of namespaces that should be treated as case insensitive",
+    )
+    create_non_snaplogic_datasets: bool = Field(
+        default=False,
+        description="Whether to create datasets for non-SnapLogic datasets (e.g., databases, S3, etc.)",
+    )
+    stateful_ingestion: Optional[StatefulStaleMetadataRemovalConfig] = None

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())