import os
import sys
from copy import deepcopy
from inspect import currentframe, signature
from logging import getLogger
from pathlib import Path
from shutil import which
from subprocess import run
from typing import TYPE_CHECKING, Literal
from airflow_pydantic import BaseModel, Dag, DagArgs, ImportPath, TaskArgs
if TYPE_CHECKING:
from airflow.providers.standard.operators.empty import EmptyOperator
else:
from airflow_pydantic.airflow import EmptyOperator
from hydra import compose, initialize_config_dir
from hydra.utils import instantiate
from pydantic import AliasChoices, BaseModel as PydanticBaseModel, Field, SerializeAsAny, model_validator
from airflow_config.exceptions import ConfigFileExtensionError, ConfigNotFoundError
from airflow_config.utils import _get_calling_dag
__all__ = (
"Configuration",
"load_config",
)
_log = getLogger(__name__)
class _Templates(BaseModel):
dag: dict[str, DagArgs] | None = Field(default=None, description="Template for DAGs")
task: dict[str, TaskArgs] | None = Field(default=None, description="Template for Tasks")
class _OperatorDefaults(BaseModel, extra="allow"):
operator: ImportPath
[docs]
class Configuration(BaseModel):
default_task_args: TaskArgs = Field(
default_factory=TaskArgs,
description="Global default default_args (task arguments)",
validation_alias=AliasChoices("default_args", "default_task_args"),
)
default_dag_args: DagArgs = Field(default_factory=DagArgs, description="Global default dag arguments")
default_operator_args: list[_OperatorDefaults] = Field(
default_factory=list, description="Default task arguments selected by task model or operator class"
)
dags: dict[str, SerializeAsAny[Dag]] = Field(default_factory=dict, description="List of dags statically configured via Pydantic")
# Templates:
templates: _Templates | None = Field(default_factory=_Templates, description="Templates for DAGs and Tasks")
# Extensions
extensions: dict[str, PydanticBaseModel] = Field(default_factory=dict, description="Any user-defined extensions")
# Generic options
env: str | None = Field(default="", description="Environment to use for this configuration")
name: str | None = Field(default="", description="Name of the configuration")
root: Path | None = Field(default=None, description="Root path")
tags: dict[str, str] | None = Field(default_factory=dict, description="Generic Tags for config. NOTE: Not related to dag tags")
@property
def default_args(self):
return self.default_task_args
@model_validator(mode="after")
def _validate_after(self):
# TODO add more validation here
if not self.default_dag_args.start_date:
_log.warning("No start_date set in default_dag_args, please set a start_date in the default_dag_args or in the dag itself.")
# Pre-apply default_args to all dags if not set
for dag in self.dags.values():
# See if we can grab the whole thing
if "default_args" not in dag.model_fields_set:
dag.default_args = self.default_args.model_copy()
else:
# Otherwise set key-by-key
for key, value in self.default_args.model_dump(exclude_unset=True).items():
assert dag.default_args is not None
if key not in dag.default_args.model_fields_set:
setattr(dag.default_args, key, deepcopy(value))
# Pre-apply config changes to all dags
for key, value in self.default_dag_args.model_dump(exclude_unset=True).items():
for dag in self.dags.values():
# If `key` is not set in `dag`, set it to the value from `default_dag_args`
if key not in dag.model_fields_set:
setattr(dag, key, deepcopy(value))
for dag in self.dags.values():
for task in (dag.tasks or {}).values():
defaults = {}
for rule in self.default_operator_args:
if isinstance(task, rule.operator) or issubclass(task.operator, rule.operator):
defaults.update(rule.model_extra or {})
for key, value in defaults.items():
if key not in task.model_fields_set:
setattr(task, key, deepcopy(value))
return self
@staticmethod
def _find_parent_config_folder(config_dir: str = "config", config_name: str = "", *, basepath: str = "", _offset: int = 2):
if config_name.endswith(".yml"):
raise ConfigFileExtensionError("Config file must be .yaml, not .yml")
if config_name and not config_name.endswith(".yaml"):
config_name = f"{config_name}.yaml"
if basepath:
if basepath.endswith((".py", ".cfg", ".yaml")):
calling_dag = Path(basepath)
else:
calling_dag = Path(basepath) / "dummy.py"
else:
calling_dag = Path(_get_calling_dag(offset=_offset))
folder = calling_dag.parent.resolve()
exists = (folder / config_dir).exists() if not config_name else (folder / config_dir / f"{config_name}").exists()
while not exists:
folder = folder.parent
if str(folder) == os.path.abspath(os.sep):
raise ConfigNotFoundError(config_dir=config_dir, dagfile=calling_dag)
exists = (folder / config_dir).exists() if not config_name else (folder / config_dir / f"{config_name}").exists()
if not exists and (folder / config_dir / f"{config_name}.yml").exists():
raise ConfigFileExtensionError(f"Config file {config_name}.yml exists in {config_dir} but must be .yaml!")
resolved_config_dir = (folder / config_dir).resolve()
if not config_name:
return folder.resolve(), resolved_config_dir, ""
elif (folder / config_dir / f"{config_name}.yml").exists():
raise ConfigFileExtensionError(f"Config file {config_name}.yml exists in {config_dir} but must be .yaml!")
return folder.resolve(), resolved_config_dir, (folder / config_dir / f"{config_name}").resolve()
[docs]
@staticmethod
def load(
config_dir: str = "config",
config_name: str = "",
overrides: list[str] | None = None,
*,
basepath: str = "",
_offset: int = 3,
) -> "Configuration":
overrides = overrides or []
with initialize_config_dir(config_dir=str(Path(__file__).resolve().parent / "hydra"), version_base=None):
if config_dir:
hydra_folder, resolved_config_dir, _ = Configuration._find_parent_config_folder(
config_dir=config_dir, config_name=config_name, basepath=basepath, _offset=_offset
)
cfg = compose(config_name="conf", overrides=[], return_hydra_config=True)
searchpaths = cfg["hydra"]["searchpath"]
searchpaths.extend([hydra_folder, resolved_config_dir])
if config_name:
overrides = [f"+config={config_name}", *overrides.copy(), f"hydra.searchpath=[{','.join(searchpaths)}]"]
else:
overrides = [*overrides.copy(), f"hydra.searchpath=[{','.join(searchpaths)}]"]
cfg = compose(config_name="conf", overrides=overrides)
config = instantiate(cfg)
if not isinstance(config, Configuration):
config = Configuration(**config)
# Populate dag ids and task ids as a convenience
for dag_id, dag in config.dags.items():
if not dag.dag_id:
dag.dag_id = dag_id
for task_id, task in (dag.tasks or {}).items():
if not task.task_id:
task.task_id = task_id
return config
[docs]
def pre_apply(self, dag, dag_kwargs):
# update options in config based on hard-coded overrides in the DAG file itself
per_dag_default_args = None
# look up per-dag options
if dag_kwargs.get("dag_id", None) in self.dags:
# first try to see if per-dag options have default_args for subtasks
per_dag_kwargs = self.dags[dag_kwargs["dag_id"]]
per_dag_default_args = per_dag_kwargs.default_args.model_dump(exclude_unset=True) if per_dag_kwargs.default_args else None
# if dag is disabled directly, quit right away
if per_dag_kwargs.enabled is False or (per_dag_kwargs.enabled is None and self.default_dag_args.enabled is False):
sys.exit(0)
for attr, val in per_dag_kwargs.model_dump(exclude_unset=True, exclude={"default_args", "enabled", "tasks"}).items():
if attr == "dag_id":
# set
per_dag_kwargs.dag_id = dag_kwargs["dag_id"]
continue
if attr not in dag_kwargs:
dag_kwargs[attr] = val
elif self.default_dag_args.enabled is False:
# if dag has no per-dag-config, but default dag args is disabled, quit right away
sys.exit(0)
if per_dag_default_args is None:
# if no per-dag-default_args, use global default_args
per_dag_default_args = self.default_args.model_dump(exclude_unset=True)
# Override with args hardcoded in the DAG file
if "default_args" in dag_kwargs:
per_dag_default_args.update(dag_kwargs.get("default_args", {}))
# update the dag_kwargs with the final default_args
dag_kwargs["default_args"] = per_dag_default_args
for attr, val in self.default_dag_args.model_dump(exclude_unset=True, exclude={"enabled"}).items():
if attr not in dag_kwargs:
dag_kwargs[attr] = val
doc_md = dag_kwargs.get("doc_md", "")
if dag_kwargs.get("dag_id", None) in self.dags and self.dags[dag_kwargs["dag_id"]].tasks:
doc_md += f"\n# Code Equivalent\n```python\n{self.dags[dag_kwargs['dag_id']].render()}\n```"
if doc_md:
dag_kwargs["doc_md"] = doc_md
[docs]
def apply(self, dag, dag_kwargs):
# update the options in the dag if necessary, instantiate tasks
if dag.dag_id in self.dags:
# Instantiate self against the dag
self.dags[dag.dag_id].instantiate(dag=dag)
def _managed_dags(self):
dags = []
for extension in (self.extensions or {}).values():
factory = getattr(extension, "managed_dags", None)
if factory:
dags.extend(factory())
return dags
def _generated_files(self, airflow_major_version: Literal[2, 3]):
generated = {}
for extension in (self.extensions or {}).values():
factory = getattr(extension, "generated_files", None)
if not factory:
continue
kwargs = {"airflow_major_version": airflow_major_version} if "airflow_major_version" in signature(factory).parameters else {}
for filename, source in factory(**kwargs).items():
if filename in generated and generated[filename] != source:
raise ValueError(f"Extensions generated conflicting files named {filename}")
generated[filename] = source
return generated
[docs]
def generate_in_mem(self, dir: Path | str | None = None, placeholder_dag_id: str = "airflow-config-generate-dags"):
from ..dag import DAG
frame = currentframe()
assert frame is not None
cur_frame = frame.f_back
assert cur_frame is not None
_log.info("Generating Placeholder DAG for in-memory generation")
placeholder_dag = DAG(dag_id=placeholder_dag_id, config=self, schedule=None)
with placeholder_dag:
EmptyOperator(task_id=f"{placeholder_dag_id}-placeholder-task")
cur_frame.f_globals[placeholder_dag_id] = placeholder_dag
dir = dir or Path.cwd()
dir_path = Path(dir)
for dag_id, dag in self.dags.items():
if dag.tasks:
_log.info(f"Generating DAG: {dag_id} in memory")
# rendered = dag.render()
dag_path = dir_path / f"{dag_id}.py"
dag_instance = DAG(dag_id=dag_id, config=self)
# dag_instance.doc_md = f"# Code\n```python\n{rendered}\n```"
_log.info(f"Updating DAG fileloc from {dag_instance.fileloc} to {dag_path!s}")
dag_instance.fileloc = str(dag_path)
cur_frame.f_globals[dag_id] = dag_instance
_log.info(f"Adding DAG {dag_id} complete")
for dag in self._managed_dags():
if dag.dag_id in cur_frame.f_globals:
raise ValueError(f"Managed DAG ID conflicts with an existing DAG: {dag.dag_id}")
dag_path = dir_path / f"{dag.dag_id}.py"
dag_instance = dag.instantiate()
dag_instance.fileloc = str(dag_path)
cur_frame.f_globals[dag.dag_id] = dag_instance
[docs]
def generate(self, dir: Path | str | None = None, *, airflow_major_version: Literal[2, 3]):
dir = dir or Path.cwd()
dir_path = Path(dir)
dir_path.mkdir(parents=True, exist_ok=True)
generated_files = self._generated_files(airflow_major_version)
configured_dag_files = {f"{dag_id}.py" for dag_id in (self.dags or {})}
conflicts = configured_dag_files & generated_files.keys()
if conflicts:
filenames = ", ".join(sorted(conflicts))
raise ValueError(f"Managed DAG files conflict with configured DAGs: {filenames}")
for dag_id, dag in self.dags.items():
if dag.tasks:
rendered = dag.render(airflow_major_version=airflow_major_version)
dag_path = dir_path / f"{dag_id}.py"
if dag_path.exists() and dag_path.read_text() == rendered:
continue
dag_path.write_text(rendered)
ruff = which("ruff")
for filename, source in generated_files.items():
generated_path = dir_path / filename
if generated_path.suffix == ".py" and ruff:
for arguments in (("check", "--fix", "--select", "I"), ("format",)):
source = run(
[ruff, *arguments, "--stdin-filename", str(generated_path), "-"],
input=source,
capture_output=True,
text=True,
check=True,
).stdout
if generated_path.exists() and generated_path.read_text() == source:
continue
generated_path.write_text(source)
load_config = Configuration.load