ferruzzi commented on code in PR #70148:
URL: https://github.com/apache/airflow/pull/70148#discussion_r3661506444
##########
airflow-core/src/airflow/models/serialized_dag.py:
##########
@@ -645,6 +646,12 @@ def write_dag(
name_updated = False
reused_deadline_data: dict[str, dict] | None = None
if dag.data.get("dag", {}).get("deadline"):
+ # The deadline handling below rewrites data["dag"]["deadline"]
from a list of
+ # encoded dicts into a list of UUID references. Work on a copy so
we never mutate
+ # the caller's LazyDeserializedDAG in place.
+ from airflow.serialization.serialized_objects import
LazyDeserializedDAG
Review Comment:
I don't think this needs to be a local import; move it to the top if you can.
##########
airflow-core/tests/unit/models/test_serialized_dag.py:
##########
@@ -1155,3 +1156,99 @@ def
test_deadline_reuse_skips_write_when_hash_matches(self, testing_dag_bundle,
alert = session.scalar(select(DAM).where(DAM.serialized_dag_id ==
orig_serdag.id))
assert alert is not None
assert alert.id == orig_alert.id
+
+ def test_write_dag_with_deadline_passes_schema_validation(self,
testing_dag_bundle, session):
+ """The persisted serialized Dag for a deadline-bearing Dag must
satisfy the JSON schema.
+
+ write_dag stores ``data["dag"]["deadline"]`` as a list of UUID strings
referencing
+ deadline_alert rows, so the schema has to accept that persisted form
and not only the
+ list-of-dicts form produced before the dict->UUID rewrite.
+ """
+ dag_id = "test_deadline_schema_valid"
+ dag = DAG(
+ dag_id=dag_id,
+ deadline=DeadlineAlert(
+ reference=DeadlineReference.DAGRUN_QUEUED_AT,
+ interval=timedelta(minutes=5),
+ callback=AsyncCallback(empty_callback_for_deadline),
+ ),
+ )
+ EmptyOperator(task_id="task1", dag=dag)
+ sync_dag_to_db(dag, session=session)
+ session.commit()
+
+ result = session.scalar(select(SDM).where(SDM.dag_id == dag_id))
+ persisted_deadline = result.data["dag"]["deadline"]
+ assert isinstance(persisted_deadline, list)
+ assert persisted_deadline
+ assert all(isinstance(ref, str) for ref in persisted_deadline)
+
+ # Must not raise: the stored UUID-reference form has to satisfy the
serialized Dag schema.
+ DagSerialization.validate_schema(result.data)
+
+ def test_write_dag_does_not_mutate_caller_deadline_data(self,
testing_dag_bundle, session):
+ """write_dag must not rewrite the caller's LazyDeserializedDAG
deadline in place.
+
+ The dict->UUID replacement in ``_generate_deadline_uuids`` has to
happen on a copy so a
+ LazyDeserializedDAG the caller still references keeps its original
list-of-dicts deadline.
+ """
+ dag_id = "test_deadline_no_mutation"
+ dag = DAG(
+ dag_id=dag_id,
+ deadline=DeadlineAlert(
+ reference=DeadlineReference.DAGRUN_QUEUED_AT,
+ interval=timedelta(minutes=5),
+ callback=AsyncCallback(empty_callback_for_deadline),
+ ),
+ )
+ EmptyOperator(task_id="task1", dag=dag)
+ sync_dag_to_db(dag, session=session)
+ session.commit()
+
+ # Change the interval so write_dag regenerates UUIDs (the dict->UUID
rewrite path)
+ # rather than reusing the existing ones.
+ dag.deadline = DeadlineAlert(
+ reference=DeadlineReference.DAGRUN_QUEUED_AT,
+ interval=timedelta(minutes=10),
+ callback=AsyncCallback(empty_callback_for_deadline),
+ )
+ lazy_dag = LazyDeserializedDAG.from_dag(dag)
+ original_deadline = copy.deepcopy(lazy_dag.data["dag"]["deadline"])
+ assert original_deadline
+ assert all(isinstance(item, dict) for item in original_deadline)
+
+ SDM.write_dag(lazy_dag, bundle_name="testing", session=session)
+ session.commit()
+
+ assert lazy_dag.data["dag"]["deadline"] == original_deadline
+
+ def test_sync_dag_to_db_returns_db_normalized_deadline_ids(self,
testing_dag_bundle, session):
+ """sync_dag_to_db must return a SerializedDAG with the DB-normalized
deadline UUIDs. Basically just verify that the UUIDs returned by sync_dag_to_db
match the persisted deadline_alert rows in the DB."""
Review Comment:
```suggestion
"""sync_dag_to_db must return a SerializedDAG with the DB-normalized
deadline UUIDs. Verify that the UUIDs returned by sync_dag_to_db match the
persisted deadline_alert rows in the DB."""
```
##########
airflow-core/src/airflow/models/serialized_dag.py:
##########
@@ -645,6 +646,12 @@ def write_dag(
name_updated = False
reused_deadline_data: dict[str, dict] | None = None
if dag.data.get("dag", {}).get("deadline"):
+ # The deadline handling below rewrites data["dag"]["deadline"]
from a list of
+ # encoded dicts into a list of UUID references. Work on a copy so
we never mutate
+ # the caller's LazyDeserializedDAG in place.
+ from airflow.serialization.serialized_objects import
LazyDeserializedDAG
+
+ dag = LazyDeserializedDAG(data=copy.deepcopy(dag.data),
last_loaded=dag.last_loaded)
Review Comment:
This is fragile. If someone ever adds a new field to `LazyDeserializedDAG`,
this will silently drop it and would be very hard to troubleshoot. Instead,
let's make a copy keeping the existing parameters and just update `data`:
```suggestion
dag = dag.model_copy(update={"data": copy.deepcopy(dag.data)})
```
Some might say the deepcopy isn't strictly required here, but this is
already an O(n) method since the hash calculation is already walking the full
dag data anyway. It's not adding any meaningful extra overhead for the safety
it provides. I like it.
In fact, if we do that, it also eliminates the new `LazyDeserializedDAG`
import above.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
