Skip to content

Commit 317c5c8

Browse files
kaxilpotiuk
authored andcommitted
[v3-1-test] Fix dag-processor crash when renaming DAG tag case on MySQL (#57113)
When a user changed only the case of a DAG tag (e.g., 'dangerous' to 'DANGEROUS'), the dag-processor would crash with a duplicate key error on MySQL due to case-insensitive collation in the PRIMARY KEY. This occurred because SQLAlchemy executed INSERT operations before DELETE operations during the flush. The fix ensures DELETE operations complete before attempting INSERT operations by explicitly flushing and refreshing the tag relationship from the database. Fixes #56940 (cherry picked from commit 5013aad) Co-authored-by: Kaxil Naik <[email protected]>
1 parent 735b89b commit 317c5c8

2 files changed

Lines changed: 55 additions & 1 deletion

File tree

‎airflow-core/src/airflow/dag_processing/collection.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -158,10 +158,32 @@ def calculate(cls, dags: dict[str, LazyDeserializedDAG], *, session: Session) ->
158158

159159
def _update_dag_tags(tag_names: set[str], dm: DagModel, *, session: Session) -> None:
160160
orm_tags = {t.name: t for t in dm.tags}
161+
tags_to_delete = []
161162
for name, orm_tag in orm_tags.items():
162163
if name not in tag_names:
163164
session.delete(orm_tag)
164-
dm.tags.extend(DagTag(name=name, dag_id=dm.dag_id) for name in tag_names.difference(orm_tags))
165+
tags_to_delete.append(orm_tag)
166+
167+
tags_to_add = tag_names.difference(orm_tags)
168+
if tags_to_delete:
169+
# Remove deleted tags from the collection to keep it in sync
170+
for tag in tags_to_delete:
171+
dm.tags.remove(tag)
172+
173+
# Check if there's a potential case-only rename on MySQL (e.g., 'tag' -> 'TAG').
174+
# MySQL uses case-insensitive collation for the (name, dag_id) primary key by default,
175+
# which can cause duplicate key errors when renaming tags with only case changes.
176+
if get_dialect_name(session) == "mysql":
177+
orm_tags_lower = {name.lower(): name for name in orm_tags}
178+
has_case_only_change = any(tag.lower() in orm_tags_lower for tag in tags_to_add)
179+
180+
if has_case_only_change:
181+
# Force DELETE operations to execute before INSERT operations.
182+
session.flush()
183+
# Refresh the tags relationship from the database to reflect the deletions.
184+
session.expire(dm, ["tags"])
185+
186+
dm.tags.extend(DagTag(name=name, dag_id=dm.dag_id) for name in tags_to_add)
165187

166188

167189
def _update_dag_owner_links(dag_owner_links: dict[str, str], dm: DagModel, *, session: Session) -> None:

‎airflow-core/tests/unit/dag_processing/test_collection.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
AssetModelOperation,
3838
DagModelOperation,
3939
_get_latest_runs_stmt,
40+
_update_dag_tags,
4041
update_dag_parsing_results_in_db,
4142
)
4243
from airflow.exceptions import SerializationError
@@ -48,6 +49,7 @@
4849
DagScheduleAssetNameReference,
4950
DagScheduleAssetUriReference,
5051
)
52+
from airflow.models.dag import DagTag
5153
from airflow.models.errors import ParseImportError
5254
from airflow.models.serialized_dag import SerializedDagModel
5355
from airflow.providers.standard.operators.empty import EmptyOperator
@@ -931,3 +933,33 @@ def test_max_consecutive_failed_dag_runs_defaults_from_conf_when_none(
931933
update_dag_parsing_results_in_db("testing", None, [dag], {}, 0.1, set(), session)
932934
orm_dag = session.get(DagModel, "dag_max_failed_runs_default")
933935
assert orm_dag.max_consecutive_failed_dag_runs == 6
936+
937+
938+
@pytest.mark.db_test
939+
class TestUpdateDagTags:
940+
@pytest.fixture(autouse=True)
941+
def setup_teardown(self, session):
942+
yield
943+
session.query(DagModel).filter(DagModel.dag_id == "test_dag").delete()
944+
session.commit()
945+
946+
@pytest.mark.parametrize(
947+
["initial_tags", "new_tags", "expected_tags"],
948+
[
949+
(["dangerous"], {"DANGEROUS"}, {"DANGEROUS"}),
950+
(["existing"], {"existing", "new"}, {"existing", "new"}),
951+
(["tag1", "tag2"], {"tag1"}, {"tag1"}),
952+
(["keep", "remove", "lowercase"], {"keep", "LOWERCASE", "new"}, {"keep", "LOWERCASE", "new"}),
953+
(["tag1", "tag2"], set(), set()),
954+
],
955+
)
956+
def test_update_dag_tags(self, testing_dag_bundle, session, initial_tags, new_tags, expected_tags):
957+
dag_model = DagModel(dag_id="test_dag", bundle_name="testing")
958+
dag_model.tags = [DagTag(name=tag, dag_id="test_dag") for tag in initial_tags]
959+
session.add(dag_model)
960+
session.commit()
961+
962+
_update_dag_tags(new_tags, dag_model, session=session)
963+
session.commit()
964+
965+
assert {t.name for t in dag_model.tags} == expected_tags

0 commit comments

Comments
 (0)