From fc246551d89a41166b90893ce4300ec3414eb03f Mon Sep 17 00:00:00 2001 From: Speculator55005 <50082482+fas89@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:19:07 +0200 Subject: [PATCH 1/3] fix(snowflake): a python task no longer loses the orchestration DAG providers/snowflake/orchestration/airflow_generator.py emitted `python_callable=` for a function the generated module never defined, so Airflow raised NameError parsing the file and dropped the whole DAG. Reproduced in a real DagBag on Airflow 3.2.2 and 2.10.5: "NameError: name 'score_customers' is not defined". The generator is a library function: no fluid command calls it (fluid export uses providers/snowflake/codegen). It now defines a stub of that callable right above its PythonOperator. The stub raises AirflowFailException naming the task and the callable, imported when it runs (airflow.sdk.exceptions on Airflow 3, airflow.exceptions on 2), so the DAG loads and that one task fails at once, without the retries a stub cannot pass, until someone writes the function. `airflow tasks test` marks it failed on both versions; a NotImplementedError left it up_for_retry. A task that names no callable gets `_func`, as before. Injection safety is unchanged: the stub's name goes through sanitize_identifier and its message is one py_str_literal. A stub `def` would rebind a name the module already uses above it, so a callable named `dag`, `DAG`, `default_args`, an operator class, an earlier task's variable or a `__dunder__` is refused with a ValueError at generation. A task that only becomes an ERROR or TODO comment binds nothing and reserves no name. Sharing its own task's name works: `x = PythonOperator(python_callable=x)` reads the stub before the rebind. Refusing python tasks outright was the alternative. A stub keeps the DAG a usable scaffold, the way this generator already marks what it cannot generate (TODO and ERROR comments). Airflow's @task.stub is for tasks an external worker runs and needs a recent standard provider, so it is not used. The new tests load the generated module against stub Airflow modules (the pattern of tests/cli/_schedule_dag_fixtures.load_dag) whose operators, like Airflow's, refuse a `dag=` that is not the DAG and `>>` on anything but a task, so the NameError and any rebinding reproduce without Airflow installed. --- CHANGELOG.md | 9 + .../orchestration/airflow_generator.py | 109 +++++++- .../test_snowflake_airflow_generator.py | 254 ++++++++++++++++++ 3 files changed, 360 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 791f1bb2..ee32c749 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -183,6 +183,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 source as text, so a contract value ran as Python when Airflow parsed the DAG. Both are now read as numbers (retries 0 to 100, the delay in whole minutes), falling back to 3 and 5. +- **A python task no longer loses the Snowflake orchestration DAG.** `generate_airflow_dag` + in `providers/snowflake/orchestration`, a library function no `fluid` command calls + (`fluid export` uses `providers/snowflake/codegen`), emitted `python_callable=` + for a function it never defined, so Airflow raised `NameError` parsing the file and + dropped the whole DAG. It now defines a stub of that callable above the task: the DAG + loads, and the task fails with `AirflowFailException`, naming itself and skipping the + retries a stub cannot pass, until the function is written. A callable whose name the DAG + already uses (`dag`, an operator class, an earlier task) or a `__dunder__` name is refused + with a `ValueError` when the DAG is generated. ## [0.21.0] — 2026-10-19 diff --git a/fluid_build/providers/snowflake/orchestration/airflow_generator.py b/fluid_build/providers/snowflake/orchestration/airflow_generator.py index 9c280342..4dbfc8ac 100644 --- a/fluid_build/providers/snowflake/orchestration/airflow_generator.py +++ b/fluid_build/providers/snowflake/orchestration/airflow_generator.py @@ -19,7 +19,7 @@ Converts FLUID provider action tasks into Airflow operators: - sf.table.ensure → SQLExecuteQueryOperator - sf.sql.execute → SQLExecuteQueryOperator -- python → PythonOperator +- python → PythonOperator, calling a stub of the callable it names - bash → BashOperator """ @@ -51,6 +51,10 @@ def generate_airflow_dag( Returns: Python code for Airflow DAG + + Raises: + ValueError: a SQL identifier or type is unsafe, or a python task's callable + name would rebind a name the generated DAG uses. """ if config is None: config = _build_config_from_contract(contract, orchestration) @@ -100,6 +104,25 @@ def _build_config_from_contract( ) +#: Every name ``_generate_dag_header`` binds, plus ``ImportError``, which each stub's +#: ``except`` names. A python task's stub ``def`` under one of these would rebind it for +#: every task below. +_DAG_MODULE_NAMES = frozenset( + { + "datetime", + "timedelta", + "DAG", + "SQLExecuteQueryOperator", + "PythonOperator", + "BashOperator", + "ExternalTaskSensor", + "default_args", + "dag", + "ImportError", + } +) + + def _generate_dag_header(config: OrchestrationConfig) -> str: """Generate DAG header with imports and configuration. @@ -163,16 +186,19 @@ def _generate_dag_header(config: OrchestrationConfig) -> str: def _generate_task_definitions(tasks: List[Dict[str, Any]], config: OrchestrationConfig) -> str: """Generate Airflow task definitions.""" code = "# Task definitions\n\n" + # The variables the tasks above bound: a python task's stub ``def`` must not rebind one. + task_vars: set = set() for task in tasks: task_type = task.get("type", "provider_action") task_name = task.get("name", "unknown") task_id = sanitize_task_id(task_name) + start = len(code) if task_type == "provider_action": code += _generate_provider_action_task(task, task_id, config) elif task_type == "python": - code += _generate_python_task(task, task_id) + code += _generate_python_task(task, task_id, task_vars) elif task_type == "bash": code += _generate_bash_task(task, task_id) elif task_type == "sensor": @@ -184,6 +210,9 @@ def _generate_task_definitions(tasks: List[Dict[str, Any]], config: Orchestratio f"# TODO: Unsupported task type {py_str_literal(task_type)} " f"for task {py_str_literal(task_name)}\n\n" ) + var = sanitize_identifier(task_id) + if f"\n{var} = " in "\n" + code[start:]: # an ERROR or TODO comment binds nothing + task_vars.add(var) return code @@ -243,23 +272,59 @@ def _generate_provider_action_task( """ -def _generate_python_task(task: Dict[str, Any], task_id: str) -> str: - """Generate PythonOperator. +def _generate_python_task(task: Dict[str, Any], task_id: str, task_vars: set) -> str: + """Generate a PythonOperator, preceded by a stub of the callable it names. - ``python_callable`` is emitted as a *bare* Python name (a reference to a - callable), so the contract-supplied ``callable`` is routed through - ``sanitize_identifier`` -- otherwise a value like ``os.system(...) or x`` - would be interpolated as a top-level expression and executed at parse - time. The LHS variable name uses the same helper so it matches the - dependency wiring; ``task_id`` is emitted as a ``repr()``-escaped literal. + The contract ships no code for the callable, and ``python_callable=`` naming + an undefined function is a ``NameError`` when Airflow parses the file, which + loses the whole DAG. The stub raises ``AirflowFailException`` naming the task: + the DAG loads, and this task fails, without the retries a stub cannot pass, + until someone writes the function. It never goes green having done nothing. + It imports the exception when it runs, from Airflow 3's ``airflow.sdk`` + first, so nothing is imported at parse time. + + The stub's name is the contract's ``callable`` routed through + ``sanitize_identifier`` (otherwise ``os.system(...) or x`` would be + interpolated as code and run at parse time), and its message is one + ``repr()``-escaped literal via ``py_str_literal``. The LHS variable name + uses the same helper so it matches the dependency wiring. """ - callable_name = task.get("callable", task_id + "_func") + named = task.get("callable") + callable_name = named or f"{task_id}_func" var = sanitize_identifier(task_id) task_id_lit = py_str_literal(task_id) callable_ident = sanitize_identifier(callable_name) - return f"""{var} = PythonOperator( + clash = _stub_clash(callable_ident, task_vars) + if clash: + what = f"callable {str(named)!r}" if named else f"its default callable {callable_ident!r}" + raise ValueError( + f"python task {task_id!r}: cannot stub {what}: {clash}. " + "Set the task's callable to another name." + ) + + source = ( + f"the contract names the callable {str(named)!r} but ships no code for it" + if named + else "the contract's python task names no callable" + ) + message_lit = py_str_literal( + f"Task {task_id!r} calls {callable_ident}(), which this generated DAG only stubs: " + f"{source}. Implement {callable_ident}() in this file, or make the task a bash " + "task that runs your code." + ) + + return f"""# Stub: the contract does not ship this callable's code. +def {callable_ident}(*args, **kwargs): + try: # Airflow 3 + from airflow.sdk.exceptions import AirflowFailException + except ImportError: # Airflow 2.x + from airflow.exceptions import AirflowFailException + raise AirflowFailException({message_lit}) + + +{var} = PythonOperator( task_id={task_id_lit}, python_callable={callable_ident}, dag=dag, @@ -268,6 +333,26 @@ def _generate_python_task(task: Dict[str, Any], task_id: str) -> str: """ +def _stub_clash(ident: str, task_vars: set) -> str: + """Why a stub ``def`` named *ident* is refused, or ``""`` when it is not. + + It would rebind a name the module defines above it. For the ``dag``, an + operator class or an earlier task's variable the DAG would then fail to load + or wire the wrong object; the header's other names are refused too, so no stub + shadows them. A ``__dunder__`` name is Python's. Sharing its own task's name is + fine: the stub sits right above the operator, and + ``x = PythonOperator(python_callable=x)`` reads the function before the + assignment rebinds ``x``. + """ + if ident.startswith("__") and ident.endswith("__"): + return "names of the form __name__ are Python's" + if ident in task_vars: + return f"an earlier task's variable is already named {ident!r}" + if ident in _DAG_MODULE_NAMES: + return f"the generated DAG already uses the name {ident!r}" + return "" + + def _generate_bash_task(task: Dict[str, Any], task_id: str) -> str: """Generate BashOperator. diff --git a/tests/providers/test_snowflake_airflow_generator.py b/tests/providers/test_snowflake_airflow_generator.py index 16618f00..8a83351f 100644 --- a/tests/providers/test_snowflake_airflow_generator.py +++ b/tests/providers/test_snowflake_airflow_generator.py @@ -28,6 +28,10 @@ from __future__ import annotations import ast +import re +import sys +import types +from typing import Any, Dict import pytest @@ -509,3 +513,253 @@ def test_unknown_action_and_unsupported_type_comments_are_inert(self) -> None: ] } _assert_no_injection(airflow_generator.generate_airflow_dag(contract, orchestration)) + + def test_python_callable_injection_is_inert(self, monkeypatch: pytest.MonkeyPatch) -> None: + payload = "x'); " + _PWN + " #" + _NL + _PWN + code = _python_dag({"name": "py", "type": "python", "callable": payload}) + _assert_no_injection(code) + dag = _load_with_stub_airflow(code, monkeypatch) + with pytest.raises(_AirflowFailException) as error: + dag["py"].kwargs["python_callable"]() + assert _PWN in str(error.value) # the contract's text survives only as data + + +class _AirflowFailException(Exception): + """Stands in for Airflow's exception that fails a task without retrying it.""" + + +def _load_with_stub_airflow(source: str, monkeypatch: pytest.MonkeyPatch) -> Dict[str, Any]: + """Run a generated DAG module against stub Airflow 3 modules; return its namespace. + + Adapted from ``tests/cli/_schedule_dag_fixtures.load_dag``: Airflow is not a + test dependency, so each operator the generator imports is a class that keeps + its kwargs. Running the module body is what Airflow's DagBag does when it + parses the file, so a name the DAG uses but never defines raises ``NameError`` + here as it would there. Like Airflow, an operator refuses a ``dag=`` that is + not the DAG and ``>>`` refuses anything but a task, so a name the module + rebinds fails here too. + """ + + class _DAG: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + class _Task: + def __init__(self, **kwargs: Any) -> None: + assert isinstance(kwargs.get("dag"), _DAG), f"dag={kwargs.get('dag')!r}" + assert callable(kwargs.get("python_callable", len)), kwargs + self.kwargs = kwargs + + def __rshift__(self, other: "_Task") -> "_Task": + assert isinstance(other, _Task), f"{self!r} >> {other!r}" + return other + + def module(name: str, **attrs: Any) -> types.ModuleType: + mod = types.ModuleType(name) + mod.__dict__.update(attrs) + return mod + + def operator(name: str) -> type: + return type(name, (_Task,), {}) + + stubs = { + name: module(name) + for name in ( + "airflow", + "airflow.providers", + "airflow.providers.common", + "airflow.providers.common.sql", + "airflow.providers.common.sql.operators", + "airflow.providers.standard", + "airflow.providers.standard.operators", + "airflow.providers.standard.sensors", + ) + } + for name, attrs in { + "airflow.sdk": {"DAG": _DAG}, + # A stub imports it when it runs, not when the DAG is parsed. + "airflow.sdk.exceptions": {"AirflowFailException": _AirflowFailException}, + "airflow.providers.common.sql.operators.sql": { + "SQLExecuteQueryOperator": operator("SQLExecuteQueryOperator") + }, + "airflow.providers.standard.operators.python": { + "PythonOperator": operator("PythonOperator") + }, + "airflow.providers.standard.operators.bash": {"BashOperator": operator("BashOperator")}, + "airflow.providers.standard.sensors.external_task": { + "ExternalTaskSensor": operator("ExternalTaskSensor") + }, + }.items(): + stubs[name] = module(name, **attrs) + for name, mod in stubs.items(): + monkeypatch.setitem(sys.modules, name, mod) + + namespace: Dict[str, Any] = {"__name__": "generated_dag"} + exec(compile(source, "generated_dag.py", "exec"), namespace) # noqa: S102 - test harness + return namespace + + +def _python_dag(*tasks: Dict[str, Any]) -> str: + """A DAG whose first task, ``load``, runs SQL, followed by *tasks*.""" + load = { + "name": "load", + "type": "provider_action", + "action": "sf.sql.execute", + "parameters": {"sql": "SELECT 1"}, + } + return airflow_generator.generate_airflow_dag({"id": "sales.orders"}, {"tasks": [load, *tasks]}) + + +class TestPythonTaskStub: + """A python task's callable is defined in the DAG it is generated into. + + It used to be emitted as ``python_callable=`` with no ``def``, so + parsing the file raised ``NameError`` and Airflow lost the whole DAG. + """ + + def test_the_dag_loads_and_the_stub_fails_its_task_by_name( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + dag = _load_with_stub_airflow( + _python_dag( + { + "name": "score", + "type": "python", + "callable": "score_customers", + "depends_on": ["load"], + }, + {"name": "notify", "type": "bash", "command": "echo done", "depends_on": ["score"]}, + ), + monkeypatch, + ) + stub = dag["score"].kwargs["python_callable"] + assert stub is dag["score_customers"] + with pytest.raises( + _AirflowFailException, match=r"Task 'score' calls score_customers\(\)" + ) as error: + stub(ds="2026-01-01") # Airflow passes the context as keywords + assert "names the callable 'score_customers' but ships no code" in str(error.value) + + def test_a_task_without_a_callable_gets_a_default_stub( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + dag = _load_with_stub_airflow(_python_dag({"name": "audit", "type": "python"}), monkeypatch) + assert dag["audit"].kwargs["python_callable"] is dag["audit_func"] + with pytest.raises(_AirflowFailException, match="names no callable"): + dag["audit_func"]() + + def test_a_stub_may_share_its_own_tasks_name(self, monkeypatch: pytest.MonkeyPatch) -> None: + dag = _load_with_stub_airflow( + _python_dag({"name": "transform", "type": "python", "callable": "transform"}), + monkeypatch, + ) + assert type(dag["transform"]).__name__ == "PythonOperator" + with pytest.raises(_AirflowFailException, match="Task 'transform'"): + dag["transform"].kwargs["python_callable"]() + + def test_tasks_sharing_a_callable_each_name_their_own_task( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + dag = _load_with_stub_airflow( + _python_dag( + {"name": "a", "type": "python", "callable": "shared"}, + {"name": "b", "type": "python", "callable": "shared"}, + ), + monkeypatch, + ) + for task in ("a", "b"): + with pytest.raises(_AirflowFailException, match=f"Task '{task}'"): + dag[task].kwargs["python_callable"]() + + def test_a_later_task_may_take_the_callables_name( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The operator holds the stub before the later task rebinds the name.""" + dag = _load_with_stub_airflow( + _python_dag( + {"name": "py", "type": "python", "callable": "notify"}, + {"name": "notify", "type": "bash", "command": "echo hi", "depends_on": ["py"]}, + ), + monkeypatch, + ) + assert type(dag["notify"]).__name__ == "BashOperator" + with pytest.raises(_AirflowFailException, match="Task 'py'"): + dag["py"].kwargs["python_callable"]() + + @pytest.mark.parametrize( + "callable_name", + # ``load`` is the task variable defined above the python task. + [ + "dag", + "DAG", + "PythonOperator", + "default_args", + "ExternalTaskSensor", + "ImportError", + "__builtins__", + "load", + ], + ) + def test_a_callable_that_would_rebind_a_name_the_dag_uses_is_refused( + self, callable_name: str + ) -> None: + with pytest.raises( + ValueError, match=re.escape(f"python task 'py': cannot stub callable '{callable_name}'") + ): + _python_dag({"name": "py", "type": "python", "callable": callable_name}) + + def test_a_default_callable_clash_names_the_default(self) -> None: + # ``audit``'s default callable is ``audit_func``, already a task above it here. + with pytest.raises( + ValueError, match="cannot stub its default callable 'audit_func': an earlier task's" + ): + _python_dag( + {"name": "audit_func", "type": "bash", "command": "echo"}, + {"name": "audit", "type": "python"}, + ) + + def test_a_task_that_binds_nothing_reserves_no_name( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """An unsupported task type becomes a TODO comment, so its name stays free.""" + dag = _load_with_stub_airflow( + _python_dag( + {"name": "notify", "type": "email"}, + {"name": "py", "type": "python", "callable": "notify"}, + ), + monkeypatch, + ) + with pytest.raises(_AirflowFailException, match="Task 'py'"): + dag["notify"]() + + def test_the_harness_catches_a_rebound_dag(self, monkeypatch: pytest.MonkeyPatch) -> None: + source = _python_dag( + {"name": "py", "type": "python", "callable": "f"}, + {"name": "later", "type": "bash", "command": "echo"}, + ) + rebound = source.replace("def f(", "def dag(").replace( + "python_callable=f,", "python_callable=dag," + ) + assert rebound != source + with pytest.raises(AssertionError, match="dag= None: + from fluid_build.providers.snowflake.orchestration.common import ( + OrchestrationConfig, + OrchestrationEngine, + ) + + header = airflow_generator._generate_dag_header( + OrchestrationConfig(engine=OrchestrationEngine.AIRFLOW, dag_id="d") + ) + bound = set() + for node in ast.walk(ast.parse(header)): + if isinstance(node, (ast.Import, ast.ImportFrom)): + bound.update((alias.asname or alias.name).split(".")[0] for alias in node.names) + elif isinstance(node, ast.Assign): + bound.update(target.id for target in node.targets if isinstance(target, ast.Name)) + assert bound, "parsed nothing" + assert bound <= airflow_generator._DAG_MODULE_NAMES, ( + bound - airflow_generator._DAG_MODULE_NAMES + ) From d2a4b2600946acac8495921cc2b48d0938a57031 Mon Sep 17 00:00:00 2001 From: Speculator55005 <50082482+fas89@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:19:20 +0200 Subject: [PATCH 2/3] fix(init): the DAG helper writes what fluid generate schedule writes, or nothing generate_dag_for_project ran `fluid generate-airflow --output-dir ...`, a flag generate-airflow never had. argparse exited 2 ("unrecognized arguments: --output-dir"), so every project got the basic fallback DAG under a warning that the command was "not available"; with no `fluid` on PATH the FileNotFoundError meant no DAG at all. Its tests mocked subprocess.run, so nothing ever parsed that argv. The helper now runs `fluid generate schedule --scheduler airflow --output --contract-path contract.fluid.yaml --overwrite`, the command generate-airflow's own deprecation notice points to. (With valid flags, generate-airflow's DAG for the shipped templates is all `echo` placeholder tasks.) It runs in-process, parsed by the CLI's own build_parser, as stage 3 (artifact_fanout) runs the same command: a flag the command lacks fails here as on the command line, and with no child process there is no PATH, cwd, sys.path or output-decoding assumption. --overwrite lets a re-run regenerate its DAGs. It writes no DAG, and says why, when: - the contract schedules nothing: no build with a cron trigger and no orchestration.tasks (fluid_apply.has_scheduled_builds/has_explicit_tasks). customer-360's build is `trigger.type: manual`; generate schedule would have synthesized a nightly DAG whose only task logs and passes. - the command refuses the contract, an invalid cron for instance. The basic validate/plan/apply DAG it used to fall back to ran on a schedule the contract never declared and hid the error; stage 3 treats the same exit as a hard failure. Both are judged on contract.fluid.yaml as the command reads it: its loader resolves `$ref` builds and the legacy singular `build:` key, which the raw YAML hides. should_generate_dag, the caller's gate, now also accepts a build with a cron trigger, so the two gates agree. dags/README.md lists each DAG this run wrote (dag_id and schedule read from the file with ast), not other files already in dags/, says to edit the contract and regenerate rather than edit the generated DAG, and says that regenerating does not delete the DAG of a renamed or removed build. create_basic_dag stays as a Python helper. Its apply task ran `fluid apply contract.fluid.yaml --auto-approve`, which fluid apply rejects (exit 2); it passes --yes now, and its project path is shell-quoted and then emitted as a py_str_literal, so a path holding a quote no longer breaks the DAG or runs what follows it. No CLI command reaches this helper with DAGs enabled today: `fluid demo` passes no_dag=True and no `fluid init` mode calls it. tests/test_init_generate_schedule_dag.py parses the helper's argv with the real build_parser() and runs the real command on shipped templates: scheduled builds get their fluid apply DAG and README, a manual or continuous build gets none, a refused contract gets none plus the command's reason, a re-run regenerates, a file already in dags/ stays out of the README, a relative project dir works, and a singular `build:` contract gets its DAG. Older tests that mocked subprocess.run now mock the in-process call. --- CHANGELOG.md | 12 + fluid_build/cli/_init_dag_helpers.py | 342 ++++++++++++++++------- tests/test_init_branches.py | 79 ++++-- tests/test_init_extra_branches.py | 52 ++-- tests/test_init_generate_schedule_dag.py | 279 ++++++++++++++++++ 5 files changed, 631 insertions(+), 133 deletions(-) create mode 100644 tests/test_init_generate_schedule_dag.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ee32c749..fbbfc1d7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -192,6 +192,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 retries a stub cannot pass, until the function is written. A callable whose name the DAG already uses (`dag`, an operator class, an earlier task) or a `__dunder__` name is refused with a `ValueError` when the DAG is generated. +- **The `fluid init` DAG helper writes what `fluid generate schedule` writes, or nothing.** + `generate_dag_for_project` ran `fluid generate-airflow --output-dir`, a flag that command + does not have, so it always wrote a basic validate/plan/apply DAG, or no DAG when `fluid` + was not on PATH. It now runs `fluid generate schedule --scheduler airflow --overwrite` + in-process through the CLI's own parser, and its `dags/README.md` names the DAGs that run + wrote and says to edit the contract, not the DAG. It writes no DAG, and says why, when the + contract schedules nothing (no build with a cron trigger and no `orchestration.tasks`, as + with customer-360's `manual` build) or when the command refuses the contract, an invalid + cron for instance, instead of a basic DAG on a schedule the contract never declared. + `create_basic_dag`, no longer a fallback, passes `fluid apply --yes` instead of the + rejected `--auto-approve` and quotes the project path. No CLI command reaches the helper + with DAGs enabled today: `fluid demo` turns them off. ## [0.21.0] — 2026-10-19 diff --git a/fluid_build/cli/_init_dag_helpers.py b/fluid_build/cli/_init_dag_helpers.py index 406866b3..7ec6471f 100644 --- a/fluid_build/cli/_init_dag_helpers.py +++ b/fluid_build/cli/_init_dag_helpers.py @@ -29,8 +29,6 @@ import logging import os -import re -import re as _re # alias used inside ``create_basic_dag`` to disambiguate from pathlib import Path from typing import Any, Dict, Optional @@ -58,13 +56,23 @@ def should_generate_dag(contract: dict, template: str = None) -> bool: Auto-generate DAGs when: 1. Contract has explicit orchestration config - 2. Template is orchestration-focused (customer-360, sales-analytics, ml-features, data-quality) - 3. Project has multiple provider actions (complex pipeline) + 2. A build has a cron trigger (what ``generate_dag_for_project`` writes a DAG for) + 3. Template is orchestration-focused (customer-360, sales-analytics, ml-features, data-quality) + 4. Project has multiple provider actions (complex pipeline) """ # Check for explicit orchestration config if "orchestration" in contract: return True + from fluid_build.schedulers.airflow import fluid_apply + + # ``build:`` is the schema's legacy singular form of ``builds:``. + build = contract.get("build") + if fluid_apply.has_scheduled_builds(contract) or ( + isinstance(build, dict) and fluid_apply.has_scheduled_builds({"builds": [build]}) + ): + return True + # Check for orchestration-focused templates orchestrated_templates = ["customer-360", "sales-analytics", "ml-features", "data-quality"] if template and template in orchestrated_templates: @@ -83,64 +91,61 @@ def generate_dag_for_project( project_dir: Path, contract: dict, logger, console, template: str = None ) -> bool: """ - Generate Airflow DAG using existing generate-airflow command. + Generate the project's Airflow DAG with ``fluid generate schedule``. Creates dags/ folder with: - - DAG Python file (contract_name_dag.py) - - README.md with usage instructions + - the DAG file(s) ``fluid generate schedule --scheduler airflow`` writes: a + ``fluid apply`` DAG per build with a cron trigger, or the scheduler's DAG + for the contract's ``orchestration.tasks`` + - README.md naming each DAG, with usage instructions + + Writes nothing and returns False, saying why, when the contract schedules + nothing (a manual or continuous build runs on no timetable, so its DAG + would only run a placeholder) or when the command refuses the contract, + an invalid cron for instance: a stand-in DAG on a schedule the contract + never declared would hide the error. Both are judged on + ``contract.fluid.yaml`` as the command reads it, not on ``contract``. """ try: - import subprocess - - # Get contract details - contract_name = contract.get("name", "my_product") - orchestration = contract.get("orchestration", {}) - - # Prepare DAG parameters — sanitize to prevent injection - schedule = orchestration.get("schedule", "@daily") - dag_id = contract_name.replace("-", "_").replace(" ", "_") - # Strict identifier validation: only alphanumeric + underscore - dag_id = re.sub(r"[^a-zA-Z0-9_]", "", dag_id) or "fluid_dag" - # Validate schedule is a plausible cron/preset string - if not re.match(r"^[@a-zA-Z0-9_ */,-]+$", schedule): - schedule = "@daily" - - # Call generate-airflow command - dag_dir = project_dir / "dags" - dag_dir.mkdir(exist_ok=True) - - # Build command - cmd = [ - "fluid", - "generate-airflow", - str(project_dir / "contract.fluid.yaml"), - "--output-dir", - str(dag_dir), - "--dag-id", - dag_id, - "--schedule", - schedule, - ] - + from fluid_build.schedulers.airflow import fluid_apply + + loaded = _load_as_the_command_does(project_dir.absolute() / "contract.fluid.yaml", logger) + if loaded is not None and not ( + fluid_apply.has_scheduled_builds(loaded) or fluid_apply.has_explicit_tasks(loaded) + ): + return _no_dag( + "the contract schedules nothing (no build has a cron trigger and there " + "are no orchestration.tasks)", + logger, + console, + ) + + dag_dir = project_dir.absolute() / "dags" if _rich_available(): console.print("\n[cyan]📅 Generating Airflow DAG...[/cyan]") - # Execute command - result = subprocess.run(cmd, capture_output=True, text=True, cwd=str(project_dir)) - - if result.returncode != 0: - # generate-airflow may not exist yet - create DAG manually. - # Resolve via ``_init`` so test patches on - # ``fluid_build.cli.init.create_basic_dag`` flow through. - logger.warning("generate-airflow command not available, creating basic DAG template") - _init.create_basic_dag(project_dir, contract, logger) - - # Create README - dag_filename = f"{dag_id}_dag.py" - create_dags_readme(dag_dir, dag_id, schedule, dag_filename) - + before = _dag_files(dag_dir) + exit_code = _run_generate_schedule(project_dir, dag_dir, logger) + if exit_code != 0: + return _no_dag( + f"fluid generate schedule refused the contract (exit {exit_code}); " + "its message above says why", + logger, + console, + warning=True, + ) + written = sorted( + path for path, stamp in _dag_files(dag_dir).items() if before.get(path) != stamp + ) + if not written: + return _no_dag("fluid generate schedule wrote no DAG file", logger, console) + + names = _write_dags_readme(dag_dir, written) if _rich_available(): - console.print(f"[green]✅ DAG created: dags/{dag_filename}[/green]") + from rich.markup import escape + + for name in names: + console.print(f"[green]✅ DAG created: dags/{escape(name)}[/green]") return True @@ -149,8 +154,159 @@ def generate_dag_for_project( return False +def _load_as_the_command_does(contract_path: Path, logger): + """The contract as ``fluid generate schedule`` reads it, or None when it does not load. + + The command's loader resolves ``$ref`` builds and the legacy singular + ``build:`` key, which the raw YAML does not show. When the contract does not + load, the command itself reports why. + """ + from fluid_build._contract_loader import load_contract_with_overlay + + try: + return load_contract_with_overlay(str(contract_path), None, logger) + except Exception: + return None + + +def _no_dag(reason: str, logger, console, warning: bool = False) -> bool: + """Say why no DAG was written; return False.""" + (logger.warning if warning else logger.info)(f"No Airflow DAG written: {reason}") + if _rich_available(): + from rich.markup import escape + + console.print( + f"[{'yellow' if warning else 'dim'}]No Airflow DAG written: {escape(reason)}[/]" + ) + return False + + +def _generate_schedule_argv(project_dir: Path, dag_dir: Path) -> list: + """The ``fluid generate schedule`` arguments that write the project's Airflow DAGs. + + Absolute paths and ``--contract-path`` make the result independent of the + current directory: a ``fluid apply`` DAG records the contract's path relative + to ``$FLUID_PROJECT_DIR`` on the worker, which is the project root here. + ``--overwrite`` lets a later run regenerate the DAGs an earlier one wrote. + """ + contract = project_dir.absolute() / "contract.fluid.yaml" + return ["generate", "schedule", str(contract), "--scheduler", "airflow"] + [ + "--output", + str(dag_dir), + "--contract-path", + "contract.fluid.yaml", + "--overwrite", + ] + + +def _run_generate_schedule(project_dir: Path, dag_dir: Path, logger) -> int: + """Run ``fluid generate schedule`` in this process; return its exit code. + + As stage 3 runs it (``forge/core/artifact_fanout.py``), but parsed by the + CLI's own ``build_parser``, so a flag the command does not have fails here + as it would on the command line. The command prints its own messages. + """ + from fluid_build.cli import build_parser + + try: + args = build_parser().parse_args(_generate_schedule_argv(project_dir, dag_dir)) + except SystemExit as exc: # argparse has printed why + return exc.code if isinstance(exc.code, int) else 2 + return args.func(args, logger) + + +def _dag_files(dag_dir: Path) -> dict: + """Each regular ``*.py`` file under *dag_dir*, with its modification stamp. + + Comparing a snapshot taken before a run with one taken after tells the + files the run wrote from ones already there. Only regular files count: + reading a FIFO named ``x.py`` would block. + """ + if not dag_dir.is_dir(): + return {} + stamps = {} + for path in dag_dir.rglob("*.py"): + if path.is_file(): + stat = path.stat() + stamps[path] = (stat.st_mtime_ns, stat.st_size) + return stamps + + +def _write_dags_readme(dag_dir: Path, dag_files: list) -> list: + """Write dags/README.md naming each DAG file; return their names.""" + names = [path.relative_to(dag_dir).as_posix() for path in dag_files] + dags = [ + (dag_id or "(see the DAG file)", schedule or "(see the DAG file)", name) + for (dag_id, schedule), name in zip(map(_dag_facts, dag_files), names, strict=True) + ] + create_dags_readme(dag_dir, *dags[0], others=dags[1:]) + return names + + +def _dag_facts(path: Path) -> tuple: + """``(dag_id, schedule)`` as a generated DAG file spells them, for the README. + + Read with ``ast``, not imported: importing needs Airflow and runs the file. + ``schedule=`` may name a module constant (the ``fluid apply`` DAG's + ``SCHEDULE``). A value not written as a string comes back as ``""``. + """ + import ast + + try: + tree = ast.parse(path.read_text(encoding="utf-8-sig")) + except (OSError, SyntaxError, ValueError, RecursionError): + return "", "" + constants = { + target.id: node.value.value + for node in tree.body + if isinstance(node, ast.Assign) and isinstance(node.value, ast.Constant) + for target in node.targets + if isinstance(target, ast.Name) + } + + def string_value(node) -> str: + if isinstance(node, ast.Constant): + value = node.value + else: + value = constants.get(getattr(node, "id", None)) + return value if isinstance(value, str) else "" + + for node in ast.walk(tree): + func = getattr(node, "func", None) + if isinstance(node, ast.Call) and "DAG" in ( + getattr(func, "id", None), + getattr(func, "attr", None), + ): + kwargs = {kw.arg: kw.value for kw in node.keywords if kw.arg} + return string_value(kwargs.get("dag_id")), string_value(kwargs.get("schedule")) + return "", "" + + +def _basic_dag_commands(project_dir: Path) -> tuple: + """The basic DAG's validate, plan and apply commands, as Python string literals. + + The project path goes into a shell command inside a Python string literal, so + it is quoted for both: a path holding a quote must not end the literal and run + what follows when Airflow parses the DAG. + """ + import shlex + + from fluid_build.providers.common.codegen_utils import py_str_literal + + in_project = f"cd {shlex.quote(str(project_dir.absolute()))} && fluid" + return tuple( + py_str_literal(f"{in_project} {step} contract.fluid.yaml{flags}") + for step, flags in (("validate", ""), ("plan", ""), ("apply", " --yes")) + ) + + def create_basic_dag(project_dir: Path, contract: dict, logger): - """Create a basic DAG template if generate-airflow is not available.""" + """Create a basic validate, plan and apply DAG for the project (Python API). + + ``generate_dag_for_project`` does not fall back to it: when ``fluid generate + schedule`` refuses a contract, a DAG on a schedule the contract never + declared would hide the error. + """ import re as _re @@ -159,6 +315,7 @@ def create_basic_dag(project_dir: Path, contract: dict, logger): airflow_standard_import, ) + validate_cmd, plan_cmd, apply_cmd = _basic_dag_commands(project_dir) contract_name = contract.get("name", "my_product") orchestration = contract.get("orchestration", {}) # Sanitize values to prevent code injection in generated Python. @@ -204,19 +361,19 @@ def create_basic_dag(project_dir: Path, contract: dict, logger): # Validate contract validate = BashOperator( task_id='validate_contract', - bash_command='cd {project_dir.absolute()} && fluid validate contract.fluid.yaml', + bash_command={validate_cmd}, ) # Plan execution plan = BashOperator( task_id='plan_execution', - bash_command='cd {project_dir.absolute()} && fluid plan contract.fluid.yaml', + bash_command={plan_cmd}, ) # Apply changes apply = BashOperator( task_id='apply_contract', - bash_command='cd {project_dir.absolute()} && fluid apply contract.fluid.yaml --auto-approve', + bash_command={apply_cmd}, ) validate >> plan >> apply @@ -232,36 +389,34 @@ def create_basic_dag(project_dir: Path, contract: dict, logger): logger.info(f"Created basic DAG template: {dag_file}") -def create_dags_readme(dag_dir: Path, dag_id: str, schedule: str, dag_filename: str): - """Create README in dags/ folder with usage instructions.""" +def create_dags_readme( + dag_dir: Path, dag_id: str, schedule: str, dag_filename: str, others: tuple = () +): + """Create README in dags/ folder with usage instructions. + + ``others`` holds one ``(dag_id, schedule, dag_filename)`` per further DAG, + when ``fluid generate schedule`` wrote more than one. + """ + dags = [(dag_id, schedule, dag_filename), *others] + listing = "\n\n".join( + f"- **DAG ID**: `{one_id}`\n- **Schedule**: `{one_schedule}`\n- **File**: `{one_file}`" + for one_id, one_schedule, one_file in dags + ) readme_content = f"""# Airflow DAG Configuration -This folder contains the Airflow DAG for your FLUID data product. +This folder contains the Airflow {"DAG" if len(dags) == 1 else "DAGs"} for your FLUID data product, \ +generated from `contract.fluid.yaml` by `fluid generate schedule`. -## Generated DAG +## Generated {"DAG" if len(dags) == 1 else "DAGs"} -- **DAG ID**: `{dag_id}` -- **Schedule**: `{schedule}` -- **File**: `{dag_filename}` +{listing} ## Usage -### Local Development - -Run the DAG locally using Airflow: - -```bash -# Start Airflow (from project root) -docker-compose --profile airflow up -d - -# Access Airflow UI -open http://localhost:8080 - -# Default credentials -# Username: admin -# Password: admin -``` +Copy the DAG file into your Airflow deployment's DAGs folder (Cloud Composer, MWAA, +Astronomer or a self-hosted Airflow), after reading its docstring: a `fluid apply` DAG +needs `FLUID_PROJECT_DIR` set to this project on the Airflow worker and `fluid` on PATH. ### Manual Execution @@ -274,28 +429,21 @@ def create_dags_readme(dag_dir: Path, dag_id: str, schedule: str, dag_filename: # Plan execution fluid plan contract.fluid.yaml -# Apply changes -fluid apply contract.fluid.yaml --auto-approve +# Apply changes and run the builds +fluid apply contract.fluid.yaml --mode amend-and-build --yes ``` -### CI/CD Integration - -This DAG can be deployed to: -- Cloud Composer (GCP) -- MWAA (AWS) -- Astronomer -- Self-hosted Airflow - -See `.jenkins/` folder for CI/CD pipeline configuration. - ## Customization -To customize the DAG: +The DAG is generated: edit the contract, not the DAG file. To change the schedule or +the tasks, update `contract.fluid.yaml` and regenerate: + +```bash +fluid generate schedule contract.fluid.yaml --scheduler airflow --output dags --overwrite +``` -1. Edit `{dag_filename}` -2. Add custom operators or sensors -3. Configure alerting and notifications -4. Update schedule interval as needed +Regenerating does not delete the DAG of a build you renamed or removed: delete that +file from `dags/` yourself. ## Next Steps diff --git a/tests/test_init_branches.py b/tests/test_init_branches.py index 80664884..1357600d 100644 --- a/tests/test_init_branches.py +++ b/tests/test_init_branches.py @@ -312,36 +312,75 @@ def test_exception(self, _mock_detect): class TestGenerateDagForProject: - @patch("subprocess.run") - def test_dag_generation_success(self, mock_run, tmp_path): + # ``orchestration.tasks`` makes the contract one that schedules something. + CONTRACT = {"name": "test-product", "orchestration": {"tasks": [{"taskId": "t"}]}} + RUN = "fluid_build.cli._init_dag_helpers._run_generate_schedule" + + @staticmethod + def _project(tmp_path, contract): + """A project holding *contract* (JSON is YAML), as the helper reads it from disk.""" + import json + + (tmp_path / "contract.fluid.yaml").write_text(json.dumps(contract)) + return tmp_path + + def test_dag_generation_success(self, tmp_path): from fluid_build.cli.init import generate_dag_for_project - mock_run.return_value = MagicMock(returncode=0) - contract = {"name": "test-product", "orchestration": {"schedule": "@daily"}} - mock_console = MagicMock() - result = generate_dag_for_project( - tmp_path, contract, logging.getLogger("test"), mock_console - ) + def write_dag(project_dir, dag_dir, logger): # what ``fluid generate schedule`` writes + dag_dir.mkdir() + (dag_dir / "p_dag.py").write_text("DAG(dag_id='p', schedule='@daily')\n") + return 0 + + with patch(self.RUN, side_effect=write_dag): + result = generate_dag_for_project( + self._project(tmp_path, self.CONTRACT), + self.CONTRACT, + logging.getLogger("test"), + MagicMock(), + ) assert result is True + assert "`p_dag.py`" in (tmp_path / "dags" / "README.md").read_text() - @patch("subprocess.run") - def test_dag_generation_fallback(self, mock_run, tmp_path): + def test_dag_generation_writes_nothing_when_nothing_is_scheduled(self, tmp_path): from fluid_build.cli.init import generate_dag_for_project - mock_run.return_value = MagicMock(returncode=1) - contract = {"name": "test-product"} - mock_console = MagicMock() - result = generate_dag_for_project( - tmp_path, contract, logging.getLogger("test"), mock_console - ) - assert result is True + with patch(self.RUN) as run: + result = generate_dag_for_project( + self._project(tmp_path, {"name": "test-product"}), + {"name": "test-product"}, + logging.getLogger("test"), + MagicMock(), + ) + assert result is False + run.assert_not_called() + assert not (tmp_path / "dags").exists() + + def test_dag_generation_refused_writes_no_dag(self, tmp_path): + from fluid_build.cli.init import generate_dag_for_project + + with ( + patch(self.RUN, return_value=1), + patch("fluid_build.cli.init.create_basic_dag") as basic, + ): + result = generate_dag_for_project( + self._project(tmp_path, self.CONTRACT), + self.CONTRACT, + logging.getLogger("test"), + MagicMock(), + ) + assert result is False + basic.assert_not_called() + assert not (tmp_path / "dags").exists() def test_dag_generation_exception(self, tmp_path): from fluid_build.cli.init import generate_dag_for_project - contract = {"name": "test-product"} - with patch("subprocess.run", side_effect=FileNotFoundError("no fluid")): + with patch(self.RUN, side_effect=OSError("disk full")): result = generate_dag_for_project( - tmp_path, contract, logging.getLogger("test"), MagicMock() + self._project(tmp_path, self.CONTRACT), + self.CONTRACT, + logging.getLogger("test"), + MagicMock(), ) assert result is False diff --git a/tests/test_init_extra_branches.py b/tests/test_init_extra_branches.py index 9d6cdfae..ce2d51e2 100644 --- a/tests/test_init_extra_branches.py +++ b/tests/test_init_extra_branches.py @@ -498,36 +498,56 @@ def test_exception_returns_1(self, _mock_dm, logger): class TestGenerateDagForProject: + # A build with a cron trigger makes the contract one that schedules something. + CONTRACT = { + "name": "ok-dag", + "builds": [{"id": "b", "execution": {"trigger": {"type": "schedule", "cron": "@hourly"}}}], + } + RUN = "fluid_build.cli._init_dag_helpers._run_generate_schedule" + + @staticmethod + def _project(tmp_path, contract): + """A project holding *contract* (JSON is YAML), as the helper reads it from disk.""" + import json + + (tmp_path / "contract.fluid.yaml").write_text(json.dumps(contract)) + return tmp_path + @patch("fluid_build.cli.init.RICH_AVAILABLE", False) - @patch("fluid_build.cli.init.create_dags_readme") @patch("fluid_build.cli.init.create_basic_dag") - def test_subprocess_failure_creates_basic(self, mock_basic, _mock_readme, tmp_path, logger): + def test_refused_contract_writes_no_basic_dag(self, mock_basic, tmp_path, logger): from fluid_build.cli.init import generate_dag_for_project - contract = {"name": "test-dag"} - with patch("subprocess.run") as mock_run: - mock_run.return_value = MagicMock(returncode=1) - result = generate_dag_for_project(tmp_path, contract, logger, None) - assert result is True - mock_basic.assert_called_once() + with patch(self.RUN, return_value=2): + result = generate_dag_for_project( + self._project(tmp_path, self.CONTRACT), self.CONTRACT, logger, None + ) + assert result is False + mock_basic.assert_not_called() @patch("fluid_build.cli.init.RICH_AVAILABLE", False) - @patch("fluid_build.cli.init.create_dags_readme") - def test_subprocess_success(self, _mock_readme, tmp_path, logger): + def test_success(self, tmp_path, logger): from fluid_build.cli.init import generate_dag_for_project - contract = {"name": "ok-dag", "orchestration": {"schedule": "@hourly"}} - with patch("subprocess.run") as mock_run: - mock_run.return_value = MagicMock(returncode=0) - result = generate_dag_for_project(tmp_path, contract, logger, None) + def write_dag(project_dir, dag_dir, logger): # what ``fluid generate schedule`` writes + dag_dir.mkdir() + (dag_dir / "ok_dag_dag.py").write_text("DAG(dag_id='ok_dag')\n") + return 0 + + with patch(self.RUN, side_effect=write_dag): + result = generate_dag_for_project( + self._project(tmp_path, self.CONTRACT), self.CONTRACT, logger, None + ) assert result is True @patch("fluid_build.cli.init.RICH_AVAILABLE", False) def test_exception_returns_false(self, tmp_path, logger): from fluid_build.cli.init import generate_dag_for_project - with patch("subprocess.run", side_effect=OSError("no subprocess")): - result = generate_dag_for_project(tmp_path, {}, logger, None) + with patch(self.RUN, side_effect=OSError("disk full")): + result = generate_dag_for_project( + self._project(tmp_path, self.CONTRACT), self.CONTRACT, logger, None + ) assert result is False diff --git a/tests/test_init_generate_schedule_dag.py b/tests/test_init_generate_schedule_dag.py new file mode 100644 index 00000000..740f5260 --- /dev/null +++ b/tests/test_init_generate_schedule_dag.py @@ -0,0 +1,279 @@ +# Copyright 2024-2026 Agentics Transformation Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""``fluid init``'s DAG helper writes what ``fluid generate schedule`` writes, or nothing. + +``generate_dag_for_project`` used to run ``fluid generate-airflow --output-dir ...``, +a flag that command never had: argparse exited 2, so every project got the basic +fallback DAG, or no DAG at all when ``fluid`` was not on PATH. Its tests mocked +``subprocess.run`` and nothing ever parsed that argv. The helper now runs +``fluid generate schedule`` in-process through the CLI's own parser, and writes no +DAG, saying why, when the contract schedules nothing or the command refuses it. +These tests parse its argv with the real parser and run the real command. +""" + +from __future__ import annotations + +import ast +import logging +import shlex +import shutil +from pathlib import Path +from typing import Dict, Tuple +from unittest.mock import patch + +import pytest +import yaml + +import fluid_build +from fluid_build.cli import build_parser +from fluid_build.cli import init as cli_init +from fluid_build.cli._init_dag_helpers import ( + _dag_facts, + _generate_schedule_argv, + create_basic_dag, + create_dags_readme, + generate_dag_for_project, + should_generate_dag, +) + +TEMPLATES = Path(fluid_build.__file__).parent / "templates" +LOG = logging.getLogger(__name__) + + +@pytest.fixture(autouse=True) +def _plain_console(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(cli_init, "RICH_AVAILABLE", False) + + +def _project(tmp_path: Path, template: str) -> Tuple[Path, dict]: + project = tmp_path / template + shutil.copytree(TEMPLATES / template, project) + return project, yaml.safe_load((project / "contract.fluid.yaml").read_text("utf-8")) + + +def _bash_commands(source: str) -> Dict[str, str]: + commands = {} + for node in ast.walk(ast.parse(source)): + if isinstance(node, ast.Call) and getattr(node.func, "id", "") == "BashOperator": + kwargs = {kw.arg: kw.value for kw in node.keywords} + commands[kwargs["task_id"].value] = kwargs["bash_command"].value + return commands + + +def test_the_helper_passes_generate_schedule_flags_the_cli_accepts(tmp_path: Path): + argv = _generate_schedule_argv(tmp_path, tmp_path / "dags") + args = build_parser().parse_args(argv) # exits 2 on a flag the command lacks + + assert args.func.__module__ == "fluid_build.cli.generate_schedule" + assert args.contract == str(tmp_path / "contract.fluid.yaml") + assert args.scheduler == "airflow" + assert args.output == str(tmp_path / "dags") + assert args.dag_contract_path == "contract.fluid.yaml" + assert args.overwrite is True + + +def test_the_argv_check_can_fail(): + """The argv init used to build. Parsing it exits 2, which the test above would catch.""" + with pytest.raises(SystemExit) as exit_info: + build_parser().parse_args( + ["generate-airflow", "contract.fluid.yaml", "--output-dir", "dags"] + ) + assert exit_info.value.code == 2 + + +@pytest.mark.parametrize( + "template, dag_file, dag_id, schedule", + [ + ( + "first-dag", + "sales_analytics_pipeline_dag.py", + "tutorial.sales_analytics_pipeline_v1__sales_analytics_pipeline", + "0 2 * * *", + ), + ( + "pipeline-orchestration", + "orchestrated_pipeline_dag.py", + "orchestration.complex_pipeline_v1__orchestrated_pipeline", + "0 3 * * *", + ), + ], +) +def test_a_scheduled_build_gets_its_fluid_apply_dag_and_a_readme_naming_it( + tmp_path: Path, template: str, dag_file: str, dag_id: str, schedule: str +): + project, contract = _project(tmp_path, template) + + assert generate_dag_for_project(project, contract, LOG, None) is True + + dags = project / "dags" + assert sorted(path.name for path in dags.iterdir()) == ["README.md", dag_file] + assert "Edit the contract, not this file" in (dags / dag_file).read_text("utf-8") + readme = (dags / "README.md").read_text("utf-8") + assert f"**DAG ID**: `{dag_id}`" in readme + assert f"**Schedule**: `{schedule}`" in readme + assert f"**File**: `{dag_file}`" in readme + assert "--overwrite" in readme # how to regenerate after editing the contract + assert "--auto-approve" not in readme + + +def test_a_manual_build_gets_no_dag(tmp_path: Path, caplog: pytest.LogCaptureFixture): + """customer-360's build is ``trigger.type: manual``, so a nightly DAG would only run a + placeholder task.""" + project, contract = _project(tmp_path, "customer-360") + + with caplog.at_level(logging.INFO): + assert generate_dag_for_project(project, contract, LOG, None) is False + + assert not (project / "dags").exists() + assert "No Airflow DAG written: the contract schedules nothing" in caplog.text + + +def test_a_continuous_build_gets_no_dag(tmp_path: Path, caplog: pytest.LogCaptureFixture): + contract = { + "fluidVersion": "0.7.6", + "kind": "DataProduct", + "id": "streams.clicks", + "name": "clicks", + "orchestration": {"engine": "airflow"}, + "builds": [ + { + "id": "sink", + "pattern": "acquisition", + "engine": "kafka-connect", + "execution": {"trigger": {"type": "continuous"}}, + } + ], + } + (tmp_path / "contract.fluid.yaml").write_text(yaml.safe_dump(contract), "utf-8") + + with caplog.at_level(logging.INFO): + assert generate_dag_for_project(tmp_path, contract, LOG, None) is False + + assert not (tmp_path / "dags").exists() + assert "No Airflow DAG written: the contract schedules nothing" in caplog.text + + +def test_a_refused_contract_gets_no_dag_and_says_why( + tmp_path: Path, caplog: pytest.LogCaptureFixture +): + """No stand-in DAG on a schedule the contract never declared: it would hide the error.""" + project, contract = _project(tmp_path, "first-dag") + path = project / "contract.fluid.yaml" + text = path.read_text("utf-8").replace("schedule: 0 2 * * *", "schedule: 61 25 * * *") + assert "61 25" in text + path.write_text(text, "utf-8") + + with patch.object(cli_init, "create_basic_dag") as basic, caplog.at_level(logging.INFO): + assert generate_dag_for_project(project, yaml.safe_load(text), LOG, None) is False + + basic.assert_not_called() + assert not (project / "dags").exists() + assert "'61' is not a valid minute field" in caplog.text # the command's own message + assert "refused the contract (exit 2)" in caplog.text + + +def test_a_second_run_regenerates_the_dag(tmp_path: Path): + project, contract = _project(tmp_path, "first-dag") + + assert generate_dag_for_project(project, contract, LOG, None) is True + assert generate_dag_for_project(project, contract, LOG, None) is True + + dags = sorted(path.name for path in (project / "dags").iterdir()) + assert dags == ["README.md", "sales_analytics_pipeline_dag.py"] + + +def test_the_readme_names_only_the_dags_this_run_wrote(tmp_path: Path): + project, contract = _project(tmp_path, "first-dag") + (project / "dags").mkdir() + (project / "dags" / "helpers.py").write_text("X = 1\n", "utf-8") + + assert generate_dag_for_project(project, contract, LOG, None) is True + + readme = (project / "dags" / "README.md").read_text("utf-8") + assert "`sales_analytics_pipeline_dag.py`" in readme + assert "helpers.py" not in readme + + +def test_a_relative_project_dir_works(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + """``demo_mode`` passes ``Path(project_name)``, relative to the current directory.""" + monkeypatch.chdir(tmp_path) + _, contract = _project(tmp_path, "first-dag") + + assert generate_dag_for_project(Path("first-dag"), contract, LOG, None) is True + + assert (tmp_path / "first-dag" / "dags" / "sales_analytics_pipeline_dag.py").is_file() + + +def test_a_singular_build_key_is_judged_as_the_command_reads_it(tmp_path: Path): + """The schema's legacy ``build:`` holds one build, and the command's loader turns it + into ``builds``: the helper must not call that contract unscheduled.""" + project, contract = _project(tmp_path, "first-dag") + (build,) = contract.pop("builds") + contract["build"] = build + (project / "contract.fluid.yaml").write_text(yaml.safe_dump(contract), "utf-8") + + assert should_generate_dag(contract) is True + assert generate_dag_for_project(project, contract, LOG, None) is True + assert (project / "dags" / "sales_analytics_pipeline_dag.py").is_file() + + +def test_should_generate_dag_for_a_scheduled_build_only(): + def build(trigger: dict) -> dict: + return {"builds": [{"id": "b", "execution": {"trigger": trigger}}]} + + assert should_generate_dag(build({"type": "schedule", "cron": "0 2 * * *"})) is True + assert should_generate_dag(build({"type": "manual"})) is False + + +def test_the_readme_lists_each_dag_on_its_own(tmp_path: Path): + create_dags_readme(tmp_path, "a", "@daily", "a_dag.py", others=[("b", "0 3 * * *", "b_dag.py")]) + + readme = (tmp_path / "README.md").read_text("utf-8") + assert "## Generated DAGs" in readme + assert "- **DAG ID**: `a`\n- **Schedule**: `@daily`\n- **File**: `a_dag.py`" in readme + assert "- **DAG ID**: `b`\n- **Schedule**: `0 3 * * *`\n- **File**: `b_dag.py`" in readme + + +def test_the_basic_dag_quotes_the_project_path_for_python_and_shell(tmp_path: Path): + project = tmp_path / 'o\'brien $(touch pwned) "data"' + project.mkdir() + + create_basic_dag(project, {"name": "orders"}, LOG) + + commands = _bash_commands((project / "dags" / "orders_dag.py").read_text("utf-8")) + assert sorted(commands) == ["apply_contract", "plan_execution", "validate_contract"] + for command in commands.values(): + assert shlex.split(command)[:4] == ["cd", str(project.absolute()), "&&", "fluid"] + assert shlex.split(commands["apply_contract"])[4:] == ["apply", "contract.fluid.yaml", "--yes"] + + +@pytest.mark.parametrize( + "source, facts", + [ + ("dag = DAG(dag_id='a', schedule='@daily')\n", ("a", "@daily")), + # The ``fluid apply`` DAG names its schedule through a module constant. + ( + "SCHEDULE = '0 2 * * *'\nwith DAG(dag_id='b', schedule=SCHEDULE):\n pass\n", + ("b", "0 2 * * *"), + ), + ("import airflow.sdk as sdk\nsdk.DAG(dag_id='c', schedule=None)\n", ("c", "")), + ("dag = DAG(dag_id='d', schedule='@hourly')\n", ("d", "@hourly")), # a BOM + ("this is not python", ("", "")), + ], +) +def test_dag_facts(tmp_path: Path, source: str, facts: Tuple[str, str]): + path = tmp_path / "x_dag.py" + path.write_text(source, "utf-8") + assert _dag_facts(path) == facts From e9bffb0e2d811d0fdb3ceeefe47fbb81698803c4 Mon Sep 17 00:00:00 2001 From: Speculator55005 <50082482+fas89@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:19:20 +0200 Subject: [PATCH 3/3] docs(examples): regenerate financial_analytics_dag.py in the Airflow 2.7+/3 shape The example was AwsProvider.export output from before the Airflow 3 work: DAG(schedule_interval=...) and bare Airflow 2 imports, which fail to import on Airflow 3, and helper bodies hollowed out by hand. Regenerated with the call examples/generate_example_dags.py makes (AwsProvider.export, engine airflow) and black-formatted like the file it replaces: schedule=, and DAG and PythonOperator imported from airflow.sdk and apache-airflow-providers-standard with the Airflow 2 fallback. It loads in a real DagBag on Airflow 3.2.2 and 2.10.5 with the amazon provider installed. --- .../generated_dags/financial_analytics_dag.py | 145 +++++++++++++----- 1 file changed, 106 insertions(+), 39 deletions(-) diff --git a/examples/generated_dags/financial_analytics_dag.py b/examples/generated_dags/financial_analytics_dag.py index 86475b88..9e55e004 100644 --- a/examples/generated_dags/financial_analytics_dag.py +++ b/examples/generated_dags/financial_analytics_dag.py @@ -8,33 +8,30 @@ AWS Region: us-east-1 Auto-generated by FLUID Forge - DO NOT EDIT MANUALLY -Generated: 2026-01-30T14:02:10.320788 +Generated: 2026-10-21T21:45:23.515208 """ -import json -import logging from datetime import datetime, timedelta from zoneinfo import ZoneInfo -from airflow import DAG -from airflow.operators.python import PythonOperator -from airflow.providers.amazon.aws.operators.athena import AthenaOperator +try: # Airflow 3 + from airflow.sdk import DAG +except ImportError: # Airflow 2.x + from airflow import DAG +try: # Airflow 3, or Airflow 2 with apache-airflow-providers-standard + from airflow.providers.standard.operators.python import PythonOperator +except ImportError: # Airflow 2.x core + from airflow.operators.python import PythonOperator +from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator from airflow.providers.amazon.aws.operators.glue import GlueJobOperator -from airflow.providers.amazon.aws.operators.lambda_function import LambdaInvokeFunctionOperator +from airflow.providers.amazon.aws.operators.athena import AthenaOperator from airflow.providers.amazon.aws.operators.redshift_data import RedshiftDataOperator -from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator +from airflow.providers.amazon.aws.operators.lambda_function import LambdaInvokeFunctionOperator +import json +import logging logger = logging.getLogger(__name__) - -def _ensure_glue_database(**kwargs): - return kwargs - - -def _ensure_glue_table(**kwargs): - return kwargs - - # Default DAG arguments default_args = { "owner": "fluid-forge", @@ -51,12 +48,100 @@ def _ensure_glue_table(**kwargs): dag_id="financial_analytics", default_args=default_args, description="Financial Analytics Pipeline", - schedule_interval="0 2 * * *", + schedule="0 2 * * *", start_date=datetime(2026, 1, 1, tzinfo=ZoneInfo("UTC")), catchup=False, tags=["fluid", "auto-generated", "financial_analytics"], ) + +# Helper functions + + +def _ensure_glue_database(database: str, region: str): + """Ensure Glue database exists.""" + import boto3 + + client = boto3.client("glue", region_name=region) + + try: + client.create_database(DatabaseInput={"Name": database}) + logger.info(f"Created Glue database: {database}") + except client.exceptions.AlreadyExistsException: + logger.info(f"Glue database already exists: {database}") + + +def _ensure_glue_table(database: str, table: str, params: dict, region: str): + """Ensure Glue table exists.""" + import boto3 + + client = boto3.client("glue", region_name=region) + + columns = params.get("columns", []) + location = params.get("location", "") + input_format = params.get("input_format", "parquet") + + FORMAT_MAP = { + "parquet": { + "InputFormat": "org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat", + "OutputFormat": "org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat", + "SerDe": "org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe", + }, + "orc": { + "InputFormat": "org.apache.hadoop.hive.ql.io.orc.OrcInputFormat", + "OutputFormat": "org.apache.hadoop.hive.ql.io.orc.OrcOutputFormat", + "SerDe": "org.apache.hadoop.hive.ql.io.orc.OrcSerde", + }, + "json": { + "InputFormat": "org.apache.hadoop.mapred.TextInputFormat", + "OutputFormat": "org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat", + "SerDe": "org.openx.data.jsonserde.JsonSerDe", + }, + "csv": { + "InputFormat": "org.apache.hadoop.mapred.TextInputFormat", + "OutputFormat": "org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat", + "SerDe": "org.apache.hadoop.hive.serde2.OpenCSVSerde", + }, + } + fmt = FORMAT_MAP.get(input_format, FORMAT_MAP["parquet"]) + + table_input = { + "Name": table, + "StorageDescriptor": { + "Columns": columns, + "Location": location, + "InputFormat": fmt["InputFormat"], + "OutputFormat": fmt["OutputFormat"], + "SerdeInfo": {"SerializationLibrary": fmt["SerDe"]}, + }, + "TableType": "EXTERNAL_TABLE", + } + description = params.get("description") + if description: + table_input["Description"] = description + + try: + client.get_table(DatabaseName=database, Name=table) + client.update_table(DatabaseName=database, TableInput=table_input) + logger.info(f"Updated Glue table: {database}.{table}") + except client.exceptions.EntityNotFoundException: + client.create_table(DatabaseName=database, TableInput=table_input) + logger.info(f"Created Glue table: {database}.{table}") + + +def _execute_provider_action(action: str, params: dict, account_id: str, region: str): + """Provider-action tasks are retired. + + forge-cli compiles the contract to OpenTofu and runs ``tofu`` — a + generated DAG no longer performs native cloud CRUD. Provision + infrastructure with ``fluid apply`` (the OpenTofu engine). + """ + raise RuntimeError( + f"provider-action task {action!r} is retired: provision infrastructure " + "with `fluid apply` (the OpenTofu engine), then re-generate this DAG" + ) + + # Task definitions create_bucket = S3CreateBucketOperator( task_id="create_bucket", @@ -79,27 +164,9 @@ def _ensure_glue_table(**kwargs): python_callable=lambda: _ensure_glue_table( database="financial_analytics", table="transactions", - params={ - "database": "financial_analytics", - "table": "transactions", - "s3_location": "s3://financial-data-lake-{{account_id}}/transactions/", - "format": "parquet", - "schema": { - "columns": [ - {"name": "transaction_id", "type": "string"}, - {"name": "customer_id", "type": "string"}, - {"name": "amount", "type": "double"}, - {"name": "currency", "type": "string"}, - {"name": "timestamp", "type": "timestamp"}, - {"name": "category", "type": "string"}, - ], - "partition_keys": [ - {"name": "year", "type": "int"}, - {"name": "month", "type": "int"}, - {"name": "day", "type": "int"}, - ], - }, - }, + params=json.loads( + '{"database": "financial_analytics", "format": "parquet", "s3_location": "s3://financial-data-lake-{{account_id}}/transactions/", "schema": {"columns": [{"name": "transaction_id", "type": "string"}, {"name": "customer_id", "type": "string"}, {"name": "amount", "type": "double"}, {"name": "currency", "type": "string"}, {"name": "timestamp", "type": "timestamp"}, {"name": "category", "type": "string"}], "partition_keys": [{"name": "year", "type": "int"}, {"name": "month", "type": "int"}, {"name": "day", "type": "int"}]}, "table": "transactions"}' + ), region="us-east-1", ), dag=dag,