amoghrajesh commented on code in PR #70298:
URL: https://github.com/apache/airflow/pull/70298#discussion_r4044097825


##########
dev/registry/extract_parameters.py:
##########
@@ -391,28 +394,159 @@ 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 _strip_comment_lines(source: str) -> str:
+    """Drop whole-line comments so they can't be mistaken for real delegation 
code.
+
+    A comment can say the opposite of what the code does (e.g. "overrides 
execute rather
+    than calling super().execute()"), and a raw-text search can't tell the two 
apart.
+    """
+    return "\n".join(line for line in source.splitlines() if not 
line.strip().startswith("#"))
+
+
+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.
+
+    Every explicit match is tried in order, since an unrelated earlier call 
(e.g.
+    `cursor.execute(...)`) can otherwise shadow the real delegation. The 
matched class is
+    resolved to whoever actually owns `execute` (it may only inherit one, e.g.
+    `GKEStartPodOperator.execute(self, context)`), the same way a `super()` 
hop is.
+    """
+    source = _strip_comment_lines(source)
+    if _SUPER_EXECUTE_RE.search(source):
+        return _next_execute_via_super(origin, current)
+
+    mro_by_name = {base.__name__: base for base in origin.__mro__}
+    for match in _EXPLICIT_EXECUTE_RE.finditer(source):
+        named_cls = mro_by_name.get(match.group(1))
+        if named_cls is None:
+            continue
+        owner = _find_owner_of_execute(origin.__mro__, 
origin.__mro__.index(named_cls))
+        if owner is not None:
+            return owner
+    return 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

Review Comment:
   Handled in 3f2104367f



##########
registry/src/css/main.css:
##########
@@ -3297,6 +3298,27 @@ main {
   max-width: 16rem;
 }
 
+.capability-filter-toggles {

Review Comment:
   Handled in 3f2104367f



##########
dev/registry/extract_parameters.py:
##########
@@ -530,6 +770,7 @@ def make_entry(
             "provider_id": provider_id,
             "provider_name": provider_name,
             "supports_durable_execution": is_durable_capable(cls_or_obj, 
resumable_mixin),
+            "supports_deferrable": supports_deferrable(cls_or_obj),

Review Comment:
   Handled in 3f2104367f



##########
dev/registry/extract_parameters.py:
##########
@@ -421,15 +555,121 @@ def is_durable_capable(cls: type, resumable_mixin: type 
| None) -> bool:
     if inspect.isabstract(cls):
         return False
 
-    execute = getattr(cls, "execute", None)
-    if execute is None:
-        return False
+    return _execute_chain_calls_resumable(cls, _MAX_DEFERRAL_WALK_DEPTH)
+
+
+def _is_terminal_block(body: list[ast.stmt]) -> bool:
+    """Return True if `body`'s last statement always exits the function (raise 
or return)."""
+    return bool(body) and isinstance(body[-1], (ast.Raise, ast.Return))
+
+
+def _strip_dead_version_guard_branches(source: str, resolve_global: 
typing.Callable[[str], object]) -> str:
+    """Blank out code after a terminal `if <flag>: raise ...` whose flag 
resolves True here.
+
+    Some operators write `if AIRFLOW_V_3_3_PLUS: raise ...` with no `else`, 
followed by a
+    pre-3.3 `self.defer(...)` fallback that never runs on the core being 
imported.
+    `resolve_global` looks up the flag by name rather than a fixed list.
+    """
+    try:
+        tree = ast.parse(textwrap.dedent(source))
+    except SyntaxError:
+        return source
+
+    if not tree.body or not isinstance(tree.body[0], (ast.FunctionDef, 
ast.AsyncFunctionDef)):
+        return source
+
+    dead_from: int | None = None
+    for stmt in tree.body[0].body:
+        if (
+            isinstance(stmt, ast.If)
+            and not stmt.orelse
+            and isinstance(stmt.test, ast.Name)
+            and _is_terminal_block(stmt.body)
+            and resolve_global(stmt.test.id) is True
+        ):
+            dead_from = stmt.end_lineno
+            break
+
+    if dead_from is None:
+        return source
+    return "".join(source.splitlines(keepends=True)[:dead_from])
+
+
+def _get_reachable_method_source(cls: type, name: str) -> str | None:
+    """Like `_get_method_source`, but with dead version-guard branches 
stripped first."""
+    method = getattr(cls, name, None)
+    if method is None:
+        return None
     try:
-        source = inspect.getsource(execute)
+        source = inspect.getsource(method)
     except (OSError, TypeError):
+        return None
+
+    # unwrap() undoes a functools.wraps() decorator, whose __globals__ would 
otherwise
+    # point at the decorator's own module instead of the method's.
+    func = inspect.unwrap(getattr(method, "__func__", method))

Review Comment:
   Handled in 3f2104367f



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