Skip to content

Commit

Permalink
Better message for exception for templated base operator fields (apac…
Browse files Browse the repository at this point in the history
…he#37668)

* Better message for exception for templated base operator fields

* fix test
  • Loading branch information
RNHTTR authored Feb 24, 2024
1 parent 2ac8abf commit 27da663
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 2 deletions.
8 changes: 7 additions & 1 deletion airflow/serialization/serialized_objects.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import weakref
from dataclasses import dataclass
from inspect import signature
from textwrap import dedent
from typing import TYPE_CHECKING, Any, Collection, Iterable, Mapping, NamedTuple, Union

import attrs
Expand Down Expand Up @@ -912,7 +913,12 @@ def _serialize_node(cls, op: BaseOperator | MappedOperator, include_deps: bool)
if op.template_fields:
for template_field in op.template_fields:
if template_field in forbidden_fields:
raise AirflowException(f"Cannot template BaseOperator field: {template_field!r}")
raise AirflowException(
dedent(
f"""Cannot template BaseOperator field:
{template_field!r} {op.__class__.__name__=} {op.template_fields=}"""
)
)
value = getattr(op, template_field, None)
if not cls._is_excluded(value, template_field, op):
serialize_op[template_field] = serialize_template_field(value)
Expand Down
12 changes: 11 additions & 1 deletion tests/serialization/test_dag_serialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,11 @@
import multiprocessing
import os
import pickle
import re
from datetime import datetime, timedelta
from glob import glob
from pathlib import Path
from textwrap import dedent
from typing import TYPE_CHECKING
from unittest import mock

Expand Down Expand Up @@ -2072,7 +2074,15 @@ def execute(self, context: Context):
task.render_template_fields(context={"test_email_list": ["[email protected]", "[email protected]"]})
assert task.email == "[email protected],[email protected]"

with pytest.raises(AirflowException, match="Cannot template BaseOperator field: 'execution_timeout'"):
with pytest.raises(
AirflowException,
match=re.escape(
dedent(
"""Failed to serialize DAG 'test_dag': Cannot template BaseOperator field:
'execution_timeout' op.__class__.__name__='TestOperator' op.template_fields=('email', 'execution_timeout')"""
)
),
):
SerializedDAG.to_dict(dag)


Expand Down

0 comments on commit 27da663

Please sign in to comment.