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 a1c2c50f4e4 Add SetVariable and DeleteVariable to the Go SDK task
client (#73084)
a1c2c50f4e4 is described below
commit a1c2c50f4e4264f8347426e6d9ed155ce632cb3d
Author: Henry Chen <[email protected]>
AuthorDate: Fri Sep 18 16:54:45 2026 +0800
Add SetVariable and DeleteVariable to the Go SDK task client (#73084)
* Add SetVariable and DeleteVariable to the Go SDK task client
Go tasks could read Airflow Variables but not write or delete them, so the
Go SDK did not meet the variable-read-write capability that language SDKs
must support. The supervisor already handles PutVariable and DeleteVariable,
so Go tasks only lacked a way to send them.
* Document that any secrets backend overrides Go SDK Variable writes on read
The note only mentioned AIRFLOW_VAR_* environment variables, but reads go
through every secrets backend configured on the API server, so any of them
can hide a value written from Go. The Java SDK docs already describe this
precedence the same way.
---
.../authoring-and-scheduling/language-sdks/go.rst | 22 +++-
.../go_sdk_tests/test_go_sdk_variable_write.py | 81 +++++++++++++++
.../cmd/airflow-go-pack/pack_integration_test.go | 3 +
go-sdk/dags/go_examples.py | 19 +++-
.../bundle/concurrentxcom/concurrentxcom_test.go | 8 ++
go-sdk/example/bundle/main.go | 7 ++
go-sdk/example/bundle/main_test.go | 10 ++
.../example/bundle/variablewrite/variablewrite.go | 51 ++++++++++
go-sdk/pkg/execution/client.go | 19 ++++
go-sdk/pkg/execution/client_test.go | 113 +++++++++++++++++++++
go-sdk/sdk/sdk.go | 19 +++-
11 files changed, 344 insertions(+), 8 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst
b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst
index 067d2d0bc6a..42b1ecf6213 100644
--- a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst
+++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst
@@ -232,8 +232,8 @@ The ``sdk.Client`` surface
``sdk.Client`` composes three smaller interfaces, so a task can depend on just
one:
-* ``VariableClient`` - ``GetVariable`` (returns the Variable as a string) and
``UnmarshalJSONVariable``
- (decodes a JSON Variable into a pointer you provide).
+* ``VariableClient`` - ``GetVariable`` (returns the Variable as a string),
``UnmarshalJSONVariable``
+ (decodes a JSON Variable into a pointer you provide), ``SetVariable``, and
``DeleteVariable``.
* ``ConnectionClient`` - ``GetConnection``, returning a ``Connection`` with
fields ``ID``, ``Type``,
``Host``, ``Port``, ``Login``, ``Password``, ``Path``, ``Extra`` (a
``map[string]any``), plus a
``GetURI()`` helper.
@@ -242,6 +242,24 @@ The ``sdk.Client`` surface
``GetXCom`` returns the stored value as an ``any``; see :ref:`go-sdk/types`
for how the stored JSON maps to
Go types.
+``SetVariable`` stores the value as a string, so encode structured data (for
example with ``json.Marshal``)
+before storing it.
+
+.. code-block:: go
+
+ if err := client.SetVariable(ctx, "process_threshold", "42", "Rows above
this count take the slow path"); err != nil {
+ return err
+ }
+ if err := client.DeleteVariable(ctx, "legacy_threshold"); err != nil {
+ return err
+ }
+
+.. 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`` with an empty
+ description clears any existing description.
+
Not-found lookups return sentinel errors - ``VariableNotFound``,
``ConnectionNotFound``, ``XComNotFound`` -
so you can branch on a missing value with ``errors.Is`` rather than parsing an
error string.
diff --git
a/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_variable_write.py
b/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_variable_write.py
new file mode 100644
index 00000000000..b4db9bd8482
--- /dev/null
+++
b/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_variable_write.py
@@ -0,0 +1,81 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""E2E test for the Go SDK ``variable_write_dag`` example.
+
+``write_and_delete_variable`` (Go, ``go-sdk/example/bundle/variablewrite``)
stores
+its run id in ``go_e2e_variable`` and deletes a scratch Variable it has just
+written, so both paths go through the supervisor and the Task Execution API.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from http import HTTPStatus
+
+import pytest
+import requests
+
+from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient
+
+_GO_TASK_TIMEOUT = 300
+
+_DAG_ID = "variable_write_dag"
+_TASK_ID = "write_and_delete_variable"
+
+
+@dataclass
+class _CompletedRun:
+ """The single ``variable_write_dag`` run shared across this module's
tests."""
+
+ client: AirflowClient
+ run_id: str
+ state: str
+ ti_states: dict[str, str]
+
+
[email protected](scope="module")
+def completed_run() -> _CompletedRun:
+ """Trigger ``variable_write_dag`` once and wait for it to finish."""
+ client = AirflowClient()
+ resp = client.trigger_dag(_DAG_ID, json={"logical_date":
datetime.now(timezone.utc).isoformat()})
+ run_id = resp["dag_run_id"]
+ state = client.wait_for_dag_run(dag_id=_DAG_ID, run_id=run_id,
timeout=_GO_TASK_TIMEOUT)
+ ti_resp = client.get_task_instances(dag_id=_DAG_ID, run_id=run_id)
+ ti_states = {ti["task_id"]: ti.get("state") for ti in
ti_resp.get("task_instances", [])}
+ return _CompletedRun(client=client, run_id=run_id, state=state,
ti_states=ti_states)
+
+
+def test_task_succeeded(completed_run: _CompletedRun):
+ assert completed_run.state == "success", (
+ f"expected the run to succeed; got {completed_run.state!r}. task
states: {completed_run.ti_states}"
+ )
+ assert completed_run.ti_states.get(_TASK_ID) == "success",
completed_run.ti_states
+
+
+def test_variable_written_by_go_task_is_readable(completed_run: _CompletedRun):
+ variable = completed_run.client.get_variable("go_e2e_variable")
+ assert variable.get("value") == completed_run.run_id, (
+ f"go_e2e_variable should hold this run's id {completed_run.run_id!r},
got {variable!r}"
+ )
+ assert variable.get("description") == "written by the Go SDK e2e test",
variable
+
+
+def test_scratch_variable_deleted_by_go_task_is_gone(completed_run:
_CompletedRun):
+ with pytest.raises(requests.HTTPError) as excinfo:
+ completed_run.client.get_variable("go_e2e_scratch")
+ assert excinfo.value.response.status_code == HTTPStatus.NOT_FOUND
diff --git a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go
b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go
index ace867461e2..f064935d6a5 100644
--- a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go
+++ b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go
@@ -166,6 +166,9 @@ dags:
- "via_flat_map"
- "via_struct_map"
- "via_plain_map"
+ variable_write_dag:
+ tasks:
+ - "write_and_delete_variable"
`
assert.Equal(t, expectedManifest, string(metadata))
diff --git a/go-sdk/dags/go_examples.py b/go-sdk/dags/go_examples.py
index d6e768a2c37..46367234e4b 100644
--- a/go-sdk/dags/go_examples.py
+++ b/go-sdk/dags/go_examples.py
@@ -17,11 +17,12 @@
"""
Python stub Dags mirroring the Go SDK example bundle
(``go-sdk/example/bundle``).
-Three Dags, all backed by the same Go bundle: ``simple_dag``
(extract/transform/
+Four Dags, all backed by the same Go bundle: ``simple_dag`` (extract/transform/
load, below), ``concurrent_xcom_dag`` (one ``pull_xcoms_concurrently`` task
-timing sequential vs goroutine XCom pulls), and ``taskflow_binding_dag`` (one
+timing sequential vs goroutine XCom pulls), ``taskflow_binding_dag`` (one
task per shape of the TaskFlow argument-binding surface; see its Dag function
-below).
+below), and ``variable_write_dag`` (one ``write_and_delete_variable`` task that
+writes and deletes Airflow Variables).
``simple_dag`` sandwiches the Go tasks between two native Python tasks so the
run exercises XCom across the language boundary, the same way
@@ -208,3 +209,15 @@ def taskflow_binding_dag():
taskflow_binding_dag()
+
+
[email protected](queue="golang")
+def write_and_delete_variable(): ...
+
+
+@dag(dag_id="variable_write_dag")
+def variable_write_dag():
+ write_and_delete_variable()
+
+
+variable_write_dag()
diff --git a/go-sdk/example/bundle/concurrentxcom/concurrentxcom_test.go
b/go-sdk/example/bundle/concurrentxcom/concurrentxcom_test.go
index 658a5c3c14c..f087fbbb3da 100644
--- a/go-sdk/example/bundle/concurrentxcom/concurrentxcom_test.go
+++ b/go-sdk/example/bundle/concurrentxcom/concurrentxcom_test.go
@@ -75,6 +75,14 @@ func (m *mockXComClient) UnmarshalJSONVariable(ctx
context.Context, key string,
panic("unimplemented")
}
+func (m *mockXComClient) SetVariable(ctx context.Context, key, value,
description string) error {
+ panic("unimplemented")
+}
+
+func (m *mockXComClient) DeleteVariable(ctx context.Context, key string) error
{
+ panic("unimplemented")
+}
+
func (m *mockXComClient) GetConnection(ctx context.Context, connID string)
(sdk.Connection, error) {
panic("unimplemented")
}
diff --git a/go-sdk/example/bundle/main.go b/go-sdk/example/bundle/main.go
index 6d656988900..e01e29a6447 100644
--- a/go-sdk/example/bundle/main.go
+++ b/go-sdk/example/bundle/main.go
@@ -28,6 +28,7 @@ import (
"github.com/apache/airflow/go-sdk/bundle/bundlev1/bundlev1server"
"github.com/apache/airflow/go-sdk/example/bundle/concurrentxcom"
"github.com/apache/airflow/go-sdk/example/bundle/taskflowbinding"
+ "github.com/apache/airflow/go-sdk/example/bundle/variablewrite"
"github.com/apache/airflow/go-sdk/sdk"
)
@@ -58,6 +59,12 @@ func (m *myBundle) RegisterDags(dagbag v1.Registry) error {
bindingDag.AddTaskWithName("via_struct_map",
taskflowbinding.ViaStructMap)
bindingDag.AddTaskWithName("via_plain_map", taskflowbinding.ViaPlainMap)
+ variableWriteDag := dagbag.AddDag("variable_write_dag")
+ variableWriteDag.AddTaskWithName(
+ "write_and_delete_variable",
+ variablewrite.WriteAndDeleteVariable,
+ )
+
return nil
}
diff --git a/go-sdk/example/bundle/main_test.go
b/go-sdk/example/bundle/main_test.go
index 16be1e22e92..0503495d90e 100644
--- a/go-sdk/example/bundle/main_test.go
+++ b/go-sdk/example/bundle/main_test.go
@@ -47,6 +47,16 @@ func (m *mockVars) UnmarshalJSONVariable(ctx
context.Context, key string, pointe
panic("unimplemented")
}
+// SetVariable implements sdk.VariableClient.
+func (m *mockVars) SetVariable(ctx context.Context, key, value, description
string) error {
+ panic("unimplemented")
+}
+
+// DeleteVariable implements sdk.VariableClient.
+func (m *mockVars) DeleteVariable(ctx context.Context, key string) error {
+ panic("unimplemented")
+}
+
var _ sdk.VariableClient = (*mockVars)(nil)
func Test_transform(t *testing.T) {
diff --git a/go-sdk/example/bundle/variablewrite/variablewrite.go
b/go-sdk/example/bundle/variablewrite/variablewrite.go
new file mode 100644
index 00000000000..6094605a45c
--- /dev/null
+++ b/go-sdk/example/bundle/variablewrite/variablewrite.go
@@ -0,0 +1,51 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Package variablewrite holds the write_and_delete_variable task, which
+// round-trips Airflow Variables through the supervisor.
+package variablewrite
+
+import (
+ "fmt"
+
+ "github.com/apache/airflow/go-sdk/sdk"
+)
+
+const (
+ // WrittenKey holds the run id of the Dag run that wrote it, so a
reader can
+ // tell this run's write apart from a previous one.
+ WrittenKey = "go_e2e_variable"
+ // WrittenDescription is stored alongside WrittenKey.
+ WrittenDescription = "written by the Go SDK e2e test"
+ // ScratchKey is written and then deleted within the same task.
+ ScratchKey = "go_e2e_scratch"
+)
+
+// WriteAndDeleteVariable stores the current run id under WrittenKey, then
+// writes and deletes ScratchKey to exercise the delete path.
+func WriteAndDeleteVariable(ctx sdk.TIRunContext, client sdk.VariableClient)
error {
+ if err := client.SetVariable(ctx, WrittenKey, ctx.DagRun().RunID,
WrittenDescription); err != nil {
+ return fmt.Errorf("setting %s: %w", WrittenKey, err)
+ }
+ if err := client.SetVariable(ctx, ScratchKey, "scratch", ""); err !=
nil {
+ return fmt.Errorf("setting %s: %w", ScratchKey, err)
+ }
+ if err := client.DeleteVariable(ctx, ScratchKey); err != nil {
+ return fmt.Errorf("deleting %s: %w", ScratchKey, err)
+ }
+ return nil
+}
diff --git a/go-sdk/pkg/execution/client.go b/go-sdk/pkg/execution/client.go
index 5fc962cf8fd..b79c6843ab9 100644
--- a/go-sdk/pkg/execution/client.go
+++ b/go-sdk/pkg/execution/client.go
@@ -125,6 +125,25 @@ func (c *CoordinatorClient) UnmarshalJSONVariable(
return json.Unmarshal([]byte(val), pointer)
}
+// SetVariable asks the supervisor to store a variable value.
+func (c *CoordinatorClient) SetVariable(
+ ctx context.Context,
+ key, value, description string,
+) error {
+ msg := genmodels.PutVariable{Key: key, Value: value}
+ if description != "" {
+ msg.Description = description
+ }
+ _, err := c.comm.Communicate(ctx, msg)
+ return err
+}
+
+// DeleteVariable asks the supervisor to delete a variable.
+func (c *CoordinatorClient) DeleteVariable(ctx context.Context, key string)
error {
+ _, err := c.comm.Communicate(ctx, genmodels.DeleteVariable{Key: key})
+ return err
+}
+
// GetConnection requests a connection from the supervisor.
func (c *CoordinatorClient) GetConnection(
ctx context.Context,
diff --git a/go-sdk/pkg/execution/client_test.go
b/go-sdk/pkg/execution/client_test.go
index 96d7020b5d2..cd72641c40a 100644
--- a/go-sdk/pkg/execution/client_test.go
+++ b/go-sdk/pkg/execution/client_test.go
@@ -161,6 +161,119 @@ func TestCoordinatorClientErrorPassThrough(t *testing.T) {
assert.Equal(t, "API_SERVER_ERROR", apiErr.Err)
}
+// TestCoordinatorClientSetVariable verifies the PutVariable frame always
+// carries description, sending null when none is given: the supervisor
+// validates PutVariable with a required description field.
+func TestCoordinatorClientSetVariable(t *testing.T) {
+ tests := []struct {
+ name string
+ description string
+ wantDescription any
+ }{
+ {
+ name: "description is sent",
+ description: "row threshold",
+ wantDescription: "row threshold",
+ },
+ {name: "empty description is sent as null", description: "",
wantDescription: nil},
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ responsePayload := encodeResponseFrame(t, 0, nil, nil)
+ var responseBuf bytes.Buffer
+ require.NoError(t, writeFrame(&responseBuf,
responsePayload))
+
+ var requestBuf bytes.Buffer
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ comm := NewCoordinatorComm(&responseBuf, &requestBuf,
logger)
+ client := NewCoordinatorClient(comm)
+
+ require.NoError(
+ t,
+ client.SetVariable(context.Background(),
"my_key", "42", tc.description),
+ )
+
+ sent, err := readFrame(&requestBuf)
+ require.NoError(t, err)
+ assert.Equal(t, map[string]any{
+ "type": "PutVariable",
+ "key": "my_key",
+ "value": "42",
+ "description": tc.wantDescription,
+ }, rawToMap(t, sent.Body))
+ })
+ }
+}
+
+// TestCoordinatorClientDeleteVariable verifies the DeleteVariable frame sent
+// to the supervisor.
+func TestCoordinatorClientDeleteVariable(t *testing.T) {
+ responsePayload := encodeResponseFrame(
+ t,
+ 0,
+ map[string]any{"type": "OKResponse", "ok": true},
+ nil,
+ )
+ var responseBuf bytes.Buffer
+ require.NoError(t, writeFrame(&responseBuf, responsePayload))
+
+ var requestBuf bytes.Buffer
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ comm := NewCoordinatorComm(&responseBuf, &requestBuf, logger)
+ client := NewCoordinatorClient(comm)
+
+ require.NoError(t, client.DeleteVariable(context.Background(),
"my_key"))
+
+ sent, err := readFrame(&requestBuf)
+ require.NoError(t, err)
+ assert.Equal(t, map[string]any{
+ "type": "DeleteVariable",
+ "key": "my_key",
+ }, rawToMap(t, sent.Body))
+}
+
+// TestCoordinatorClientVariableWriteErrors verifies SetVariable and
+// DeleteVariable surface a supervisor ErrorResponse to the task.
+func TestCoordinatorClientVariableWriteErrors(t *testing.T) {
+ tests := []struct {
+ name string
+ call func(client *CoordinatorClient) error
+ }{
+ {
+ name: "SetVariable",
+ call: func(client *CoordinatorClient) error {
+ return client.SetVariable(context.Background(),
"my_key", "v", "")
+ },
+ },
+ {
+ name: "DeleteVariable",
+ call: func(client *CoordinatorClient) error {
+ return
client.DeleteVariable(context.Background(), "my_key")
+ },
+ },
+ }
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ responsePayload := encodeResponseFrame(t, 0, nil,
map[string]any{
+ "type": "ErrorResponse",
+ "error": "API_SERVER_ERROR",
+ "detail": map[string]any{"status_code": 403},
+ })
+ var responseBuf bytes.Buffer
+ require.NoError(t, writeFrame(&responseBuf,
responsePayload))
+
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ comm := NewCoordinatorComm(&responseBuf, io.Discard,
logger)
+ client := NewCoordinatorClient(comm)
+
+ var apiErr *ApiError
+ require.ErrorAs(t, tc.call(client), &apiErr)
+ assert.Equal(t, "API_SERVER_ERROR", apiErr.Err)
+ })
+ }
+}
+
// TestCoordinatorClientGetConnectionPreservesEmptyCredentials verifies the
// coordinator client forwards an explicitly empty login/password as a
// pointer-to-"" on sdk.Connection rather than nil. Connections that use
diff --git a/go-sdk/sdk/sdk.go b/go-sdk/sdk/sdk.go
index 746b04d7db7..e378527b01d 100644
--- a/go-sdk/sdk/sdk.go
+++ b/go-sdk/sdk/sdk.go
@@ -35,7 +35,7 @@ const (
XComReturnValueKey = "return_value"
)
-// VariableClient reads Airflow Variables.
+// VariableClient reads, writes, and deletes Airflow Variables.
//
// Go has no function overloading, so the "give me the raw string" and
// "give me a decoded struct" cases are split into two methods rather
@@ -67,6 +67,19 @@ type VariableClient interface {
//
// pointer must be a non-nil pointer, as required by encoding/json.
UnmarshalJSONVariable(ctx context.Context, key string, pointer any)
error
+
+ // SetVariable stores value under key, creating the Variable or
replacing
+ // an existing one. An empty description is sent as null, which clears
any
+ // description the Variable already had.
+ //
+ // The value is stored as-is: encode structured data (for example with
+ // json.Marshal) before storing it. A value supplied by a secrets
backend
+ // (for example an AIRFLOW_VAR_<KEY> environment variable) still takes
+ // precedence over the stored value when the Variable is read back.
+ SetVariable(ctx context.Context, key, value, description string) error
+
+ // DeleteVariable removes the Variable stored under key.
+ DeleteVariable(ctx context.Context, key string) error
}
// ConnectionClient reads Airflow Connections.
@@ -107,8 +120,8 @@ type XComClient interface {
PushXCom(ctx context.Context, ti TaskInstance, key string, value any)
error
}
-// Client is the full task-facing API: read Variables and Connections, and
-// read/write XCom. A task that declares an sdk.Client parameter is handed one
+// Client is the full task-facing API: read/write Variables, read Connections,
+// and read/write XCom. A task that declares an sdk.Client parameter is handed
one
// by the runtime. If a task needs only one capability, ask for the narrower
// VariableClient, ConnectionClient, or XComClient instead.
type Client interface {
