Andrushika commented on code in PR #73464:
URL: https://github.com/apache/airflow/pull/73464#discussion_r4193839571
##########
java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt:
##########
@@ -215,6 +226,63 @@ class Client internal constructor(
}
}
+/**
+ * Key-value state scoped to one task instance, shared across its retries
+ * within the same Dag run.
+ *
+ * Values must be JSON-serializable. Keys expire at the time given when they
+ * were stored; a key stored without a retention never expires and is skipped
+ * by Airflow's periodic garbage collection.
+ *
+ * Values are stored in the metadata database as-is; the `[workers]
+ * state_store_backend` used by Python tasks is not applied here.
+ */
+class TaskStateStore internal constructor(
+ private val details: StartupDetails,
+ private val impl: Client,
+) {
+ /**
+ * Reads the value stored under [key].
+ *
+ * @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.getTaskStateStore(details.ti.id,
key)?.value
+
+ /**
+ * Stores [value] under [key], replacing any existing value.
+ *
+ * @param key State key.
+ * @param value Value to store. Must be JSON-serializable.
+ * @param retention How long to keep the key, or `null` to keep it until
deleted.
+ * @throws ApiError if the API call fails.
+ */
+ @JvmOverloads fun set(
+ key: String,
+ value: Any,
+ retention: Duration? = null,
+ ) = impl.setTaskStateStore(
+ tiId = details.ti.id,
+ key = key,
+ value = value,
+ expiresAt = retention?.let { OffsetDateTime.now(ZoneOffset.UTC).plus(it) },
+ )
Review Comment:
Thanks! Fixed these:
- Added `TaskStateStore.NEVER_EXPIRE`
- `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
- Updated `java.rst` and the PR description
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
