diff --git a/dagfactory/dagbuilder.py b/dagfactory/dagbuilder.py index b9c4058e..1980244d 100644 --- a/dagfactory/dagbuilder.py +++ b/dagfactory/dagbuilder.py @@ -280,6 +280,26 @@ def get_dag_params(self) -> Dict[str, Any]: return dag_params + @staticmethod + def _resolve_operator_alias(operator: str, operator_aliases: List[Dict[str, str]]) -> str: + """ + Resolves a short operator alias to its full dotted module path. + + If ``operator`` matches an ``operator_alias`` entry in the provided list, + the corresponding ``operator_path`` is returned. Otherwise the original + string is returned unchanged, allowing full dotted paths to pass through. + + :param operator: the operator string from the task config (alias or full path) + :param operator_aliases: list of dicts with ``operator_alias`` and ``operator_path`` keys + :returns: resolved full dotted operator path + """ + aliases_map: Dict[str, str] = { + entry["operator_alias"]: entry["operator_path"] + for entry in operator_aliases + if "operator_alias" in entry and "operator_path" in entry + } + return aliases_map.get(operator, operator) + @staticmethod def _resolve_user_defined_macros(macros: Dict[str, Any], path: str = "user_defined_macros") -> Dict[str, Any]: """ @@ -974,6 +994,9 @@ def build(self) -> Dict[str, Union[str, DAG]]: if "operator" in task_conf: operator: str = task_conf["operator"] + operator = DagBuilder._resolve_operator_alias( + operator, self.default_config.get("operator_aliases", []) + ) if task_conf.get("expand"): task_conf = self.replace_expand_values(task_conf, tasks_dict) diff --git a/dagfactory/dagfactory.py b/dagfactory/dagfactory.py index 5918c36b..05264772 100644 --- a/dagfactory/dagfactory.py +++ b/dagfactory/dagfactory.py @@ -233,6 +233,9 @@ def build_dags(self) -> Tuple[Dict[str, DAG], Optional[Tuple[str, Exception]]]: default_config["default_args"] = self._merge_default_args_from_list_configs( [global_default_args, default_config] ) + # Propagate operator_aliases from defaults.yml unless the dag YAML already defines them + if "operator_aliases" not in default_config and "operator_aliases" in global_default_args: + default_config["operator_aliases"] = global_default_args["operator_aliases"] dags: Dict[str, Any] = {} diff --git a/dev/dags/defaults.yml b/dev/dags/defaults.yml index e406ea95..1f8e4b59 100644 --- a/dev/dags/defaults.yml +++ b/dev/dags/defaults.yml @@ -1,3 +1,9 @@ default_args: start_date: "2025-01-01" owner: "global_owner" + +operator_aliases: + - operator_alias: BashOperator + operator_path: airflow.providers.standard.operators.bash.BashOperator + - operator_alias: PythonOperator + operator_path: airflow.providers.standard.operators.python.PythonOperator diff --git a/tests/test_dagbuilder.py b/tests/test_dagbuilder.py index 5bb9b033..fd139f4c 100644 --- a/tests/test_dagbuilder.py +++ b/tests/test_dagbuilder.py @@ -400,6 +400,47 @@ def test_make_task_bad_operator(): td.make_task(operator, task_params) +def test_resolve_operator_alias_returns_full_path(): + aliases = [{"operator_alias": "BashOperator", "operator_path": get_bash_operator_path()}] + resolved = DagBuilder._resolve_operator_alias("BashOperator", aliases) + assert resolved == get_bash_operator_path() + + +def test_resolve_operator_alias_passthrough_full_path(): + full_path = get_bash_operator_path() + resolved = DagBuilder._resolve_operator_alias(full_path, []) + assert resolved == full_path + + +def test_resolve_operator_alias_unknown_alias_passthrough(): + aliases = [{"operator_alias": "BashOperator", "operator_path": get_bash_operator_path()}] + resolved = DagBuilder._resolve_operator_alias("UnknownOperator", aliases) + assert resolved == "UnknownOperator" + + +def test_resolve_operator_alias_empty_aliases(): + resolved = DagBuilder._resolve_operator_alias("BashOperator", []) + assert resolved == "BashOperator" + + +def test_build_dag_with_operator_alias(): + aliases = [{"operator_alias": "BashOperator", "operator_path": get_bash_operator_path()}] + default_config_with_aliases = {**DEFAULT_CONFIG, "operator_aliases": aliases} + dag_config_with_alias = { + **DAG_CONFIG, + "tasks": [ + { + "task_id": "task_1", + "operator": "BashOperator", + "bash_command": "echo 1", + } + ], + } + td = dagbuilder.DagBuilder("test_dag", dag_config_with_alias, default_config_with_aliases) + result = td.build() + assert "task_1" in result["dag"].task_ids + + def test_make_task_missing_required_param(): td = dagbuilder.DagBuilder("test_dag", DAG_CONFIG, DEFAULT_CONFIG) operator = get_bash_operator_path() diff --git a/tests/test_dagfactory.py b/tests/test_dagfactory.py index 6107f9d7..7383415b 100644 --- a/tests/test_dagfactory.py +++ b/tests/test_dagfactory.py @@ -633,6 +633,40 @@ def test_build_dag_with_global_dag_level_defaults(): assert "global_tag" in dags["test_dag_override"].tags +def test_build_dags_with_operator_aliases_from_global_defaults(): + """Test that operator_aliases in global defaults.yml propagate to DagBuilder and resolve correctly.""" + global_defaults = { + "default_args": { + "owner": "test_owner", + "start_date": "2025-01-01", + }, + "operator_aliases": [ + {"operator_alias": "BashOperator", "operator_path": get_bash_operator_path()}, + ], + } + + config = { + "test_dag": { + "tasks": [ + { + "task_id": "task_1", + "operator": "BashOperator", + "bash_command": "echo 1", + } + ] + } + } + + td = _DagFactory(config_dict=config) + with pytest.MonkeyPatch.context() as m: + m.setattr(td, "_global_default_args", lambda: global_defaults) + dags, build_error = td.build_dags() + + assert build_error is None + assert "task_1" in dags["test_dag"].task_ids + assert type(dags["test_dag"].get_task("task_1")).__name__ == "BashOperator" + + def test_build_dag_with_global_default_dict(): dags, _ = _DagFactory( config_dict=DAG_FACTORY_CONFIG,