datahub-yaml-source 0.1.0__py3-none-any.whl

datahub_yaml_source/__init__.py

@@ -0,0 +1,11 @@
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _version
+
+from datahub_yaml_source.yaml_source import YamlSource
+
+__all__ = ["YamlSource"]
+
+try:
+    __version__ = _version("datahub-yaml-source")
+except PackageNotFoundError:  # running from a source tree with no install
+    __version__ = "0.0.0"

datahub_yaml_source/builders/application.py

@@ -0,0 +1,68 @@
+from collections.abc import Iterable
+
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.metadata.schema_classes import (
+    ApplicationLineageClass,
+    ApplicationPropertiesClass,
+    EdgeClass,
+)
+
+from datahub_yaml_source.builders.common import common_aspect_mcps, mcp_workunit
+from datahub_yaml_source.models import ApplicationDoc, ApplicationLineageEdgeDoc
+from datahub_yaml_source.urns import ReferenceIndex, api_urn, application_urn, dataset_urn
+from datahub_yaml_source.yaml_source_report import YamlSourceReport
+
+
+def _application_lineage_edges(
+    edges: list[ApplicationLineageEdgeDoc] | None, direction: str, context: str
+) -> list[EdgeClass] | None:
+    if not edges:
+        return None
+    result = []
+    for i, edge in enumerate(edges):
+        if bool(edge.api) == bool(edge.dataset):
+            raise ValueError(
+                f"{context} applicationLineage.{direction}[{i}] must set "
+                f"exactly one of 'api' or 'dataset'"
+            )
+        if edge.api:
+            destination_urn = api_urn(edge.api)
+        else:
+            assert edge.dataset is not None  # guaranteed by the exactly-one check above
+            destination_urn = dataset_urn(edge.dataset)
+        result.append(EdgeClass(destinationUrn=destination_urn))
+    return result
+
+
+def build_application(
+    doc: ApplicationDoc, index: ReferenceIndex, report: YamlSourceReport
+) -> Iterable[MetadataWorkUnit]:
+    context = f"APPLICATION '{doc.id}'"
+    entity_urn = application_urn(doc.id)
+
+    # Resolve (and validate) applicationLineage before yielding anything --
+    # `get_workunits_internal` streams work units straight to the sink as they're
+    # yielded, so raising after an earlier yield would leave a partially-emitted
+    # entity instead of cleanly skipping the whole malformed document (see
+    # `build_document`'s equivalent native/external check for the same reason).
+    input_edges = output_edges = None
+    if doc.applicationLineage:
+        input_edges = _application_lineage_edges(
+            doc.applicationLineage.consumes, "consumes", context
+        )
+        output_edges = _application_lineage_edges(
+            doc.applicationLineage.produces, "produces", context
+        )
+
+    yield mcp_workunit(
+        entity_urn,
+        ApplicationPropertiesClass(name=doc.name, description=doc.description),
+    )
+
+    if input_edges or output_edges:
+        yield mcp_workunit(
+            entity_urn,
+            ApplicationLineageClass(inputEdges=input_edges, outputEdges=output_edges),
+        )
+
+    yield from common_aspect_mcps(entity_urn, doc, index, report, context)

datahub_yaml_source/builders/assertion.py

@@ -0,0 +1,254 @@
+import re
+from collections.abc import Iterable
+from typing import Any
+
+from datahub.emitter.mce_builder import make_schema_field_urn
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.metadata.schema_classes import (
+    AssertionActionClass,
+    AssertionActionsClass,
+    AssertionInfoClass,
+    AssertionNoteClass,
+    AssertionSourceClass,
+    AssertionStdParameterClass,
+    AssertionStdParametersClass,
+    AuditStampClass,
+    CustomAssertionInfoClass,
+    FieldAssertionInfoClass,
+    FieldMetricAssertionClass,
+    FieldValuesAssertionClass,
+    FieldValuesFailThresholdClass,
+    FreshnessAssertionInfoClass,
+    FreshnessAssertionScheduleClass,
+    FreshnessCronScheduleClass,
+    OtherSchemaClass,
+    RowCountChangeClass,
+    RowCountTotalClass,
+    SchemaAssertionInfoClass,
+    SchemaFieldClass,
+    SchemaFieldSpecClass,
+    SchemaMetadataClass,
+    SqlAssertionInfoClass,
+    VolumeAssertionInfoClass,
+)
+from datahub.metadata.urns import AssertionUrn
+
+from datahub_yaml_source.builders.common import (
+    DEFAULT_ACTOR_URN,
+    common_aspect_mcps,
+    mcp_workunit,
+    schema_field_data_type,
+    stringify_custom_properties,
+)
+from datahub_yaml_source.models import AssertionActionsDoc, AssertionDoc
+from datahub_yaml_source.urns import ReferenceIndex
+from datahub_yaml_source.yaml_source_report import YamlSourceReport
+
+
+def _build_freshness_assertion(doc: AssertionDoc) -> FreshnessAssertionInfoClass:
+    a = doc.assertion
+    schedule = FreshnessAssertionScheduleClass(
+        type=a.scheduleType or "CRON",
+        cron=FreshnessCronScheduleClass(cron=a.cron or "", timezone=a.timezone or "UTC"),
+    )
+    return FreshnessAssertionInfoClass(
+        type=a.freshnessType or "DATASET_CHANGE", entity=a.entityUrn, schedule=schedule
+    )
+
+
+def _build_sql_assertion(doc: AssertionDoc) -> SqlAssertionInfoClass:
+    a = doc.assertion
+    return SqlAssertionInfoClass(
+        type=a.sqlType or "METRIC",
+        entity=a.entityUrn,
+        statement=a.statement or "",
+        operator=a.operator or "GREATER_THAN",
+        parameters=AssertionStdParametersClass(
+            value=AssertionStdParameterClass(value=str(a.value), type="NUMBER")
+        ),
+        changeType=a.changeType,
+    )
+
+
+def _build_field_assertion(doc: AssertionDoc) -> FieldAssertionInfoClass:
+    a = doc.assertion
+    field_spec = SchemaFieldSpecClass(
+        path=a.fieldPath or "",
+        type=a.dataType or "STRING",
+        nativeType=a.nativeDataType or "",
+    )
+
+    field_values_assertion = None
+    field_metric_assertion = None
+    if a.fieldType == "FIELD_METRIC":
+        field_metric_assertion = FieldMetricAssertionClass(
+            field=field_spec,
+            metric=a.metric or "UNIQUE_PERCENTAGE",
+            operator=a.metricOperator or "EQUAL_TO",
+            parameters=(
+                AssertionStdParametersClass(
+                    value=AssertionStdParameterClass(value=str(a.metricValue), type="NUMBER")
+                )
+                if a.metricValue is not None
+                else None
+            ),
+        )
+    else:
+        field_values_assertion = FieldValuesAssertionClass(
+            field=field_spec,
+            operator=a.operator or "NOT_NULL",
+            failThreshold=FieldValuesFailThresholdClass(type="COUNT", value=0),
+            excludeNulls=a.excludeNulls,
+        )
+
+    return FieldAssertionInfoClass(
+        type=a.fieldType or "FIELD_VALUES",
+        entity=a.entityUrn,
+        fieldValuesAssertion=field_values_assertion,
+        fieldMetricAssertion=field_metric_assertion,
+    )
+
+
+_DATASET_URN_PLATFORM_RE = re.compile(r"^urn:li:dataset:\((?P<platform>urn:li:dataPlatform:[^,]+),")
+
+
+def _platform_urn_from_dataset_urn(entity_urn: str) -> str:
+    match = _DATASET_URN_PLATFORM_RE.match(entity_urn)
+    return match.group("platform") if match else "urn:li:dataPlatform:unknown"
+
+
+def _build_volume_assertion(doc: AssertionDoc) -> VolumeAssertionInfoClass:
+    a = doc.assertion
+    if a.value is None:
+        raise ValueError(f"VOLUME assertion '{doc.id}' is missing a 'value' to compare against.")
+    # `RowCountTotalClass.parameters` / `RowCountChangeClass.parameters` are
+    # required in the Avro schema despite the generated stub typing them
+    # Optional -- passing None parses fine but fails MCP validation at emit
+    # time (same trap as `DeprecationClass.note`, see build_deprecation_aspect()).
+    parameters = AssertionStdParametersClass(
+        value=AssertionStdParameterClass(value=str(a.value), type="NUMBER")
+    )
+    if a.volumeType == "ROW_COUNT_CHANGE":
+        return VolumeAssertionInfoClass(
+            type="ROW_COUNT_CHANGE",
+            entity=a.entityUrn,
+            rowCountChange=RowCountChangeClass(
+                type=a.changeType or "ABSOLUTE",
+                operator=a.operator or "GREATER_THAN_OR_EQUAL_TO",
+                parameters=parameters,
+            ),
+        )
+    return VolumeAssertionInfoClass(
+        type="ROW_COUNT_TOTAL",
+        entity=a.entityUrn,
+        rowCountTotal=RowCountTotalClass(
+            operator=a.operator or "GREATER_THAN_OR_EQUAL_TO", parameters=parameters
+        ),
+    )
+
+
+def _build_data_schema_assertion(doc: AssertionDoc) -> SchemaAssertionInfoClass:
+    a = doc.assertion
+    fields = [
+        SchemaFieldClass(
+            fieldPath=f.path,
+            type=schema_field_data_type(f.type),
+            nativeDataType=f.nativeType or f.type,
+        )
+        for f in (a.schemaFields or [])
+    ]
+    expected_schema = SchemaMetadataClass(
+        schemaName=f"{doc.id}-expected-schema",
+        platform=_platform_urn_from_dataset_urn(a.entityUrn),
+        version=0,
+        hash="",
+        platformSchema=OtherSchemaClass(rawSchema=""),
+        fields=fields,
+    )
+    return SchemaAssertionInfoClass(
+        entity=a.entityUrn, schema=expected_schema, compatibility=a.compatibility
+    )
+
+
+def _build_custom_assertion(doc: AssertionDoc) -> CustomAssertionInfoClass:
+    a = doc.assertion
+    # `CustomAssertionInfoClass.field` is a schemaField URN string, not a
+    # SchemaFieldSpec -- build the URN from the dataset URN + the column path.
+    field_urn = make_schema_field_urn(a.entityUrn, a.fieldPath) if a.fieldPath else None
+    return CustomAssertionInfoClass(
+        type=a.customType or "CUSTOM", entity=a.entityUrn, field=field_urn, logic=a.logic or ""
+    )
+
+
+_BUILDERS = {
+    "FRESHNESS": (_build_freshness_assertion, "freshnessAssertion"),
+    "VOLUME": (_build_volume_assertion, "volumeAssertion"),
+    "SQL": (_build_sql_assertion, "sqlAssertion"),
+    "FIELD": (_build_field_assertion, "fieldAssertion"),
+    "DATA_SCHEMA": (_build_data_schema_assertion, "schemaAssertion"),
+    "CUSTOM": (_build_custom_assertion, "customAssertion"),
+}
+
+
+def _build_assertion_actions(actions: AssertionActionsDoc | None) -> AssertionActionsClass | None:
+    if actions is None:
+        return None
+    return AssertionActionsClass(
+        onSuccess=[AssertionActionClass(type=a.type) for a in (actions.onSuccess or [])] or None,
+        onFailure=[AssertionActionClass(type=a.type) for a in (actions.onFailure or [])] or None,
+    )
+
+
+def build_assertion(
+    doc: AssertionDoc, index: ReferenceIndex, report: YamlSourceReport
+) -> Iterable[MetadataWorkUnit]:
+    builder = _BUILDERS.get(doc.assertion.type)
+    if builder is None:
+        raise ValueError(
+            f"Unsupported assertion type '{doc.assertion.type}' for assertion '{doc.id}'. "
+            f"Supported types: {sorted(_BUILDERS)}"
+        )
+    build_fn, kwarg_name = builder
+    sub_assertion = build_fn(doc)
+
+    entity_urn = AssertionUrn(doc.id).urn()
+
+    custom_properties = None
+    if doc.properties:
+        custom_properties = stringify_custom_properties(
+            {
+                k: v
+                for k, v in {
+                    "name": doc.properties.name,
+                    "dbt_test": doc.properties.dbt_test,
+                }.items()
+                if v is not None
+            }
+        )
+
+    # The sub-assertion kwarg name is chosen from `_BUILDERS` at runtime; `Any`
+    # is the honest element type for that dispatch.
+    sub_assertion_kwarg: dict[str, Any] = {kwarg_name: sub_assertion}
+    aspect = AssertionInfoClass(
+        type=doc.assertion.type,
+        description=doc.description,
+        customProperties=custom_properties,
+        source=AssertionSourceClass(type=doc.sourceType),
+        **sub_assertion_kwarg,
+    )
+    yield mcp_workunit(entity_urn, aspect)
+
+    if doc.assertionNote:
+        yield mcp_workunit(
+            entity_urn,
+            AssertionNoteClass(
+                content=doc.assertionNote,
+                lastModified=AuditStampClass(time=0, actor=DEFAULT_ACTOR_URN),
+            ),
+        )
+
+    actions_aspect = _build_assertion_actions(doc.assertionActions)
+    if actions_aspect:
+        yield mcp_workunit(entity_urn, actions_aspect)
+
+    yield from common_aspect_mcps(entity_urn, doc, index, report, f"ASSERTION '{doc.id}'")

datahub_yaml_source/builders/chart.py

@@ -0,0 +1,41 @@
+from collections.abc import Iterable
+
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.sdk.chart import Chart
+
+from datahub_yaml_source.builders.common import common_sdk_kwargs, stringify_custom_properties
+from datahub_yaml_source.models import ChartDoc
+from datahub_yaml_source.urns import ReferenceIndex, container_key, dataset_urn
+from datahub_yaml_source.yaml_source_report import YamlSourceReport
+
+
+def build_chart(
+    doc: ChartDoc, index: ReferenceIndex, report: YamlSourceReport
+) -> Iterable[MetadataWorkUnit]:
+    context = f"CHART '{doc.name}'"
+    common = common_sdk_kwargs(doc, index, report, context)
+
+    # `parent_container=` defaults to a private `Unset` sentinel, not `None`
+    # -- passing `None` explicitly still triggers `_set_container(None)`,
+    # which emits a (harmless but unwanted-here) empty `browsePathsV2`. Only
+    # pass the kwarg at all when there's an actual container (same trap as
+    # `build_dataset()`/`build_data_flow()`).
+    if doc.container is not None:
+        if not index.has_container(doc.container):
+            report.report_dangling_reference(f"{context} references an undeclared container")
+        common["parent_container"] = container_key(doc.container)
+
+    chart = Chart(
+        platform=doc.platform,
+        name=doc.name,
+        platform_instance=doc.instance,
+        display_name=doc.displayName,
+        description=doc.description,
+        external_url=doc.externalUrl,
+        chart_url=doc.chartUrl,
+        chart_type=doc.chartType,
+        custom_properties=stringify_custom_properties(doc.properties),
+        input_datasets=[dataset_urn(ref) for ref in (doc.inputDatasets or [])] or None,
+        **common,
+    )
+    yield from chart.as_workunits()