diff --git a/CHANGELOG.md b/CHANGELOG.md index 791f1bb2..fbbfc1d7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -183,6 +183,27 @@ 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. +- **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/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, 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/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 + ) 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