[v3-3-test] Add support for operators that achieve durable execution manually in registry (#70291)

dev/registry/extract_parameters.py

@@ -386,11 +386,29 @@ def load_resumable_job_mixin() -> type | None:
 
 
 def is_durable_capable(cls: type, resumable_mixin: type | None) -> bool:
-    """Return True if a class fully implements ResumableJobMixin's crash-recovery contract.
-
-    Inheriting the mixin is not sufficient: a complete override is inert unless
-    execute() actually calls execute_resumable().
+    """Return True if a class implements durable/crash-safe execution.
+
+    Two ways to qualify:
+    1. A class-level `__supports_durable_execution = True`
+    declaration (for operators that implement this directly against
+    task_state_store, without ResumableJobMixin -- e.g. KubernetesPodOperator,
+    AgentOperator).
+    2. Genuinely implementing ResumableJobMixin's contract.
+
+    The first path deliberately looks up the class prefixed attribute
+    (`_{ClassName}__supports_durable_execution`) rather than a fixed string.
+    A subclass that overrides execute() itself (e.g. SparkKubernetesOperator)
+    may not preserve the parent's task_state_store reconnect behavior, so the
+    declaration must not be inherited -- only the exact class that wrote
+    `__supports_durable_execution` in its own body qualifies this way.
+
+    Inheriting the mixin alone is not sufficient for the second path: a
+    complete override is inert unless execute() actually calls
+    execute_resumable().
     """
+    if getattr(cls, f"_{cls.__name__}__supports_durable_execution", None) is True:
+        return True
+
     if resumable_mixin is None or resumable_mixin not in cls.__mro__:
         return False
 

dev/registry/tests/test_extract_parameters.py

@@ -225,6 +225,24 @@ def execute(self, context):
         return None
 
 
+class ManuallyDurableOperator:
+    """Implements durable execution directly (e.g. via task_state_store), without
+    ResumableJobMixin -- mirrors KubernetesPodOperator/AgentOperator."""
+
+    __supports_durable_execution = True
+
+    def execute(self, context):
+        return None
+
+
+class ManuallyDurableSubclass(ManuallyDurableOperator):
+    """Overrides execute() itself -- must NOT inherit the parent's declaration,
+    since nothing here verifies it preserves the reconnect behavior."""
+
+    def execute(self, context):
+        return "something else entirely"
+
+
 class TestIsDurableCapable:
     def test_fully_implemented_and_wired_qualifies(self):
         assert is_durable_capable(FullyImplementedResumableOperator, FakeResumableJobMixin) is True
@@ -241,6 +259,12 @@ def test_no_mixin_in_mro_disqualifies(self):
     def test_mixin_unavailable_disqualifies(self):
         assert is_durable_capable(FullyImplementedResumableOperator, None) is False
 
+    def test_manual_durable_marker_qualifies_without_mixin(self):
+        assert is_durable_capable(ManuallyDurableOperator, None) is True
+
+    def test_subclass_not_redeclaring_marker_disqualifies(self):
+        assert is_durable_capable(ManuallyDurableSubclass, FakeResumableJobMixin) is False
+
 
 # ---------------------------------------------------------------------------
 # Module dataclass