vatsrahul1001 commented on code in PR #69243:
URL: https://github.com/apache/airflow/pull/69243#discussion_r3518382727


##########
airflow-core/tests/unit/serialization/test_serialized_objects.py:
##########
@@ -1309,6 +1309,59 @@ def 
test_has_retry_policy_flag_false_when_no_policy(self):
         task = deserialized.task_dict["op_no_policy"]
         assert task.has_retry_policy is False
 
+    def test_mapped_task_retry_policy_serializes_as_flag(self):
+        """A mapped task's retry_policy must serialize as has_retry_policy, 
not the object."""
+        import json

Review Comment:
   Hoisted them — the four `retry_policy` imports are now a single module-level 
import. Kept `from airflow.sdk import DAG` local in each test, since (per your 
earlier note) the module-level `DAG` is `airflow.models.dag.DAG` and hoisting 
the SDK one would shadow it; happy to alias it at top instead (`from 
airflow.sdk import DAG as ...`) if you prefer.



##########
airflow-core/src/airflow/serialization/serialized_objects.py:
##########
@@ -974,6 +974,15 @@ def serialize_mapped_operator(cls, op: MappedOperator) -> 
dict[str, Any]:
                 if cls._is_excluded(v, k, op):
                     continue
 
+                if k == "retry_policy":

Review Comment:
   Done — shrunk to a one-liner pointing at `_HAS_FLAG_FIELDS`.



##########
airflow-core/tests/unit/serialization/test_serialized_objects.py:
##########
@@ -1309,6 +1309,109 @@ def 
test_has_retry_policy_flag_false_when_no_policy(self):
         task = deserialized.task_dict["op_no_policy"]
         assert task.has_retry_policy is False
 
+    def test_mapped_task_retry_policy_serializes_as_flag(self):
+        """A mapped task's retry_policy must serialize as has_retry_policy, 
not the object."""
+        from airflow.sdk import DAG  # module-level DAG is 
airflow.models.dag.DAG
+        from airflow.sdk.definitions.retry_policy import ExceptionRetryPolicy, 
RetryAction, RetryRule
+
+        policy = ExceptionRetryPolicy(
+            rules=[RetryRule(exception=ValueError, action=RetryAction.FAIL, 
reason="bad data")],
+        )
+
+        with DAG(dag_id="test_mapped_retry_policy_ser", 
start_date=DEFAULT_DATE) as dag:
+
+            @task(retries=3, retry_policy=policy)
+            def mapped(x):
+                return x
+
+            mapped.expand(x=[1, 2, 3])
+
+        serialized = DagSerialization.serialize_dag(dag)
+        # The RetryPolicy object must never be embedded -- str(obj) leaks a 
memory address.
+        assert "ExceptionRetryPolicy object at 0x" not in 
json.dumps(serialized)
+
+        deserialized = DagSerialization.deserialize_dag(serialized)
+        assert deserialized.task_dict["mapped"].has_retry_policy is True
+
+    def test_mapped_task_retry_policy_serialization_is_deterministic(self):
+        """Serializing the same mapped-task-with-policy DAG twice yields 
identical output.
+
+        Regression test: the RetryPolicy object was serialized via str(obj), 
embedding a
+        per-process memory address, so every re-parse produced a different 
serialized DAG
+        (and a spurious new DagVersion).
+        """
+        from airflow.sdk import DAG  # module-level DAG is 
airflow.models.dag.DAG
+        from airflow.sdk.definitions.retry_policy import ExceptionRetryPolicy, 
RetryAction, RetryRule
+
+        def build():
+            policy = ExceptionRetryPolicy(
+                rules=[RetryRule(exception=ValueError, 
action=RetryAction.FAIL)],
+            )
+            with DAG(dag_id="test_mapped_retry_policy_determinism", 
start_date=DEFAULT_DATE) as dag:
+
+                @task(retries=3, retry_policy=policy)
+                def mapped(x):
+                    return x
+
+                mapped.expand(x=[1, 2, 3])
+            return dag
+
+        assert DagSerialization.serialize_dag(build()) == 
DagSerialization.serialize_dag(build())
+
+    def test_dag_default_args_retry_policy_serializes_as_flag(self):
+        """A retry_policy in DAG default_args must not leak the object into 
serialized default_args.
+
+        Regression test: serialize_dag serializes the raw default_args dict, 
so a RetryPolicy
+        there hit the str(obj) fallback (memory address -> new DagVersion 
every parse) even
+        though each task's has_retry_policy was set correctly.
+        """
+        from airflow.sdk import DAG  # module-level DAG is 
airflow.models.dag.DAG
+        from airflow.sdk.definitions.retry_policy import ExceptionRetryPolicy, 
RetryAction, RetryRule
+
+        def build():
+            policy = ExceptionRetryPolicy(
+                rules=[RetryRule(exception=ValueError, 
action=RetryAction.FAIL)],
+            )
+            with DAG(
+                dag_id="test_default_args_retry_policy",
+                start_date=DEFAULT_DATE,
+                default_args={"retry_policy": policy},
+            ) as dag:
+
+                @task
+                def plain():
+                    return 1
+
+                plain()
+            return dag
+
+        serialized = DagSerialization.serialize_dag(build())
+        # The RetryPolicy object must never be embedded in serialized 
default_args.
+        assert "ExceptionRetryPolicy object at 0x" not in 
json.dumps(serialized)
+        # Deterministic across independent parses (no embedded memory address).
+        assert DagSerialization.serialize_dag(build()) == 
DagSerialization.serialize_dag(build())

Review Comment:
   Good point — added your two asserts (`has_retry_policy is True`, 
`retry_policy` absent from the serialized `default_args`).



-- 
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]

Reply via email to