From 589c00cca0dac20bc304acb2233a241ac9cffee7 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Wed, 8 Jul 2026 16:09:29 +0530 Subject: [PATCH 1/2] Making ResumableJobMixin a abstract class and its methods abstractmethods --- .../airflow/sdk/bases/resumablejobmixin.py | 11 +- .../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 27533dbe8409c..ef1629da97503 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 @@ 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 @@ def execute_resumable(self, context: Context) -> Any: 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 @@ def submit_job(self, context: Context) -> JsonValue: """ 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 @@ def get_job_status(self, external_id: JsonValue, context: Context) -> str: """ 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 @@ def is_job_active(self, status: str) -> bool: """ raise NotImplementedError + @abstractmethod def is_job_succeeded(self, status: str) -> bool: """ Return True if the job completed successfully. @@ -247,10 +252,12 @@ def is_job_succeeded(self, status: str) -> bool: """ 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 8e962583e992d..af446473bb7cc 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 @@ def test_log_fields_for_stored_job(self, status, event_fragment, log_level, extr 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, match=f"abstract method '{missing_method}'"): + partial_cls(task_id="test_task") From 22b8867de40f35c46545a5703a640cfebeff6096 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Thu, 9 Jul 2026 11:15:16 +0530 Subject: [PATCH 2/2] fixing broken tests --- task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py b/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py index af446473bb7cc..1f84baa2d89e4 100644 --- a/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py +++ b/task-sdk/tests/task_sdk/bases/test_resumablejobmixin.py @@ -552,5 +552,5 @@ def test_missing_all_methods_reports_every_one(self): ) def test_missing_single_method_fails_at_construction(self, partial_cls, missing_method): assert partial_cls.__abstractmethods__ == frozenset({missing_method}) - with pytest.raises(TypeError, match=f"abstract method '{missing_method}'"): + with pytest.raises(TypeError): partial_cls(task_id="test_task")