dagster-datacontract 0.1.0__py3-none-any.whl

dagster_datacontract/__init__.py

@@ -0,0 +1,139 @@
+import json
+import textwrap
+from datetime import timedelta
+from typing import Any
+
+import dagster as dg
+from dagster import TableColumnLineage, TableSchema
+from datacontract.data_contract import DataContract
+from datacontract.model.run import ResultEnum
+
+
+class DataContractLoader:
+    def __init__(
+        self,
+        asset_name: str,
+        data_contract: DataContract,
+    ):
+        self.asset_name = asset_name
+        self.asset_key = dg.AssetKey(path=self.asset_name)
+        self.data_contract = data_contract
+        self.data_contract_specification = (
+            self.data_contract.get_data_contract_specification()
+        )
+        self.metadata = self._load_metadata()
+        self.tags = self._load_tags()
+        self.description = self._load_description()
+        self.owner = self._load_owner()
+        self.version = self._load_version()
+        self.cron_schedule = self._load_cron_schedule()
+
+    def _load_metadata(
+        self,
+    ) -> dict[str, TableColumnLineage | TableSchema | Any] | None:
+        fields = self.data_contract_specification.models.get(self.asset_name).fields
+
+        columns = []
+        deps_by_column = {}
+
+        for column_name, column_field in fields.items():
+            columns.append(
+                dg.TableColumn(
+                    name=column_name,
+                    type=column_field.type,
+                    description=column_field.description,
+                )
+            )
+
+            lineage = json.loads(column_field.model_dump_json()).get("lineage")
+            if not lineage:
+                deps_by_column[column_name] = []
+            else:
+                lineage_entries = lineage.get("inputFields")
+
+                deps_by_column[column_name] = [
+                    dg.TableColumnDep(
+                        asset_key=dg.AssetKey(lineage_entry["name"]),
+                        column_name=lineage_entry["field"],
+                    )
+                    for lineage_entry in lineage_entries
+                ]
+
+        return {
+            "dagster/column_schema": dg.TableSchema(columns=columns),
+            "dagster/column_lineage": dg.TableColumnLineage(
+                deps_by_column=deps_by_column
+            ),
+        }
+
+    def _load_tags(self) -> dict[str, str]:
+        tags = {
+            item.split(":")[0].strip(): item.split(":")[1].strip()
+            if ":" in item
+            else ""
+            for item in self.data_contract_specification.tags
+        }
+
+        return tags
+
+    def _load_description(self) -> str | None:
+        model_description = self.data_contract_specification.models.get(
+            self.asset_name
+        ).description.replace("\n", "\n\n")
+        info_description = self.data_contract_specification.info.description.replace(
+            "\n", "\n\n"
+        )
+
+        if model_description and info_description:
+            return f"{model_description}\n\n{info_description}"
+        elif model_description:
+            return textwrap.dedent(model_description)
+        elif info_description:
+            return textwrap.dedent(info_description)
+
+        return None
+
+    def _load_owner(self) -> list[str] | None:
+        owner = self.data_contract_specification.info.owner
+
+        return [f"team:{owner}"] if owner else None
+
+    def _load_version(self) -> str | None:
+        version = self.data_contract_specification.info.version
+
+        return version
+
+    def _load_cron_schedule(self) -> str | None:
+        try:
+            cron_schedule = (
+                self.data_contract_specification.servicelevels.frequency.cron
+            )
+            return cron_schedule
+        except AttributeError:
+            return None
+
+    def load_data_quality_checks(self) -> dg.AssetChecksDefinition:
+        @dg.asset_check(
+            asset=self.asset_key,
+            blocking=True,
+        )
+        def check_asset():
+            run = self.data_contract.test()
+
+            return dg.AssetCheckResult(
+                passed=run.result == ResultEnum.passed,
+                metadata={
+                    "quality check": run.pretty(),
+                },
+            )
+
+        return check_asset
+
+    def load_freshness_checks(self, lower_bound_delta: timedelta):
+        freshness_checks = dg.build_last_update_freshness_checks(
+            assets=[self.asset_name],
+            lower_bound_delta=lower_bound_delta,
+            deadline_cron=self.cron_schedule,
+        )
+
+        return freshness_checks

dagster_datacontract-0.1.0.dist-info/METADATA

@@ -0,0 +1,102 @@
+Metadata-Version: 2.4
+Name: dagster-datacontract
+Version: 0.1.0
+Summary: Load metadata and asset check spesifications from data contracts.
+Author-email: Fredrik Bakken <fredrik@dataheim.io>
+Requires-Python: >=3.11.10
+Description-Content-Type: text/markdown
+Requires-Dist: dagster>=1.10.9
+Requires-Dist: dagster-webserver>=1.10.9
+Requires-Dist: datacontract-cli>=0.10.23
+
+# dagster-datacontract
+
+[Dagster](https://github.com/dagster-io/dagster) [Datacontract](https://github.com/datacontract/datacontract-cli) implementation for loading asset metadata and context information from data contracts to Dagster assets.
+
+## Installation
+
+Install `dagster-datacontract` using [`uv`](https://github.com/astral-sh/uv):
+
+```shell
+uv add dagster-datacontract
+```
+
+## Simple Example
+
+The following example can be found inside the [example](https://github.com/dataheim-io/dagster-datacontract/tree/main/example) directory:
+
+```python
+from datetime import timedelta
+from pathlib import Path
+
+import dagster as dg
+import requests
+import polars as pl
+from datacontract.data_contract import DataContract
+
+from dagster_datacontract import DataContractLoader
+
+
+asset_name = "yellow_taxi_trip_records"
+data_contract = DataContractLoader(
+    asset_name=asset_name,
+    data_contract=DataContract(
+        data_contract_file="./example/datacontract.yml",
+    ),
+)
+
+
+@dg.asset(
+    name=asset_name,
+    metadata=data_contract.metadata,
+    tags=data_contract.tags,
+    description=data_contract.description,
+    owners=data_contract.owner,
+    code_version=data_contract.version,
+)
+def yellow_taxi_trip_records(
+    context: dg.AssetExecutionContext,
+) -> None:
+    download_path = "./example/data"
+    Path(download_path).mkdir(parents=True, exist_ok=True)
+
+    url = "https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2025-01.parquet"
+    response = requests.get(url=url)
+
+    file_path = f"{download_path}/yellow_tripdata_2025-01.parquet"
+    context.log.info(f"Reading data from '{url}' and writing to '{file_path}'.")
+    with open(file_path, "wb") as f:
+        f.write(response.content)
+
+    df = pl.read_parquet(file_path)
+    context.log.info(f"File contents downloaded:\n{df}")
+
+
+asset_check_yellow_taxi_trip_records = data_contract.load_data_quality_checks()
+
+freshness_checks = data_contract.load_freshness_checks(
+    lower_bound_delta=timedelta(minutes=5)
+)
+freshness_checks_sensor = dg.build_sensor_for_freshness_checks(
+    freshness_checks=freshness_checks,
+    default_status=dg.DefaultSensorStatus.RUNNING,
+)
+
+job = dg.define_asset_job(
+    name="monthly_taxi_trips",
+    selection=[asset_name],
+)
+schedule = dg.ScheduleDefinition(
+    job=job,
+    cron_schedule=data_contract.cron_schedule,
+    default_status=dg.DefaultScheduleStatus.RUNNING,
+)
+
+
+defs = dg.Definitions(
+    assets=[yellow_taxi_trip_records],
+    asset_checks=[asset_check_yellow_taxi_trip_records, *freshness_checks],
+    schedules=[schedule],
+    sensors=[freshness_checks_sensor],
+)
+```

dagster_datacontract-0.1.0.dist-info/RECORD

@@ -0,0 +1,5 @@
+dagster_datacontract/__init__.py,sha256=AzD8f6UK7F4QmvtfUPrgpcvXPtNT-x_bT_CXTiM29OU,4459
+dagster_datacontract-0.1.0.dist-info/METADATA,sha256=O8yiXq6f5G-cN1maR_mI5w1reVP3MTn_Tc3nB_dIAeg,2969
+dagster_datacontract-0.1.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+dagster_datacontract-0.1.0.dist-info/top_level.txt,sha256=_HUQ6OJ50Q0VZxEkdocTtxk1QkJpztb1QY7A0rcvtCE,21
+dagster_datacontract-0.1.0.dist-info/RECORD,,

dagster_datacontract-0.1.0.dist-info/WHEEL

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

dagster_datacontract-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+dagster_datacontract