Skip to content

Commit 52afaf0

Browse files
committed
Add asset state store API to the Java SDK
Signed-off-by: PoAn Yang <payang@apache.org>
1 parent f74a685 commit 52afaf0

17 files changed

Lines changed: 732 additions & 7 deletions

File tree

‎airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -816,6 +816,59 @@ structured data (for example to JSON) before storing it.
816816
takes precedence over the stored value when the Variable is read back. Calling ``setVariable``
817817
without a description clears any existing description.
818818

819+
.. _java-sdk/asset-state-store:
820+
821+
Asset state store
822+
-----------------
823+
824+
``client.assetStateStore`` reads and writes the :doc:`asset state store </core-concepts/asset-state-store>`,
825+
a key-value store that belongs to an asset rather than to a Dag run. A value that one run stores, such as an
826+
incremental-load watermark, is still there in later runs.
827+
828+
The Java SDK does not receive the inlets and outlets declared on the ``@task.stub``. The task therefore looks
829+
up the asset itself, by name with ``byName`` or by URI with ``byUri``. Airflow keeps state only for assets
830+
that a Dag references, so also declare the asset as an inlet of the stub:
831+
832+
.. code-block:: python
833+
834+
orders = Asset(name="orders", uri="s3://warehouse/orders")
835+
836+
837+
@task.stub(queue="java", inlets=[orders])
838+
def load_orders(): ...
839+
840+
.. code-block:: java
841+
842+
@Builder.TaskHandler(dag = "orders_pipeline", task = "load_orders")
843+
public void loadOrders(Client client) {
844+
var orders = client.assetStateStore.byName("orders");
845+
var watermark = (String) orders.get("watermark"); // null on the first run
846+
// Your own code: load the rows created after the watermark and return the newest creation time.
847+
String newest = loadRowsCreatedAfter(watermark);
848+
orders.set("watermark", newest);
849+
}
850+
851+
``get`` returns ``null`` when the key is not set. ``set`` replaces any existing value. Entries do not expire.
852+
``delete`` removes one key, and ``clear`` removes every key of the asset.
853+
854+
``set`` sends the value as JSON. Store strings, booleans, integers that fit in a ``long``, finite ``double``
855+
values, lists, and maps with string keys. Jackson converts other objects first. A bean becomes a map of its
856+
properties, and an ``Instant`` or a ``UUID`` becomes a string. ``null``, ``BigDecimal``, and ``Duration``
857+
fail with an exception before anything is sent. Never store a ``byte[]``. The supervisor cannot decode it and
858+
never answers, so the task hangs. Encode binary data as a string instead, for example with Base64.
859+
860+
``get`` returns the Java types listed in :ref:`java-sdk/types`. For example, a value stored as an ``Integer``
861+
comes back as a ``Long``.
862+
863+
If no Dag references an asset with the name or URI that the task passed, ``get`` returns ``null`` and the
864+
other methods throw ``ApiError``. Airflow normalizes asset URIs when it parses a Dag, so pass ``byUri`` the
865+
normalized form. For example, ``s3://warehouse/orders/`` is stored as ``s3://warehouse/orders``.
866+
867+
The Java SDK does not apply a ``[workers] state_store_backend``. It sends each value to Airflow as-is. A key
868+
that a Python task stored through such a backend comes back to Java as a map that holds only the reference
869+
under ``__airflow_state_ref__``, not the stored value. Deleting or clearing such a key from Java removes only
870+
that reference, and the stored value stays in the backend.
871+
819872
.. _java-sdk/build:
820873

821874
Building and packaging

‎airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java‎

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,15 @@
2020
package org.apache.airflow.e2e;
2121

2222
import java.util.List;
23+
import java.util.Map;
24+
import java.util.Objects;
2325
import org.apache.airflow.sdk.*;
2426
import org.jetbrains.annotations.NotNull;
2527

2628
/**
2729
* Bundle for the runner-behaviour E2E tests: deliberately broken task classes that exercise
28-
* instantiation failures, and a task that round-trips Airflow Variables through the supervisor.
30+
* instantiation failures, and tasks that round-trip Airflow Variables and asset state through the
31+
* supervisor.
2932
*/
3033
public class TestBundleBuilder implements BundleBuilder {
3134
public static class MissingNoArgConstructor implements Task {
@@ -60,6 +63,46 @@ public void execute(@NotNull Context context, Client client) {
6063
}
6164
}
6265

66+
/**
67+
* Calls every asset state store method on one asset, through a lookup by name and a lookup by
68+
* URI. Each step changes the state through one lookup and checks the result through the other,
69+
* so the task fails unless both lookups reach the same asset. The E2E test then reads the
70+
* remaining keys through the REST API.
71+
*/
72+
public static class UseAssetStateStore implements Task {
73+
public void execute(@NotNull Context context, Client client) {
74+
var byName = client.assetStateStore.byName("java_e2e_orders");
75+
var byUri = client.assetStateStore.byUri("x-java-e2e://orders");
76+
77+
byName.set("scratch", "first");
78+
expect("first", byUri.get("scratch"), "set() through the name");
79+
byUri.clear();
80+
expect(null, byName.get("scratch"), "clear() through the URI");
81+
82+
byUri.set("scratch", "second");
83+
expect("second", byName.get("scratch"), "set() through the URI");
84+
byName.clear();
85+
expect(null, byUri.get("scratch"), "clear() through the name");
86+
87+
byName.set("deleted", "temporary");
88+
byUri.delete("deleted");
89+
expect(null, byName.get("deleted"), "delete() through the URI");
90+
byUri.set("deleted", "temporary");
91+
byName.delete("deleted");
92+
expect(null, byUri.get("deleted"), "delete() through the name");
93+
94+
var summary = Map.of("run_id", context.dagRun.runId, "rows", 3L);
95+
byName.set("summary", summary);
96+
expect(summary, byUri.get("summary"), "set() through the name, with a map value");
97+
}
98+
99+
private static void expect(Object expected, Object actual, String step) {
100+
if (!Objects.equals(expected, actual)) {
101+
throw new IllegalStateException(step + ": expected " + expected + " but read " + actual);
102+
}
103+
}
104+
}
105+
63106
@NotNull
64107
@Override
65108
public Iterable<DagDef> getDags() {
@@ -68,7 +111,9 @@ public Iterable<DagDef> getDags() {
68111
uninstantiable.addTask("non_static_inner", NonStaticInner.class);
69112
var variableWrite = new DagDef("java_variable_write");
70113
variableWrite.addTask("write_and_delete", WriteAndDeleteVariable.class);
71-
return List.of(uninstantiable, variableWrite);
114+
var assetStateStore = new DagDef("java_asset_state_store");
115+
assetStateStore.addTask("use_asset_state_store", UseAssetStateStore.class);
116+
return List.of(uninstantiable, variableWrite, assetStateStore);
72117
}
73118

74119
public static void main(String[] args) {

‎airflow-e2e-tests/java-test-bundle/src/resources/dags/java_test_dags.py‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818

1919
from __future__ import annotations
2020

21-
from airflow.sdk import dag, task
21+
from airflow.sdk import Asset, dag, task
2222

2323

2424
@task.stub(queue="java-test")
@@ -48,3 +48,20 @@ def java_variable_write():
4848

4949

5050
java_variable_write()
51+
52+
53+
# The Java task cannot see this inlet. Declaring it keeps the asset active, and the
54+
# Execution API only reads and writes the state of active assets.
55+
java_e2e_orders = Asset(name="java_e2e_orders", uri="x-java-e2e://orders")
56+
57+
58+
@task.stub(queue="java-test", inlets=[java_e2e_orders])
59+
def use_asset_state_store(): ...
60+
61+
62+
@dag(dag_id="java_asset_state_store")
63+
def java_asset_state_store():
64+
use_asset_state_store()
65+
66+
67+
java_asset_state_store()

‎airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,14 @@ def get_variable(self, key: str):
163163
"""Get an Airflow Variable via API."""
164164
return self._make_request(method="GET", endpoint=f"variables/{key}")
165165

166+
def get_asset_state_store_entry(self, asset_name: str, key: str):
167+
"""Get one key of an asset's state store via API, looking up the asset by its name."""
168+
# name_pattern matches substrings, so keep only the asset named exactly asset_name.
169+
assets = self._make_request(method="GET", endpoint="assets", params={"name_pattern": asset_name})
170+
asset_ids = [asset["id"] for asset in assets["assets"] if asset["name"] == asset_name]
171+
assert len(asset_ids) == 1, f"expected one asset named {asset_name!r}, got {assets!r}"
172+
return self._make_request(method="GET", endpoint=f"assets/{asset_ids[0]}/state-store/{key}")
173+
166174
def trigger_dag_and_wait(self, dag_id: str, json=None):
167175
"""Trigger a DAG and wait for it to complete."""
168176
self.un_pause_dag(dag_id)

‎airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,9 @@
9595
_SCALA_SPARK_DAG_ID = "scala_spark_example"
9696
_VARIABLE_WRITE_DAG_ID = "java_variable_write"
9797
_VARIABLE_WRITE_DESCRIPTION = "written by the Java SDK e2e test"
98+
_ASSET_STATE_STORE_DAG_ID = "java_asset_state_store"
99+
_ASSET_STATE_STORE_TASK_ID = "use_asset_state_store"
100+
_ASSET_STATE_STORE_ASSET_NAME = "java_e2e_orders"
98101

99102

100103
@dataclass
@@ -205,6 +208,12 @@ def variable_write_run() -> _CompletedRun:
205208
return _trigger_and_wait_for_dag(_VARIABLE_WRITE_DAG_ID, _JAVA_TASK_TIMEOUT)
206209

207210

211+
@pytest.fixture(scope="module")
212+
def asset_state_store_run() -> _CompletedRun:
213+
"""Trigger the asset state store Dag once for all of its assertions."""
214+
return _trigger_and_wait_for_dag(_ASSET_STATE_STORE_DAG_ID, _JAVA_TASK_TIMEOUT)
215+
216+
208217
class TestJavaSDKAnnotationExample:
209218
"""Verify the annotation-based Java SDK example executes correctly."""
210219

@@ -406,6 +415,56 @@ def test_scratch_variable_deleted_by_java_task_is_gone(self, variable_write_run:
406415
)
407416

408417

418+
class TestJavaSDKAssetStateStore:
419+
"""Verify a Java task can read and write the state store of an asset.
420+
421+
The task lives in the java-test-bundle fixture project (served on the
422+
dedicated "java-test" queue). It looks up the ``java_e2e_orders`` asset both
423+
by name and by URI and calls every asset state store method. Each step
424+
changes the state through one lookup and checks it through the other, so the
425+
task fails unless both lookups reach the same asset. It leaves the
426+
``summary`` key behind for the REST API to read.
427+
"""
428+
429+
def _assert_task_succeeded(self, run: _CompletedRun) -> None:
430+
ti = run.get_task_instance(_ASSET_STATE_STORE_TASK_ID)
431+
assert ti.get("state") == "success", (
432+
f"Java {_ASSET_STATE_STORE_TASK_ID!r} task did not succeed.\n"
433+
f" task state : {ti.get('state')!r}\n"
434+
f" dag state : {run.state!r}\n"
435+
f" all tasks : {run.ti_states}"
436+
)
437+
438+
def test_value_written_by_java_task_is_readable(self, asset_state_store_run: _CompletedRun):
439+
"""The value set from Java is stored on the asset and attributed to the Java task."""
440+
self._assert_task_succeeded(asset_state_store_run)
441+
442+
entry = asset_state_store_run.client.get_asset_state_store_entry(
443+
_ASSET_STATE_STORE_ASSET_NAME, "summary"
444+
)
445+
assert entry["value"] == {"run_id": asset_state_store_run.run_id, "rows": 3}, (
446+
f"summary should hold the map the Java task stored in this run, got {entry!r}"
447+
)
448+
assert entry["last_updated_by"] == {
449+
"kind": "task",
450+
"dag_id": _ASSET_STATE_STORE_DAG_ID,
451+
"run_id": asset_state_store_run.run_id,
452+
"task_id": _ASSET_STATE_STORE_TASK_ID,
453+
"map_index": -1,
454+
}, f"summary should be attributed to the Java task, got {entry!r}"
455+
456+
def test_key_deleted_by_java_task_is_gone(self, asset_state_store_run: _CompletedRun):
457+
"""A key written and then deleted from Java no longer exists."""
458+
self._assert_task_succeeded(asset_state_store_run)
459+
460+
with pytest.raises(requests.HTTPError) as excinfo:
461+
asset_state_store_run.client.get_asset_state_store_entry(_ASSET_STATE_STORE_ASSET_NAME, "deleted")
462+
assert excinfo.value.response.status_code == HTTPStatus.NOT_FOUND, (
463+
f"the 'deleted' key should have been deleted by the Java task, "
464+
f"got HTTP {excinfo.value.response.status_code}"
465+
)
466+
467+
409468
class TestJavaSDKInterfaceExample:
410469
"""Verify tasks written against ``InputTask`` receive their bound arguments."""
411470

‎java-sdk/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -615,7 +615,7 @@ prek hook regenerate it.
615615
| capability: `self-contained-bundle` | MUST | ✓ | 3.3 | Airflow metadata embedded in the jar artifact |
616616
| capability: `retry-policy` | MAY | ✗ | – | no task-facing retry-policy API yet |
617617
| capability: `task-state-store` | MAY | ✗ | – | no task-facing state-store API yet |
618-
| capability: `asset-state-store` | MAY | ✗ | – | no task-facing state-store API yet |
618+
| capability: `asset-state-store` | MAY | ✓ | 3.4 | Client.assetStateStore, by asset name or URI |
619619
| capability: `asset-event-emit` | MAY | ✗ | – | runtime does not emit asset events yet |
620620
| capability: `asset-event-read` | MAY | ✗ | – | no task-facing asset-event API yet |
621621
| **Native-Dag authoring** | | | | |

‎java-sdk/capabilities.yaml‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -86,8 +86,9 @@ capabilities:
8686
supported: false
8787
note: "no task-facing state-store API yet"
8888
asset-state-store:
89-
supported: false
90-
note: "no task-facing state-store API yet"
89+
supported: true
90+
since: "3.4"
91+
note: "Client.assetStateStore, by asset name or URI"
9192
asset-event-emit:
9293
supported: false
9394
note: "runtime does not emit asset events yet"

‎java-sdk/sdk/module.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ meaning of each dimension is defined in the
5252
| capability: `self-contained-bundle` | MUST | ✓ | 3.3 | Airflow metadata embedded in the jar artifact |
5353
| capability: `retry-policy` | MAY | ✗ | – | no task-facing retry-policy API yet |
5454
| capability: `task-state-store` | MAY | ✗ | – | no task-facing state-store API yet |
55-
| capability: `asset-state-store` | MAY | ✗ | – | no task-facing state-store API yet |
55+
| capability: `asset-state-store` | MAY | ✓ | 3.4 | Client.assetStateStore, by asset name or URI |
5656
| capability: `asset-event-emit` | MAY | ✗ | – | runtime does not emit asset events yet |
5757
| capability: `asset-event-read` | MAY | ✗ | – | no task-facing asset-event API yet |
5858
| **Native-Dag authoring** | | | | |

0 commit comments

Comments
 (0)