Shared source-neutral discovery AST (#61)

src/flowx/bundler/dab_writer.py

@@ -30,16 +30,21 @@
 from flowx.models.ir import (
     Activity,
     AppendVariableActivity,
+    ControlEdge,
     CopyActivity,
+    DataAsset,
+    DataEdge,
     DbtFactoryActivity,
     DeleteActivity,
     Dependency,
     ExecutePipelineActivity,
     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

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))