Skip to content

Add task state store API to the Java SDK - #73464

Merged
jason810496 merged 5 commits into
apache:mainfrom
Andrushika:java-sdk-task-state-store
Oct 6, 2026
Merged

jason810496 merged 5 commits into
apache:mainfrom
Andrushika:java-sdk-task-state-store

Conversation

@Andrushika

@Andrushika Andrushika commented Sep 21, 2026 •

Copy link
Copy Markdown
Contributor

Why

Airflow 3.3 added the task state store (AIP-103). A task can save a value under a key, and the value survives retries.

The supervisor already handles the four *TaskStateStore messages and the Java SDK schema.json already has them. But nothing on the Kotlin side sends them. So Java tasks cannot use the store and capabilities.yaml lists task-state-store as unsupported.

What

Add client.getTaskStateStore() with get, set, delete and clear. It is wired the same way as variables. get returns null when the key is not there. set takes an optional java.time.Duration retention, or TaskStateStore.NEVER_EXPIRE. Without a retention it uses [state_store] default_retention_days, read from AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS the same way as the Go SDK in #73420 (0 means never expire, a missing variable fails).

SetTaskStateStore is added to REQUIRED_NULLABLE_REQUESTS. Without it Jackson drops a null expires_at and the supervisor rejects the message. I checked this with the supervisor decoder.

One difference from Python is noted in java.rst: [workers] state_store_backend is not supported. The backend hooks run inside the Python task process today, so supporting it in the Java SDK needs a further design discussion.

related: #73420


Was generative AI tooling used to co-author this PR?
  • Yes — Claude Code (Fable 5.1)

Generated-by: Claude Code (Fable 5.1) following the guidelines

@potiuk

potiuk commented Sep 25, 2026

Copy link
Copy Markdown
Member

Hello @Andrushika - thank you for your contributions to Apache Airflow!

The Airflow community has introduced a limit of 5 open pull requests at a time for contributors without write access to the repository. You currently have 7 open pull requests, so - as a one-time step of introducing the limit - we closed the ones where maintainers have not engaged yet:

These pull requests stay open because maintainers are already engaged in them - they count towards your limit:

This is not a judgement of you or of your changes. We never told contributors before that opening many pull requests at once was a problem, so there is nothing to feel bad about - and nothing is lost: your branches, commits and the review history stay where they are.

What we ask you to do is to make your first prioritization decision: choose which of the pull requests above matter most to you, and reopen them (up to 5 open at a time, including the ones still open) with the "Reopen pull request" button or gh pr reopen <PR_NUMBER> --repo apache/airflow. Reopen the ones you are ready to follow through - keep them rebased, respond to review comments and fix failing checks.

While your pull requests are waiting for review, the most valuable thing you can do is help in other ways - reviewing other contributors' pull requests, helping with issues, and taking part in the discussions on the devlist and Slack.

Why we introduced the limit, what it means for you and how to reopen or restore a pull request is explained in https://github.com/apache/airflow/blob/main/contributing-docs/32_open_pull_request_limit.rst.


Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting

@potiuk potiuk closed this Sep 25, 2026
@jason810496 jason810496 reopened this Sep 27, 2026
@Andrushika
Andrushika force-pushed the java-sdk-task-state-store branch from 249ca7e to 1b0ca65 Compare September 27, 2026 14:47
Comment thread airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst Outdated
Comment thread java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt Outdated

@FrankYang0529 FrankYang0529 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR. Leave one comment.

Comment thread java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt Outdated
@Andrushika
Andrushika force-pushed the java-sdk-task-state-store branch from 1b0ca65 to 4b9c632 Compare October 5, 2026 02:58

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice, thanks! LGTM overall.

Let's also add a TODO before the serialization to respect the state_store/max_value_storage_bytes that will be passed from the coordinator side.

Comment thread java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt Outdated
Airflow 3.3 added a task state store (AIP-103) and the supervisor already
handles the Get/Set/Delete/ClearTaskStateStore messages, but the Java SDK
had no task-facing API for it and capabilities.yaml listed it as
unsupported. Java tasks could not keep state such as an external job ID
across retries.
Review found that a @JvmField on Client cannot be stubbed by Mockito, so a
Java unit test of a task that reads the store hit a NullPointerException.
The docs also said entries survive later runs, but the store is scoped to
one task instance, so they only survive retries within the same Dag run.
A key stored without a retention never expired, while Python applies
[state_store] default_retention_days. The coordinator passes that setting
to the JVM as AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS, so read it
the same way the Go SDK does: absent means 30 days, 0 means never expire,
and a malformed value fails instead of retaining for a different period.
TaskStateStore.NEVER_EXPIRE replaces the old "no retention" meaning, and a
zero or negative retention is rejected.
@Andrushika
Andrushika force-pushed the java-sdk-task-state-store branch from 4b9c632 to f775faa Compare October 6, 2026 08:54
@Andrushika
Andrushika requested a review from jason810496 October 6, 2026 09:49
Review on the Go and TS task state store PRs asked to drop the hardcoded
30-day fallback: the coordinator always passes
AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS, so a missing variable means
the JVM was not launched by the coordinator, and guessing a period there
would drift from the deployment's config. Also leave a TODO for the
max_value_storage_bytes warning the Python accessor emits.
@Andrushika

Copy link
Copy Markdown
Contributor Author

Also applied the two points from #73420 and #73618 that apply here:

  • Removed the 30-day fallback. set without a retention now fails with a clear error when AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS is absent.
  • Added a TODO before sending for the max_value_storage_bytes warning.

One note on ordering: the coordinator only passes that variable once #73420 is merged. If this PR goes in first, set without a retention will fail until then.

@jason810496

jason810496 commented Oct 6, 2026 •

Copy link
Copy Markdown
Member

Also applied the two points from #73420 and #73618 that apply here:

  • Removed the 30-day fallback. set without a retention now fails with a clear error when AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS is absent.
  • Added a TODO before sending for the max_value_storage_bytes warning.

One note on ordering: the coordinator only passes that variable once #73420 is merged. If this PR goes in first, set without a retention will fail until then.

Would it be better to add the explicit NEVER_EXPIRE sentinel instead of using 0 as the never expire?

@Andrushika

Copy link
Copy Markdown
Contributor Author

Yes, we already have NEVER_EXPIRE for the Client API.

set without a retention now reads AIRFLOW__STATE_STORE__DEFAULT_RETENTION_DAYS, with the same rules as the Go SDK in #73420: absent means 30 days, 0 means never expire, and a malformed value fails. A zero or negative retention is rejected too

The 0 I described in the resolved comment is the config value (default_retention_days) passed from the coordinator, not a value users pass to set.

Sorry for the misleading wording :(

@jason810496 jason810496 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice, thanks.

…-store

# Conflicts:
#	airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:java-sdk closed because of open PR limit Closed as a one-time step of introducing the open pull request limit kind:documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants