diff --git a/CHANGELOG.md b/CHANGELOG.md index c150fed7..791f1bb2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -64,6 +64,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - An environment overlay cannot turn a scheduled build continuous: overlays cannot delete the base's `cron` or `schedule`, which a continuous trigger refuses. Declare the trigger in each environment's overlay instead of in the base. +- Generated Airflow DAGs run Snowflake SQL through `SQLExecuteQueryOperator` + (apache-airflow-providers-common-sql, which the snowflake provider installs) with + `conn_id`, because apache-airflow-providers-snowflake 6.0 removed `SnowflakeOperator`; + a DAG generated before needs regenerating to load. A GCP `create_table` task uses + `BigQueryCreateTableOperator` where the google provider has it (13+; 17.0 removed the + old one), and `BigQueryCreateEmptyTableOperator` before that. +- Both new operators template more of their arguments than the ones they replace: + `SQLExecuteQueryOperator` renders `conn_id` and `hook_params` (the warehouse and + database), and `BigQueryCreateTableOperator` renders `table_resource` (the schema). A + literal `{{`, `{%` or `{#` in a warehouse or database name, or in a schema field such + as a description, is now rendered as Jinja when the task runs, and fails it. Escape it + as `{{ '{{' }}`, `{{ '{%' }}` or `{{ '{#' }}`, or wrap the value in `{% raw %}` ... + `{% endraw %}`. ### Added @@ -149,6 +162,28 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 so an escape sequence in a contract or an error message can no longer clear the terminal or plant a link. JSON and JUnit output are unchanged. +### Fixed + +- **Generated Airflow DAGs load on Airflow 3.** `fluid export --engine airflow`, `fluid + generate schedule`, `fluid init` and the AWS, GCP and Snowflake code generators emitted + `DAG(schedule_interval=...)` and `start_date=days_ago(1)`, which Airflow 3 removed, so + their DAGs failed to import. They now pass `schedule=` with a fixed `start_date` and + `catchup=False`. Every generator, `fluid scaffold-composer` and `fluid schedule-sync` + included, imports `DAG`, `dag`/`task` and the core operators and sensors from their + Airflow 3 homes (`airflow.sdk`, apache-airflow-providers-standard), falling back to the + Airflow 2 paths, so the same DAG loads on Airflow 2.7+ and 3.x without the deprecation + warnings Airflow 3 gives the old import paths. +- The AWS TaskFlow DAG wired its `@task` functions (`load >> publish`), which raised + `TypeError` when Airflow parsed it; it now wires each task's result, and an operator + task as itself. +- The GCP DAG imports the Dataflow operator only when a task uses it: its module needs + apache-airflow-providers-apache-beam, and the unused import stopped every GCP DAG from + loading where that provider was absent. +- `fluid init`'s basic DAG spliced `orchestration.retries` and `retry_delay` into the DAG + 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. + ## [0.21.0] — 2026-10-19 ### Upgrade notes diff --git a/fluid_build/cli/_acquisition_stage_ext.py b/fluid_build/cli/_acquisition_stage_ext.py index 936ea534..e3d7bd77 100644 --- a/fluid_build/cli/_acquisition_stage_ext.py +++ b/fluid_build/cli/_acquisition_stage_ext.py @@ -59,6 +59,7 @@ from fluid_build.build_runners._ids import ( validate_identifier as _validate_identifier, ) +from fluid_build.providers.common.codegen_utils import airflow_dag_import, airflow_standard_import from fluid_build.util.contract import is_unscheduled_build LOG = logging.getLogger("fluid.acquire.stage_ext") @@ -583,9 +584,10 @@ def _render_airflow_dag( contract_path = f"contracts/{contract_id}.fluid.yaml" return ( '"""Auto-generated by fluid schedule-sync. Edit the contract, not this file."""\n\n' - "from datetime import datetime, timedelta\n" - "from airflow import DAG\n" - "from airflow.operators.bash import BashOperator\n\n" + "from datetime import datetime, timedelta\n\n" + + airflow_dag_import() + + airflow_standard_import("BashOperator") + + "\n" "default_args = {\n" ' "owner": "fluid",\n' ' "depends_on_past": False,\n' diff --git a/fluid_build/cli/_init_dag_helpers.py b/fluid_build/cli/_init_dag_helpers.py index 1ed48d4c..406866b3 100644 --- a/fluid_build/cli/_init_dag_helpers.py +++ b/fluid_build/cli/_init_dag_helpers.py @@ -154,12 +154,25 @@ def create_basic_dag(project_dir: Path, contract: dict, logger): import re as _re + from fluid_build.providers.common.codegen_utils import ( + airflow_dag_import, + airflow_standard_import, + ) + contract_name = contract.get("name", "my_product") orchestration = contract.get("orchestration", {}) # Sanitize values to prevent code injection in generated Python. dag_id = _re.sub(r"[^a-zA-Z0-9_]", "_", contract_name)[:128] schedule = _re.sub(r"[^a-zA-Z0-9@*/, _-]", "", orchestration.get("schedule", "@daily"))[:64] contract_name = _re.sub(r"[^a-zA-Z0-9 _.-]", "_", contract_name)[:128] + # Numbers, never text: both are interpolated into the DAG's Python source, so a + # contract value such as `3, __import__('os').system('...')` would run at parse time. + try: + retries = max(0, min(int(orchestration.get("retries", 3)), 100)) + except (TypeError, ValueError, OverflowError): # OverflowError: `retries: .inf` + retries = 3 + delay = _re.fullmatch(r"\s*(\d{1,4})\s*m?\s*", str(orchestration.get("retry_delay", "5m"))) + retry_minutes = int(delay.group(1)) if delay else 5 dag_content = f'''""" Airflow DAG for FLUID contract: {contract_name} @@ -167,24 +180,23 @@ def create_basic_dag(project_dir: Path, contract: dict, logger): """ from datetime import datetime, timedelta -from airflow import DAG -from airflow.operators.bash import BashOperator +{airflow_dag_import()}{airflow_standard_import("BashOperator")} default_args = {{ 'owner': 'fluid', 'depends_on_past': False, 'start_date': datetime(2024, 1, 1), 'email_on_failure': False, 'email_on_retry': False, - 'retries': {orchestration.get("retries", 3)}, - 'retry_delay': timedelta(minutes={orchestration.get("retry_delay", "5m").replace("m", "")}), + 'retries': {retries}, + 'retry_delay': timedelta(minutes={retry_minutes}), }} with DAG( dag_id='{dag_id}', default_args=default_args, description='FLUID data product: {contract_name}', - schedule_interval='{schedule}', + schedule='{schedule}', catchup=False, tags=['fluid', 'data-product'], ) as dag: diff --git a/fluid_build/cli/scaffold_composer.py b/fluid_build/cli/scaffold_composer.py index a7553e9f..b45b616d 100644 --- a/fluid_build/cli/scaffold_composer.py +++ b/fluid_build/cli/scaffold_composer.py @@ -22,7 +22,12 @@ from typing import Any, Dict, List from .._console import cprint -from ..providers.common.codegen_utils import py_str_literal, sanitize_identifier +from ..providers.common.codegen_utils import ( + airflow_dag_import, + airflow_standard_import, + py_str_literal, + sanitize_identifier, +) from ..util.contract import has_only_unscheduled_builds, is_unscheduled_build from ._acquisition_stage_ext import _validate_cron as _validate_cron_grammar from ._common import CLIError, load_contract_with_overlay @@ -39,10 +44,13 @@ def register(subparsers: argparse._SubParsersAction): p.set_defaults(cmd=COMMAND, func=run) -DAG_TMPL = """from __future__ import annotations -from airflow import DAG -from airflow.operators.bash import BashOperator -from datetime import datetime +# Airflow 2.7+ and 3.x: see codegen_utils.airflow_dag_import. The import blocks hold +# no braces, so the template stays safe for str.format. +DAG_TMPL = ( + "from __future__ import annotations\n" + + airflow_dag_import() + + airflow_standard_import("BashOperator") + + """from datetime import datetime with DAG( dag_id={dag_id}, start_date=datetime(2024,1,1), @@ -65,6 +73,7 @@ def register(subparsers: argparse._SubParsersAction): ) validate >> plan >> apply """ +) def _validate_cron(raw: object) -> str: diff --git a/fluid_build/providers/aws/codegen/airflow.py b/fluid_build/providers/aws/codegen/airflow.py index da448947..63aadd77 100644 --- a/fluid_build/providers/aws/codegen/airflow.py +++ b/fluid_build/providers/aws/codegen/airflow.py @@ -32,6 +32,8 @@ from typing import Any, Dict, List, Optional from fluid_build.providers.common.codegen_utils import ( + airflow_dag_import, + airflow_standard_import, escape_for_docstring, json_literal, py_str_literal, @@ -41,11 +43,20 @@ logger = logging.getLogger(__name__) -# Module-level names the generated DAG binds itself. A task variable must not -# collide with any of them. +# Module-level names the generated DAG binds itself, including every name its +# imports bind (the TaskFlow ``dag`` / ``task`` decorators and the operator +# classes). A task variable must not collide with any of them: ``taskId: task`` +# emitted ``def task():`` inside the DAG function, which shadows the decorator. _EMITTED_MODULE_NAMES = frozenset( { "dag", + "task", + "PythonOperator", + "S3CreateBucketOperator", + "GlueJobOperator", + "AthenaOperator", + "RedshiftDataOperator", + "LambdaInvokeFunctionOperator", "default_args", "logger", "json", @@ -187,11 +198,14 @@ def _generate_dag_header( def _generate_dag_imports() -> str: """Generate import statements.""" - return """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.s3 import S3CreateBucketOperator + return ( + ( + "from datetime import datetime, timedelta\n" + "from zoneinfo import ZoneInfo\n" + + airflow_dag_import() + + airflow_standard_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.athena import AthenaOperator from airflow.providers.amazon.aws.operators.redshift_data import RedshiftDataOperator @@ -200,6 +214,7 @@ def _generate_dag_imports() -> str: import logging logger = logging.getLogger(__name__)""" + ) def _generate_dag_definition( @@ -227,7 +242,7 @@ def _generate_dag_definition( dag_id={py_str_literal(_sanitize_dag_id(contract_id))}, default_args=default_args, description={py_str_literal(contract_name)}, - schedule_interval={py_str_literal(airflow_schedule)}, + schedule={py_str_literal(airflow_schedule)}, start_date=datetime(2026, 1, 1, tzinfo=ZoneInfo({py_str_literal(timezone)})), catchup=False, tags={tags!r}, @@ -686,7 +701,10 @@ def _generate_taskflow_imports() -> str: """Generate import statements for TaskFlow API.""" return """from datetime import datetime, timedelta from zoneinfo import ZoneInfo -from airflow.decorators import dag, task +try: # Airflow 3 + from airflow.sdk import dag, task +except ImportError: # Airflow 2.x + from airflow.decorators import dag, task from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator from airflow.providers.amazon.aws.operators.athena import AthenaOperator from airflow.providers.amazon.aws.operators.redshift_data import RedshiftDataOperator @@ -736,6 +754,20 @@ def _generate_taskflow_dag( # Build per-task variable names for dependency wiring task_vars: Dict[str, str] = _unique_task_identifiers(tasks) + # What a dependency edge connects: the operator itself, or a @task's call result + # (``_result``). A @task function is not a task: ``fn >> fn`` raises TypeError. + nodes: Dict[str, str] = {} + # A @task's result gets a name of its own, allocated against every task variable and + # every other result: a task named ``chk_result`` and ``chk``'s result must not share one. + taken = set(task_vars.values()) | set(_EMITTED_MODULE_NAMES) + + def _result_name(var: str) -> str: + candidate, suffix = f"{var}_result", 2 + while candidate in taken: + candidate = f"{var}_result_{suffix}" + suffix += 1 + taken.add(candidate) + return candidate for t in tasks: tid = t.get("taskId", "") @@ -755,6 +787,7 @@ def _generate_taskflow_dag( lines.append(f" region_name={py_str_literal(params.get('region', region))},") lines.append(" aws_conn_id='aws_default',") lines.append(" )") + nodes[tid] = var elif service == "athena" and operation == "execute_query": query = params.get("query", "") db = params.get("database", "default") @@ -769,6 +802,7 @@ def _generate_taskflow_dag( lines.append(f" output_location={py_str_literal(out)},") lines.append(" aws_conn_id='aws_default',") lines.append(" )") + nodes[tid] = var elif service == "glue": # Glue operations use @task decorators lines.append("") @@ -792,7 +826,9 @@ def _generate_taskflow_dag( f"json.loads({json_literal(params)}), {py_str_literal(account_id)}, " f"{py_str_literal(region)})" ) - lines.append(f" {var}_result = {var}()") + result = _result_name(var) + lines.append(f" {result} = {var}()") + nodes[tid] = result else: # Generic @task fallback lines.append("") @@ -803,17 +839,19 @@ def _generate_taskflow_dag( f"json.loads({json_literal(params)}), {py_str_literal(account_id)}, " f"{py_str_literal(region)})" ) - lines.append(f" {var}_result = {var}()") + result = _result_name(var) + lines.append(f" {result} = {var}()") + nodes[tid] = result # Dependencies for t in tasks: tid = t.get("taskId", "") depends_on = t.get("dependsOn", []) if depends_on: - var = task_vars[tid] + node = nodes[tid] for dep in depends_on: - dep_var = task_vars.get(dep) - if dep_var is None: + dep_node = nodes.get(dep) + if dep_node is None: logger.warning( "aws airflow: task %r dependsOn %r, which is not a " "provider_action task in this DAG — edge skipped", @@ -822,7 +860,7 @@ def _generate_taskflow_dag( ) lines.append(f" # dependency on non-DAG task {dep!r} skipped") continue - lines.append(f" {dep_var} >> {var}") + lines.append(f" {dep_node} >> {node}") lines.append("") lines.append("") diff --git a/fluid_build/providers/common/codegen_utils.py b/fluid_build/providers/common/codegen_utils.py index 5f70f4be..ea165c95 100644 --- a/fluid_build/providers/common/codegen_utils.py +++ b/fluid_build/providers/common/codegen_utils.py @@ -166,7 +166,7 @@ def convert_schedule_to_cron(schedule: str) -> str: def convert_schedule_to_airflow(schedule: str) -> str: """ - Convert FLUID schedule to Airflow schedule_interval format. + Convert FLUID schedule to an Airflow ``schedule`` value. Airflow accepts both cron and special keywords like @daily. @@ -174,7 +174,7 @@ def convert_schedule_to_airflow(schedule: str) -> str: schedule: FLUID schedule string Returns: - Airflow schedule_interval value + Airflow ``schedule`` value (``DAG(schedule=...)``) """ schedule_lower = schedule.lower().strip() @@ -190,6 +190,45 @@ def convert_schedule_to_airflow(schedule: str) -> str: return "@daily" +# ── Generated DAGs run on Airflow 2.7+ and 3.x ───────────────────────────── +# +# Airflow 3 removed ``DAG(schedule_interval=...)`` and ``airflow.utils.dates.days_ago``, +# moved the core operators and sensors to apache-airflow-providers-standard, and +# exports ``DAG`` from ``airflow.sdk``. A generated DAG therefore passes +# ``schedule=`` (Airflow 2.4+), a fixed timezone-aware ``start_date`` with +# ``catchup=False``, and imports each name from its Airflow 3 home first, falling +# back to the Airflow 2 core path: the same shape schedulers/airflow/fluid_apply.py +# already emits. The import paths are Airflow's own deprecation map +# (airflow-core/src/airflow/operators/__init__.py and sensors/__init__.py at 3.2.2). +_AIRFLOW_STANDARD_MODULES = { + "BashOperator": "operators.bash", + "PythonOperator": "operators.python", + "ExternalTaskSensor": "sensors.external_task", +} + + +def airflow_dag_import(indent: str = "") -> str: + """``DAG`` from ``airflow.sdk`` on Airflow 3, from ``airflow`` on Airflow 2.""" + return ( + f"{indent}try: # Airflow 3\n" + f"{indent} from airflow.sdk import DAG\n" + f"{indent}except ImportError: # Airflow 2.x\n" + f"{indent} from airflow import DAG\n" + ) + + +def airflow_standard_import(name: str, indent: str = "") -> str: + """``name`` from apache-airflow-providers-standard (Airflow 3, or Airflow 2 with the + provider installed), else from Airflow 2's core module.""" + module = _AIRFLOW_STANDARD_MODULES[name] + return ( + f"{indent}try: # Airflow 3, or Airflow 2 with apache-airflow-providers-standard\n" + f"{indent} from airflow.providers.standard.{module} import {name}\n" + f"{indent}except ImportError: # Airflow 2.x core\n" + f"{indent} from airflow.{module} import {name}\n" + ) + + def json_literal(value: Any) -> str: """Emit ``value`` as a Python *string literal* holding its JSON encoding. diff --git a/fluid_build/providers/gcp/codegen/airflow.py b/fluid_build/providers/gcp/codegen/airflow.py index a42c9106..c556f8f0 100644 --- a/fluid_build/providers/gcp/codegen/airflow.py +++ b/fluid_build/providers/gcp/codegen/airflow.py @@ -21,9 +21,11 @@ from datetime import datetime from datetime import timezone as dt_timezone -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from fluid_build.providers.common.codegen_utils import ( + airflow_dag_import, + airflow_standard_import, convert_schedule_to_airflow, generate_file_header, generate_task_dependencies_code, @@ -72,7 +74,7 @@ def generate_airflow_dag(contract: Dict[str, Any], project: str, region: str) -> gcp_region=region, ) dag_code += "\n\n" - dag_code += _generate_dag_imports() + dag_code += _generate_dag_imports(provider_tasks) dag_code += "\n\n" dag_code += _generate_dag_definition(contract_id, contract_name, schedule, timezone) dag_code += "\n\n" @@ -101,17 +103,33 @@ def _generate_dag_header( """''' -def _generate_dag_imports() -> str: - """Generate import statements.""" - return """from datetime import datetime, timedelta -from zoneinfo import ZoneInfo -from airflow import DAG -from airflow.operators.python import PythonOperator -from airflow.providers.google.cloud.operators.bigquery import ( +def _is_dataflow_template(task: Dict[str, Any]) -> bool: + """The test ``_generate_single_task`` dispatches on: ``.dataflow.run_template``.""" + parts = str(task.get("action", "")).split(".") + return len(parts) >= 3 and parts[1] == "dataflow" and parts[2] == "run_template" + + +def _generate_dag_imports(tasks: Optional[List[Dict[str, Any]]] = None) -> str: + """Generate import statements. + + The Dataflow operator is imported only when a task uses it: its module needs + apache-airflow-providers-apache-beam, and an unused import that fails stops the + whole DAG from loading.""" + dataflow = any(_is_dataflow_template(t) for t in tasks or []) + return ( + "from datetime import datetime, timedelta\n" + "from zoneinfo import ZoneInfo\n" + + airflow_dag_import() + + airflow_standard_import("PythonOperator") + + """from airflow.providers.google.cloud.operators.bigquery import ( BigQueryCreateEmptyDatasetOperator, - BigQueryCreateEmptyTableOperator, BigQueryInsertJobOperator, ) +try: # apache-airflow-providers-google 13+ (17.0 removed BigQueryCreateEmptyTableOperator) + from airflow.providers.google.cloud.operators.bigquery import BigQueryCreateTableOperator +except ImportError: # older google providers (Airflow 2) + BigQueryCreateTableOperator = None + from airflow.providers.google.cloud.operators.bigquery import BigQueryCreateEmptyTableOperator from airflow.providers.google.cloud.operators.gcs import ( GCSCreateBucketOperator, GCSDeleteBucketOperator, @@ -120,13 +138,19 @@ def _generate_dag_imports() -> str: PubSubCreateTopicOperator, PubSubCreateSubscriptionOperator, ) -from airflow.providers.google.cloud.operators.dataflow import ( - DataflowTemplatedJobStartOperator, -) -import json +""" + + ( + "from airflow.providers.google.cloud.operators.dataflow import (\n" + " DataflowTemplatedJobStartOperator,\n" + ")\n" + if dataflow + else "" + ) + + """import json import logging logger = logging.getLogger(__name__)""" + ) def _generate_dag_definition( @@ -161,7 +185,7 @@ def _generate_dag_definition( dag_id={dag_id_lit}, default_args=default_args, description={contract_name_lit}, - schedule_interval={airflow_schedule_lit}, + schedule={airflow_schedule_lit}, start_date=datetime(2026, 1, 1, tzinfo=ZoneInfo({timezone_lit})), catchup=False, tags={tags!r}, @@ -233,14 +257,26 @@ def _generate_bigquery_task( dataset_lit = py_str_literal(params.get("dataset_id", "unknown_dataset")) table_lit = py_str_literal(params.get("table_id", "unknown_table")) schema_fields = params.get("schema", []) - return f"""{var} = BigQueryCreateEmptyTableOperator( - task_id={task_id_lit}, - dataset_id={dataset_lit}, - table_id={table_lit}, - project_id={project_lit}, - schema_fields={schema_fields!r}, - dag=dag, -)""" + # BigQueryCreateTableOperator where the google provider has it (13+), the + # BigQueryCreateEmptyTableOperator it replaced (removed in 17.0) before that. + return f"""if BigQueryCreateTableOperator is not None: + {var} = BigQueryCreateTableOperator( + task_id={task_id_lit}, + dataset_id={dataset_lit}, + table_id={table_lit}, + project_id={project_lit}, + table_resource={{"schema": {{"fields": {schema_fields!r}}}}}, + dag=dag, + ) +else: + {var} = BigQueryCreateEmptyTableOperator( + task_id={task_id_lit}, + dataset_id={dataset_lit}, + table_id={table_lit}, + project_id={project_lit}, + schema_fields={schema_fields!r}, + dag=dag, + )""" elif operation == "query" or operation == "run_query": query_sql = params.get("query", "SELECT 1") diff --git a/fluid_build/providers/snowflake/codegen/airflow.py b/fluid_build/providers/snowflake/codegen/airflow.py index d46e639e..9e0874af 100644 --- a/fluid_build/providers/snowflake/codegen/airflow.py +++ b/fluid_build/providers/snowflake/codegen/airflow.py @@ -30,6 +30,8 @@ from fluid_build.providers._sql_safety import validate_ident, validate_sql_type_name from fluid_build.providers.common.codegen_utils import ( + airflow_dag_import, + airflow_standard_import, escape_for_docstring, py_str_literal, sanitize_identifier, @@ -125,16 +127,20 @@ def _generate_dag_header( def _generate_dag_imports() -> str: """Generate import statements.""" - return """from datetime import datetime, timedelta -from zoneinfo import ZoneInfo -from airflow import DAG -from airflow.operators.python import PythonOperator -from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator + return ( + ( + "from datetime import datetime, timedelta\n" + "from zoneinfo import ZoneInfo\n" + + airflow_dag_import() + + airflow_standard_import("PythonOperator") + ) + + """from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook import json import logging logger = logging.getLogger(__name__)""" + ) def _generate_dag_definition( @@ -169,7 +175,7 @@ def _generate_dag_definition( dag_id={dag_id_lit}, default_args=default_args, description={description_lit}, - schedule_interval={schedule_lit}, + schedule={schedule_lit}, start_date=datetime(2026, 1, 1, tzinfo=ZoneInfo({timezone_lit})), catchup=False, tags={tags!r}, @@ -218,7 +224,7 @@ def _generate_snowflake_sql_task( """Generate Snowflake SQL task code. SQL-injection safety: ``create_database`` / ``create_schema`` / - ``create_table`` build DDL that the generated ``SnowflakeOperator`` runs at + ``create_table`` build DDL that the generated ``SQLExecuteQueryOperator`` runs at DAG *runtime*, so ``py_str_literal`` (which only neutralises the Python layer) is not enough on its own. Every IDENTIFIER (database / schema / table / column name) is routed through ``validate_ident`` and every column @@ -268,12 +274,13 @@ def _generate_snowflake_sql_task( warehouse_lit = py_str_literal(warehouse) database_lit = py_str_literal(database) - return f"""{var} = SnowflakeOperator( + # SQLExecuteQueryOperator (apache-airflow-providers-common-sql), not SnowflakeOperator, + # which apache-airflow-providers-snowflake 6.0 removed. + return f"""{var} = SQLExecuteQueryOperator( task_id={task_id_lit}, - snowflake_conn_id=SNOWFLAKE_CONN_ID, + conn_id=SNOWFLAKE_CONN_ID, sql={sql_lit}, - warehouse={warehouse_lit}, - database={database_lit}, + hook_params={{"warehouse": {warehouse_lit}, "database": {database_lit}}}, dag=dag, )""" @@ -345,7 +352,7 @@ def _generate_task_dependencies(tasks: List[Dict[str, Any]]) -> str: def _convert_schedule(schedule: str) -> str: - """Convert FLUID schedule to Airflow schedule_interval.""" + """Convert FLUID schedule to an Airflow ``schedule`` value.""" schedule_lower = schedule.lower() if schedule_lower in ("@hourly", "@daily", "@weekly", "@monthly"): return schedule_lower diff --git a/fluid_build/providers/snowflake/orchestration/airflow_generator.py b/fluid_build/providers/snowflake/orchestration/airflow_generator.py index b39504cf..9c280342 100644 --- a/fluid_build/providers/snowflake/orchestration/airflow_generator.py +++ b/fluid_build/providers/snowflake/orchestration/airflow_generator.py @@ -17,8 +17,8 @@ Airflow DAG generation from FLUID 0.7.1 orchestration specifications. Converts FLUID provider action tasks into Airflow operators: -- sf.table.ensure → SnowflakeOperator -- sf.sql.execute → SnowflakeOperator +- sf.table.ensure → SQLExecuteQueryOperator +- sf.sql.execute → SQLExecuteQueryOperator - python → PythonOperator - bash → BashOperator """ @@ -26,7 +26,12 @@ from typing import Any, Dict, List, Optional from ..._sql_safety import quote_string_literal, validate_ident, validate_sql_type_name -from ...common.codegen_utils import py_str_literal, sanitize_identifier +from ...common.codegen_utils import ( + airflow_dag_import, + airflow_standard_import, + py_str_literal, + sanitize_identifier, +) from ..registry import SnowflakeActionRegistry from .common import OrchestrationConfig, OrchestrationEngine, extract_dependencies, sanitize_task_id @@ -128,12 +133,8 @@ def _generate_dag_header(config: OrchestrationConfig) -> str: return f"""{docstring_lit} from datetime import datetime, timedelta -from airflow import DAG -from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator -from airflow.operators.python import PythonOperator -from airflow.operators.bash import BashOperator -from airflow.sensors.external_task import ExternalTaskSensor - +{airflow_dag_import()}from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator +{airflow_standard_import("PythonOperator")}{airflow_standard_import("BashOperator")}{airflow_standard_import("ExternalTaskSensor")} # DAG default arguments default_args = {{ 'owner': {owner_lit}, @@ -150,7 +151,7 @@ def _generate_dag_header(config: OrchestrationConfig) -> str: dag_id={dag_id_lit}, default_args=default_args, description={description_lit}, - schedule_interval={schedule_lit}, + schedule={schedule_lit}, start_date=datetime(2025, 1, 1), catchup=False, tags={config.tags!r}, @@ -190,7 +191,7 @@ def _generate_task_definitions(tasks: List[Dict[str, Any]], config: Orchestratio def _generate_provider_action_task( task: Dict[str, Any], task_id: str, config: OrchestrationConfig ) -> str: - """Generate SnowflakeOperator for provider action.""" + """Generate a SQLExecuteQueryOperator for a provider action.""" action = task.get("action", "") params = task.get("parameters", {}) task.get("description", "") @@ -230,9 +231,11 @@ def _generate_provider_action_task( conn_lit = py_str_literal(config.snowflake_conn_id) sql_lit = py_str_literal(sql) - return f"""{var} = SnowflakeOperator( + # SQLExecuteQueryOperator (common-sql), not SnowflakeOperator: snowflake provider 6.0 + # removed it. + return f"""{var} = SQLExecuteQueryOperator( task_id={task_id_lit}, - snowflake_conn_id={conn_lit}, + conn_id={conn_lit}, sql={sql_lit}, dag=dag, ) @@ -339,7 +342,7 @@ def _generate_create_table_sql(params: Dict[str, Any]) -> str: """Generate CREATE TABLE SQL. SQL-injection safety: this SQL is executed by the generated - ``SnowflakeOperator`` at DAG *runtime*, so every IDENTIFIER (database / + ``SQLExecuteQueryOperator`` at DAG *runtime*, so every IDENTIFIER (database / schema / table / column / cluster-by column) is routed through ``validate_ident`` before the existing double-quote wrapping, and every column TYPE through ``validate_sql_type_name``. A contract-supplied value diff --git a/fluid_build/runtimes/airflow_provider_actions.py b/fluid_build/runtimes/airflow_provider_actions.py index 1303cb65..5ca54512 100644 --- a/fluid_build/runtimes/airflow_provider_actions.py +++ b/fluid_build/runtimes/airflow_provider_actions.py @@ -26,6 +26,8 @@ from typing import Any, Dict, List, Optional from ..providers.common.codegen_utils import ( + airflow_dag_import, + airflow_standard_import, escape_for_docstring, py_str_literal, sanitize_identifier, @@ -279,11 +281,7 @@ def _generate_dag_header(self, dag_id: str, schedule: str, contract: Dict[str, A Domain: {d_domain} Description: {d_description} """ -from airflow import DAG -from airflow.operators.bash import BashOperator -from airflow.operators.python import PythonOperator -from airflow.utils.dates import days_ago -from datetime import datetime, timedelta +{airflow_dag_import()}{airflow_standard_import("BashOperator")}{airflow_standard_import("PythonOperator")}from datetime import datetime, timedelta, timezone # DAG configuration default_args = {{ @@ -299,8 +297,8 @@ def _generate_dag_header(self, dag_id: str, schedule: str, contract: Dict[str, A dag = DAG( dag_id={py_str_literal(dag_id)}, description={py_str_literal(description)}, - schedule_interval={py_str_literal(schedule)}, - start_date=days_ago(1), + schedule={py_str_literal(schedule)}, + start_date=datetime(2026, 1, 1, tzinfo=timezone.utc), catchup=False, tags=["fluid", "data-product", {py_str_literal(kind)}, {py_str_literal(domain)}], default_args=default_args @@ -454,14 +452,13 @@ def _generate_empty_dag( return f'''""" Empty Airflow DAG - No provider actions found """ -from airflow import DAG -from airflow.operators.bash import BashOperator -from airflow.utils.dates import days_ago +from datetime import datetime, timezone +{airflow_dag_import()}{airflow_standard_import("BashOperator")} dag = DAG( dag_id={py_str_literal(dag_id)}, - schedule_interval={py_str_literal(schedule)}, - start_date=days_ago(1), + schedule={py_str_literal(schedule)}, + start_date=datetime(2026, 1, 1, tzinfo=timezone.utc), catchup=False, tags=["fluid", "placeholder"] ) diff --git a/fluid_build/schedulers/airflow/__init__.py b/fluid_build/schedulers/airflow/__init__.py index 3185dac9..2c7760fb 100644 --- a/fluid_build/schedulers/airflow/__init__.py +++ b/fluid_build/schedulers/airflow/__init__.py @@ -25,6 +25,8 @@ from fluid_build.engines.base import Severity, ValidationIssue from fluid_build.providers.common.codegen_utils import ( + airflow_dag_import, + airflow_standard_import, convert_schedule_to_airflow, generate_file_header, generate_task_dependencies_code, @@ -148,11 +150,13 @@ def validate( def _generate_imports(provider: Optional[str]) -> str: """Generate import statements based on provider.""" - base = """from datetime import datetime, timedelta -from zoneinfo import ZoneInfo -from airflow import DAG -from airflow.operators.python import PythonOperator -import json + base = ( + "from datetime import datetime, timedelta\n" + "from zoneinfo import ZoneInfo\n" + + airflow_dag_import() + + airflow_standard_import("PythonOperator") + ) + base += """import json import logging logger = logging.getLogger(__name__)""" @@ -162,7 +166,6 @@ def _generate_imports(provider: Optional[str]) -> str: base += """ from airflow.providers.google.cloud.operators.bigquery import ( BigQueryCreateEmptyDatasetOperator, - BigQueryCreateEmptyTableOperator, BigQueryInsertJobOperator, ) from airflow.providers.google.cloud.operators.gcs import ( @@ -176,7 +179,7 @@ def _generate_imports(provider: Optional[str]) -> str: from airflow.providers.amazon.aws.operators.step_function import StepFunctionStartExecutionOperator""" elif provider_lower == "snowflake": base += """ -from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator""" +from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator""" return base @@ -214,7 +217,7 @@ def _generate_dag_definition( dag_id={py_str_literal(sanitize_identifier(contract_id))}, default_args=default_args, description={py_str_literal(contract_name)}, - schedule_interval={py_str_literal(airflow_schedule)}, + schedule={py_str_literal(airflow_schedule)}, start_date=datetime(2026, 1, 1, tzinfo=ZoneInfo({py_str_literal(timezone)})), catchup=False, tags={tags!r}, @@ -371,19 +374,20 @@ def _generate_snowflake_task(task: Dict[str, Any], config: Dict[str, Any]) -> st # repr() escapes embedded quotes / triple-quotes / newlines, so SQL # cannot terminate the literal and inject a top-level statement. sql_lit = py_str_literal(params.get("sql", params.get("query", "SELECT 1"))) - return f"""{var} = SnowflakeOperator( + return f"""{var} = SQLExecuteQueryOperator( task_id={task_id_lit}, sql={sql_lit}, - snowflake_conn_id={conn_lit}, + conn_id={conn_lit}, dag=dag, )""" - # Fallback: wrap action in SnowflakeOperator SQL call + # Fallback: wrap action in a SQLExecuteQueryOperator SQL call (snowflake provider 6.0 + # removed SnowflakeOperator) sql_lit = py_str_literal(params.get("sql", f"-- TODO: implement {action}")) - return f"""{var} = SnowflakeOperator( + return f"""{var} = SQLExecuteQueryOperator( task_id={task_id_lit}, sql={sql_lit}, - snowflake_conn_id={conn_lit}, + conn_id={conn_lit}, dag=dag, )""" diff --git a/fluid_build/templates/first-dag/README.md b/fluid_build/templates/first-dag/README.md index 52795c02..8feb0246 100644 --- a/fluid_build/templates/first-dag/README.md +++ b/fluid_build/templates/first-dag/README.md @@ -74,10 +74,17 @@ outputs: ```python # dags/sales_analytics_pipeline_dag.py (auto-generated!) -from airflow import DAG -from airflow.operators.python import PythonOperator from datetime import datetime, timedelta +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 + default_args = { 'owner': 'analytics-team', 'depends_on_past': False, @@ -89,7 +96,7 @@ default_args = { with DAG( 'sales_analytics_pipeline', default_args=default_args, - schedule_interval='0 2 * * *', + schedule='0 2 * * *', catchup=False, tags=['sales', 'analytics', 'auto-generated'], ) as dag: diff --git a/tests/test_airflow3_dag_compat.py b/tests/test_airflow3_dag_compat.py new file mode 100644 index 00000000..12ea12eb --- /dev/null +++ b/tests/test_airflow3_dag_compat.py @@ -0,0 +1,473 @@ +# 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. + +"""Every Airflow DAG forge-cli generates parses on Airflow 2.7+ and on Airflow 3. + +Airflow 3 removed ``DAG(schedule_interval=...)`` (``TypeError``) and +``airflow.utils.dates.days_ago`` (``ImportError``), so a DAG that uses either fails to +import. It also moved the core operators and sensors to apache-airflow-providers-standard +and exports ``DAG``, ``dag`` and ``task`` from ``airflow.sdk``; the old paths still +resolve on 3.2 with a deprecation warning. Every generator (``fluid export --engine +airflow``, ``fluid generate schedule``, ``fluid init``, ``fluid scaffold-composer``, +``fluid schedule-sync``, the AWS, GCP and Snowflake code generators) emits the shape +both versions accept without either: ``schedule=``, a fixed ``start_date``, and each +moved name imported from its Airflow 3 home with an ``except ImportError`` fallback to +the Airflow 2 path (``providers/common/codegen_utils.py``). + +This renders every generator and checks the source, so it needs no Airflow +install. The author also loaded each rendered DAG into a real Airflow 3.2.2 and +2.10.5 ``DagBag`` (see the PR). +""" + +from __future__ import annotations + +import ast +import logging +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +# Imported only in an ``except ImportError`` fallback, behind its Airflow 3 path. +AIRFLOW2_ONLY = ("airflow.operators", "airflow.sensors", "airflow.decorators") +AIRFLOW2_DAG_MODULES = {"airflow", "airflow.models", "airflow.models.dag"} +AIRFLOW3_HOMES = ("airflow.sdk", "airflow.providers.standard.") + +_TASKS = [ + { + "taskId": "load", + "type": "provider_action", + "action": "aws.athena.query", + "params": {"query": "SELECT 1"}, + }, + { + "taskId": "publish", + "type": "provider_action", + "action": "aws.athena.query", + "params": {"query": "SELECT 2"}, + "dependsOn": ["load"], + }, +] + + +def _contract(action_prefix: str = "aws.athena") -> dict: + tasks = [{**t, "action": f"{action_prefix}.query"} for t in _TASKS] + return { + "id": "sales.orders", + "name": "Orders", + "orchestration": {"schedule": "0 2 * * *", "timezone": "UTC", "tasks": tasks}, + } + + +def _provider_actions() -> list: + from fluid_build.forge.core.provider_actions import ActionType + + def action(action_id, kind, provider, params, depends_on=None): + return SimpleNamespace( + action_id=action_id, + action_type=ActionType(kind), + provider=provider, + params=params, + depends_on=depends_on or [], + description=None, + ) + + return [ + action( + "provision-dataset", + "provisionDataset", + "gcp", + {"exposeId": "orders", "binding": {"location": {"project": "p", "dataset": "d"}}}, + ), + action( + "schedule-build", + "scheduleTask", + "gcp", + {"engine": "dbt", "script": "orders_model", "buildId": "b1"}, + ["provision-dataset"], + ), + ] + + +def _runtime_dag(actions: list) -> str: + from fluid_build.runtimes.airflow_provider_actions import AirflowDAGGenerator + + parser = MagicMock() + parser.parse.return_value = actions + with patch("fluid_build.forge.core.provider_actions.ProviderActionParser", return_value=parser): + return AirflowDAGGenerator(logger=logging.getLogger(__name__)).generate_dag( + {"id": "sales.orders", "name": "Orders", "domain": "sales"} + ) + + +def _scheduler(provider: str) -> str: + from fluid_build.schedulers.airflow import AirflowScheduler + + prefix = {"aws": "aws.athena", "gcp": "gcp.bigquery", "snowflake": "sf.snowflake"}[provider] + (code,) = AirflowScheduler().generate(_contract(prefix), provider=provider).values() + return code + + +def _aws(taskflow: bool = False, mixed: bool = False) -> str: + from fluid_build.providers.aws.codegen import airflow + + contract = _contract() + if mixed: # an operator task upstream of a @task + contract["orchestration"]["tasks"][0].update( + action="aws.s3.ensure_bucket", params={"bucket": "b"} + ) + render = airflow.generate_airflow_dag_taskflow if taskflow else airflow.generate_airflow_dag + return render(contract, "123456789012", "eu-west-1") + + +def _gcp(dataflow: bool = False) -> str: + from fluid_build.providers.gcp.codegen.airflow import generate_airflow_dag + + contract = _contract("gcp.bigquery") + contract["orchestration"]["tasks"].append( + { + "taskId": "create_orders", + "type": "provider_action", + "action": "gcp.bigquery.create_table", + "params": { + "dataset_id": "sales", + "table_id": "orders", + "schema": [{"name": "id", "type": "INT64"}], + }, + } + ) + if dataflow: + contract["orchestration"]["tasks"].append( + { + "taskId": "flow", + "type": "provider_action", + "action": "gcp.dataflow.run_template", + "params": {"template": "gs://t", "job_name": "j"}, + } + ) + return generate_airflow_dag(contract, "my-project", "europe-west1") + + +def _snowflake() -> str: + """The SQL tasks (``sf.query.*``, ``sf.table.*``): ``sf.snowflake.*`` falls back to a + PythonOperator and never reaches SQLExecuteQueryOperator.""" + from fluid_build.providers.snowflake.codegen.airflow import generate_airflow_dag + + contract = _contract("sf.snowflake") + contract["orchestration"]["tasks"] = [ + { + "taskId": "load", + "type": "provider_action", + "action": "sf.query.run_query", + "params": {"sql": "SELECT 1"}, + }, + { + "taskId": "publish", + "type": "provider_action", + "action": "sf.table.create_table", + "params": {"table": "ORDERS", "columns": [{"name": "ID", "type": "INTEGER"}]}, + "dependsOn": ["load"], + }, + ] + return generate_airflow_dag(contract, account="acct", database="DB") + + +def _snowflake_orchestration() -> str: + """Every operator it imports, in its own task shape (tasks keyed by ``name``).""" + from fluid_build.providers.snowflake.orchestration.airflow_generator import generate_airflow_dag + + orchestration = { + "schedule": "0 2 * * *", + "tasks": [ + { + "name": "load", + "type": "provider_action", + "action": "sf.sql.execute", + "params": {"sql": "SELECT 1"}, + }, + {"name": "notify", "type": "bash", "command": "echo done"}, + { + "name": "upstream", + "type": "sensor", + "external_dag_id": "raw", + "external_task_id": "land", + }, + ], + } + return generate_airflow_dag({"id": "sales.orders", "name": "Orders"}, orchestration) + + +def _init_dag(tmp_path: Path) -> str: + from fluid_build.cli.init import create_basic_dag # re-exported; the helper module imports init + + create_basic_dag( + tmp_path, {"name": "orders", "orchestration": {"schedule": "@daily"}}, MagicMock() + ) + (dag,) = tmp_path.rglob("*.py") + return dag.read_text() + + +def _scaffold_composer() -> str: + from fluid_build.cli.scaffold_composer import DAG_TMPL + + return DAG_TMPL.format( + dag_id="'orders'", + cron="'0 2 * * *'", + validate_cmd="'true'", + plan_cmd="'true'", + apply_cmd="'true'", + ) + + +def _schedule_sync() -> str: + from fluid_build.cli._acquisition_stage_ext import _render_airflow_dag + + return _render_airflow_dag("sales.orders", "b1", {}, "0 2 * * *") + + +RENDERS = { + "runtime with actions": lambda tmp: _runtime_dag(_provider_actions()), + "runtime without actions": lambda tmp: _runtime_dag([]), + "scheduler aws": lambda tmp: _scheduler("aws"), + "scheduler gcp": lambda tmp: _scheduler("gcp"), + "scheduler snowflake": lambda tmp: _scheduler("snowflake"), + "aws codegen": lambda tmp: _aws(), + "aws codegen taskflow": lambda tmp: _aws(taskflow=True), + "aws codegen taskflow with an operator": lambda tmp: _aws(taskflow=True, mixed=True), + "gcp codegen": lambda tmp: _gcp(), + "gcp codegen with dataflow": lambda tmp: _gcp(dataflow=True), + "snowflake codegen": lambda tmp: _snowflake(), + "snowflake orchestration": lambda tmp: _snowflake_orchestration(), + "fluid init": _init_dag, + "scaffold-composer": lambda tmp: _scaffold_composer(), + "schedule-sync": lambda tmp: _schedule_sync(), +} + + +def _under(module: str, prefixes) -> bool: + return any(module == p or module.startswith(p + ".") for p in prefixes) + + +def _imported_names(node: ast.AST) -> set[str]: + return ( + {a.asname or a.name for a in node.names} + if isinstance(node, (ast.Import, ast.ImportFrom)) + else set() + ) + + +def _airflow2_only(node: ast.AST) -> bool: + if isinstance(node, ast.Import): + return any(_under(a.name, AIRFLOW2_ONLY) for a in node.names) + if isinstance(node, ast.ImportFrom): + module = node.module or "" + names = {a.name for a in node.names} + return _under(module, AIRFLOW2_ONLY) or ( + module in AIRFLOW2_DAG_MODULES + and bool(names & {"DAG", "dag", "task", "operators", "sensors"}) + ) + return False + + +def _catches_import_error(handler: ast.ExceptHandler) -> bool: + kinds = handler.type.elts if isinstance(handler.type, ast.Tuple) else [handler.type] + return any( + isinstance(k, ast.Name) and k.id in ("ImportError", "ModuleNotFoundError") for k in kinds + ) + + +def _airflow2_only_imports_outside_a_fallback(tree: ast.AST) -> list[str]: + """Each Airflow-2-only import (``from x import y`` or ``import x``), unless it sits in an + ``except ImportError`` whose ``try`` imports the same names from their Airflow 3 home.""" + fallbacks: set[int] = set() + for node in ast.walk(tree): + if not isinstance(node, ast.Try): + continue + airflow3 = [ + n + for stmt in node.body + for n in ast.walk(stmt) + if isinstance(n, ast.ImportFrom) and (n.module or "").startswith(AIRFLOW3_HOMES) + ] + provided = set().union(*(_imported_names(n) for n in airflow3)) if airflow3 else set() + for handler in node.handlers: + if not (airflow3 and _catches_import_error(handler)): + continue + for stmt in handler.body: + for n in ast.walk(stmt): + if _imported_names(n) and _imported_names(n) <= provided: + fallbacks.add(id(n)) + return [ + f"line {node.lineno}: {ast.unparse(node)}" + for node in ast.walk(tree) + if _airflow2_only(node) and id(node) not in fallbacks + ] + + +@pytest.mark.parametrize("name", sorted(RENDERS)) +def test_the_generated_dag_parses_on_airflow_2_and_3(name: str, tmp_path: Path): + source = RENDERS[name](tmp_path) + tree = ast.parse(source) # it is Python at all + assert ( + "schedule_interval" not in source + ), "Airflow 3 removed DAG(schedule_interval=); use schedule=" + assert "days_ago" not in source, "Airflow 3 removed airflow.utils.dates.days_ago" + assert not _airflow2_only_imports_outside_a_fallback( + tree + ), _airflow2_only_imports_outside_a_fallback(tree) + assert any( + isinstance(n, ast.keyword) and n.arg == "schedule" for n in ast.walk(tree) + ), "no schedule=" + + +@pytest.mark.parametrize( + "source", + [ + "from airflow import DAG\n", + "from airflow.operators.bash import BashOperator\n", + "import airflow.operators.bash\n", + "from airflow.operators import bash\n", + "from airflow.models import DAG\n", + "from airflow.decorators import dag, task\n", + # A fallback with no Airflow 3 import in its ``try`` is no fallback ... + "try:\n import json\nexcept ImportError:\n from airflow.operators.bash import BashOperator\n", + # ... nor one that imports a different name than its ``try``. + "try:\n from airflow.providers.standard.operators.bash import BashOperator\n" + "except ImportError:\n from airflow.operators.python import PythonOperator\n", + ], +) +def test_the_import_check_can_fail(source: str): + assert len(_airflow2_only_imports_outside_a_fallback(ast.parse(source))) == 1 + + +def test_the_import_check_accepts_the_fallbacks_generators_emit(): + tuple_handler = ( + "try:\n from airflow.sdk import DAG\n" + "except (ImportError, ModuleNotFoundError):\n from airflow import DAG\n" + ) + assert _airflow2_only_imports_outside_a_fallback(ast.parse(tuple_handler)) == [] + from fluid_build.providers.common.codegen_utils import ( + airflow_dag_import, + airflow_standard_import, + ) + + shimmed = airflow_dag_import() + airflow_standard_import("BashOperator") + assert _airflow2_only_imports_outside_a_fallback(ast.parse(shimmed)) == [] + + +def test_the_dataflow_import_comes_only_with_a_dataflow_task(tmp_path: Path): + """Its module needs the apache-beam provider; an unused failing import stops the DAG.""" + assert "DataflowTemplatedJobStartOperator" not in RENDERS["gcp codegen"](tmp_path) + assert "DataflowTemplatedJobStartOperator" in RENDERS["gcp codegen with dataflow"](tmp_path) + + +def test_no_generator_emits_an_operator_its_provider_removed(tmp_path: Path): + """snowflake provider 6.0 removed SnowflakeOperator; google provider 17.0 removed + BigQueryCreateEmptyTableOperator, which survives only behind the newer operator.""" + for name, render in RENDERS.items(): + workdir = tmp_path / name.replace(" ", "_") + workdir.mkdir() + source = render(workdir) + assert "SnowflakeOperator" not in source, name + if "BigQueryCreateEmptyTableOperator" in source: + assert "BigQueryCreateTableOperator = None" in source, name + + +def test_taskflow_wires_task_results_not_task_functions(tmp_path: Path): + source = RENDERS["aws codegen taskflow"](tmp_path) + assert "load_result >> publish_result" in source + assert " load >> publish" not in source + # An operator task is wired as itself; only a @task is wired by its result. + assert " load >> publish_result" in RENDERS["aws codegen taskflow with an operator"]( + tmp_path + ) + + +def test_snowflake_sql_runs_through_sql_execute_query_with_its_hook_params(tmp_path: Path): + source = RENDERS["snowflake codegen"](tmp_path) + assert source.count("SQLExecuteQueryOperator(") == 2 + assert "conn_id=SNOWFLAKE_CONN_ID" in source and "hook_params={" in source + + +def test_a_dataflow_import_follows_the_dispatch_not_the_prefix(): + from fluid_build.providers.gcp.codegen.airflow import _generate_dag_imports + + def imports(action: str) -> str: + return _generate_dag_imports([{"action": action}]) + + assert "DataflowTemplatedJobStartOperator" in imports("dataflow.dataflow.run_template") + assert "DataflowTemplatedJobStartOperator" not in imports("gcp.dataflow.launch_flex") + # Before the standard library imports, not after `logger = ...` (E402 in the DAG). + flow = imports("gcp.dataflow.run_template") + assert flow.index("DataflowTemplatedJobStartOperator") < flow.index("import json") + + +def test_a_task_named_like_another_tasks_result_keeps_its_own_edge(): + """``chk_result`` (an operator) and ``chk`` (a @task) must not share a variable, or the + edge from ``chk_result`` silently starts at ``chk``'s result instead.""" + from fluid_build.providers.aws.codegen.airflow import generate_airflow_dag_taskflow + + contract = _contract() + contract["orchestration"]["tasks"] = [ + { + "taskId": "chk_result", + "type": "provider_action", + "action": "aws.s3.ensure_bucket", + "params": {"bucket": "b"}, + }, + { + "taskId": "chk", + "type": "provider_action", + "action": "aws.glue.ensure_database", + "params": {"database": "d"}, + }, + { + "taskId": "after", + "type": "provider_action", + "action": "aws.glue.ensure_database", + "params": {"database": "e"}, + "dependsOn": ["chk_result"], + }, + ] + source = generate_airflow_dag_taskflow(contract, "123456789012", "eu-west-1") + assert " chk_result_2 = chk()" in source + assert " chk_result >> after_result" in source + + +@pytest.mark.parametrize("task_id", ["task", "dag", "AthenaOperator", "S3CreateBucketOperator"]) +def test_a_task_named_like_an_imported_name_does_not_shadow_it(task_id: str): + """``taskId: task`` emitted ``def task():`` inside the DAG function, shadowing the + ``@task`` decorator (UnboundLocalError when Airflow parsed it).""" + from fluid_build.providers.aws.codegen.airflow import generate_airflow_dag_taskflow + + contract = _contract() + contract["orchestration"]["tasks"] = [ + { + "taskId": task_id, + "type": "provider_action", + "action": "aws.glue.ensure_database", + "params": {"database": "d"}, + } + ] + tree = ast.parse(generate_airflow_dag_taskflow(contract, "123456789012", "eu-west-1")) + bound = {n.name for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)} + bound |= { + t.id + for n in ast.walk(tree) + if isinstance(n, ast.Assign) + for t in n.targets + if isinstance(t, ast.Name) + } + assert task_id not in bound, f"{task_id!r} rebinds an imported name" diff --git a/tests/test_airflow_codegen_injection.py b/tests/test_airflow_codegen_injection.py index b9565892..d40c4753 100644 --- a/tests/test_airflow_codegen_injection.py +++ b/tests/test_airflow_codegen_injection.py @@ -286,7 +286,9 @@ def test_benign_contract_still_produces_expected_operators(): } (code,) = AirflowScheduler().generate(contract, provider="snowflake").values() assert_inert(code) - assert "load = SnowflakeOperator(" in code + assert ( + "load = SQLExecuteQueryOperator(" in code + ) # snowflake provider 6.0 removed SnowflakeOperator assert "sql='SELECT 1'" in code assert "task_id='load'" in code @@ -982,3 +984,71 @@ def test_aws_task_id_cannot_rebind_emitted_module_names(): m = _unique_task_identifiers([{"taskId": "dag"}, {"taskId": "_execute_provider_action"}]) assert m["dag"] != "dag" assert m["_execute_provider_action"] != "_execute_provider_action" + + +# --------------------------------------------------------------------------- +# GCP create_table schema, and `fluid init`'s basic DAG +# --------------------------------------------------------------------------- + + +def test_gcp_create_table_schema_is_injection_safe(): + """The schema reaches both create-table operators as a literal: as ``schema_fields`` + and inside ``table_resource``.""" + from fluid_build.providers.gcp.codegen.airflow import generate_airflow_dag + + hostile = [{"name": POC_SQL, "type": POC_SQUOTE, "description": POC_IDENT, POC_SQL: 1}] + contract = { + "id": "p", + "name": "p", + "orchestration": { + "tasks": [ + { + "taskId": POC_IDENT, + "type": "provider_action", + "action": "gcp.bigquery.create_table", + "params": {"dataset_id": POC_SQL, "table_id": POC_SQUOTE, "schema": hostile}, + } + ] + }, + } + code = generate_airflow_dag(contract, "my-project", "us-central1") + assert_inert(code) + assert "BigQueryCreateTableOperator(" in code and "BigQueryCreateEmptyTableOperator(" in code + + +def test_init_basic_dag_retries_are_numbers_not_code(tmp_path): + """``orchestration.retries`` and ``retry_delay`` used to be spliced into the DAG source + as text, so a contract value ran at DAG-parse time.""" + import logging + + from fluid_build.cli.init import create_basic_dag + + contract = { + "name": "p", + "orchestration": { + "schedule": "@daily", + "retries": "3, __import__('os').system('touch /tmp/PWNED4')", + "retry_delay": "5) or __import__('os').system('touch /tmp/PWNED5') or (5", + }, + } + create_basic_dag(tmp_path, contract, logging.getLogger(__name__)) + (dag,) = (tmp_path / "dags").glob("*.py") + code = dag.read_text() + assert_inert(code) + assert "'retries': 3," in code and "timedelta(minutes=5)" in code + assert "PWNED" not in code + + +def test_init_basic_dag_infinite_retries_fall_back(tmp_path): + """``int(float('inf'))`` raises OverflowError, which skipped the DAG instead.""" + import logging + + from fluid_build.cli.init import create_basic_dag + + create_basic_dag( + tmp_path, + {"name": "p", "orchestration": {"retries": float("inf")}}, + logging.getLogger(__name__), + ) + (dag,) = (tmp_path / "dags").glob("*.py") + assert "'retries': 3," in dag.read_text() diff --git a/tests/test_airflow_extended.py b/tests/test_airflow_extended.py index 425aa17f..58d02fb2 100644 --- a/tests/test_airflow_extended.py +++ b/tests/test_airflow_extended.py @@ -191,7 +191,9 @@ def test_dependencies_present(self): ] ) result = generate_airflow_dag_taskflow(contract, "123", "us-east-1") - assert "a >> b" in result + # The S3 operator to the Glue @task's result: the @task function itself is not + # a task, and `a >> b` raised TypeError when Airflow parsed the DAG. + assert " a >> b_result" in result def test_athena_uses_classic_operator(self): contract = _make_contract( diff --git a/tests/test_scheduler_airflow.py b/tests/test_scheduler_airflow.py index 0ba87909..f18ec701 100644 --- a/tests/test_scheduler_airflow.py +++ b/tests/test_scheduler_airflow.py @@ -191,7 +191,9 @@ def test_snowflake_generates_dag(self, scheduler, snowflake_contract): content = list(files.values())[0] assert "from airflow import DAG" in content - assert "SnowflakeOperator" in content + assert ( + "SQLExecuteQueryOperator" in content + ) # snowflake provider 6.0 removed SnowflakeOperator assert "load_customers" in content assert "sf_conn" in content