From 4ff0653a5116bb2f4fb798060c6f48b5ffac6c28 Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Wed, 14 Oct 2026 16:23:15 +0800 Subject: [PATCH] Add asset state store API to the Java SDK Signed-off-by: PoAn Yang --- .../language-sdks/java.rst | 53 +++++ .../apache/airflow/e2e/TestBundleBuilder.java | 49 ++++- .../src/resources/dags/java_test_dags.py | 19 +- .../e2e_test_utils/clients.py | 8 + .../java_sdk_tests/test_java_sdk_dag.py | 66 +++++++ java-sdk/README.md | 2 +- java-sdk/capabilities.yaml | 5 +- java-sdk/sdk/module.md | 2 +- .../org/apache/airflow/sdk/AssetStateStore.kt | 116 +++++++++++ .../kotlin/org/apache/airflow/sdk/Client.kt | 16 ++ .../apache/airflow/sdk/execution/Client.kt | 118 +++++++++++ .../org/apache/airflow/sdk/ArgTestSupport.kt | 20 ++ .../org/apache/airflow/sdk/ClientTest.kt | 41 ++++ .../apache/airflow/sdk/execution/CommTest.kt | 185 ++++++++++++++++++ .../apache/airflow/sdk/execution/TaskTest.kt | 36 ++++ .../sdk/execution/TriggerRunnerTest.kt | 18 ++ .../airflow/sdk/internal/ArgValuesTest.kt | 20 ++ 17 files changed, 766 insertions(+), 8 deletions(-) create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/AssetStateStore.kt diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst index e8e63811c9617..6bb336fba15d9 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -1099,6 +1099,59 @@ structured data (for example to JSON) before storing it. takes precedence over the stored value when the Variable is read back. Calling ``setVariable`` without a description clears any existing description. +.. _java-sdk/asset-state-store: + +Asset state store +----------------- + +``client.getAssetStateStore()`` reads and writes the :doc:`asset state store `, +a key-value store that belongs to an asset rather than to a Dag run. A value that one run stores, such as an +incremental-load watermark, is still there in later runs. + +The Java SDK does not receive the inlets and outlets declared on the ``@task.stub``. The task therefore looks +up the asset itself, by name with ``byName`` or by URI with ``byUri``. Airflow keeps state only for assets +that a Dag references, so also declare the asset as an inlet of the stub: + +.. code-block:: python + + orders = Asset(name="orders", uri="s3://warehouse/orders") + + + @task.stub(queue="java", inlets=[orders]) + def load_orders(): ... + +.. code-block:: java + + @Airflow.TaskHandler(dag = "orders_pipeline", task = "load_orders") + public void loadOrders(Client client) { + var orders = client.getAssetStateStore().byName("orders"); + var watermark = (String) orders.get("watermark"); // null on the first run + // Your own code: load the rows created after the watermark and return the newest creation time. + String newest = loadRowsCreatedAfter(watermark); + orders.set("watermark", newest); + } + +``get`` returns ``null`` when the key is not set. ``set`` replaces any existing value. Entries do not expire. +``delete`` removes one key, and ``clear`` removes every key of the asset. + +``set`` sends the value as JSON. Store strings, booleans, integers that fit in a ``long``, finite ``double`` +values, lists, and maps with string keys. Jackson converts other objects first. A bean becomes a map of its +properties, and an ``Instant`` or a ``UUID`` becomes a string. ``null``, ``BigDecimal``, and ``Duration`` +fail with an exception before anything is sent. Never store a ``byte[]``. The supervisor cannot decode it and +never answers, so the task hangs. Encode binary data as a string instead, for example with Base64. + +``get`` returns the Java types listed in :ref:`java-sdk/types`. For example, a value stored as an ``Integer`` +comes back as a ``Long``. + +If no Dag references an asset with the name or URI that the task passed, ``get`` returns ``null`` and the +other methods throw ``ApiError``. Airflow normalizes asset URIs when it parses a Dag, so pass ``byUri`` the +normalized form. For example, ``s3://warehouse/orders/`` is stored as ``s3://warehouse/orders``. + +The Java SDK does not apply a ``[workers] state_store_backend``. It sends each value to Airflow as-is. A key +that a Python task stored through such a backend comes back to Java as a map that holds only the reference +under ``__airflow_state_ref__``, not the stored value. Deleting or clearing such a key from Java removes only +that reference, and the stored value stays in the backend. + .. _java-sdk/build: Building and packaging 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 39c4bd46afc30..02e8017bd557b 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 @@ -21,13 +21,14 @@ import java.util.List; import java.util.Map; +import java.util.Objects; 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 tasks that round-trip Airflow Variables and task state through the - * supervisor. + * instantiation failures, and tasks that round-trip Airflow Variables, task state and asset state + * through the supervisor. */ public class TestBundleBuilder implements BundleBuilder { public static class MissingNoArgConstructor implements Task { @@ -90,6 +91,46 @@ public void execute(@NotNull Context context, Client client) { } } + /** + * Calls every asset state store method on one asset, through a lookup by name and a lookup by + * URI. Each step changes the state through one lookup and checks the result through the other, + * so the task fails unless both lookups reach the same asset. The E2E test then reads the + * remaining keys through the REST API. + */ + public static class UseAssetStateStore implements Task { + public void execute(@NotNull Context context, Client client) { + var byName = client.getAssetStateStore().byName("java_e2e_orders"); + var byUri = client.getAssetStateStore().byUri("x-java-e2e://orders"); + + byName.set("scratch", "first"); + expect("first", byUri.get("scratch"), "set() through the name"); + byUri.clear(); + expect(null, byName.get("scratch"), "clear() through the URI"); + + byUri.set("scratch", "second"); + expect("second", byName.get("scratch"), "set() through the URI"); + byName.clear(); + expect(null, byUri.get("scratch"), "clear() through the name"); + + byName.set("deleted", "temporary"); + byUri.delete("deleted"); + expect(null, byName.get("deleted"), "delete() through the URI"); + byUri.set("deleted", "temporary"); + byName.delete("deleted"); + expect(null, byUri.get("deleted"), "delete() through the name"); + + var summary = Map.of("run_id", context.dagRun.runId, "rows", 3L); + byName.set("summary", summary); + expect(summary, byUri.get("summary"), "set() through the name, with a map value"); + } + + private static void expect(Object expected, Object actual, String step) { + if (!Objects.equals(expected, actual)) { + throw new IllegalStateException(step + ": expected " + expected + " but read " + actual); + } + } + } + @NotNull @Override public Iterable getDags() { @@ -100,7 +141,9 @@ public Iterable getDags() { variableWrite.addTask("write_and_delete", WriteAndDeleteVariable.class); var taskState = new DagDef("java_task_state"); taskState.addTask("roundtrip_task_state", RoundtripTaskState.class); - return List.of(uninstantiable, variableWrite, taskState); + var assetStateStore = new DagDef("java_asset_state_store"); + assetStateStore.addTask("use_asset_state_store", UseAssetStateStore.class); + return List.of(uninstantiable, variableWrite, taskState, assetStateStore); } 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 8a4818923f040..f78a8823b5a12 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 @@ -18,7 +18,7 @@ from __future__ import annotations -from airflow.sdk import dag, task +from airflow.sdk import Asset, dag, task @task.stub(queue="java-test") @@ -60,3 +60,20 @@ def java_task_state(): java_task_state() + + +# The Java task cannot see this inlet. Declaring it keeps the asset active, and the +# Execution API only reads and writes the state of active assets. +java_e2e_orders = Asset(name="java_e2e_orders", uri="x-java-e2e://orders") + + +@task.stub(queue="java-test", inlets=[java_e2e_orders]) +def use_asset_state_store(): ... + + +@dag(dag_id="java_asset_state_store") +def java_asset_state_store(): + use_asset_state_store() + + +java_asset_state_store() diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py b/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py index 821765ad21202..74e01961f9f89 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py @@ -230,6 +230,14 @@ def get_dag_run(self, dag_id: str, run_id: str): """Get a Dag run, with its state, run type and conf.""" return self._make_request(method="GET", endpoint=f"dags/{dag_id}/dagRuns/{run_id}") + def get_asset_state_store_entry(self, asset_name: str, key: str): + """Get one key of an asset's state store via API, looking up the asset by its name.""" + # name_pattern matches substrings, so keep only the asset named exactly asset_name. + assets = self._make_request(method="GET", endpoint="assets", params={"name_pattern": asset_name}) + asset_ids = [asset["id"] for asset in assets["assets"] if asset["name"] == asset_name] + assert len(asset_ids) == 1, f"expected one asset named {asset_name!r}, got {assets!r}" + return self._make_request(method="GET", endpoint=f"assets/{asset_ids[0]}/state-store/{key}") + def trigger_dag_and_wait(self, dag_id: str, json=None): """Trigger a DAG and wait for it to complete.""" self.un_pause_dag(dag_id) 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 ec2f1408f910a..44d65a138d017 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 @@ -98,6 +98,9 @@ _VARIABLE_WRITE_DESCRIPTION = "written by the Java SDK e2e test" _TASK_STATE_DAG_ID = "java_task_state" _TASK_STATE_TASK_ID = "roundtrip_task_state" +_ASSET_STATE_STORE_DAG_ID = "java_asset_state_store" +_ASSET_STATE_STORE_TASK_ID = "use_asset_state_store" +_ASSET_STATE_STORE_ASSET_NAME = "java_e2e_orders" @dataclass @@ -226,6 +229,12 @@ def task_state_run() -> _CompletedRun: return _trigger_and_wait_for_dag(_TASK_STATE_DAG_ID, _JAVA_TASK_TIMEOUT) +@pytest.fixture(scope="module") +def asset_state_store_run() -> _CompletedRun: + """Trigger the asset state store Dag once for all of its assertions.""" + return _trigger_and_wait_for_dag(_ASSET_STATE_STORE_DAG_ID, _JAVA_TASK_TIMEOUT) + + class TestJavaSDKAnnotationExample: """Verify the annotation-based Java SDK example executes correctly.""" @@ -427,6 +436,63 @@ def test_scratch_variable_deleted_by_java_task_is_gone(self, variable_write_run: ) +class TestJavaSDKAssetStateStore: + """Verify a Java task can read and write the state store of an asset. + + The task lives in the java-test-bundle fixture project (served on the + dedicated "java-test" queue). It looks up the ``java_e2e_orders`` asset both + by name and by URI and calls every asset state store method. Each step + changes the state through one lookup and checks it through the other, so the + task fails unless both lookups reach the same asset. It leaves the + ``summary`` key behind for the REST API to read. + """ + + def _assert_task_succeeded(self, run: _CompletedRun) -> None: + ti = run.get_task_instance(_ASSET_STATE_STORE_TASK_ID) + assert ti.get("state") == "success", ( + f"Java {_ASSET_STATE_STORE_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_value_written_by_java_task_is_readable(self, asset_state_store_run: _CompletedRun): + """The value set from Java is stored on the asset and attributed to the Java task.""" + self._assert_task_succeeded(asset_state_store_run) + + entry = asset_state_store_run.client.get_asset_state_store_entry( + _ASSET_STATE_STORE_ASSET_NAME, "summary" + ) + assert entry["value"] == {"run_id": asset_state_store_run.run_id, "rows": 3}, ( + f"summary should hold the map the Java task stored in this run, got {entry!r}" + ) + ti = asset_state_store_run.get_task_instance(_ASSET_STATE_STORE_TASK_ID) + expected_writer = { + "kind": "task", + "dag_id": _ASSET_STATE_STORE_DAG_ID, + "run_id": asset_state_store_run.run_id, + "task_id": _ASSET_STATE_STORE_TASK_ID, + "map_index": -1, + "task_instance_id": ti["id"], + "try_number": ti["try_number"], + } + writer = entry["last_updated_by"] + assert {key: writer.get(key) for key in expected_writer} == expected_writer, ( + f"summary should be attributed to the Java task, got {entry!r}" + ) + + def test_key_deleted_by_java_task_is_gone(self, asset_state_store_run: _CompletedRun): + """A key written and then deleted from Java no longer exists.""" + self._assert_task_succeeded(asset_state_store_run) + + with pytest.raises(requests.HTTPError) as excinfo: + asset_state_store_run.client.get_asset_state_store_entry(_ASSET_STATE_STORE_ASSET_NAME, "deleted") + assert excinfo.value.response.status_code == HTTPStatus.NOT_FOUND, ( + f"the 'deleted' key should have been deleted by the Java task, " + f"got HTTP {excinfo.value.response.status_code}" + ) + + class TestJavaSDKInterfaceExample: """Verify tasks written against ``InputTask`` receive their bound arguments.""" diff --git a/java-sdk/README.md b/java-sdk/README.md index 11f326e3a8e46..cdc59d2af54ce 100644 --- a/java-sdk/README.md +++ b/java-sdk/README.md @@ -635,7 +635,7 @@ prek hook regenerate it. | capability: `self-contained-bundle` | MUST | ✓ | 3.3 | Airflow metadata embedded in the jar artifact | | capability: `retry-policy` | MAY | ✗ | – | no task-facing retry-policy API yet | | capability: `task-state-store` | MAY | ✓ | 3.3 | Client.getTaskStateStore() get/set/delete/clear | -| capability: `asset-state-store` | MAY | ✗ | – | no task-facing state-store API yet | +| capability: `asset-state-store` | MAY | ✓ | 3.4 | Client.getAssetStateStore(), by asset name or URI | | capability: `asset-event-emit` | MAY | ✗ | – | runtime does not emit asset events yet | | capability: `asset-event-read` | MAY | ✗ | – | no task-facing asset-event API yet | | **Native-Dag authoring** | | | | | diff --git a/java-sdk/capabilities.yaml b/java-sdk/capabilities.yaml index cf9ec1c5fa758..b50921e699d57 100644 --- a/java-sdk/capabilities.yaml +++ b/java-sdk/capabilities.yaml @@ -89,8 +89,9 @@ capabilities: since: "3.3" note: "Client.getTaskStateStore() get/set/delete/clear" asset-state-store: - supported: false - note: "no task-facing state-store API yet" + supported: true + since: "3.4" + note: "Client.getAssetStateStore(), by asset name or URI" asset-event-emit: supported: false note: "runtime does not emit asset events yet" diff --git a/java-sdk/sdk/module.md b/java-sdk/sdk/module.md index 4e9f368e47512..b6536b4103b69 100644 --- a/java-sdk/sdk/module.md +++ b/java-sdk/sdk/module.md @@ -52,7 +52,7 @@ meaning of each dimension is defined in the | capability: `self-contained-bundle` | MUST | ✓ | 3.3 | Airflow metadata embedded in the jar artifact | | capability: `retry-policy` | MAY | ✗ | – | no task-facing retry-policy API yet | | capability: `task-state-store` | MAY | ✓ | 3.3 | Client.getTaskStateStore() get/set/delete/clear | -| capability: `asset-state-store` | MAY | ✗ | – | no task-facing state-store API yet | +| capability: `asset-state-store` | MAY | ✓ | 3.4 | Client.getAssetStateStore(), by asset name or URI | | capability: `asset-event-emit` | MAY | ✗ | – | runtime does not emit asset events yet | | capability: `asset-event-read` | MAY | ✗ | – | no task-facing asset-event API yet | | **Native-Dag authoring** | | | | | diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/AssetStateStore.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/AssetStateStore.kt new file mode 100644 index 0000000000000..be50e594f0120 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/AssetStateStore.kt @@ -0,0 +1,116 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk + +import org.apache.airflow.sdk.execution.AssetRef +import org.apache.airflow.sdk.execution.Client + +/** + * Looks up the state store of an asset by the asset's name or URI. + * + * The Java SDK does not receive the inlets and outlets declared on the + * `@task.stub`, so it cannot know which assets a task uses. The task code + * passes the name or URI of the asset it wants instead. + */ +class AssetStateStores internal constructor( + private val impl: Client, +) { + /** + * Returns the state store of the asset whose name is [name]. + * + * @throws IllegalArgumentException if [name] is empty. + */ + fun byName(name: String): AssetStateStore { + require(name.isNotEmpty()) { "Asset name must not be empty" } + return AssetStateStore(AssetRef.Name(name), impl) + } + + /** + * Returns the state store of the asset whose URI is [uri]. + * + * @throws IllegalArgumentException if [uri] is empty. + */ + fun byUri(uri: String): AssetStateStore { + require(uri.isNotEmpty()) { "Asset URI must not be empty" } + return AssetStateStore(AssetRef.Uri(uri), impl) + } +} + +/** + * Holds the key-value state of one asset. Every task and Dag run that uses the + * asset shares the same entries. + * + * Entries do not expire. They stay until they are deleted, or until no Dag + * references the asset anymore. + * + * If no active asset matches the name or URI used to look up this store, [get] + * returns `null` and the other methods throw [ApiError]. + * + * The SDK sends values to Airflow as-is and does not apply the `[workers] + * state_store_backend` that Python tasks use. + */ +class AssetStateStore internal constructor( + private val asset: AssetRef, + private val impl: Client, +) { + /** + * Reads the value stored under [key]. + * + * Integers come back as `Long`, decimals as `Double`, JSON objects as `Map`, + * and JSON arrays as `List`. So a value stored as an `Integer` comes back as + * a `Long`. + * + * @return The stored value, or `null` if the key is not set. + * @throws ApiError if the API call fails. + */ + fun get(key: String): Any? = impl.getAssetStateStore(asset, key)?.value + + /** + * Stores [value] under [key], replacing any existing value. + * + * The SDK sends the value as JSON. Do not store a `ByteArray`. The supervisor + * cannot decode it and never answers, so the task hangs. Encode binary data + * as a string instead, for example with Base64. + * + * @param key State key. + * @param value Value to store. Must be JSON-serializable. + * @throws IllegalArgumentException if the SDK cannot send the type of + * [value], for example a `BigDecimal`. + * @throws ApiError if the API call fails. + */ + fun set( + key: String, + value: Any, + ) = impl.setAssetStateStore(asset, key, value) + + /** + * Deletes the value stored under [key]. Does nothing if the key is not set. + * + * @throws ApiError if the API call fails. + */ + fun delete(key: String) = impl.deleteAssetStateStore(asset, key) + + /** + * Deletes every key stored for this asset. + * + * @throws ApiError if the API call fails. + */ + fun clear() = impl.clearAssetStateStore(asset) +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt index c394e270436d8..2ef806a57b4d6 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt @@ -79,6 +79,22 @@ class Client internal constructor( */ val taskStateStore: TaskStateStore = TaskStateStore(details, impl, env) + /** + * Gives the task the state store of each asset, which it looks up by the + * asset's name or URI. + * + * Entries belong to the asset, not to a task instance or a Dag run. A value + * that one Dag run stores, such as an incremental-load watermark, is still + * there in later runs: + * + * ```java + * var orders = client.getAssetStateStore().byName("orders"); + * var watermark = (String) orders.get("watermark"); // null until a task sets it + * orders.set("watermark", "2026-10-01T00:00:00Z"); + * ``` + */ + val assetStateStore: AssetStateStores = AssetStateStores(impl) + /** * Retrieves a connection from the Airflow connection store. * diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt index 317dac7ebb0a6..6252a8f044df9 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt @@ -20,13 +20,20 @@ package org.apache.airflow.sdk.execution import kotlinx.coroutines.runBlocking +import org.apache.airflow.sdk.execution.comm.AssetStateStoreResult +import org.apache.airflow.sdk.execution.comm.ClearAssetStateStoreByName +import org.apache.airflow.sdk.execution.comm.ClearAssetStateStoreByUri import org.apache.airflow.sdk.execution.comm.ClearTaskStateStore import org.apache.airflow.sdk.execution.comm.ConnectionResult import org.apache.airflow.sdk.execution.comm.DagResult import org.apache.airflow.sdk.execution.comm.DagRunStateResult +import org.apache.airflow.sdk.execution.comm.DeleteAssetStateStoreByName +import org.apache.airflow.sdk.execution.comm.DeleteAssetStateStoreByUri import org.apache.airflow.sdk.execution.comm.DeleteTaskStateStore import org.apache.airflow.sdk.execution.comm.DeleteVariable import org.apache.airflow.sdk.execution.comm.ErrorResponse +import org.apache.airflow.sdk.execution.comm.GetAssetStateStoreByName +import org.apache.airflow.sdk.execution.comm.GetAssetStateStoreByUri import org.apache.airflow.sdk.execution.comm.GetConnection import org.apache.airflow.sdk.execution.comm.GetDag import org.apache.airflow.sdk.execution.comm.GetDagRunState @@ -35,6 +42,8 @@ import org.apache.airflow.sdk.execution.comm.GetVariable import org.apache.airflow.sdk.execution.comm.GetXCom import org.apache.airflow.sdk.execution.comm.OKResponse import org.apache.airflow.sdk.execution.comm.PutVariable +import org.apache.airflow.sdk.execution.comm.SetAssetStateStoreByName +import org.apache.airflow.sdk.execution.comm.SetAssetStateStoreByUri import org.apache.airflow.sdk.execution.comm.SetTaskStateStore import org.apache.airflow.sdk.execution.comm.SetXCom import org.apache.airflow.sdk.execution.comm.SkipDownstreamTasks @@ -45,6 +54,22 @@ import org.apache.airflow.sdk.execution.comm.XComResult import java.time.OffsetDateTime import java.util.UUID +/** + * @suppress + * + * Names the asset whose state store a request reads or writes. The supervisor + * has a separate message for each way of naming an asset. + */ +sealed interface AssetRef { + data class Name( + val name: String, + ) : AssetRef + + data class Uri( + val uri: String, + ) : AssetRef +} + /** * @suppress * @@ -141,6 +166,25 @@ interface Client { /** Whether the Dag is paused, which a task may refuse to trigger. */ fun isDagPaused(dagId: String): Boolean + + /** Returns `null` when the key is not stored, or no active asset matches [asset]. */ + fun getAssetStateStore( + asset: AssetRef, + key: String, + ): AssetStateStoreResult? + + fun setAssetStateStore( + asset: AssetRef, + key: String, + value: Any, + ) + + fun deleteAssetStateStore( + asset: AssetRef, + key: String, + ) + + fun clearAssetStateStore(asset: AssetRef) } /** @@ -319,4 +363,78 @@ class CoordinatorClient( override fun isDagPaused(dagId: String): Boolean = runBlocking { exec.communicate(GetDag().also { it.dagId = dagId }) }.isPaused == true + + override fun getAssetStateStore( + asset: AssetRef, + key: String, + ): AssetStateStoreResult? { + val message = + when (asset) { + is AssetRef.Name -> + GetAssetStateStoreByName().also { + it.name = asset.name + it.key = key + } + is AssetRef.Uri -> + GetAssetStateStoreByUri().also { + it.uri = asset.uri + it.key = key + } + } + return runBlocking { + exec.communicateOrNullIf(message, ErrorResponse.ErrorType.ASSET_STORE_NOT_FOUND) + } + } + + override fun setAssetStateStore( + asset: AssetRef, + key: String, + value: Any, + ) { + val message = + when (asset) { + is AssetRef.Name -> + SetAssetStateStoreByName().also { + it.name = asset.name + it.key = key + it.value = value + } + is AssetRef.Uri -> + SetAssetStateStoreByUri().also { + it.uri = asset.uri + it.key = key + it.value = value + } + } + runBlocking { exec.communicate(message) } + } + + override fun deleteAssetStateStore( + asset: AssetRef, + key: String, + ) { + val message = + when (asset) { + is AssetRef.Name -> + DeleteAssetStateStoreByName().also { + it.name = asset.name + it.key = key + } + is AssetRef.Uri -> + DeleteAssetStateStoreByUri().also { + it.uri = asset.uri + it.key = key + } + } + runBlocking { exec.communicate(message) } + } + + override fun clearAssetStateStore(asset: AssetRef) { + val message = + when (asset) { + is AssetRef.Name -> ClearAssetStateStoreByName().also { it.name = asset.name } + is AssetRef.Uri -> ClearAssetStateStoreByUri().also { it.uri = asset.uri } + } + runBlocking { exec.communicate(message) } + } } diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt index 95edc43bbcfc1..9d1fb4f7373b0 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt @@ -19,6 +19,8 @@ package org.apache.airflow.sdk +import org.apache.airflow.sdk.execution.AssetRef +import org.apache.airflow.sdk.execution.comm.AssetStateStoreResult import org.apache.airflow.sdk.execution.comm.ConnectionResult import org.apache.airflow.sdk.execution.comm.StartupDetails import org.apache.airflow.sdk.execution.comm.TIRunContext @@ -109,6 +111,24 @@ internal class FakeXComTransport( ): String = throw NotImplementedError() override fun isDagPaused(dagId: String): Boolean = throw NotImplementedError() + + override fun getAssetStateStore( + asset: AssetRef, + key: String, + ): AssetStateStoreResult? = throw NotImplementedError() + + override fun setAssetStateStore( + asset: AssetRef, + key: String, + value: Any, + ) = throw NotImplementedError() + + override fun deleteAssetStateStore( + asset: AssetRef, + key: String, + ) = throw NotImplementedError() + + override fun clearAssetStateStore(asset: AssetRef) = throw NotImplementedError() } internal fun startupDetails(argBindings: List>?): StartupDetails = diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt index 239e2bbf95496..159cf11dfda7f 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt @@ -19,6 +19,8 @@ package org.apache.airflow.sdk +import org.apache.airflow.sdk.execution.AssetRef +import org.apache.airflow.sdk.execution.comm.AssetStateStoreResult import org.apache.airflow.sdk.execution.comm.ConnectionResult import org.apache.airflow.sdk.execution.comm.StartupDetails import org.apache.airflow.sdk.execution.comm.TaskInstance @@ -123,6 +125,24 @@ private class FakeTransport( ): String = throw NotImplementedError() override fun isDagPaused(dagId: String): Boolean = throw NotImplementedError() + + override fun getAssetStateStore( + asset: AssetRef, + key: String, + ): AssetStateStoreResult? = throw NotImplementedError() + + override fun setAssetStateStore( + asset: AssetRef, + key: String, + value: Any, + ) = throw NotImplementedError() + + override fun deleteAssetStateStore( + asset: AssetRef, + key: String, + ) = throw NotImplementedError() + + override fun clearAssetStateStore(asset: AssetRef) = throw NotImplementedError() } class ClientTest { @@ -299,6 +319,27 @@ class ClientTest { ) } + @Test + @DisplayName("assetStateStore is exposed to Java as a getter so mocking frameworks can stub it") + fun assetStateStoreIsExposedAsGetter() { + val getter = Client::class.java.getMethod("getAssetStateStore") + + Assertions.assertEquals(AssetStateStores::class.java, getter.returnType) + Assertions.assertTrue(Client::class.java.fields.none { it.name == "assetStateStore" }) + } + + @Test + @DisplayName("assetStateStore rejects an empty asset name or URI before anything is sent") + fun assetStateStoreRejectsEmptyNameOrUri() { + val client = clientWith(ConnectionResult()) + + val byName = Assertions.assertThrows(IllegalArgumentException::class.java) { client.assetStateStore.byName("") } + val byUri = Assertions.assertThrows(IllegalArgumentException::class.java) { client.assetStateStore.byUri("") } + + Assertions.assertEquals("Asset name must not be empty", byName.message) + Assertions.assertEquals("Asset URI must not be empty", byUri.message) + } + @Test @DisplayName("MissingXComException builds the full message naming the task and parameter") fun missingXComExceptionBuildsFullMessage() { diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt index 20fba213948f5..4a5371a6c1d29 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt @@ -777,4 +777,189 @@ class CommsTest { ) comm.close() } + + private fun assetStateStoreResultFrame(id: Int): ByteArray { + val out = ByteArrayOutputStream() + MessagePack.newDefaultPacker(out).use { packer -> + packer.packArrayHeader(3) + packer.packInt(id) + packer.packMapHeader(2) + packer.packString("type") + packer.packString("AssetStateStoreResult") + packer.packString("value") + packer.packString("2026-10-01T00:00:00Z") + packer.packNil() + } + return out.toByteArray() + } + + // The supervisor's handler returns this ErrorResponse as its result, so it arrives in the body slot. + private fun assetStoreNotFoundFrame(id: Int): ByteArray { + val out = ByteArrayOutputStream() + MessagePack.newDefaultPacker(out).use { packer -> + packer.packArrayHeader(3) + packer.packInt(id) + packer.packMapHeader(3) + packer.packString("type") + packer.packString("ErrorResponse") + packer.packString("error") + packer.packString("ASSET_STORE_NOT_FOUND") + packer.packString("detail") + packer.packMapHeader(1) + packer.packString("key") + packer.packString("watermark") + packer.packNil() + } + return out.toByteArray() + } + + // The supervisor reports a failed Execution API call in the error slot. + private fun apiServerErrorFrame( + id: Int, + statusCode: Int, + ): ByteArray { + val out = ByteArrayOutputStream() + MessagePack.newDefaultPacker(out).use { packer -> + packer.packArrayHeader(3) + packer.packInt(id) + packer.packNil() + packer.packMapHeader(3) + packer.packString("type") + packer.packString("ErrorResponse") + packer.packString("error") + packer.packString("API_SERVER_ERROR") + packer.packString("detail") + packer.packMapHeader(1) + packer.packString("status_code") + packer.packInt(statusCode) + } + return out.toByteArray() + } + + @Test + @DisplayName("assetStateStore.get sends the asset name or URI with the key and unwraps the stored value") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun assetStateStoreGetUnwrapsResult() { + var byName: Any? = null + val (byNameBody, byNameFailure) = + roundTrip(::assetStateStoreResultFrame) { byName = it.assetStateStore.byName("orders").get("watermark") } + var byUri: Any? = null + val (byUriBody, byUriFailure) = + roundTrip(::assetStateStoreResultFrame) { byUri = it.assetStateStore.byUri("s3://bucket/orders").get("watermark") } + + Assertions.assertNull(byNameFailure, "get should return normally on AssetStateStoreResult, got $byNameFailure") + Assertions.assertEquals( + mapOf("type" to "GetAssetStateStoreByName", "name" to "orders", "key" to "watermark"), + byNameBody, + ) + Assertions.assertEquals("2026-10-01T00:00:00Z", byName) + + Assertions.assertNull(byUriFailure, "get should return normally on AssetStateStoreResult, got $byUriFailure") + Assertions.assertEquals( + mapOf("type" to "GetAssetStateStoreByUri", "uri" to "s3://bucket/orders", "key" to "watermark"), + byUriBody, + ) + Assertions.assertEquals("2026-10-01T00:00:00Z", byUri) + } + + @Test + @DisplayName("assetStateStore.get returns null on ASSET_STORE_NOT_FOUND instead of raising") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun assetStateStoreGetReturnsNullWhenNotFound() { + var value: Any? = "unset" + val (_, failure) = roundTrip(::assetStoreNotFoundFrame) { value = it.assetStateStore.byName("orders").get("watermark") } + + Assertions.assertNull(failure, "get should return normally on ASSET_STORE_NOT_FOUND, got $failure") + Assertions.assertNull(value) + } + + @Test + @DisplayName("assetStateStore.get raises ApiError instead of returning null when the API server fails") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun assetStateStoreGetRaisesApiErrorOnOtherErrors() { + val (_, failure) = + roundTrip({ apiServerErrorFrame(it, 503) }) { it.assetStateStore.byName("orders").get("watermark") } + + Assertions.assertInstanceOf(ApiError::class.java, failure) + Assertions.assertTrue(failure!!.message!!.startsWith("[API_SERVER_ERROR]"), "unexpected message: ${failure.message}") + } + + @Test + @DisplayName("assetStateStore.set sends the asset name or URI, key and value and accepts the supervisor's OK response") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun assetStateStoreSetAcceptsOkResponse() { + val (byNameBody, byNameFailure) = + roundTrip(::okResponseFrame) { it.assetStateStore.byName("orders").set("watermark", "2026-10-01T00:00:00Z") } + val (byUriBody, byUriFailure) = + roundTrip(::okResponseFrame) { + it.assetStateStore.byUri("s3://bucket/orders").set("summary", mapOf("rows" to 3, "region" to "emea")) + } + + Assertions.assertNull(byNameFailure, "set should return normally on OKResponse, got $byNameFailure") + Assertions.assertEquals( + mapOf( + "type" to "SetAssetStateStoreByName", + "name" to "orders", + "key" to "watermark", + "value" to "2026-10-01T00:00:00Z", + ), + byNameBody, + ) + + Assertions.assertNull(byUriFailure, "set should return normally on OKResponse, got $byUriFailure") + Assertions.assertEquals( + mapOf( + "type" to "SetAssetStateStoreByUri", + "uri" to "s3://bucket/orders", + "key" to "summary", + // The msgpack decoder yields Long for wire integers. + "value" to mapOf("rows" to 3L, "region" to "emea"), + ), + byUriBody, + ) + } + + @Test + @DisplayName("assetStateStore.delete and clear send the asset name or URI and accept the supervisor's OK response") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun assetStateStoreDeleteAndClearAcceptOkResponse() { + val sent = + listOf<(PublicClient) -> Unit>( + { it.assetStateStore.byName("orders").delete("watermark") }, + { it.assetStateStore.byUri("s3://bucket/orders").delete("watermark") }, + { it.assetStateStore.byName("orders").clear() }, + { it.assetStateStore.byUri("s3://bucket/orders").clear() }, + ).map { call -> + val (body, failure) = roundTrip(::okResponseFrame, call = call) + Assertions.assertNull(failure, "the call should return normally on OKResponse, got $failure") + body + } + + Assertions.assertEquals( + listOf( + mapOf("type" to "DeleteAssetStateStoreByName", "name" to "orders", "key" to "watermark"), + mapOf("type" to "DeleteAssetStateStoreByUri", "uri" to "s3://bucket/orders", "key" to "watermark"), + mapOf("type" to "ClearAssetStateStoreByName", "name" to "orders"), + mapOf("type" to "ClearAssetStateStoreByUri", "uri" to "s3://bucket/orders"), + ), + sent, + ) + } + + @Test + @DisplayName("assetStateStore.set, delete and clear raise ApiError when the supervisor reports an error") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun assetStateStoreWritesRaiseApiErrorOnErrorResponse() { + val calls = + listOf<(PublicClient) -> Unit>( + { it.assetStateStore.byName("orders").set("watermark", "2026-10-01T00:00:00Z") }, + { it.assetStateStore.byName("orders").delete("watermark") }, + { it.assetStateStore.byName("orders").clear() }, + ) + + for (call in calls) { + val (_, failure) = roundTrip({ apiServerErrorFrame(it, 404) }, call = call) + Assertions.assertInstanceOf(ApiError::class.java, failure) + } + } } diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt index a18907a4e25e8..298caa014da5e 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt @@ -450,6 +450,24 @@ class TaskTest { ): String = throw UnsupportedOperationException("not used in test") override fun isDagPaused(dagId: String): Boolean = throw UnsupportedOperationException("not used in test") + + override fun getAssetStateStore( + asset: AssetRef, + key: String, + ) = throw UnsupportedOperationException("not used in test") + + override fun setAssetStateStore( + asset: AssetRef, + key: String, + value: Any, + ): Unit = throw UnsupportedOperationException("not used in test") + + override fun deleteAssetStateStore( + asset: AssetRef, + key: String, + ): Unit = throw UnsupportedOperationException("not used in test") + + override fun clearAssetStateStore(asset: AssetRef): Unit = throw UnsupportedOperationException("not used in test") }, ) @@ -591,6 +609,24 @@ class TaskTest { ): String = throw UnsupportedOperationException("not used in test") override fun isDagPaused(dagId: String): Boolean = throw UnsupportedOperationException("not used in test") + + override fun getAssetStateStore( + asset: AssetRef, + key: String, + ) = throw UnsupportedOperationException("not used in test") + + override fun setAssetStateStore( + asset: AssetRef, + key: String, + value: Any, + ): Unit = throw UnsupportedOperationException("not used in test") + + override fun deleteAssetStateStore( + asset: AssetRef, + key: String, + ): Unit = throw UnsupportedOperationException("not used in test") + + override fun clearAssetStateStore(asset: AssetRef): Unit = throw UnsupportedOperationException("not used in test") } class SuccessTask : Task { diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TriggerRunnerTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TriggerRunnerTest.kt index ad0b629fa2a24..b782fb51cbbda 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TriggerRunnerTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TriggerRunnerTest.kt @@ -136,6 +136,24 @@ private class FakeTrigger( override fun clearTaskStateStore(tiId: UUID): Unit = throw UnsupportedOperationException("not used in test") override fun skipDownstreamTasks(taskIds: List): Unit = throw UnsupportedOperationException("not used in test") + + override fun getAssetStateStore( + asset: AssetRef, + key: String, + ) = throw UnsupportedOperationException("not used in test") + + override fun setAssetStateStore( + asset: AssetRef, + key: String, + value: Any, + ): Unit = throw UnsupportedOperationException("not used in test") + + override fun deleteAssetStateStore( + asset: AssetRef, + key: String, + ): Unit = throw UnsupportedOperationException("not used in test") + + override fun clearAssetStateStore(asset: AssetRef): Unit = throw UnsupportedOperationException("not used in test") } /** The trigger time the tests run at, so a generated run ID is the same on every run. */ diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt index bcb254547be33..3d5877c387ca9 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt @@ -32,6 +32,8 @@ import org.apache.airflow.sdk.Task import org.apache.airflow.sdk.TaskDef import org.apache.airflow.sdk.TaskInstance import org.apache.airflow.sdk.TaskRef +import org.apache.airflow.sdk.execution.AssetRef +import org.apache.airflow.sdk.execution.comm.AssetStateStoreResult import org.apache.airflow.sdk.execution.comm.ConnectionResult import org.apache.airflow.sdk.execution.comm.StartupDetails import org.apache.airflow.sdk.execution.comm.TaskStateStoreResult @@ -147,6 +149,24 @@ internal class ArgValuesTest { ): String = throw NotImplementedError() override fun isDagPaused(dagId: String): Boolean = throw NotImplementedError() + + override fun getAssetStateStore( + asset: AssetRef, + key: String, + ): AssetStateStoreResult? = throw NotImplementedError() + + override fun setAssetStateStore( + asset: AssetRef, + key: String, + value: Any, + ): Unit = throw NotImplementedError() + + override fun deleteAssetStateStore( + asset: AssetRef, + key: String, + ): Unit = throw NotImplementedError() + + override fun clearAssetStateStore(asset: AssetRef): Unit = throw NotImplementedError() }, )