amoghrajesh commented on code in PR #70298:
URL: https://github.com/apache/airflow/pull/70298#discussion_r4012645933
##########
dev/registry/extract_parameters.py:
##########
@@ -391,28 +392,138 @@ def load_resumable_job_mixin() -> type | None:
return None
+# Matches an actual self.defer() call or self.deferrable attribute read, but
not
+# self.defer_for_approval(). TaskDeferred catches operators that raise it
directly instead of
+# calling self.defer() (e.g. VespaIngestOperator).
+_DEFERRAL_TOKEN_RE =
re.compile(r"self\.defer\(|self\.deferrable\b|TaskDeferred\b")
+_SELF_CALL_RE = re.compile(r"self\.([A-Za-z_][A-Za-z0-9_]*)\(")
+_SUPER_EXECUTE_RE = re.compile(r"super\(\)\.execute\(")
+# Matches the @task.* decorator idiom of naming the parent class directly
instead of using
+# super() (e.g. `AgentOperator.execute(self, context)` in common.ai's
@task.agent).
+_EXPLICIT_EXECUTE_RE = re.compile(r"\b([A-Za-z_][A-Za-z0-9_]*)\.execute\(")
+
+# To prevent infinite looping, most cases in the repo are 1-2 hops away.
+_MAX_DEFERRAL_WALK_DEPTH = 6
+
+
+def _get_method_source(cls: type, name: str) -> str | None:
+ method = getattr(cls, name, None)
+ if method is None:
+ return None
+ try:
+ return inspect.getsource(method)
+ except (OSError, TypeError):
+ return None
+
+
+def _find_owner_of_execute(mro: tuple[type, ...], start_idx: int) -> type |
None:
+ """Return the first class in `mro[start_idx:]` whose own `__dict__`
defines `execute`."""
+ for cls in mro[start_idx:]:
+ if "execute" in cls.__dict__:
+ return cls
+ return None
+
+
+def _next_execute_via_super(origin: type, current: type) -> type | None:
+ """Find what `super().execute()` resolves to from `current`, per
`origin`'s MRO.
+
+ Must use `origin`'s MRO, not `current`'s own: with multiple inheritance
(e.g. a
+ `@task.kubernetes`-built class) they diverge, and a mixin's own MRO may
have no
+ relationship to the class the chain actually needs to reach.
+ """
+ try:
+ idx = origin.__mro__.index(current)
+ except ValueError:
+ return None
+ return _find_owner_of_execute(origin.__mro__, idx + 1)
+
+
+def _next_execute_hop(origin: type, current: type, source: str) -> type | None:
+ """Find the next class in `source`'s delegation chain, via
`super().execute()` or an
+ explicit `ParentClass.execute(...)` call."""
+ if _SUPER_EXECUTE_RE.search(source):
+ return _next_execute_via_super(origin, current)
+ match = _EXPLICIT_EXECUTE_RE.search(source)
+ if match is None:
+ return None
+ name = match.group(1)
+ return next((cls for cls in origin.__mro__ if cls.__name__ == name), None)
+
+
+def _find_marker_declaring_class(cls: type) -> type | None:
+ """Return the class in `cls`'s MRO whose own body sets
`__supports_durable_execution = True`.
+
+ The lookup is per-class (`_{base.__name__}__supports_durable_execution`),
not a fixed
+ string, and only matches a class whose own `__dict__` carries the
(mangled) name;
+ inheriting the attribute value from a base doesn't count, only writing it
yourself does.
+ Leading underscores in the class name are stripped first, matching
Python's own name
+ mangling rule (`_Foo` mangles to `_Foo__x`, not `__Foo__x`).
+ """
+ for base in cls.__mro__:
+ mangled = f"_{base.__name__.lstrip('_')}__supports_durable_execution"
+ if base.__dict__.get(mangled) is True:
+ return base
+ return None
+
+
+def _delegates_execute_to(cls: type, target: type, depth: int) -> bool:
+ """Return True if `cls`'s resolved `execute()` chain reaches
`target.execute`.
+
+ Covers a class that never overrides `execute` (inherits `target.execute`
directly, e.g.
+ GKEStartPodOperator), one whose override ends in
`super().execute(context)` (e.g.
+ EksPodOperator, SparkKubernetesOperator), and one that names the parent
class directly
+ instead (e.g. `AgentOperator.execute(self, context)` in `@task.agent`).
+ """
+ owner = _find_owner_of_execute(cls.__mro__, 0)
+ remaining = depth
+ while owner is not None:
+ if owner is target:
+ return True
+ if remaining <= 0:
+ return False
+ source = _get_method_source(owner, "execute")
+ if source is None:
+ return False
+ owner = _next_execute_hop(cls, owner, source)
+ remaining -= 1
+ return False
+
+
+def _execute_chain_calls_resumable(cls: type, depth: int) -> bool:
+ """Return True if some class along `cls`'s resolved `execute()` chain
calls execute_resumable().
+
+ Same walk as `_delegates_execute_to`: a delegating override's own source
may not mention
+ `execute_resumable` even though the class it hands off to does.
+ """
+ owner = _find_owner_of_execute(cls.__mro__, 0)
+ remaining = depth
+ while owner is not None:
+ source = _get_method_source(owner, "execute")
+ if source is None:
+ return False
+ if "execute_resumable" in source:
+ return True
+ if remaining <= 0:
+ return False
+ owner = _next_execute_hop(cls, owner, source)
+ remaining -= 1
+ return False
+
+
def is_durable_capable(cls: type, resumable_mixin: type | None) -> bool:
"""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().
+ 1. A class-level `__supports_durable_execution = True` declaration (for
operators like
+ KubernetesPodOperator/AgentOperator that implement this directly against
+ task_state_store, without ResumableJobMixin). Inherited by a subclass that
hasn't
+ replaced the declaring class's `execute()`, whether by not overriding it
at all, or by
+ delegating back via `super().execute()`.
Review Comment:
re the widening scope, keeping this fix in the current PR since it fixes a
real issue either way. Also worth noting: the reasoning for
`SparkKubernetesOperator` is coincidental rather than solid; the answer is
correct but not for the stated reason. Added a docstring note explaining that
so it is not mistaken for a load-bearing invariant. Leaving the durable count
in the description as is for now, since it will change again before this
settles.
--
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]
