This is an automated email from the ASF dual-hosted git repository.
amoghrajesh pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f2f02d2bab3 Making ResumableJobMixin a abstract class and its methods
abstractmethods (#69607)
f2f02d2bab3 is described below
commit f2f02d2bab36acb3fe8a20af6e92301c07308ee0
Author: Amogh Desai <[email protected]>
AuthorDate: Thu Jul 9 12:44:55 2026 +0530
Making ResumableJobMixin a abstract class and its methods abstractmethods
(#69607)
---
.../src/airflow/sdk/bases/resumablejobmixin.py | 11 +-
.../tests/task_sdk/bases/test_resumablejobmixin.py | 152 +++++++++++++++++++++
2 files changed, 161 insertions(+), 2 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/bases/resumablejobmixin.py
b/task-sdk/src/airflow/sdk/bases/resumablejobmixin.py
index 27533dbe840..ef1629da975 100644
--- a/task-sdk/src/airflow/sdk/bases/resumablejobmixin.py
+++ b/task-sdk/src/airflow/sdk/bases/resumablejobmixin.py
@@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations
+from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any
from opentelemetry import trace
@@ -32,7 +33,7 @@ if TYPE_CHECKING:
tracer = trace.get_tracer(__name__)
-class ResumableJobMixin:
+class ResumableJobMixin(ABC):
"""
Mixin for operators that submit one long-running job to an external system
and poll for completion.
@@ -52,7 +53,7 @@ class ResumableJobMixin:
Usage: call ``execute_resumable(context)`` from the operator's
``execute()`` when reconnection
is supported.
- Subclasses must implement the methods specific to their external system.
The mixin owns
+ Subclasses must implement all the methods specific to their external
system. The mixin owns
only ``execute_resumable()`` and the task_state_store read/write logic.
Example::
@@ -209,6 +210,7 @@ class ResumableJobMixin:
self.poll_until_complete(external_id, context)
return self.get_job_result(external_id, context)
+ @abstractmethod
def submit_job(self, context: Context) -> JsonValue:
"""
Submit the job to the external system. Return its external ID.
@@ -218,6 +220,7 @@ class ResumableJobMixin:
"""
raise NotImplementedError
+ @abstractmethod
def get_job_status(self, external_id: JsonValue, context: Context) -> str:
"""
Query the external system for the current job status.
@@ -229,6 +232,7 @@ class ResumableJobMixin:
"""
raise NotImplementedError
+ @abstractmethod
def is_job_active(self, status: str) -> bool:
"""
Return True if the job is still running and can be reconnected to.
@@ -238,6 +242,7 @@ class ResumableJobMixin:
"""
raise NotImplementedError
+ @abstractmethod
def is_job_succeeded(self, status: str) -> bool:
"""
Return True if the job completed successfully.
@@ -247,10 +252,12 @@ class ResumableJobMixin:
"""
raise NotImplementedError
+ @abstractmethod
def poll_until_complete(self, external_id: JsonValue, context: Context) ->
None:
"""Block until the job reaches a terminal state. Raise on failure."""
raise NotImplementedError
+ @abstractmethod
def get_job_result(self, external_id: JsonValue, context: Context) -> Any:
"""Return the job result after completion. Return None if not
applicable."""
raise NotImplementedError
diff --git a/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py
b/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py
index 8e962583e99..1f84baa2d89 100644
--- a/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py
+++ b/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py
@@ -402,3 +402,155 @@ class TestLogging:
assert entry["external_id"] == "job-001"
for key, val in extra_fields.items():
assert entry[key] == val
+
+
+class _MissingSubmitJob(ResumableJobMixin, BaseOperator):
+ def get_job_status(self, external_id, context) -> str:
+ return "RUNNING"
+
+ def is_job_active(self, status: str) -> bool:
+ return True
+
+ def is_job_succeeded(self, status: str) -> bool:
+ return False
+
+ def poll_until_complete(self, external_id, context) -> None:
+ return None
+
+ def get_job_result(self, external_id, context):
+ return None
+
+
+class _MissingGetJobStatus(ResumableJobMixin, BaseOperator):
+ def submit_job(self, context):
+ return "id"
+
+ def is_job_active(self, status: str) -> bool:
+ return True
+
+ def is_job_succeeded(self, status: str) -> bool:
+ return False
+
+ def poll_until_complete(self, external_id, context) -> None:
+ return None
+
+ def get_job_result(self, external_id, context):
+ return None
+
+
+class _MissingIsJobActive(ResumableJobMixin, BaseOperator):
+ def submit_job(self, context):
+ return "id"
+
+ def get_job_status(self, external_id, context) -> str:
+ return "RUNNING"
+
+ def is_job_succeeded(self, status: str) -> bool:
+ return False
+
+ def poll_until_complete(self, external_id, context) -> None:
+ return None
+
+ def get_job_result(self, external_id, context):
+ return None
+
+
+class _MissingIsJobSucceeded(ResumableJobMixin, BaseOperator):
+ def submit_job(self, context):
+ return "id"
+
+ def get_job_status(self, external_id, context) -> str:
+ return "RUNNING"
+
+ def is_job_active(self, status: str) -> bool:
+ return True
+
+ def poll_until_complete(self, external_id, context) -> None:
+ return None
+
+ def get_job_result(self, external_id, context):
+ return None
+
+
+class _MissingPollUntilComplete(ResumableJobMixin, BaseOperator):
+ def submit_job(self, context):
+ return "id"
+
+ def get_job_status(self, external_id, context) -> str:
+ return "RUNNING"
+
+ def is_job_active(self, status: str) -> bool:
+ return True
+
+ def is_job_succeeded(self, status: str) -> bool:
+ return False
+
+ def get_job_result(self, external_id, context):
+ return None
+
+
+class _MissingGetJobResult(ResumableJobMixin, BaseOperator):
+ def submit_job(self, context):
+ return "id"
+
+ def get_job_status(self, external_id, context) -> str:
+ return "RUNNING"
+
+ def is_job_active(self, status: str) -> bool:
+ return True
+
+ def is_job_succeeded(self, status: str) -> bool:
+ return False
+
+ def poll_until_complete(self, external_id, context) -> None:
+ return None
+
+
+class _MissingEverything(ResumableJobMixin, BaseOperator):
+ pass
+
+
+class TestAbstractMethodEnforcement:
+ def test_mixin_itself_cannot_be_instantiated(self):
+ with pytest.raises(TypeError, match="Can't instantiate abstract class
ResumableJobMixin"):
+ ResumableJobMixin()
+
+ def test_fully_implemented_subclass_has_no_abstract_methods(self):
+ assert ConcreteResumableOperator.__abstractmethods__ == frozenset()
+
+ def
test_fully_implemented_subclass_instantiates_and_behaves_normally(self):
+ op = ConcreteResumableOperator(task_id="test_task")
+ assert isinstance(op, ResumableJobMixin)
+
+ op.execute_resumable(make_context(FakeTaskState()))
+ assert op.submitted_ids == ["job-001"]
+
+ def test_missing_all_methods_reports_every_one(self):
+ assert _MissingEverything.__abstractmethods__ == frozenset(
+ {
+ "submit_job",
+ "get_job_status",
+ "is_job_active",
+ "is_job_succeeded",
+ "poll_until_complete",
+ "get_job_result",
+ }
+ )
+ with pytest.raises(TypeError, match="Can't instantiate abstract class
_MissingEverything"):
+ _MissingEverything(task_id="test_task")
+
+ @pytest.mark.parametrize(
+ ("partial_cls", "missing_method"),
+ [
+ pytest.param(_MissingSubmitJob, "submit_job", id="submit_job"),
+ pytest.param(_MissingGetJobStatus, "get_job_status",
id="get_job_status"),
+ pytest.param(_MissingIsJobActive, "is_job_active",
id="is_job_active"),
+ pytest.param(_MissingIsJobSucceeded, "is_job_succeeded",
id="is_job_succeeded"),
+ pytest.param(_MissingPollUntilComplete, "poll_until_complete",
id="poll_until_complete"),
+ pytest.param(_MissingGetJobResult, "get_job_result",
id="get_job_result"),
+ ],
+ )
+ def test_missing_single_method_fails_at_construction(self, partial_cls,
missing_method):
+ assert partial_cls.__abstractmethods__ == frozenset({missing_method})
+ with pytest.raises(TypeError):
+ partial_cls(task_id="test_task")
