Airflow: map discovery captures to shared source graphs

[peterpark-db]

Summary

Closes #63.

Maps Airflow's existing static capture model onto the shared SourceGraph contract introduced by #84 without replacing the Airflow parser or changing the current convert/package path.

• reuses one _DagVisitor and source audit pass for discovery graph projection and existing IR lowering;
• shares collision-safe task-key allocation between discovery and lowering;
• preserves DAG metadata, operator FQNs, source spans/raw arguments, trigger rules, policies, stable task/edge capture identities, TaskFlow calls, TaskGroups, dynamic mapping, findings, transformations, and audit candidates;
• represents unsupported and unclaimed constructs as explicit GapNode values;
• derives source-neutral data/control lineage for TaskFlow values, declared assets, known file/table/SQL arguments, TriggerDagRunOperator, and Databricks run-now operators without treating ExternalTaskSensor as an invocation;
• writes metadata/source_graphs.json through the shared discovery serde;
• projects the Airflow inventory through the source-neutral inventory emitter, then layers the existing Airflow audit/reconciliation/coverage fields back unchanged;
• keeps structural TaskGroup containers out of activity counts while retaining their hierarchy in the full source graph.

This PR is intentionally stacked on #84 (discovery/shared-ast-standard). Issue #86 remains a separate Phase 2 PR that will consume the persisted source graph during conversion.

Verification

• make fmt
• make test: 1,231 passed, 2 skipped
• make integration: 141 passed, 16 deselected, 6 expected xfails
• all 39 DAG declarations in the repository Airflow corpus produce JSON-serializable, serde-round-trippable source graphs
• all 39 current Airflow Pipeline reports are byte-identical between the reused-capture path and a fresh legacy lowering pass

[peterpark-db]

Code review

Found 11 issues:

1. The source-order allocator changes existing mixed classic/TaskFlow collision ownership and therefore renames deployed bundle tasks, despite this PR promising the current conversion path remains unchanged.

flowx/src/flowx/sources/airflow/loader/graph.py

Lines 19 to 39 in 4290cf9

	"""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 in capture_ids:
	task_id = task_ids[variable]
	base = _sanitize_task_key(task_id)
	if variable in groups:

2. Dynamic TriggerDagRun and Databricks RunNow targets omit the invocation marker entirely, so resolvable call sites disappear instead of producing unresolved control-lineage edges.

flowx/src/flowx/sources/airflow/discovery_mapping.py

Lines 262 to 276 in 4290cf9

	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"}:

3. Mapped TaskFlow arguments are not included in XCom reads, so process.expand(value=upstream) keeps a control dependency but loses its data-lineage edge.

flowx/src/flowx/sources/airflow/discovery_mapping.py

Lines 347 to 359 in 4290cf9

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

4. Context-managed Cosmos groups are registered without capture source/task metadata; source-graph discovery then raises KeyError for a supported with DbtTaskGroup(...) form.

flowx/src/flowx/sources/airflow/discovery_mapping.py

Lines 232 to 255 in 4290cf9

	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)

5. Named module-level Asset/Dataset references in inlets/outlets are not resolved, causing empty reads/writes and missing lineage for the standard outlets=[orders] form.

flowx/src/flowx/sources/airflow/discovery_mapping.py

Lines 544 to 589 in 4290cf9

	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

6. Gap nodes are appended after all captured tasks and outside TaskGroup nesting, losing both source order and group scope.

flowx/src/flowx/sources/airflow/discovery_mapping.py

Lines 129 to 141 in 4290cf9

	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()},
	}
	if declaration.factory is not None:
	declaration_raw["factory_definition"] = _definition_payload(declaration.factory, source)

7. Persisted callable definitions omit referenced helpers, classes, constants, and imports, so Phase 2 cannot recreate current callable notebooks from source_graphs.json without reparsing the DAG.

flowx/src/flowx/sources/airflow/discovery_mapping.py

Lines 251 to 258 in 4290cf9

	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)

8. Structural TaskGroup keys bypass collision allocation and can normalize to the same key as an executable task, violating SourceGraph's uniqueness contract.

flowx/src/flowx/sources/airflow/discovery_mapping.py

Lines 440 to 461 in 4290cf9

	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": []},

9. Task-group definitions are retrieved from a module-wide last-writer-wins name map, so a nested same-named group can replace the definition persisted for a different lexical invocation.

flowx/src/flowx/sources/airflow/discovery_mapping.py

Lines 408 to 413 in 4290cf9

	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(

10. New public discovery APIs omit the Google-format Attributes/Args/Returns sections required by src/AGENTS.md.

flowx/src/flowx/sources/airflow/discovery_mapping.py

Lines 67 to 103 in 4290cf9

	@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."""

11. _captured_node is a 200-line multi-branch dispatcher, exceeding the explicit function-shape guidance to extract helpers after roughly 40 lines or three nesting levels.

flowx/src/flowx/sources/airflow/discovery_mapping.py

Lines 222 to 250 in 4290cf9

	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

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}