From 97daac7f55b3045838b86f76a021c0308ef7e22b Mon Sep 17 00:00:00 2001 From: peterpark-db <127158477+peterpark-db@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:48:04 -0700 Subject: [PATCH 1/3] Map Airflow discovery to shared source graphs --- README.md | 3 +- docs/content/docs/guide.mdx | 1 + src/flowx/discovery_inventory.py | 100 +++ src/flowx/sources/airflow/discover.py | 93 ++- .../sources/airflow/discovery_mapping.py | 680 ++++++++++++++++++ src/flowx/sources/airflow/loader/__init__.py | 4 + src/flowx/sources/airflow/loader/api.py | 143 +++- src/flowx/sources/airflow/loader/captures.py | 1 + src/flowx/sources/airflow/loader/graph.py | 28 + src/flowx/sources/airflow/loader/lowering.py | 36 +- src/flowx/sources/airflow/loader/visitor.py | 5 + tests/unit/test_airflow_adapter_reporting.py | 10 + tests/unit/test_airflow_discovery_mapping.py | 261 +++++++ tests/unit/test_discovery_inventory.py | 85 +++ 14 files changed, 1384 insertions(+), 66 deletions(-) create mode 100644 src/flowx/discovery_inventory.py create mode 100644 src/flowx/sources/airflow/discovery_mapping.py create mode 100644 tests/unit/test_airflow_discovery_mapping.py create mode 100644 tests/unit/test_discovery_inventory.py diff --git a/README.md b/README.md index ea3e7e8..02ad813 100644 --- a/README.md +++ b/README.md @@ -228,7 +228,7 @@ to the supported static subset; flowx never imports or executes DAG modules. ## How It Works ### Phase 1: Discover -Parses the source into typed nodes and classifies each activity/operator as deterministic, agentic, or unsupported — ADF JSON from Unity Catalog volumes (or a `/Workspace` Git folder, normalizing ARM template format), or Airflow DAG `.py` modules read statically with `ast`. Airflow inventory includes audited/deterministic/agentic/failed/excluded counts, reconciliation status, stable finding fingerprints, translation-path coverage, and deterministic coverage. Produces `metadata/inventory.json` and a per-pipeline complexity report at `metadata/profile_report.csv`. +Parses the source into typed nodes and classifies each activity/operator as deterministic, agentic, or unsupported — ADF JSON from Unity Catalog volumes (or a `/Workspace` Git folder, normalizing ARM template format), or Airflow DAG `.py` modules read statically with `ast`. Airflow discovery also persists the source-faithful shared graph at `metadata/source_graphs.json`, including source identities, dependencies, policies, groups, mappings, gaps, and lineage before target lowering. Airflow inventory includes audited/deterministic/agentic/failed/excluded counts, reconciliation status, stable finding fingerprints, translation-path coverage, and deterministic coverage. Produces `metadata/inventory.json` and a per-pipeline complexity report at `metadata/profile_report.csv`. ### Phase 2: Convert Applies deterministic translators (ADF activity registry / Airflow operator mapping), resolves dependencies, and records unresolved gaps. ADF supports its guided agentic translation workflow. Airflow supports a fingerprint-bound, explicitly reviewed leaf-gap workflow whose constrained provider output is replayed against an immutable deterministic baseline before packaging. Produces the shared Pipeline IR consumed unchanged by the package phase. @@ -258,6 +258,7 @@ flowx_output/ metadata/ inventory.json # discover: activity inventory profile_report.csv # discover: per-pipeline complexity report + source_graphs.json # Airflow discover: source-faithful shared graphs and lineage .arm.json # discover: verbatim original ADF/ARM source configuration.json # modify: collected configuration answers .work/ # transient intermediates (translation report, IR, gaps.json); pruned by package diff --git a/docs/content/docs/guide.mdx b/docs/content/docs/guide.mdx index 821e9da..736d817 100644 --- a/docs/content/docs/guide.mdx +++ b/docs/content/docs/guide.mdx @@ -79,6 +79,7 @@ flowx_output/ ├── metadata/ │ ├── inventory.json # discover: activity inventory │ ├── profile_report.csv # profile: per-pipeline complexity report +│ ├── source_graphs.json # Airflow discover: source-faithful shared graphs │ ├── .arm.json # discover: verbatim original ADF/ARM pipeline source │ └── configuration.json # modify: the collected configuration answers └── .work/ # transient intermediates (translation report, IR, gaps.json); pruned by prepare diff --git a/src/flowx/discovery_inventory.py b/src/flowx/discovery_inventory.py new file mode 100644 index 0000000..b1efc46 --- /dev/null +++ b/src/flowx/discovery_inventory.py @@ -0,0 +1,100 @@ +"""Source-neutral projection of shared discovery graphs to ``inventory.json``.""" + +from __future__ import annotations + +from typing import Any + +from flowx.ir_serde import lineage_to_dict +from flowx.models.discovery import ContainerNode, SourceGraph, SourceNode + +STRATEGY_PROPERTY = "strategy" +INVENTORY_VISIBLE_PROPERTY = "inventory_visible" + +_DETERMINISTIC = "deterministic" +_AGENTIC = "agentic" + + +def build_source_inventory( + graphs: list[SourceGraph], + *, + source: str, + source_dir: str, + include_empty_pipelines: bool = True, +) -> dict[str, Any]: + """Projects source graphs into the common discovery inventory shape.""" + pipeline_entries: list[dict[str, Any]] = [] + deterministic = 0 + agentic = 0 + unsupported = 0 + + for graph in graphs: + flattened = [ + node for node in _flatten_nodes(graph.tasks) if node.properties.get(INVENTORY_VISIBLE_PROPERTY, True) + ] + for node in flattened: + strategy = node.properties.get(STRATEGY_PROPERTY) + if strategy == _DETERMINISTIC: + deterministic += 1 + elif strategy == _AGENTIC: + agentic += 1 + else: + unsupported += 1 + if flattened or include_empty_pipelines: + entry: dict[str, Any] = { + "name": graph.name, + "activities": [_activity_entry(node) for node in flattened], + } + if graph.lineage is not None: + entry["lineage"] = lineage_to_dict(graph.lineage) + pipeline_entries.append(entry) + + total = deterministic + agentic + unsupported + coverage_pct = round((deterministic + agentic) / total * 100, 1) if total else 0.0 + return { + "source": source, + "source_dir": source_dir, + "pipelines": pipeline_entries, + "summary": { + "pipeline_count": len(graphs), + "activity_count": total, + "deterministic_count": deterministic, + "agentic_count": agentic, + "unsupported_count": unsupported, + "coverage_pct": coverage_pct, + }, + } + + +def _flatten_nodes(nodes: list[SourceNode]) -> list[SourceNode]: + """Flattens container branches in source order.""" + flattened: list[SourceNode] = [] + for node in nodes: + flattened.append(node) + if isinstance(node, ContainerNode): + for children in node.branches.values(): + flattened.extend(_flatten_nodes(children)) + return flattened + + +def _activity_entry(node: SourceNode) -> dict[str, Any]: + """Builds one additive per-activity inventory entry.""" + entry: dict[str, Any] = { + "name": node.name if node.name is not None else node.task_key, + "type": node.native_type, + "strategy": node.properties.get(STRATEGY_PROPERTY), + } + upstream_names = [dependency.upstream for dependency in node.dependencies] + if upstream_names: + entry["depends_on"] = upstream_names + entry["original_type"] = node.native_type + entry["dependencies"] = [ + { + "upstream": dependency.upstream, + "conditions": list(dependency.conditions), + "resolved": dependency.resolved, + } + for dependency in node.dependencies + ] + if node.raw is not None: + entry["raw"] = node.raw + return entry diff --git a/src/flowx/sources/airflow/discover.py b/src/flowx/sources/airflow/discover.py index e64f52f..b325581 100644 --- a/src/flowx/sources/airflow/discover.py +++ b/src/flowx/sources/airflow/discover.py @@ -1,9 +1,9 @@ """Airflow discover phase: parse DAGs into a classified inventory. -Mirrors the ADF discover contract: writes ``metadata/inventory.json`` and -``metadata/profile_report.csv`` under the shared output dir. Independently audited -task candidates drive deterministic, agentic, failed, and excluded counts; emitted -IR tasks remain available for the per-task inventory. +Mirrors the ADF discover contract: writes ``metadata/inventory.json``, +``metadata/source_graphs.json``, and ``metadata/profile_report.csv`` under the shared +output dir. Shared source graphs drive per-task inventory while independently audited +task candidates drive deterministic, agentic, failed, and excluded counts. Exposes ``main(argv)`` so the adapter runs it in-process, like the ADF loader. """ @@ -17,9 +17,13 @@ from typing import Any from flowx.adapter.predicates import walk_activities +from flowx.discovery_inventory import build_source_inventory +from flowx.discovery_lineage import walk_nodes +from flowx.discovery_serde import source_graph_to_dict +from flowx.models.discovery import SourceGraph from flowx.models.ir import NotebookActivity, Pipeline, PlaceholderActivity from flowx.sources.adf.loader import clear_stale_outputs -from flowx.sources.airflow.loader import load_pipelines +from flowx.sources.airflow.loader import load_discovery_results logger = logging.getLogger(__name__) @@ -42,12 +46,34 @@ def _classify(pipeline: Pipeline) -> list[dict[str, str]]: return items -def build_inventory_dict(pipelines: list[Pipeline], source_dir: str) -> dict[str, Any]: +def build_inventory_dict( + pipelines: list[Pipeline], + source_dir: str, + *, + graphs: list[SourceGraph] | None = None, +) -> dict[str, Any]: """Builds the inventory.json payload matching the ADF discover shape.""" + if graphs is not None and len(graphs) != len(pipelines): + raise ValueError("Airflow inventory requires exactly one source graph per pipeline") + base: dict[str, Any] = ( + build_source_inventory(graphs, source="airflow", source_dir=source_dir) + if graphs is not None + else {"source": "airflow", "source_dir": source_dir, "pipelines": [], "summary": {}} + ) + base_entries = list(base["pipelines"]) pipeline_entries: list[dict[str, Any]] = [] audited = deterministic = agentic = failed = excluded = 0 for pipeline in pipelines: - items = _classify(pipeline) + base_entry = base_entries[len(pipeline_entries)] if len(base_entries) > len(pipeline_entries) else None + items = list(base_entry["activities"]) if base_entry is not None else _classify(pipeline) + if graphs is not None and base_entry is not None: + graph_nodes = [ + node + for node in walk_nodes(graphs[len(pipeline_entries)].tasks) + if node.properties.get("inventory_visible", True) + ] + for item, node in zip(items, graph_nodes, strict=True): + item["task_key"] = node.task_key pipeline_audited = int(pipeline.audit.get("audited_activity_count", len(items))) pipeline_deterministic = int( pipeline.audit.get("deterministic_count", sum(1 for item in items if item["strategy"] == "deterministic")) @@ -70,7 +96,8 @@ def build_inventory_dict(pipelines: list[Pipeline], source_dir: str) -> dict[str agentic += pipeline_agentic failed += pipeline_failed excluded += pipeline_excluded - pipeline_entries.append( + entry = dict(base_entry or {}) + entry.update( { "name": pipeline.name, "activities": items, @@ -87,6 +114,7 @@ def build_inventory_dict(pipelines: list[Pipeline], source_dir: str) -> dict[str "transformations": pipeline.audit.get("transformations", []), } ) + pipeline_entries.append(entry) coverage = round(100.0 * (deterministic + agentic) / audited, 1) if audited else 0.0 deterministic_coverage = round(100.0 * deterministic / audited, 1) if audited else 0.0 reconciliation_status = ( @@ -98,24 +126,25 @@ def build_inventory_dict(pipelines: list[Pipeline], source_dir: str) -> dict[str if pipelines and all(pipeline.migration_status == "excluded" for pipeline in pipelines) else "verified" ) - return { - "source": "airflow", - "source_dir": source_dir, - "pipelines": pipeline_entries, - "summary": { - "pipeline_count": len(pipelines), - "activity_count": audited, - "audited_activity_count": audited, - "deterministic_count": deterministic, - "agentic_count": agentic, - "unsupported_count": 0, - "failed_count": failed, - "excluded_count": excluded, - "coverage_pct": coverage, - "deterministic_coverage_pct": deterministic_coverage, - "reconciliation_status": reconciliation_status, - }, - } + base.update( + { + "pipelines": pipeline_entries, + "summary": { + "pipeline_count": len(pipelines), + "activity_count": audited, + "audited_activity_count": audited, + "deterministic_count": deterministic, + "agentic_count": agentic, + "unsupported_count": 0, + "failed_count": failed, + "excluded_count": excluded, + "coverage_pct": coverage, + "deterministic_coverage_pct": deterministic_coverage, + "reconciliation_status": reconciliation_status, + }, + } + ) + return base # Full profile column set the shared reporting.coverage / dashboard consume. Airflow has no @@ -187,7 +216,9 @@ def main(argv: list[str] | None = None) -> int: logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") - pipelines = load_pipelines(args.source_dir, pipeline=args.pipeline, exclude_dags=set(args.exclude_dag)) + results = load_discovery_results(args.source_dir, pipeline=args.pipeline, exclude_dags=set(args.exclude_dag)) + pipelines = [result.pipeline for result in results] + graphs = [result.graph for result in results] if not pipelines: logger.error("No Airflow DAGs found under %s (or none matched --pipeline).", args.source_dir) return 1 @@ -198,8 +229,14 @@ def main(argv: list[str] | None = None) -> int: metadata_dir = output_dir / "metadata" metadata_dir.mkdir(parents=True, exist_ok=True) - inventory = build_inventory_dict(pipelines, str(args.source_dir)) + inventory = build_inventory_dict(pipelines, str(args.source_dir), graphs=graphs) (metadata_dir / "inventory.json").write_text(json.dumps(inventory, indent=2), encoding="utf-8") + source_graphs = { + "contract_version": "1", + "source": "airflow", + "graphs": [source_graph_to_dict(graph) for graph in graphs], + } + (metadata_dir / "source_graphs.json").write_text(json.dumps(source_graphs, indent=2), encoding="utf-8") _write_profile_csv(pipelines, metadata_dir / "profile_report.csv") summary = inventory["summary"] diff --git a/src/flowx/sources/airflow/discovery_mapping.py b/src/flowx/sources/airflow/discovery_mapping.py new file mode 100644 index 0000000..5716e8a --- /dev/null +++ b/src/flowx/sources/airflow/discovery_mapping.py @@ -0,0 +1,680 @@ +"""Map Airflow source captures onto the shared discovery graph.""" + +from __future__ import annotations + +import ast +import json +from collections.abc import Iterable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from flowx.discovery_lineage import INVOKES_WAIT_PROPERTY, INVOKES_WORKFLOW_PROPERTY, walk_nodes, with_graph_lineage +from flowx.models.discovery import ( + CONCEPT_BRANCH, + CONCEPT_GAP, + CONCEPT_GROUP, + CONCEPT_LOOP, + CONCEPT_NOTEBOOK, + CONCEPT_QUERY, + CONCEPT_RUN_WORKFLOW, + CONCEPT_SCRIPT, + CONCEPT_WAIT, + SOURCE_AIRFLOW, + ContainerNode, + GapNode, + ParameterSpec, + PolicySpec, + ScheduleSpec, + SourceDependency, + SourceGraph, + SourceNode, +) +from flowx.models.ir import Activity, DataAsset, ForEachActivity, Pipeline, PlaceholderActivity +from flowx.sources.airflow import operators as ops +from flowx.sources.airflow import templating +from flowx.sources.airflow.audit import SourceAudit +from flowx.sources.airflow.loader.captures import DagDeclaration, SourceSpan +from flowx.sources.airflow.loader.graph import _allocate_task_keys, _expand_group_edges +from flowx.sources.airflow.loader.policy import _job_timeout_seconds +from flowx.sources.airflow.loader.visitor import _DagVisitor + +_QUERY_OPERATORS = frozenset( + { + "DatabricksSqlOperator", + "DatabricksSQLStatementsOperator", + "SQLExecuteQueryOperator", + "PostgresOperator", + "MySqlOperator", + "HiveOperator", + "DatabricksCopyIntoOperator", + } +) +_SCRIPT_OPERATORS = frozenset( + { + "BashOperator", + "SSHOperator", + "PythonOperator", + "PythonVirtualenvOperator", + "ExternalPythonOperator", + "SparkSubmitOperator", + } +) +_BRANCH_OPERATORS = frozenset({"BranchPythonOperator", "ShortCircuitOperator"}) + + +@dataclass(slots=True, kw_only=True) +class AirflowDiscoveryResult: + """One Airflow DAG represented as both current IR and source-faithful discovery graph.""" + + pipeline: Pipeline + graph: SourceGraph + + +def sync_graph_translation_metadata(graph: SourceGraph, pipeline: Pipeline) -> None: + """Refreshes target classifications after exclusions or cross-DAG rewrites.""" + activities = _activity_index(pipeline.tasks) + for node in walk_nodes(graph.tasks): + if node.properties.get("structural_only"): + continue + if node.task_key in activities: + node.properties["strategy"] = _strategy_for(node.task_key, activities, pipeline.reconciliation_status) + graph.properties.update( + { + "reconciliation_status": pipeline.reconciliation_status, + "migration_status": pipeline.migration_status, + "findings": list(pipeline.not_translatable), + "transformations": list(pipeline.audit.get("transformations", [])), + } + ) + + +def build_airflow_source_graph( + *, + dag_path: Path, + source_file: str, + source: str, + declaration: DagDeclaration, + visitor: _DagVisitor, + audit: SourceAudit, + pipeline: Pipeline, +) -> SourceGraph: + """Projects one captured DAG onto ``SourceGraph`` without reparsing it.""" + task_keys = _allocate_task_keys( + visitor.operators, + {variable: task.task_id for variable, task in visitor.taskflow_tasks.items()}, + {variable: task_id for variable, (task_id, _, _) in visitor.taskgroup_calls.items()}, + visitor.groups, + ) + expanded_edges = _expand_group_edges(visitor.edges, visitor.groups, visitor.group_vars) + upstreams: dict[str, list[str]] = {capture_id: [] for capture_id in task_keys} + for upstream, downstream in expanded_edges: + if upstream in task_keys and downstream in upstreams: + upstreams[downstream].append(upstream) + + activities = _activity_index(pipeline.tasks) + nodes_by_capture: dict[str, SourceNode] = {} + for capture_id, task_key in task_keys.items(): + nodes_by_capture[capture_id] = _captured_node( + capture_id=capture_id, + task_key=task_key, + upstreams=upstreams.get(capture_id, []), + task_keys=task_keys, + visitor=visitor, + source=source, + activities=activities, + ) + + tasks = _nest_task_groups(nodes_by_capture, visitor, source) + tasks.extend(_gap_nodes(visitor, source)) + declaration_raw = { + "capture_id": declaration.capture_id, + "kind": declaration.kind, + "source_file": source_file, + "source_span": _span_dict(declaration.span), + "source": ast.get_source_segment(source, declaration.node) or ast.unparse(declaration.node), + "dag_arguments": {name: _expression_payload(value, source) for name, value in visitor.dag_kwargs.items()}, + } + graph = SourceGraph( + name=pipeline.name, + source=SOURCE_AIRFLOW, + description=visitor.dag_description, + parameters={name: ParameterSpec(default=value) for name, value in visitor.dag_params.items()}, + schedule=_source_schedule(visitor, source), + default_policy=_policy(visitor.default_args), + run_timeout_seconds=_job_timeout_seconds(visitor), + tags=list(visitor.dag_user_tags), + tasks=tasks, + properties={ + "reconciliation_status": pipeline.reconciliation_status, + "migration_status": pipeline.migration_status, + "findings": list(pipeline.not_translatable), + "transformations": list(pipeline.audit.get("transformations", [])), + }, + extensions={ + "source_file": source_file, + "source_path": dag_path.as_posix(), + "airflow_generation": visitor.airflow_generation, + "declaration_capture_id": declaration.capture_id, + "edge_captures": [ + { + "upstream_capture_id": edge.upstream_id, + "downstream_capture_id": edge.downstream_id, + "source_span": _span_dict(edge.span), + } + for edge in visitor.edge_captures + ], + "audit": { + "tasks": [_audit_candidate_payload(candidate) for candidate in audit.tasks], + "edges": [_audit_candidate_payload(candidate) for candidate in audit.edges], + "settings": [_audit_candidate_payload(candidate) for candidate in audit.settings], + "unresolved": [_audit_candidate_payload(candidate) for candidate in audit.unresolved], + }, + }, + raw=declaration_raw, + ) + return with_graph_lineage(graph) + + +def failed_declaration_source_graph( + *, + dag_path: Path, + source_file: str, + source: str, + declaration: DagDeclaration, + pipeline: Pipeline, +) -> SourceGraph: + """Builds a reportable graph for a DAG declaration that static capture rejected.""" + raw_source = ast.get_source_segment(source, declaration.node) or ast.unparse(declaration.node) + gap = GapNode( + source_id=declaration.capture_id, + task_key=f"__flowx_gap_{declaration.span.line}_{declaration.span.column}", + source=SOURCE_AIRFLOW, + name=declaration.variable or dag_path.stem, + native_type="DagDeclaration", + reason=declaration.unsupported_reason, + properties={"strategy": "unsupported"}, + raw={"source": raw_source, "source_span": _span_dict(declaration.span)}, + ) + return SourceGraph( + name=pipeline.name, + source=SOURCE_AIRFLOW, + tasks=[gap], + properties={ + "reconciliation_status": pipeline.reconciliation_status, + "findings": list(pipeline.not_translatable), + }, + extensions={"source_file": source_file, "source_path": dag_path.as_posix()}, + raw={ + "capture_id": declaration.capture_id, + "kind": declaration.kind, + "source_file": source_file, + "source_span": _span_dict(declaration.span), + "source": raw_source, + }, + ) + + +def _captured_node( + *, + capture_id: str, + task_key: str, + upstreams: list[str], + task_keys: dict[str, str], + visitor: _DagVisitor, + source: str, + activities: dict[str, Activity], +) -> SourceNode: + dependencies = [ + SourceDependency(upstream=task_keys[upstream], resolved=True) + for upstream in dict.fromkeys(upstreams) + if upstream in task_keys + ] + source_node = visitor.capture_source_nodes[capture_id] + span = SourceSpan( + line=getattr(source_node, "lineno", 0), + column=getattr(source_node, "col_offset", 0), + end_line=getattr(source_node, "end_lineno", getattr(source_node, "lineno", 0)), + end_column=getattr(source_node, "end_col_offset", getattr(source_node, "col_offset", 0)), + ) + raw: dict[str, Any] = { + "source": ast.get_source_segment(source, source_node) or ast.unparse(source_node), + "source_span": _span_dict(span), + } + properties: dict[str, Any] = {"strategy": _strategy_for(task_key, activities, pipeline_status=None)} + run_condition: str | None = None + + if capture_id in visitor.operators: + task_id, operator, kwargs = visitor.operators[capture_id] + raw["arguments"] = {name: _expression_payload(value, source) for name, value in kwargs.items()} + raw["operator_fqn"] = visitor.task_captures[capture_id].operator_fqn + raw["argument_disposition"] = ops.argument_classification(operator, kwargs) + run_condition = ops.literal_str(kwargs.get("trigger_rule")) + policy = _policy(kwargs) + concept = _operator_concept(operator, capture_id in visitor.mapped) + reads, writes = _operator_assets(operator, kwargs) + if operator == "TriggerDagRunOperator": + target = ops.literal_str(kwargs.get("trigger_dag_id")) + if target is not None: + properties[INVOKES_WORKFLOW_PROPERTY] = target + properties[INVOKES_WAIT_PROPERTY] = bool(ops.literal_value(kwargs.get("wait_for_completion"))) + elif operator in {"DatabricksRunNowOperator", "DatabricksRunNowDeferrableOperator"}: + job_id = ops.literal_value(kwargs.get("job_id")) + if job_id is not None: + properties[INVOKES_WORKFLOW_PROPERTY] = f"databricks-job:{job_id}" + elif operator in {"ExternalTaskSensor", "ExternalTaskSensorAsync"}: + properties["external_workflow_wait"] = { + "dag_id": ops.literal_str(kwargs.get("external_dag_id")), + "task_id": ops.literal_str(kwargs.get("external_task_id")), + } + if capture_id in visitor.mapped: + properties["mapping"] = { + "expand_arguments": list(visitor.expand_kwargs.get(capture_id, [])), + "has_partial": capture_id in visitor.partial_mapped, + } + return ContainerNode( + source_id=capture_id, + task_key=task_key, + concept=CONCEPT_LOOP, + source=SOURCE_AIRFLOW, + name=task_id, + native_type=operator, + dependencies=dependencies, + run_condition=run_condition, + policy=policy, + data_reads=reads, + data_writes=writes, + properties=properties, + raw=raw, + branches={"body": []}, + ) + activity = activities.get(task_key) + if isinstance(activity, PlaceholderActivity) or concept == CONCEPT_GAP: + reason = activity.comment if isinstance(activity, PlaceholderActivity) else None + return GapNode( + source_id=capture_id, + task_key=task_key, + source=SOURCE_AIRFLOW, + name=task_id, + native_type=operator, + dependencies=dependencies, + run_condition=run_condition, + policy=policy, + data_reads=reads, + data_writes=writes, + properties=properties, + raw=raw, + reason=reason or f"Airflow operator {operator!r} has no deterministic mapping.", + ) + return SourceNode( + source_id=capture_id, + task_key=task_key, + concept=concept, + source=SOURCE_AIRFLOW, + name=task_id, + native_type=operator, + dependencies=dependencies, + run_condition=run_condition, + policy=policy, + data_reads=reads, + data_writes=writes, + properties=properties, + raw=raw, + ) + + if capture_id in visitor.taskflow_tasks: + task = visitor.taskflow_tasks[capture_id] + raw.update( + { + "decorator": task.decorator, + "callable": task.def_name, + "source_reference": task.source_reference, + "positional_arguments": dict(task.positional_values), + "keyword_arguments": dict(task.keyword_values), + "unresolved_arguments": list(task.unresolved_arguments), + } + ) + reads = [ + DataAsset(signature=f"xcom:{task_keys[upstream]}", asset_type="value") + for upstream in dict.fromkeys(upstreams) + if upstream in task_keys + ] + writes = [DataAsset(signature=f"xcom:{task_key}", asset_type="value")] + if capture_id in visitor.mapped: + properties["mapping"] = { + "expand_argument": task.expand_kwarg, + "expand_items_json": task.expand_items_json, + } + return ContainerNode( + source_id=capture_id, + task_key=task_key, + concept=CONCEPT_LOOP, + source=SOURCE_AIRFLOW, + name=task.task_id, + native_type=task.decorator, + dependencies=dependencies, + data_reads=reads, + data_writes=writes, + properties=properties, + raw=raw, + branches={"body": []}, + ) + activity = activities.get(task_key) + if isinstance(activity, PlaceholderActivity): + return GapNode( + source_id=capture_id, + task_key=task_key, + source=SOURCE_AIRFLOW, + name=task.task_id, + native_type=task.decorator, + dependencies=dependencies, + data_reads=reads, + data_writes=writes, + properties=properties, + raw=raw, + reason=activity.comment or "TaskFlow invocation requires manual migration.", + ) + return SourceNode( + source_id=capture_id, + task_key=task_key, + concept=CONCEPT_SCRIPT, + source=SOURCE_AIRFLOW, + name=task.task_id, + native_type=task.decorator, + dependencies=dependencies, + data_reads=reads, + data_writes=writes, + properties=properties, + raw=raw, + ) + + task_id, definition, mapped = visitor.taskgroup_calls[capture_id] + raw.update({"task_group_callable": definition, "mapped": mapped}) + return ContainerNode( + source_id=capture_id, + task_key=task_key, + concept=CONCEPT_GROUP, + source=SOURCE_AIRFLOW, + name=task_id, + native_type="task_group", + dependencies=dependencies, + properties=properties, + raw=raw, + branches={"group": []}, + ) + + +def _nest_task_groups( + nodes_by_capture: dict[str, SourceNode], + visitor: _DagVisitor, + source: str, +) -> list[SourceNode]: + roots: list[SourceNode] = [] + containers: dict[str, ContainerNode] = {} + + def container(path: str) -> ContainerNode: + existing = containers.get(path) + if existing is not None: + return existing + parent_path, _, leaf = path.rpartition("__") + group_node = visitor.group_source_nodes.get(path) + raw = None + if group_node is not None: + raw = { + "source": ast.get_source_segment(source, group_node) or ast.unparse(group_node), + "source_span": { + "line": getattr(group_node, "lineno", 0), + "column": getattr(group_node, "col_offset", 0), + "end_line": getattr(group_node, "end_lineno", 0), + "end_column": getattr(group_node, "end_col_offset", 0), + }, + } + created = ContainerNode( + source_id=f"task-group:{path}", + task_key=path, + concept=CONCEPT_GROUP, + source=SOURCE_AIRFLOW, + name=leaf, + native_type="TaskGroup", + properties={"inventory_visible": False, "structural_only": True}, + raw=raw, + branches={"group": []}, + ) + containers[path] = created + if parent_path: + container(parent_path).branches["group"].append(created) + else: + roots.append(created) + return created + + for capture_id, node in nodes_by_capture.items(): + group_path = visitor.groups.get(capture_id) + if group_path: + container(group_path).branches["group"].append(node) + else: + roots.append(node) + return roots + + +def _gap_nodes(visitor: _DagVisitor, source: str) -> list[GapNode]: + entries: list[tuple[str, ast.AST, str]] = [] + entries.extend( + ("unclaimed_task_call", node, "Airflow task call was not captured by the static subset.") + for node in visitor.unclaimed_task_calls + ) + entries.extend( + ("unclaimed_statement", node, "DAG-body statement was not claimed by the static subset.") + for node in visitor.unclaimed_statements + ) + entries.extend(("unresolved_construct", node, reason) for reason, node in visitor.unresolved_constructs) + gaps: list[GapNode] = [] + seen: set[tuple[str, int, int, int, int]] = set() + for code, node, reason in entries: + key = ( + code, + getattr(node, "lineno", 0), + getattr(node, "col_offset", 0), + getattr(node, "end_lineno", 0), + getattr(node, "end_col_offset", 0), + ) + if key in seen: + continue + seen.add(key) + source_id = f"{code}:{key[1]}:{key[2]}:{len(gaps) + 1}" + gaps.append( + GapNode( + source_id=source_id, + task_key=f"__flowx_gap_{key[1]}_{key[2]}_{len(gaps) + 1}", + source=SOURCE_AIRFLOW, + name=code, + native_type=type(node).__name__, + reason=reason, + properties={"strategy": "unsupported"}, + raw={ + "source": ast.get_source_segment(source, node) or ast.unparse(node), + "source_span": { + "line": key[1], + "column": key[2], + "end_line": key[3], + "end_column": key[4], + }, + }, + ) + ) + return gaps + + +def _operator_concept(operator: str, mapped: bool) -> str: + if mapped: + return CONCEPT_LOOP + if operator in _QUERY_OPERATORS: + return CONCEPT_QUERY + if operator in _SCRIPT_OPERATORS: + return CONCEPT_SCRIPT + if operator in _BRANCH_OPERATORS: + return CONCEPT_BRANCH + if operator == "TriggerDagRunOperator" or operator.startswith("DatabricksRunNow"): + return CONCEPT_RUN_WORKFLOW + if operator.endswith("Sensor") or operator.endswith("SensorAsync"): + return CONCEPT_WAIT + if operator in ops.OPERATOR_REGISTRY: + return CONCEPT_NOTEBOOK + return CONCEPT_GAP + + +def _operator_assets(operator: str, kwargs: dict[str, ast.expr]) -> tuple[list[DataAsset], list[DataAsset]]: + reads = _declared_assets(kwargs.get("inlets"), direction="read") + writes = _declared_assets(kwargs.get("outlets"), direction="write") + if operator in ops.FILE_SENSORS: + path = ops.file_sensor_path(kwargs) + if path: + reads.append(DataAsset(signature=path, identity=path, asset_type="file")) + if operator in ops.TABLE_SENSORS: + table = ops.literal_str(kwargs.get("table_name")) + if table: + reads.append(DataAsset(signature=table, identity=table, asset_type="table")) + if operator == "DatabricksCopyIntoOperator": + location = ops.literal_str(kwargs.get("file_location")) + table = ops.literal_str(kwargs.get("table_name")) + if location: + reads.append(DataAsset(signature=location, identity=location, asset_type="file")) + if table: + writes.append(DataAsset(signature=table, identity=table, asset_type="table")) + sql = ops.literal_str(kwargs.get("sql")) or ops.literal_str(kwargs.get("hql")) + if sql: + reads.append(DataAsset(signature=sql.strip(), asset_type="query", properties={"role": "source_sql"})) + return _deduplicate_assets(reads), _deduplicate_assets(writes) + + +def _declared_assets(node: ast.expr | None, *, direction: str) -> list[DataAsset]: + if node is None: + return [] + candidates = list(node.elts) if isinstance(node, (ast.List, ast.Tuple, ast.Set)) else [node] + assets: list[DataAsset] = [] + for candidate in candidates: + value: str | None = None + if isinstance(candidate, ast.Call) and candidate.args: + value = ops.literal_str(candidate.args[0]) + else: + value = ops.literal_str(candidate) + if value is not None: + assets.append( + DataAsset( + signature=value, + identity=value if "://" in value else None, + asset_type="logical", + properties={"airflow_direction": direction}, + ) + ) + return assets + + +def _deduplicate_assets(assets: Iterable[DataAsset]) -> list[DataAsset]: + result: list[DataAsset] = [] + seen: set[tuple[str, str | None, str | None]] = set() + for asset in assets: + key = (asset.signature, asset.identity, asset.asset_type) + if key not in seen: + seen.add(key) + result.append(asset) + return result + + +def _activity_index(tasks: list[Activity]) -> dict[str, Activity]: + result: dict[str, Activity] = {} + stack = list(tasks) + while stack: + activity = stack.pop(0) + result.setdefault(activity.task_key, activity) + if isinstance(activity, ForEachActivity): + stack[0:0] = activity.inner_activities + return result + + +def _strategy_for(task_key: str, activities: dict[str, Activity], pipeline_status: str | None) -> str: + activity = activities.get(task_key) + if isinstance(activity, PlaceholderActivity): + return "agentic" + if isinstance(activity, ForEachActivity): + nested = _activity_index(activity.inner_activities) + if any(isinstance(item, PlaceholderActivity) for item in nested.values()): + return "agentic" + if activity is not None: + return "deterministic" + return "unsupported" if pipeline_status == "failed" else "deterministic" + + +def _policy(arguments: dict[str, ast.expr]) -> PolicySpec | None: + retries = ops.literal_value(arguments.get("retries")) + max_retries = retries if isinstance(retries, int) and not isinstance(retries, bool) else None + retry_interval = templating.timedelta_seconds(arguments.get("retry_delay")) + timeout = templating.timedelta_seconds(arguments.get("execution_timeout")) + extensions: dict[str, Any] = {} + for name in ("depends_on_past", "email", "email_on_failure", "email_on_retry"): + if name in arguments: + extensions[name] = _json_safe(ops.literal_value(arguments[name]), fallback=ast.unparse(arguments[name])) + if max_retries is None and retry_interval is None and timeout is None and not extensions: + return None + return PolicySpec( + timeout_seconds=timeout, + max_retries=max_retries, + retry_interval_seconds=retry_interval, + extensions=extensions, + ) + + +def _source_schedule(visitor: _DagVisitor, source: str) -> ScheduleSpec | None: + node = visitor.schedule_node + if node is None or (isinstance(node, ast.Constant) and node.value is None): + return None + literal = ops.literal_value(node) + source_expression = ast.get_source_segment(source, node) or ast.unparse(node) + expression = source_expression if literal is None else _json_safe(literal, fallback=source_expression) + kind = "schedule" if isinstance(literal, str) else "interval" + is_asset_expression = isinstance(node, ast.Call) and "Asset" in ast.unparse(node) + if isinstance(node, (ast.List, ast.Tuple, ast.Set)) or is_asset_expression: + kind = "asset" + return ScheduleSpec( + kind=kind, + expression=expression, + timezone=visitor.timezone, + extensions={"source_expression": source_expression}, + ) + + +def _expression_payload(node: ast.expr, source: str) -> dict[str, Any]: + value = ops.literal_value(node) + return { + "source": ast.get_source_segment(source, node) or ast.unparse(node), + "value": _json_safe(value, fallback=None), + } + + +def _json_safe(value: Any, *, fallback: Any) -> Any: + try: + json.dumps(value) + except (TypeError, ValueError): + return fallback + return value + + +def _span_dict(span: SourceSpan) -> dict[str, int]: + return { + "line": span.line, + "column": span.column, + "end_line": span.end_line, + "end_column": span.end_column, + } + + +def _audit_candidate_payload(candidate: Any) -> dict[str, Any]: + return { + "kind": candidate.kind, + "code": candidate.code, + "line": candidate.line, + "column": candidate.column, + "end_line": candidate.end_line, + "end_column": candidate.end_column, + "occurrence": candidate.occurrence, + "details": dict(candidate.details), + } diff --git a/src/flowx/sources/airflow/loader/__init__.py b/src/flowx/sources/airflow/loader/__init__.py index 3e72ad3..1c69e42 100644 --- a/src/flowx/sources/airflow/loader/__init__.py +++ b/src/flowx/sources/airflow/loader/__init__.py @@ -6,14 +6,18 @@ detect_hosts, discover_dags, load_airflow_dag, + load_airflow_dag_results, load_airflow_dags, + load_discovery_results, load_pipelines, ) __all__ = [ "detect_hosts", "discover_dags", + "load_discovery_results", "load_airflow_dag", + "load_airflow_dag_results", "load_airflow_dags", "load_pipelines", ] diff --git a/src/flowx/sources/airflow/loader/api.py b/src/flowx/sources/airflow/loader/api.py index 21cb889..7041c98 100644 --- a/src/flowx/sources/airflow/loader/api.py +++ b/src/flowx/sources/airflow/loader/api.py @@ -4,14 +4,22 @@ import ast import re +from dataclasses import dataclass from pathlib import Path +from flowx.models.discovery import SourceGraph from flowx.models.ir import ( Pipeline, PlaceholderActivity, RunJobActivity, ) from flowx.sources.airflow import audit as source_audit +from flowx.sources.airflow.discovery_mapping import ( + AirflowDiscoveryResult, + build_airflow_source_graph, + failed_declaration_source_graph, + sync_graph_translation_metadata, +) from flowx.sources.airflow.loader.ast_utils import _expand_top_level_loops from flowx.sources.airflow.loader.dag_discovery import ( _failed_dag_declaration_pipeline, @@ -19,11 +27,18 @@ _top_level_dag_declarations, ) from flowx.sources.airflow.loader.lowering import _load_airflow_module +from flowx.sources.airflow.loader.visitor import _DagVisitor from flowx.utils import normalize_task_key _HOST_PATTERN = re.compile(r"https://([A-Za-z0-9._-]*(?:azuredatabricks\.net|databricks\.com|cloud\.databricks\.com))") +@dataclass(slots=True, kw_only=True) +class _LoadedDag: + pipeline: Pipeline + graph: SourceGraph | None + + def load_airflow_dag(dag_path: Path, *, dbt_mode: str = "static") -> Pipeline: """Parses the first Airflow DAG in a file into a flowx Pipeline IR.""" pipelines = load_airflow_dags(dag_path, dbt_mode=dbt_mode) @@ -39,31 +54,94 @@ def load_airflow_dags( source_file: str | None = None, ) -> list[Pipeline]: """Parses every independently declared Airflow DAG in a Python file.""" + return [ + result.pipeline + for result in _load_airflow_dag_results( + dag_path, + dbt_mode=dbt_mode, + source_file=source_file, + include_graph=False, + ) + ] + + +def load_airflow_dag_results( + dag_path: Path, + *, + dbt_mode: str = "static", + source_file: str | None = None, +) -> list[AirflowDiscoveryResult]: + """Parses every DAG into current IR and its shared source graph in one capture pass.""" + results = _load_airflow_dag_results( + dag_path, + dbt_mode=dbt_mode, + source_file=source_file, + include_graph=True, + ) + return [ + AirflowDiscoveryResult(pipeline=result.pipeline, graph=result.graph) + for result in results + if result.graph is not None + ] + + +def _load_airflow_dag_results( + dag_path: Path, + *, + dbt_mode: str, + source_file: str | None, + include_graph: bool, +) -> list[_LoadedDag]: source = Path(dag_path).read_text(encoding="utf-8") module = _expand_top_level_loops(ast.parse(source)) declarations = _top_level_dag_declarations(module) - pipelines: list[Pipeline] = [] + results: list[_LoadedDag] = [] + label = source_file or dag_path.name for declaration in declarations: if declaration.unsupported_reason is not None: - pipelines.append( - _failed_dag_declaration_pipeline( - dag_path, - declaration, - source_file=source_file or dag_path.name, + pipeline = _failed_dag_declaration_pipeline(dag_path, declaration, source_file=label) + graph = ( + failed_declaration_source_graph( + dag_path=dag_path, + source_file=label, + source=source, + declaration=declaration, + pipeline=pipeline, ) + if include_graph + else None ) + results.append(_LoadedDag(pipeline=pipeline, graph=graph)) continue - pipelines.append( - _load_airflow_module( - dag_path, - source, - _module_for_dag(module, declaration, declarations), - dbt_mode=dbt_mode, - target_dag_variable=declaration.target_dag_variable, - source_file=source_file or dag_path.name, + isolated = _module_for_dag(module, declaration, declarations) + audit = source_audit.audit_module(isolated, target_dag_variable=declaration.target_dag_variable) + visitor = _DagVisitor(isolated, target_dag_variable=declaration.target_dag_variable) + visitor.visit(isolated) + pipeline = _load_airflow_module( + dag_path, + source, + isolated, + dbt_mode=dbt_mode, + target_dag_variable=declaration.target_dag_variable, + source_file=label, + captured_audit=audit, + captured_visitor=visitor, + ) + graph = ( + build_airflow_source_graph( + dag_path=dag_path, + source_file=label, + source=source, + declaration=declaration, + visitor=visitor, + audit=audit, + pipeline=pipeline, ) + if include_graph + else None ) - return pipelines + results.append(_LoadedDag(pipeline=pipeline, graph=graph)) + return results def load_pipelines( @@ -97,6 +175,40 @@ def load_pipelines( if pipeline is not None: pipelines = [p for p in pipelines if p.name == pipeline] excluded = set(exclude_dags or ()) + _apply_exclusions(pipelines, excluded) + return pipelines + + +def load_discovery_results( + source_path: Path, + pipeline: str | None = None, + *, + dbt_mode: str = "static", + exclude_dags: set[str] | None = None, +) -> list[AirflowDiscoveryResult]: + """Loads Airflow DAGs into paired Pipeline IR and shared source graphs.""" + root = source_path if source_path.is_dir() else source_path.parent + results = [ + result + for dag_path in discover_dags(source_path) + for result in load_airflow_dag_results( + dag_path, + dbt_mode=dbt_mode, + source_file=source_audit.source_label(dag_path, root), + ) + ] + if pipeline is not None: + results = [result for result in results if result.pipeline.name == pipeline] + excluded = set(exclude_dags or ()) + pipelines = [result.pipeline for result in results] + _apply_exclusions(pipelines, excluded) + for result in results: + sync_graph_translation_metadata(result.graph, result.pipeline) + return results + + +def _apply_exclusions(pipelines: list[Pipeline], excluded: set[str]) -> None: + """Marks excluded DAGs and replaces included references to them.""" for loaded in pipelines: if loaded.name in excluded: loaded.migration_status = "excluded" @@ -111,7 +223,6 @@ def load_pipelines( ) if excluded: _replace_excluded_dag_references(pipelines, excluded) - return pipelines def _replace_excluded_dag_references(pipelines: list[Pipeline], excluded: set[str]) -> None: diff --git a/src/flowx/sources/airflow/loader/captures.py b/src/flowx/sources/airflow/loader/captures.py index f750987..af9d160 100644 --- a/src/flowx/sources/airflow/loader/captures.py +++ b/src/flowx/sources/airflow/loader/captures.py @@ -41,6 +41,7 @@ class TaskCapture: variable: str task_id: str operator: str + operator_fqn: str call: ast.Call span: SourceSpan diff --git a/src/flowx/sources/airflow/loader/graph.py b/src/flowx/sources/airflow/loader/graph.py index 10edb7e..a29d764 100644 --- a/src/flowx/sources/airflow/loader/graph.py +++ b/src/flowx/sources/airflow/loader/graph.py @@ -5,6 +5,34 @@ import ast from flowx.sources.airflow import operators as ops +from flowx.sources.airflow.loader.ast_utils import _sanitize_task_key + + +def _allocate_task_keys( + operators: dict[str, tuple[str, str, dict[str, ast.expr]]], + taskflow_task_ids: dict[str, str], + taskgroup_task_ids: dict[str, str], + groups: dict[str, str], +) -> dict[str, str]: + """Allocates stable, collision-free task keys in capture order.""" + task_ids = {variable: task_id for variable, (task_id, _, _) in operators.items()} + task_ids.update(taskflow_task_ids) + task_ids.update(taskgroup_task_ids) + + allocated: dict[str, str] = {} + used: set[str] = set() + for variable, task_id in task_ids.items(): + base = _sanitize_task_key(task_id) + if variable in groups: + base = f"{groups[variable]}__{base}" + candidate = base + suffix = 2 + while candidate in used: + candidate = f"{base}__{suffix}" + suffix += 1 + used.add(candidate) + allocated[variable] = candidate + return allocated def _expand_group_edges( diff --git a/src/flowx/sources/airflow/loader/lowering.py b/src/flowx/sources/airflow/loader/lowering.py index aad4a6d..1f26722 100644 --- a/src/flowx/sources/airflow/loader/lowering.py +++ b/src/flowx/sources/airflow/loader/lowering.py @@ -16,15 +16,17 @@ from flowx.sources.airflow import audit as source_audit from flowx.sources.airflow import operators as ops from flowx.sources.airflow import templating +from flowx.sources.airflow.audit import SourceAudit from flowx.sources.airflow.loader import reconcile from flowx.sources.airflow.loader.activity_templates import ( _convert_activity_templates, _declared_param_default, _unresolved_activity_templates, ) -from flowx.sources.airflow.loader.ast_utils import _sanitize_task_key, _span +from flowx.sources.airflow.loader.ast_utils import _span from flowx.sources.airflow.loader.dbt import _build_dbt_factory from flowx.sources.airflow.loader.graph import ( + _allocate_task_keys, _expand_group_edges, _rewire_dropped, _root_trigger_sensor, @@ -47,6 +49,8 @@ def _load_airflow_module( dbt_mode: str = "static", target_dag_variable: str | None = None, source_file: str | None = None, + captured_audit: SourceAudit | None = None, + captured_visitor: _DagVisitor | None = None, ) -> Pipeline: """Parses one isolated DAG declaration into a flowx Pipeline IR. @@ -63,33 +67,23 @@ def _load_airflow_module( sensors remain explicit placeholders; unmapped operators become a PlaceholderActivity. """ - audit = source_audit.audit_module(module, target_dag_variable=target_dag_variable) - visitor = _DagVisitor(module, target_dag_variable=target_dag_variable) - visitor.visit(module) + audit = captured_audit or source_audit.audit_module(module, target_dag_variable=target_dag_variable) + visitor = captured_visitor or _DagVisitor(module, target_dag_variable=target_dag_variable) + if captured_visitor is None: + visitor.visit(module) functions = visitor.functions() - # Prefix TaskGroup member keys with the group id (e.g. extract__run) so two tasks named - # `run` in different groups don't collide. - def _task_key(var: str, task_id: str) -> str: - key = _sanitize_task_key(task_id) - return f"{visitor.groups[var]}__{key}" if var in visitor.groups else key - # TaskFlow @task instances share the task table with classic operators (both are just tasks with # a task_key and dependency edges downstream). var_task_ids: dict[str, str] = {var: tid for var, (tid, _, _) in visitor.operators.items()} var_task_ids.update({var: tf.task_id for var, tf in visitor.taskflow_tasks.items()}) var_task_ids.update({var: task_id for var, (task_id, _, _) in visitor.taskgroup_calls.items()}) - var_to_task_key: dict[str, str] = {} - used_task_keys: set[str] = set() - for var, task_id in var_task_ids.items(): - base = _task_key(var, task_id) - candidate = base - suffix = 2 - while candidate in used_task_keys: - candidate = f"{base}__{suffix}" - suffix += 1 - used_task_keys.add(candidate) - var_to_task_key[var] = candidate + var_to_task_key = _allocate_task_keys( + visitor.operators, + {var: task.task_id for var, task in visitor.taskflow_tasks.items()}, + {var: task_id for var, (task_id, _, _) in visitor.taskgroup_calls.items()}, + visitor.groups, + ) # Expand group-level edges (`group_a >> group_b`, `task >> group`, ...) into edges between the # groups' boundary tasks: leaves of the upstream group -> roots of the downstream group, matching diff --git a/src/flowx/sources/airflow/loader/visitor.py b/src/flowx/sources/airflow/loader/visitor.py index 1f0828c..51c2772 100644 --- a/src/flowx/sources/airflow/loader/visitor.py +++ b/src/flowx/sources/airflow/loader/visitor.py @@ -12,6 +12,7 @@ _UNRESOLVED, _airflow_generation, _bind_constants, + _canonical_name, _construct_name, _import_aliases, _index_lexical_functions, @@ -96,6 +97,7 @@ def __init__(self, module: ast.Module, *, target_dag_variable: str | None = None # task variable name -> TaskGroup id prefix (for task-key namespacing) self.groups: dict[str, str] = {} self._group_stack: list[str] = [] + self.group_source_nodes: dict[str, ast.Call] = {} # `with TaskGroup(...) as tg:` binding -> the group's prefix, so a group-level edge # (tg >> other) can expand to edges between the groups' boundary tasks. self.group_vars: dict[str, str] = {} @@ -263,6 +265,7 @@ def _register_operator_call(self, node: ast.Call, var: str, *, binding: str | No call = direct or (mapped[0] if mapped is not None else None) if call is None: return False + operator_fqn = _canonical_name(call.func, self._aliases) construct = _construct_name(call.func, self._aliases) kwargs = {kw.arg: _bind_constants(kw.value, self._constants) for kw in call.keywords if kw.arg} dag_node = kwargs.get("dag") @@ -285,6 +288,7 @@ def _register_operator_call(self, node: ast.Call, var: str, *, binding: str | No variable=binding or var, task_id=task_id, operator=construct, + operator_fqn=operator_fqn, call=call, span=_span(node), ) @@ -609,6 +613,7 @@ def visit_With(self, node: ast.With) -> None: or "group" ) self._group_stack.append(_sanitize_task_key(group_id)) + self.group_source_nodes["__".join(self._group_stack)] = call pushed_group = True # Record the `as tg` binding (with the full nested prefix) so a group-level # edge on `tg` resolves to the group's member tasks. diff --git a/tests/unit/test_airflow_adapter_reporting.py b/tests/unit/test_airflow_adapter_reporting.py index 7502673..90b4243 100644 --- a/tests/unit/test_airflow_adapter_reporting.py +++ b/tests/unit/test_airflow_adapter_reporting.py @@ -2,12 +2,14 @@ from __future__ import annotations +import json import tempfile from pathlib import Path import pytest from flowx.adapter.session import MigrationInputSession +from flowx.discovery_serde import source_graph_from_dict from flowx.models.ir import NotebookActivity, Pipeline from flowx.reporting.coverage import COVERAGE_METRIC_COLUMNS, build_coverage_rows from flowx.sources.airflow.discover import _profile_row, build_inventory_dict @@ -67,6 +69,14 @@ def test_airflow_profile_csv_has_all_coverage_columns(): assert row["databricks_native_activities"] == 1 # the PythonOperator assert row["other_activities"] == 1 # the placeholder assert row["complexity_score"] == 4 # 1*1 + 1*3 + source_graphs = json.loads((out / "metadata" / "source_graphs.json").read_text(encoding="utf-8")) + assert source_graphs["contract_version"] == "1" + assert source_graphs["source"] == "airflow" + graph = source_graph_from_dict(source_graphs["graphs"][0]) + assert graph.name == "cov" + inventory = json.loads((out / "metadata" / "inventory.json").read_text(encoding="utf-8")) + assert [activity["task_key"] for activity in inventory["pipelines"][0]["activities"]] == ["a", "b"] + assert "lineage" in inventory["pipelines"][0] def test_airflow_inventory_persists_audit_status_counts_and_findings() -> None: diff --git a/tests/unit/test_airflow_discovery_mapping.py b/tests/unit/test_airflow_discovery_mapping.py new file mode 100644 index 0000000..5c70e7a --- /dev/null +++ b/tests/unit/test_airflow_discovery_mapping.py @@ -0,0 +1,261 @@ +from __future__ import annotations + +from pathlib import Path + +from flowx.discovery_inventory import build_source_inventory +from flowx.discovery_lineage import walk_nodes +from flowx.discovery_serde import source_graph_from_dict, source_graph_to_dict +from flowx.models.discovery import CONCEPT_GROUP, CONCEPT_LOOP, ContainerNode, GapNode +from flowx.sources.airflow.loader import load_airflow_dag_results, load_discovery_results + + +def _load(tmp_path: Path, source: str): + dag_path = tmp_path / "dag.py" + dag_path.write_text(source, encoding="utf-8") + results = load_airflow_dag_results(dag_path, source_file="dags/dag.py") + assert len(results) == 1 + return results[0] + + +def test_maps_assigned_dag_metadata_tasks_edges_and_collision_safe_keys(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from datetime import timedelta +from airflow import DAG +from airflow.operators.bash import BashOperator + +dag = DAG( + dag_id="source_graph", + schedule_interval="0 3 * * *", + default_args={"retries": 3, "retry_delay": timedelta(minutes=4)}, + params={"environment": "prod"}, + dagrun_timeout=timedelta(hours=2), + tags=["migration", "airflow"], + description="Source faithful", +) +first = BashOperator(task_id="load.data", bash_command="echo first", dag=dag) +second = BashOperator( + task_id="load_data", + bash_command="echo second", + trigger_rule="all_done", + retries=5, + dag=dag, +) +first >> second +""", + ) + + graph = result.graph + assert graph.name == "source_graph" + assert graph.description == "Source faithful" + assert graph.parameters["environment"].default == "prod" + assert graph.schedule is not None + assert graph.schedule.expression == "0 3 * * *" + assert graph.default_policy is not None + assert graph.default_policy.max_retries == 3 + assert graph.default_policy.retry_interval_seconds == 240 + assert graph.run_timeout_seconds == 7200 + assert graph.tags == ["migration", "airflow"] + assert graph.raw is not None + assert graph.raw["capture_id"].startswith("dag:") + assert graph.extensions["source_file"] == "dags/dag.py" + assert graph.extensions["edge_captures"][0]["upstream_capture_id"] == "first" + assert graph.extensions["edge_captures"][0]["downstream_capture_id"] == "second" + + nodes = {node.source_id: node for node in walk_nodes(graph.tasks) if not node.properties.get("structural_only")} + assert nodes["first"].task_key == "load_data" + assert nodes["second"].task_key == "load_data__2" + assert [dependency.upstream for dependency in nodes["second"].dependencies] == ["load_data"] + assert nodes["second"].run_condition == "all_done" + assert nodes["second"].policy is not None + assert nodes["second"].policy.max_retries == 5 + assert nodes["first"].raw is not None + assert nodes["first"].raw["operator_fqn"] == "airflow.operators.bash.BashOperator" + assert nodes["first"].raw["source_span"]["line"] > 0 + assert nodes["first"].raw["arguments"]["bash_command"]["value"] == "echo first" + + assert source_graph_from_dict(source_graph_to_dict(graph)) == graph + + +def test_maps_task_groups_and_dynamic_mapping_as_containers(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow import DAG +from airflow.operators.bash import BashOperator +from airflow.utils.task_group import TaskGroup + +with DAG(dag_id="groups", schedule=None) as dag: + start = BashOperator(task_id="start", bash_command="echo start") + with TaskGroup(group_id="processing") as processing: + mapped = BashOperator.partial(task_id="work", bash_command="echo {{ params.item }}").expand( + params=[{"item": "a"}, {"item": "b"}] + ) + end = BashOperator(task_id="end", bash_command="echo end") + start >> processing >> end +""", + ) + + group = next( + node for node in result.graph.tasks if isinstance(node, ContainerNode) and node.concept == CONCEPT_GROUP + ) + assert group.properties == {"inventory_visible": False, "structural_only": True} + mapped = group.branches["group"][0] + assert isinstance(mapped, ContainerNode) + assert mapped.concept == CONCEPT_LOOP + assert mapped.task_key == "processing__work" + assert mapped.properties["mapping"] == {"expand_arguments": ["params"], "has_partial": True} + end = next(node for node in result.graph.tasks if node.name == "end") + assert [dependency.upstream for dependency in end.dependencies] == ["processing__work"] + inventory = build_source_inventory([result.graph], source="airflow", source_dir="/dags") + assert inventory["summary"]["activity_count"] == 3 + assert [activity["name"] for activity in inventory["pipelines"][0]["activities"]] == [ + "start", + "work", + "end", + ] + + +def test_maps_taskflow_xcom_and_cross_dag_control_lineage(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow.decorators import dag, task +from airflow.operators.trigger_dagrun import TriggerDagRunOperator + +@dag(dag_id="lineage", schedule=None) +def build(): + @task + def extract(): + return 1 + + @task + def consume(value): + return value + + raw = extract() + done = consume(raw) + trigger = TriggerDagRunOperator(task_id="trigger_child", trigger_dag_id="child_dag") + done >> trigger + +build() +""", + ) + + nodes = {node.task_key: node for node in walk_nodes(result.graph.tasks)} + assert nodes["raw"].data_writes[0].signature == "xcom:raw" + assert nodes["done"].data_reads[0].signature == "xcom:raw" + assert nodes["trigger_child"].properties["invokes_workflow"] == "child_dag" + assert result.graph.lineage is not None + assert len(result.graph.lineage.data_edges) == 1 + assert result.graph.lineage.data_edges[0].source_task_key == "raw" + assert result.graph.lineage.data_edges[0].target_task_key == "done" + assert len(result.graph.lineage.control_edges) == 1 + assert result.graph.lineage.control_edges[0].target_workflow == "child_dag" + + +def test_unclaimed_comprehension_is_an_explicit_gap(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow import DAG +from airflow.operators.bash import BashOperator + +with DAG(dag_id="gap", schedule=None) as dag: + head = BashOperator(task_id="head", bash_command="echo head") + fanout = [BashOperator(task_id=f"work_{index}", bash_command="echo work") for index in range(3)] +""", + ) + + gaps = [node for node in result.graph.tasks if isinstance(node, GapNode)] + assert gaps + assert any("not captured" in (gap.reason or "") or "not claimed" in (gap.reason or "") for gap in gaps) + assert all(gap.properties["strategy"] == "unsupported" for gap in gaps) + + +def test_maps_explicit_airflow_assets_without_guessing_logical_identity(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow import DAG, Dataset +from airflow.operators.bash import BashOperator +from airflow.providers.databricks.operators.databricks import DatabricksCopyIntoOperator + +with DAG(dag_id="assets", schedule=[Dataset("logical.orders")]) as dag: + produce = BashOperator( + task_id="produce", + bash_command="echo ready", + outlets=[Dataset("logical.orders")], + ) + copy = DatabricksCopyIntoOperator( + task_id="copy", + file_location="s3://landing/orders", + table_name="main.bronze.orders", + ) + produce >> copy +""", + ) + + nodes = {node.task_key: node for node in walk_nodes(result.graph.tasks)} + assert nodes["produce"].data_writes[0].signature == "logical.orders" + assert nodes["produce"].data_writes[0].identity is None + assert nodes["copy"].data_reads[0].identity == "s3://landing/orders" + assert nodes["copy"].data_writes[0].identity == "main.bronze.orders" + assert result.graph.schedule is not None + assert result.graph.schedule.kind == "asset" + + +def test_maps_multiple_dag_declarations_independently(tmp_path: Path) -> None: + dag_path = tmp_path / "multiple.py" + dag_path.write_text( + """ +from airflow import DAG +from airflow.operators.bash import BashOperator as ShellTask + +first_dag = DAG(dag_id="first", schedule="@daily") +first_task = ShellTask(task_id="run", bash_command="echo first", dag=first_dag) + +second_dag = DAG(dag_id="second", schedule="@hourly") +second_task = ShellTask(task_id="run", bash_command="echo second", dag=second_dag) +""", + encoding="utf-8", + ) + + results = load_airflow_dag_results(dag_path, source_file="dags/multiple.py") + + assert [result.graph.name for result in results] == ["first", "second"] + assert [result.graph.schedule.expression for result in results if result.graph.schedule is not None] == [ + "@daily", + "@hourly", + ] + for result in results: + task = next(node for node in walk_nodes(result.graph.tasks) if not node.properties.get("structural_only")) + assert task.raw is not None + assert task.raw["operator_fqn"] == "airflow.operators.bash.BashOperator" + + +def test_exclusion_metadata_is_reflected_in_persisted_graphs(tmp_path: Path) -> None: + dag_path = tmp_path / "cross_dag.py" + dag_path.write_text( + """ +from airflow import DAG +from airflow.operators.bash import BashOperator +from airflow.operators.trigger_dagrun import TriggerDagRunOperator + +parent = DAG(dag_id="parent", schedule=None) +trigger = TriggerDagRunOperator(task_id="trigger", trigger_dag_id="child", dag=parent) + +child = DAG(dag_id="child", schedule=None) +work = BashOperator(task_id="work", bash_command="echo child", dag=child) +""", + encoding="utf-8", + ) + + results = load_discovery_results(dag_path, exclude_dags={"child"}) + by_name = {result.graph.name: result for result in results} + + assert by_name["child"].graph.properties["migration_status"] == "excluded" + assert by_name["parent"].graph.properties["reconciliation_status"] == "verified_with_gaps" + parent_task = next(iter(by_name["parent"].graph.tasks)) + assert parent_task.properties["strategy"] == "agentic" diff --git a/tests/unit/test_discovery_inventory.py b/tests/unit/test_discovery_inventory.py new file mode 100644 index 0000000..2f15d68 --- /dev/null +++ b/tests/unit/test_discovery_inventory.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import ast + +import flowx.discovery_inventory as discovery_inventory +from flowx.discovery_inventory import INVENTORY_VISIBLE_PROPERTY, STRATEGY_PROPERTY, build_source_inventory +from flowx.models.discovery import CONCEPT_GROUP, CONCEPT_NOTEBOOK, ContainerNode, SourceGraph, SourceNode +from flowx.models.ir import ControlEdge, Lineage + + +def _node(task_key: str, strategy: str) -> SourceNode: + return SourceNode( + source_id=task_key, + task_key=task_key, + concept=CONCEPT_NOTEBOOK, + source="unit", + name=task_key, + native_type="Notebook", + properties={STRATEGY_PROPERTY: strategy}, + raw={"task_key": task_key}, + ) + + +def test_emitter_has_no_source_specific_imports() -> None: + module_path = (discovery_inventory.__file__ or "").rstrip("c") + tree = ast.parse(open(module_path, encoding="utf-8").read()) + imports = [node.module for node in ast.walk(tree) if isinstance(node, ast.ImportFrom) and node.module is not None] + assert not any(name.startswith("flowx.sources") for name in imports) + + +def test_inventory_shape_counts_and_structural_containers() -> None: + structural_group = ContainerNode( + source_id="group:etl", + task_key="etl", + concept=CONCEPT_GROUP, + source="unit", + name="etl", + native_type="TaskGroup", + properties={INVENTORY_VISIBLE_PROPERTY: False, "structural_only": True}, + branches={"group": [_node("extract", "deterministic"), _node("load", "agentic")]}, + ) + graph = SourceGraph(name="workflow", source="unit", tasks=[structural_group]) + + inventory = build_source_inventory([graph], source="unit", source_dir="/source") + + assert set(inventory) == {"source", "source_dir", "pipelines", "summary"} + assert [item["name"] for item in inventory["pipelines"][0]["activities"]] == ["extract", "load"] + assert inventory["summary"] == { + "pipeline_count": 1, + "activity_count": 2, + "deterministic_count": 1, + "agentic_count": 1, + "unsupported_count": 0, + "coverage_pct": 100.0, + } + + +def test_inventory_emits_lineage_and_additive_source_fields() -> None: + graph = SourceGraph( + name="workflow", + source="unit", + tasks=[_node("task", "deterministic")], + lineage=Lineage( + control_edges=[ + ControlEdge( + source_workflow="workflow", + target_workflow="child", + via_task_key="task", + wait_for_completion=False, + ) + ] + ), + ) + + entry = build_source_inventory([graph], source="unit", source_dir="/source")["pipelines"][0] + + assert entry["activities"][0] == { + "name": "task", + "type": "Notebook", + "strategy": "deterministic", + "original_type": "Notebook", + "dependencies": [], + "raw": {"task_key": "task"}, + } + assert entry["lineage"]["control_edges"][0]["target_workflow"] == "child" From 4290cf930c0aefb5b38755cb31a3afbf6fb9401f Mon Sep 17 00:00:00 2001 From: peterpark-db <127158477+peterpark-db@users.noreply.github.com> Date: Mon, 21 Sep 2026 13:26:26 -0700 Subject: [PATCH 2/3] Fix Airflow source graph fidelity --- .../sources/airflow/discovery_mapping.py | 40 ++++- src/flowx/sources/airflow/loader/graph.py | 15 +- src/flowx/sources/airflow/loader/lowering.py | 6 + src/flowx/sources/airflow/loader/visitor.py | 9 +- tests/unit/test_airflow_discovery_mapping.py | 148 ++++++++++++++++++ 5 files changed, 211 insertions(+), 7 deletions(-) diff --git a/src/flowx/sources/airflow/discovery_mapping.py b/src/flowx/sources/airflow/discovery_mapping.py index 5716e8a..c3f6c43 100644 --- a/src/flowx/sources/airflow/discovery_mapping.py +++ b/src/flowx/sources/airflow/discovery_mapping.py @@ -37,6 +37,7 @@ from flowx.sources.airflow.loader.captures import DagDeclaration, SourceSpan from flowx.sources.airflow.loader.graph import _allocate_task_keys, _expand_group_edges from flowx.sources.airflow.loader.policy import _job_timeout_seconds +from flowx.sources.airflow.loader.schedule import _asset_expression from flowx.sources.airflow.loader.visitor import _DagVisitor _QUERY_OPERATORS = frozenset( @@ -105,6 +106,7 @@ def build_airflow_source_graph( {variable: task.task_id for variable, task in visitor.taskflow_tasks.items()}, {variable: task_id for variable, (task_id, _, _) in visitor.taskgroup_calls.items()}, visitor.groups, + visitor.capture_source_nodes, ) expanded_edges = _expand_group_edges(visitor.edges, visitor.groups, visitor.group_vars) upstreams: dict[str, list[str]] = {capture_id: [] for capture_id in task_keys} @@ -135,6 +137,8 @@ def build_airflow_source_graph( "source": ast.get_source_segment(source, declaration.node) or ast.unparse(declaration.node), "dag_arguments": {name: _expression_payload(value, source) for name, value in visitor.dag_kwargs.items()}, } + if declaration.factory is not None: + declaration_raw["factory_definition"] = _definition_payload(declaration.factory, source) graph = SourceGraph( name=pipeline.name, source=SOURCE_AIRFLOW, @@ -249,6 +253,9 @@ def _captured_node( raw["arguments"] = {name: _expression_payload(value, source) for name, value in kwargs.items()} raw["operator_fqn"] = visitor.task_captures[capture_id].operator_fqn raw["argument_disposition"] = ops.argument_classification(operator, kwargs) + callable_definition = visitor.resolved_callable_for(capture_id) + if callable_definition is not None: + raw["callable_definition"] = _definition_payload(callable_definition, source) run_condition = ops.literal_str(kwargs.get("trigger_rule")) policy = _policy(kwargs) concept = _operator_concept(operator, capture_id in visitor.mapped) @@ -257,7 +264,11 @@ def _captured_node( target = ops.literal_str(kwargs.get("trigger_dag_id")) if target is not None: properties[INVOKES_WORKFLOW_PROPERTY] = target - properties[INVOKES_WAIT_PROPERTY] = bool(ops.literal_value(kwargs.get("wait_for_completion"))) + wait_node = kwargs.get("wait_for_completion") + wait_value = ops.literal_value(wait_node) + properties[INVOKES_WAIT_PROPERTY] = ( + False if wait_node is None else wait_value if isinstance(wait_value, bool) else None + ) elif operator in {"DatabricksRunNowOperator", "DatabricksRunNowDeferrableOperator"}: job_id = ops.literal_value(kwargs.get("job_id")) if job_id is not None: @@ -334,9 +345,15 @@ def _captured_node( "unresolved_arguments": list(task.unresolved_arguments), } ) + taskflow_definition = visitor.taskflow_defs[task.def_name][0] + raw["callable_definition"] = _definition_payload(taskflow_definition, source) + data_upstreams = [ + *[task.positional_deps[position] for position in sorted(task.positional_deps)], + *task.keyword_deps.values(), + ] reads = [ DataAsset(signature=f"xcom:{task_keys[upstream]}", asset_type="value") - for upstream in dict.fromkeys(upstreams) + for upstream in dict.fromkeys(data_upstreams) if upstream in task_keys ] writes = [DataAsset(signature=f"xcom:{task_key}", asset_type="value")] @@ -390,6 +407,9 @@ def _captured_node( task_id, definition, mapped = visitor.taskgroup_calls[capture_id] raw.update({"task_group_callable": definition, "mapped": mapped}) + taskgroup_definition = visitor.taskgroup_defs.get(definition) + if taskgroup_definition is not None: + raw["callable_definition"] = _definition_payload(taskgroup_definition, source) return ContainerNode( source_id=capture_id, task_key=task_key, @@ -631,8 +651,8 @@ def _source_schedule(visitor: _DagVisitor, source: str) -> ScheduleSpec | None: source_expression = ast.get_source_segment(source, node) or ast.unparse(node) expression = source_expression if literal is None else _json_safe(literal, fallback=source_expression) kind = "schedule" if isinstance(literal, str) else "interval" - is_asset_expression = isinstance(node, ast.Call) and "Asset" in ast.unparse(node) - if isinstance(node, (ast.List, ast.Tuple, ast.Set)) or is_asset_expression: + is_asset_expression = _asset_expression(node, visitor._aliases, visitor.asset_definitions) is not None + if is_asset_expression: kind = "asset" return ScheduleSpec( kind=kind, @@ -650,6 +670,18 @@ def _expression_payload(node: ast.expr, source: str) -> dict[str, Any]: } +def _definition_payload(node: ast.AST, source: str) -> dict[str, Any]: + return { + "source": ast.get_source_segment(source, node) or ast.unparse(node), + "source_span": { + "line": getattr(node, "lineno", 0), + "column": getattr(node, "col_offset", 0), + "end_line": getattr(node, "end_lineno", getattr(node, "lineno", 0)), + "end_column": getattr(node, "end_col_offset", getattr(node, "col_offset", 0)), + }, + } + + def _json_safe(value: Any, *, fallback: Any) -> Any: try: json.dumps(value) diff --git a/src/flowx/sources/airflow/loader/graph.py b/src/flowx/sources/airflow/loader/graph.py index a29d764..a815d8a 100644 --- a/src/flowx/sources/airflow/loader/graph.py +++ b/src/flowx/sources/airflow/loader/graph.py @@ -3,6 +3,7 @@ from __future__ import annotations import ast +from collections.abc import Mapping from flowx.sources.airflow import operators as ops from flowx.sources.airflow.loader.ast_utils import _sanitize_task_key @@ -13,15 +14,27 @@ def _allocate_task_keys( taskflow_task_ids: dict[str, str], taskgroup_task_ids: dict[str, str], groups: dict[str, str], + capture_source_nodes: Mapping[str, ast.AST], ) -> dict[str, str]: """Allocates stable, collision-free task keys in capture order.""" task_ids = {variable: task_id for variable, (task_id, _, _) in operators.items()} task_ids.update(taskflow_task_ids) task_ids.update(taskgroup_task_ids) + capture_indexes = {capture_id: index for index, capture_id in enumerate(capture_source_nodes)} + capture_ids = sorted( + task_ids, + key=lambda capture_id: ( + getattr(capture_source_nodes.get(capture_id), "lineno", 0), + getattr(capture_source_nodes.get(capture_id), "col_offset", 0), + capture_indexes.get(capture_id, len(capture_indexes)), + ), + ) + allocated: dict[str, str] = {} used: set[str] = set() - for variable, task_id in task_ids.items(): + for variable in capture_ids: + task_id = task_ids[variable] base = _sanitize_task_key(task_id) if variable in groups: base = f"{groups[variable]}__{base}" diff --git a/src/flowx/sources/airflow/loader/lowering.py b/src/flowx/sources/airflow/loader/lowering.py index 1f26722..1021e80 100644 --- a/src/flowx/sources/airflow/loader/lowering.py +++ b/src/flowx/sources/airflow/loader/lowering.py @@ -83,6 +83,7 @@ def _load_airflow_module( {var: task.task_id for var, task in visitor.taskflow_tasks.items()}, {var: task_id for var, (task_id, _, _) in visitor.taskgroup_calls.items()}, visitor.groups, + visitor.capture_source_nodes, ) # Expand group-level edges (`group_a >> group_b`, `task >> group`, ...) into edges between the @@ -205,10 +206,12 @@ def _dep(upstream_var: str, outcome: str | None) -> str: return dbt_key_remap.get(key, key) tasks: list[Activity] = [] + task_capture_ids: dict[int, str] = {} placeholder_capture_ids: dict[int, str] = {} helper_expansion_ids = {str(item["capture_id"]) for item in visitor.helper_expansions} def append_task(activity: Activity, capture_id: str) -> None: + task_capture_ids[id(activity)] = capture_id for placeholder in _iter_placeholders([activity]): placeholder_capture_ids[id(placeholder)] = capture_id tasks.append(activity) @@ -573,6 +576,9 @@ def append_task(activity: Activity, capture_id: str) -> None: placeholder.depends_on = depends_on append_task(placeholder, var) + capture_order = {capture_id: index for index, capture_id in enumerate(var_to_task_key)} + tasks.sort(key=lambda activity: capture_order[task_capture_ids[id(activity)]]) + # Declare every job parameter -- those referenced in templates plus any from the DAG's # params={...} -- each with a default (Databricks requires one): the params={...} default when # present; a reserved logical-date parameter its schedule-aware time ref so a native backfill can diff --git a/src/flowx/sources/airflow/loader/visitor.py b/src/flowx/sources/airflow/loader/visitor.py index 51c2772..c02252c 100644 --- a/src/flowx/sources/airflow/loader/visitor.py +++ b/src/flowx/sources/airflow/loader/visitor.py @@ -114,7 +114,7 @@ def __init__(self, module: ast.Module, *, target_dag_variable: str | None = None self.taskflow_defs: dict[str, tuple[ast.FunctionDef | ast.AsyncFunctionDef, str]] = {} # @task_group def names -- a group is a sub-pipeline, not a single renderable task, so an # invocation routes to a placeholder + gap rather than being expanded here. - self.taskgroup_defs: set[str] = set() + self.taskgroup_defs: dict[str, ast.FunctionDef | ast.AsyncFunctionDef] = {} for fn in _iter_functions(module): decorator = next( ( @@ -127,7 +127,7 @@ def __init__(self, module: ast.Module, *, target_dag_variable: str | None = None if decorator is not None: self.taskflow_defs[fn.name] = (fn, decorator) elif _has_decorator(fn, _TASK_GROUP_DECORATORS, self._aliases): - self.taskgroup_defs.add(fn.name) + self.taskgroup_defs[fn.name] = fn # TaskFlow task instances: var name -> _TaskFlowTask (id, def-name, decorator, arg bindings). self.taskflow_tasks: dict[str, _TaskFlowTask] = {} # @task_group invocations: var name -> (task_id, def-name, is_mapped). @@ -157,6 +157,11 @@ def functions_for(self, task_var: str) -> dict[str, ast.FunctionDef]: functions[name] = definition return functions + def resolved_callable_for(self, task_var: str) -> ast.FunctionDef | None: + """Returns the callable definition resolved for a classic operator task.""" + resolved = self._resolved_callables.get(task_var) + return resolved[1] if resolved is not None else None + def visit_FunctionDef(self, node: ast.FunctionDef) -> None: # A @task- or @task_group-decorated function defines a task / sub-pipeline from its body, # which is internal logic rather than DAG structure, so don't descend. @dag marks the diff --git a/tests/unit/test_airflow_discovery_mapping.py b/tests/unit/test_airflow_discovery_mapping.py index 5c70e7a..902e39d 100644 --- a/tests/unit/test_airflow_discovery_mapping.py +++ b/tests/unit/test_airflow_discovery_mapping.py @@ -155,6 +155,129 @@ def consume(value): assert result.graph.lineage.control_edges[0].target_workflow == "child_dag" +def test_persists_factory_and_task_callable_definitions(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow.decorators import dag, task, task_group +from airflow.operators.python import PythonOperator + +def classic_callable(): + return "classic" + +@task_group +def grouped(): + @task + def nested(): + return "nested" + nested() + +@dag(dag_id="callables", schedule=None) +def build(): + @task + def taskflow_callable(): + return "taskflow" + + classic = PythonOperator(task_id="classic", python_callable=classic_callable) + taskflow = taskflow_callable() + group = grouped() + classic >> taskflow >> group + +build() +""", + ) + + assert result.graph.raw is not None + assert "def build():" in result.graph.raw["factory_definition"]["source"] + nodes = {node.source_id: node for node in walk_nodes(result.graph.tasks)} + assert "def classic_callable():" in nodes["classic"].raw["callable_definition"]["source"] + assert "def taskflow_callable():" in nodes["taskflow"].raw["callable_definition"]["source"] + assert "def grouped():" in nodes["group"].raw["callable_definition"]["source"] + + +def test_ordering_only_taskflow_dependency_does_not_create_xcom_lineage(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow.decorators import dag, task + +@dag(dag_id="ordering_only", schedule=None) +def build(): + @task + def first(): + return 1 + + @task + def second(): + return 2 + + first_task = first() + second_task = second() + first_task >> second_task + +build() +""", + ) + + nodes = {node.source_id: node for node in walk_nodes(result.graph.tasks)} + assert [dependency.upstream for dependency in nodes["second_task"].dependencies] == ["first_task"] + assert nodes["second_task"].data_reads == [] + assert result.graph.lineage is not None + assert result.graph.lineage.data_edges == [] + + +def test_preserves_unknown_trigger_dag_run_wait_semantics(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow import DAG +from airflow.models import Variable +from airflow.operators.trigger_dagrun import TriggerDagRunOperator + +with DAG(dag_id="dynamic_wait", schedule=None) as dag: + trigger = TriggerDagRunOperator( + task_id="trigger", + trigger_dag_id="child", + wait_for_completion=Variable.get("WAIT_FOR_CHILD") == "true", + ) +""", + ) + + trigger = result.graph.tasks[0] + assert trigger.properties["invokes_wait"] is None + assert result.graph.lineage is not None + assert result.graph.lineage.control_edges[0].wait_for_completion is None + + +def test_preserves_mixed_task_kind_source_order_and_collision_allocation(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow.decorators import dag, task +from airflow.operators.bash import BashOperator + +@task +def taskflow_callable(): + return 1 + +@dag(dag_id="mixed_order", schedule=None) +def build(): + same = taskflow_callable() + middle = BashOperator(task_id="same", bash_command="echo middle") + last = taskflow_callable() + +build() +""", + ) + + assert [(node.source_id, node.task_key) for node in result.graph.tasks] == [ + ("same", "same"), + ("middle", "same__2"), + ("last", "last"), + ] + assert [task.task_key for task in result.pipeline.tasks] == ["same", "same__2", "last"] + + def test_unclaimed_comprehension_is_an_explicit_gap(tmp_path: Path) -> None: result = _load( tmp_path, @@ -206,6 +329,31 @@ def test_maps_explicit_airflow_assets_without_guessing_logical_identity(tmp_path assert result.graph.schedule.kind == "asset" +def test_maps_named_composed_airflow_three_asset_schedule(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow.sdk import Asset, DAG +from airflow.providers.standard.operators.bash import BashOperator + +orders = Asset("x-databricks-table://main.raw.orders") +customers = Asset("x-databricks-table://main.raw.customers") + +with DAG(dag_id="asset_expression", schedule=orders & customers) as dag: + run = BashOperator(task_id="run", bash_command="echo ready") +""", + ) + + assert result.graph.schedule is not None + assert result.graph.schedule.kind == "asset" + assert result.pipeline.schedule == { + "kind": "table_update", + "table_names": ["main.raw.orders", "main.raw.customers"], + "condition": "ALL_UPDATED", + "pause_status": "UNPAUSED", + } + + def test_maps_multiple_dag_declarations_independently(tmp_path: Path) -> None: dag_path = tmp_path / "multiple.py" dag_path.write_text( From 29207be41bdbb790bf720310c6ee527df0424faf Mon Sep 17 00:00:00 2001 From: peterpark-db <127158477+peterpark-db@users.noreply.github.com> Date: Mon, 21 Sep 2026 14:13:50 -0700 Subject: [PATCH 3/3] Fix Airflow source graph review findings --- src/flowx/discovery_inventory.py | 12 +- .../sources/airflow/callable_notebook.py | 39 +- .../sources/airflow/discovery_mapping.py | 621 ++++++++++++------ src/flowx/sources/airflow/loader/api.py | 27 +- src/flowx/sources/airflow/loader/captures.py | 1 + src/flowx/sources/airflow/loader/graph.py | 55 +- src/flowx/sources/airflow/loader/lowering.py | 22 +- src/flowx/sources/airflow/loader/policy.py | 22 +- src/flowx/sources/airflow/loader/reconcile.py | 4 +- src/flowx/sources/airflow/loader/schedule.py | 18 + src/flowx/sources/airflow/loader/visitor.py | 86 ++- tests/unit/test_airflow_discovery_mapping.py | 257 +++++++- 12 files changed, 889 insertions(+), 275 deletions(-) diff --git a/src/flowx/discovery_inventory.py b/src/flowx/discovery_inventory.py index b1efc46..e44d3a8 100644 --- a/src/flowx/discovery_inventory.py +++ b/src/flowx/discovery_inventory.py @@ -21,7 +21,17 @@ def build_source_inventory( source_dir: str, include_empty_pipelines: bool = True, ) -> dict[str, Any]: - """Projects source graphs into the common discovery inventory shape.""" + """Projects source graphs into the common discovery inventory shape. + + Args: + graphs: Source-faithful workflow graphs to summarize. + source: Source-system identifier written to the inventory. + source_dir: Source directory label written to the inventory. + include_empty_pipelines: Whether workflows without visible activities remain in the inventory. + + Returns: + The serializable inventory payload with per-pipeline activities and aggregate coverage. + """ pipeline_entries: list[dict[str, Any]] = [] deterministic = 0 agentic = 0 diff --git a/src/flowx/sources/airflow/callable_notebook.py b/src/flowx/sources/airflow/callable_notebook.py index 510cf0d..7e9f7e5 100644 --- a/src/flowx/sources/airflow/callable_notebook.py +++ b/src/flowx/sources/airflow/callable_notebook.py @@ -19,7 +19,10 @@ _AIRFLOW_IMPORT_ROOTS: frozenset[str] = frozenset({"airflow", "cosmos", "airflow_dbt"}) -def _enclosing_statements(module: ast.Module, func: ast.FunctionDef) -> list[ast.stmt]: +def _enclosing_statements( + module: ast.Module, + func: ast.FunctionDef | ast.AsyncFunctionDef, +) -> list[ast.stmt]: """Returns safe statements visible from the callable's enclosing function scopes.""" scopes = [ node @@ -89,7 +92,9 @@ def _names_used(node: ast.AST) -> set[str]: def _closure( - func: ast.FunctionDef, defs: dict[str, ast.stmt], assigns: dict[str, ast.stmt] + func: ast.FunctionDef | ast.AsyncFunctionDef, + defs: dict[str, ast.stmt], + assigns: dict[str, ast.stmt], ) -> tuple[list[str], set[str]]: """Returns transitively-referenced module symbols (defs+assigns) in source order, plus all names used. @@ -118,6 +123,36 @@ def _closure( return ordered, all_names +def render_source_closure(func: ast.FunctionDef | ast.AsyncFunctionDef, source: str) -> str: + """Renders a callable and every statically resolved source dependency it uses. + + Args: + func: Callable definition whose source closure is required. + source: Complete DAG module source used to resolve and slice dependencies. + + Returns: + Source containing referenced imports, helpers, classes, constants, and the callable. + """ + module = ast.parse(source) + enclosing_statements = _enclosing_statements(module, func) + definitions, assignments = _module_symbols(module, enclosing_statements) + imports = _import_bindings(module, enclosing_statements) + dependency_names, used_names = _closure(func, definitions, assignments) + + import_nodes: dict[int, ast.stmt] = {} + for name in used_names: + binding = imports.get(name) + if binding is not None: + import_nodes[id(binding[0])] = binding[0] + + nodes = [ + *sorted(import_nodes.values(), key=lambda node: node.lineno), + *(definitions.get(name) or assignments[name] for name in dependency_names), + func, + ] + return "\n\n".join(segment for node in nodes if (segment := ast.get_source_segment(source, node)) is not None) + + def render_definitions(func: ast.FunctionDef, source: str, *, note: str) -> str: """Renders the callable's ``def`` plus its transitive deps as a notebook prelude (no invocation). diff --git a/src/flowx/sources/airflow/discovery_mapping.py b/src/flowx/sources/airflow/discovery_mapping.py index c3f6c43..412a2fd 100644 --- a/src/flowx/sources/airflow/discovery_mapping.py +++ b/src/flowx/sources/airflow/discovery_mapping.py @@ -7,7 +7,7 @@ from collections.abc import Iterable from dataclasses import dataclass from pathlib import Path -from typing import Any +from typing import Any, TypedDict from flowx.discovery_lineage import INVOKES_WAIT_PROPERTY, INVOKES_WORKFLOW_PROPERTY, walk_nodes, with_graph_lineage from flowx.models.discovery import ( @@ -31,14 +31,14 @@ SourceNode, ) from flowx.models.ir import Activity, DataAsset, ForEachActivity, Pipeline, PlaceholderActivity +from flowx.sources.airflow import callable_notebook, templating from flowx.sources.airflow import operators as ops -from flowx.sources.airflow import templating from flowx.sources.airflow.audit import SourceAudit from flowx.sources.airflow.loader.captures import DagDeclaration, SourceSpan -from flowx.sources.airflow.loader.graph import _allocate_task_keys, _expand_group_edges -from flowx.sources.airflow.loader.policy import _job_timeout_seconds -from flowx.sources.airflow.loader.schedule import _asset_expression -from flowx.sources.airflow.loader.visitor import _DagVisitor +from flowx.sources.airflow.loader.graph import allocate_task_keys, expand_group_edges +from flowx.sources.airflow.loader.policy import job_timeout_seconds +from flowx.sources.airflow.loader.schedule import asset_expression +from flowx.sources.airflow.loader.visitor import DagVisitor _QUERY_OPERATORS = frozenset( { @@ -64,16 +64,43 @@ _BRANCH_OPERATORS = frozenset({"BranchPythonOperator", "ShortCircuitOperator"}) +class _SourceNodeArguments(TypedDict): + """Carries the fields shared by source-node variants.""" + + source_id: str + task_key: str + source: str + name: str | None + native_type: str | None + dependencies: list[SourceDependency] + run_condition: str | None + policy: PolicySpec | None + data_reads: list[DataAsset] + data_writes: list[DataAsset] + properties: dict[str, Any] + raw: dict[str, Any] | None + + @dataclass(slots=True, kw_only=True) class AirflowDiscoveryResult: - """One Airflow DAG represented as both current IR and source-faithful discovery graph.""" + """One Airflow DAG represented as current IR and a source-faithful graph. + + Attributes: + pipeline: Existing Airflow-to-Databricks intermediate representation. + graph: Shared discovery graph projected from the same static capture. + """ pipeline: Pipeline graph: SourceGraph def sync_graph_translation_metadata(graph: SourceGraph, pipeline: Pipeline) -> None: - """Refreshes target classifications after exclusions or cross-DAG rewrites.""" + """Refreshes target classifications after exclusions or cross-DAG rewrites. + + Args: + graph: Discovery graph whose translation metadata needs synchronization. + pipeline: Final pipeline carrying exclusion and reconciliation changes. + """ activities = _activity_index(pipeline.tasks) for node in walk_nodes(graph.tasks): if node.properties.get("structural_only"): @@ -96,19 +123,31 @@ def build_airflow_source_graph( source_file: str, source: str, declaration: DagDeclaration, - visitor: _DagVisitor, + visitor: DagVisitor, audit: SourceAudit, pipeline: Pipeline, ) -> SourceGraph: - """Projects one captured DAG onto ``SourceGraph`` without reparsing it.""" - task_keys = _allocate_task_keys( + """Projects one captured DAG onto the shared discovery graph. + + Args: + dag_path: Path to the original DAG module. + source_file: Stable source label relative to the discovery root. + source: Complete DAG module source. + declaration: Selected DAG declaration within the module. + visitor: Completed static capture for the declaration. + audit: Independent source audit for the declaration. + pipeline: Existing lowered pipeline produced from the same capture. + + Returns: + A source-faithful graph carrying audit, lineage, and translation metadata. + """ + task_keys = allocate_task_keys( visitor.operators, {variable: task.task_id for variable, task in visitor.taskflow_tasks.items()}, {variable: task_id for variable, (task_id, _, _) in visitor.taskgroup_calls.items()}, visitor.groups, - visitor.capture_source_nodes, ) - expanded_edges = _expand_group_edges(visitor.edges, visitor.groups, visitor.group_vars) + expanded_edges = expand_group_edges(visitor.edges, visitor.groups, visitor.group_vars) upstreams: dict[str, list[str]] = {capture_id: [] for capture_id in task_keys} for upstream, downstream in expanded_edges: if upstream in task_keys and downstream in upstreams: @@ -116,7 +155,8 @@ def build_airflow_source_graph( activities = _activity_index(pipeline.tasks) nodes_by_capture: dict[str, SourceNode] = {} - for capture_id, task_key in task_keys.items(): + for capture_id in _capture_ids_in_source_order(task_keys, visitor.capture_source_nodes): + task_key = task_keys[capture_id] nodes_by_capture[capture_id] = _captured_node( capture_id=capture_id, task_key=task_key, @@ -127,8 +167,14 @@ def build_airflow_source_graph( activities=activities, ) - tasks = _nest_task_groups(nodes_by_capture, visitor, source) - tasks.extend(_gap_nodes(visitor, source)) + gap_nodes, gap_group_paths = _gap_nodes(visitor, source) + ordered_nodes = dict( + sorted( + [*nodes_by_capture.items(), *gap_nodes.items()], + key=lambda item: _source_node_sort_key(item[1]), + ) + ) + tasks = _nest_task_groups(ordered_nodes, {**visitor.groups, **gap_group_paths}, visitor, source) declaration_raw = { "capture_id": declaration.capture_id, "kind": declaration.kind, @@ -146,7 +192,7 @@ def build_airflow_source_graph( parameters={name: ParameterSpec(default=value) for name, value in visitor.dag_params.items()}, schedule=_source_schedule(visitor, source), default_policy=_policy(visitor.default_args), - run_timeout_seconds=_job_timeout_seconds(visitor), + run_timeout_seconds=job_timeout_seconds(visitor), tags=list(visitor.dag_user_tags), tasks=tasks, properties={ @@ -188,7 +234,18 @@ def failed_declaration_source_graph( declaration: DagDeclaration, pipeline: Pipeline, ) -> SourceGraph: - """Builds a reportable graph for a DAG declaration that static capture rejected.""" + """Builds a reportable graph for a rejected DAG declaration. + + Args: + dag_path: Path to the original DAG module. + source_file: Stable source label relative to the discovery root. + source: Complete DAG module source. + declaration: DAG declaration rejected by static capture. + pipeline: Failed pipeline carrying reconciliation findings. + + Returns: + A graph containing one explicit declaration gap. + """ raw_source = ast.get_source_segment(source, declaration.node) or ast.unparse(declaration.node) gap = GapNode( source_id=declaration.capture_id, @@ -225,10 +282,65 @@ def _captured_node( task_key: str, upstreams: list[str], task_keys: dict[str, str], - visitor: _DagVisitor, + visitor: DagVisitor, source: str, activities: dict[str, Activity], ) -> SourceNode: + """Builds the source node for one captured operator, TaskFlow task, or task group.""" + dependencies, properties, raw = _captured_node_metadata( + capture_id=capture_id, + task_key=task_key, + upstreams=upstreams, + task_keys=task_keys, + visitor=visitor, + source=source, + activities=activities, + ) + if capture_id in visitor.operators: + return _operator_node( + capture_id=capture_id, + task_key=task_key, + dependencies=dependencies, + visitor=visitor, + source=source, + activities=activities, + properties=properties, + raw=raw, + ) + if capture_id in visitor.taskflow_tasks: + return _taskflow_node( + capture_id=capture_id, + task_key=task_key, + dependencies=dependencies, + task_keys=task_keys, + visitor=visitor, + source=source, + activities=activities, + properties=properties, + raw=raw, + ) + return _taskgroup_node( + capture_id=capture_id, + task_key=task_key, + dependencies=dependencies, + visitor=visitor, + source=source, + properties=properties, + raw=raw, + ) + + +def _captured_node_metadata( + *, + capture_id: str, + task_key: str, + upstreams: list[str], + task_keys: dict[str, str], + visitor: DagVisitor, + source: str, + activities: dict[str, Activity], +) -> tuple[list[SourceDependency], dict[str, Any], dict[str, Any]]: + """Builds the dependencies and source metadata shared by captured node variants.""" dependencies = [ SourceDependency(upstream=task_keys[upstream], resolved=True) for upstream in dict.fromkeys(upstreams) @@ -246,170 +358,179 @@ def _captured_node( "source_span": _span_dict(span), } properties: dict[str, Any] = {"strategy": _strategy_for(task_key, activities, pipeline_status=None)} - run_condition: str | None = None + return dependencies, properties, raw - if capture_id in visitor.operators: - task_id, operator, kwargs = visitor.operators[capture_id] - raw["arguments"] = {name: _expression_payload(value, source) for name, value in kwargs.items()} - raw["operator_fqn"] = visitor.task_captures[capture_id].operator_fqn - raw["argument_disposition"] = ops.argument_classification(operator, kwargs) - callable_definition = visitor.resolved_callable_for(capture_id) - if callable_definition is not None: - raw["callable_definition"] = _definition_payload(callable_definition, source) - run_condition = ops.literal_str(kwargs.get("trigger_rule")) - policy = _policy(kwargs) - concept = _operator_concept(operator, capture_id in visitor.mapped) - reads, writes = _operator_assets(operator, kwargs) - if operator == "TriggerDagRunOperator": - target = ops.literal_str(kwargs.get("trigger_dag_id")) - if target is not None: - properties[INVOKES_WORKFLOW_PROPERTY] = target - wait_node = kwargs.get("wait_for_completion") - wait_value = ops.literal_value(wait_node) - properties[INVOKES_WAIT_PROPERTY] = ( - False if wait_node is None else wait_value if isinstance(wait_value, bool) else None - ) - elif operator in {"DatabricksRunNowOperator", "DatabricksRunNowDeferrableOperator"}: - job_id = ops.literal_value(kwargs.get("job_id")) - if job_id is not None: - properties[INVOKES_WORKFLOW_PROPERTY] = f"databricks-job:{job_id}" - elif operator in {"ExternalTaskSensor", "ExternalTaskSensorAsync"}: - properties["external_workflow_wait"] = { - "dag_id": ops.literal_str(kwargs.get("external_dag_id")), - "task_id": ops.literal_str(kwargs.get("external_task_id")), - } - if capture_id in visitor.mapped: - properties["mapping"] = { - "expand_arguments": list(visitor.expand_kwargs.get(capture_id, [])), - "has_partial": capture_id in visitor.partial_mapped, - } - return ContainerNode( - source_id=capture_id, - task_key=task_key, - concept=CONCEPT_LOOP, - source=SOURCE_AIRFLOW, - name=task_id, - native_type=operator, - dependencies=dependencies, - run_condition=run_condition, - policy=policy, - data_reads=reads, - data_writes=writes, - properties=properties, - raw=raw, - branches={"body": []}, - ) - activity = activities.get(task_key) - if isinstance(activity, PlaceholderActivity) or concept == CONCEPT_GAP: - reason = activity.comment if isinstance(activity, PlaceholderActivity) else None - return GapNode( - source_id=capture_id, - task_key=task_key, - source=SOURCE_AIRFLOW, - name=task_id, - native_type=operator, - dependencies=dependencies, - run_condition=run_condition, - policy=policy, - data_reads=reads, - data_writes=writes, - properties=properties, - raw=raw, - reason=reason or f"Airflow operator {operator!r} has no deterministic mapping.", - ) - return SourceNode( - source_id=capture_id, - task_key=task_key, - concept=concept, - source=SOURCE_AIRFLOW, - name=task_id, - native_type=operator, - dependencies=dependencies, - run_condition=run_condition, - policy=policy, - data_reads=reads, - data_writes=writes, - properties=properties, - raw=raw, - ) - if capture_id in visitor.taskflow_tasks: - task = visitor.taskflow_tasks[capture_id] - raw.update( - { - "decorator": task.decorator, - "callable": task.def_name, - "source_reference": task.source_reference, - "positional_arguments": dict(task.positional_values), - "keyword_arguments": dict(task.keyword_values), - "unresolved_arguments": list(task.unresolved_arguments), - } +def _operator_node( + *, + capture_id: str, + task_key: str, + dependencies: list[SourceDependency], + visitor: DagVisitor, + source: str, + activities: dict[str, Activity], + properties: dict[str, Any], + raw: dict[str, Any], +) -> SourceNode: + """Builds a source node for one classic operator capture.""" + task_id, operator, arguments = visitor.operators[capture_id] + raw["arguments"] = {name: _expression_payload(value, source) for name, value in arguments.items()} + raw["operator_fqn"] = visitor.task_captures[capture_id].operator_fqn + raw["argument_disposition"] = ops.argument_classification(operator, arguments) + callable_definition = visitor.resolved_callable_for(capture_id) + if callable_definition is not None: + raw["callable_definition"] = _definition_payload( + callable_definition, + source, + closure_source=callable_notebook.render_source_closure(callable_definition, source), ) - taskflow_definition = visitor.taskflow_defs[task.def_name][0] - raw["callable_definition"] = _definition_payload(taskflow_definition, source) - data_upstreams = [ - *[task.positional_deps[position] for position in sorted(task.positional_deps)], - *task.keyword_deps.values(), - ] - reads = [ - DataAsset(signature=f"xcom:{task_keys[upstream]}", asset_type="value") - for upstream in dict.fromkeys(data_upstreams) - if upstream in task_keys - ] - writes = [DataAsset(signature=f"xcom:{task_key}", asset_type="value")] - if capture_id in visitor.mapped: - properties["mapping"] = { - "expand_argument": task.expand_kwarg, - "expand_items_json": task.expand_items_json, - } - return ContainerNode( - source_id=capture_id, - task_key=task_key, - concept=CONCEPT_LOOP, - source=SOURCE_AIRFLOW, - name=task.task_id, - native_type=task.decorator, - dependencies=dependencies, - data_reads=reads, - data_writes=writes, - properties=properties, - raw=raw, - branches={"body": []}, - ) - activity = activities.get(task_key) - if isinstance(activity, PlaceholderActivity): - return GapNode( - source_id=capture_id, - task_key=task_key, - source=SOURCE_AIRFLOW, - name=task.task_id, - native_type=task.decorator, - dependencies=dependencies, - data_reads=reads, - data_writes=writes, - properties=properties, - raw=raw, - reason=activity.comment or "TaskFlow invocation requires manual migration.", - ) - return SourceNode( - source_id=capture_id, - task_key=task_key, - concept=CONCEPT_SCRIPT, - source=SOURCE_AIRFLOW, - name=task.task_id, - native_type=task.decorator, - dependencies=dependencies, - data_reads=reads, - data_writes=writes, - properties=properties, - raw=raw, + + run_condition = ops.literal_str(arguments.get("trigger_rule")) + policy = _policy(arguments) + concept = _operator_concept(operator, capture_id in visitor.mapped) + reads, writes = _operator_assets(operator, arguments, asset_definitions=visitor.asset_definitions) + _add_operator_control_properties(operator, arguments, properties) + common: _SourceNodeArguments = { + "source_id": capture_id, + "task_key": task_key, + "source": SOURCE_AIRFLOW, + "name": task_id, + "native_type": operator, + "dependencies": dependencies, + "run_condition": run_condition, + "policy": policy, + "data_reads": reads, + "data_writes": writes, + "properties": properties, + "raw": raw, + } + if capture_id in visitor.mapped: + properties["mapping"] = { + "expand_arguments": list(visitor.expand_kwargs.get(capture_id, [])), + "has_partial": capture_id in visitor.partial_mapped, + } + return ContainerNode(concept=CONCEPT_LOOP, branches={"body": []}, **common) + + activity = activities.get(task_key) + if isinstance(activity, PlaceholderActivity) or concept == CONCEPT_GAP: + reason = activity.comment if isinstance(activity, PlaceholderActivity) else None + return GapNode(reason=reason or f"Airflow operator {operator!r} has no deterministic mapping.", **common) + return SourceNode(concept=concept, **common) + + +def _add_operator_control_properties( + operator: str, + arguments: dict[str, ast.expr], + properties: dict[str, Any], +) -> None: + """Adds cross-workflow and external-wait metadata for control operators.""" + if operator == "TriggerDagRunOperator": + properties[INVOKES_WORKFLOW_PROPERTY] = ops.literal_str(arguments.get("trigger_dag_id")) or "" + wait_node = arguments.get("wait_for_completion") + wait_value = ops.literal_value(wait_node) + properties[INVOKES_WAIT_PROPERTY] = ( + False if wait_node is None else wait_value if isinstance(wait_value, bool) else None ) + return + if operator in {"DatabricksRunNowOperator", "DatabricksRunNowDeferrableOperator"}: + job_id = ops.literal_value(arguments.get("job_id")) + properties[INVOKES_WORKFLOW_PROPERTY] = f"databricks-job:{job_id}" if job_id is not None else "" + return + if operator in {"ExternalTaskSensor", "ExternalTaskSensorAsync"}: + properties["external_workflow_wait"] = { + "dag_id": ops.literal_str(arguments.get("external_dag_id")), + "task_id": ops.literal_str(arguments.get("external_task_id")), + } + + +def _taskflow_node( + *, + capture_id: str, + task_key: str, + dependencies: list[SourceDependency], + task_keys: dict[str, str], + visitor: DagVisitor, + source: str, + activities: dict[str, Activity], + properties: dict[str, Any], + raw: dict[str, Any], +) -> SourceNode: + """Builds a source node for one TaskFlow invocation.""" + task = visitor.taskflow_tasks[capture_id] + raw.update( + { + "decorator": task.decorator, + "callable": task.def_name, + "source_reference": task.source_reference, + "positional_arguments": dict(task.positional_values), + "keyword_arguments": dict(task.keyword_values), + "unresolved_arguments": list(task.unresolved_arguments), + } + ) + taskflow_definition = visitor.taskflow_defs[task.def_name][0] + raw["callable_definition"] = _definition_payload( + taskflow_definition, + source, + closure_source=callable_notebook.render_source_closure(taskflow_definition, source), + ) + data_upstreams = [ + *[task.positional_deps[position] for position in sorted(task.positional_deps)], + *task.keyword_deps.values(), + *task.mapped_deps, + ] + reads = [ + DataAsset(signature=f"xcom:{task_keys[upstream]}", asset_type="value") + for upstream in dict.fromkeys(data_upstreams) + if upstream in task_keys + ] + writes = [DataAsset(signature=f"xcom:{task_key}", asset_type="value")] + common: _SourceNodeArguments = { + "source_id": capture_id, + "task_key": task_key, + "source": SOURCE_AIRFLOW, + "name": task.task_id, + "native_type": task.decorator, + "dependencies": dependencies, + "run_condition": None, + "policy": None, + "data_reads": reads, + "data_writes": writes, + "properties": properties, + "raw": raw, + } + if capture_id in visitor.mapped: + properties["mapping"] = { + "expand_argument": task.expand_kwarg, + "expand_items_json": task.expand_items_json, + } + return ContainerNode(concept=CONCEPT_LOOP, branches={"body": []}, **common) + activity = activities.get(task_key) + if isinstance(activity, PlaceholderActivity): + return GapNode(reason=activity.comment or "TaskFlow invocation requires manual migration.", **common) + return SourceNode(concept=CONCEPT_SCRIPT, **common) + + +def _taskgroup_node( + *, + capture_id: str, + task_key: str, + dependencies: list[SourceDependency], + visitor: DagVisitor, + source: str, + properties: dict[str, Any], + raw: dict[str, Any], +) -> ContainerNode: + """Builds a source container for one decorated task-group invocation.""" task_id, definition, mapped = visitor.taskgroup_calls[capture_id] raw.update({"task_group_callable": definition, "mapped": mapped}) - taskgroup_definition = visitor.taskgroup_defs.get(definition) + taskgroup_definition = visitor.taskgroup_definition_for(capture_id) if taskgroup_definition is not None: - raw["callable_definition"] = _definition_payload(taskgroup_definition, source) + raw["callable_definition"] = _definition_payload( + taskgroup_definition, + source, + closure_source=callable_notebook.render_source_closure(taskgroup_definition, source), + ) return ContainerNode( source_id=capture_id, task_key=task_key, @@ -426,11 +547,13 @@ def _captured_node( def _nest_task_groups( nodes_by_capture: dict[str, SourceNode], - visitor: _DagVisitor, + group_paths: dict[str, str], + visitor: DagVisitor, source: str, ) -> list[SourceNode]: roots: list[SourceNode] = [] containers: dict[str, ContainerNode] = {} + used_task_keys = {node.task_key for node in nodes_by_capture.values()} def container(path: str) -> ContainerNode: existing = containers.get(path) @@ -449,9 +572,10 @@ def container(path: str) -> ContainerNode: "end_column": getattr(group_node, "end_col_offset", 0), }, } + task_key = _allocate_structural_task_key(path, used_task_keys) created = ContainerNode( source_id=f"task-group:{path}", - task_key=path, + task_key=task_key, concept=CONCEPT_GROUP, source=SOURCE_AIRFLOW, name=leaf, @@ -468,7 +592,7 @@ def container(path: str) -> ContainerNode: return created for capture_id, node in nodes_by_capture.items(): - group_path = visitor.groups.get(capture_id) + group_path = group_paths.get(capture_id) if group_path: container(group_path).branches["group"].append(node) else: @@ -476,18 +600,62 @@ def container(path: str) -> ContainerNode: return roots -def _gap_nodes(visitor: _DagVisitor, source: str) -> list[GapNode]: - entries: list[tuple[str, ast.AST, str]] = [] - entries.extend( - ("unclaimed_task_call", node, "Airflow task call was not captured by the static subset.") - for node in visitor.unclaimed_task_calls +def _capture_ids_in_source_order( + task_keys: dict[str, str], + source_nodes: dict[str, ast.Call], +) -> list[str]: + capture_indexes = {capture_id: index for index, capture_id in enumerate(source_nodes)} + return sorted( + task_keys, + key=lambda capture_id: ( + getattr(source_nodes.get(capture_id), "lineno", 0), + getattr(source_nodes.get(capture_id), "col_offset", 0), + capture_indexes.get(capture_id, len(capture_indexes)), + ), + ) + + +def _source_node_sort_key(node: SourceNode) -> tuple[int, int, int, int]: + span = (node.raw or {}).get("source_span", {}) + return ( + int(span.get("line", 0)), + int(span.get("column", 0)), + int(span.get("end_line", 0)), + int(span.get("end_column", 0)), ) + + +def _allocate_structural_task_key(path: str, used_task_keys: set[str]) -> str: + base = path + candidate = base + suffix = 2 + while candidate in used_task_keys: + candidate = f"{base}__{suffix}" + suffix += 1 + used_task_keys.add(candidate) + return candidate + + +def _gap_nodes(visitor: DagVisitor, source: str) -> tuple[dict[str, GapNode], dict[str, str]]: + entries: list[tuple[str, ast.AST, str]] = [] entries.extend( ("unclaimed_statement", node, "DAG-body statement was not claimed by the static subset.") for node in visitor.unclaimed_statements ) + calls_with_unclaimed_statements = { + id(candidate) + for statement in visitor.unclaimed_statements + for candidate in ast.walk(statement) + if isinstance(candidate, ast.Call) + } + entries.extend( + ("unclaimed_task_call", node, "Airflow task call was not captured by the static subset.") + for node in visitor.unclaimed_task_calls + if id(node) not in calls_with_unclaimed_statements + ) entries.extend(("unresolved_construct", node, reason) for reason, node in visitor.unresolved_constructs) - gaps: list[GapNode] = [] + gaps: dict[str, GapNode] = {} + group_paths: dict[str, str] = {} seen: set[tuple[str, int, int, int, int]] = set() for code, node, reason in entries: key = ( @@ -501,27 +669,28 @@ def _gap_nodes(visitor: _DagVisitor, source: str) -> list[GapNode]: continue seen.add(key) source_id = f"{code}:{key[1]}:{key[2]}:{len(gaps) + 1}" - gaps.append( - GapNode( - source_id=source_id, - task_key=f"__flowx_gap_{key[1]}_{key[2]}_{len(gaps) + 1}", - source=SOURCE_AIRFLOW, - name=code, - native_type=type(node).__name__, - reason=reason, - properties={"strategy": "unsupported"}, - raw={ - "source": ast.get_source_segment(source, node) or ast.unparse(node), - "source_span": { - "line": key[1], - "column": key[2], - "end_line": key[3], - "end_column": key[4], - }, + gaps[source_id] = GapNode( + source_id=source_id, + task_key=f"__flowx_gap_{key[1]}_{key[2]}_{len(gaps) + 1}", + source=SOURCE_AIRFLOW, + name=code, + native_type=type(node).__name__, + reason=reason, + properties={"strategy": "unsupported"}, + raw={ + "source": ast.get_source_segment(source, node) or ast.unparse(node), + "source_span": { + "line": key[1], + "column": key[2], + "end_line": key[3], + "end_column": key[4], }, - ) + }, ) - return gaps + group_path = visitor.gap_group_paths.get(id(node)) + if group_path is not None: + group_paths[source_id] = group_path + return gaps, group_paths def _operator_concept(operator: str, mapped: bool) -> str: @@ -542,9 +711,14 @@ def _operator_concept(operator: str, mapped: bool) -> str: return CONCEPT_GAP -def _operator_assets(operator: str, kwargs: dict[str, ast.expr]) -> tuple[list[DataAsset], list[DataAsset]]: - reads = _declared_assets(kwargs.get("inlets"), direction="read") - writes = _declared_assets(kwargs.get("outlets"), direction="write") +def _operator_assets( + operator: str, + kwargs: dict[str, ast.expr], + *, + asset_definitions: dict[str, ast.Call], +) -> tuple[list[DataAsset], list[DataAsset]]: + reads = _declared_assets(kwargs.get("inlets"), direction="read", asset_definitions=asset_definitions) + writes = _declared_assets(kwargs.get("outlets"), direction="write", asset_definitions=asset_definitions) if operator in ops.FILE_SENSORS: path = ops.file_sensor_path(kwargs) if path: @@ -566,17 +740,20 @@ def _operator_assets(operator: str, kwargs: dict[str, ast.expr]) -> tuple[list[D return _deduplicate_assets(reads), _deduplicate_assets(writes) -def _declared_assets(node: ast.expr | None, *, direction: str) -> list[DataAsset]: +def _declared_assets( + node: ast.expr | None, + *, + direction: str, + asset_definitions: dict[str, ast.Call], +) -> list[DataAsset]: if node is None: return [] candidates = list(node.elts) if isinstance(node, (ast.List, ast.Tuple, ast.Set)) else [node] assets: list[DataAsset] = [] for candidate in candidates: - value: str | None = None - if isinstance(candidate, ast.Call) and candidate.args: - value = ops.literal_str(candidate.args[0]) - else: - value = ops.literal_str(candidate) + if isinstance(candidate, ast.Name): + candidate = asset_definitions.get(candidate.id, candidate) + value = _asset_uri(candidate) if value is not None: assets.append( DataAsset( @@ -589,6 +766,15 @@ def _declared_assets(node: ast.expr | None, *, direction: str) -> list[DataAsset return assets +def _asset_uri(node: ast.expr) -> str | None: + if not isinstance(node, ast.Call): + return ops.literal_str(node) + if node.args: + return ops.literal_str(node.args[0]) + uri = next((keyword.value for keyword in node.keywords if keyword.arg == "uri"), None) + return ops.literal_str(uri) + + def _deduplicate_assets(assets: Iterable[DataAsset]) -> list[DataAsset]: result: list[DataAsset] = [] seen: set[tuple[str, str | None, str | None]] = set() @@ -643,7 +829,7 @@ def _policy(arguments: dict[str, ast.expr]) -> PolicySpec | None: ) -def _source_schedule(visitor: _DagVisitor, source: str) -> ScheduleSpec | None: +def _source_schedule(visitor: DagVisitor, source: str) -> ScheduleSpec | None: node = visitor.schedule_node if node is None or (isinstance(node, ast.Constant) and node.value is None): return None @@ -651,7 +837,7 @@ def _source_schedule(visitor: _DagVisitor, source: str) -> ScheduleSpec | None: source_expression = ast.get_source_segment(source, node) or ast.unparse(node) expression = source_expression if literal is None else _json_safe(literal, fallback=source_expression) kind = "schedule" if isinstance(literal, str) else "interval" - is_asset_expression = _asset_expression(node, visitor._aliases, visitor.asset_definitions) is not None + is_asset_expression = asset_expression(node, visitor._aliases, visitor.asset_definitions) is not None if is_asset_expression: kind = "asset" return ScheduleSpec( @@ -670,8 +856,8 @@ def _expression_payload(node: ast.expr, source: str) -> dict[str, Any]: } -def _definition_payload(node: ast.AST, source: str) -> dict[str, Any]: - return { +def _definition_payload(node: ast.AST, source: str, *, closure_source: str | None = None) -> dict[str, Any]: + payload = { "source": ast.get_source_segment(source, node) or ast.unparse(node), "source_span": { "line": getattr(node, "lineno", 0), @@ -680,6 +866,9 @@ def _definition_payload(node: ast.AST, source: str) -> dict[str, Any]: "end_column": getattr(node, "end_col_offset", getattr(node, "col_offset", 0)), }, } + if closure_source is not None: + payload["closure_source"] = closure_source + return payload def _json_safe(value: Any, *, fallback: Any) -> Any: diff --git a/src/flowx/sources/airflow/loader/api.py b/src/flowx/sources/airflow/loader/api.py index 7041c98..dee1860 100644 --- a/src/flowx/sources/airflow/loader/api.py +++ b/src/flowx/sources/airflow/loader/api.py @@ -27,7 +27,7 @@ _top_level_dag_declarations, ) from flowx.sources.airflow.loader.lowering import _load_airflow_module -from flowx.sources.airflow.loader.visitor import _DagVisitor +from flowx.sources.airflow.loader.visitor import DagVisitor from flowx.utils import normalize_task_key _HOST_PATTERN = re.compile(r"https://([A-Za-z0-9._-]*(?:azuredatabricks\.net|databricks\.com|cloud\.databricks\.com))") @@ -71,7 +71,16 @@ def load_airflow_dag_results( dbt_mode: str = "static", source_file: str | None = None, ) -> list[AirflowDiscoveryResult]: - """Parses every DAG into current IR and its shared source graph in one capture pass.""" + """Parses every DAG into current IR and its shared source graph in one capture pass. + + Args: + dag_path: Python file containing one or more Airflow DAG declarations. + dbt_mode: dbt-factory render mode, either ``"static"`` or ``"pydabs"``. + source_file: Stable source label stored in discovery metadata. Defaults to the file name. + + Returns: + One paired pipeline and source graph for each DAG declaration in the file. + """ results = _load_airflow_dag_results( dag_path, dbt_mode=dbt_mode, @@ -115,7 +124,7 @@ def _load_airflow_dag_results( continue isolated = _module_for_dag(module, declaration, declarations) audit = source_audit.audit_module(isolated, target_dag_variable=declaration.target_dag_variable) - visitor = _DagVisitor(isolated, target_dag_variable=declaration.target_dag_variable) + visitor = DagVisitor(isolated, target_dag_variable=declaration.target_dag_variable) visitor.visit(isolated) pipeline = _load_airflow_module( dag_path, @@ -186,7 +195,17 @@ def load_discovery_results( dbt_mode: str = "static", exclude_dags: set[str] | None = None, ) -> list[AirflowDiscoveryResult]: - """Loads Airflow DAGs into paired Pipeline IR and shared source graphs.""" + """Loads Airflow DAGs into paired Pipeline IR and shared source graphs. + + Args: + source_path: A DAG Python file or directory searched recursively for DAG files. + pipeline: Optional DAG identifier used to select one pipeline. + dbt_mode: dbt-factory render mode, either ``"static"`` or ``"pydabs"``. + exclude_dags: DAG identifiers retained in reporting but excluded from migration output. + + Returns: + Paired pipeline and source-graph results after filtering and exclusion rewrites. + """ root = source_path if source_path.is_dir() else source_path.parent results = [ result diff --git a/src/flowx/sources/airflow/loader/captures.py b/src/flowx/sources/airflow/loader/captures.py index af9d160..9d36b39 100644 --- a/src/flowx/sources/airflow/loader/captures.py +++ b/src/flowx/sources/airflow/loader/captures.py @@ -77,6 +77,7 @@ class _TaskFlowTask: is_async: bool = False positional_deps: dict[int, str] = field(default_factory=dict) keyword_deps: dict[str, str] = field(default_factory=dict) + mapped_deps: list[str] = field(default_factory=list) positional_values: dict[int, str] = field(default_factory=dict) keyword_values: dict[str, str] = field(default_factory=dict) unresolved_arguments: list[str] = field(default_factory=list) diff --git a/src/flowx/sources/airflow/loader/graph.py b/src/flowx/sources/airflow/loader/graph.py index a815d8a..e3e83b4 100644 --- a/src/flowx/sources/airflow/loader/graph.py +++ b/src/flowx/sources/airflow/loader/graph.py @@ -3,7 +3,6 @@ from __future__ import annotations import ast -from collections.abc import Mapping from flowx.sources.airflow import operators as ops from flowx.sources.airflow.loader.ast_utils import _sanitize_task_key @@ -14,27 +13,15 @@ def _allocate_task_keys( taskflow_task_ids: dict[str, str], taskgroup_task_ids: dict[str, str], groups: dict[str, str], - capture_source_nodes: Mapping[str, ast.AST], ) -> dict[str, str]: - """Allocates stable, collision-free task keys in capture order.""" + """Allocates stable, collision-free task keys using the legacy lowering order.""" task_ids = {variable: task_id for variable, (task_id, _, _) in operators.items()} task_ids.update(taskflow_task_ids) task_ids.update(taskgroup_task_ids) - capture_indexes = {capture_id: index for index, capture_id in enumerate(capture_source_nodes)} - capture_ids = sorted( - task_ids, - key=lambda capture_id: ( - getattr(capture_source_nodes.get(capture_id), "lineno", 0), - getattr(capture_source_nodes.get(capture_id), "col_offset", 0), - capture_indexes.get(capture_id, len(capture_indexes)), - ), - ) - allocated: dict[str, str] = {} used: set[str] = set() - for variable in capture_ids: - task_id = task_ids[variable] + for variable, task_id in task_ids.items(): base = _sanitize_task_key(task_id) if variable in groups: base = f"{groups[variable]}__{base}" @@ -191,3 +178,41 @@ def _trigger_from_sensor(operator: str, kwargs: dict[str, ast.expr]) -> dict[str "pause_status": "UNPAUSED", } return None + + +def allocate_task_keys( + operators: dict[str, tuple[str, str, dict[str, ast.expr]]], + taskflow_task_ids: dict[str, str], + taskgroup_task_ids: dict[str, str], + groups: dict[str, str], +) -> dict[str, str]: + """Allocates stable, collision-free keys for all executable task kinds. + + Args: + operators: Classic operator captures keyed by source identity. + taskflow_task_ids: TaskFlow task identifiers keyed by source identity. + taskgroup_task_ids: Decorated task-group identifiers keyed by source identity. + groups: Enclosing TaskGroup paths keyed by source identity. + + Returns: + Allocated task keys keyed by source identity. + """ + return _allocate_task_keys(operators, taskflow_task_ids, taskgroup_task_ids, groups) + + +def expand_group_edges( + edges: list[tuple[str, str]], + groups: dict[str, str], + group_vars: dict[str, str], +) -> list[tuple[str, str]]: + """Expands TaskGroup dependency endpoints to their boundary tasks. + + Args: + edges: Dependency edges expressed in source identities. + groups: Enclosing TaskGroup paths keyed by task source identity. + group_vars: TaskGroup paths keyed by context-manager variables. + + Returns: + Dependency edges whose TaskGroup endpoints have been replaced by boundary tasks. + """ + return _expand_group_edges(edges, groups, group_vars) diff --git a/src/flowx/sources/airflow/loader/lowering.py b/src/flowx/sources/airflow/loader/lowering.py index 1021e80..fee94bb 100644 --- a/src/flowx/sources/airflow/loader/lowering.py +++ b/src/flowx/sources/airflow/loader/lowering.py @@ -26,17 +26,17 @@ from flowx.sources.airflow.loader.ast_utils import _span from flowx.sources.airflow.loader.dbt import _build_dbt_factory from flowx.sources.airflow.loader.graph import ( - _allocate_task_keys, - _expand_group_edges, _rewire_dropped, _root_trigger_sensor, _trigger_from_sensor, + allocate_task_keys, + expand_group_edges, ) from flowx.sources.airflow.loader.policy import _job_email_notifications, _job_timeout_seconds from flowx.sources.airflow.loader.reconcile import _iter_placeholders, _semantic_finding from flowx.sources.airflow.loader.schedule import _asset_schedule_from_node, _schedule_from_interval from flowx.sources.airflow.loader.taskflow import _build_taskflow_task, _wrap_in_for_each, _wrap_taskflow_in_for_each -from flowx.sources.airflow.loader.visitor import _DagVisitor +from flowx.sources.airflow.loader.visitor import DagVisitor _DATABRICKS_JOB_TAG_LIMIT = 25 @@ -50,7 +50,7 @@ def _load_airflow_module( target_dag_variable: str | None = None, source_file: str | None = None, captured_audit: SourceAudit | None = None, - captured_visitor: _DagVisitor | None = None, + captured_visitor: DagVisitor | None = None, ) -> Pipeline: """Parses one isolated DAG declaration into a flowx Pipeline IR. @@ -68,7 +68,7 @@ def _load_airflow_module( PlaceholderActivity. """ audit = captured_audit or source_audit.audit_module(module, target_dag_variable=target_dag_variable) - visitor = captured_visitor or _DagVisitor(module, target_dag_variable=target_dag_variable) + visitor = captured_visitor or DagVisitor(module, target_dag_variable=target_dag_variable) if captured_visitor is None: visitor.visit(module) functions = visitor.functions() @@ -78,18 +78,17 @@ def _load_airflow_module( var_task_ids: dict[str, str] = {var: tid for var, (tid, _, _) in visitor.operators.items()} var_task_ids.update({var: tf.task_id for var, tf in visitor.taskflow_tasks.items()}) var_task_ids.update({var: task_id for var, (task_id, _, _) in visitor.taskgroup_calls.items()}) - var_to_task_key = _allocate_task_keys( + var_to_task_key = allocate_task_keys( visitor.operators, {var: task.task_id for var, task in visitor.taskflow_tasks.items()}, {var: task_id for var, (task_id, _, _) in visitor.taskgroup_calls.items()}, visitor.groups, - visitor.capture_source_nodes, ) # Expand group-level edges (`group_a >> group_b`, `task >> group`, ...) into edges between the # groups' boundary tasks: leaves of the upstream group -> roots of the downstream group, matching # Airflow's TaskGroup dependency semantics. A non-group var resolves to itself. - edges = _expand_group_edges(visitor.edges, visitor.groups, visitor.group_vars) + edges = expand_group_edges(visitor.edges, visitor.groups, visitor.group_vars) # Build the upstream adjacency in dependency terms, then drop structural nodes # (Dummy/Empty and lifted root sensors) by rewiring their downstreams to their upstreams. @@ -128,7 +127,7 @@ def _load_airflow_module( }, } elif schedule_gap is not None: - visitor.unresolved_constructs.append((schedule_gap, schedule_node)) + visitor.add_unresolved_construct(schedule_gap, schedule_node) has_schedule = schedule is not None # Dummy/Empty operators are structural and can be removed after dependency rewiring. @@ -206,12 +205,10 @@ def _dep(upstream_var: str, outcome: str | None) -> str: return dbt_key_remap.get(key, key) tasks: list[Activity] = [] - task_capture_ids: dict[int, str] = {} placeholder_capture_ids: dict[int, str] = {} helper_expansion_ids = {str(item["capture_id"]) for item in visitor.helper_expansions} def append_task(activity: Activity, capture_id: str) -> None: - task_capture_ids[id(activity)] = capture_id for placeholder in _iter_placeholders([activity]): placeholder_capture_ids[id(placeholder)] = capture_id tasks.append(activity) @@ -576,9 +573,6 @@ def append_task(activity: Activity, capture_id: str) -> None: placeholder.depends_on = depends_on append_task(placeholder, var) - capture_order = {capture_id: index for index, capture_id in enumerate(var_to_task_key)} - tasks.sort(key=lambda activity: capture_order[task_capture_ids[id(activity)]]) - # Declare every job parameter -- those referenced in templates plus any from the DAG's # params={...} -- each with a default (Databricks requires one): the params={...} default when # present; a reserved logical-date parameter its schedule-aware time ref so a native backfill can diff --git a/src/flowx/sources/airflow/loader/policy.py b/src/flowx/sources/airflow/loader/policy.py index bf9b67f..a4a2831 100644 --- a/src/flowx/sources/airflow/loader/policy.py +++ b/src/flowx/sources/airflow/loader/policy.py @@ -6,7 +6,7 @@ from flowx.sources.airflow import operators as ops from flowx.sources.airflow import templating -from flowx.sources.airflow.loader.visitor import _DagVisitor +from flowx.sources.airflow.loader.visitor import DagVisitor _RECOGNIZED_DAG_SETTINGS = frozenset( { @@ -41,11 +41,11 @@ _NON_EXECUTION_DAG_SETTINGS = frozenset({"tags", "description", "doc_md", "dag_display_name", "default_args.owner"}) -def _job_timeout_seconds(visitor: _DagVisitor) -> int | None: +def _job_timeout_seconds(visitor: DagVisitor) -> int | None: return templating.timedelta_seconds(visitor.dag_kwargs.get("dagrun_timeout")) -def _job_email_notifications(visitor: _DagVisitor) -> dict[str, list[str]]: +def _job_email_notifications(visitor: DagVisitor) -> dict[str, list[str]]: recipients = templating.literal_email_recipients(visitor.default_args.get("email")) on_failure = visitor.default_args.get("email_on_failure") failure_enabled = on_failure is None or (isinstance(on_failure, ast.Constant) and on_failure.value is True) @@ -54,7 +54,7 @@ def _job_email_notifications(visitor: _DagVisitor) -> dict[str, list[str]]: return {} -def _retry_email_is_active(visitor: _DagVisitor) -> bool: +def _retry_email_is_active(visitor: DagVisitor) -> bool: """Returns whether any captured task can emit an Airflow retry email.""" for _, _, kwargs in visitor.operators.values(): retry_node = kwargs.get("retries", visitor.default_args.get("retries")) @@ -70,7 +70,7 @@ def _retry_email_is_active(visitor: _DagVisitor) -> bool: return False -def _dag_setting_disposition(name: str, visitor: _DagVisitor) -> dict[str, str] | None: +def _dag_setting_disposition(name: str, visitor: DagVisitor) -> dict[str, str] | None: """Classifies recognized DAG settings as mapped, intentional no-ops, or runtime gaps.""" if name not in _RECOGNIZED_DAG_SETTINGS: return { @@ -239,3 +239,15 @@ def _dag_setting_disposition(name: str, visitor: _DagVisitor) -> dict[str, str] "rationale": "preserved_as_databricks_job_failure_notification", } return None + + +def job_timeout_seconds(visitor: DagVisitor) -> int | None: + """Returns the statically resolved DAG run timeout in seconds. + + Args: + visitor: Completed capture containing the DAG arguments. + + Returns: + The DAG run timeout in seconds, or ``None`` when it cannot be resolved. + """ + return _job_timeout_seconds(visitor) diff --git a/src/flowx/sources/airflow/loader/reconcile.py b/src/flowx/sources/airflow/loader/reconcile.py index 4a0e744..2384fd1 100644 --- a/src/flowx/sources/airflow/loader/reconcile.py +++ b/src/flowx/sources/airflow/loader/reconcile.py @@ -21,7 +21,7 @@ _RECOGNIZED_DAG_SETTINGS, _dag_setting_disposition, ) -from flowx.sources.airflow.loader.visitor import _DagVisitor +from flowx.sources.airflow.loader.visitor import DagVisitor def _semantic_finding( @@ -78,7 +78,7 @@ def _reconcile_pipeline( pipeline: Pipeline, *, audit: source_audit.SourceAudit, - visitor: _DagVisitor, + visitor: DagVisitor, source_file: str, var_to_task_key: dict[str, str], dropped: set[str], diff --git a/src/flowx/sources/airflow/loader/schedule.py b/src/flowx/sources/airflow/loader/schedule.py index b67f0b8..732a37e 100644 --- a/src/flowx/sources/airflow/loader/schedule.py +++ b/src/flowx/sources/airflow/loader/schedule.py @@ -285,3 +285,21 @@ def _asset_schedule_from_node( }, None, ) + + +def asset_expression( + node: ast.expr, + aliases: dict[str, str], + definitions: dict[str, ast.Call], +) -> tuple[list[str], str, str | None] | None: + """Resolves an Airflow Asset or Dataset schedule expression. + + Args: + node: Schedule expression to resolve. + aliases: Canonical import names keyed by local aliases. + definitions: Named Asset or Dataset constructors keyed by variable name. + + Returns: + Table names, combination mode, and an optional error code, or ``None`` when the expression is not asset-based. + """ + return _asset_expression(node, aliases, definitions) diff --git a/src/flowx/sources/airflow/loader/visitor.py b/src/flowx/sources/airflow/loader/visitor.py index c02252c..9297edf 100644 --- a/src/flowx/sources/airflow/loader/visitor.py +++ b/src/flowx/sources/airflow/loader/visitor.py @@ -43,8 +43,20 @@ _EDGE_MODIFIER_CONSTRUCTS = frozenset({"Label"}) -class _DagVisitor(ast.NodeVisitor): - """Collects operator calls, dependency edges, and the DAG's schedule.""" +class DagVisitor(ast.NodeVisitor): + """Collects source-faithful tasks, dependencies, settings, and unsupported constructs. + + Args: + module: Isolated Python module containing one DAG declaration. + target_dag_variable: Assigned DAG variable selected from a multi-DAG module. + + Attributes: + task_captures: Classic operator captures keyed by stable source identity. + edge_captures: Dependency declarations retained with source spans. + taskflow_tasks: TaskFlow invocations keyed by stable source identity. + taskgroup_calls: Decorated task-group invocations keyed by stable source identity. + unresolved_constructs: Unsupported constructs paired with their source nodes. + """ def __init__(self, module: ast.Module, *, target_dag_variable: str | None = None) -> None: self._aliases = _import_aliases(module) @@ -70,6 +82,7 @@ def __init__(self, module: ast.Module, *, target_dag_variable: str | None = None self.unclaimed_task_calls: list[ast.Call] = [] self.unclaimed_statements: list[ast.stmt] = [] self.unresolved_constructs: list[tuple[str, ast.AST]] = [] + self.gap_group_paths: dict[int, str] = {} self._claimed_task_call_ids: set[int] = set() self._claimed_statement_ids: set[int] = set() self._dag_scope_depth = 0 @@ -114,7 +127,8 @@ def __init__(self, module: ast.Module, *, target_dag_variable: str | None = None self.taskflow_defs: dict[str, tuple[ast.FunctionDef | ast.AsyncFunctionDef, str]] = {} # @task_group def names -- a group is a sub-pipeline, not a single renderable task, so an # invocation routes to a placeholder + gap rather than being expanded here. - self.taskgroup_defs: dict[str, ast.FunctionDef | ast.AsyncFunctionDef] = {} + self.taskgroup_defs: set[str] = set() + self.taskgroup_definitions_by_capture: dict[str, ast.FunctionDef | ast.AsyncFunctionDef] = {} for fn in _iter_functions(module): decorator = next( ( @@ -127,7 +141,7 @@ def __init__(self, module: ast.Module, *, target_dag_variable: str | None = None if decorator is not None: self.taskflow_defs[fn.name] = (fn, decorator) elif _has_decorator(fn, _TASK_GROUP_DECORATORS, self._aliases): - self.taskgroup_defs[fn.name] = fn + self.taskgroup_defs.add(fn.name) # TaskFlow task instances: var name -> _TaskFlowTask (id, def-name, decorator, arg bindings). self.taskflow_tasks: dict[str, _TaskFlowTask] = {} # @task_group invocations: var name -> (task_id, def-name, is_mapped). @@ -162,6 +176,19 @@ def resolved_callable_for(self, task_var: str) -> ast.FunctionDef | None: resolved = self._resolved_callables.get(task_var) return resolved[1] if resolved is not None else None + def taskgroup_definition_for(self, task_var: str) -> ast.FunctionDef | ast.AsyncFunctionDef | None: + """Returns the lexical task-group definition resolved for an invocation.""" + return self.taskgroup_definitions_by_capture.get(task_var) + + def add_unresolved_construct(self, reason: str, node: ast.AST) -> None: + """Records an unsupported construct and its enclosing TaskGroup path.""" + self.unresolved_constructs.append((reason, node)) + self._record_gap_group(node) + + def _record_gap_group(self, node: ast.AST) -> None: + if self._group_stack: + self.gap_group_paths[id(node)] = "__".join(self._group_stack) + def visit_FunctionDef(self, node: ast.FunctionDef) -> None: # A @task- or @task_group-decorated function defines a task / sub-pipeline from its body, # which is internal logic rather than DAG structure, so don't descend. @dag marks the @@ -493,6 +520,8 @@ def _register_taskflow_call(self, call: ast.Call, var: str, *, source_reference: for mapped_arg in _mapping_chain_args(call): dep = self._resolve_taskflow_arg(mapped_arg) if dep is not None and dep != var: + if dep not in task.mapped_deps: + task.mapped_deps.append(dep) self._add_edges([dep], [var], call) if self._group_stack: self.groups[var] = "__".join(self._group_stack) @@ -569,16 +598,43 @@ def _register_taskgroup_call(self, call: ast.Call, var: str | None) -> bool: if not (isinstance(func, ast.Name) and func.id in self.taskgroup_defs): return False def_name = func.id + definition = self._resolve_taskgroup_definition(def_name, call) + if definition is None: + return False if var is None: self._taskgroup_counter += 1 var = f"{def_name}__tg{self._taskgroup_counter}" self.taskgroup_calls[var] = (var, def_name, mapped) + self.taskgroup_definitions_by_capture[var] = definition self.capture_source_nodes[var] = call self._claimed_task_call_ids.add(id(call)) if self._group_stack: self.groups[var] = "__".join(self._group_stack) return True + def _resolve_taskgroup_definition( + self, + name: str, + reference: ast.AST, + ) -> ast.FunctionDef | ast.AsyncFunctionDef | None: + definition = self._resolve_lexical_function(name, reference) + if definition is not None and _has_decorator(definition, _TASK_GROUP_DECORATORS, self._aliases): + return definition + module = self._scope_stack[0] + if any(self._lexical_functions.get(id(scope), {}).get(name) for scope in self._scope_stack[1:]): + return None + events = self._lexical_functions.get(id(module), {}).get(name, []) + line = getattr(reference, "lineno", 0) + if any(event_line <= line for event_line, _conditional, _candidate in events): + return None + if any(conditional for _event_line, conditional, _candidate in events): + return None + candidates = [candidate for _line, conditional, candidate in events if not conditional] + candidate = candidates[-1] if candidates else None + if candidate is not None and _has_decorator(candidate, _TASK_GROUP_DECORATORS, self._aliases): + return candidate + return None + def _resolve_taskflow_arg(self, arg: ast.expr) -> str | None: """Returns the upstream task var an argument refers to, else None (a literal / unknown). @@ -628,10 +684,7 @@ def visit_With(self, node: ast.With) -> None: # `with DbtTaskGroup(...) as g:` — a cosmos group bound to a name. if isinstance(item.optional_vars, ast.Name): var = item.optional_vars.id - kwargs = {kw.arg: kw.value for kw in call.keywords if kw.arg} - task_id = ops.literal_str(kwargs.get("group_id")) or var - self.operators[var] = (task_id, construct, kwargs) - self.calls[var] = call + self._register_operator_call(call, var, binding=var) if opens_dag_scope: self._dag_scope_depth += 1 try: @@ -667,6 +720,7 @@ def _visit_dag_statement(self, statement: ast.stmt) -> None: self._claimed_statement_ids.add(id(statement)) return self.unclaimed_statements.append(statement) + self._record_gap_group(statement) def visit_Call(self, node: ast.Call) -> None: """Fails closed when a task-producing call in a DAG scope was not captured.""" @@ -700,6 +754,7 @@ def visit_Call(self, node: ast.Call) -> None: or self._helper_factory_return(node) ): self.unclaimed_task_calls.append(node) + self._record_gap_group(node) self.generic_visit(node) def _helper_targets_assigned_dag(self, call: ast.Call) -> bool: @@ -736,7 +791,7 @@ def _read_dag_kwargs(self, call: ast.Call) -> None: self.dag_id = ops.literal_str(kwargs.get("dag_id")) or positional_dag_id self._apply_dag_kwargs(kwargs) if self.airflow_generation == "1.10" and not {"schedule", "schedule_interval"} & kwargs.keys(): - self.unresolved_constructs.append(("ambiguous_airflow_1_10_default_schedule", call)) + self.add_unresolved_construct("ambiguous_airflow_1_10_default_schedule", call) def _apply_dag_kwargs(self, kwargs: dict[str, ast.expr]) -> None: self.dag_kwargs.update(kwargs) @@ -769,7 +824,7 @@ def _apply_dag_kwargs(self, kwargs: dict[str, ast.expr]) -> None: for key, val in zip(params.keys, params.values): if isinstance(key, ast.Constant) and isinstance(key.value, str): if key.value.startswith(templating.FLOWX_INTERNAL_PARAMETER_PREFIX): - self.unresolved_constructs.append(("reserved_airflow_parameter_name", key)) + self.add_unresolved_construct("reserved_airflow_parameter_name", key) continue self.dag_params[key.value] = _param_default(val) @@ -839,7 +894,7 @@ def visit_Expr(self, node: ast.Expr) -> None: def visit_For(self, node: ast.For) -> None: """Executes bounded literal/range loops with Python name rebinding semantics.""" if not isinstance(node.target, ast.Name): - self.unresolved_constructs.append(("dynamic_loop_target", node)) + self.add_unresolved_construct("dynamic_loop_target", node) self._claimed_statement_ids.add(id(node)) return items = _static_iteration_nodes(node.iter, self._constants) @@ -849,7 +904,7 @@ def visit_For(self, node: ast.For) -> None: if isinstance(node.iter, (ast.List, ast.Tuple)): items = list(node.iter.elts) else: - self.unresolved_constructs.append(("dynamic_loop_iterable", node)) + self.add_unresolved_construct("dynamic_loop_iterable", node) self._claimed_statement_ids.add(id(node)) return for item in items: @@ -860,7 +915,7 @@ def visit_For(self, node: ast.For) -> None: else: value = _safe_static_value(item, self._constants) if value is _UNRESOLVED: - self.unresolved_constructs.append(("dynamic_loop_value", item)) + self.add_unresolved_construct("dynamic_loop_value", item) self._claimed_statement_ids.add(id(node)) return self._constants[node.target.id] = value @@ -877,7 +932,7 @@ def visit_If(self, node: ast.If) -> None: if value is _UNRESOLVED and isinstance(node.test, ast.Name) and node.test.id in self._task_bindings: value = True if value is _UNRESOLVED: - self.unresolved_constructs.append(("ambiguous_condition", node)) + self.add_unresolved_construct("ambiguous_condition", node) self._claimed_statement_ids.add(id(node)) return branch = node.body if bool(value) else node.orelse @@ -968,3 +1023,6 @@ def _collect_set_dependency(self, call: ast.Call) -> None: self._add_edges(this_names, others, call) elif func.attr == "set_upstream": self._add_edges(others, this_names, call) + + +_DagVisitor = DagVisitor diff --git a/tests/unit/test_airflow_discovery_mapping.py b/tests/unit/test_airflow_discovery_mapping.py index 902e39d..dac6fcf 100644 --- a/tests/unit/test_airflow_discovery_mapping.py +++ b/tests/unit/test_airflow_discovery_mapping.py @@ -271,11 +271,264 @@ def build(): ) assert [(node.source_id, node.task_key) for node in result.graph.tasks] == [ - ("same", "same"), - ("middle", "same__2"), + ("same", "same__2"), + ("middle", "same"), ("last", "last"), ] assert [task.task_key for task in result.pipeline.tasks] == ["same", "same__2", "last"] + assert [task.name for task in result.pipeline.tasks] == ["same", "same", "last"] + + +def test_records_unresolved_cross_workflow_invocations(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow import DAG +from airflow.models import Variable +from airflow.operators.trigger_dagrun import TriggerDagRunOperator +from airflow.providers.databricks.operators.databricks import DatabricksRunNowOperator + +with DAG(dag_id="dynamic_targets", schedule=None) as dag: + trigger = TriggerDagRunOperator(task_id="trigger", trigger_dag_id=Variable.get("CHILD_DAG")) + run_now = DatabricksRunNowOperator(task_id="run_now", job_id=Variable.get("JOB_ID")) +""", + ) + + nodes = {node.task_key: node for node in result.graph.tasks} + assert nodes["trigger"].properties["invokes_workflow"] == "" + assert nodes["run_now"].properties["invokes_workflow"] == "" + assert result.graph.lineage is not None + control_edges = [ + (edge.via_task_key, edge.target_workflow, edge.resolved) for edge in result.graph.lineage.control_edges + ] + assert control_edges == [ + ("trigger", "", False), + ("run_now", "", False), + ] + + +def test_preserves_mapped_taskflow_xcom_lineage(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow.decorators import dag, task + +@dag(dag_id="mapped_xcom", schedule=None) +def build(): + @task + def extract(): + return [1, 2] + + @task + def process(value, fixed=None): + return value + + upstream = extract() + mapped = process.expand(value=upstream) + partial_mapped = process.partial(fixed=upstream).expand(value=[1, 2]) + +build() +""", + ) + + nodes = {node.source_id: node for node in result.graph.tasks} + assert [asset.signature for asset in nodes["mapped"].data_reads] == ["xcom:upstream"] + assert [asset.signature for asset in nodes["partial_mapped"].data_reads] == ["xcom:upstream"] + assert result.graph.lineage is not None + assert {(edge.source_task_key, edge.target_task_key) for edge in result.graph.lineage.data_edges} == { + ("upstream", "mapped"), + ("upstream", "partial_mapped"), + } + + +def test_maps_context_managed_cosmos_task_group(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow import DAG +from cosmos import DbtTaskGroup, ProfileConfig, ProjectConfig + +with DAG(dag_id="cosmos_context", schedule=None) as dag: + with DbtTaskGroup( + group_id="transform", + project_config=ProjectConfig("/opt/dbt"), + profile_config=ProfileConfig(profile_name="analytics", target_name="prod"), + ) as transform: + pass +""", + ) + + transform = next(node for node in walk_nodes(result.graph.tasks) if node.source_id == "transform") + assert transform.native_type == "DbtTaskGroup" + assert transform.raw is not None + assert transform.raw["operator_fqn"] == "cosmos.DbtTaskGroup" + + +def test_resolves_named_task_assets_and_lineage(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow import DAG, Dataset +from airflow.operators.bash import BashOperator +from airflow.sdk import Asset + +orders = Dataset("s3://warehouse/orders") +customers = Asset(uri="s3://warehouse/customers") + +with DAG(dag_id="named_assets", schedule=None) as dag: + produce_orders = BashOperator(task_id="produce_orders", bash_command="echo orders", outlets=[orders]) + consume_orders = BashOperator(task_id="consume_orders", bash_command="echo orders", inlets=[orders]) + produce_customers = BashOperator(task_id="produce_customers", bash_command="echo customers", outlets=[customers]) + consume_customers = BashOperator(task_id="consume_customers", bash_command="echo customers", inlets=[customers]) +""", + ) + + nodes = {node.task_key: node for node in result.graph.tasks} + assert nodes["produce_orders"].data_writes[0].identity == "s3://warehouse/orders" + assert nodes["consume_orders"].data_reads[0].identity == "s3://warehouse/orders" + assert nodes["produce_customers"].data_writes[0].identity == "s3://warehouse/customers" + assert nodes["consume_customers"].data_reads[0].identity == "s3://warehouse/customers" + assert result.graph.lineage is not None + assert {(edge.source_task_key, edge.target_task_key) for edge in result.graph.lineage.data_edges} == { + ("produce_orders", "consume_orders"), + ("produce_customers", "consume_customers"), + } + + +def test_preserves_gap_source_order_and_task_group_scope(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow import DAG +from airflow.operators.bash import BashOperator +from airflow.utils.task_group import TaskGroup + +with DAG(dag_id="ordered_gaps", schedule=None) as dag: + first = BashOperator(task_id="first", bash_command="echo first") + fanout = [BashOperator(task_id=f"work_{index}", bash_command="echo work") for index in range(3)] + last = BashOperator(task_id="last", bash_command="echo last") + with TaskGroup(group_id="nested") as nested: + inner = BashOperator(task_id="inner", bash_command="echo inner") + grouped_fanout = [BashOperator(task_id=f"grouped_{index}", bash_command="echo work") for index in range(2)] +""", + ) + + root_names = [node.name for node in result.graph.tasks] + assert root_names[:3] == ["first", "unclaimed_task_call", "last"] + nested = next(node for node in result.graph.tasks if isinstance(node, ContainerNode) and node.name == "nested") + assert [node.name for node in nested.branches["group"]][:2] == ["inner", "unclaimed_task_call"] + assert any(isinstance(node, GapNode) for node in nested.branches["group"]) + + +def test_persists_complete_callable_dependency_closure(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +import math +from airflow import DAG +from airflow.operators.python import PythonOperator + +SCALE = 3 + +class Multiplier: + def apply(self, value): + return value * SCALE + +def helper(value): + return Multiplier().apply(math.floor(value)) + +def callable_task(): + return helper(2.5) + +with DAG(dag_id="callable_closure", schedule=None) as dag: + run = PythonOperator(task_id="run", python_callable=callable_task) +""", + ) + + callable_definition = result.graph.tasks[0].raw["callable_definition"] + closure = callable_definition["closure_source"] + assert "import math" in closure + assert "SCALE = 3" in closure + assert "class Multiplier:" in closure + assert "def helper(value):" in closure + assert "def callable_task():" in closure + + +def test_allocates_unique_structural_task_group_keys(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow import DAG +from airflow.operators.bash import BashOperator +from airflow.utils.task_group import TaskGroup + +with DAG(dag_id="group_key_collision", schedule=None) as dag: + collision = BashOperator(task_id="processing", bash_command="echo collision") + with TaskGroup(group_id="processing") as processing: + inner = BashOperator(task_id="inner", bash_command="echo inner") +""", + ) + + keys = [node.task_key for node in walk_nodes(result.graph.tasks)] + assert len(keys) == len(set(keys)) + group = next(node for node in result.graph.tasks if isinstance(node, ContainerNode)) + assert group.task_key == "processing__2" + + +def test_resolves_task_group_definitions_lexically(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow.decorators import dag, task_group + +@task_group +def grouped(): + module_marker = "module" + +@dag(dag_id="lexical_groups", schedule=None) +def build(): + @task_group + def grouped(): + nested_marker = "nested" + + selected = grouped() + +build() +""", + ) + + selected = next(node for node in result.graph.tasks if node.source_id == "selected") + definition = selected.raw["callable_definition"] + assert "nested_marker" in definition["source"] + assert "module_marker" not in definition["source"] + + +def test_routes_conditionally_ambiguous_task_group_definition_to_gap(tmp_path: Path) -> None: + result = _load( + tmp_path, + """ +from airflow.decorators import dag, task_group +from airflow.models import Variable + +@task_group +def grouped(): + module_marker = "module" + +@dag(dag_id="ambiguous_group", schedule=None) +def build(): + if Variable.get("USE_LOCAL"): + @task_group + def grouped(): + conditional_marker = "conditional" + + selected = grouped() + +build() +""", + ) + + assert not any(node.source_id == "selected" and node.concept == CONCEPT_GROUP for node in result.graph.tasks) + assert any(isinstance(node, GapNode) for node in result.graph.tasks) def test_unclaimed_comprehension_is_an_explicit_gap(tmp_path: Path) -> None: