Skip to content

Commit 4ff0653

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

17 files changed

Lines changed: 766 additions & 8 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
@@ -1099,6 +1099,59 @@ structured data (for example to JSON) before storing it.
10991099
takes precedence over the stored value when the Variable is read back. Calling ``setVariable``
11001100
without a description clears any existing description.
11011101

1102+
.. _java-sdk/asset-state-store:
1103+
1104+
Asset state store
1105+
-----------------
1106+
1107+
``client.getAssetStateStore()`` reads and writes the :doc:`asset state store </core-concepts/asset-state-store>`,
1108+
a key-value store that belongs to an asset rather than to a Dag run. A value that one run stores, such as an
1109+
incremental-load watermark, is still there in later runs.
1110+
1111+
The Java SDK does not receive the inlets and outlets declared on the ``@task.stub``. The task therefore looks
1112+
up the asset itself, by name with ``byName`` or by URI with ``byUri``. Airflow keeps state only for assets
1113+
that a Dag references, so also declare the asset as an inlet of the stub:
1114+
1115+
.. code-block:: python
1116+
1117+
orders = Asset(name="orders", uri="s3://warehouse/orders")
1118+
1119+
1120+
@task.stub(queue="java", inlets=[orders])
1121+
def load_orders(): ...
1122+
1123+
.. code-block:: java
1124+
1125+
@Airflow.TaskHandler(dag = "orders_pipeline", task = "load_orders")
1126+
public void loadOrders(Client client) {
1127+
var orders = client.getAssetStateStore().byName("orders");
1128+
var watermark = (String) orders.get("watermark"); // null on the first run
1129+
// Your own code: load the rows created after the watermark and return the newest creation time.
1130+
String newest = loadRowsCreatedAfter(watermark);
1131+
orders.set("watermark", newest);
1132+
}
1133+
1134+
``get`` returns ``null`` when the key is not set. ``set`` replaces any existing value. Entries do not expire.
1135+
``delete`` removes one key, and ``clear`` removes every key of the asset.
1136+
1137+
``set`` sends the value as JSON. Store strings, booleans, integers that fit in a ``long``, finite ``double``
1138+
values, lists, and maps with string keys. Jackson converts other objects first. A bean becomes a map of its
1139+
properties, and an ``Instant`` or a ``UUID`` becomes a string. ``null``, ``BigDecimal``, and ``Duration``
1140+
fail with an exception before anything is sent. Never store a ``byte[]``. The supervisor cannot decode it and
1141+
never answers, so the task hangs. Encode binary data as a string instead, for example with Base64.
1142+
1143+
``get`` returns the Java types listed in :ref:`java-sdk/types`. For example, a value stored as an ``Integer``
1144+
comes back as a ``Long``.
1145+
1146+
If no Dag references an asset with the name or URI that the task passed, ``get`` returns ``null`` and the
1147+
other methods throw ``ApiError``. Airflow normalizes asset URIs when it parses a Dag, so pass ``byUri`` the
1148+
normalized form. For example, ``s3://warehouse/orders/`` is stored as ``s3://warehouse/orders``.
1149+
1150+
The Java SDK does not apply a ``[workers] state_store_backend``. It sends each value to Airflow as-is. A key
1151+
that a Python task stored through such a backend comes back to Java as a map that holds only the reference
1152+
under ``__airflow_state_ref__``, not the stored value. Deleting or clearing such a key from Java removes only
1153+
that reference, and the stored value stays in the backend.
1154+
11021155
.. _java-sdk/build:
11031156

11041157
Building and packaging

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

Lines changed: 46 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,14 @@
2121

2222
import java.util.List;
2323
import java.util.Map;
24+
import java.util.Objects;
2425
import org.apache.airflow.sdk.*;
2526
import org.jetbrains.annotations.NotNull;
2627

2728
/**
2829
* Bundle for the runner-behaviour E2E tests: deliberately broken task classes that exercise
29-
* instantiation failures, and tasks that round-trip Airflow Variables and task state through the
30-
* supervisor.
30+
* instantiation failures, and tasks that round-trip Airflow Variables, task state and asset state
31+
* through the supervisor.
3132
*/
3233
public class TestBundleBuilder implements BundleBuilder {
3334
public static class MissingNoArgConstructor implements Task {
@@ -90,6 +91,46 @@ public void execute(@NotNull Context context, Client client) {
9091
}
9192
}
9293

94+
/**
95+
* Calls every asset state store method on one asset, through a lookup by name and a lookup by
96+
* URI. Each step changes the state through one lookup and checks the result through the other,
97+
* so the task fails unless both lookups reach the same asset. The E2E test then reads the
98+
* remaining keys through the REST API.
99+
*/
100+
public static class UseAssetStateStore implements Task {
101+
public void execute(@NotNull Context context, Client client) {
102+
var byName = client.getAssetStateStore().byName("java_e2e_orders");
103+
var byUri = client.getAssetStateStore().byUri("x-java-e2e://orders");
104+
105+
byName.set("scratch", "first");
106+
expect("first", byUri.get("scratch"), "set() through the name");
107+
byUri.clear();
108+
expect(null, byName.get("scratch"), "clear() through the URI");
109+
110+
byUri.set("scratch", "second");
111+
expect("second", byName.get("scratch"), "set() through the URI");
112+
byName.clear();
113+
expect(null, byUri.get("scratch"), "clear() through the name");
114+
115+
byName.set("deleted", "temporary");
116+
byUri.delete("deleted");
117+
expect(null, byName.get("deleted"), "delete() through the URI");
118+
byUri.set("deleted", "temporary");
119+
byName.delete("deleted");
120+
expect(null, byUri.get("deleted"), "delete() through the name");
121+
122+
var summary = Map.of("run_id", context.dagRun.runId, "rows", 3L);
123+
byName.set("summary", summary);
124+
expect(summary, byUri.get("summary"), "set() through the name, with a map value");
125+
}
126+
127+
private static void expect(Object expected, Object actual, String step) {
128+
if (!Objects.equals(expected, actual)) {
129+
throw new IllegalStateException(step + ": expected " + expected + " but read " + actual);
130+
}
131+
}
132+
}
133+
93134
@NotNull
94135
@Override
95136
public Iterable<DagDef> getDags() {
@@ -100,7 +141,9 @@ public Iterable<DagDef> getDags() {
100141
variableWrite.addTask("write_and_delete", WriteAndDeleteVariable.class);
101142
var taskState = new DagDef("java_task_state");
102143
taskState.addTask("roundtrip_task_state", RoundtripTaskState.class);
103-
return List.of(uninstantiable, variableWrite, taskState);
144+
var assetStateStore = new DagDef("java_asset_state_store");
145+
assetStateStore.addTask("use_asset_state_store", UseAssetStateStore.class);
146+
return List.of(uninstantiable, variableWrite, taskState, assetStateStore);
104147
}
105148

106149
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")
@@ -60,3 +60,20 @@ def java_task_state():
6060

6161

6262
java_task_state()
63+
64+
65+
# The Java task cannot see this inlet. Declaring it keeps the asset active, and the
66+
# Execution API only reads and writes the state of active assets.
67+
java_e2e_orders = Asset(name="java_e2e_orders", uri="x-java-e2e://orders")
68+
69+
70+
@task.stub(queue="java-test", inlets=[java_e2e_orders])
71+
def use_asset_state_store(): ...
72+
73+
74+
@dag(dag_id="java_asset_state_store")
75+
def java_asset_state_store():
76+
use_asset_state_store()
77+
78+
79+
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
@@ -230,6 +230,14 @@ def get_dag_run(self, dag_id: str, run_id: str):
230230
"""Get a Dag run, with its state, run type and conf."""
231231
return self._make_request(method="GET", endpoint=f"dags/{dag_id}/dagRuns/{run_id}")
232232

233+
def get_asset_state_store_entry(self, asset_name: str, key: str):
234+
"""Get one key of an asset's state store via API, looking up the asset by its name."""
235+
# name_pattern matches substrings, so keep only the asset named exactly asset_name.
236+
assets = self._make_request(method="GET", endpoint="assets", params={"name_pattern": asset_name})
237+
asset_ids = [asset["id"] for asset in assets["assets"] if asset["name"] == asset_name]
238+
assert len(asset_ids) == 1, f"expected one asset named {asset_name!r}, got {assets!r}"
239+
return self._make_request(method="GET", endpoint=f"assets/{asset_ids[0]}/state-store/{key}")
240+
233241
def trigger_dag_and_wait(self, dag_id: str, json=None):
234242
"""Trigger a DAG and wait for it to complete."""
235243
self.un_pause_dag(dag_id)

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

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,9 @@
9898
_VARIABLE_WRITE_DESCRIPTION = "written by the Java SDK e2e test"
9999
_TASK_STATE_DAG_ID = "java_task_state"
100100
_TASK_STATE_TASK_ID = "roundtrip_task_state"
101+
_ASSET_STATE_STORE_DAG_ID = "java_asset_state_store"
102+
_ASSET_STATE_STORE_TASK_ID = "use_asset_state_store"
103+
_ASSET_STATE_STORE_ASSET_NAME = "java_e2e_orders"
101104

102105

103106
@dataclass
@@ -226,6 +229,12 @@ def task_state_run() -> _CompletedRun:
226229
return _trigger_and_wait_for_dag(_TASK_STATE_DAG_ID, _JAVA_TASK_TIMEOUT)
227230

228231

232+
@pytest.fixture(scope="module")
233+
def asset_state_store_run() -> _CompletedRun:
234+
"""Trigger the asset state store Dag once for all of its assertions."""
235+
return _trigger_and_wait_for_dag(_ASSET_STATE_STORE_DAG_ID, _JAVA_TASK_TIMEOUT)
236+
237+
229238
class TestJavaSDKAnnotationExample:
230239
"""Verify the annotation-based Java SDK example executes correctly."""
231240

@@ -427,6 +436,63 @@ def test_scratch_variable_deleted_by_java_task_is_gone(self, variable_write_run:
427436
)
428437

429438

439+
class TestJavaSDKAssetStateStore:
440+
"""Verify a Java task can read and write the state store of an asset.
441+
442+
The task lives in the java-test-bundle fixture project (served on the
443+
dedicated "java-test" queue). It looks up the ``java_e2e_orders`` asset both
444+
by name and by URI and calls every asset state store method. Each step
445+
changes the state through one lookup and checks it through the other, so the
446+
task fails unless both lookups reach the same asset. It leaves the
447+
``summary`` key behind for the REST API to read.
448+
"""
449+
450+
def _assert_task_succeeded(self, run: _CompletedRun) -> None:
451+
ti = run.get_task_instance(_ASSET_STATE_STORE_TASK_ID)
452+
assert ti.get("state") == "success", (
453+
f"Java {_ASSET_STATE_STORE_TASK_ID!r} task did not succeed.\n"
454+
f" task state : {ti.get('state')!r}\n"
455+
f" dag state : {run.state!r}\n"
456+
f" all tasks : {run.ti_states}"
457+
)
458+
459+
def test_value_written_by_java_task_is_readable(self, asset_state_store_run: _CompletedRun):
460+
"""The value set from Java is stored on the asset and attributed to the Java task."""
461+
self._assert_task_succeeded(asset_state_store_run)
462+
463+
entry = asset_state_store_run.client.get_asset_state_store_entry(
464+
_ASSET_STATE_STORE_ASSET_NAME, "summary"
465+
)
466+
assert entry["value"] == {"run_id": asset_state_store_run.run_id, "rows": 3}, (
467+
f"summary should hold the map the Java task stored in this run, got {entry!r}"
468+
)
469+
ti = asset_state_store_run.get_task_instance(_ASSET_STATE_STORE_TASK_ID)
470+
expected_writer = {
471+
"kind": "task",
472+
"dag_id": _ASSET_STATE_STORE_DAG_ID,
473+
"run_id": asset_state_store_run.run_id,
474+
"task_id": _ASSET_STATE_STORE_TASK_ID,
475+
"map_index": -1,
476+
"task_instance_id": ti["id"],
477+
"try_number": ti["try_number"],
478+
}
479+
writer = entry["last_updated_by"]
480+
assert {key: writer.get(key) for key in expected_writer} == expected_writer, (
481+
f"summary should be attributed to the Java task, got {entry!r}"
482+
)
483+
484+
def test_key_deleted_by_java_task_is_gone(self, asset_state_store_run: _CompletedRun):
485+
"""A key written and then deleted from Java no longer exists."""
486+
self._assert_task_succeeded(asset_state_store_run)
487+
488+
with pytest.raises(requests.HTTPError) as excinfo:
489+
asset_state_store_run.client.get_asset_state_store_entry(_ASSET_STATE_STORE_ASSET_NAME, "deleted")
490+
assert excinfo.value.response.status_code == HTTPStatus.NOT_FOUND, (
491+
f"the 'deleted' key should have been deleted by the Java task, "
492+
f"got HTTP {excinfo.value.response.status_code}"
493+
)
494+
495+
430496
class TestJavaSDKInterfaceExample:
431497
"""Verify tasks written against ``InputTask`` receive their bound arguments."""
432498

‎java-sdk/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -635,7 +635,7 @@ prek hook regenerate it.
635635
| capability: `self-contained-bundle` | MUST | ✓ | 3.3 | Airflow metadata embedded in the jar artifact |
636636
| capability: `retry-policy` | MAY | ✗ | – | no task-facing retry-policy API yet |
637637
| capability: `task-state-store` | MAY | ✓ | 3.3 | Client.getTaskStateStore() get/set/delete/clear |
638-
| capability: `asset-state-store` | MAY | ✗ | – | no task-facing state-store API yet |
638+
| capability: `asset-state-store` | MAY | ✓ | 3.4 | Client.getAssetStateStore(), by asset name or URI |
639639
| capability: `asset-event-emit` | MAY | ✗ | – | runtime does not emit asset events yet |
640640
| capability: `asset-event-read` | MAY | ✗ | – | no task-facing asset-event API yet |
641641
| **Native-Dag authoring** | | | | |

‎java-sdk/capabilities.yaml‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,8 +89,9 @@ capabilities:
8989
since: "3.3"
9090
note: "Client.getTaskStateStore() get/set/delete/clear"
9191
asset-state-store:
92-
supported: false
93-
note: "no task-facing state-store API yet"
92+
supported: true
93+
since: "3.4"
94+
note: "Client.getAssetStateStore(), by asset name or URI"
9495
asset-event-emit:
9596
supported: false
9697
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 | ✓ | 3.3 | Client.getTaskStateStore() get/set/delete/clear |
55-
| capability: `asset-state-store` | MAY | ✗ | – | no task-facing state-store API yet |
55+
| capability: `asset-state-store` | MAY | ✓ | 3.4 | Client.getAssetStateStore(), 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)