Adding run_after field in task SDK datamodels

task_sdk/src/airflow/sdk/api/datamodels/_generated.py

@@ -283,6 +283,7 @@ class DagRun(BaseModel):
     logical_date: Annotated[datetime, Field(title="Logical Date")]
     data_interval_start: Annotated[datetime | None, Field(title="Data Interval Start")] = None
     data_interval_end: Annotated[datetime | None, Field(title="Data Interval End")] = None
+    run_after: Annotated[datetime, Field(title="Run After")]
     start_date: Annotated[datetime, Field(title="Start Date")]
     end_date: Annotated[datetime | None, Field(title="End Date")] = None
     run_type: DagRunType

task_sdk/tests/conftest.py

@@ -154,7 +154,9 @@ def __call__(
         data_interval_start: str | datetime = ...,
         data_interval_end: str | datetime = ...,
         start_date: str | datetime = ...,
+        run_after: str | datetime = ...,
         run_type: str = ...,
+        conf=None,
     ) -> TIRunContext: ...
 
 
@@ -167,7 +169,9 @@ def __call__(
         data_interval_start: str | datetime = ...,
         data_interval_end: str | datetime = ...,
         start_date: str | datetime = ...,
+        run_after: str | datetime = ...,
         run_type: str = ...,
+        conf=None,
     ) -> dict[str, Any]: ...
 
 
@@ -183,6 +187,7 @@ def _make_context(
         data_interval_start: str | datetime = "2024-12-01T00:00:00Z",
         data_interval_end: str | datetime = "2024-12-01T01:00:00Z",
         start_date: str | datetime = "2024-12-01T01:00:00Z",
+        run_after: str | datetime = "2024-12-01T01:00:00Z",
         run_type: str = "manual",
         conf=None,
     ) -> TIRunContext:
@@ -195,6 +200,7 @@ def _make_context(
                 data_interval_end=data_interval_end,  # type: ignore
                 start_date=start_date,  # type: ignore
                 run_type=run_type,  # type: ignore
+                run_after=run_after,  # type: ignore
                 conf=conf,
             ),
             max_tries=0,
@@ -214,7 +220,9 @@ def _make_context_dict(
         data_interval_start: str | datetime = "2024-12-01T00:00:00Z",
         data_interval_end: str | datetime = "2024-12-01T01:00:00Z",
         start_date: str | datetime = "2024-12-01T00:00:00Z",
+        run_after: str | datetime = "2024-12-01T00:00:00Z",
         run_type: str = "manual",
+        conf=None,
     ) -> dict[str, Any]:
         context = make_ti_context(
             dag_id=dag_id,
@@ -223,7 +231,9 @@ def _make_context_dict(
             data_interval_start=data_interval_start,
             data_interval_end=data_interval_end,
             start_date=start_date,
+            run_after=run_after,
             run_type=run_type,
+            conf=conf,
         )
         return context.model_dump(exclude_unset=True, mode="json")
 

task_sdk/tests/execution_time/conftest.py

@@ -132,6 +132,7 @@ def _create_task_instance(
         data_interval_start: str | datetime = "2024-12-01T00:00:00Z",
         data_interval_end: str | datetime = "2024-12-01T01:00:00Z",
         start_date: str | datetime = "2024-12-01T01:00:00Z",
+        run_after: str | datetime = "2024-12-01T01:00:00Z",
         run_type: str = "manual",
         try_number: int = 1,
         conf=None,
@@ -148,6 +149,7 @@ def _create_task_instance(
             data_interval_end=data_interval_end,
             start_date=start_date,
             run_type=run_type,
+            run_after=run_after,
             conf=conf,
         )
 

task_sdk/tests/execution_time/test_supervisor.py

@@ -238,7 +238,7 @@ def subprocess_main():
         assert "Last chance exception handler" in captured.err
         assert "RuntimeError: Fake syntax error" in captured.err
 
-    def test_regular_heartbeat(self, spy_agency: kgb.SpyAgency, monkeypatch):
+    def test_regular_heartbeat(self, spy_agency: kgb.SpyAgency, monkeypatch, mocker, make_ti_context):
         """Test that the WatchedSubprocess class regularly sends heartbeat requests, up to a certain frequency"""
         import airflow.sdk.execution_time.supervisor
 
@@ -252,6 +252,8 @@ def subprocess_main():
                 sleep(0.05)
 
         ti_id = uuid7()
+        _ = mocker.patch.object(sdk_client.TaskInstanceOperations, "start", return_value=make_ti_context())
+
         spy = spy_agency.spy_on(sdk_client.TaskInstanceOperations.heartbeat)
         proc = ActivitySubprocess.start(
             dag_rel_path=os.devnull,
@@ -271,7 +273,7 @@ def subprocess_main():
         # The exact number we get will depend on timing behaviour, so be a little lenient
         assert 1 <= len(spy.calls) <= 4
 
-    def test_run_simple_dag(self, test_dags_dir, captured_logs, time_machine):
+    def test_run_simple_dag(self, test_dags_dir, captured_logs, time_machine, mocker, make_ti_context):
         """Test running a simple DAG in a subprocess and capturing the output."""
 
         instant = tz.datetime(2024, 11, 7, 12, 34, 56, 78901)
@@ -285,6 +287,12 @@ def test_run_simple_dag(self, test_dags_dir, captured_logs, time_machine):
             run_id="c",
             try_number=1,
         )
+
+        # Create a mock client to assert calls to the client
+        # We assume the implementation of the client is correct and only need to check the calls
+        mock_client = mocker.Mock(spec=sdk_client.Client)
+        mock_client.task_instances.start.return_value = make_ti_context()
+
         bundle_info = BundleInfo(name="my-bundle", version=None)
         with patch.dict(os.environ, local_dag_bundle_cfg(test_dags_dir, bundle_info.name)):
             exit_code = supervise(
@@ -293,6 +301,7 @@ def test_run_simple_dag(self, test_dags_dir, captured_logs, time_machine):
                 token="",
                 server="",
                 dry_run=True,
+                client=mock_client,
                 bundle_info=bundle_info,
             )
             assert exit_code == 0, captured_logs

task_sdk/tests/execution_time/test_task_runner.py

@@ -108,7 +108,7 @@ def test_recv_StartupDetails(self):
             b'"id": "4d828a62-a417-4936-a7a6-2b3fabacecab", "task_id": "a", "try_number": 1, "run_id": "b", "dag_id": "c" }, '
             b'"ti_context":{"dag_run":{"dag_id":"c","run_id":"b","logical_date":"2024-12-01T01:00:00Z",'
             b'"data_interval_start":"2024-12-01T00:00:00Z","data_interval_end":"2024-12-01T01:00:00Z",'
-            b'"start_date":"2024-12-01T01:00:00Z","end_date":null,"run_type":"manual","conf":null},'
+            b'"start_date":"2024-12-01T01:00:00Z","run_after":"2024-12-01T01:00:00Z","end_date":null,"run_type":"manual","conf":null},'
             b'"max_tries":0,"variables":null,"connections":null},"file": "/dev/null", "dag_rel_path": "/dev/null", "bundle_info": {"name": '
             b'"any-name", "version": "any-version"}, "requests_fd": '
             + str(w2.fileno()).encode("ascii")