Making ResumableJobMixin a abstract class and its methods abstractmethods

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

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):
+            partial_cls(task_id="test_task")