Skip to content

[v2-10-test] Fix task id validation in BaseOperator #44938

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
potiuk merged 3 commits into from
Dec 15, 2024
Merged

[v2-10-test] Fix task id validation in BaseOperator #44938

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
abb7922
Fix task_id validation in baseoperator
gopidesupavan Dec 14, 2024
1a5946d
Fix task_id validation in baseoperator
gopidesupavan Dec 14, 2024
94b103e
add additional tests to check task id length
gopidesupavan Dec 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion airflow/models/baseoperator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -965,12 +965,14 @@ def __init__(
category=RemovedInAirflow3Warning,
stacklevel=3,
)
validate_key(task_id)

dag = dag or DagContext.get_current_dag()
task_group = task_group or TaskGroupContext.get_current_task_group(dag)

self.task_id = task_group.child_id(task_id) if task_group else task_id

validate_key(self.task_id)

if not self.__from_mapped and task_group:
task_group.add(self)

Expand Down
27 changes: 27 additions & 0 deletions tests/models/test_baseoperator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -564,6 +564,33 @@ def test_chain(self):
assert [op2] == tgop3.get_direct_relatives(upstream=False)
assert [op2] == tgop4.get_direct_relatives(upstream=False)

def test_baseoperator_raises_exception_when_task_id_plus_taskgroup_id_exceeds_250_chars(self):
"""Test exception is raised when operator task id + taskgroup id > 250 chars."""
dag = DAG(dag_id="foo", schedule=None, start_date=datetime.now())

tg1 = TaskGroup("A" * 20, dag=dag)
with pytest.raises(AirflowException, match="The key has to be less than 250 characters"):
BaseOperator(task_id="1" * 250, task_group=tg1, dag=dag)

def test_baseoperator_with_task_id_and_taskgroup_id_less_than_250_chars(self):
"""Test exception is not raised when operator task id + taskgroup id < 250 chars."""
dag = DAG(dag_id="foo", schedule=None, start_date=datetime.now())

tg1 = TaskGroup("A" * 10, dag=dag)
try:
BaseOperator(task_id="1" * 239, task_group=tg1, dag=dag)
except Exception as e:
pytest.fail(f"Exception raised: {e}")

def test_baseoperator_with_task_id_less_than_250_chars(self):
"""Test exception is not raised when operator task id < 250 chars."""
dag = DAG(dag_id="foo", schedule=None, start_date=datetime.now())

try:
BaseOperator(task_id="1" * 249, dag=dag)
except Exception as e:
pytest.fail(f"Exception raised: {e}")

def test_chain_linear(self):
dag = DAG(dag_id="test_chain_linear", schedule=None, start_date=datetime.now())

Expand Down