This is an automated email from the ASF dual-hosted git repository.
jason810496 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 fd4969bd3c7 Add an e2e test for the Java SDK task state store (#74403)
fd4969bd3c7 is described below
commit fd4969bd3c773906e0f084c36087812f4e7685bd
Author: Andrew Chang <[email protected]>
AuthorDate: Tue Oct 13 17:04:50 2026 +0800
Add an e2e test for the Java SDK task state store (#74403)
The Java SDK task state store had unit tests only. Nothing checked
that a real Java task reaches the database through the supervisor, or
that the coordinator passes [state_store] default_retention_days to
the JVM. The Go SDK already has an e2e test for both.
---
.../org/apache/airflow/e2e/TestBundleBuilder.java | 36 ++++++++++-
.../src/resources/dags/java_test_dags.py | 12 ++++
.../tests/airflow_e2e_tests/conftest.py | 17 +++--
.../tests/airflow_e2e_tests/constants.py | 6 +-
.../go_sdk_tests/test_go_sdk_task_state.py | 6 +-
.../java_sdk_tests/test_java_sdk_dag.py | 72 +++++++++++++++++++++-
6 files changed, 134 insertions(+), 15 deletions(-)
diff --git
a/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java
b/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java
index ad33a15708a..39c4bd46afc 100644
---
a/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java
+++
b/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java
@@ -20,12 +20,14 @@
package org.apache.airflow.e2e;
import java.util.List;
+import java.util.Map;
import org.apache.airflow.sdk.*;
import org.jetbrains.annotations.NotNull;
/**
* Bundle for the runner-behaviour E2E tests: deliberately broken task classes
that exercise
- * instantiation failures, and a task that round-trips Airflow Variables
through the supervisor.
+ * instantiation failures, and tasks that round-trip Airflow Variables and
task state through the
+ * supervisor.
*/
public class TestBundleBuilder implements BundleBuilder {
public static class MissingNoArgConstructor implements Task {
@@ -60,6 +62,34 @@ public class TestBundleBuilder implements BundleBuilder {
}
}
+ /**
+ * Writes task state keys the E2E test reads back through the REST API, and
reads them back
+ * here to exercise the get path. No clear: the E2E test reads the keys
after the task ends.
+ */
+ public static class RoundtripTaskState implements Task {
+ public void execute(@NotNull Context context, Client client) {
+ var store = client.getTaskStateStore();
+ var runId = context.dagRun.runId;
+ store.set("java_e2e_run_id", runId, TaskStateStore.NEVER_EXPIRE);
+ store.set("java_e2e_counter", Map.of("processed", 3, "cursor",
"abc-123"));
+ store.set("java_e2e_retained", "retained");
+ store.set("java_e2e_scratch", "scratch");
+ store.delete("java_e2e_scratch");
+
+ var readBack = store.get("java_e2e_run_id");
+ if (!runId.equals(readBack)) {
+ throw new IllegalStateException("java_e2e_run_id: got " + readBack +
", want " + runId);
+ }
+ var counter = store.get("java_e2e_counter");
+ if (!Map.of("processed", 3L, "cursor", "abc-123").equals(counter)) {
+ throw new IllegalStateException("java_e2e_counter: got " + counter);
+ }
+ if (store.get("java_e2e_scratch") != null) {
+ throw new IllegalStateException("java_e2e_scratch survived its
delete");
+ }
+ }
+ }
+
@NotNull
@Override
public Iterable<DagDef> getDags() {
@@ -68,7 +98,9 @@ public class TestBundleBuilder implements BundleBuilder {
uninstantiable.addTask("non_static_inner", NonStaticInner.class);
var variableWrite = new DagDef("java_variable_write");
variableWrite.addTask("write_and_delete", WriteAndDeleteVariable.class);
- return List.of(uninstantiable, variableWrite);
+ var taskState = new DagDef("java_task_state");
+ taskState.addTask("roundtrip_task_state", RoundtripTaskState.class);
+ return List.of(uninstantiable, variableWrite, taskState);
}
public static void main(String[] args) {
diff --git
a/airflow-e2e-tests/java-test-bundle/src/resources/dags/java_test_dags.py
b/airflow-e2e-tests/java-test-bundle/src/resources/dags/java_test_dags.py
index d4ee554d66e..8a4818923f0 100644
--- a/airflow-e2e-tests/java-test-bundle/src/resources/dags/java_test_dags.py
+++ b/airflow-e2e-tests/java-test-bundle/src/resources/dags/java_test_dags.py
@@ -48,3 +48,15 @@ def java_variable_write():
java_variable_write()
+
+
[email protected](queue="java-test")
+def roundtrip_task_state(): ...
+
+
+@dag(dag_id="java_task_state")
+def java_task_state():
+ roundtrip_task_state()
+
+
+java_task_state()
diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py
b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py
index beb1376e7ea..38ea7bd8141 100644
--- a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py
+++ b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py
@@ -47,7 +47,6 @@ from airflow_e2e_tests.constants import (
GO_SDK_NATIVE_BUNDLE_NAME,
GO_SDK_NATIVE_BUNDLE_PKG,
GO_SDK_ROOT_PATH,
- GO_SDK_STATE_STORE_RETENTION_DAYS,
JAVA_COMPOSE_PATH,
JAVA_DOCKERFILE_PATH,
JAVA_NATIVE_BUNDLE_LIBS_PATH,
@@ -61,6 +60,7 @@ from airflow_e2e_tests.constants import (
JAVA_TEST_BUNDLE_ROOT_PATH,
KAFKA_DIR_PATH,
LANG_SDK_NATIVE_TOOLCHAIN,
+ LANG_SDK_STATE_STORE_RETENTION_DAYS,
LOCALSTACK_PATH,
LOGS_FOLDER,
NODE_IMAGE,
@@ -314,6 +314,14 @@ def _build_dag_bundle_config(artifact_bundles: dict[str,
str]) -> str:
return json.dumps(bundles)
+def _write_state_store_retention_config(tmp_dir):
+ # Config file, not an env var: the supervisor copies the worker
environment into the
+ # language SDK subprocess, so an env var would reach it even without the
propagation under test.
+ (tmp_dir / "config" / "airflow.cfg").write_text(
+ f"[state_store]\ndefault_retention_days =
{LANG_SDK_STATE_STORE_RETENTION_DAYS}\n"
+ )
+
+
def _run_java_sdk_gradle(workdir, *gradle_argv, capture_output=False,
native=False):
"""Run the Java SDK Gradle wrapper natively or inside the pinned JDK
container.
@@ -556,6 +564,7 @@ def _setup_java_sdk_integration(dot_env_file, tmp_dir):
"AIRFLOW_VAR_MY_VARIABLE=test_value\n"
)
os.environ["ENV_FILE_PATH"] = str(dot_env_file)
+ _write_state_store_retention_config(tmp_dir)
def _run_go_sdk_pack(package, output_path, *, capture_output=False,
native=False):
@@ -713,11 +722,7 @@ def _setup_go_sdk_integration(dot_env_file, tmp_dir):
)
os.environ["ENV_FILE_PATH"] = str(dot_env_file)
- # Config file, not an env var: the supervisor copies the worker
environment into the Go
- # subprocess, so an env var would reach Go even without the propagation
under test.
- (tmp_dir / "config" / "airflow.cfg").write_text(
- f"[state_store]\ndefault_retention_days =
{GO_SDK_STATE_STORE_RETENTION_DAYS}\n"
- )
+ _write_state_store_retention_config(tmp_dir)
def _setup_openlineage_integration(dot_env_file, tmp_dir, compose_file_names):
diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py
b/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py
index 024a90fb01b..7fd1323cc3f 100644
--- a/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py
+++ b/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py
@@ -62,6 +62,9 @@ OPENLINEAGE_COMPAT_DB_COMPOSE_PATH = (
# CI sets this (the same switch the lang-SDK k8s job uses) to build the
lang-SDK
# artifacts with the host toolchain instead of ephemeral toolchain containers.
LANG_SDK_NATIVE_TOOLCHAIN = os.environ.get("LANG_SDK_NATIVE_TOOLCHAIN",
"").lower() == "true"
+# Far from the supervisor's own 30-day config fallback, so the e2e test can
tell a
+# propagated value from it.
+LANG_SDK_STATE_STORE_RETENTION_DAYS = 7
# Java SDK E2E test paths
JAVA_SDK_ROOT_PATH = AIRFLOW_ROOT_PATH / "java-sdk"
@@ -103,9 +106,6 @@ GO_SDK_NATIVE_BUNDLE_NAME = "go_native_dags"
# Where airflow-go-pack writes the packed bundle inside the repo (go-sdk/bin
is gitignored).
GO_SDK_BIN_PATH = GO_SDK_ROOT_PATH / "bin"
GO_COMPOSE_PATH = AIRFLOW_ROOT_PATH / "airflow-e2e-tests" / "docker" / "go.yml"
-# Far from the supervisor's own 30-day config fallback, so the e2e test can
tell a
-# propagated value from it.
-GO_SDK_STATE_STORE_RETENTION_DAYS = 7
# Go toolchain image used to build the bundle in the containerized path (i.e.
unless
# LANG_SDK_NATIVE_TOOLCHAIN is set); must satisfy go-sdk/go.mod's toolchain.
# The Alpine variant is ~7x smaller than the Debian one and is safe here
because the
diff --git
a/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_task_state.py
b/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_task_state.py
index 3a8f2a81afb..d0b63455016 100644
---
a/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_task_state.py
+++
b/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_task_state.py
@@ -29,7 +29,7 @@ from http import HTTPStatus
import pytest
import requests
-from airflow_e2e_tests.constants import GO_SDK_STATE_STORE_RETENTION_DAYS
+from airflow_e2e_tests.constants import LANG_SDK_STATE_STORE_RETENTION_DAYS
from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
_GO_TASK_TIMEOUT = 300
@@ -106,9 +106,9 @@ def test_default_retention_applied(completed_run:
_CompletedRun):
entry = completed_run.state_store("go_e2e_retained")
assert entry.get("expires_at") is not None, entry
gap = _parse_timestamp(entry["expires_at"]) -
_parse_timestamp(entry["updated_at"])
- expected = timedelta(days=GO_SDK_STATE_STORE_RETENTION_DAYS)
+ expected = timedelta(days=LANG_SDK_STATE_STORE_RETENTION_DAYS)
assert abs(gap - expected) <= timedelta(minutes=5), (
- f"expected ~{GO_SDK_STATE_STORE_RETENTION_DAYS} days, got {gap /
timedelta(days=1):.1f} days; "
+ f"expected ~{LANG_SDK_STATE_STORE_RETENTION_DAYS} days, got {gap /
timedelta(days=1):.1f} days; "
"if ~30, the supervisor fell back to its own config default, so it did
not see "
f"[state_store] default_retention_days. entry: {entry!r}"
)
diff --git
a/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py
b/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py
index fcd8d8d01fa..ec2f1408f91 100644
---
a/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py
+++
b/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py
@@ -69,13 +69,14 @@ from __future__ import annotations
import time
from dataclasses import dataclass, field
-from datetime import UTC, datetime
+from datetime import UTC, datetime, timedelta
from http import HTTPStatus
from typing import TYPE_CHECKING
import pytest
import requests
+from airflow_e2e_tests.constants import LANG_SDK_STATE_STORE_RETENTION_DAYS
from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
if TYPE_CHECKING:
@@ -95,6 +96,8 @@ _XCOM_CASTING_DAG_ID = "java_xcom_casting_example"
_SCALA_SPARK_DAG_ID = "scala_spark_example"
_VARIABLE_WRITE_DAG_ID = "java_variable_write"
_VARIABLE_WRITE_DESCRIPTION = "written by the Java SDK e2e test"
+_TASK_STATE_DAG_ID = "java_task_state"
+_TASK_STATE_TASK_ID = "roundtrip_task_state"
@dataclass
@@ -128,6 +131,14 @@ class _CompletedRun:
key=key,
).get("value")
+ def get_task_state_store(self, task_id: str, key: str) -> dict:
+ self._raise_if_failed()
+ if self.run_id is None:
+ raise RuntimeError(f"Dag {self.dag_id!r} did not produce a run ID")
+ return self.client.get_task_state_store(
+ dag_id=self.dag_id, run_id=self.run_id, task_id=task_id, key=key
+ )
+
def wait_for_log_record(
self, task_id: str, try_number: int, match: Callable[[dict], bool]
) -> tuple[dict | None, list[dict]]:
@@ -175,6 +186,10 @@ def _trigger_and_wait_for_dag(dag_id: str, timeout: int)
-> _CompletedRun:
)
+def _parse_timestamp(value: str) -> datetime:
+ return datetime.fromisoformat(value.replace("Z", "+00:00"))
+
+
@pytest.fixture(scope="module")
def annotation_example_run() -> _CompletedRun:
"""Trigger the annotation example once for all of its assertions."""
@@ -205,6 +220,12 @@ def variable_write_run() -> _CompletedRun:
return _trigger_and_wait_for_dag(_VARIABLE_WRITE_DAG_ID,
_JAVA_TASK_TIMEOUT)
[email protected](scope="module")
+def task_state_run() -> _CompletedRun:
+ """Trigger the task state Dag once for all of its assertions."""
+ return _trigger_and_wait_for_dag(_TASK_STATE_DAG_ID, _JAVA_TASK_TIMEOUT)
+
+
class TestJavaSDKAnnotationExample:
"""Verify the annotation-based Java SDK example executes correctly."""
@@ -480,3 +501,52 @@ class TestJavaSDKScalaSparkExample:
assert load_xcom == _SPARK_EXPECTED_TOTAL_REVENUE, (
f"Expected spark_load to return total revenue
{_SPARK_EXPECTED_TOTAL_REVENUE}, got {load_xcom!r}"
)
+
+
+class TestJavaSDKTaskStateStore:
+ """Verify a Java task can set, get and delete task state store keys.
+
+ The task lives in the java-test-bundle fixture project (served on the
+ dedicated "java-test" queue). It also reads its keys back and fails if a
+ value does not match, so ``success`` covers the get path. Assertions here
+ read the store through the REST API, so the data is proven to have reached
+ the database.
+ """
+
+ def _assert_task_succeeded(self, run: _CompletedRun) -> None:
+ ti = run.get_task_instance(_TASK_STATE_TASK_ID)
+ assert ti.get("state") == "success", (
+ f"Java {_TASK_STATE_TASK_ID!r} task did not succeed.\n"
+ f" task state : {ti.get('state')!r}\n"
+ f" dag state : {run.state!r}\n"
+ f" all tasks : {run.ti_states}"
+ )
+
+ def test_never_expire_key_holds_run_id(self, task_state_run:
_CompletedRun):
+ self._assert_task_succeeded(task_state_run)
+ entry = task_state_run.get_task_state_store(_TASK_STATE_TASK_ID,
"java_e2e_run_id")
+ assert entry.get("value") == task_state_run.run_id, entry
+ assert entry.get("expires_at") is None, entry
+
+ def test_structured_value_roundtrips(self, task_state_run: _CompletedRun):
+ self._assert_task_succeeded(task_state_run)
+ entry = task_state_run.get_task_state_store(_TASK_STATE_TASK_ID,
"java_e2e_counter")
+ assert entry.get("value") == {"processed": 3, "cursor": "abc-123"},
entry
+
+ def test_deleted_key_is_gone(self, task_state_run: _CompletedRun):
+ self._assert_task_succeeded(task_state_run)
+ with pytest.raises(requests.HTTPError) as excinfo:
+ task_state_run.get_task_state_store(_TASK_STATE_TASK_ID,
"java_e2e_scratch")
+ assert excinfo.value.response.status_code == HTTPStatus.NOT_FOUND
+
+ def test_default_retention_applied(self, task_state_run: _CompletedRun):
+ self._assert_task_succeeded(task_state_run)
+ entry = task_state_run.get_task_state_store(_TASK_STATE_TASK_ID,
"java_e2e_retained")
+ assert entry.get("expires_at") is not None, entry
+ gap = _parse_timestamp(entry["expires_at"]) -
_parse_timestamp(entry["updated_at"])
+ expected = timedelta(days=LANG_SDK_STATE_STORE_RETENTION_DAYS)
+ assert abs(gap - expected) <= timedelta(minutes=5), (
+ f"expected ~{LANG_SDK_STATE_STORE_RETENTION_DAYS} days, got {gap /
timedelta(days=1):.1f} days; "
+ "if ~30, the coordinator passed Airflow's config default, so the
JVM did not see "
+ f"[state_store] default_retention_days. entry: {entry!r}"
+ )
