This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 2678bc8e3ae Java SDK: Add setVariable and deleteVariable to the task
Client (#72676)
2678bc8e3ae is described below
commit 2678bc8e3ae55b4568d4a0e456658be830b8af9f
Author: PoAn Yang <[email protected]>
AuthorDate: Mon Sep 14 16:39:25 2026 +0900
Java SDK: Add setVariable and deleteVariable to the task Client (#72676)
* Java SDK: Add setVariable and deleteVariable to the task Client
Signed-off-by: PoAn Yang <[email protected]>
* Remove null description case
Signed-off-by: PoAn Yang <[email protected]>
---------
Signed-off-by: PoAn Yang <[email protected]>
---
.../language-sdks/java.rst | 18 +++
.../org/apache/airflow/e2e/TestBundleBuilder.java | 26 ++++-
.../src/resources/dags/java_test_dags.py | 12 ++
.../airflow_e2e_tests/e2e_test_utils/clients.py | 4 +
.../java_sdk_tests/test_java_sdk_dag.py | 55 +++++++++
java-sdk/README.md | 2 +-
java-sdk/capabilities.yaml | 4 +-
java-sdk/sdk/module.md | 2 +-
.../main/kotlin/org/apache/airflow/sdk/Client.kt | 25 +++++
.../org/apache/airflow/sdk/execution/Client.kt | 29 +++++
.../org/apache/airflow/sdk/execution/Frame.kt | 14 +++
.../kotlin/org/apache/airflow/sdk/ClientTest.kt | 8 ++
.../org/apache/airflow/sdk/execution/CommTest.kt | 123 +++++++++++++++++++++
.../org/apache/airflow/sdk/execution/TaskTest.kt | 8 ++
14 files changed, 321 insertions(+), 9 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
index 4bae1652230..27ad40ab675 100644
--- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
+++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
@@ -619,6 +619,24 @@ represented as Java objects when read back via ``getXCom``.
``null`` and the task fails with ``MissingXComException``. Declare the
parameter with a
boxed type when the upstream XCom may be absent.
+Variables
+---------
+
+``Client`` reads, writes, and deletes Airflow Variables. Values are stored as
strings. Serialize
+structured data (for example to JSON) before storing it.
+
+.. code-block:: java
+
+ var threshold = (String) client.getVariable("process_threshold");
+ client.setVariable("process_threshold", "42", "Rows above this count take
the slow path");
+ client.deleteVariable("legacy_threshold");
+
+.. note::
+
+ A value supplied by a secrets backend (for example an ``AIRFLOW_VAR_*``
environment variable) still
+ takes precedence over the stored value when the Variable is read back.
Calling ``setVariable``
+ without a description clears any existing description.
+
.. _java-sdk/build:
Building and packaging
diff --git
a/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java
b/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java
index 5986f9d57b0..ad33a15708a 100644
---
a/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java
+++
b/airflow-e2e-tests/java-test-bundle/src/java/org/apache/airflow/e2e/TestBundleBuilder.java
@@ -24,7 +24,8 @@ import org.apache.airflow.sdk.*;
import org.jetbrains.annotations.NotNull;
/**
- * Bundle of deliberately broken task classes for the runner-behaviour E2E
tests.
+ * Bundle for the runner-behaviour E2E tests: deliberately broken task classes
that exercise
+ * instantiation failures, and a task that round-trips Airflow Variables
through the supervisor.
*/
public class TestBundleBuilder implements BundleBuilder {
public static class MissingNoArgConstructor implements Task {
@@ -46,13 +47,28 @@ public class TestBundleBuilder implements BundleBuilder {
}
}
+ /**
+ * Stores this run's id where the E2E test can read it back through the REST
API, then writes
+ * and deletes a scratch variable to exercise the delete path.
+ */
+ public static class WriteAndDeleteVariable implements Task {
+ public void execute(@NotNull Context context, Client client) {
+ client.setVariable(
+ "java_e2e_variable", context.dagRun.runId, "written by the Java SDK
e2e test");
+ client.setVariable("java_e2e_scratch", "scratch");
+ client.deleteVariable("java_e2e_scratch");
+ }
+ }
+
@NotNull
@Override
public Iterable<DagDef> getDags() {
- var dag = new DagDef("java_uninstantiable");
- dag.addTask("missing_no_arg_constructor", MissingNoArgConstructor.class);
- dag.addTask("non_static_inner", NonStaticInner.class);
- return List.of(dag);
+ var uninstantiable = new DagDef("java_uninstantiable");
+ uninstantiable.addTask("missing_no_arg_constructor",
MissingNoArgConstructor.class);
+ uninstantiable.addTask("non_static_inner", NonStaticInner.class);
+ var variableWrite = new DagDef("java_variable_write");
+ variableWrite.addTask("write_and_delete", WriteAndDeleteVariable.class);
+ return List.of(uninstantiable, variableWrite);
}
public static void main(String[] args) {
diff --git
a/airflow-e2e-tests/java-test-bundle/src/resources/dags/java_test_dags.py
b/airflow-e2e-tests/java-test-bundle/src/resources/dags/java_test_dags.py
index 09274b77cd9..d4ee554d66e 100644
--- a/airflow-e2e-tests/java-test-bundle/src/resources/dags/java_test_dags.py
+++ b/airflow-e2e-tests/java-test-bundle/src/resources/dags/java_test_dags.py
@@ -36,3 +36,15 @@ def java_uninstantiable():
java_uninstantiable()
+
+
[email protected](queue="java-test")
+def write_and_delete(): ...
+
+
+@dag(dag_id="java_variable_write")
+def java_variable_write():
+ write_and_delete()
+
+
+java_variable_write()
diff --git
a/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py
b/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py
index 0c1166cfc8f..d56cecae3d7 100644
--- a/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py
+++ b/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py
@@ -159,6 +159,10 @@ class AirflowClient:
endpoint=f"dags/{dag_id}/dagRuns/{run_id}/taskInstances/{task_id}/xcomEntries/{key}?map_index={map_index}",
)
+ def get_variable(self, key: str):
+ """Get an Airflow Variable via API."""
+ return self._make_request(method="GET", endpoint=f"variables/{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)
diff --git
a/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py
b/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py
index 16fa1b61f28..07180d2c8a0 100644
---
a/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py
+++
b/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py
@@ -61,9 +61,11 @@ from __future__ import annotations
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
+from http import HTTPStatus
from typing import TYPE_CHECKING
import pytest
+import requests
from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
@@ -81,6 +83,8 @@ _LOG_FETCH_TIMEOUT = 60
_ANNOTATION_DAG_ID = "java_annotation_example"
_XCOM_CASTING_DAG_ID = "java_xcom_casting_example"
_SCALA_SPARK_DAG_ID = "scala_spark_example"
+_VARIABLE_WRITE_DAG_ID = "java_variable_write"
+_VARIABLE_WRITE_DESCRIPTION = "written by the Java SDK e2e test"
@dataclass
@@ -179,6 +183,12 @@ def scala_spark_example_run() -> _CompletedRun:
return _trigger_and_wait_for_dag(_SCALA_SPARK_DAG_ID, _SPARK_TASK_TIMEOUT)
[email protected](scope="module")
+def variable_write_run() -> _CompletedRun:
+ """Trigger the variable write Dag once for all of its assertions."""
+ return _trigger_and_wait_for_dag(_VARIABLE_WRITE_DAG_ID,
_JAVA_TASK_TIMEOUT)
+
+
class TestJavaSDKAnnotationExample:
"""Verify the annotation-based Java SDK example executes correctly."""
@@ -322,6 +332,51 @@ class TestJavaSDKUninstantiableTask:
)
+class TestJavaSDKVariableWrite:
+ """Verify a Java task can write and delete Airflow Variables.
+
+ The task lives in the java-test-bundle fixture project (served on the
+ dedicated "java-test" queue). It stores its own run id in
+ ``java_e2e_variable`` and deletes a scratch Variable it has just written,
+ so both paths go through the supervisor and the Task Execution API.
+ """
+
+ def test_variable_written_by_java_task_is_readable(self,
variable_write_run: _CompletedRun):
+ """The value and description set from Java are visible through the
REST API."""
+ ti = variable_write_run.get_task_instance("write_and_delete")
+ assert ti.get("state") == "success", (
+ "Java 'write_and_delete' task did not succeed.\n"
+ f" task state : {ti.get('state')!r}\n"
+ f" dag state : {variable_write_run.state!r}\n"
+ f" all tasks : {variable_write_run.ti_states}"
+ )
+
+ variable = variable_write_run.client.get_variable("java_e2e_variable")
+ assert variable.get("value") == variable_write_run.run_id, (
+ f"java_e2e_variable should hold this run's id
{variable_write_run.run_id!r}, got {variable!r}"
+ )
+ assert variable.get("description") == _VARIABLE_WRITE_DESCRIPTION, (
+ f"java_e2e_variable should carry the description set from Java,
got {variable!r}"
+ )
+
+ def test_scratch_variable_deleted_by_java_task_is_gone(self,
variable_write_run: _CompletedRun):
+ """A Variable written and then deleted from Java no longer exists."""
+ ti = variable_write_run.get_task_instance("write_and_delete")
+ assert ti.get("state") == "success", (
+ "Java 'write_and_delete' task did not succeed.\n"
+ f" task state : {ti.get('state')!r}\n"
+ f" dag state : {variable_write_run.state!r}\n"
+ f" all tasks : {variable_write_run.ti_states}"
+ )
+
+ with pytest.raises(requests.HTTPError) as excinfo:
+ variable_write_run.client.get_variable("java_e2e_scratch")
+ assert excinfo.value.response.status_code == HTTPStatus.NOT_FOUND, (
+ f"java_e2e_scratch should have been deleted by the Java task, "
+ f"got HTTP {excinfo.value.response.status_code}"
+ )
+
+
class TestJavaSDKXComCastingExample:
"""Verify numeric XCom values are cast across declared Java types."""
diff --git a/java-sdk/README.md b/java-sdk/README.md
index ef998487127..372426e31dc 100644
--- a/java-sdk/README.md
+++ b/java-sdk/README.md
@@ -608,7 +608,7 @@ prek hook regenerate it.
| capability: `task-logging` | MUST | ✓ | 3.3 | SLF4J + JPL bridged to the
task log |
| capability: `xcom-read-write` | MUST | ✓ | 3.3 | |
| capability: `connection-read` | MUST | ✓ | 3.3 | |
-| capability: `variable-read-write` | MUST | ✗ | – | getVariable only; no
write over the comm socket yet |
+| capability: `variable-read-write` | MUST | ✓ | 3.3 | |
| 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 | ✗ | – | no task-facing state-store
API yet |
diff --git a/java-sdk/capabilities.yaml b/java-sdk/capabilities.yaml
index a65a3f343c7..50dc38c9528 100644
--- a/java-sdk/capabilities.yaml
+++ b/java-sdk/capabilities.yaml
@@ -70,8 +70,8 @@ capabilities:
supported: true
since: "3.3"
variable-read-write:
- supported: false
- note: "getVariable only; no write over the comm socket yet"
+ supported: true
+ since: "3.3"
self-contained-bundle:
supported: true
since: "3.3"
diff --git a/java-sdk/sdk/module.md b/java-sdk/sdk/module.md
index b1189de7cbe..fd402bb048b 100644
--- a/java-sdk/sdk/module.md
+++ b/java-sdk/sdk/module.md
@@ -47,7 +47,7 @@ meaning of each dimension is defined in the
| capability: `task-logging` | MUST | ✓ | 3.3 | SLF4J + JPL bridged to the
task log |
| capability: `xcom-read-write` | MUST | ✓ | 3.3 | |
| capability: `connection-read` | MUST | ✓ | 3.3 | |
-| capability: `variable-read-write` | MUST | ✗ | – | getVariable only; no
write over the comm socket yet |
+| capability: `variable-read-write` | MUST | ✓ | 3.3 | |
| 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 | ✗ | – | no task-facing state-store
API yet |
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt
index 59aa7832a37..be7e582eef7 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt
@@ -95,6 +95,31 @@ class Client internal constructor(
*/
fun getVariable(key: String): Any? = impl.getVariable(key).value
+ /**
+ * Stores an Airflow variable, replacing any existing value.
+ *
+ * The value is stored as-is. Serialize structured data (for example to
+ * JSON) before storing it.
+ *
+ * @param key Variable key.
+ * @param value Value to store.
+ * @param description Description of the variable.
+ * @throws ApiError if the API call fails.
+ */
+ @JvmOverloads fun setVariable(
+ key: String,
+ value: String,
+ description: String? = null,
+ ) = impl.setVariable(key = key, value = value, description = description)
+
+ /**
+ * Deletes an Airflow variable.
+ *
+ * @param key Variable key.
+ * @throws ApiError if the API call fails.
+ */
+ fun deleteVariable(key: String) = impl.deleteVariable(key)
+
/**
* Reads an XCom value pushed by another task.
*
diff --git
a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt
index 2a2610efddb..d01b151dd36 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt
@@ -21,9 +21,12 @@ package org.apache.airflow.sdk.execution
import kotlinx.coroutines.runBlocking
import org.apache.airflow.sdk.execution.comm.ConnectionResult
+import org.apache.airflow.sdk.execution.comm.DeleteVariable
import org.apache.airflow.sdk.execution.comm.GetConnection
import org.apache.airflow.sdk.execution.comm.GetVariable
import org.apache.airflow.sdk.execution.comm.GetXCom
+import org.apache.airflow.sdk.execution.comm.OKResponse
+import org.apache.airflow.sdk.execution.comm.PutVariable
import org.apache.airflow.sdk.execution.comm.SetXCom
import org.apache.airflow.sdk.execution.comm.VariableResult
import org.apache.airflow.sdk.execution.comm.XComResult
@@ -46,6 +49,14 @@ interface Client {
fun getVariable(key: String): VariableResult
+ fun setVariable(
+ key: String,
+ value: String,
+ description: String?,
+ )
+
+ fun deleteVariable(key: String)
+
fun getXCom(
key: String,
dagId: String,
@@ -89,6 +100,24 @@ class CoordinatorClient(
exec.communicate<VariableResult>(GetVariable().also { it.key = key })
}
+ override fun setVariable(
+ key: String,
+ value: String,
+ description: String?,
+ ) {
+ val message =
+ PutVariable().also {
+ it.key = key
+ it.value = value
+ it.description = description
+ }
+ runBlocking { exec.communicate<Unit>(message) }
+ }
+
+ override fun deleteVariable(key: String) {
+ runBlocking { exec.communicate<OKResponse>(DeleteVariable().also { it.key
= key }) }
+ }
+
override fun setXCom(
key: String,
value: Any,
diff --git
a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Frame.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Frame.kt
index 035cfb89114..ca61ddd28c2 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Frame.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Frame.kt
@@ -19,12 +19,14 @@
package org.apache.airflow.sdk.execution
+import com.fasterxml.jackson.annotation.JsonInclude
import com.fasterxml.jackson.databind.DeserializationFeature
import com.fasterxml.jackson.databind.ObjectMapper
import com.fasterxml.jackson.databind.SerializationFeature
import com.fasterxml.jackson.databind.util.StdDateFormat
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule
import org.apache.airflow.sdk.execution.comm.Discriminator
+import org.apache.airflow.sdk.execution.comm.PutVariable
import org.msgpack.core.MessagePack
import org.msgpack.core.MessageUnpacker
import org.msgpack.core.buffer.ArrayBufferInput
@@ -37,6 +39,17 @@ data class RawFrame(
val rawError: Any?,
)
+/**
+ * jsonschema2pojo stamps every model with `@JsonInclude(NON_NULL)`, so
+ * Jackson drops null fields from the wire map. The supervisor requires these
+ * fields to be present. When one is missing it fails validation, logs the
+ * frame and never replies, so the caller blocks forever.
+ */
+private val REQUIRED_NULLABLE_REQUESTS = setOf(PutVariable::class.java)
+
+@JsonInclude(JsonInclude.Include.ALWAYS)
+private abstract class KeepNullFields
+
object Frame {
internal const val MAX_FRAME_LENGTH = 0xFFFF_FFFFL
@@ -47,6 +60,7 @@ object Frame {
registerModule(JavaTimeModule())
registerModule(TimestampToJavaOffsetDateTimeModule())
setDateFormat(StdDateFormat().withColonInTimeZone(true))
+ REQUIRED_NULLABLE_REQUESTS.forEach { addMixIn(it,
KeepNullFields::class.java) }
}
fun encodeRequest(
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt
b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt
index 6428ec224ee..bd39991abb2 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt
@@ -34,6 +34,14 @@ private class FakeTransport(
override fun getVariable(key: String): VariableResult = throw
NotImplementedError()
+ override fun setVariable(
+ key: String,
+ value: String,
+ description: String?,
+ ) = throw NotImplementedError()
+
+ override fun deleteVariable(key: String) = throw NotImplementedError()
+
override fun getXCom(
key: String,
dagId: String,
diff --git
a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt
b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt
index ef0459b0922..0c2ef844f6c 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt
@@ -35,6 +35,7 @@ import org.junit.jupiter.api.DisplayName
import org.junit.jupiter.api.Test
import org.junit.jupiter.api.Timeout
import org.msgpack.core.MessagePack
+import org.msgpack.core.buffer.ArrayBufferInput
import java.io.ByteArrayOutputStream
import java.time.OffsetDateTime
import java.time.ZoneOffset
@@ -136,6 +137,84 @@ class CommsTest {
return out.toByteArray()
}
+ private fun okResponseFrame(id: Int): ByteArray {
+ val out = ByteArrayOutputStream()
+ MessagePack.newDefaultPacker(out).use { packer ->
+ packer.packArrayHeader(3)
+ packer.packInt(id)
+ packer.packMapHeader(2)
+ packer.packString("type")
+ packer.packString("OKResponse")
+ packer.packString("ok")
+ packer.packBoolean(true)
+ packer.packNil()
+ }
+ return out.toByteArray()
+ }
+
+ // The supervisor answers a request that yields no result with a bare `[id]`
frame.
+ private fun emptyResponseFrame(id: Int): ByteArray {
+ val out = ByteArrayOutputStream()
+ MessagePack.newDefaultPacker(out).use { packer ->
+ packer.packArrayHeader(1)
+ packer.packInt(id)
+ }
+ return out.toByteArray()
+ }
+
+ private fun errorResponseFrame(id: Int): ByteArray {
+ val out = ByteArrayOutputStream()
+ MessagePack.newDefaultPacker(out).use { packer ->
+ packer.packArrayHeader(3)
+ packer.packInt(id)
+ packer.packNil()
+ packer.packMapHeader(2)
+ packer.packString("type")
+ packer.packString("ErrorResponse")
+ packer.packString("detail")
+ packer.packMapHeader(1)
+ packer.packString("status_code")
+ packer.packInt(500)
+ }
+ return out.toByteArray()
+ }
+
+ private fun readRequest(fromClient: ByteChannel): RawFrame =
+ runBlocking {
+ val prefix = fromClient.readByteArray(4)
+ val payload =
fromClient.readByteArray(Frame.parseLengthPrefix(prefix).toInt())
+ Frame.decodeRaw(ArrayBufferInput(payload))
+ }
+
+ /**
+ * Run one call on the public client against a fake supervisor that answers
+ * its single request with [response]. Returns the raw request body as sent
+ * on the wire and the exception the call threw, if any.
+ */
+ private fun roundTrip(
+ response: (Int) -> ByteArray,
+ call: (PublicClient) -> Unit,
+ ): Pair<Map<*, *>, Throwable?> {
+ val toClient = ByteChannel(autoFlush = true)
+ val fromClient = ByteChannel(autoFlush = true)
+ val comm = CoordinatorComm(toClient, fromClient)
+ val client = PublicClient(StartupDetails(), CoordinatorClient(comm))
+
+ val requests = ConcurrentLinkedQueue<RawFrame>()
+ val server =
+ Thread {
+ val request = readRequest(fromClient)
+ requests.add(request)
+ runBlocking { toClient.writeFrame(response(request.id)) }
+ }
+ server.start()
+ val failure = runCatching { call(client) }.exceptionOrNull()
+ server.join()
+ comm.close()
+
+ return (requests.single().rawBody as Map<*, *>) to failure
+ }
+
private suspend fun ByteChannel.writeFrame(payload: ByteArray) {
writeByteArray(Frame.lengthPrefix(payload.size.toUInt()))
writeByteArray(payload)
@@ -263,6 +342,50 @@ class CommsTest {
comm.close()
}
+ @Test
+ @DisplayName("setVariable keeps a null description on the wire so the
supervisor accepts the request")
+ @Timeout(value = 30, unit = TimeUnit.SECONDS)
+ fun setVariableKeepsNullDescriptionOnTheWire() {
+ val (body, failure) = roundTrip(::emptyResponseFrame) {
it.setVariable("k", "v") }
+
+ Assertions.assertNull(failure, "setVariable should return normally on an
empty response, got $failure")
+ Assertions.assertEquals("PutVariable", body["type"])
+ Assertions.assertEquals("k", body["key"])
+ Assertions.assertEquals("v", body["value"])
+ Assertions.assertTrue(body.containsKey("description"), "description must
be sent even when null: $body")
+ Assertions.assertNull(body["description"])
+ }
+
+ @Test
+ @DisplayName("setVariable sends the description when one is given")
+ @Timeout(value = 30, unit = TimeUnit.SECONDS)
+ fun setVariableSendsDescription() {
+ val (body, failure) = roundTrip(::emptyResponseFrame) {
it.setVariable("k", "v", "why") }
+
+ Assertions.assertNull(failure, "setVariable should return normally on an
empty response, got $failure")
+ Assertions.assertEquals("why", body["description"])
+ }
+
+ @Test
+ @DisplayName("deleteVariable sends the key and accepts the supervisor's OK
response")
+ @Timeout(value = 30, unit = TimeUnit.SECONDS)
+ fun deleteVariableAcceptsOkResponse() {
+ val (body, failure) = roundTrip(::okResponseFrame) {
it.deleteVariable("k") }
+
+ Assertions.assertNull(failure, "deleteVariable should return normally on
OKResponse, got $failure")
+ Assertions.assertEquals("DeleteVariable", body["type"])
+ Assertions.assertEquals("k", body["key"])
+ }
+
+ @Test
+ @DisplayName("deleteVariable raises ApiError when the supervisor reports an
error")
+ @Timeout(value = 30, unit = TimeUnit.SECONDS)
+ fun deleteVariableRaisesApiErrorOnErrorResponse() {
+ val (_, failure) = roundTrip(::errorResponseFrame) {
it.deleteVariable("k") }
+
+ Assertions.assertInstanceOf(ApiError::class.java, failure)
+ }
+
@Test
@DisplayName("Should fail a pending call when the coordinator socket closes")
@Timeout(value = 30, unit = TimeUnit.SECONDS)
diff --git
a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt
b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt
index 1b1fd288ed7..5d16eda1ccb 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt
@@ -214,6 +214,14 @@ class TaskTest {
override fun getVariable(key: String) = throw
UnsupportedOperationException("not used in test")
+ override fun setVariable(
+ key: String,
+ value: String,
+ description: String?,
+ ): Unit = throw UnsupportedOperationException("not used in test")
+
+ override fun deleteVariable(key: String): Unit = throw
UnsupportedOperationException("not used in test")
+
override fun getXCom(
key: String,
dagId: String,
