From 573d28c165ee39735df3245498b541b0dc6a5d34 Mon Sep 17 00:00:00 2001 From: Matthew Moorcroft Date: Thu, 10 Sep 2026 15:59:37 +0100 Subject: [PATCH 1/7] Add shared source-agnostic lineage substrate (#61) Introduces the source-neutral lineage substrate (PR1 of the deterministic- lineage plan). Substrate only: no source is wired, no translator or extractor is touched, inventory.json is untouched (follow-ups #62/#63). - models/ir.py: new source-neutral IR types DataAsset, ControlEdge, DataEdge, Lineage, MotifAnnotation. DataEdge carries the two-tier join tags from #36 (match_kind identity|signature, match_key; identity nullable). Field vocabulary stays neutral for both ADF (pipelines/ExecutePipeline) and Airflow (DAGs/RunJobActivity) -- no producer_pipeline/producer_activity. Adds data_reads / data_writes / motif_id (all defaulted) to the Activity base, and a lineage block to Pipeline. - R1: motif_id now lives on the Activity base; MotifActivity redeclares it as required so it stays typed str and existing callers are unaffected. Rehydration drops the explicit motif_id= kwarg (it arrives via _common_activity_kwargs), avoiding the "multiple values for keyword argument 'motif_id'" TypeError. A MotifActivity round-trip test proves it. - ir_serde.py: serialises the new base fields and a top-level lineage block; dab_writer.py rehydrates them (data assets, edges, motifs) through the existing round-trip so nothing silently round-trips to None. - lineage.py: pure build_control_edges / build_data_edges (+ build_lineage, build_motif_annotations, with_lineage) over IR primitives only, no source imports. Handles fan-out, nested/Switch recursion, identity-vs-signature tiers, and emits no self-edges and no duplicate edges. Nothing is mutated; with_lineage returns a new Pipeline. - tests: unit coverage for the derivation and ir_serde round-trips (incl. the MotifActivity R1 case). Co-authored-by: Isaac --- src/flowx/bundler/dab_writer.py | 66 +++- src/flowx/ir_serde.py | 68 +++- src/flowx/lineage.py | 251 +++++++++++++++ src/flowx/models/ir.py | 142 +++++++++ tests/unit/test_lineage_substrate.py | 454 +++++++++++++++++++++++++++ 5 files changed, 979 insertions(+), 2 deletions(-) create mode 100644 src/flowx/lineage.py create mode 100644 tests/unit/test_lineage_substrate.py diff --git a/src/flowx/bundler/dab_writer.py b/src/flowx/bundler/dab_writer.py index 436900d..8cdf20b 100644 --- a/src/flowx/bundler/dab_writer.py +++ b/src/flowx/bundler/dab_writer.py @@ -30,7 +30,10 @@ from flowx.models.ir import ( Activity, AppendVariableActivity, + ControlEdge, CopyActivity, + DataAsset, + DataEdge, DbtFactoryActivity, DeleteActivity, Dependency, @@ -38,8 +41,10 @@ FilterActivity, ForEachActivity, IfConditionActivity, + Lineage, LookupActivity, MotifActivity, + MotifAnnotation, NotebookActivity, Pipeline, PlaceholderActivity, @@ -2037,6 +2042,7 @@ def pipeline_dict_to_ir(pipeline_dict: dict[str, Any]) -> tuple[Pipeline, list[d reconciliation_status=pipeline_dict.get("reconciliation_status"), migration_status=pipeline_dict.get("migration_status", "included"), audit=dict(pipeline_dict.get("audit") or {}), + lineage=_reconstruct_lineage(pipeline_dict.get("lineage")), ) return pipeline, parameters @@ -2258,9 +2264,10 @@ def _reconstruct_ir(task_ir: dict[str, Any]) -> Activity: bridge_required_parameters=dict(task_ir.get("bridge_required_parameters") or {}), ) if task_type == "MotifActivity": + # motif_id arrives via ``base`` (Activity now owns the field); passing it again here + # would raise "multiple values for keyword argument 'motif_id'". return MotifActivity( **base, - motif_id=task_ir.get("motif_id", "unknown"), display_name=task_ir.get("display_name", base["name"]), databricks_replacement=task_ir.get("databricks_replacement", "notebook"), matched_activity_names=list(task_ir.get("matched_activity_names", [])), @@ -2311,9 +2318,66 @@ def _common_activity_kwargs(task_ir: dict[str, Any]) -> dict[str, Any]: "required_parameters": dict(task_ir.get("required_parameters") or {}), "compute_mode": task_ir.get("compute_mode"), "notifications": task_ir.get("notifications"), + "motif_id": task_ir.get("motif_id"), + "data_reads": _reconstruct_data_assets(task_ir.get("data_reads")), + "data_writes": _reconstruct_data_assets(task_ir.get("data_writes")), } +def _reconstruct_data_assets(raw: list[dict[str, Any]] | None) -> list[DataAsset]: + """Rehydrates serialised DataAsset dicts into typed :class:`DataAsset` nodes.""" + if not raw: + return [] + return [ + DataAsset( + signature=asset.get("signature", ""), + identity=asset.get("identity"), + asset_type=asset.get("asset_type"), + properties=dict(asset.get("properties") or {}), + ) + for asset in raw + ] + + +def _reconstruct_lineage(raw: dict[str, Any] | None) -> Lineage | None: + """Rehydrates a serialised lineage block into a typed :class:`Lineage`, or ``None``.""" + if not raw: + return None + return Lineage( + control_edges=[ + ControlEdge( + source_workflow=edge.get("source_workflow", ""), + target_workflow=edge.get("target_workflow", ""), + via_task_key=edge.get("via_task_key", ""), + wait_for_completion=edge.get("wait_for_completion"), + resolved=bool(edge.get("resolved", True)), + ) + for edge in raw.get("control_edges") or [] + ], + data_edges=[ + DataEdge( + source_task_key=edge.get("source_task_key", ""), + target_task_key=edge.get("target_task_key", ""), + match_kind=edge.get("match_kind", ""), + match_key=edge.get("match_key", ""), + identity=edge.get("identity"), + asset_type=edge.get("asset_type"), + ) + for edge in raw.get("data_edges") or [] + ], + motifs=[ + MotifAnnotation( + motif_id=motif.get("motif_id", ""), + member_task_keys=list(motif.get("member_task_keys") or []), + display_name=motif.get("display_name"), + databricks_replacement=motif.get("databricks_replacement"), + notes=list(motif.get("notes") or []), + ) + for motif in raw.get("motifs") or [] + ], + ) + + def _reconstruct_dependencies(raw: list[dict[str, Any]] | None) -> list[Dependency] | None: if not raw: return None diff --git a/src/flowx/ir_serde.py b/src/flowx/ir_serde.py index 51b3d32..c49ba6a 100644 --- a/src/flowx/ir_serde.py +++ b/src/flowx/ir_serde.py @@ -23,12 +23,14 @@ Activity, AppendVariableActivity, CopyActivity, + DataAsset, DbtFactoryActivity, DeleteActivity, ExecutePipelineActivity, FilterActivity, ForEachActivity, IfConditionActivity, + Lineage, LookupActivity, MotifActivity, NotebookActivity, @@ -80,9 +82,67 @@ def pipeline_to_dict(pipeline: Pipeline) -> dict[str, Any]: } if pipeline.translation_configuration is not None: result["translation_configuration"] = configuration_to_dict(pipeline.translation_configuration) + if pipeline.lineage is not None: + result["lineage"] = lineage_to_dict(pipeline.lineage) return result +def data_asset_to_dict(asset: DataAsset) -> dict[str, Any]: + """Serialise a :class:`DataAsset` to a JSON-friendly dictionary. + + ``signature`` always appears; the optional ``identity`` / ``asset_type`` / + ``properties`` are emitted only when set so reports stay compact. + """ + result: dict[str, Any] = {"signature": asset.signature} + if asset.identity is not None: + result["identity"] = asset.identity + if asset.asset_type is not None: + result["asset_type"] = asset.asset_type + if asset.properties: + result["properties"] = dict(asset.properties) + return result + + +def lineage_to_dict(lineage: Lineage) -> dict[str, Any]: + """Serialise a :class:`Lineage` block to a JSON-friendly dictionary. + + Edge lists are always emitted (empty, never ``None``) for stable diffs. + """ + return { + "control_edges": [ + { + "source_workflow": edge.source_workflow, + "target_workflow": edge.target_workflow, + "via_task_key": edge.via_task_key, + "wait_for_completion": edge.wait_for_completion, + "resolved": edge.resolved, + } + for edge in lineage.control_edges + ], + "data_edges": [ + { + "source_task_key": edge.source_task_key, + "target_task_key": edge.target_task_key, + "match_kind": edge.match_kind, + "match_key": edge.match_key, + "identity": edge.identity, + "asset_type": edge.asset_type, + } + for edge in lineage.data_edges + ], + "motifs": [ + { + "motif_id": motif.motif_id, + "member_task_keys": list(motif.member_task_keys), + "display_name": motif.display_name, + "databricks_replacement": motif.databricks_replacement, + "notes": list(motif.notes), + } + for motif in lineage.motifs + ], + } + + def configuration_to_dict(configuration: Any) -> dict[str, Any]: """Serialise a TranslationConfiguration instance to a JSON-friendly dictionary. @@ -143,6 +203,12 @@ def activity_to_dict(task: Activity) -> dict[str, Any]: task_dict["libraries"] = task.libraries if task.parameter_approximations: task_dict["parameter_approximations"] = task.parameter_approximations + if task.motif_id: + task_dict["motif_id"] = task.motif_id + if task.data_reads: + task_dict["data_reads"] = [data_asset_to_dict(asset) for asset in task.data_reads] + if task.data_writes: + task_dict["data_writes"] = [data_asset_to_dict(asset) for asset in task.data_writes] extra = activity_extra_fields(task) task_dict.update(extra) @@ -337,7 +403,7 @@ def activity_extra_fields(activity: Activity) -> dict[str, Any]: if activity.job_parameters: extra["job_parameters"] = activity.job_parameters case MotifActivity(): - extra["motif_id"] = activity.motif_id + # motif_id now lives on the Activity base and is serialised by activity_to_dict. extra["display_name"] = activity.display_name extra["databricks_replacement"] = activity.databricks_replacement extra["matched_activity_names"] = activity.matched_activity_names diff --git a/src/flowx/lineage.py b/src/flowx/lineage.py new file mode 100644 index 0000000..a417718 --- /dev/null +++ b/src/flowx/lineage.py @@ -0,0 +1,251 @@ +"""Source-neutral lineage derivation over the flowx Pipeline IR. + +These functions turn an already-translated :class:`~flowx.models.ir.Pipeline` +into its :class:`~flowx.models.ir.Lineage` block. They operate on IR primitives +only -- ``Activity`` subclasses, ``DataAsset``, ``task_key`` -- and import nothing +from ``sources/adf`` or ``sources/airflow`` so both front-ends share one code +path once they populate ``data_reads`` / ``data_writes`` / ``motif_id``. + +Everything here is pure: the functions read the pipeline and return new edge +lists / a new :class:`Lineage`; nothing is mutated. :func:`with_lineage` attaches +a block by returning a *new* ``Pipeline`` rather than mutating the input, unlike +the in-place dependency rewrite in ``motifs/collapser.py``. +""" + +from __future__ import annotations + +import dataclasses +from collections.abc import Iterator + +from flowx.models.ir import ( + Activity, + ControlEdge, + DataAsset, + DataEdge, + ExecutePipelineActivity, + ForEachActivity, + IfConditionActivity, + Lineage, + MotifActivity, + MotifAnnotation, + Pipeline, + RunJobActivity, + SwitchActivity, +) + + +def walk_activities(activities: list[Activity]) -> Iterator[Activity]: + """Yield every activity in *activities*, descending into control-flow containers. + + Recurses into ForEach inner activities, both If-condition branches, and every + Switch case plus its default branch, so a nested ExecutePipeline or a data + asset buried inside a Switch case is still reached. Motif ``original_activities`` + are intentionally not traversed: they are the pre-collapse originals kept for + reference, not live graph members. + + Args: + activities: Top-level (or already-nested) activity list to walk. + + Yields: + Each activity, container nodes included, in depth-first order. + """ + for activity in activities: + yield activity + match activity: + case ForEachActivity(): + yield from walk_activities(activity.inner_activities) + case IfConditionActivity(): + yield from walk_activities(activity.if_true_activities) + yield from walk_activities(activity.if_false_activities) + case SwitchActivity(): + for case_branch in activity.cases: + yield from walk_activities(case_branch.activities) + yield from walk_activities(activity.default_activities) + + +def build_control_edges(pipeline: Pipeline) -> list[ControlEdge]: + """Derive cross-workflow invocation edges for a pipeline. + + Emits one :class:`ControlEdge` per invoking activity -- an + ``ExecutePipelineActivity`` (ADF) or a ``RunJobActivity`` (Airflow) -- found + anywhere in the pipeline, including inside ForEach / If / Switch containers + (fan-out is preserved: each call site is its own edge). Edges whose callee + equals the caller are dropped (no self-edges), and identical edges are + collapsed (no duplicates). An unresolved callee is recorded with + ``resolved=False`` rather than dropped. + + Args: + pipeline: The translated pipeline IR. + + Returns: + Deduplicated list of control edges, in first-seen order. + """ + edges: list[ControlEdge] = [] + seen: set[tuple[str, str, str]] = set() + for activity in walk_activities(pipeline.tasks): + target: str | None + wait: bool | None + match activity: + case ExecutePipelineActivity(): + target = activity.pipeline_name + wait = activity.wait_on_completion + case RunJobActivity(): + target = activity.job_name + wait = None + case _: + continue + target_name = target or "" + if target_name and target_name == pipeline.name: + continue + key = (pipeline.name, target_name, activity.task_key) + if key in seen: + continue + seen.add(key) + edges.append( + ControlEdge( + source_workflow=pipeline.name, + target_workflow=target_name, + via_task_key=activity.task_key, + wait_for_completion=wait, + resolved=bool(target_name), + ) + ) + return edges + + +def _match_assets(producer: DataAsset, consumer: DataAsset) -> tuple[str, str, str | None] | None: + """Decide whether a written asset hands off to a read asset, and how. + + Two tiers, per #36: + + - **identity** -- when both sides resolved to a physical identity, they match + only if those identities are equal. Two differently-resolved identities do + *not* fall through to a signature match; that is what manufactured the + spurious edges #36 removed. + - **signature** -- when at least one identity is unresolved, fall back to the + neutral descriptor and match on equal, non-empty signatures. + + Returns: + ``(match_kind, match_key, identity)`` when the pair matches, else ``None``. + """ + if producer.identity is not None and consumer.identity is not None: + if producer.identity == consumer.identity: + return "identity", producer.identity, producer.identity + return None + if producer.signature and producer.signature == consumer.signature: + return "signature", producer.signature, producer.identity or consumer.identity + return None + + +def build_data_edges(pipeline: Pipeline) -> list[DataEdge]: + """Derive proven producer -> consumer data hand-offs for a pipeline. + + A producer is any activity with a ``data_writes`` asset; a consumer any + activity with a ``data_reads`` asset, gathered across the whole pipeline + (ForEach / If / Switch bodies included). Each producer asset is joined against + each consumer asset via :func:`_match_assets`, tagging the edge as an + ``identity`` or ``signature`` match. An activity never hands off to itself + (no self-edges), and identical edges are collapsed (no duplicates). + + Args: + pipeline: The translated pipeline IR. + + Returns: + Deduplicated list of data edges, in first-seen order. + """ + activities = list(walk_activities(pipeline.tasks)) + producers = [(activity.task_key, asset) for activity in activities for asset in activity.data_writes] + consumers = [(activity.task_key, asset) for activity in activities for asset in activity.data_reads] + + edges: list[DataEdge] = [] + seen: set[tuple[str, str, str, str]] = set() + for producer_key, producer_asset in producers: + for consumer_key, consumer_asset in consumers: + if producer_key == consumer_key: + continue + matched = _match_assets(producer_asset, consumer_asset) + if matched is None: + continue + match_kind, match_key, identity = matched + dedupe_key = (producer_key, consumer_key, match_kind, match_key) + if dedupe_key in seen: + continue + seen.add(dedupe_key) + edges.append( + DataEdge( + source_task_key=producer_key, + target_task_key=consumer_key, + match_kind=match_kind, + match_key=match_key, + identity=identity, + asset_type=producer_asset.asset_type or consumer_asset.asset_type, + ) + ) + return edges + + +def build_motif_annotations(pipeline: Pipeline) -> list[MotifAnnotation]: + """Derive motif annotations from the collapsed motif activities in a pipeline. + + One annotation per :class:`MotifActivity`, listing the task keys it spans + (the motif task itself plus any member activities that carry the same + ``motif_id`` tag). Deduplicated by ``motif_id`` in first-seen order. + + Args: + pipeline: The translated pipeline IR. + + Returns: + List of motif annotations. + """ + annotations: list[MotifAnnotation] = [] + seen: set[str] = set() + activities = list(walk_activities(pipeline.tasks)) + for activity in activities: + if not isinstance(activity, MotifActivity): + continue + if activity.motif_id in seen: + continue + seen.add(activity.motif_id) + members = [activity.task_key] + members.extend( + other.task_key for other in activities if other is not activity and other.motif_id == activity.motif_id + ) + annotations.append( + MotifAnnotation( + motif_id=activity.motif_id, + member_task_keys=members, + display_name=activity.display_name, + databricks_replacement=activity.databricks_replacement, + notes=list(activity.confidence_notes), + ) + ) + return annotations + + +def build_lineage(pipeline: Pipeline) -> Lineage: + """Compose the full source-neutral lineage block for a pipeline. + + Args: + pipeline: The translated pipeline IR. + + Returns: + A :class:`Lineage` with control edges, data edges, and motif annotations. + """ + return Lineage( + control_edges=build_control_edges(pipeline), + data_edges=build_data_edges(pipeline), + motifs=build_motif_annotations(pipeline), + ) + + +def with_lineage(pipeline: Pipeline, lineage: Lineage) -> Pipeline: + """Return a *new* pipeline carrying *lineage*, leaving the input untouched. + + Args: + pipeline: The pipeline to copy. + lineage: The lineage block to attach. + + Returns: + A shallow copy of *pipeline* with ``lineage`` set. + """ + return dataclasses.replace(pipeline, lineage=lineage) diff --git a/src/flowx/models/ir.py b/src/flowx/models/ir.py index 600e8c7..94def2f 100644 --- a/src/flowx/models/ir.py +++ b/src/flowx/models/ir.py @@ -55,6 +55,134 @@ class Dependency: outcome: str | None = None +@dataclass(slots=True, kw_only=True) +class DataAsset: + """A physical data source or sink an activity reads from or writes to. + + Deliberately source-neutral: the same shape describes an ADF dataset, an + Airflow dataset/hook target, or any other front-end's data reference, so the + lineage substrate never has to know which source produced it. + + Two-tier identity (from #36): ``identity`` is the resolved physical location + (``schema.table`` or a concrete storage path) and is the strong join key. + It is ``None`` when it cannot be resolved deterministically -- never a guess. + ``signature`` is always present and carries the neutral fallback descriptor + (a dataset name, a normalised reference, or an expression) so two assets can + still be compared when neither side resolved to a physical identity. + + Attributes: + signature: Neutral, always-present descriptor used as the weak join key + (e.g. dataset name or normalised expression). + identity: Resolved physical identity used as the strong join key, or + ``None`` when it could not be resolved deterministically. + asset_type: Neutral kind of the asset (``"table"`` / ``"file"`` / + ``"volume"`` / ...), or ``None`` when unknown. + properties: Free-form extra attributes carried through verbatim. + """ + + signature: str + identity: str | None = None + asset_type: str | None = None + properties: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True, kw_only=True) +class ControlEdge: + """A control-flow invocation between two workflows. + + Emitted for cross-workflow calls -- an ADF ``ExecutePipeline`` invoking a + child pipeline, an Airflow ``RunJobActivity`` triggering another job -- so + the field names stay neutral (``source_workflow`` / ``target_workflow``) + rather than baking in either source's vocabulary. + + Attributes: + source_workflow: Name of the calling workflow (the pipeline/DAG the + invoking activity lives in). + target_workflow: Name of the invoked workflow (child pipeline / job). + via_task_key: Task key of the activity that performs the invocation. + wait_for_completion: Whether the caller blocks on the callee, when the + source expresses it; ``None`` when the source has no such notion. + resolved: ``False`` when the callee could not be resolved from a partial + export (recorded, not dropped); ``True`` otherwise. + """ + + source_workflow: str + target_workflow: str + via_task_key: str + wait_for_completion: bool | None = None + resolved: bool = True + + +@dataclass(slots=True, kw_only=True) +class DataEdge: + """A proven data hand-off between a producing and a consuming task. + + A producer writes a :class:`DataAsset` that a consumer later reads. The two + tiers from #36 are recorded on the edge itself: ``match_kind`` says whether + the two assets were joined on resolved physical ``identity`` or on their + neutral ``signature``, and ``match_key`` is the value they matched on. + + Attributes: + source_task_key: Task key of the producer (writes the asset). + target_task_key: Task key of the consumer (reads the asset). + match_kind: ``"identity"`` when joined on resolved physical identity, + ``"signature"`` when joined on the neutral fallback descriptor. + match_key: The value the two assets matched on. + identity: Resolved physical identity of the hand-off, or ``None`` when + the match was signature-only / the identity was unresolvable. + asset_type: Neutral kind of the handed-off asset, when known. + """ + + source_task_key: str + target_task_key: str + match_kind: str + match_key: str + identity: str | None = None + asset_type: str | None = None + + +@dataclass(slots=True, kw_only=True) +class MotifAnnotation: + """Describes a detected motif spanning one or more activities. + + Source-neutral record tying a motif id to the tasks that belong to it, so + the lineage block can report motifs without depending on how any particular + source detects them. Each member activity also carries the same + :attr:`Activity.motif_id` tag. + + Attributes: + motif_id: Identifier of the matched motif definition. + member_task_keys: Task keys of the activities the motif spans. + display_name: Human-readable motif name, if any. + databricks_replacement: Target Databricks construct the motif maps to. + notes: Detector notes explaining the match rationale. + """ + + motif_id: str + member_task_keys: list[str] = field(default_factory=list) + display_name: str | None = None + databricks_replacement: str | None = None + notes: list[str] = field(default_factory=list) + + +@dataclass(slots=True, kw_only=True) +class Lineage: + """Source-neutral lineage block attached to a translated pipeline. + + Edge lists are always concrete lists (empty, never ``None``) so serialised + reports produce stable golden diffs. + + Attributes: + control_edges: Cross-workflow invocation edges. + data_edges: Proven producer -> consumer data hand-offs. + motifs: Detected motif annotations. + """ + + control_edges: list[ControlEdge] = field(default_factory=list) + data_edges: list[DataEdge] = field(default_factory=list) + motifs: list[MotifAnnotation] = field(default_factory=list) + + @dataclass(slots=True, kw_only=True) class Activity: """Base class for all translated pipeline activities. @@ -94,6 +222,13 @@ class Activity: compute_mode: str | None = None # Collapsed activity_and_notify spec set by the adapter: {destination, events, args, destination_name}. notifications: dict[str, Any] | None = None + # Lineage substrate (#61): source-neutral data assets this activity reads from and writes to, + # populated by per-source extractors in follow-up work. Always lists, never None. + data_reads: list[DataAsset] = field(default_factory=list) + data_writes: list[DataAsset] = field(default_factory=list) + # Id of the motif this activity was folded into (or belongs to); None when it is part of no motif. + # Owned here so every activity type -- not only MotifActivity -- can carry the tag. + motif_id: str | None = None @dataclass(slots=True, kw_only=True) @@ -570,6 +705,10 @@ class PlaceholderActivity(Activity): class MotifActivity(Activity): """Activity produced by collapsing a detected motif pattern. + Redeclares :attr:`Activity.motif_id` as required (the base owns the field so + every activity type can carry the tag and it round-trips through one code + path, but a motif activity always has one). + Attributes: motif_id: Identifier of the matched motif definition. display_name: Human-readable motif name. @@ -624,6 +763,8 @@ class Pipeline: reconciliation_status: Source-audit result for this pipeline. migration_status: Whether the pipeline is included or explicitly excluded. audit: Source-audit counts and transformation ledger. + lineage: Source-neutral lineage block (control/data edges + motif + annotations), or ``None`` when lineage has not been derived. """ name: str @@ -640,6 +781,7 @@ class Pipeline: audit: dict[str, Any] = field(default_factory=dict) translation_configuration: TranslationConfiguration | None = None bundle_variables: dict[str, dict[str, Any]] = field(default_factory=dict) + lineage: Lineage | None = None @dataclass(frozen=True, slots=True) diff --git a/tests/unit/test_lineage_substrate.py b/tests/unit/test_lineage_substrate.py new file mode 100644 index 0000000..1f1d8a8 --- /dev/null +++ b/tests/unit/test_lineage_substrate.py @@ -0,0 +1,454 @@ +"""Unit tests for the source-neutral lineage substrate (#61). + +Covers the pure derivation (:mod:`flowx.lineage`) -- control fan-out, nested and +Switch recursion, the identity-vs-signature join tiers, no self-edges, no +duplicate edges -- and the ``ir_serde`` round-trip for the new IR types plus the +new ``Activity`` base fields, including a MotifActivity round-trip that proves the +R1 ``motif_id`` collision is handled. +""" + +from __future__ import annotations + +import json + +from flowx.bundler.dab_writer import pipeline_dict_to_ir +from flowx.ir_serde import pipeline_to_dict +from flowx.lineage import ( + build_control_edges, + build_data_edges, + build_lineage, + build_motif_annotations, + with_lineage, +) +from flowx.models.ir import ( + ControlEdge, + DataAsset, + DataEdge, + ExecutePipelineActivity, + ForEachActivity, + IfConditionActivity, + Lineage, + MotifActivity, + MotifAnnotation, + NotebookActivity, + Pipeline, + RunJobActivity, + SwitchActivity, + SwitchCase, + WaitActivity, +) + + +def _notebook(task_key: str, *, reads=None, writes=None, motif_id=None) -> NotebookActivity: + return NotebookActivity( + name=task_key, + task_key=task_key, + notebook_path=f"/Shared/{task_key}", + data_reads=list(reads or []), + data_writes=list(writes or []), + motif_id=motif_id, + ) + + +def _execute(task_key: str, callee: str, *, wait: bool = True) -> ExecutePipelineActivity: + return ExecutePipelineActivity(name=task_key, task_key=task_key, pipeline_name=callee, wait_on_completion=wait) + + +# --------------------------------------------------------------------------- # +# Control-edge derivation +# --------------------------------------------------------------------------- # + + +def test_control_edges_fan_out_and_nested_switch_recursion(): + """ExecutePipeline calls are found at top level and inside ForEach/If/Switch.""" + pipeline = Pipeline( + name="parent", + tasks=[ + _execute("call_a", "child_a"), + ForEachActivity( + name="fe", + task_key="fe", + items_expression="@x", + inner_activities=[_execute("call_b", "child_b")], + ), + IfConditionActivity( + name="cond", + task_key="cond", + op="equals", + left="@a", + right="@b", + if_true_activities=[_execute("call_c", "child_c")], + if_false_activities=[_execute("call_d", "child_d")], + ), + SwitchActivity( + name="sw", + task_key="sw", + on_expression="@e", + cases=[SwitchCase(value="one", activities=[_execute("call_e", "child_e")])], + default_activities=[_execute("call_f", "child_f")], + ), + ], + ) + + edges = build_control_edges(pipeline) + + targets = sorted(edge.target_workflow for edge in edges) + assert targets == ["child_a", "child_b", "child_c", "child_d", "child_e", "child_f"] + assert all(edge.source_workflow == "parent" for edge in edges) + # Each call site keeps its own via_task_key (fan-out preserved). + assert {edge.via_task_key for edge in edges} == { + "call_a", + "call_b", + "call_c", + "call_d", + "call_e", + "call_f", + } + + +def test_control_edges_run_job_activity_is_source_neutral(): + """A RunJobActivity (Airflow) produces a control edge just like ExecutePipeline.""" + pipeline = Pipeline( + name="dag_main", + tasks=[RunJobActivity(name="run", task_key="run", job_name="downstream_job")], + ) + + edges = build_control_edges(pipeline) + + assert len(edges) == 1 + assert edges[0].source_workflow == "dag_main" + assert edges[0].target_workflow == "downstream_job" + assert edges[0].via_task_key == "run" + assert edges[0].wait_for_completion is None + assert edges[0].resolved is True + + +def test_control_edges_unresolved_callee_is_recorded_not_dropped(): + """An empty callee is kept with resolved=False rather than silently dropped.""" + pipeline = Pipeline(name="parent", tasks=[_execute("call", "")]) + + edges = build_control_edges(pipeline) + + assert len(edges) == 1 + assert edges[0].target_workflow == "" + assert edges[0].resolved is False + + +def test_control_edges_no_self_edge(): + """A pipeline invoking itself produces no edge.""" + pipeline = Pipeline(name="loop", tasks=[_execute("call", "loop")]) + + assert build_control_edges(pipeline) == [] + + +def test_control_edges_no_duplicate_from_recursion(): + """A single call site nested in a container is emitted exactly once.""" + pipeline = Pipeline( + name="parent", + tasks=[ + ForEachActivity( + name="fe", + task_key="fe", + items_expression="@x", + inner_activities=[_execute("call", "child")], + ) + ], + ) + + edges = build_control_edges(pipeline) + + assert len(edges) == 1 + assert edges[0].target_workflow == "child" + + +# --------------------------------------------------------------------------- # +# Data-edge derivation: identity vs signature tiers +# --------------------------------------------------------------------------- # + + +def test_data_edges_identity_tier_joins_across_different_signatures(): + """Two assets with the same resolved identity match even when their names differ.""" + pipeline = Pipeline( + name="p", + tasks=[ + _notebook("writer", writes=[DataAsset(signature="ds_out", identity="curated.orders")]), + _notebook("reader", reads=[DataAsset(signature="ds_in_other_name", identity="curated.orders")]), + ], + ) + + edges = build_data_edges(pipeline) + + assert len(edges) == 1 + assert edges[0].source_task_key == "writer" + assert edges[0].target_task_key == "reader" + assert edges[0].match_kind == "identity" + assert edges[0].match_key == "curated.orders" + assert edges[0].identity == "curated.orders" + + +def test_data_edges_signature_tier_when_identity_unresolved(): + """When identity is unresolvable, matching falls back to the neutral signature.""" + pipeline = Pipeline( + name="p", + tasks=[ + _notebook("writer", writes=[DataAsset(signature="shared_ds")]), + _notebook("reader", reads=[DataAsset(signature="shared_ds")]), + ], + ) + + edges = build_data_edges(pipeline) + + assert len(edges) == 1 + assert edges[0].match_kind == "signature" + assert edges[0].match_key == "shared_ds" + assert edges[0].identity is None + + +def test_data_edges_distinct_identities_do_not_fall_back_to_signature(): + """Two resolved-but-different identities never manufacture a signature edge (#36).""" + pipeline = Pipeline( + name="p", + tasks=[ + _notebook("writer", writes=[DataAsset(signature="shared", identity="a.first")]), + _notebook("reader", reads=[DataAsset(signature="shared", identity="b.second")]), + ], + ) + + assert build_data_edges(pipeline) == [] + + +def test_data_edges_fan_out_one_writer_many_readers(): + """One producer handing off to several consumers yields one edge each.""" + pipeline = Pipeline( + name="p", + tasks=[ + _notebook("writer", writes=[DataAsset(signature="ds", identity="x.y")]), + _notebook("reader_one", reads=[DataAsset(signature="ds", identity="x.y")]), + _notebook("reader_two", reads=[DataAsset(signature="ds", identity="x.y")]), + ], + ) + + edges = build_data_edges(pipeline) + + assert sorted(edge.target_task_key for edge in edges) == ["reader_one", "reader_two"] + assert all(edge.source_task_key == "writer" for edge in edges) + + +def test_data_edges_no_self_edge(): + """An activity that both writes and reads the same asset does not edge to itself.""" + pipeline = Pipeline( + name="p", + tasks=[ + _notebook( + "roundtrip", + writes=[DataAsset(signature="ds", identity="x.y")], + reads=[DataAsset(signature="ds", identity="x.y")], + ) + ], + ) + + assert build_data_edges(pipeline) == [] + + +def test_data_edges_no_duplicate_from_repeated_asset(): + """A producer listing the same asset twice still yields a single edge.""" + pipeline = Pipeline( + name="p", + tasks=[ + _notebook( + "writer", + writes=[DataAsset(signature="ds", identity="x.y"), DataAsset(signature="ds", identity="x.y")], + ), + _notebook("reader", reads=[DataAsset(signature="ds", identity="x.y")]), + ], + ) + + edges = build_data_edges(pipeline) + + assert len(edges) == 1 + + +def test_data_edges_nested_switch_recursion(): + """A producer buried in a Switch case hands off to a top-level consumer.""" + pipeline = Pipeline( + name="p", + tasks=[ + SwitchActivity( + name="sw", + task_key="sw", + on_expression="@e", + cases=[ + SwitchCase( + value="one", + activities=[_notebook("writer", writes=[DataAsset(signature="ds", identity="x.y")])], + ) + ], + default_activities=[], + ), + _notebook("reader", reads=[DataAsset(signature="ds", identity="x.y")]), + ], + ) + + edges = build_data_edges(pipeline) + + assert len(edges) == 1 + assert edges[0].source_task_key == "writer" + assert edges[0].target_task_key == "reader" + + +# --------------------------------------------------------------------------- # +# Motif annotations + composition + purity +# --------------------------------------------------------------------------- # + + +def test_build_motif_annotations_groups_members_by_tag(): + """A MotifActivity plus tagged members become one annotation over their task keys.""" + pipeline = Pipeline( + name="p", + tasks=[ + MotifActivity( + name="motif", + task_key="motif_auto_loader", + motif_id="auto_loader", + display_name="Auto Loader", + databricks_replacement="auto_loader", + matched_activity_names=["Copy A", "Copy B"], + confidence_notes=["matched on file source"], + ), + _notebook("member", motif_id="auto_loader"), + _notebook("unrelated"), + ], + ) + + annotations = build_motif_annotations(pipeline) + + assert len(annotations) == 1 + assert annotations[0].motif_id == "auto_loader" + assert annotations[0].member_task_keys == ["motif_auto_loader", "member"] + assert annotations[0].display_name == "Auto Loader" + assert annotations[0].databricks_replacement == "auto_loader" + + +def test_with_lineage_is_pure(): + """with_lineage returns a new pipeline and never mutates the input.""" + pipeline = Pipeline(name="p", tasks=[_notebook("n")]) + lineage = build_lineage(pipeline) + + updated = with_lineage(pipeline, lineage) + + assert pipeline.lineage is None + assert updated is not pipeline + assert updated.lineage is lineage + + +# --------------------------------------------------------------------------- # +# ir_serde round-trips +# --------------------------------------------------------------------------- # + + +def test_serde_round_trip_new_activity_fields_and_lineage_block(): + """data_reads/data_writes/motif_id and the lineage block survive JSON round-trip.""" + pipeline = Pipeline( + name="p", + tasks=[ + _notebook( + "writer", + writes=[DataAsset(signature="ds_out", identity="curated.orders", asset_type="table")], + motif_id="auto_loader", + ), + _notebook( + "reader", + reads=[ + DataAsset( + signature="ds_in", + identity="curated.orders", + asset_type="table", + properties={"format": "delta"}, + ) + ], + ), + WaitActivity(name="pause", task_key="pause", wait_time_seconds=5), + ], + lineage=Lineage( + control_edges=[ + ControlEdge( + source_workflow="p", + target_workflow="child", + via_task_key="writer", + wait_for_completion=True, + resolved=True, + ) + ], + data_edges=[ + DataEdge( + source_task_key="writer", + target_task_key="reader", + match_kind="identity", + match_key="curated.orders", + identity="curated.orders", + asset_type="table", + ) + ], + motifs=[ + MotifAnnotation( + motif_id="auto_loader", + member_task_keys=["writer"], + display_name="Auto Loader", + databricks_replacement="auto_loader", + notes=["note"], + ) + ], + ), + ) + + reloaded, _ = pipeline_dict_to_ir(json.loads(json.dumps(pipeline_to_dict(pipeline)))) + + writer = reloaded.tasks[0] + reader = reloaded.tasks[1] + assert writer.motif_id == "auto_loader" + assert writer.data_writes == [DataAsset(signature="ds_out", identity="curated.orders", asset_type="table")] + assert reader.data_reads == [ + DataAsset(signature="ds_in", identity="curated.orders", asset_type="table", properties={"format": "delta"}) + ] + # A task without lineage fields rehydrates to empty lists / None, never missing. + assert reloaded.tasks[2].data_reads == [] + assert reloaded.tasks[2].data_writes == [] + assert reloaded.tasks[2].motif_id is None + + assert reloaded.lineage == pipeline.lineage + + +def test_serde_round_trip_motif_activity_no_kwarg_collision_r1(): + """A MotifActivity round-trips without the R1 'multiple values for motif_id' TypeError.""" + pipeline = Pipeline( + name="p", + tasks=[ + MotifActivity( + name="motif", + task_key="motif_auto_loader", + motif_id="auto_loader", + display_name="Auto Loader", + databricks_replacement="auto_loader", + matched_activity_names=["Copy A", "Copy B"], + data_reads=[DataAsset(signature="src", identity="raw.src")], + ) + ], + ) + + reloaded, _ = pipeline_dict_to_ir(json.loads(json.dumps(pipeline_to_dict(pipeline)))) + + task = reloaded.tasks[0] + assert isinstance(task, MotifActivity) + assert task.motif_id == "auto_loader" + assert task.display_name == "Auto Loader" + assert task.matched_activity_names == ["Copy A", "Copy B"] + assert task.data_reads == [DataAsset(signature="src", identity="raw.src")] + + +def test_lineage_block_always_emits_lists_never_null(): + """An attached empty lineage serialises its edge collections as lists, not null.""" + pipeline = with_lineage(Pipeline(name="p", tasks=[_notebook("n")]), Lineage()) + + serialised = pipeline_to_dict(pipeline) + + assert serialised["lineage"] == {"control_edges": [], "data_edges": [], "motifs": []} From 781a07e5d8f04f666235ca2f84b903924c2b90fd Mon Sep 17 00:00:00 2001 From: Matthew Moorcroft Date: Mon, 14 Sep 2026 17:07:01 +0100 Subject: [PATCH 2/7] Define the shared source-neutral discovery AST (#61) Adds the source-faithful, standardized model that both ADF and Airflow map onto for discovery -- distinct from the lossy Databricks IR target model. This is the contract the per-source mappers (#62 ADF, #63 Airflow) align to. No source is wired, convert is untouched, ir_serde's convert->package contract is unchanged, and inventory.json emission is untouched. - models/discovery.py: SourceGraph (workflow container), SourceNode (base task) + ContainerNode (labelled control-flow branches) + GapNode (couldn't classify, carries raw + reason), SourceDependency (conditions kept as a LIST), ScheduleSpec (promotes today's Pipeline.schedule dict), ParameterSpec, PolicySpec. Shape = shared CORE + explicit EXTENSION SEAM: every graph/node carries verbatim `raw` and a free-form `properties`/`extensions` bag, so nothing platform-specific is flattened away. Typed field set kept minimal; no typed Connection/linked-service field yet (rides in properties -- a marked expansion point). - Reuses the #61 lineage primitive DataAsset for data_reads/data_writes rather than duplicating it, so the discovery AST and lineage substrate share one physical-asset vocabulary. - discovery_serde.py: separate serialize<->deserialize round-trip for the model (NOT ir_serde -- kept apart so the two contracts evolve independently). Reuses data_asset_to_dict / new data_asset_from_dict from ir_serde for the DataAsset shape. Node dicts carry a node_type discriminator so Container/Gap subclasses rehydrate to the right class; lists always emit as lists. - ir_serde.py: add data_asset_from_dict (canonical inverse of the existing data_asset_to_dict); additive, does not touch the convert->package contract. - Grounded read-only against both sources (models/adf_ast.py and Airflow's loader captures/visitor on origin/main) to confirm every core field is populatable from both; platform-specific bits have a home in raw/extensions. - tests: construction + exact round-trip incl. ContainerNode branches, a node with data_reads/data_writes + extensions, GapNode, schedule, and parameters. Co-authored-by: Isaac --- src/flowx/discovery_serde.py | 212 +++++++++++++++++++++++ src/flowx/ir_serde.py | 15 ++ src/flowx/models/discovery.py | 261 +++++++++++++++++++++++++++++ tests/unit/test_discovery_model.py | 179 ++++++++++++++++++++ 4 files changed, 667 insertions(+) create mode 100644 src/flowx/discovery_serde.py create mode 100644 src/flowx/models/discovery.py create mode 100644 tests/unit/test_discovery_model.py diff --git a/src/flowx/discovery_serde.py b/src/flowx/discovery_serde.py new file mode 100644 index 0000000..7e50fb6 --- /dev/null +++ b/src/flowx/discovery_serde.py @@ -0,0 +1,212 @@ +"""JSON serialisation for the shared discovery AST (:mod:`flowx.models.discovery`). + +Kept separate from :mod:`flowx.ir_serde` on purpose: ``ir_serde`` owns the +``translation_report.json`` convert->package contract for the Databricks IR, and +the discovery AST is a different model with a different lifecycle. This module is +its own round-trip pair so evolving one shape never disturbs the other. + +The DataAsset (de)serialisers are reused from ``ir_serde`` (``data_asset_to_dict`` +/ ``data_asset_from_dict``) so the physical-asset shape has a single definition +shared by the lineage substrate and the discovery AST. + +Every node dict carries a ``node_type`` discriminator (the dataclass name) so a +:class:`~flowx.models.discovery.ContainerNode` or +:class:`~flowx.models.discovery.GapNode` rehydrates to the right class. Lists are +always emitted as lists (never ``None``) for stable golden diffs. +""" + +from __future__ import annotations + +from typing import Any + +from flowx.ir_serde import data_asset_from_dict, data_asset_to_dict +from flowx.models.discovery import ( + ContainerNode, + GapNode, + ParameterSpec, + PolicySpec, + ScheduleSpec, + SourceDependency, + SourceGraph, + SourceNode, +) + + +def source_graph_to_dict(graph: SourceGraph) -> dict[str, Any]: + """Serialise a :class:`SourceGraph` to a JSON-friendly dictionary.""" + result: dict[str, Any] = { + "name": graph.name, + "source": graph.source, + "parameters": {name: _parameter_to_dict(spec) for name, spec in graph.parameters.items()}, + "variables": {name: _parameter_to_dict(spec) for name, spec in graph.variables.items()}, + "tags": list(graph.tags), + "tasks": [_node_to_dict(node) for node in graph.tasks], + } + if graph.description is not None: + result["description"] = graph.description + if graph.schedule is not None: + result["schedule"] = _schedule_to_dict(graph.schedule) + if graph.properties: + result["properties"] = dict(graph.properties) + if graph.extensions: + result["extensions"] = dict(graph.extensions) + if graph.raw is not None: + result["raw"] = graph.raw + return result + + +def source_graph_from_dict(raw: dict[str, Any]) -> SourceGraph: + """Rehydrate a :class:`SourceGraph` from the dict :func:`source_graph_to_dict` emits.""" + schedule = raw.get("schedule") + return SourceGraph( + name=raw.get("name", ""), + source=raw.get("source", ""), + description=raw.get("description"), + parameters={name: _parameter_from_dict(spec) for name, spec in (raw.get("parameters") or {}).items()}, + variables={name: _parameter_from_dict(spec) for name, spec in (raw.get("variables") or {}).items()}, + schedule=_schedule_from_dict(schedule) if schedule else None, + tags=list(raw.get("tags") or []), + tasks=[_node_from_dict(node) for node in raw.get("tasks") or []], + properties=dict(raw.get("properties") or {}), + extensions=dict(raw.get("extensions") or {}), + raw=raw.get("raw"), + ) + + +def _parameter_to_dict(spec: ParameterSpec) -> dict[str, Any]: + result: dict[str, Any] = {} + if spec.type is not None: + result["type"] = spec.type + if spec.default is not None: + result["default"] = spec.default + return result + + +def _parameter_from_dict(raw: dict[str, Any]) -> ParameterSpec: + return ParameterSpec(type=raw.get("type"), default=raw.get("default")) + + +def _schedule_to_dict(schedule: ScheduleSpec) -> dict[str, Any]: + result: dict[str, Any] = {"kind": schedule.kind} + if schedule.quartz_cron_expression is not None: + result["quartz_cron_expression"] = schedule.quartz_cron_expression + if schedule.timezone_id is not None: + result["timezone_id"] = schedule.timezone_id + if schedule.pause_status is not None: + result["pause_status"] = schedule.pause_status + if schedule.extensions: + result["extensions"] = dict(schedule.extensions) + return result + + +def _schedule_from_dict(raw: dict[str, Any]) -> ScheduleSpec: + return ScheduleSpec( + kind=raw.get("kind", ""), + quartz_cron_expression=raw.get("quartz_cron_expression"), + timezone_id=raw.get("timezone_id"), + pause_status=raw.get("pause_status"), + extensions=dict(raw.get("extensions") or {}), + ) + + +def _policy_to_dict(policy: PolicySpec) -> dict[str, Any]: + result: dict[str, Any] = {} + if policy.timeout_seconds is not None: + result["timeout_seconds"] = policy.timeout_seconds + if policy.max_retries is not None: + result["max_retries"] = policy.max_retries + if policy.retry_interval_seconds is not None: + result["retry_interval_seconds"] = policy.retry_interval_seconds + if policy.extensions: + result["extensions"] = dict(policy.extensions) + return result + + +def _policy_from_dict(raw: dict[str, Any]) -> PolicySpec: + return PolicySpec( + timeout_seconds=raw.get("timeout_seconds"), + max_retries=raw.get("max_retries"), + retry_interval_seconds=raw.get("retry_interval_seconds"), + extensions=dict(raw.get("extensions") or {}), + ) + + +def _dependency_to_dict(dependency: SourceDependency) -> dict[str, Any]: + return { + "upstream": dependency.upstream, + "conditions": list(dependency.conditions), + "resolved": dependency.resolved, + } + + +def _dependency_from_dict(raw: dict[str, Any]) -> SourceDependency: + return SourceDependency( + upstream=raw.get("upstream", ""), + conditions=list(raw.get("conditions") or []), + resolved=bool(raw.get("resolved", True)), + ) + + +def _node_to_dict(node: SourceNode) -> dict[str, Any]: + """Serialise any SourceNode (including Container/Gap subclasses) to a dict. + + The ``node_type`` discriminator is the dataclass name so the matching + subclass is rebuilt on the way back. + """ + result: dict[str, Any] = { + "node_type": type(node).__name__, + "source_id": node.source_id, + "task_key": node.task_key, + "concept": node.concept, + "source": node.source, + "dependencies": [_dependency_to_dict(dependency) for dependency in node.dependencies], + "data_reads": [data_asset_to_dict(asset) for asset in node.data_reads], + "data_writes": [data_asset_to_dict(asset) for asset in node.data_writes], + } + if node.name is not None: + result["name"] = node.name + if node.native_type is not None: + result["native_type"] = node.native_type + if node.policy is not None: + result["policy"] = _policy_to_dict(node.policy) + if node.properties: + result["properties"] = dict(node.properties) + if node.raw is not None: + result["raw"] = node.raw + if isinstance(node, ContainerNode): + result["branches"] = { + label: [_node_to_dict(child) for child in children] for label, children in node.branches.items() + } + if isinstance(node, GapNode) and node.reason is not None: + result["reason"] = node.reason + return result + + +def _node_from_dict(raw: dict[str, Any]) -> SourceNode: + """Rehydrate any SourceNode from its dict, dispatching on ``node_type``.""" + common: dict[str, Any] = { + "source_id": raw.get("source_id", ""), + "task_key": raw.get("task_key", ""), + "concept": raw.get("concept", ""), + "source": raw.get("source", ""), + "name": raw.get("name"), + "native_type": raw.get("native_type"), + "dependencies": [_dependency_from_dict(dependency) for dependency in raw.get("dependencies") or []], + "policy": _policy_from_dict(raw["policy"]) if raw.get("policy") else None, + "data_reads": [data_asset_from_dict(asset) for asset in raw.get("data_reads") or []], + "data_writes": [data_asset_from_dict(asset) for asset in raw.get("data_writes") or []], + "properties": dict(raw.get("properties") or {}), + "raw": raw.get("raw"), + } + node_type = raw.get("node_type", "SourceNode") + if node_type == "ContainerNode": + return ContainerNode( + **common, + branches={ + label: [_node_from_dict(child) for child in children] + for label, children in (raw.get("branches") or {}).items() + }, + ) + if node_type == "GapNode": + return GapNode(**common, reason=raw.get("reason")) + return SourceNode(**common) diff --git a/src/flowx/ir_serde.py b/src/flowx/ir_serde.py index c49ba6a..72b00a6 100644 --- a/src/flowx/ir_serde.py +++ b/src/flowx/ir_serde.py @@ -103,6 +103,21 @@ def data_asset_to_dict(asset: DataAsset) -> dict[str, Any]: return result +def data_asset_from_dict(raw: dict[str, Any]) -> DataAsset: + """Rehydrate a :class:`DataAsset` from the dict :func:`data_asset_to_dict` emits. + + The canonical inverse of :func:`data_asset_to_dict`, so every consumer of the + serialised DataAsset shape (the lineage substrate and the discovery AST) + reads it the same way. + """ + return DataAsset( + signature=raw.get("signature", ""), + identity=raw.get("identity"), + asset_type=raw.get("asset_type"), + properties=dict(raw.get("properties") or {}), + ) + + def lineage_to_dict(lineage: Lineage) -> dict[str, Any]: """Serialise a :class:`Lineage` block to a JSON-friendly dictionary. diff --git a/src/flowx/models/discovery.py b/src/flowx/models/discovery.py new file mode 100644 index 0000000..0412955 --- /dev/null +++ b/src/flowx/models/discovery.py @@ -0,0 +1,261 @@ +"""Source-faithful shared discovery AST. + +The Databricks IR in :mod:`flowx.models.ir` is a *target* model: it keeps only +what Databricks needs to run a workflow, so it is deliberately lossy about the +source. The discovery layer needs the opposite -- a source-*faithful*, +standardised model that both Azure Data Factory and Apache Airflow map onto +without either being flattened to a lowest common denominator. This module +defines that model. It is the contract the per-source mappers (#62 for ADF, #63 +for Airflow) align to. + +Design shape: a small shared **core** of genuinely common concepts as typed +fields, plus an explicit **extension seam** so nothing platform-specific is +lost. Every graph and node carries: + +* ``raw`` -- the verbatim source dict (the ADF activity/pipeline JSON, the + Airflow capture), so a mapper can always fall back to the original; and +* ``properties`` / ``extensions`` -- a free-form bag for platform-specific + attributes that have no shared typed field yet. + +The typed field set is intentionally minimal. Concepts get promoted to typed +fields only once they are genuinely shared; everything else rides in +``raw`` / ``extensions`` until a later PR promotes it. In particular there is +**no typed Connection / linked-service field yet** -- connection and +linked-service details ride in ``properties`` / ``extensions`` for now. That is +a deliberate expansion point, not an oversight. + +The lineage primitives from #61 (:class:`~flowx.models.ir.DataAsset`) are +reused here rather than duplicated: a node's reads and writes are lists of +``DataAsset``, so the discovery AST and the lineage substrate share one +vocabulary for physical data references. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from flowx.models.ir import DataAsset + +# --------------------------------------------------------------------------- # +# Well-known discriminators and concepts (open vocabularies). +# +# ``source`` and ``concept`` are plain strings, not enums, so a source can +# introduce a new value without churning this module -- the extension seam +# applies to the vocabulary too. The constants below name the values the core +# already understands; anything a mapper cannot classify becomes ``CONCEPT_GAP``. +# --------------------------------------------------------------------------- # + +SOURCE_ADF = "adf" +SOURCE_AIRFLOW = "airflow" + +# Neutral node concepts shared across sources. Not exhaustive by design. +CONCEPT_NOTEBOOK = "notebook" +CONCEPT_SCRIPT = "script" +CONCEPT_COPY_DATA = "copy_data" +CONCEPT_QUERY = "query" +CONCEPT_SET_VARIABLE = "set_variable" +CONCEPT_WAIT = "wait" +CONCEPT_RUN_WORKFLOW = "run_workflow" +CONCEPT_BRANCH = "branch" +CONCEPT_LOOP = "loop" +CONCEPT_SWITCH = "switch" +CONCEPT_GROUP = "group" +CONCEPT_GAP = "gap" + + +@dataclass(slots=True, kw_only=True) +class ScheduleSpec: + """A workflow schedule, promoted from the untyped ``Pipeline.schedule`` dict. + + The shared core is the trigger ``kind`` plus the fields a cron-style + schedule needs; everything source-specific (tumbling-window frequency, + Airflow interval/unit, file-arrival url/events, approximation notes) rides + in :attr:`extensions`. + + Attributes: + kind: Trigger kind -- ``"schedule"`` / ``"periodic"`` / ``"continuous"`` + / ``"file_arrival"`` / ``"manual_setup"`` / ... + quartz_cron_expression: Quartz cron string for cron-style schedules. + timezone_id: IANA timezone id, when the source resolves one. + pause_status: ``"PAUSED"`` / ``"UNPAUSED"`` when the source expresses it. + extensions: Source-specific schedule fields with no shared typed home. + """ + + kind: str + quartz_cron_expression: str | None = None + timezone_id: str | None = None + pause_status: str | None = None + extensions: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True, kw_only=True) +class ParameterSpec: + """A parameter or variable declaration: an optional type and a default. + + Reused for both graph parameters and graph variables since both sources + describe them the same way (a declared type plus an optional default). + + Attributes: + type: Declared type string when the source has one (ADF ``String`` / + ``Int`` / ...), else ``None`` (Airflow params carry only a default). + default: Default / initial value, or ``None``. + """ + + type: str | None = None + default: Any = None + + +@dataclass(slots=True, kw_only=True) +class PolicySpec: + """Retry / timeout policy on a node, normalised to seconds. + + Only the genuinely shared retry/timeout knobs are typed; source-specific + policy (ADF ``secure_input`` / ``secure_output``, Airflow retry-delay + shapes, email-on-failure) rides in :attr:`extensions`. + + Attributes: + timeout_seconds: Execution timeout in seconds, when resolvable. + max_retries: Maximum retry count. + retry_interval_seconds: Delay between retries in seconds. + extensions: Source-specific policy fields with no shared typed home. + """ + + timeout_seconds: int | None = None + max_retries: int | None = None + retry_interval_seconds: int | None = None + extensions: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True, kw_only=True) +class SourceDependency: + """A dependency edge from a node to one upstream node. + + Conditions stay a **list** -- ADF dependency edges carry one or more + outcome conditions (``["Succeeded", "Skipped"]``), so collapsing them to a + single outcome would lose information. Airflow's unconditional edges use an + empty list (or a single normalised condition). + + Attributes: + upstream: Task key of the upstream node this edge depends on. + conditions: Required upstream outcome(s); empty when unconditional. + resolved: ``False`` when the upstream could not be resolved (e.g. a + partial export); recorded, not dropped. + """ + + upstream: str + conditions: list[str] = field(default_factory=list) + resolved: bool = True + + +@dataclass(slots=True, kw_only=True) +class SourceNode: + """A single task in a source workflow, standardised but source-faithful. + + Attributes: + source_id: Stable identifier from the source (ADF activity name, an + Airflow capture id) used to resolve dependency edges before task + keys are allocated. + task_key: Normalised task key unique within the graph. + concept: Neutral, classified kind (see the ``CONCEPT_*`` constants); + ``CONCEPT_GAP`` when the source could not classify it. + source: Source discriminator (``SOURCE_ADF`` / ``SOURCE_AIRFLOW``). + name: Human-readable display name, when the source has one distinct + from ``task_key``. + native_type: The source's own type string, preserved verbatim (ADF + ``"Copy"`` / ``"DatabricksNotebook"``, Airflow ``"BashOperator"`` / + a TaskFlow decorator). + dependencies: Upstream dependency edges. + policy: Retry / timeout policy, when the source declares one. + data_reads: Physical data assets this node reads (reuses #61 + :class:`~flowx.models.ir.DataAsset`). + data_writes: Physical data assets this node writes. + properties: Free-form bag for platform-specific attributes with no + shared typed field yet. Connection / linked-service details live + here for now -- a typed connection field is a deliberate future + expansion point. + raw: Verbatim source dict for this node, for lossless fallback. + """ + + source_id: str + task_key: str + concept: str + source: str + name: str | None = None + native_type: str | None = None + dependencies: list[SourceDependency] = field(default_factory=list) + policy: PolicySpec | None = None + data_reads: list[DataAsset] = field(default_factory=list) + data_writes: list[DataAsset] = field(default_factory=list) + properties: dict[str, Any] = field(default_factory=dict) + raw: dict[str, Any] | None = None + + +@dataclass(slots=True, kw_only=True) +class ContainerNode(SourceNode): + """A control-flow container that nests child nodes under labelled branches. + + One shape covers every source's control flow: an ADF ``ForEach`` / + ``Until`` becomes ``{"body": [...]}``, an ``IfCondition`` becomes + ``{"true": [...], "false": [...]}``, a ``Switch`` becomes + ``{"": [...], "default": [...]}``, and an Airflow ``TaskGroup`` + becomes ``{"group": [...]}``. The branch label is the source's own, so no + control-flow structure is flattened away. + + Attributes: + branches: Branch label -> ordered child nodes. + """ + + branches: dict[str, list[SourceNode]] = field(default_factory=dict) + + +@dataclass(slots=True, kw_only=True) +class GapNode(SourceNode): + """A node the source could not classify, kept with its raw payload + reason. + + Its :attr:`concept` defaults to ``CONCEPT_GAP`` so gaps are uniform across + sources. The unclassifiable payload is preserved in ``raw`` (inherited) and + the human-readable cause in :attr:`reason`, so discovery can report and + later resolve it rather than dropping it. + + Attributes: + reason: Why the node could not be classified. + """ + + concept: str = CONCEPT_GAP + reason: str | None = None + + +@dataclass(slots=True, kw_only=True) +class SourceGraph: + """A source workflow (an ADF pipeline, an Airflow DAG), standardised. + + Attributes: + name: Workflow name (ADF pipeline name, Airflow ``dag_id``). + source: Source discriminator (``SOURCE_ADF`` / ``SOURCE_AIRFLOW``). + description: Human-readable description, when the source has one. + parameters: Parameter declarations keyed by name. + variables: Variable declarations keyed by name. Empty for sources with + no graph-scoped variable concept (Airflow); such sources' global + variables ride in :attr:`extensions`. + schedule: Workflow schedule, when one is declared. + tags: Free-form label list (ADF annotations, Airflow user tags). + tasks: Top-level nodes; control flow nests further nodes via + :class:`ContainerNode`. + properties: Free-form bag for graph-level platform-specific attributes. + extensions: Alias-free overflow for anything else source-specific + (e.g. ADF ``folder``, Airflow global Variables). + raw: Verbatim source dict for the whole workflow. + """ + + name: str + source: str + description: str | None = None + parameters: dict[str, ParameterSpec] = field(default_factory=dict) + variables: dict[str, ParameterSpec] = field(default_factory=dict) + schedule: ScheduleSpec | None = None + tags: list[str] = field(default_factory=list) + tasks: list[SourceNode] = field(default_factory=list) + properties: dict[str, Any] = field(default_factory=dict) + extensions: dict[str, Any] = field(default_factory=dict) + raw: dict[str, Any] | None = None diff --git a/tests/unit/test_discovery_model.py b/tests/unit/test_discovery_model.py new file mode 100644 index 0000000..ca08ad6 --- /dev/null +++ b/tests/unit/test_discovery_model.py @@ -0,0 +1,179 @@ +"""Unit tests for the shared discovery AST (models/discovery.py + discovery_serde.py). + +Covers construction of the node set and an exact serialize<->deserialize round +trip, including a ContainerNode with labelled branches, a node carrying +data_reads/data_writes and properties, a GapNode, a schedule, and parameters. +""" + +from __future__ import annotations + +import json + +from flowx.discovery_serde import source_graph_from_dict, source_graph_to_dict +from flowx.models.discovery import ( + CONCEPT_COPY_DATA, + CONCEPT_GAP, + CONCEPT_LOOP, + CONCEPT_NOTEBOOK, + SOURCE_ADF, + SOURCE_AIRFLOW, + ContainerNode, + GapNode, + ParameterSpec, + PolicySpec, + ScheduleSpec, + SourceDependency, + SourceGraph, + SourceNode, +) +from flowx.models.ir import DataAsset + + +def test_gap_node_defaults_to_gap_concept(): + """A GapNode is a gap without the caller having to restate the concept.""" + gap = GapNode(source_id="a1", task_key="a1", source=SOURCE_ADF, reason="unmapped ExecuteDataFlow") + + assert gap.concept == CONCEPT_GAP + assert gap.reason == "unmapped ExecuteDataFlow" + + +def test_container_node_holds_labelled_branches(): + """A ContainerNode nests child nodes under source-named branch labels.""" + container = ContainerNode( + source_id="fe", + task_key="for_each", + concept=CONCEPT_LOOP, + source=SOURCE_ADF, + native_type="ForEach", + branches={"body": [SourceNode(source_id="c", task_key="copy", concept=CONCEPT_COPY_DATA, source=SOURCE_ADF)]}, + ) + + assert list(container.branches) == ["body"] + assert container.branches["body"][0].task_key == "copy" + + +def _sample_graph() -> SourceGraph: + return SourceGraph( + name="ingest_orders", + source=SOURCE_ADF, + description="Loads orders and fans out per region", + parameters={"region": ParameterSpec(type="String", default="us")}, + variables={"batch": ParameterSpec(type="String")}, + schedule=ScheduleSpec( + kind="schedule", + quartz_cron_expression="0 0 * * * ?", + timezone_id="UTC", + pause_status="UNPAUSED", + extensions={"note": "approximated"}, + ), + tags=["prod", "orders"], + tasks=[ + SourceNode( + source_id="copy_orders", + task_key="copy_orders", + concept=CONCEPT_COPY_DATA, + source=SOURCE_ADF, + name="Copy Orders", + native_type="Copy", + policy=PolicySpec(timeout_seconds=3600, max_retries=2, extensions={"secure_output": True}), + data_reads=[DataAsset(signature="ds_raw_orders", asset_type="file")], + data_writes=[ + DataAsset( + signature="ds_curated_orders", + identity="curated.orders", + asset_type="table", + properties={"format": "delta"}, + ) + ], + properties={"linked_service": "AzureSqlDatabase1"}, + raw={"type": "Copy", "name": "Copy Orders"}, + ), + ContainerNode( + source_id="per_region", + task_key="per_region", + concept=CONCEPT_LOOP, + source=SOURCE_ADF, + native_type="ForEach", + dependencies=[SourceDependency(upstream="copy_orders", conditions=["Succeeded", "Skipped"])], + branches={ + "body": [ + SourceNode( + source_id="run_region", + task_key="run_region", + concept=CONCEPT_NOTEBOOK, + source=SOURCE_ADF, + native_type="DatabricksNotebook", + ) + ] + }, + ), + GapNode( + source_id="dataflow", + task_key="mapping_dataflow", + source=SOURCE_ADF, + native_type="ExecuteDataFlow", + reason="ExecuteDataFlow has no deterministic mapping", + raw={"type": "ExecuteDataFlow"}, + ), + ], + properties={"folder": "ingest"}, + extensions={"annotations": ["team:data"]}, + raw={"name": "ingest_orders"}, + ) + + +def test_source_graph_round_trip_is_exact(): + """A fully-populated graph survives to_dict -> JSON -> from_dict unchanged.""" + graph = _sample_graph() + + reloaded = source_graph_from_dict(json.loads(json.dumps(source_graph_to_dict(graph)))) + + assert reloaded == graph + + +def test_round_trip_preserves_container_branches_and_gap(): + """Subclass identity (Container/Gap) and their extra fields survive the round trip.""" + graph = _sample_graph() + + reloaded = source_graph_from_dict(json.loads(json.dumps(source_graph_to_dict(graph)))) + + container = reloaded.tasks[1] + assert isinstance(container, ContainerNode) + assert container.branches["body"][0].task_key == "run_region" + assert container.dependencies[0].conditions == ["Succeeded", "Skipped"] + + gap = reloaded.tasks[2] + assert isinstance(gap, GapNode) + assert gap.concept == CONCEPT_GAP + assert gap.reason == "ExecuteDataFlow has no deterministic mapping" + + +def test_round_trip_preserves_data_assets_and_extension_bags(): + """data_reads/data_writes (reused DataAsset) and the extension bags round-trip.""" + graph = _sample_graph() + + reloaded = source_graph_from_dict(json.loads(json.dumps(source_graph_to_dict(graph)))) + + copy_node = reloaded.tasks[0] + assert copy_node.data_reads == [DataAsset(signature="ds_raw_orders", asset_type="file")] + assert copy_node.data_writes == [ + DataAsset( + signature="ds_curated_orders", identity="curated.orders", asset_type="table", properties={"format": "delta"} + ) + ] + assert copy_node.properties == {"linked_service": "AzureSqlDatabase1"} + assert reloaded.extensions == {"annotations": ["team:data"]} + assert reloaded.schedule is not None + assert reloaded.schedule.extensions == {"note": "approximated"} + + +def test_empty_graph_emits_lists_and_dicts_never_null(): + """A minimal graph serialises its collections as empty containers, not null.""" + serialised = source_graph_to_dict(SourceGraph(name="empty", source=SOURCE_AIRFLOW)) + + assert serialised["tasks"] == [] + assert serialised["tags"] == [] + assert serialised["parameters"] == {} + assert serialised["variables"] == {} + # Airflow has no graph-scoped variables; the field stays empty rather than absent. + assert source_graph_from_dict(serialised).variables == {} From 872ef58c8f7378f95f5812c49e3aeeab5620c578 Mon Sep 17 00:00:00 2001 From: Matthew Moorcroft Date: Mon, 14 Sep 2026 17:23:30 +0100 Subject: [PATCH 3/7] Address review: GapNode default + source-faithful ScheduleSpec (#61) FIX 1 (bug): discovery_serde._node_from_dict forced concept="" on every node, so a partial GapNode dict (no `concept` key) rehydrated with "" instead of the model default CONCEPT_GAP. Now `concept` is passed through only when the dict carries a non-empty value, letting each class's model default apply; other omitted optional fields already fall back to their model defaults. Adds tests for a partial GapNode dict (-> CONCEPT_GAP) and a partial SourceNode dict (optional fields -> model defaults, not empty strings). FIX 2 (design): ScheduleSpec no longer bakes in Databricks target-shaped fields (quartz_cron_expression, timezone_id, pause_status). Reworked to a minimal source-faithful form -- neutral `kind` + `expression` holding the source schedule as-given (Airflow schedule_interval/cron string, ADF recurrence payload). Target normalisation is a convert concern; if needed for now it rides in `extensions`. discovery_serde + tests updated to match. Co-authored-by: Isaac --- src/flowx/discovery_serde.py | 33 +++++++++++++--------- src/flowx/models/discovery.py | 33 ++++++++++++---------- tests/unit/test_discovery_model.py | 45 ++++++++++++++++++++++++++---- 3 files changed, 77 insertions(+), 34 deletions(-) diff --git a/src/flowx/discovery_serde.py b/src/flowx/discovery_serde.py index 7e50fb6..c2f5b9c 100644 --- a/src/flowx/discovery_serde.py +++ b/src/flowx/discovery_serde.py @@ -88,12 +88,8 @@ def _parameter_from_dict(raw: dict[str, Any]) -> ParameterSpec: def _schedule_to_dict(schedule: ScheduleSpec) -> dict[str, Any]: result: dict[str, Any] = {"kind": schedule.kind} - if schedule.quartz_cron_expression is not None: - result["quartz_cron_expression"] = schedule.quartz_cron_expression - if schedule.timezone_id is not None: - result["timezone_id"] = schedule.timezone_id - if schedule.pause_status is not None: - result["pause_status"] = schedule.pause_status + if schedule.expression is not None: + result["expression"] = schedule.expression if schedule.extensions: result["extensions"] = dict(schedule.extensions) return result @@ -102,9 +98,7 @@ def _schedule_to_dict(schedule: ScheduleSpec) -> dict[str, Any]: def _schedule_from_dict(raw: dict[str, Any]) -> ScheduleSpec: return ScheduleSpec( kind=raw.get("kind", ""), - quartz_cron_expression=raw.get("quartz_cron_expression"), - timezone_id=raw.get("timezone_id"), - pause_status=raw.get("pause_status"), + expression=raw.get("expression"), extensions=dict(raw.get("extensions") or {}), ) @@ -183,11 +177,17 @@ def _node_to_dict(node: SourceNode) -> dict[str, Any]: def _node_from_dict(raw: dict[str, Any]) -> SourceNode: - """Rehydrate any SourceNode from its dict, dispatching on ``node_type``.""" + """Rehydrate any SourceNode from its dict, dispatching on ``node_type``. + + Optional fields absent from the dict fall back to each class's model + default rather than being forced to ``""`` / empty. In particular + ``concept`` is passed through only when the dict carries a non-empty value, + so a partial :class:`GapNode` dict (no ``concept`` key) rehydrates with the + model default ``CONCEPT_GAP`` instead of being overridden with ``""``. + """ common: dict[str, Any] = { "source_id": raw.get("source_id", ""), "task_key": raw.get("task_key", ""), - "concept": raw.get("concept", ""), "source": raw.get("source", ""), "name": raw.get("name"), "native_type": raw.get("native_type"), @@ -198,7 +198,16 @@ def _node_from_dict(raw: dict[str, Any]) -> SourceNode: "properties": dict(raw.get("properties") or {}), "raw": raw.get("raw"), } + concept = raw.get("concept") + if concept: + common["concept"] = concept + node_type = raw.get("node_type", "SourceNode") + if node_type == "GapNode": + # Leave `concept` unset when absent so GapNode's CONCEPT_GAP default wins. + return GapNode(**common, reason=raw.get("reason")) + # SourceNode / ContainerNode require `concept`; only a malformed dict omits it. + common.setdefault("concept", "") if node_type == "ContainerNode": return ContainerNode( **common, @@ -207,6 +216,4 @@ def _node_from_dict(raw: dict[str, Any]) -> SourceNode: for label, children in (raw.get("branches") or {}).items() }, ) - if node_type == "GapNode": - return GapNode(**common, reason=raw.get("reason")) return SourceNode(**common) diff --git a/src/flowx/models/discovery.py b/src/flowx/models/discovery.py index 0412955..392e952 100644 --- a/src/flowx/models/discovery.py +++ b/src/flowx/models/discovery.py @@ -66,26 +66,29 @@ @dataclass(slots=True, kw_only=True) class ScheduleSpec: - """A workflow schedule, promoted from the untyped ``Pipeline.schedule`` dict. - - The shared core is the trigger ``kind`` plus the fields a cron-style - schedule needs; everything source-specific (tumbling-window frequency, - Airflow interval/unit, file-arrival url/events, approximation notes) rides - in :attr:`extensions`. + """A workflow schedule kept in the SOURCE's own shape. + + Deliberately source-faithful, not Databricks-normalised. ``kind`` is a + neutral trigger category and ``expression`` holds the schedule exactly as + the source gives it -- an Airflow ``schedule_interval`` cron string or + preset, an ADF trigger recurrence payload. Databricks-*target* + normalisation (a Quartz cron string, ``pause_status``, a resolved timezone + id) is a convert/target concern and is intentionally **not** typed here; a + mapper that needs to stash such derived values for now puts them in + :attr:`extensions`, never as first-class fields. Kept minimal on purpose -- + fields get promoted only once genuinely shared. Attributes: - kind: Trigger kind -- ``"schedule"`` / ``"periodic"`` / ``"continuous"`` - / ``"file_arrival"`` / ``"manual_setup"`` / ... - quartz_cron_expression: Quartz cron string for cron-style schedules. - timezone_id: IANA timezone id, when the source resolves one. - pause_status: ``"PAUSED"`` / ``"UNPAUSED"`` when the source expresses it. - extensions: Source-specific schedule fields with no shared typed home. + kind: Neutral trigger category (e.g. ``"schedule"`` / ``"interval"`` / + ``"file_arrival"`` / ``"continuous"`` / ``"manual"``), ``""`` when + unknown. + expression: The source schedule as-given -- a cron / interval / preset + string, or a structured recurrence payload. + extensions: Overflow for any other source-specific schedule detail. """ kind: str - quartz_cron_expression: str | None = None - timezone_id: str | None = None - pause_status: str | None = None + expression: Any = None extensions: dict[str, Any] = field(default_factory=dict) diff --git a/tests/unit/test_discovery_model.py b/tests/unit/test_discovery_model.py index ca08ad6..c4af00c 100644 --- a/tests/unit/test_discovery_model.py +++ b/tests/unit/test_discovery_model.py @@ -9,7 +9,7 @@ import json -from flowx.discovery_serde import source_graph_from_dict, source_graph_to_dict +from flowx.discovery_serde import _node_from_dict, source_graph_from_dict, source_graph_to_dict from flowx.models.discovery import ( CONCEPT_COPY_DATA, CONCEPT_GAP, @@ -37,6 +37,39 @@ def test_gap_node_defaults_to_gap_concept(): assert gap.reason == "unmapped ExecuteDataFlow" +def test_partial_gap_node_dict_rehydrates_with_gap_concept(): + """A GapNode dict with no `concept` key falls back to the CONCEPT_GAP default.""" + partial = {"node_type": "GapNode", "source_id": "d1", "task_key": "dataflow", "source": SOURCE_ADF} + + node = _node_from_dict(partial) + + assert isinstance(node, GapNode) + assert node.concept == CONCEPT_GAP + + +def test_partial_source_node_dict_falls_back_to_model_defaults(): + """Omitted optional fields on a plain SourceNode dict use model defaults, not empty strings.""" + partial = { + "node_type": "SourceNode", + "source_id": "n1", + "task_key": "run", + "concept": CONCEPT_NOTEBOOK, + "source": SOURCE_AIRFLOW, + } + + node = _node_from_dict(partial) + + assert type(node) is SourceNode + assert node.name is None + assert node.native_type is None + assert node.policy is None + assert node.dependencies == [] + assert node.data_reads == [] + assert node.data_writes == [] + assert node.properties == {} + assert node.raw is None + + def test_container_node_holds_labelled_branches(): """A ContainerNode nests child nodes under source-named branch labels.""" container = ContainerNode( @@ -61,10 +94,8 @@ def _sample_graph() -> SourceGraph: variables={"batch": ParameterSpec(type="String")}, schedule=ScheduleSpec( kind="schedule", - quartz_cron_expression="0 0 * * * ?", - timezone_id="UTC", - pause_status="UNPAUSED", - extensions={"note": "approximated"}, + expression={"frequency": "Day", "interval": 1}, + extensions={"timezone": "UTC"}, ), tags=["prod", "orders"], tasks=[ @@ -164,7 +195,9 @@ def test_round_trip_preserves_data_assets_and_extension_bags(): assert copy_node.properties == {"linked_service": "AzureSqlDatabase1"} assert reloaded.extensions == {"annotations": ["team:data"]} assert reloaded.schedule is not None - assert reloaded.schedule.extensions == {"note": "approximated"} + # Source-faithful schedule: the ADF recurrence rides verbatim in `expression`. + assert reloaded.schedule.expression == {"frequency": "Day", "interval": 1} + assert reloaded.schedule.extensions == {"timezone": "UTC"} def test_empty_graph_emits_lists_and_dicts_never_null(): From 2725a7161e17c6220305eab786d886691555847e Mon Sep 17 00:00:00 2001 From: Matthew Moorcroft Date: Mon, 14 Sep 2026 19:01:25 +0100 Subject: [PATCH 4/7] Make ScheduleSpec timezone a first-class source-faithful field (#61) Follow-up to the review fixes: the source-declared timezone (ADF recurrence `timeZone`, an Airflow DAG timezone) is faithful to the source, so promote it from the extensions bag to a typed ScheduleSpec.timezone field. Databricks- target normalisation (Quartz cron, pause_status) stays out of the typed fields. discovery_serde + tests updated. Co-authored-by: Isaac --- src/flowx/discovery_serde.py | 3 +++ src/flowx/models/discovery.py | 13 ++++++++----- tests/unit/test_discovery_model.py | 10 +++++++--- 3 files changed, 18 insertions(+), 8 deletions(-) diff --git a/src/flowx/discovery_serde.py b/src/flowx/discovery_serde.py index c2f5b9c..9dc77fb 100644 --- a/src/flowx/discovery_serde.py +++ b/src/flowx/discovery_serde.py @@ -90,6 +90,8 @@ def _schedule_to_dict(schedule: ScheduleSpec) -> dict[str, Any]: result: dict[str, Any] = {"kind": schedule.kind} if schedule.expression is not None: result["expression"] = schedule.expression + if schedule.timezone is not None: + result["timezone"] = schedule.timezone if schedule.extensions: result["extensions"] = dict(schedule.extensions) return result @@ -99,6 +101,7 @@ def _schedule_from_dict(raw: dict[str, Any]) -> ScheduleSpec: return ScheduleSpec( kind=raw.get("kind", ""), expression=raw.get("expression"), + timezone=raw.get("timezone"), extensions=dict(raw.get("extensions") or {}), ) diff --git a/src/flowx/models/discovery.py b/src/flowx/models/discovery.py index 392e952..1590c0b 100644 --- a/src/flowx/models/discovery.py +++ b/src/flowx/models/discovery.py @@ -69,11 +69,12 @@ class ScheduleSpec: """A workflow schedule kept in the SOURCE's own shape. Deliberately source-faithful, not Databricks-normalised. ``kind`` is a - neutral trigger category and ``expression`` holds the schedule exactly as - the source gives it -- an Airflow ``schedule_interval`` cron string or - preset, an ADF trigger recurrence payload. Databricks-*target* - normalisation (a Quartz cron string, ``pause_status``, a resolved timezone - id) is a convert/target concern and is intentionally **not** typed here; a + neutral trigger category, ``expression`` holds the schedule exactly as the + source gives it -- an Airflow ``schedule_interval`` cron string or preset, + an ADF trigger recurrence payload -- and ``timezone`` carries the timezone + the source itself declares (ADF ``timeZone``, an Airflow DAG timezone). + Databricks-*target* normalisation -- a Quartz cron string, ``pause_status`` + -- is a convert/target concern and is intentionally **not** typed here; a mapper that needs to stash such derived values for now puts them in :attr:`extensions`, never as first-class fields. Kept minimal on purpose -- fields get promoted only once genuinely shared. @@ -84,11 +85,13 @@ class ScheduleSpec: unknown. expression: The source schedule as-given -- a cron / interval / preset string, or a structured recurrence payload. + timezone: The timezone the source declares, verbatim, or ``None``. extensions: Overflow for any other source-specific schedule detail. """ kind: str expression: Any = None + timezone: str | None = None extensions: dict[str, Any] = field(default_factory=dict) diff --git a/tests/unit/test_discovery_model.py b/tests/unit/test_discovery_model.py index c4af00c..2e1191e 100644 --- a/tests/unit/test_discovery_model.py +++ b/tests/unit/test_discovery_model.py @@ -95,7 +95,8 @@ def _sample_graph() -> SourceGraph: schedule=ScheduleSpec( kind="schedule", expression={"frequency": "Day", "interval": 1}, - extensions={"timezone": "UTC"}, + timezone="UTC", + extensions={"runtime_state": "Started"}, ), tags=["prod", "orders"], tasks=[ @@ -195,9 +196,12 @@ def test_round_trip_preserves_data_assets_and_extension_bags(): assert copy_node.properties == {"linked_service": "AzureSqlDatabase1"} assert reloaded.extensions == {"annotations": ["team:data"]} assert reloaded.schedule is not None - # Source-faithful schedule: the ADF recurrence rides verbatim in `expression`. + # Source-faithful schedule: the ADF recurrence rides verbatim in `expression`, + # the source-declared timezone is a typed field, and no Databricks-target + # shape (Quartz cron / pause_status) is baked in. assert reloaded.schedule.expression == {"frequency": "Day", "interval": 1} - assert reloaded.schedule.extensions == {"timezone": "UTC"} + assert reloaded.schedule.timezone == "UTC" + assert reloaded.schedule.extensions == {"runtime_state": "Started"} def test_empty_graph_emits_lists_and_dicts_never_null(): From da7e81006db42dccb395e6e5084b89b4e67016c1 Mon Sep 17 00:00:00 2001 From: Matthew Moorcroft Date: Tue, 15 Sep 2026 11:18:59 +0100 Subject: [PATCH 5/7] Make lineage part of the shared discovery-AST standard (#61) Wire the source-neutral lineage layer up into the shared AST that #66 defines, so the standard is complete and self-contained (what #62/#63 align to). All additions are source-neutral: the neutral modules import only shared IR / discovery types, never sources/*. - models/discovery.py: add SourceGraph.lineage (Lineage | None), reusing the #61 Lineage type so the discovery AST and IR substrate share one vocabulary. - discovery_lineage.py (new): source-neutral graph lineage derivation over a SourceGraph -- Switch/ForEach/If-aware walk, control edges from a neutral invocation marker in node.properties, data edges from node data_reads/ data_writes via the shared two-tier match. Motifs stay empty (convert-time). - lineage.py: extract the primitive cores control_edges_from_calls / data_edges_from_endpoints (task_key strings + DataAsset only) so both the IR and discovery-AST derivations join through one implementation. IR-facing build_control_edges / build_data_edges signatures and behavior unchanged. - discovery_serde.py: serialize/rehydrate SourceGraph.lineage (forward via ir_serde.lineage_to_dict; inverse _lineage_from_dict here). - tests: source-neutral graph-lineage derivation (tiers, fan-out, no self/dup edges, Switch/ForEach recursion, control markers) + a lineage round-trip in the discovery-model serde tests. Per-source POPULATION of data_reads/data_writes (ADF datasets, Airflow SQL/operators) and the ADF invocation markers remain the per-source work in #62/#63; no sources/ code is wired here, convert and ir_serde's convert->package contract are untouched. Co-authored-by: Isaac --- src/flowx/discovery_lineage.py | 136 ++++++++++++++++++++ src/flowx/discovery_serde.py | 53 +++++++- src/flowx/lineage.py | 143 ++++++++++++++------- src/flowx/models/discovery.py | 8 +- tests/unit/test_discovery_lineage.py | 179 +++++++++++++++++++++++++++ tests/unit/test_discovery_model.py | 37 +++++- 6 files changed, 509 insertions(+), 47 deletions(-) create mode 100644 src/flowx/discovery_lineage.py create mode 100644 tests/unit/test_discovery_lineage.py diff --git a/src/flowx/discovery_lineage.py b/src/flowx/discovery_lineage.py new file mode 100644 index 0000000..1e9d0c0 --- /dev/null +++ b/src/flowx/discovery_lineage.py @@ -0,0 +1,136 @@ +"""Source-neutral lineage derivation over the shared discovery AST. + +The parallel of :mod:`flowx.lineage`, which derives a +:class:`~flowx.models.ir.Lineage` block over the Databricks IR. This module +derives the same block over the shared discovery AST +(:mod:`flowx.models.discovery`) instead, so the discover phase can attach lineage +to a :class:`~flowx.models.discovery.SourceGraph` before any IR translation +exists. + +It reuses :mod:`flowx.lineage`'s primitive cores -- :func:`control_edges_from_calls` +and :func:`data_edges_from_endpoints` -- so the two-tier match, the self-edge drop, +and the dedup live in exactly one place and both phases behave identically. It +imports nothing from ``sources/*``: the walk is over the neutral +:class:`~flowx.models.discovery.ContainerNode` branch shape, so an ADF or an +Airflow graph derives lineage through this one code path once its nodes carry +``data_reads`` / ``data_writes`` and (for control edges) the +:data:`INVOKES_WORKFLOW_PROPERTY` marker. +""" + +from __future__ import annotations + +import dataclasses +from collections.abc import Iterator + +from flowx.lineage import control_edges_from_calls, data_edges_from_endpoints +from flowx.models.discovery import ContainerNode, SourceGraph, SourceNode +from flowx.models.ir import ControlEdge, DataEdge, Lineage + +# Neutral node-property key under which a mapper records the workflow a node +# invokes (an ADF ``ExecutePipeline`` callee, an Airflow triggered job). Kept in +# the free-form ``properties`` seam because the invocation target is a per-source +# detail with no shared typed field; :func:`build_graph_control_edges` reads it +# here so the control-edge derivation stays source-agnostic. +INVOKES_WORKFLOW_PROPERTY = "invokes_workflow" +# Companion key: whether the caller waits for the invoked workflow to complete +# (``True`` / ``False``), or absent when the source has no such notion. +INVOKES_WAIT_PROPERTY = "invokes_wait" + + +def walk_nodes(nodes: list[SourceNode]) -> Iterator[SourceNode]: + """Yield every node depth-first, descending into every container branch. + + Recurses through :class:`ContainerNode` branches in their insertion order, so + a Switch's cases *and* its ``default`` branch, a ForEach / Until ``body``, and + both sides of an IfCondition are all reached -- a data asset or an invocation + buried inside a Switch case is still found. + + Args: + nodes: Top-level (or already-nested) node list to walk. + + Yields: + Each node, container nodes included, in depth-first order. + """ + for node in nodes: + yield node + if isinstance(node, ContainerNode): + for children in node.branches.values(): + yield from walk_nodes(children) + + +def build_graph_control_edges(graph: SourceGraph) -> list[ControlEdge]: + """Derive cross-workflow invocation edges for a discovery graph. + + One edge per node that carries the :data:`INVOKES_WORKFLOW_PROPERTY` marker, + found anywhere in the graph (fan-out inside ForEach / If / Switch preserved). + Delegates the self-edge drop, dedup, and unresolved-callee recording to the + shared :func:`~flowx.lineage.control_edges_from_calls`. + + Args: + graph: The source graph to derive control edges for. + + Returns: + Deduplicated control edges, in first-seen order. + """ + + def _calls() -> Iterator[tuple[str, bool | None, str]]: + for node in walk_nodes(graph.tasks): + if INVOKES_WORKFLOW_PROPERTY not in node.properties: + continue + target = node.properties.get(INVOKES_WORKFLOW_PROPERTY) or "" + wait = node.properties.get(INVOKES_WAIT_PROPERTY) + yield str(target), wait, node.task_key + + return control_edges_from_calls(graph.name, _calls()) + + +def build_graph_data_edges(graph: SourceGraph) -> list[DataEdge]: + """Derive proven producer -> consumer data hand-offs for a discovery graph. + + A producer is any node with a ``data_writes`` asset; a consumer any node with a + ``data_reads`` asset, gathered across the whole graph (every container branch + included). Delegates the two-tier match, self-edge drop, and dedup to the shared + :func:`~flowx.lineage.data_edges_from_endpoints`. + + Args: + graph: The source graph to derive data edges for. + + Returns: + Deduplicated data edges, in first-seen order. + """ + nodes = list(walk_nodes(graph.tasks)) + producers = [(node.task_key, asset) for node in nodes for asset in node.data_writes] + consumers = [(node.task_key, asset) for node in nodes for asset in node.data_reads] + return data_edges_from_endpoints(producers, consumers) + + +def build_graph_lineage(graph: SourceGraph) -> Lineage: + """Compose the source-neutral lineage block for a discovery graph. + + Motif annotations are a convert-time IR concern (motifs are detected during + translation, not discovery), so the discovery lineage block leaves them empty. + + Args: + graph: The source graph to derive lineage for. + + Returns: + A :class:`Lineage` with control edges and data edges (motifs empty). + """ + return Lineage( + control_edges=build_graph_control_edges(graph), + data_edges=build_graph_data_edges(graph), + ) + + +def with_graph_lineage(graph: SourceGraph) -> SourceGraph: + """Return a *new* graph carrying its derived lineage, leaving the input untouched. + + Mirrors :func:`flowx.lineage.with_lineage` for the discovery AST. + + Args: + graph: The source graph to copy. + + Returns: + A shallow copy of *graph* with :attr:`SourceGraph.lineage` populated. + """ + return dataclasses.replace(graph, lineage=build_graph_lineage(graph)) diff --git a/src/flowx/discovery_serde.py b/src/flowx/discovery_serde.py index 9dc77fb..ee7a881 100644 --- a/src/flowx/discovery_serde.py +++ b/src/flowx/discovery_serde.py @@ -7,7 +7,10 @@ The DataAsset (de)serialisers are reused from ``ir_serde`` (``data_asset_to_dict`` / ``data_asset_from_dict``) so the physical-asset shape has a single definition -shared by the lineage substrate and the discovery AST. +shared by the lineage substrate and the discovery AST. A graph's derived +:class:`~flowx.models.ir.Lineage` block is serialised through ``ir_serde``'s +``lineage_to_dict`` for the same reason; its inverse (:func:`_lineage_from_dict`) +lives here because ``ir_serde`` ships only the forward direction. Every node dict carries a ``node_type`` discriminator (the dataclass name) so a :class:`~flowx.models.discovery.ContainerNode` or @@ -19,7 +22,7 @@ from typing import Any -from flowx.ir_serde import data_asset_from_dict, data_asset_to_dict +from flowx.ir_serde import data_asset_from_dict, data_asset_to_dict, lineage_to_dict from flowx.models.discovery import ( ContainerNode, GapNode, @@ -30,6 +33,7 @@ SourceGraph, SourceNode, ) +from flowx.models.ir import ControlEdge, DataEdge, Lineage, MotifAnnotation def source_graph_to_dict(graph: SourceGraph) -> dict[str, Any]: @@ -46,6 +50,8 @@ def source_graph_to_dict(graph: SourceGraph) -> dict[str, Any]: result["description"] = graph.description if graph.schedule is not None: result["schedule"] = _schedule_to_dict(graph.schedule) + if graph.lineage is not None: + result["lineage"] = lineage_to_dict(graph.lineage) if graph.properties: result["properties"] = dict(graph.properties) if graph.extensions: @@ -58,6 +64,7 @@ def source_graph_to_dict(graph: SourceGraph) -> dict[str, Any]: def source_graph_from_dict(raw: dict[str, Any]) -> SourceGraph: """Rehydrate a :class:`SourceGraph` from the dict :func:`source_graph_to_dict` emits.""" schedule = raw.get("schedule") + lineage = raw.get("lineage") return SourceGraph( name=raw.get("name", ""), source=raw.get("source", ""), @@ -67,12 +74,54 @@ def source_graph_from_dict(raw: dict[str, Any]) -> SourceGraph: schedule=_schedule_from_dict(schedule) if schedule else None, tags=list(raw.get("tags") or []), tasks=[_node_from_dict(node) for node in raw.get("tasks") or []], + lineage=_lineage_from_dict(lineage) if lineage else None, properties=dict(raw.get("properties") or {}), extensions=dict(raw.get("extensions") or {}), raw=raw.get("raw"), ) +def _lineage_from_dict(raw: dict[str, Any]) -> Lineage: + """Rehydrate a :class:`Lineage` block from the dict ``ir_serde.lineage_to_dict`` emits. + + The inverse of that forward serialiser (which ``ir_serde`` does not itself + ship), so a discovery graph's lineage round-trips through this module. + """ + return Lineage( + control_edges=[ + ControlEdge( + source_workflow=edge.get("source_workflow", ""), + target_workflow=edge.get("target_workflow", ""), + via_task_key=edge.get("via_task_key", ""), + wait_for_completion=edge.get("wait_for_completion"), + resolved=bool(edge.get("resolved", True)), + ) + for edge in raw.get("control_edges") or [] + ], + data_edges=[ + DataEdge( + source_task_key=edge.get("source_task_key", ""), + target_task_key=edge.get("target_task_key", ""), + match_kind=edge.get("match_kind", ""), + match_key=edge.get("match_key", ""), + identity=edge.get("identity"), + asset_type=edge.get("asset_type"), + ) + for edge in raw.get("data_edges") or [] + ], + motifs=[ + MotifAnnotation( + motif_id=motif.get("motif_id", ""), + member_task_keys=list(motif.get("member_task_keys") or []), + display_name=motif.get("display_name"), + databricks_replacement=motif.get("databricks_replacement"), + notes=list(motif.get("notes") or []), + ) + for motif in raw.get("motifs") or [] + ], + ) + + def _parameter_to_dict(spec: ParameterSpec) -> dict[str, Any]: result: dict[str, Any] = {} if spec.type is not None: diff --git a/src/flowx/lineage.py b/src/flowx/lineage.py index a417718..447e265 100644 --- a/src/flowx/lineage.py +++ b/src/flowx/lineage.py @@ -6,6 +6,16 @@ from ``sources/adf`` or ``sources/airflow`` so both front-ends share one code path once they populate ``data_reads`` / ``data_writes`` / ``motif_id``. +The tier-matching, self-edge drop, and dedup rules are factored into two +primitive-level cores -- :func:`control_edges_from_calls` and +:func:`data_edges_from_endpoints` -- that take only ``task_key`` strings and +:class:`DataAsset` values, never a ``Pipeline``. The IR entry points +(:func:`build_control_edges` / :func:`build_data_edges`) gather those primitives +from a pipeline and delegate, and the source-neutral discovery-AST derivation in +:mod:`flowx.discovery_lineage` gathers the same primitives from a +:class:`~flowx.models.discovery.SourceGraph` and delegates too, so both phases +join edges through exactly one implementation. + Everything here is pure: the functions read the pipeline and return new edge lists / a new :class:`Lineage`; nothing is mutated. :func:`with_lineage` attaches a block by returning a *new* ``Pipeline`` rather than mutating the input, unlike @@ -15,7 +25,7 @@ from __future__ import annotations import dataclasses -from collections.abc import Iterator +from collections.abc import Iterable, Iterator from flowx.models.ir import ( Activity, @@ -63,49 +73,43 @@ def walk_activities(activities: list[Activity]) -> Iterator[Activity]: yield from walk_activities(activity.default_activities) -def build_control_edges(pipeline: Pipeline) -> list[ControlEdge]: - """Derive cross-workflow invocation edges for a pipeline. +def control_edges_from_calls( + source_workflow: str, + calls: Iterable[tuple[str, bool | None, str]], +) -> list[ControlEdge]: + """Assemble deduplicated control edges from raw invocation primitives. - Emits one :class:`ControlEdge` per invoking activity -- an - ``ExecutePipelineActivity`` (ADF) or a ``RunJobActivity`` (Airflow) -- found - anywhere in the pipeline, including inside ForEach / If / Switch containers - (fan-out is preserved: each call site is its own edge). Edges whose callee - equals the caller are dropped (no self-edges), and identical edges are - collapsed (no duplicates). An unresolved callee is recorded with - ``resolved=False`` rather than dropped. + The shared core behind :func:`build_control_edges` (IR) and the discovery-AST + control-edge derivation: it owns the self-edge drop, the dedup, and the + unresolved-callee recording so those rules live in exactly one place and both + phases behave identically. Args: - pipeline: The translated pipeline IR. + source_workflow: Name of the calling workflow (pipeline / DAG). + calls: One ``(target_workflow, wait_for_completion, via_task_key)`` triple + per call site, in the order they should be considered. ``target_workflow`` + may be empty when the callee could not be resolved from a partial export. Returns: - Deduplicated list of control edges, in first-seen order. + Deduplicated control edges in first-seen order. A call whose target equals + ``source_workflow`` is dropped (no self-edge); an empty target is kept with + ``resolved=False`` rather than dropped. """ edges: list[ControlEdge] = [] seen: set[tuple[str, str, str]] = set() - for activity in walk_activities(pipeline.tasks): - target: str | None - wait: bool | None - match activity: - case ExecutePipelineActivity(): - target = activity.pipeline_name - wait = activity.wait_on_completion - case RunJobActivity(): - target = activity.job_name - wait = None - case _: - continue + for target, wait, via_task_key in calls: target_name = target or "" - if target_name and target_name == pipeline.name: + if target_name and target_name == source_workflow: continue - key = (pipeline.name, target_name, activity.task_key) + key = (source_workflow, target_name, via_task_key) if key in seen: continue seen.add(key) edges.append( ControlEdge( - source_workflow=pipeline.name, + source_workflow=source_workflow, target_workflow=target_name, - via_task_key=activity.task_key, + via_task_key=via_task_key, wait_for_completion=wait, resolved=bool(target_name), ) @@ -113,6 +117,35 @@ def build_control_edges(pipeline: Pipeline) -> list[ControlEdge]: return edges +def build_control_edges(pipeline: Pipeline) -> list[ControlEdge]: + """Derive cross-workflow invocation edges for a pipeline. + + Emits one :class:`ControlEdge` per invoking activity -- an + ``ExecutePipelineActivity`` (ADF) or a ``RunJobActivity`` (Airflow) -- found + anywhere in the pipeline, including inside ForEach / If / Switch containers + (fan-out is preserved: each call site is its own edge). Edges whose callee + equals the caller are dropped (no self-edges), and identical edges are + collapsed (no duplicates). An unresolved callee is recorded with + ``resolved=False`` rather than dropped. + + Args: + pipeline: The translated pipeline IR. + + Returns: + Deduplicated list of control edges, in first-seen order. + """ + + def _calls() -> Iterator[tuple[str, bool | None, str]]: + for activity in walk_activities(pipeline.tasks): + match activity: + case ExecutePipelineActivity(): + yield activity.pipeline_name or "", activity.wait_on_completion, activity.task_key + case RunJobActivity(): + yield activity.job_name or "", None, activity.task_key + + return control_edges_from_calls(pipeline.name, _calls()) + + def _match_assets(producer: DataAsset, consumer: DataAsset) -> tuple[str, str, str | None] | None: """Decide whether a written asset hands off to a read asset, and how. @@ -137,30 +170,32 @@ def _match_assets(producer: DataAsset, consumer: DataAsset) -> tuple[str, str, s return None -def build_data_edges(pipeline: Pipeline) -> list[DataEdge]: - """Derive proven producer -> consumer data hand-offs for a pipeline. +def data_edges_from_endpoints( + producers: Iterable[tuple[str, DataAsset]], + consumers: Iterable[tuple[str, DataAsset]], +) -> list[DataEdge]: + """Join producer endpoints to consumer endpoints via the two-tier match. - A producer is any activity with a ``data_writes`` asset; a consumer any - activity with a ``data_reads`` asset, gathered across the whole pipeline - (ForEach / If / Switch bodies included). Each producer asset is joined against - each consumer asset via :func:`_match_assets`, tagging the edge as an - ``identity`` or ``signature`` match. An activity never hands off to itself - (no self-edges), and identical edges are collapsed (no duplicates). + The shared core behind :func:`build_data_edges` (IR) and the discovery-AST + data-edge derivation: it owns the :func:`_match_assets` tier logic, the + no-self-edge rule, and the dedup, so both phases join identically. Args: - pipeline: The translated pipeline IR. + producers: ``(task_key, written asset)`` pairs, in first-seen order. + consumers: ``(task_key, read asset)`` pairs, in first-seen order. Returns: - Deduplicated list of data edges, in first-seen order. + Deduplicated data edges in first-seen order. A producer never hands off to + a consumer sharing its ``task_key`` (no self-edge), and identical + ``(producer, consumer, match_kind, match_key)`` edges are collapsed. """ - activities = list(walk_activities(pipeline.tasks)) - producers = [(activity.task_key, asset) for activity in activities for asset in activity.data_writes] - consumers = [(activity.task_key, asset) for activity in activities for asset in activity.data_reads] + producer_list = list(producers) + consumer_list = list(consumers) edges: list[DataEdge] = [] seen: set[tuple[str, str, str, str]] = set() - for producer_key, producer_asset in producers: - for consumer_key, consumer_asset in consumers: + for producer_key, producer_asset in producer_list: + for consumer_key, consumer_asset in consumer_list: if producer_key == consumer_key: continue matched = _match_assets(producer_asset, consumer_asset) @@ -184,6 +219,28 @@ def build_data_edges(pipeline: Pipeline) -> list[DataEdge]: return edges +def build_data_edges(pipeline: Pipeline) -> list[DataEdge]: + """Derive proven producer -> consumer data hand-offs for a pipeline. + + A producer is any activity with a ``data_writes`` asset; a consumer any + activity with a ``data_reads`` asset, gathered across the whole pipeline + (ForEach / If / Switch bodies included). Each producer asset is joined against + each consumer asset via :func:`_match_assets`, tagging the edge as an + ``identity`` or ``signature`` match. An activity never hands off to itself + (no self-edges), and identical edges are collapsed (no duplicates). + + Args: + pipeline: The translated pipeline IR. + + Returns: + Deduplicated list of data edges, in first-seen order. + """ + activities = list(walk_activities(pipeline.tasks)) + producers = [(activity.task_key, asset) for activity in activities for asset in activity.data_writes] + consumers = [(activity.task_key, asset) for activity in activities for asset in activity.data_reads] + return data_edges_from_endpoints(producers, consumers) + + def build_motif_annotations(pipeline: Pipeline) -> list[MotifAnnotation]: """Derive motif annotations from the collapsed motif activities in a pipeline. diff --git a/src/flowx/models/discovery.py b/src/flowx/models/discovery.py index 1590c0b..971d89f 100644 --- a/src/flowx/models/discovery.py +++ b/src/flowx/models/discovery.py @@ -35,7 +35,7 @@ from dataclasses import dataclass, field from typing import Any -from flowx.models.ir import DataAsset +from flowx.models.ir import DataAsset, Lineage # --------------------------------------------------------------------------- # # Well-known discriminators and concepts (open vocabularies). @@ -248,6 +248,11 @@ class SourceGraph: tags: Free-form label list (ADF annotations, Airflow user tags). tasks: Top-level nodes; control flow nests further nodes via :class:`ContainerNode`. + lineage: Source-neutral lineage block (control/data edges) derived over + this graph, or ``None`` when lineage has not been derived. Reuses the + #61 :class:`~flowx.models.ir.Lineage` type so the discovery AST and + the IR lineage substrate share one vocabulary; populated by + :mod:`flowx.discovery_lineage` in the discover phase. properties: Free-form bag for graph-level platform-specific attributes. extensions: Alias-free overflow for anything else source-specific (e.g. ADF ``folder``, Airflow global Variables). @@ -262,6 +267,7 @@ class SourceGraph: schedule: ScheduleSpec | None = None tags: list[str] = field(default_factory=list) tasks: list[SourceNode] = field(default_factory=list) + lineage: Lineage | None = None properties: dict[str, Any] = field(default_factory=dict) extensions: dict[str, Any] = field(default_factory=dict) raw: dict[str, Any] | None = None diff --git a/tests/unit/test_discovery_lineage.py b/tests/unit/test_discovery_lineage.py new file mode 100644 index 0000000..2d35c85 --- /dev/null +++ b/tests/unit/test_discovery_lineage.py @@ -0,0 +1,179 @@ +"""Tests for lineage over the shared discovery AST (:mod:`flowx.discovery_lineage`). + +The source-neutral derivation itself -- the identity vs signature join tiers, no +self-edges, no duplicates, fan-out, and Switch / ForEach / If recursion -- driven +straight off :class:`SourceGraph` / :class:`SourceNode` values. Per-source +population of a node's reads / writes / invocation markers (the ADF and Airflow +mappers) is exercised in the per-source test suites (#62 / #63); this file stays +free of any ``sources/*`` coupling. +""" + +from __future__ import annotations + +from flowx.discovery_lineage import ( + INVOKES_WAIT_PROPERTY, + INVOKES_WORKFLOW_PROPERTY, + build_graph_lineage, + walk_nodes, + with_graph_lineage, +) +from flowx.models.discovery import ( + SOURCE_ADF, + ContainerNode, + SourceGraph, + SourceNode, +) +from flowx.models.ir import DataAsset + +# --------------------------------------------------------------------------- # +# Helpers for the source-neutral layer +# --------------------------------------------------------------------------- # + + +def _node(task_key: str, *, reads=None, writes=None, invokes=None, wait=None) -> SourceNode: + properties: dict = {} + if invokes is not None: + properties[INVOKES_WORKFLOW_PROPERTY] = invokes + properties[INVOKES_WAIT_PROPERTY] = wait + return SourceNode( + source_id=task_key, + task_key=task_key, + concept="x", + source=SOURCE_ADF, + data_reads=list(reads or []), + data_writes=list(writes or []), + properties=properties, + ) + + +def _graph(*tasks: SourceNode, name: str = "pl") -> SourceGraph: + return SourceGraph(name=name, source=SOURCE_ADF, tasks=list(tasks)) + + +# --------------------------------------------------------------------------- # +# Data-edge tiers (source-neutral) +# --------------------------------------------------------------------------- # + + +def test_identity_tier_joins_across_different_signatures() -> None: + graph = _graph( + _node("writer", writes=[DataAsset(signature="ds_out", identity="curated.orders")]), + _node("reader", reads=[DataAsset(signature="ds_in_other", identity="curated.orders")]), + ) + edges = build_graph_lineage(graph).data_edges + assert len(edges) == 1 + assert (edges[0].source_task_key, edges[0].target_task_key) == ("writer", "reader") + assert edges[0].match_kind == "identity" + assert edges[0].identity == "curated.orders" + + +def test_signature_tier_when_identity_unresolved() -> None: + graph = _graph( + _node("writer", writes=[DataAsset(signature="FP[wm|slots=1]/FN[v.txt|slots=0]")]), + _node("reader", reads=[DataAsset(signature="FP[wm|slots=1]/FN[v.txt|slots=0]")]), + ) + edges = build_graph_lineage(graph).data_edges + assert len(edges) == 1 + assert edges[0].match_kind == "signature" + assert edges[0].identity is None + + +def test_distinct_identities_do_not_fall_back_to_signature() -> None: + """Two resolved-but-different identities never manufacture a signature edge (#36).""" + graph = _graph( + _node("writer", writes=[DataAsset(signature="shared", identity="a.first")]), + _node("reader", reads=[DataAsset(signature="shared", identity="b.second")]), + ) + assert build_graph_lineage(graph).data_edges == [] + + +def test_fan_out_one_writer_many_readers() -> None: + graph = _graph( + _node("writer", writes=[DataAsset(signature="ds", identity="x.y")]), + _node("reader_one", reads=[DataAsset(signature="ds", identity="x.y")]), + _node("reader_two", reads=[DataAsset(signature="ds", identity="x.y")]), + ) + edges = build_graph_lineage(graph).data_edges + assert {edge.target_task_key for edge in edges} == {"reader_one", "reader_two"} + + +def test_no_self_edge_and_no_duplicates() -> None: + graph = _graph( + _node( + "both", + writes=[DataAsset(signature="ds", identity="x.y"), DataAsset(signature="ds", identity="x.y")], + reads=[DataAsset(signature="ds", identity="x.y")], + ), + _node("reader", reads=[DataAsset(signature="ds", identity="x.y")]), + ) + edges = build_graph_lineage(graph).data_edges + assert [(edge.source_task_key, edge.target_task_key) for edge in edges] == [("both", "reader")] + + +def test_data_edges_recurse_into_switch_and_foreach_branches() -> None: + """A writer buried in a Switch case hands off to a reader in a ForEach body.""" + writer = _node("writer", writes=[DataAsset(signature="ds", identity="x.y")]) + reader = _node("reader", reads=[DataAsset(signature="ds", identity="x.y")]) + switch = ContainerNode( + source_id="sw", + task_key="sw", + concept="switch", + source=SOURCE_ADF, + branches={"caseA": [writer], "default": []}, + ) + loop = ContainerNode( + source_id="fe", + task_key="fe", + concept="loop", + source=SOURCE_ADF, + branches={"body": [reader]}, + ) + edges = build_graph_lineage(_graph(switch, loop)).data_edges + assert [(edge.source_task_key, edge.target_task_key) for edge in edges] == [("writer", "reader")] + + +def test_walk_nodes_visits_every_branch_including_default() -> None: + switch = ContainerNode( + source_id="sw", + task_key="sw", + concept="switch", + source=SOURCE_ADF, + branches={"caseA": [_node("in_case")], "default": [_node("in_default")]}, + ) + keys = [node.task_key for node in walk_nodes([switch])] + assert keys == ["sw", "in_case", "in_default"] + + +# --------------------------------------------------------------------------- # +# Control edges (source-neutral) +# --------------------------------------------------------------------------- # + + +def test_control_edge_from_invocation_marker() -> None: + graph = _graph(_node("call", invokes="child", wait=True), name="parent") + edges = build_graph_lineage(graph).control_edges + assert len(edges) == 1 + assert (edges[0].source_workflow, edges[0].target_workflow) == ("parent", "child") + assert edges[0].wait_for_completion is True + assert edges[0].resolved is True + + +def test_control_edge_unresolved_callee_recorded_not_dropped() -> None: + graph = _graph(_node("call", invokes="", wait=True), name="parent") + edges = build_graph_lineage(graph).control_edges + assert len(edges) == 1 + assert edges[0].target_workflow == "" + assert edges[0].resolved is False + + +def test_control_edge_self_call_is_dropped() -> None: + graph = _graph(_node("call", invokes="parent", wait=True), name="parent") + assert build_graph_lineage(graph).control_edges == [] + + +def test_with_graph_lineage_returns_new_graph_and_leaves_input_untouched() -> None: + graph = _graph(_node("call", invokes="child", wait=False), name="parent") + result = with_graph_lineage(graph) + assert graph.lineage is None + assert result is not graph + assert len(result.lineage.control_edges) == 1 diff --git a/tests/unit/test_discovery_model.py b/tests/unit/test_discovery_model.py index 2e1191e..df3cec6 100644 --- a/tests/unit/test_discovery_model.py +++ b/tests/unit/test_discovery_model.py @@ -26,7 +26,7 @@ SourceGraph, SourceNode, ) -from flowx.models.ir import DataAsset +from flowx.models.ir import ControlEdge, DataAsset, DataEdge, Lineage def test_gap_node_defaults_to_gap_concept(): @@ -214,3 +214,38 @@ def test_empty_graph_emits_lists_and_dicts_never_null(): assert serialised["variables"] == {} # Airflow has no graph-scoped variables; the field stays empty rather than absent. assert source_graph_from_dict(serialised).variables == {} + + +def test_lineage_block_round_trips_and_is_absent_when_none(): + """A graph's derived lineage survives serialise<->deserialise; None stays absent.""" + graph = SourceGraph( + name="pl", + source=SOURCE_ADF, + lineage=Lineage( + control_edges=[ + ControlEdge( + source_workflow="pl", + target_workflow="child", + via_task_key="Run Child", + wait_for_completion=False, + ) + ], + data_edges=[ + DataEdge( + source_task_key="writer", + target_task_key="reader", + match_kind="identity", + match_key="curated.orders", + identity="curated.orders", + asset_type="table", + ) + ], + ), + ) + reloaded = source_graph_from_dict(json.loads(json.dumps(source_graph_to_dict(graph)))) + assert reloaded == graph + + # A graph with no derived lineage omits the key entirely and rehydrates to None. + bare = source_graph_to_dict(SourceGraph(name="bare", source=SOURCE_ADF)) + assert "lineage" not in bare + assert source_graph_from_dict(bare).lineage is None From bcd1472a20e64888a6700f51432e2b1e401dbc93 Mon Sep 17 00:00:00 2001 From: Matthew Moorcroft Date: Tue, 15 Sep 2026 12:03:16 +0100 Subject: [PATCH 6/7] Make shared AST fit Airflow: node run_condition, graph default policy, open best-effort data lineage (#61) Three additive Airflow-neutrality fixes to the shared discovery-AST standard so it fits Airflow as well as ADF before #63 is built against it. All fields default to None/empty, so ADF and the approved #61 substrate are unaffected. FIX 1 -- node-level run condition. Airflow's `trigger_rule` is one rule on the downstream node quantified over all its upstreams together, so it cannot be represented per-edge. Add `SourceNode.run_condition: str | None = None` for it; `SourceDependency.conditions` stays the per-edge outcome list (the ADF shape). Docstrings state which source each covers. FIX 2 -- graph-level default policy. Airflow declares retries/retry_delay/ timeout once in DAG `default_args` (cascading to every task) plus a whole-run `dagrun_timeout`. Add `SourceGraph.default_policy: PolicySpec | None = None` (cascade: applies unless a node sets its own policy) and `SourceGraph.run_timeout_seconds: int | None = None`. ADF leaves both None. FIX 3 -- data lineage is an open, best-effort standard. Document on DataAsset (models/ir.py) and the data_reads/data_writes fields (models/discovery.py) that these are a general best-effort description of what a task reads/writes: identity is the strong physical id when resolvable (else None, never guessed), signature is the always-present weak descriptor, and asset_type is an OPEN kind that explicitly allows non-physical/logical/value hand-offs (e.g. an Airflow XCom -> "value", a logical dataset -> "logical"), not physical-only. Population is best-effort and not required to be complete for either source; empty is valid. No change to the two-tier match derivation -- vocabulary + semantics only. Serialize/deserialize all new fields in discovery_serde.py. Tests: run_condition round-trip + default-None, graph default_policy + run_timeout_seconds round-trip + default-None, and a non-physical asset_type ("value") round-trip with empty reads/writes proven valid. Co-authored-by: Isaac --- src/flowx/discovery_serde.py | 10 ++++ src/flowx/models/discovery.py | 53 ++++++++++++++++----- src/flowx/models/ir.py | 31 ++++++++---- tests/unit/test_discovery_model.py | 76 +++++++++++++++++++++++++++++- 4 files changed, 149 insertions(+), 21 deletions(-) diff --git a/src/flowx/discovery_serde.py b/src/flowx/discovery_serde.py index ee7a881..56563d7 100644 --- a/src/flowx/discovery_serde.py +++ b/src/flowx/discovery_serde.py @@ -50,6 +50,10 @@ def source_graph_to_dict(graph: SourceGraph) -> dict[str, Any]: result["description"] = graph.description if graph.schedule is not None: result["schedule"] = _schedule_to_dict(graph.schedule) + if graph.default_policy is not None: + result["default_policy"] = _policy_to_dict(graph.default_policy) + if graph.run_timeout_seconds is not None: + result["run_timeout_seconds"] = graph.run_timeout_seconds if graph.lineage is not None: result["lineage"] = lineage_to_dict(graph.lineage) if graph.properties: @@ -64,6 +68,7 @@ def source_graph_to_dict(graph: SourceGraph) -> dict[str, Any]: def source_graph_from_dict(raw: dict[str, Any]) -> SourceGraph: """Rehydrate a :class:`SourceGraph` from the dict :func:`source_graph_to_dict` emits.""" schedule = raw.get("schedule") + default_policy = raw.get("default_policy") lineage = raw.get("lineage") return SourceGraph( name=raw.get("name", ""), @@ -72,6 +77,8 @@ def source_graph_from_dict(raw: dict[str, Any]) -> SourceGraph: parameters={name: _parameter_from_dict(spec) for name, spec in (raw.get("parameters") or {}).items()}, variables={name: _parameter_from_dict(spec) for name, spec in (raw.get("variables") or {}).items()}, schedule=_schedule_from_dict(schedule) if schedule else None, + default_policy=_policy_from_dict(default_policy) if default_policy else None, + run_timeout_seconds=raw.get("run_timeout_seconds"), tags=list(raw.get("tags") or []), tasks=[_node_from_dict(node) for node in raw.get("tasks") or []], lineage=_lineage_from_dict(lineage) if lineage else None, @@ -213,6 +220,8 @@ def _node_to_dict(node: SourceNode) -> dict[str, Any]: result["name"] = node.name if node.native_type is not None: result["native_type"] = node.native_type + if node.run_condition is not None: + result["run_condition"] = node.run_condition if node.policy is not None: result["policy"] = _policy_to_dict(node.policy) if node.properties: @@ -243,6 +252,7 @@ def _node_from_dict(raw: dict[str, Any]) -> SourceNode: "source": raw.get("source", ""), "name": raw.get("name"), "native_type": raw.get("native_type"), + "run_condition": raw.get("run_condition"), "dependencies": [_dependency_from_dict(dependency) for dependency in raw.get("dependencies") or []], "policy": _policy_from_dict(raw["policy"]) if raw.get("policy") else None, "data_reads": [data_asset_from_dict(asset) for asset in raw.get("data_reads") or []], diff --git a/src/flowx/models/discovery.py b/src/flowx/models/discovery.py index 971d89f..dbd3460 100644 --- a/src/flowx/models/discovery.py +++ b/src/flowx/models/discovery.py @@ -114,7 +114,12 @@ class ParameterSpec: @dataclass(slots=True, kw_only=True) class PolicySpec: - """Retry / timeout policy on a node, normalised to seconds. + """Retry / timeout policy, normalised to seconds. + + Used both per node (:attr:`SourceNode.policy`) and as a cascading graph + default (:attr:`SourceGraph.default_policy`), e.g. an Airflow DAG's + ``default_args`` retries/timeout that apply to every task unless a task + overrides them. Only the genuinely shared retry/timeout knobs are typed; source-specific policy (ADF ``secure_input`` / ``secure_output``, Airflow retry-delay @@ -137,14 +142,18 @@ class PolicySpec: class SourceDependency: """A dependency edge from a node to one upstream node. - Conditions stay a **list** -- ADF dependency edges carry one or more - outcome conditions (``["Succeeded", "Skipped"]``), so collapsing them to a - single outcome would lose information. Airflow's unconditional edges use an - empty list (or a single normalised condition). + ``conditions`` is a **per-edge** outcome list -- the natural ADF shape, + where each individual edge carries the upstream outcome(s) that must hold + (``["Succeeded", "Skipped"]``). This is distinct from Airflow's + ``trigger_rule``, which is a single rule on the *downstream* node quantified + over all its upstreams together and so cannot live on an edge; that rides in + :attr:`SourceNode.run_condition` instead. Attributes: upstream: Task key of the upstream node this edge depends on. - conditions: Required upstream outcome(s); empty when unconditional. + conditions: Required upstream outcome(s) for THIS edge; empty when + unconditional. Per-edge (ADF); see :attr:`SourceNode.run_condition` + for Airflow's node-level ``trigger_rule``. resolved: ``False`` when the upstream could not be resolved (e.g. a partial export); recorded, not dropped. """ @@ -171,11 +180,22 @@ class SourceNode: native_type: The source's own type string, preserved verbatim (ADF ``"Copy"`` / ``"DatabricksNotebook"``, Airflow ``"BashOperator"`` / a TaskFlow decorator). - dependencies: Upstream dependency edges. - policy: Retry / timeout policy, when the source declares one. - data_reads: Physical data assets this node reads (reuses #61 - :class:`~flowx.models.ir.DataAsset`). - data_writes: Physical data assets this node writes. + dependencies: Upstream dependency edges, each with its own per-edge + outcome conditions (the ADF shape). + run_condition: Node-level aggregate run rule quantified over all this + node's upstreams together -- Airflow's ``trigger_rule`` (e.g. + ``"all_success"``, ``"none_failed_min_one_success"``, + ``"one_failed"``). ``None`` when the source has no such notion + (ADF, which expresses outcomes per-edge on :attr:`dependencies`). + policy: Retry / timeout policy, when the source declares one on the + node; falls back to :attr:`SourceGraph.default_policy` otherwise. + data_reads: Best-effort description of what this node reads, as + :class:`~flowx.models.ir.DataAsset` values (reuses #61). Population + is best-effort and open -- physical tables/files, logical datasets, + or value hand-offs (an Airflow XCom) alike -- and an empty list is + valid where the source records nothing. + data_writes: Best-effort description of what this node writes; same + open, best-effort semantics as :attr:`data_reads`. properties: Free-form bag for platform-specific attributes with no shared typed field yet. Connection / linked-service details live here for now -- a typed connection field is a deliberate future @@ -190,6 +210,7 @@ class SourceNode: name: str | None = None native_type: str | None = None dependencies: list[SourceDependency] = field(default_factory=list) + run_condition: str | None = None policy: PolicySpec | None = None data_reads: list[DataAsset] = field(default_factory=list) data_writes: list[DataAsset] = field(default_factory=list) @@ -245,6 +266,14 @@ class SourceGraph: no graph-scoped variable concept (Airflow); such sources' global variables ride in :attr:`extensions`. schedule: Workflow schedule, when one is declared. + default_policy: Cascading default retry / timeout policy applied to + every node that does not set its own :attr:`SourceNode.policy` -- + Airflow's DAG ``default_args`` (retries / retry_delay / timeout). + ``None`` when the source has no graph-level default (ADF, which + declares policy per activity). + run_timeout_seconds: Whole-run timeout for one execution of the + workflow (Airflow's ``dagrun_timeout``), or ``None`` when the source + declares none. tags: Free-form label list (ADF annotations, Airflow user tags). tasks: Top-level nodes; control flow nests further nodes via :class:`ContainerNode`. @@ -265,6 +294,8 @@ class SourceGraph: parameters: dict[str, ParameterSpec] = field(default_factory=dict) variables: dict[str, ParameterSpec] = field(default_factory=dict) schedule: ScheduleSpec | None = None + default_policy: PolicySpec | None = None + run_timeout_seconds: int | None = None tags: list[str] = field(default_factory=list) tasks: list[SourceNode] = field(default_factory=list) lineage: Lineage | None = None diff --git a/src/flowx/models/ir.py b/src/flowx/models/ir.py index 94def2f..b6df8c6 100644 --- a/src/flowx/models/ir.py +++ b/src/flowx/models/ir.py @@ -57,26 +57,39 @@ class Dependency: @dataclass(slots=True, kw_only=True) class DataAsset: - """A physical data source or sink an activity reads from or writes to. + """A general, best-effort description of something a task reads or writes. Deliberately source-neutral: the same shape describes an ADF dataset, an - Airflow dataset/hook target, or any other front-end's data reference, so the - lineage substrate never has to know which source produced it. + Airflow dataset/hook target, an XCom / TaskFlow return value, or any other + front-end's data reference, so the lineage substrate never has to know which + source produced it. + + Population is **best-effort and open**, not physical-only. Each source fills + in what it can -- fully, partially, or not at all -- and an empty + ``data_reads`` / ``data_writes`` is valid and expected in places. The asset + need not be a physically-resolved table or file: a logical dataset or a pure + value hand-off (e.g. an Airflow XCom) is an equally valid asset a source MAY + record. Two-tier identity (from #36): ``identity`` is the resolved physical location (``schema.table`` or a concrete storage path) and is the strong join key. It is ``None`` when it cannot be resolved deterministically -- never a guess. ``signature`` is always present and carries the neutral fallback descriptor - (a dataset name, a normalised reference, or an expression) so two assets can - still be compared when neither side resolved to a physical identity. + (a dataset name, a normalised reference, a value key, or an expression) so + two assets can still be compared when neither side resolved to a physical + identity. Attributes: signature: Neutral, always-present descriptor used as the weak join key - (e.g. dataset name or normalised expression). + (e.g. dataset name, value/XCom key, or normalised expression). identity: Resolved physical identity used as the strong join key, or - ``None`` when it could not be resolved deterministically. - asset_type: Neutral kind of the asset (``"table"`` / ``"file"`` / - ``"volume"`` / ...), or ``None`` when unknown. + ``None`` when it could not be resolved deterministically (or the + asset is non-physical). Never guessed. + asset_type: Open, neutral kind of the asset -- physical kinds like + ``"table"`` / ``"file"`` / ``"volume"`` as well as non-physical / + logical ones like ``"value"`` (an XCom / TaskFlow return) or + ``"logical"`` (a named logical dataset). The vocabulary is not + closed; ``None`` when unknown. properties: Free-form extra attributes carried through verbatim. """ diff --git a/tests/unit/test_discovery_model.py b/tests/unit/test_discovery_model.py index df3cec6..39b1eac 100644 --- a/tests/unit/test_discovery_model.py +++ b/tests/unit/test_discovery_model.py @@ -9,7 +9,7 @@ import json -from flowx.discovery_serde import _node_from_dict, source_graph_from_dict, source_graph_to_dict +from flowx.discovery_serde import _node_from_dict, _node_to_dict, source_graph_from_dict, source_graph_to_dict from flowx.models.discovery import ( CONCEPT_COPY_DATA, CONCEPT_GAP, @@ -249,3 +249,77 @@ def test_lineage_block_round_trips_and_is_absent_when_none(): bare = source_graph_to_dict(SourceGraph(name="bare", source=SOURCE_ADF)) assert "lineage" not in bare assert source_graph_from_dict(bare).lineage is None + + +def test_node_run_condition_round_trips_and_defaults_none(): + """Airflow's node-level trigger_rule round-trips; ADF-style nodes leave it None.""" + node = SourceNode( + source_id="join", + task_key="join", + concept=CONCEPT_NOTEBOOK, + source=SOURCE_AIRFLOW, + run_condition="none_failed_min_one_success", + ) + graph = SourceGraph(name="dag", source=SOURCE_AIRFLOW, tasks=[node]) + + reloaded = source_graph_from_dict(json.loads(json.dumps(source_graph_to_dict(graph)))) + assert reloaded == graph + assert reloaded.tasks[0].run_condition == "none_failed_min_one_success" + + # ADF/substrate default: a node that sets no run_condition omits the key and stays None. + plain = SourceNode(source_id="a", task_key="a", concept=CONCEPT_NOTEBOOK, source=SOURCE_ADF) + assert "run_condition" not in _node_to_dict(plain) + assert _node_from_dict(_node_to_dict(plain)).run_condition is None + + +def test_graph_default_policy_and_run_timeout_round_trip_and_default_none(): + """Airflow DAG default_args cascade + dagrun_timeout round-trip; ADF leaves both None.""" + graph = SourceGraph( + name="dag", + source=SOURCE_AIRFLOW, + default_policy=PolicySpec(max_retries=3, retry_interval_seconds=300, extensions={"owner": "data"}), + run_timeout_seconds=7200, + ) + + reloaded = source_graph_from_dict(json.loads(json.dumps(source_graph_to_dict(graph)))) + assert reloaded == graph + assert reloaded.default_policy == PolicySpec( + max_retries=3, retry_interval_seconds=300, extensions={"owner": "data"} + ) + assert reloaded.run_timeout_seconds == 7200 + + # ADF/substrate default: no graph-level policy or run timeout -> keys absent, fields None. + bare = source_graph_to_dict(SourceGraph(name="pl", source=SOURCE_ADF)) + assert "default_policy" not in bare + assert "run_timeout_seconds" not in bare + rehydrated = source_graph_from_dict(bare) + assert rehydrated.default_policy is None + assert rehydrated.run_timeout_seconds is None + + +def test_non_physical_asset_type_and_empty_reads_writes_are_valid(): + """A value/logical asset_type round-trips, and empty reads/writes are valid (best-effort).""" + producer = SourceNode( + source_id="extract", + task_key="extract", + concept=CONCEPT_NOTEBOOK, + source=SOURCE_AIRFLOW, + # An Airflow XCom / TaskFlow return value: no physical identity, an open non-physical kind. + data_writes=[DataAsset(signature="extract:return_value", asset_type="value")], + ) + # Best-effort population: a node may record nothing at all. + consumer = SourceNode( + source_id="load", + task_key="load", + concept=CONCEPT_NOTEBOOK, + source=SOURCE_AIRFLOW, + ) + graph = SourceGraph(name="dag", source=SOURCE_AIRFLOW, tasks=[producer, consumer]) + + reloaded = source_graph_from_dict(json.loads(json.dumps(source_graph_to_dict(graph)))) + assert reloaded == graph + assert reloaded.tasks[0].data_writes == [DataAsset(signature="extract:return_value", asset_type="value")] + assert reloaded.tasks[0].data_writes[0].identity is None + # Empty reads/writes survive as empty lists, never None. + assert reloaded.tasks[1].data_reads == [] + assert reloaded.tasks[1].data_writes == [] From 23032e5f47226809144d362299d20be10b4991fb Mon Sep 17 00:00:00 2001 From: Matthew Moorcroft Date: Tue, 15 Sep 2026 12:13:02 +0100 Subject: [PATCH 7/7] Fix stale physical-only wording in shared data-asset docstrings (#61) Docstring-only cleanup. The module docstrings for discovery.py and discovery_serde.py still described the shared data shape as physical-only, contradicting the open/best-effort data-lineage decision the amendment made. Reword both to match the already-updated DataAsset docstring: a general, best-effort description of what a task reads/writes -- identity is the strong physical id when resolvable (else None, never guessed), signature is the always-present weak descriptor, and asset_type is an open kind that also covers non-physical / logical / value hand-offs (e.g. an Airflow XCom). No code, fields, or logic changed. Co-authored-by: Isaac --- src/flowx/discovery_serde.py | 8 ++++++-- src/flowx/models/discovery.py | 7 ++++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/src/flowx/discovery_serde.py b/src/flowx/discovery_serde.py index 56563d7..0587c7f 100644 --- a/src/flowx/discovery_serde.py +++ b/src/flowx/discovery_serde.py @@ -6,8 +6,12 @@ its own round-trip pair so evolving one shape never disturbs the other. The DataAsset (de)serialisers are reused from ``ir_serde`` (``data_asset_to_dict`` -/ ``data_asset_from_dict``) so the physical-asset shape has a single definition -shared by the lineage substrate and the discovery AST. A graph's derived +/ ``data_asset_from_dict``) so the data-asset shape has a single definition +shared by the lineage substrate and the discovery AST. That shape is a general, +best-effort description of what a task reads/writes -- not physical-only: a +resolvable physical ``identity`` when there is one (else ``None``), an +always-present ``signature``, and an open ``asset_type`` that also covers +non-physical / logical / value hand-offs (e.g. an Airflow XCom). A graph's derived :class:`~flowx.models.ir.Lineage` block is serialised through ``ir_serde``'s ``lineage_to_dict`` for the same reason; its inverse (:func:`_lineage_from_dict`) lives here because ``ir_serde`` ships only the forward direction. diff --git a/src/flowx/models/discovery.py b/src/flowx/models/discovery.py index dbd3460..2a4a354 100644 --- a/src/flowx/models/discovery.py +++ b/src/flowx/models/discovery.py @@ -27,7 +27,12 @@ The lineage primitives from #61 (:class:`~flowx.models.ir.DataAsset`) are reused here rather than duplicated: a node's reads and writes are lists of ``DataAsset``, so the discovery AST and the lineage substrate share one -vocabulary for physical data references. +vocabulary for describing data references. That description is general and +best-effort, not physical-only -- ``identity`` is the strong physical id when +resolvable (else ``None``, never guessed), ``signature`` is the always-present +weak descriptor, and ``asset_type`` is an open kind that also covers +non-physical / logical / value hand-offs (e.g. an Airflow XCom). Each source +populates what it can; empty reads/writes are valid. """ from __future__ import annotations