acryl-datahub-dagster-plugin 0.0.0.dev0__py3-none-any.whl

acryl_datahub_dagster_plugin-0.0.0.dev0.dist-info/METADATA

@@ -0,0 +1,115 @@
+Metadata-Version: 2.1
+Name: acryl-datahub-dagster-plugin
+Version: 0.0.0.dev0
+Summary: Datahub Dagster plugin to capture executions and send to Datahub
+Home-page: https://datahubproject.io/
+License: Apache License 2.0
+Project-URL: Documentation, https://datahubproject.io/docs/
+Project-URL: Source, https://github.com/datahub-project/datahub
+Project-URL: Changelog, https://github.com/datahub-project/datahub/releases
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Information Technology
+Classifier: Intended Audience :: System Administrators
+Classifier: License :: OSI Approved
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: Unix
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Environment :: Console
+Classifier: Environment :: MacOS X
+Classifier: Topic :: Software Development
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: acryl-datahub[datahub-rest]
+Requires-Dist: dagster >=1.3.3
+Requires-Dist: requests-file
+Requires-Dist: requests
+Requires-Dist: dagit >=1.3.3
+Provides-Extra: dev
+Requires-Dist: dagster >=1.3.3 ; extra == 'dev'
+Requires-Dist: types-freezegun ; extra == 'dev'
+Requires-Dist: types-cachetools ; extra == 'dev'
+Requires-Dist: tox ; extra == 'dev'
+Requires-Dist: requests-file ; extra == 'dev'
+Requires-Dist: types-click ==0.1.12 ; extra == 'dev'
+Requires-Dist: flake8 >=6.0.0 ; extra == 'dev'
+Requires-Dist: packaging ; extra == 'dev'
+Requires-Dist: jsonpickle ; extra == 'dev'
+Requires-Dist: types-dataclasses ; extra == 'dev'
+Requires-Dist: black ==22.12.0 ; extra == 'dev'
+Requires-Dist: build ; extra == 'dev'
+Requires-Dist: flake8-tidy-imports >=4.3.0 ; extra == 'dev'
+Requires-Dist: types-requests ; extra == 'dev'
+Requires-Dist: mypy >=1.4.0 ; extra == 'dev'
+Requires-Dist: pytest-cov >=2.8.1 ; extra == 'dev'
+Requires-Dist: types-toml ; extra == 'dev'
+Requires-Dist: acryl-datahub[datahub-rest] ; extra == 'dev'
+Requires-Dist: types-pytz ; extra == 'dev'
+Requires-Dist: types-PyYAML ; extra == 'dev'
+Requires-Dist: requests-mock ; extra == 'dev'
+Requires-Dist: pytest >=6.2.2 ; extra == 'dev'
+Requires-Dist: requests ; extra == 'dev'
+Requires-Dist: pytest-asyncio >=0.16.0 ; extra == 'dev'
+Requires-Dist: types-tabulate ; extra == 'dev'
+Requires-Dist: deepdiff ; extra == 'dev'
+Requires-Dist: dagit >=1.3.3 ; extra == 'dev'
+Requires-Dist: freezegun ; extra == 'dev'
+Requires-Dist: flake8-bugbear ==23.3.12 ; extra == 'dev'
+Requires-Dist: types-pkg-resources ; extra == 'dev'
+Requires-Dist: isort >=5.7.0 ; extra == 'dev'
+Requires-Dist: twine ; extra == 'dev'
+Requires-Dist: types-python-dateutil ; extra == 'dev'
+Requires-Dist: types-six ; extra == 'dev'
+Requires-Dist: sqlalchemy-stubs ; extra == 'dev'
+Requires-Dist: coverage >=5.1 ; extra == 'dev'
+Requires-Dist: pydantic !=1.10.3,>=1.10.0 ; extra == 'dev'
+Provides-Extra: ignore
+Provides-Extra: integration-tests
+Requires-Dist: dagster >=1.3.3 ; extra == 'integration-tests'
+Requires-Dist: types-freezegun ; extra == 'integration-tests'
+Requires-Dist: types-cachetools ; extra == 'integration-tests'
+Requires-Dist: tox ; extra == 'integration-tests'
+Requires-Dist: requests-file ; extra == 'integration-tests'
+Requires-Dist: types-click ==0.1.12 ; extra == 'integration-tests'
+Requires-Dist: flake8 >=6.0.0 ; extra == 'integration-tests'
+Requires-Dist: packaging ; extra == 'integration-tests'
+Requires-Dist: jsonpickle ; extra == 'integration-tests'
+Requires-Dist: types-dataclasses ; extra == 'integration-tests'
+Requires-Dist: black ==22.12.0 ; extra == 'integration-tests'
+Requires-Dist: build ; extra == 'integration-tests'
+Requires-Dist: flake8-tidy-imports >=4.3.0 ; extra == 'integration-tests'
+Requires-Dist: types-requests ; extra == 'integration-tests'
+Requires-Dist: mypy >=1.4.0 ; extra == 'integration-tests'
+Requires-Dist: pytest-cov >=2.8.1 ; extra == 'integration-tests'
+Requires-Dist: types-toml ; extra == 'integration-tests'
+Requires-Dist: acryl-datahub[datahub-rest] ; extra == 'integration-tests'
+Requires-Dist: types-pytz ; extra == 'integration-tests'
+Requires-Dist: types-PyYAML ; extra == 'integration-tests'
+Requires-Dist: requests-mock ; extra == 'integration-tests'
+Requires-Dist: pytest >=6.2.2 ; extra == 'integration-tests'
+Requires-Dist: requests ; extra == 'integration-tests'
+Requires-Dist: pytest-asyncio >=0.16.0 ; extra == 'integration-tests'
+Requires-Dist: types-tabulate ; extra == 'integration-tests'
+Requires-Dist: deepdiff ; extra == 'integration-tests'
+Requires-Dist: dagit >=1.3.3 ; extra == 'integration-tests'
+Requires-Dist: freezegun ; extra == 'integration-tests'
+Requires-Dist: flake8-bugbear ==23.3.12 ; extra == 'integration-tests'
+Requires-Dist: types-pkg-resources ; extra == 'integration-tests'
+Requires-Dist: isort >=5.7.0 ; extra == 'integration-tests'
+Requires-Dist: twine ; extra == 'integration-tests'
+Requires-Dist: types-python-dateutil ; extra == 'integration-tests'
+Requires-Dist: types-six ; extra == 'integration-tests'
+Requires-Dist: sqlalchemy-stubs ; extra == 'integration-tests'
+Requires-Dist: coverage >=5.1 ; extra == 'integration-tests'
+Requires-Dist: pydantic !=1.10.3,>=1.10.0 ; extra == 'integration-tests'
+
+# Datahub Dagster Plugin
+
+See the DataHub Dagster docs for details.
+

acryl_datahub_dagster_plugin-0.0.0.dev0.dist-info/RECORD

@@ -0,0 +1,11 @@
+datahub_dagster_plugin/__init__.py,sha256=KN5nSxBdJqXJ-oGx6T8Hp0ta2pGtYDF1FD5UWFmCMTs,530
+datahub_dagster_plugin/datahub_dagster_plugin.py,sha256=szsT7ddxs16eE56mhCdNXDtlakbNLmuhd2wUbkba5Xc,64
+datahub_dagster_plugin/client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+datahub_dagster_plugin/client/dagster_generator.py,sha256=yi40Fs4kY-Ba9g5eeH6kxI9VGRB1uiRHum1SMGIvAdU,18594
+datahub_dagster_plugin/sensors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+datahub_dagster_plugin/sensors/datahub_sensors.py,sha256=oWAgOMlLeiJs37WzijfTK8W-vwnmGE0pf4IArFgt4YM,16652
+acryl_datahub_dagster_plugin-0.0.0.dev0.dist-info/METADATA,sha256=PXm9UvKsGUNH8lNgMvKArW94GfHZihPs_ghq78iV9P8,5609
+acryl_datahub_dagster_plugin-0.0.0.dev0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+acryl_datahub_dagster_plugin-0.0.0.dev0.dist-info/entry_points.txt,sha256=-CtPxtYb1u-zR36QnUQrvJJ6qbf1eDw9SruA22XBPZw,116
+acryl_datahub_dagster_plugin-0.0.0.dev0.dist-info/top_level.txt,sha256=JS8QiBAj3eMTcMNcstm_EXGAcziiXVNT2nzOcfhdEMc,23
+acryl_datahub_dagster_plugin-0.0.0.dev0.dist-info/RECORD,,

acryl_datahub_dagster_plugin-0.0.0.dev0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.43.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

acryl_datahub_dagster_plugin-0.0.0.dev0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[dagster.plugins]
+acryl-datahub-dagster-plugin = datahub_dagster_plugin.datahub_dagster_plugin:DatahubDagsterPlugin

acryl_datahub_dagster_plugin-0.0.0.dev0.dist-info/top_level.txt

@@ -0,0 +1 @@
+datahub_dagster_plugin

datahub_dagster_plugin/__init__.py

@@ -0,0 +1,21 @@
+# Published at https://pypi.org/project/acryl-datahub/.
+__package_name__ = "acryl-datahub-dagster-plugin"
+__version__ = "0.0.0.dev0"
+
+
+def is_dev_mode() -> bool:
+    return __version__.endswith("dev0")
+
+
+def nice_version_name() -> str:
+    if is_dev_mode():
+        return "unavailable (installed in develop mode)"
+    return __version__
+
+
+def get_provider_info():
+    return {
+        "package-name": f"{__package_name__}",
+        "name": f"{__package_name__}",
+        "description": "Datahub metadata collector plugin",
+    }

datahub_dagster_plugin/client/dagster_generator.py

@@ -0,0 +1,504 @@
+from dataclasses import dataclass
+from logging import Logger
+from typing import Any, Callable, Dict, List, NamedTuple, Optional, Sequence, Set
+from urllib.parse import urlsplit
+
+import pydantic
+from dagster import DagsterRunStatus, PathMetadataValue, RunStatusSensorContext
+from dagster._core.execution.stats import RunStepKeyStatsSnapshot, StepEventStatus
+from dagster._core.snap import JobSnapshot
+from dagster._core.snap.node import OpDefSnap
+from dagster._core.storage.dagster_run import DagsterRun, DagsterRunStatsSnapshot
+from datahub.api.entities.datajob import DataFlow, DataJob
+from datahub.api.entities.dataprocess.dataprocess_instance import (
+    DataProcessInstance,
+    InstanceRunResult,
+)
+from datahub.api.entities.dataset.dataset import Dataset
+from datahub.configuration.source_common import DatasetSourceConfigMixin
+from datahub.emitter.mce_builder import (
+    make_data_platform_urn,
+    make_dataplatform_instance_urn,
+)
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph
+from datahub.metadata.schema_classes import DataPlatformInstanceClass, SubTypesClass
+from datahub.utilities.urns.data_flow_urn import DataFlowUrn
+from datahub.utilities.urns.data_job_urn import DataJobUrn
+from datahub.utilities.urns.dataset_urn import DatasetUrn
+
+
+class Constant:
+    """
+    keys used in dagster plugin
+    """
+
+    ORCHESTRATOR = "dagster"
+
+    # Default config constants
+    DEFAULT_DATAHUB_REST_URL = "http://localhost:8080"
+
+    # Environment variable contants
+    DATAHUB_REST_URL = "DATAHUB_REST_URL"
+    DATAHUB_ENV = "DATAHUB_ENV"
+    DATAHUB_PLATFORM_INSTANCE = "DATAHUB_PLATFORM_INSTANCE"
+    DAGSTER_UI_URL = "DAGSTER_UI_URL"
+
+    # Datahub inputs/outputs constant
+    DATAHUB_INPUTS = "datahub.inputs"
+    DATAHUB_OUTPUTS = "datahub.outputs"
+
+    # Job run constant
+    JOB_SNAPSHOT_ID = "job_snapshot_id"
+    EXECUTION_PLAN_SNAPSHOT_ID = "execution_plan_snapshot_id"
+    ROOT_RUN_ID = "root_run_id"
+    PARENT_RUN_ID = "parent_run_id"
+    HAS_REPOSITORY_LOAD_DATA = "has_repository_load_data"
+    TAGS = "tags"
+    STEPS_SUCCEEDED = "steps_succeeded"
+    STEPS_FAILED = "steps_failed"
+    MATERIALIZATIONS = "materializations"
+    EXPECTATIONS = "expectations"
+    ENQUEUED_TIME = "enqueued_time"
+    LAUNCH_TIME = "launch_time"
+    START_TIME = "start_time"
+    END_TIME = "end_time"
+
+    # Op run contants
+    STEP_KEY = "step_key"
+    ATTEMPTS = "attempts"
+
+
+class DatasetLineage(NamedTuple):
+    inputs: Set[str]
+    outputs: Set[str]
+
+
+class DatahubDagsterSourceConfig(DatasetSourceConfigMixin):
+    datahub_client_config: DatahubClientConfig = pydantic.Field(
+        default=DatahubClientConfig(),
+        description="Datahub client config",
+    )
+
+    dagster_url: Optional[str] = pydantic.Field(
+        default=None,
+        description="Dagster UI URL. Like: https://myDagsterCloudEnvironment.dagster.cloud/prod",
+    )
+
+    capture_asset_materialization: bool = pydantic.Field(
+        default=True,
+        description="Whether to capture asset keys as Dataset on AssetMaterialization event",
+    )
+
+    capture_input_output: bool = pydantic.Field(
+        default=False,
+        description="Whether to capture and try to parse input and output from HANDLED_OUTPUT, LOADED_INPUT event. (currently only filepathvalue metadata supported",
+    )
+
+    asset_lineage_extractor: Optional[
+        Callable[
+            [RunStatusSensorContext, "DagsterGenerator", DataHubGraph],
+            Dict[str, DatasetLineage],
+        ]
+    ] = pydantic.Field(
+        default=None,
+        description="Custom asset lineage extractor function. See details at [https://datahubproject.io/docs/lineage/dagster/#define-your-custom-logic-to-capture-asset-lineage-information]",
+    )
+
+
+def _str_urn_to_dataset_urn(urns: List[str]) -> List[DatasetUrn]:
+    return [DatasetUrn.create_from_string(urn) for urn in urns]
+
+
+@dataclass
+class DagsterEnvironment:
+    repository: Optional[str]
+    is_cloud: bool = True
+    is_branch_deployment: bool = False
+    branch: Optional[str] = "prod"
+    module: Optional[str] = None
+
+
+def job_url_generator(dagster_url: str, dagster_environment: DagsterEnvironment) -> str:
+    if dagster_environment.is_cloud:
+        base_url = f"{dagster_url}/{dagster_environment.branch}"
+    else:
+        base_url = dagster_url
+
+    if dagster_environment.module:
+        base_url = f"{base_url}/locations/{dagster_environment.module}"
+
+    return base_url
+
+
+class DagsterGenerator:
+    def __init__(
+        self,
+        logger: Logger,
+        config: DatahubDagsterSourceConfig,
+        dagster_environment: DagsterEnvironment,
+    ):
+        self.logger = logger
+        self.config = config
+        self.dagster_environment = dagster_environment
+
+    def path_metadata_resolver(self, value: PathMetadataValue) -> Optional[DatasetUrn]:
+        """
+        Resolve path metadata to dataset urn
+        """
+        path = value.value
+        if not path:
+            return None
+
+        if "://" in path:
+            url = urlsplit(path)
+            scheme = url.scheme
+
+            # Need to adjust some these schemes
+            if scheme in ["s3a", "s3n"]:
+                scheme = "s3"
+            elif scheme in ["gs"]:
+                scheme = "gcs"
+
+            return DatasetUrn(platform=scheme, name=url.path)
+        else:
+            return DatasetUrn(platform="file", name=path)
+
+    def metadata_resolver(self, metadata: Any) -> Optional[DatasetUrn]:
+        """
+        Resolve metadata to dataset urn
+        """
+        if isinstance(metadata, PathMetadataValue):
+            return self.path_metadata_resolver(metadata)
+        else:
+            self.logger.info(f"Unknown Metadata: {metadata} of type {type(metadata)}")
+        return None
+
+    def generate_dataflow(
+        self,
+        job_snapshot: JobSnapshot,
+        env: str,
+        platform_instance: Optional[str] = None,
+    ) -> DataFlow:
+        """
+        Generates a Dataflow object from an Dagster Job Snapshot
+        :param job_snapshot: JobSnapshot - Job snapshot object
+        :param env: str
+        :param platform_instance: Optional[str]
+        :return: DataFlow - Data generated dataflow
+        """
+        if self.dagster_environment.is_cloud:
+            id = f"{self.dagster_environment.branch}/{self.dagster_environment.module}/{job_snapshot.name}"
+        else:
+            id = f"{self.dagster_environment.module}/{job_snapshot.name}"
+
+        dataflow = DataFlow(
+            orchestrator=Constant.ORCHESTRATOR,
+            id=id,
+            env=env,
+            name=job_snapshot.name,
+            platform_instance=platform_instance,
+        )
+        dataflow.description = job_snapshot.description
+        dataflow.tags = set(job_snapshot.tags.keys())
+        if self.config.dagster_url:
+            dataflow.url = f"{job_url_generator(dagster_url=self.config.dagster_url, dagster_environment=self.dagster_environment)}/jobs/{job_snapshot.name}"
+        flow_property_bag: Dict[str, str] = {}
+        for key in job_snapshot.metadata.keys():
+            flow_property_bag[key] = str(job_snapshot.metadata[key])
+        dataflow.properties = flow_property_bag
+        return dataflow
+
+    def generate_datajob(
+        self,
+        job_snapshot: JobSnapshot,
+        step_deps: Dict[str, List],
+        op_def_snap: OpDefSnap,
+        env: str,
+        input_datasets: Dict[str, Set[DatasetUrn]],
+        output_datasets: Dict[str, Set[DatasetUrn]],
+        platform_instance: Optional[str] = None,
+    ) -> DataJob:
+        """
+        Generates a Datajob object from an Dagster op snapshot
+        :param job_snapshot: JobSnapshot - Job snapshot object
+        :param op_def_snap: OpDefSnap - Op def snapshot object
+        :param env: str
+        :param platform_instance: Optional[str]
+        :param output_datasets: dict[str, Set[DatasetUrn]] - output datasets for each op
+        :return: DataJob - Data generated datajob
+        """
+
+        if self.dagster_environment.is_cloud:
+            flow_id = f"{self.dagster_environment.branch}/{self.dagster_environment.module}/{job_snapshot.name}"
+            job_id = f"{self.dagster_environment.branch}/{self.dagster_environment.module}/{op_def_snap.name}"
+        else:
+            flow_id = f"{self.dagster_environment.module}/{job_snapshot.name}"
+            job_id = f"{self.dagster_environment.module}/{op_def_snap.name}"
+
+        dataflow_urn = DataFlowUrn.create_from_ids(
+            orchestrator=Constant.ORCHESTRATOR,
+            flow_id=flow_id,
+            env=env,
+            platform_instance=platform_instance,
+        )
+        datajob = DataJob(
+            id=job_id,
+            flow_urn=dataflow_urn,
+            name=op_def_snap.name,
+        )
+
+        if self.config.dagster_url:
+            datajob.url = f"{job_url_generator(dagster_url=self.config.dagster_url, dagster_environment=self.dagster_environment)}/jobs/{job_snapshot.name}/{op_def_snap.name}"
+
+        datajob.description = op_def_snap.description
+        datajob.tags = set(op_def_snap.tags.keys())
+
+        # Add upstream dependencies for this op
+        for upstream_op_name in step_deps[op_def_snap.name]:
+            if self.dagster_environment.is_cloud:
+                upstream_job_id = f"{self.dagster_environment.branch}/{self.dagster_environment.module}/{upstream_op_name}"
+            else:
+                upstream_job_id = (
+                    f"{self.dagster_environment.module}/{upstream_op_name}"
+                )
+            upstream_op_urn = DataJobUrn.create_from_ids(
+                data_flow_urn=str(dataflow_urn),
+                job_id=upstream_job_id,
+            )
+            datajob.upstream_urns.extend([upstream_op_urn])
+        job_property_bag: Dict[str, str] = {}
+        if input_datasets:
+            self.logger.info(
+                f"Input datasets for {op_def_snap.name} are { list(input_datasets.get(op_def_snap.name, []))}"
+            )
+            datajob.inlets = list(input_datasets.get(op_def_snap.name, []))
+
+        if output_datasets:
+            self.logger.info(
+                f"Output datasets for {op_def_snap.name} are { list(output_datasets.get(op_def_snap.name, []))}"
+            )
+            datajob.outlets = list(output_datasets.get(op_def_snap.name, []))
+
+        # For all op inputs/outputs:
+        # Add input/output details like its type, description, metadata etc in datajob properties.
+        # Also, add datahub inputs/outputs if present in input/output metatdata.
+        for input_def_snap in op_def_snap.input_def_snaps:
+            job_property_bag[f"input.{input_def_snap.name}"] = str(
+                input_def_snap._asdict()
+            )
+            if Constant.DATAHUB_INPUTS in input_def_snap.metadata:
+                datajob.inlets.extend(
+                    _str_urn_to_dataset_urn(
+                        input_def_snap.metadata[Constant.DATAHUB_INPUTS].value  # type: ignore
+                    )
+                )
+
+        for output_def_snap in op_def_snap.output_def_snaps:
+            job_property_bag[f"output_{output_def_snap.name}"] = str(
+                output_def_snap._asdict()
+            )
+            if Constant.DATAHUB_OUTPUTS in output_def_snap.metadata:
+                datajob.outlets.extend(
+                    _str_urn_to_dataset_urn(
+                        output_def_snap.metadata[Constant.DATAHUB_OUTPUTS].value  # type: ignore
+                    )
+                )
+
+        datajob.properties = job_property_bag
+
+        return datajob
+
+    def emit_job_run(
+        self,
+        graph: DataHubGraph,
+        dataflow: DataFlow,
+        run: DagsterRun,
+        run_stats: DagsterRunStatsSnapshot,
+    ) -> None:
+        """
+        Emit a latest job run
+        :param graph: DatahubRestEmitter
+        :param dataflow: DataFlow - DataFlow object
+        :param run: DagsterRun - Dagster Run object
+        :param run_stats: DagsterRunStatsSnapshot - latest job run stats
+        """
+        dpi = DataProcessInstance.from_dataflow(dataflow=dataflow, id=run_stats.run_id)
+        if self.config.dagster_url:
+            if self.dagster_environment.is_cloud:
+                dpi.url = f"{self.config.dagster_url}/{self.dagster_environment.branch}/runs/{run.run_id}"
+            else:
+                dpi.url = f"{self.config.dagster_url}/runs/{run.run_id}"
+
+        # Add below details in dpi properties
+        dpi_property_bag: Dict[str, str] = {}
+        allowed_job_run_keys = [
+            Constant.JOB_SNAPSHOT_ID,
+            Constant.EXECUTION_PLAN_SNAPSHOT_ID,
+            Constant.ROOT_RUN_ID,
+            Constant.PARENT_RUN_ID,
+            Constant.HAS_REPOSITORY_LOAD_DATA,
+            Constant.TAGS,
+            Constant.STEPS_SUCCEEDED,
+            Constant.STEPS_FAILED,
+            Constant.MATERIALIZATIONS,
+            Constant.EXPECTATIONS,
+            Constant.ENQUEUED_TIME,
+            Constant.LAUNCH_TIME,
+            Constant.START_TIME,
+            Constant.END_TIME,
+        ]
+        for key in allowed_job_run_keys:
+            if hasattr(run, key) and getattr(run, key) is not None:
+                dpi_property_bag[key] = str(getattr(run, key))
+            if hasattr(run_stats, key) and getattr(run_stats, key) is not None:
+                dpi_property_bag[key] = str(getattr(run_stats, key))
+        dpi.properties.update(dpi_property_bag)
+
+        status_result_map = {
+            DagsterRunStatus.SUCCESS: InstanceRunResult.SUCCESS,
+            DagsterRunStatus.FAILURE: InstanceRunResult.FAILURE,
+            DagsterRunStatus.CANCELED: InstanceRunResult.SKIPPED,
+        }
+
+        if run.status not in status_result_map:
+            raise Exception(
+                f"Job run status should be either complete, failed or cancelled and it was "
+                f"{run.status }"
+            )
+
+        if run_stats.start_time is not None:
+            dpi.emit_process_start(
+                emitter=graph,
+                start_timestamp_millis=int(run_stats.start_time * 1000),
+            )
+
+        if run_stats.end_time is not None:
+            dpi.emit_process_end(
+                emitter=graph,
+                end_timestamp_millis=int(run_stats.end_time * 1000),
+                result=status_result_map[run.status],
+                result_type=Constant.ORCHESTRATOR,
+            )
+
+    def emit_op_run(
+        self,
+        graph: DataHubGraph,
+        datajob: DataJob,
+        run_step_stats: RunStepKeyStatsSnapshot,
+    ) -> None:
+        """
+        Emit an op run
+        :param graph: DataHubGraph
+        :param datajob: DataJob - DataJob object
+        :param run_step_stats: RunStepKeyStatsSnapshot - step(op) run stats
+        """
+        dpi = DataProcessInstance.from_datajob(
+            datajob=datajob,
+            id=f"{run_step_stats.run_id}.{datajob.id}",
+            clone_inlets=True,
+            clone_outlets=True,
+        )
+        if self.config.dagster_url:
+            dpi.url = f"{self.config.dagster_url}/runs/{run_step_stats.run_id}"
+            if self.dagster_environment.is_cloud:
+                dpi.url = f"{self.config.dagster_url}/{self.dagster_environment.branch}/runs/{run_step_stats.run_id}"
+            else:
+                dpi.url = f"{self.config.dagster_url}/runs/{run_step_stats.run_id}"
+
+        # Add below details in dpi properties
+        dpi_property_bag: Dict[str, str] = {}
+        allowed_op_run_keys = [
+            Constant.STEP_KEY,
+            Constant.ATTEMPTS,
+            Constant.START_TIME,
+            Constant.END_TIME,
+        ]
+        for key in allowed_op_run_keys:
+            if (
+                hasattr(run_step_stats, key)
+                and getattr(run_step_stats, key) is not None
+            ):
+                dpi_property_bag[key] = str(getattr(run_step_stats, key))
+        dpi.properties.update(dpi_property_bag)
+
+        status_result_map = {
+            StepEventStatus.SUCCESS: InstanceRunResult.SUCCESS,
+            StepEventStatus.FAILURE: InstanceRunResult.FAILURE,
+            StepEventStatus.SKIPPED: InstanceRunResult.SKIPPED,
+        }
+
+        if run_step_stats.status not in status_result_map:
+            raise Exception(
+                f"Step run status should be either complete, failed or cancelled and it was "
+                f"{run_step_stats.status }"
+            )
+
+        if run_step_stats.start_time is not None:
+            dpi.emit_process_start(
+                emitter=graph,
+                start_timestamp_millis=int(run_step_stats.start_time * 1000),
+            )
+
+        if run_step_stats.end_time is not None:
+            dpi.emit_process_end(
+                emitter=graph,
+                end_timestamp_millis=int(run_step_stats.end_time * 1000),
+                result=status_result_map[run_step_stats.status],
+                result_type=Constant.ORCHESTRATOR,
+            )
+
+    def dataset_urn_from_asset(self, asset_key: Sequence[str]) -> DatasetUrn:
+        """
+        Generate dataset urn from asset key
+        """
+        return DatasetUrn(
+            platform="dagster", env=self.config.env, name="/".join(asset_key)
+        )
+
+    def emit_asset(
+        self,
+        graph: DataHubGraph,
+        asset_key: Sequence[str],
+        description: Optional[str],
+        properties: Optional[Dict[str, str]],
+    ) -> str:
+        """
+        Emit asset to datahub
+        """
+        dataset_urn = self.dataset_urn_from_asset(asset_key)
+        dataset = Dataset(
+            id=None,
+            urn=dataset_urn.urn(),
+            platform="dagster",
+            name=asset_key[-1],
+            schema=None,
+            downstreams=None,
+            subtype="Asset",
+            subtypes=None,
+            description=description,
+            env=self.config.env,
+            properties=properties,
+        )
+        for mcp in dataset.generate_mcp():
+            graph.emit_mcp(mcp)
+
+        mcp = MetadataChangeProposalWrapper(
+            entityUrn=dataset_urn.urn(),
+            aspect=SubTypesClass(typeNames=["Asset"]),
+        )
+        graph.emit_mcp(mcp)
+
+        if self.config.platform_instance:
+            mcp = MetadataChangeProposalWrapper(
+                entityUrn=dataset_urn.urn(),
+                aspect=DataPlatformInstanceClass(
+                    instance=make_dataplatform_instance_urn(
+                        instance=self.config.platform_instance,
+                        platform="dagster",
+                    ),
+                    platform=make_data_platform_urn("dagster"),
+                ),
+            )
+            graph.emit_mcp(mcp)
+        return dataset_urn.urn()

datahub_dagster_plugin/datahub_dagster_plugin.py

@@ -0,0 +1,2 @@
+class DatahubDagsterPlugin:
+    name = "datahub_dagster_plugin"

datahub_dagster_plugin/sensors/datahub_sensors.py

@@ -0,0 +1,439 @@
+import os
+import traceback
+from typing import Dict, List, Optional, Sequence, Set, Tuple
+
+from dagster import (
+    DagsterRunStatus,
+    EventLogEntry,
+    RunStatusSensorContext,
+    SensorDefinition,
+    SkipReason,
+    run_status_sensor,
+    sensor,
+)
+from dagster._core.definitions.asset_selection import CoercibleToAssetSelection
+from dagster._core.definitions.sensor_definition import (
+    DefaultSensorStatus,
+    RawSensorEvaluationFunctionReturn,
+)
+from dagster._core.definitions.target import ExecutableDefinition
+from dagster._core.events import DagsterEventType, HandledOutputData, LoadedInputData
+from dagster._core.execution.stats import RunStepKeyStatsSnapshot
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.metadata.schema_classes import SubTypesClass
+
+from datahub_dagster_plugin.client.dagster_generator import (
+    DagsterEnvironment,
+    DagsterGenerator,
+    DatahubDagsterSourceConfig,
+)
+
+
+def make_datahub_sensor(
+    config: DatahubDagsterSourceConfig,
+    name: Optional[str] = None,
+    minimum_interval_seconds: Optional[int] = None,
+    description: Optional[str] = None,
+    job: Optional[ExecutableDefinition] = None,
+    jobs: Optional[Sequence[ExecutableDefinition]] = None,
+    default_status: DefaultSensorStatus = DefaultSensorStatus.STOPPED,
+    asset_selection: Optional[CoercibleToAssetSelection] = None,
+    required_resource_keys: Optional[Set[str]] = None,
+) -> SensorDefinition:
+    """Create a sensor on job status change emit lineage to DataHub.
+
+    Args:
+        config (DatahubDagsterSourceConfig): DataHub Sensor config
+        name: (Optional[str]): The name of the sensor. Defaults to "datahub_sensor".
+        minimum_interval_seconds: (Optional[int]): The minimum number of seconds that will elapse
+            between sensor evaluations.
+        default_status (DefaultSensorStatus): Whether the sensor starts as running or not. The default
+            status can be overridden from Dagit or via the GraphQL API.
+
+    Examples:
+        .. code-block:: python
+
+            datahub_sensor = make_datahub_sensor(
+               config
+            )
+
+            @repository
+            def my_repo():
+                return [my_job + datahub_sensor]
+
+
+    """
+
+    @sensor(
+        name=name,
+        minimum_interval_seconds=minimum_interval_seconds,
+        description=description,
+        job=job,
+        jobs=jobs,
+        default_status=default_status,
+        asset_selection=asset_selection,
+        required_resource_keys=required_resource_keys,
+    )
+    def datahub_sensor(context):
+        """
+        Sensor which instigate all run status sensors and trigger them based upon run status
+        """
+        for each in DatahubSensors(config).sensors:
+            each.evaluate_tick(context)
+        return SkipReason("Trigger run status sensors if any new runs present...")
+
+    return datahub_sensor
+
+
+class DatahubSensors:
+    def __init__(self, config: Optional[DatahubDagsterSourceConfig] = None):
+        """
+        Set dagster source configurations and initialize datahub emitter and dagster run status sensors
+        """
+        if config:
+            self.config = config
+        else:
+            self.config = DatahubDagsterSourceConfig()
+        self.graph = DataHubGraph(
+            self.config.datahub_client_config,
+        )
+
+        self.graph.test_connection()
+        self.sensors: List[SensorDefinition] = []
+        self.sensors.append(
+            run_status_sensor(
+                name="datahub_success_sensor", run_status=DagsterRunStatus.SUCCESS
+            )(self._emit_metadata)
+        )
+
+        self.sensors.append(
+            run_status_sensor(
+                name="datahub_failure_sensor", run_status=DagsterRunStatus.FAILURE
+            )(self._emit_metadata)
+        )
+
+        self.sensors.append(
+            run_status_sensor(
+                name="datahub_canceled_sensor", run_status=DagsterRunStatus.CANCELED
+            )(self._emit_metadata)
+        )
+
+    def get_dagster_environment(
+        self, context: RunStatusSensorContext
+    ) -> Optional[DagsterEnvironment]:
+        if (
+            context.dagster_run.job_code_origin
+            and context.dagster_run.job_code_origin.repository_origin
+            and context.dagster_run.job_code_origin.repository_origin.code_pointer
+        ):
+
+            code_pointer = (
+                context.dagster_run.job_code_origin.repository_origin.code_pointer
+            )
+            context.log.debug(f"code_pointer: {code_pointer}")
+
+            if hasattr(code_pointer, "attribute"):
+                repository = code_pointer.attribute
+            else:
+                repository = None
+
+            if hasattr(code_pointer, "module"):
+                module = code_pointer.module
+            else:
+                context.log.error("Unable to get Module")
+                return None
+
+            dagster_environment = DagsterEnvironment(
+                is_cloud=os.getenv("DAGSTER_CLOUD_IS_BRANCH_DEPLOYMENT", None)
+                is not None,
+                is_branch_deployment=(
+                    True
+                    if os.getenv("DAGSTER_CLOUD_IS_BRANCH_DEPLOYMENT", False) == 1
+                    else False
+                ),
+                branch=os.getenv("DAGSTER_CLOUD_DEPLOYMENT_NAME", "prod"),
+                module=module,
+                repository=repository,
+            )
+            return dagster_environment
+        else:
+            context.log.error("Unable to get Dagster Environment...")
+            return None
+
+    def process_asset_logs(
+        self,
+        dagster_generator: DagsterGenerator,
+        log: EventLogEntry,
+        dataset_inputs: Dict[str, set],
+        dataset_outputs: Dict[str, set],
+    ) -> None:
+
+        if not log.dagster_event or not log.step_key:
+            return
+
+        if log.dagster_event.event_type == DagsterEventType.ASSET_MATERIALIZATION:
+            if log.step_key not in dataset_outputs:
+                dataset_outputs[log.step_key] = set()
+
+            materialization = log.asset_materialization
+            if not materialization:
+                return
+
+            properties = {
+                key: str(value) for (key, value) in materialization.metadata.items()
+            }
+            asset_key = materialization.asset_key.path
+            dataset_urn = dagster_generator.emit_asset(
+                self.graph, asset_key, materialization.description, properties
+            )
+            dataset_outputs[log.step_key].add(dataset_urn)
+
+        elif log.dagster_event.event_type == DagsterEventType.ASSET_OBSERVATION:
+            if log.step_key not in dataset_inputs:
+                dataset_inputs[log.step_key] = set()
+            asset_observation = log.asset_observation
+            if not asset_observation:
+                return
+
+            properties = {
+                key: str(value)
+                for (key, value) in asset_observation.metadata.items()  # type: ignore
+            }
+            asset_key = asset_observation.asset_key.path  # type: ignore
+            dataset_urn = dagster_generator.emit_asset(
+                self.graph,
+                asset_key,
+                asset_observation.description,
+                properties,  # type: ignore
+            )
+            dataset_inputs[log.step_key].add(dataset_urn)
+
+    def process_handle_input_output(
+        self,
+        context: RunStatusSensorContext,
+        log: EventLogEntry,
+        dagster_generator: DagsterGenerator,
+        dataset_inputs: Dict[str, set],
+        dataset_outputs: Dict[str, set],
+    ) -> None:
+        if not log.dagster_event or not log.step_key:
+            return
+
+        if (
+            self.config.capture_input_output
+            and log.dagster_event.event_type == DagsterEventType.HANDLED_OUTPUT
+        ):
+            if log.step_key not in dataset_outputs:
+                dataset_outputs[log.step_key] = set()
+
+            event_specific_data = log.dagster_event.event_specific_data
+            if isinstance(event_specific_data, HandledOutputData):
+                context.log.debug(
+                    f"Output Path: {event_specific_data.metadata.get('path')}"
+                )
+                metadata = event_specific_data.metadata.get("path")
+                context.log.debug(f"Metadata: {metadata}")
+                if not metadata:
+                    return
+                urn = dagster_generator.metadata_resolver(metadata)
+                if urn:
+                    context.log.debug(f"Output Urn: {urn}")
+                    dataset_outputs[log.step_key].add(urn)
+        elif (
+            self.config.capture_input_output
+            and log.dagster_event.event_type == DagsterEventType.LOADED_INPUT
+        ):
+            if log.step_key not in dataset_inputs:
+                dataset_inputs[log.step_key] = set()
+            event_specific_data = log.dagster_event.event_specific_data
+            if isinstance(event_specific_data, LoadedInputData):
+                context.log.debug(
+                    f"Input Path: {event_specific_data.metadata.get('path')}"
+                )
+                metadata = event_specific_data.metadata.get("path")
+                context.log.debug(f"Metadata: {metadata}")
+                if not metadata:
+                    return
+                urn = dagster_generator.metadata_resolver(metadata)
+                if urn:
+                    context.log.debug(f"Input Urn: {urn}")
+                    dataset_inputs[log.step_key].add(urn)
+
+    def process_dagster_logs(
+        self, context: RunStatusSensorContext, dagster_generator: DagsterGenerator
+    ) -> Tuple[Dict[str, set], Dict[str, set]]:
+        dataset_outputs: Dict[str, set] = {}
+        dataset_inputs: Dict[str, set] = {}
+
+        logs = context.instance.all_logs(
+            context.dagster_run.run_id,
+            {
+                DagsterEventType.ASSET_MATERIALIZATION,
+                DagsterEventType.ASSET_OBSERVATION,
+                DagsterEventType.HANDLED_OUTPUT,
+                DagsterEventType.LOADED_INPUT,
+            },
+        )
+
+        for log in logs:
+            if not log.dagster_event or not log.step_key:
+                continue
+            context.log.debug(f"Log: {log.step_key} - {log.dagster_event}")
+            context.log.debug(f"Event Type: {log.dagster_event.event_type}")
+            if self.config.capture_input_output:
+                self.process_handle_input_output(
+                    context=context,
+                    log=log,
+                    dagster_generator=dagster_generator,
+                    dataset_inputs=dataset_inputs,
+                    dataset_outputs=dataset_outputs,
+                )
+
+            if self.config.capture_asset_materialization:
+                self.process_asset_logs(
+                    dagster_generator=dagster_generator,
+                    log=log,
+                    dataset_inputs=dataset_inputs,
+                    dataset_outputs=dataset_outputs,
+                )
+
+        return dataset_inputs, dataset_outputs
+
+    @staticmethod
+    def merge_dicts(dict1: Dict[str, Set], dict2: Dict[str, Set]) -> Dict[str, Set]:
+        """
+        Merge two dictionaries
+        """
+        for key, value in dict2.items():
+            if key in dict1:
+                dict1[key] = dict1[key].union(value)
+            else:
+                dict1[key] = value
+        return dict1
+
+    def _emit_metadata(
+        self, context: RunStatusSensorContext
+    ) -> RawSensorEvaluationFunctionReturn:
+        """
+        Function to emit metadata for datahub rest.
+        """
+        try:
+            context.log.info("Emitting metadata...")
+
+            assert context.dagster_run.job_snapshot_id
+            assert context.dagster_run.execution_plan_snapshot_id
+
+            dagster_environment = self.get_dagster_environment(context)
+            context.log.debug(f"dagster enivronment: {dagster_environment}")
+            if not dagster_environment:
+                return SkipReason(
+                    "Unable to get Dagster Environment from DataHub Sensor"
+                )
+
+            context.log.debug(f"Dagster Environment: {dagster_environment}")
+
+            dagster_generator = DagsterGenerator(
+                logger=context.log,
+                config=self.config,
+                dagster_environment=dagster_environment,
+            )
+
+            job_snapshot = context.instance.get_job_snapshot(
+                snapshot_id=context.dagster_run.job_snapshot_id
+            )
+
+            dataset_inputs: Dict[str, Set] = {}
+            dataset_outputs: Dict[str, Set] = {}
+
+            if self.config.asset_lineage_extractor:
+                asset_lineages = self.config.asset_lineage_extractor(
+                    context, dagster_generator, self.graph
+                )
+                for key, value in asset_lineages.items():
+                    dataset_inputs[key] = dataset_inputs.get(key, set()).union(
+                        value.inputs
+                    )
+                    dataset_outputs[key] = dataset_outputs.get(key, set()).union(
+                        value.outputs
+                    )
+
+            (
+                dataset_inputs_from_logs,
+                dataset_outputs_from_logs,
+            ) = self.process_dagster_logs(context, dagster_generator)
+
+            dataset_inputs = DatahubSensors.merge_dicts(
+                dataset_inputs, dataset_inputs_from_logs
+            )
+            dataset_outputs = DatahubSensors.merge_dicts(
+                dataset_outputs, dataset_outputs_from_logs
+            )
+
+            context.log.debug(f"Outputs: {dataset_outputs}")
+            # Emit dagster job entity which get mapped with datahub dataflow entity
+            dataflow = dagster_generator.generate_dataflow(
+                job_snapshot=job_snapshot,
+                env=self.config.env,
+                platform_instance=self.config.platform_instance,
+            )
+            dataflow.emit(self.graph)
+
+            # Emit dagster job run which get mapped with datahub data process instance entity
+            dagster_generator.emit_job_run(
+                graph=self.graph,
+                dataflow=dataflow,
+                run=context.dagster_run,
+                run_stats=context.instance.get_run_stats(context.dagster_run.run_id),
+            )
+
+            # Execution plan snapshot contains all steps(ops) dependency.
+            execution_plan_snapshot = context.instance.get_execution_plan_snapshot(
+                snapshot_id=context.dagster_run.execution_plan_snapshot_id
+            )
+
+            # Map step key with its run step stats
+            run_step_stats: Dict[str, RunStepKeyStatsSnapshot] = {
+                run_step_stat.step_key: run_step_stat
+                for run_step_stat in context.instance.get_run_step_stats(
+                    context.dagster_run.run_id
+                )
+            }
+
+            # For all dagster ops present in job:
+            # Emit op entity which get mapped with datahub datajob entity.
+            # Emit op run which get mapped with datahub data process instance entity.
+            for op_def_snap in job_snapshot.node_defs_snapshot.op_def_snaps:
+                datajob = dagster_generator.generate_datajob(
+                    job_snapshot=job_snapshot,
+                    step_deps=execution_plan_snapshot.step_deps,
+                    op_def_snap=op_def_snap,
+                    env=self.config.env,
+                    platform_instance=self.config.platform_instance,
+                    output_datasets=dataset_outputs,
+                    input_datasets=dataset_inputs,
+                )
+                context.log.info(f"Generated Datajob: {datajob}")
+                datajob.emit(self.graph)
+
+                self.graph.emit_mcp(
+                    mcp=MetadataChangeProposalWrapper(
+                        entityUrn=str(datajob.urn),
+                        aspect=SubTypesClass(
+                            typeNames=["Op"],
+                        ),
+                    )
+                )
+
+                dagster_generator.emit_op_run(
+                    graph=self.graph,
+                    datajob=datajob,
+                    run_step_stats=run_step_stats[op_def_snap.name],
+                )
+
+            return SkipReason("Pipeline metadata is emitted to DataHub")
+        except Exception as e:
+            context.log.error(
+                f"Error in emitting metadata to DataHub: {e}. Traceback: {traceback.format_exc()}"
+            )
+            return SkipReason("Error in emitting metadata to DataHub")