Add asset state store API to the Java SDK

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 </core-concepts/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

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<DagDef> getDags() {
@@ -100,7 +141,9 @@ public Iterable<DagDef> 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) {

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()

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)

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."""
 

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** |  |  |  |  |

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"

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** |  |  |  |  |