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/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 new file mode 100644 index 0000000..0587c7f --- /dev/null +++ b/src/flowx/discovery_serde.py @@ -0,0 +1,285 @@ +"""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 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. + +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, lineage_to_dict +from flowx.models.discovery import ( + ContainerNode, + GapNode, + ParameterSpec, + PolicySpec, + ScheduleSpec, + SourceDependency, + SourceGraph, + SourceNode, +) +from flowx.models.ir import ControlEdge, DataEdge, Lineage, MotifAnnotation + + +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.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: + 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") + default_policy = raw.get("default_policy") + lineage = raw.get("lineage") + 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, + 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, + 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: + 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.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 + + +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 {}), + ) + + +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.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: + 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``. + + 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", ""), + "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 []], + "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"), + } + 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, + branches={ + label: [_node_from_dict(child) for child in children] + for label, children in (raw.get("branches") or {}).items() + }, + ) + return SourceNode(**common) diff --git a/src/flowx/ir_serde.py b/src/flowx/ir_serde.py index 51b3d32..72b00a6 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,82 @@ 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 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. + + 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 +218,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 +418,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..447e265 --- /dev/null +++ b/src/flowx/lineage.py @@ -0,0 +1,308 @@ +"""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``. + +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 +the in-place dependency rewrite in ``motifs/collapser.py``. +""" + +from __future__ import annotations + +import dataclasses +from collections.abc import Iterable, 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 control_edges_from_calls( + source_workflow: str, + calls: Iterable[tuple[str, bool | None, str]], +) -> list[ControlEdge]: + """Assemble deduplicated control edges from raw invocation primitives. + + 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: + 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 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 target, wait, via_task_key in calls: + target_name = target or "" + if target_name and target_name == source_workflow: + continue + key = (source_workflow, target_name, via_task_key) + if key in seen: + continue + seen.add(key) + edges.append( + ControlEdge( + source_workflow=source_workflow, + target_workflow=target_name, + via_task_key=via_task_key, + wait_for_completion=wait, + resolved=bool(target_name), + ) + ) + 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. + + 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 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. + + 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: + producers: ``(task_key, written asset)`` pairs, in first-seen order. + consumers: ``(task_key, read asset)`` pairs, in first-seen order. + + Returns: + 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. + """ + 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 producer_list: + for consumer_key, consumer_asset in consumer_list: + 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_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. + + 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/discovery.py b/src/flowx/models/discovery.py new file mode 100644 index 0000000..2a4a354 --- /dev/null +++ b/src/flowx/models/discovery.py @@ -0,0 +1,309 @@ +"""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 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 + +from dataclasses import dataclass, field +from typing import Any + +from flowx.models.ir import DataAsset, Lineage + +# --------------------------------------------------------------------------- # +# 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 kept in the SOURCE's own shape. + + Deliberately source-faithful, not Databricks-normalised. ``kind`` is 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. + + Attributes: + 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. + 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) + + +@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, 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 + 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`` 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) 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. + """ + + 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, 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 + 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) + 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) + 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. + 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`. + 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). + 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 + 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 + 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/src/flowx/models/ir.py b/src/flowx/models/ir.py index 600e8c7..b6df8c6 100644 --- a/src/flowx/models/ir.py +++ b/src/flowx/models/ir.py @@ -55,6 +55,147 @@ class Dependency: outcome: str | None = None +@dataclass(slots=True, kw_only=True) +class DataAsset: + """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, 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, 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, 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 (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. + """ + + 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 +235,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 +718,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 +776,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 +794,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_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 new file mode 100644 index 0000000..39b1eac --- /dev/null +++ b/tests/unit/test_discovery_model.py @@ -0,0 +1,325 @@ +"""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 _node_from_dict, _node_to_dict, 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 ControlEdge, DataAsset, DataEdge, Lineage + + +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_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( + 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", + expression={"frequency": "Day", "interval": 1}, + timezone="UTC", + extensions={"runtime_state": "Started"}, + ), + 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 + # 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.timezone == "UTC" + assert reloaded.schedule.extensions == {"runtime_state": "Started"} + + +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 == {} + + +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 + + +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 == [] 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": []}