Merge branch 'main' into fix/stackdriver-read-run-id

Author: goingforstudying-ctrl

airflow-core/src/airflow/config_templates/config.yml

@@ -3226,7 +3226,7 @@ state_store:
       description: |
         Number of days to retain task state after their last update.
         Rows older than this are removed when cleanup is triggered.
-        This config does not affect asset_state rows.
+        This config does not affect asset_state_store rows.
         Set to 0 to disable time-based cleanup entirely.
       version_added: 3.3.0
       type: integer
@@ -3235,7 +3235,7 @@ state_store:
     state_cleanup_batch_size:
       description: |
         Number of rows deleted per batch during cleanup. Defaults to 0 (no batching).
-        Tune this on deployments with large task_state tables to improve performance per transaction.
+        Tune this on deployments with large task_state_store tables to improve performance per transaction.
       version_added: 3.3.0
       type: integer
       example: "10000"

airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_asset_state_store.py

@@ -42,7 +42,7 @@ def _create_asset(session) -> AssetModel:
     return asset
 
 
-def _create_asset_state(session, asset_id: int, key: str, value: str) -> None:
+def _create_asset_state_store_row(session, asset_id: int, key: str, value: str) -> None:
     row = AssetStateStoreModel(asset_id=asset_id, key=key, value=json.dumps(value))
     session.add(row)
     session.flush()
@@ -91,8 +91,8 @@ def test_returns_empty_list_when_no_state(self, test_client):
         assert response.json() == {"asset_state_store": [], "total_entries": 0}
 
     def test_returns_all_keys(self, test_client):
-        _create_asset_state(self._session, self.asset.id, "watermark", "2026-05-01")
-        _create_asset_state(self._session, self.asset.id, "file_count", "42")
+        _create_asset_state_store_row(self._session, self.asset.id, "watermark", "2026-05-01")
+        _create_asset_state_store_row(self._session, self.asset.id, "file_count", "42")
         self._session.commit()
 
         response = test_client.get(self._base_url)
@@ -103,7 +103,7 @@ def test_returns_all_keys(self, test_client):
         assert keys == {"watermark": "2026-05-01", "file_count": "42"}
 
     def test_returns_metadata_fields(self, test_client):
-        _create_asset_state(self._session, self.asset.id, "watermark", "2026-05-01")
+        _create_asset_state_store_row(self._session, self.asset.id, "watermark", "2026-05-01")
         self._session.commit()
 
         item = test_client.get(self._base_url).json()["asset_state_store"][0]
@@ -134,7 +134,7 @@ def test_last_updated_by_returned_when_set(self, test_client, create_task_instan
 
     def test_pagination_limit(self, test_client):
         for k in ("watermark", "file_count", "last_run"):
-            _create_asset_state(self._session, self.asset.id, k, "v")
+            _create_asset_state_store_row(self._session, self.asset.id, k, "v")
         self._session.commit()
 
         response = test_client.get(f"{self._base_url}?limit=2")
@@ -144,7 +144,7 @@ def test_pagination_limit(self, test_client):
 
     def test_pagination_offset(self, test_client):
         for k in ("watermark", "file_count", "last_run"):
-            _create_asset_state(self._session, self.asset.id, k, "v")
+            _create_asset_state_store_row(self._session, self.asset.id, k, "v")
         self._session.commit()
 
         response = test_client.get(f"{self._base_url}?limit=2&offset=2")
@@ -158,7 +158,7 @@ def test_unauthorized_returns_401(self, unauthenticated_test_client):
 
 class TestGetAssetState(TestAssetStateEndpoint):
     def test_returns_value(self, test_client):
-        _create_asset_state(self._session, self.asset.id, "watermark", "2026-05-01")
+        _create_asset_state_store_row(self._session, self.asset.id, "watermark", "2026-05-01")
         self._session.commit()
 
         response = test_client.get(f"{self._base_url}/watermark")
@@ -196,7 +196,7 @@ def test_missing_key_returns_404(self, test_client):
 
     def test_key_with_slash_is_supported(self, test_client):
         """Keys containing slashes must work — route uses {key:path}."""
-        _create_asset_state(self._session, self.asset.id, "partition/date", "2026-05-01")
+        _create_asset_state_store_row(self._session, self.asset.id, "partition/date", "2026-05-01")
         self._session.commit()
 
         response = test_client.get(f"{self._base_url}/partition/date")
@@ -274,7 +274,7 @@ def test_core_api_write_read_roundtrip(self, test_client, value):
     @pytest.mark.parametrize("value", [42, True, {"rows": 100}, [1, "two"], "hello"])
     def test_worker_write_core_api_read_roundtrip(self, test_client, value):
         """Worker write (json.dumps in DB) then Core API read returns native value."""
-        _create_asset_state(self._session, self.asset.id, "k", value)
+        _create_asset_state_store_row(self._session, self.asset.id, "k", value)
         self._session.commit()
         assert test_client.get(f"{self._base_url}/k").json()["value"] == value
 
@@ -292,7 +292,7 @@ def test_unauthorized_returns_401(self, unauthenticated_test_client):
 
 class TestDeleteAssetState(TestAssetStateEndpoint):
     def test_deletes_key(self, test_client):
-        _create_asset_state(self._session, self.asset.id, "watermark", "2026-05-01")
+        _create_asset_state_store_row(self._session, self.asset.id, "watermark", "2026-05-01")
         self._session.commit()
 
         assert test_client.delete(f"{self._base_url}/watermark").status_code == 204
@@ -302,8 +302,8 @@ def test_delete_noop_for_missing_key(self, test_client):
         assert test_client.delete(f"{self._base_url}/nonexistent").status_code == 204
 
     def test_only_deletes_target_key(self, test_client):
-        _create_asset_state(self._session, self.asset.id, "watermark", "a")
-        _create_asset_state(self._session, self.asset.id, "file_count", "b")
+        _create_asset_state_store_row(self._session, self.asset.id, "watermark", "a")
+        _create_asset_state_store_row(self._session, self.asset.id, "file_count", "b")
         self._session.commit()
 
         test_client.delete(f"{self._base_url}/watermark")
@@ -312,7 +312,7 @@ def test_only_deletes_target_key(self, test_client):
         assert test_client.get(f"{self._base_url}/file_count").json()["value"] == "b"
 
     def test_key_with_slash_is_supported(self, test_client):
-        _create_asset_state(self._session, self.asset.id, "partition/date", "v")
+        _create_asset_state_store_row(self._session, self.asset.id, "partition/date", "v")
         self._session.commit()
 
         assert test_client.delete(f"{self._base_url}/partition/date").status_code == 204
@@ -325,7 +325,7 @@ def test_unauthorized_returns_401(self, unauthenticated_test_client):
 class TestClearAssetState(TestAssetStateEndpoint):
     def test_clears_all_keys(self, test_client):
         for k, v in [("watermark", "a"), ("file_count", "b"), ("last_run", "c")]:
-            _create_asset_state(self._session, self.asset.id, k, v)
+            _create_asset_state_store_row(self._session, self.asset.id, k, v)
         self._session.commit()
 
         assert test_client.delete(self._base_url).status_code == 204
@@ -338,8 +338,8 @@ def test_clear_does_not_affect_other_assets(self, test_client):
         other_asset = AssetModel(uri="s3://other/asset", name="other_asset", group="test")
         self._session.add(other_asset)
         self._session.flush()
-        _create_asset_state(self._session, self.asset.id, "watermark", "mine")
-        _create_asset_state(self._session, other_asset.id, "watermark", "theirs")
+        _create_asset_state_store_row(self._session, self.asset.id, "watermark", "mine")
+        _create_asset_state_store_row(self._session, other_asset.id, "watermark", "theirs")
         self._session.commit()
 
         test_client.delete(self._base_url)
@@ -348,7 +348,7 @@ def test_clear_does_not_affect_other_assets(self, test_client):
         assert test_client.get(f"{other_url}/watermark").json()["value"] == "theirs"
 
     def test_clears_slash_keyed_entries(self, test_client):
-        _create_asset_state(self._session, self.asset.id, "partition/date", "v")
+        _create_asset_state_store_row(self._session, self.asset.id, "partition/date", "v")
         self._session.commit()
 
         assert test_client.delete(self._base_url).status_code == 204
@@ -372,7 +372,7 @@ class TestRoutesNeverCallCustomBackend(TestAssetStateEndpoint):
         ],
     )
     def test_route_never_calls_get_state_backend(self, test_client, method, path_suffix, kwargs):
-        _create_asset_state(self._session, self.asset.id, "watermark", "v1")
+        _create_asset_state_store_row(self._session, self.asset.id, "watermark", "v1")
         self._session.commit()
 
         with patch("airflow.state.get_state_backend") as mock_get_backend:

airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_state_store.py

@@ -55,7 +55,7 @@ def _create_dag_run(dag_maker, session):
     session.commit()
 
 
-def _create_task_state(session, key: str, value: str, dag_run: DagRun) -> None:
+def _create_task_state_store_row(session, key: str, value: str, dag_run: DagRun) -> None:
     row = TaskStateStoreModel(
         dag_run_id=dag_run.id,
         dag_id=DAG_ID,
@@ -94,8 +94,8 @@ def test_returns_empty_list_when_no_state(self, test_client):
         assert response.json() == {"task_state_store": [], "total_entries": 0}
 
     def test_returns_all_keys(self, test_client):
-        _create_task_state(self._session, "job_id", "spark_001", self.dag_run)
-        _create_task_state(self._session, "checkpoint", "step_3", self.dag_run)
+        _create_task_state_store_row(self._session, "job_id", "spark_001", self.dag_run)
+        _create_task_state_store_row(self._session, "checkpoint", "step_3", self.dag_run)
         self._session.commit()
 
         response = test_client.get(BASE_URL)
@@ -106,7 +106,7 @@ def test_returns_all_keys(self, test_client):
         assert keys == {"job_id": "spark_001", "checkpoint": "step_3"}
 
     def test_returns_state_metadata_fields(self, test_client):
-        _create_task_state(self._session, "job_id", "spark_001", self.dag_run)
+        _create_task_state_store_row(self._session, "job_id", "spark_001", self.dag_run)
         self._session.commit()
 
         response = test_client.get(BASE_URL)
@@ -133,7 +133,7 @@ def test_map_index_isolation(self, test_client):
 
     def test_pagination_limit(self, test_client):
         for k in ("a", "b", "c"):
-            _create_task_state(self._session, k, "v", self.dag_run)
+            _create_task_state_store_row(self._session, k, "v", self.dag_run)
         self._session.commit()
 
         response = test_client.get(f"{BASE_URL}?limit=2")
@@ -143,7 +143,7 @@ def test_pagination_limit(self, test_client):
 
     def test_pagination_offset(self, test_client):
         for k in ("a", "b", "c"):
-            _create_task_state(self._session, k, "v", self.dag_run)
+            _create_task_state_store_row(self._session, k, "v", self.dag_run)
         self._session.commit()
 
         response = test_client.get(f"{BASE_URL}?limit=2&offset=2")
@@ -157,7 +157,7 @@ def test_unauthorized_returns_401(self, unauthenticated_test_client):
 
 class TestGetTaskState(TestTaskStateEndpoint):
     def test_returns_value(self, test_client):
-        _create_task_state(self._session, "job_id", "spark_001", self.dag_run)
+        _create_task_state_store_row(self._session, "job_id", "spark_001", self.dag_run)
         self._session.commit()
 
         response = test_client.get(f"{BASE_URL}/job_id")
@@ -172,7 +172,7 @@ def test_missing_key_returns_404(self, test_client):
 
     def test_key_with_slash_is_supported(self, test_client):
         """Keys containing slashes must work — route uses {key:path}."""
-        _create_task_state(self._session, "workflow/step_1", "v", self.dag_run)
+        _create_task_state_store_row(self._session, "workflow/step_1", "v", self.dag_run)
         self._session.commit()
 
         response = test_client.get(f"{BASE_URL}/workflow/step_1")
@@ -254,7 +254,7 @@ def test_core_api_write_read_roundtrip(self, test_client, value):
     @pytest.mark.parametrize("value", [42, True, {"rows": 100}, [1, "two"], "hello"])
     def test_worker_write_core_api_read_roundtrip(self, test_client, value):
         """Worker write (json.dumps in DB) then Core API read returns native value."""
-        _create_task_state(self._session, "k", value, self.dag_run)
+        _create_task_state_store_row(self._session, "k", value, self.dag_run)
         self._session.commit()
         assert test_client.get(f"{BASE_URL}/k").json()["value"] == value
 
@@ -314,7 +314,7 @@ def test_unauthorized_returns_401(self, unauthenticated_test_client):
 
 class TestPatchTaskState(TestTaskStateEndpoint):
     def test_patch_updates_value(self, test_client):
-        _create_task_state(self._session, "job_id", "v1", self.dag_run)
+        _create_task_state_store_row(self._session, "job_id", "v1", self.dag_run)
         self._session.commit()
 
         assert test_client.patch(f"{BASE_URL}/job_id", json={"value": "v2"}).status_code == 200
@@ -332,12 +332,12 @@ def test_patch_missing_key_returns_404(self, test_client):
         assert test_client.patch(f"{BASE_URL}/nonexistent", json={"value": "v"}).status_code == 404
 
     def test_patch_empty_body_returns_422(self, test_client):
-        _create_task_state(self._session, "job_id", "v", self.dag_run)
+        _create_task_state_store_row(self._session, "job_id", "v", self.dag_run)
         self._session.commit()
         assert test_client.patch(f"{BASE_URL}/job_id", json={}).status_code == 422
 
     def test_patch_null_value_returns_422(self, test_client):
-        _create_task_state(self._session, "job_id", "v", self.dag_run)
+        _create_task_state_store_row(self._session, "job_id", "v", self.dag_run)
         self._session.commit()
         assert test_client.patch(f"{BASE_URL}/job_id", json={"value": None}).status_code == 422
 
@@ -356,7 +356,7 @@ def test_patch_non_finite_float_rejected_by_validator(self, bad_value):
         ],
     )
     def test_patch_stores_json_encoded_value(self, test_client, value, expected_db):
-        _create_task_state(self._session, "job_id", "initial", self.dag_run)
+        _create_task_state_store_row(self._session, "job_id", "initial", self.dag_run)
         self._session.commit()
         test_client.patch(f"{BASE_URL}/job_id", json={"value": value})
         row = self._session.scalar(
@@ -376,7 +376,7 @@ def test_unauthorized_returns_401(self, unauthenticated_test_client):
 
 class TestDeleteTaskState(TestTaskStateEndpoint):
     def test_deletes_key(self, test_client):
-        _create_task_state(self._session, "job_id", "spark_001", self.dag_run)
+        _create_task_state_store_row(self._session, "job_id", "spark_001", self.dag_run)
         self._session.commit()
 
         assert test_client.delete(f"{BASE_URL}/job_id").status_code == 204
@@ -386,8 +386,8 @@ def test_delete_noop_for_missing_key(self, test_client):
         assert test_client.delete(f"{BASE_URL}/nonexistent").status_code == 204
 
     def test_only_deletes_target_key(self, test_client):
-        _create_task_state(self._session, "job_id", "a", self.dag_run)
-        _create_task_state(self._session, "checkpoint", "b", self.dag_run)
+        _create_task_state_store_row(self._session, "job_id", "a", self.dag_run)
+        _create_task_state_store_row(self._session, "checkpoint", "b", self.dag_run)
         self._session.commit()
 
         test_client.delete(f"{BASE_URL}/job_id")
@@ -396,7 +396,7 @@ def test_only_deletes_target_key(self, test_client):
         assert test_client.get(f"{BASE_URL}/checkpoint").json()["value"] == "b"
 
     def test_key_with_slash_is_supported(self, test_client):
-        _create_task_state(self._session, "workflow/step_1", "v", self.dag_run)
+        _create_task_state_store_row(self._session, "workflow/step_1", "v", self.dag_run)
         self._session.commit()
 
         assert test_client.delete(f"{BASE_URL}/workflow/step_1").status_code == 204
@@ -409,7 +409,7 @@ def test_unauthorized_returns_401(self, unauthenticated_test_client):
 class TestClearTaskState(TestTaskStateEndpoint):
     def test_clears_all_keys(self, test_client):
         for k, v in [("job_id", "a"), ("checkpoint", "b"), ("retry_count", "c")]:
-            _create_task_state(self._session, k, v, self.dag_run)
+            _create_task_state_store_row(self._session, k, v, self.dag_run)
         self._session.commit()
 
         assert test_client.delete(BASE_URL).status_code == 204
@@ -442,7 +442,7 @@ def test_all_map_indices_clears_across_mapped_instances(self, test_client):
         assert test_client.get(f"{BASE_URL}?map_index=1").json()["total_entries"] == 0
 
     def test_key_with_slash_is_supported(self, test_client):
-        _create_task_state(self._session, "workflow/step_1", "v", self.dag_run)
+        _create_task_state_store_row(self._session, "workflow/step_1", "v", self.dag_run)
         self._session.commit()
 
         assert test_client.delete(BASE_URL).status_code == 204
@@ -467,7 +467,7 @@ class TestRoutesNeverCallCustomBackend(TestTaskStateEndpoint):
         ],
     )
     def test_route_never_calls_get_state_backend(self, test_client, method, path, kwargs):
-        _create_task_state(self._session, "job_id", "v1", self.dag_run)
+        _create_task_state_store_row(self._session, "job_id", "v1", self.dag_run)
         self._session.commit()
 
         with patch("airflow.state.get_state_backend") as mock_get_backend:

providers/apache/spark/docs/operators.rst

@@ -190,7 +190,7 @@ independently on the cluster. If the Airflow worker dies while the Spark job is
 Airflow loses track of it and the behaviour to submit a brand new job would be wasting
 the compute already done or even cause conflicts if the Spark job itself is not designed to be idempotent.
 
-Now, the ``SparkSubmitOperator`` solves this by persisting the driver ID to ``task_state`` immediately after
+Now, the ``SparkSubmitOperator`` solves this by persisting the driver ID to ``task_state_store`` immediately after
 submission. On retry, it reads the ID back and reconnects to the already-running driver instead of
 resubmitting.
 
@@ -212,7 +212,7 @@ The reconnection polling calls the Spark standalone REST API
 See :doc:`connections/spark-submit` for how to configure these fields.
 
 .. note::
-    Crash recovery in cluster mode requires Airflow 3.3+ (``task_state`` support). On earlier
+    Crash recovery in cluster mode requires Airflow 3.3+ (``task_state_store`` support). On earlier
     versions the operator falls back to the previous behavior of always submitting fresh.
 
 Tracking driver status via Kubernetes API

providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py

@@ -37,7 +37,7 @@
     # ResumableJobMixin does not exist in Airflow 2, so we need to add a stub to make it
     # behave as before
     class ResumableJobMixin:  # type: ignore[no-redef]
-        """Airflow 2 stub — no task_state, always submits fresh."""
+        """Airflow 2 stub — no task_state_store, always submits fresh."""
 
         external_id_key: str = "remote_job_id"
 
@@ -264,7 +264,7 @@ def execute(self, context: Context) -> None:
         if hook._should_track_driver_status:
             if self.reconnect_on_retry:
                 return self.execute_resumable(context)
-            # reconnect_on_retry=False: still submit-and-poll, just skip task_state persistence.
+            # reconnect_on_retry=False: still submit-and-poll, just skip task_state_store persistence.
             driver_id = self.submit_job(context)
             self.poll_until_complete(driver_id, context)
             return self.get_job_result(driver_id, context)
@@ -284,7 +284,7 @@ def execute(self, context: Context) -> None:
                 hook._validate_yarn_track_via_rm_api_config()
                 if self.reconnect_on_retry:
                     return self.execute_resumable(context)
-                # reconnect_on_retry=False: still submit-and-poll, just skip task_state persistence.
+                # reconnect_on_retry=False: still submit-and-poll, just skip task_state_store persistence.
                 driver_id = self.submit_job(context)
                 self.poll_until_complete(driver_id, context)
                 return self.get_job_result(driver_id, context)