diff --git a/src/flowx/sources/airflow/audit.py b/src/flowx/sources/airflow/audit.py index ffe9ac43..98771d08 100644 --- a/src/flowx/sources/airflow/audit.py +++ b/src/flowx/sources/airflow/audit.py @@ -11,6 +11,29 @@ _EDGE_MODIFIER_CONSTRUCTS = frozenset({"Label"}) +def _is_literal_string_statement(statement: ast.stmt) -> bool: + """Returns whether *statement* is a literal string expression, such as a docstring.""" + return ( + isinstance(statement, ast.Expr) + and isinstance(statement.value, ast.Constant) + and isinstance(statement.value.value, str) + ) + + +def _is_wiring_statement(statement: ast.stmt) -> bool: + """Returns whether *statement* only expresses task dependency wiring.""" + if not isinstance(statement, ast.Expr): + return False + value = statement.value + if isinstance(value, ast.BinOp) and isinstance(value.op, (ast.RShift, ast.LShift)): + return True + if not isinstance(value, ast.Call): + return False + if isinstance(value.func, ast.Name) and value.func.id in ("chain", "cross_downstream"): + return True + return isinstance(value.func, ast.Attribute) and value.func.attr in ("set_upstream", "set_downstream") + + @dataclass(frozen=True, slots=True, kw_only=True) class AuditCandidate: """One source construct that must be accounted for by capture and translation.""" @@ -110,6 +133,11 @@ def __init__(self, module: ast.Module, *, target_dag_variable: str | None) -> No for node in module.body if isinstance(node, ast.FunctionDef) and _single_operator_return(node, self.aliases) } + self.called_names = { + node.func.id + for node in ast.walk(module) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } def _candidate(self, kind: str, code: str, node: ast.AST, **details: Any) -> AuditCandidate: key = (kind, getattr(node, "lineno", 0), getattr(node, "col_offset", 0)) @@ -136,6 +164,34 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> None: for statement in node.body: if not isinstance(statement, ast.FunctionDef): self.visit(statement) + elif self._is_static_wiring_helper(node): + # A statically invoked wiring helper applies its edges when Airflow runs it, so the + # audit follows the same rule as the capture pass and both report the same edges. + for statement in node.body: + self.visit(statement) + + def _is_static_wiring_helper(self, node: ast.FunctionDef) -> bool: + """Returns whether *node* is a statically invoked, wiring-only helper (the visitor's rule). + + The guards here must match ``_DagVisitor._is_static_wiring_helper`` so that reconciliation + compares identical edge sets on both sides and fails closed on any shape one side skips. + """ + if node.name not in self.called_names: + return False + if node.decorator_list: + return False + if ( + node.args.args + or node.args.posonlyargs + or node.args.kwonlyargs + or node.args.vararg + or node.args.kwarg + ): + return False + body = list(node.body) + if body and _is_literal_string_statement(body[0]): + body = body[1:] + return all(_is_wiring_statement(statement) for statement in body) def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None: """Treats decorated async callables as TaskFlow definitions, not DAG-body statements.""" diff --git a/src/flowx/sources/airflow/loader/visitor.py b/src/flowx/sources/airflow/loader/visitor.py index 1f0828c0..a8c0530d 100644 --- a/src/flowx/sources/airflow/loader/visitor.py +++ b/src/flowx/sources/airflow/loader/visitor.py @@ -42,6 +42,15 @@ _EDGE_MODIFIER_CONSTRUCTS = frozenset({"Label"}) +def _is_literal_string_statement(statement: ast.stmt) -> bool: + """Returns whether *statement* is a literal string expression, such as a docstring.""" + return ( + isinstance(statement, ast.Expr) + and isinstance(statement.value, ast.Constant) + and isinstance(statement.value.value, str) + ) + + class _DagVisitor(ast.NodeVisitor): """Collects operator calls, dependency edges, and the DAG's schedule.""" @@ -110,6 +119,13 @@ def __init__(self, module: ast.Module, *, target_dag_variable: str | None = None # TaskFlow: function name -> (definition, decorator dotted-name) for @task-decorated defs. # Pre-scanned so a @task def defined after the @dag body that uses it is still resolved. self.taskflow_defs: dict[str, tuple[ast.FunctionDef | ast.AsyncFunctionDef, str]] = {} + # Bare static invocations anywhere in the module (`helper()`), so a wiring-only helper + # that Airflow executes once can be followed by both the capture pass and the audit. + self._statically_called_names = { + node.func.id + for node in ast.walk(module) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } # @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() @@ -170,6 +186,9 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> None: if not is_dag_definition: if self._dag_scope_depth: self._claimed_statement_ids.add(id(node)) + return + if self._is_static_wiring_helper(node): + self._descend_wiring_helper(node) return if is_dag_definition: self.is_taskflow_dag = True @@ -194,6 +213,70 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> None: self._dag_scope_depth -= 1 self._scope_stack.pop() + def _is_static_wiring_helper(self, node: ast.FunctionDef) -> bool: + """Returns whether *node* is a statically invoked helper whose body only wires tasks. + + Airflow applies the helper's statements only when it runs. Both this capture pass and the + independent source audit follow the same rule -- module-level definition, no decorators, no + parameters, wiring-only statements (shift chains, ``chain``/``cross_downstream``, and + ``set_upstream``/``set_downstream``), and at least one static invocation -- so reconciliation + compares the same edges on both sides. A helper that cannot be followed this way stays under + the existing rules instead of being silently reduced to a claimless body. + """ + if node.name not in self._statically_called_names: + return False + if node.decorator_list: + return False + if ( + node.args.args + or node.args.posonlyargs + or node.args.kwonlyargs + or node.args.vararg + or node.args.kwarg + ): + return False + body = list(node.body) + if body and _is_literal_string_statement(body[0]): + body = body[1:] + return all(self._is_wiring_statement(statement) for statement in body) + + def _is_wiring_statement(self, statement: ast.stmt) -> bool: + if not isinstance(statement, ast.Expr): + return False + value = statement.value + if isinstance(value, ast.BinOp) and isinstance(value.op, (ast.RShift, ast.LShift)): + return True + if not isinstance(value, ast.Call): + return False + if isinstance(value.func, ast.Name) and value.func.id in ("chain", "cross_downstream"): + return True + return isinstance(value.func, ast.Attribute) and value.func.attr in ("set_upstream", "set_downstream") + + def _descend_wiring_helper(self, node: ast.FunctionDef) -> None: + """Follows a wiring helper the way Airflow runs it: the captured edges all apply. + + Each statement must resolve to a capture; anything else lands in ``unclaimed_statements`` + so reconciliation fails closed instead of verifying with a missing dependency. + """ + self._scope_stack.append(node) + try: + for statement in node.body: + unclaimed_calls_before = len(self.unclaimed_task_calls) + unresolved_before = len(self.unresolved_constructs) + self.visit(statement) + if id(statement) in self._claimed_statement_ids: + continue + if ( + len(self.unclaimed_task_calls) > unclaimed_calls_before + or len(self.unresolved_constructs) > unresolved_before + ): + continue + if _is_literal_string_statement(statement): + continue + self.unclaimed_statements.append(statement) + finally: + self._scope_stack.pop() + def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None: """Claims native async TaskFlow definitions without treating their bodies as DAG structure.""" if self._dag_scope_depth: @@ -266,6 +349,8 @@ def _register_operator_call(self, node: ast.Call, var: str, *, binding: str | No 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") + if not self._dag_scope_depth and not isinstance(dag_node, ast.Name): + return False if self._target_dag_variable is not None and not ( isinstance(dag_node, ast.Name) and dag_node.id == self._target_dag_variable ): diff --git a/tests/unit/test_airflow_reconciliation.py b/tests/unit/test_airflow_reconciliation.py index 74476487..58b6d0c9 100644 --- a/tests/unit/test_airflow_reconciliation.py +++ b/tests/unit/test_airflow_reconciliation.py @@ -569,6 +569,90 @@ def test_static_classic_dag_factory_preserves_dag_identity_and_tasks(tmp_path: P assert "echo hello" in (pipeline.tasks[0].generated_source or "") +def test_classic_dag_default_args_do_not_leak_to_unwired_operators(tmp_path: Path) -> None: + dag_path = tmp_path / "dag_context.py" + dag_path.write_text( + "from airflow import DAG\n" + "from airflow.operators.bash import BashOperator\n" + "with DAG(dag_id='context', schedule=None, default_args={'retries': 4}) as dag:\n" + " BashOperator(task_id='inside', bash_command='echo inside')\n" + "outside = BashOperator(task_id='outside', bash_command='echo outside')\n", + encoding="utf-8", + ) + + pipeline = airflow_loader.load_pipelines(dag_path)[0] + + assert [(task.task_key, task.max_retries) for task in pipeline.tasks] == [("inside", 4)] + + +def test_wiring_helper_edges_are_followed(tmp_path: Path) -> None: + dag_path = tmp_path / "wiring_helper.py" + dag_path.write_text( + "from airflow import DAG\n" + "from airflow.operators.bash import BashOperator\n" + "dag = DAG('helper_edges', schedule=None)\n" + "start = BashOperator(task_id='start', bash_command='x', dag=dag)\n" + "t = BashOperator(task_id='t', bash_command='x', dag=dag)\n" + "def wire():\n" + " start >> t\n" + "wire()\n", + encoding="utf-8", + ) + + pipeline = airflow_loader.load_pipelines(dag_path)[0] + tasks = {task.task_key: task for task in pipeline.tasks} + + assert pipeline.reconciliation_status == "verified" + assert [dependency.task_key for dependency in tasks["t"].depends_on or []] == ["start"] + assert pipeline.audit["audited_edge_count"] == 1 + assert pipeline.audit["captured_edge_count"] == 1 + + +def test_dag_context_wiring_helper_edges_are_followed(tmp_path: Path) -> None: + dag_path = tmp_path / "context_wiring_helper.py" + dag_path.write_text( + "from airflow import DAG\n" + "from airflow.operators.bash import BashOperator\n" + "with DAG('context_helper_edges', schedule=None) as dag:\n" + " start = BashOperator(task_id='start', bash_command='x')\n" + " t = BashOperator(task_id='t', bash_command='x')\n" + "def wire():\n" + " start >> t\n" + "wire()\n", + encoding="utf-8", + ) + + pipeline = airflow_loader.load_airflow_dag(dag_path) + tasks = {task.task_key: task for task in pipeline.tasks} + + assert pipeline.reconciliation_status == "verified" + assert [dependency.task_key for dependency in tasks["t"].depends_on or []] == ["start"] + + +def test_uncalled_wiring_helper_adds_no_edges(tmp_path: Path) -> None: + # Airflow only applies the helper's wiring when it runs, so an uncalled helper + # must leave the DAG exactly as Airflow sees it (no dependency). + dag_path = tmp_path / "uncalled_wiring_helper.py" + dag_path.write_text( + "from airflow import DAG\n" + "from airflow.operators.bash import BashOperator\n" + "dag = DAG('uncalled_helper', schedule=None)\n" + "start = BashOperator(task_id='start', bash_command='x', dag=dag)\n" + "t = BashOperator(task_id='t', bash_command='x', dag=dag)\n" + "def wire():\n" + " start >> t\n", + encoding="utf-8", + ) + + pipeline = airflow_loader.load_pipelines(dag_path)[0] + tasks = {task.task_key: task for task in pipeline.tasks} + + assert pipeline.reconciliation_status == "verified" + assert tasks["t"].depends_on in (None, []) + assert pipeline.audit["audited_edge_count"] == 0 + assert pipeline.audit["captured_edge_count"] == 0 + + def test_decorated_dag_factory_override_emits_each_invocation(tmp_path: Path) -> None: dag_path = tmp_path / "decorated_factory.py" dag_path.write_text(