Go SDK: Refresh Go SDK execution API tokens

go-sdk/pkg/api/client.go

@@ -20,13 +20,16 @@ package api
 import (
 	"fmt"
 	"maps"
+	"sync/atomic"
 
 	"github.com/google/uuid"
 	"resty.dev/v3"
 )
 
 const API_VERSION = "2025-05-20"
 
+const refreshedAPITokenHeader = "Refreshed-API-Token"
+
 //go:generate -command openapi-gen go run github.com/ashb/oapi-resty-codegen@latest --config oapi-codegen.yml
 
 //go:generate openapi-gen https://airflow.staged.apache.org/schemas/execution-api/2025-05-20.json
@@ -59,10 +62,19 @@ func (c *Client) WithBearerToken(token string) (ClientInterface, error) {
 	rc.SetDebug(c.Client.IsDebug())
 	rc.SetLogger(c.Client.Logger())
 
-	// We don't use SetAuthToken/SetAuthScheme, as that produces a (valid, but annoying) warning about using Auth
-	// over HTTP: "Using sensitive credentials in HTTP mode is not secure." It's a time-limited-token though, so we
-	// can reasonably ignore that here and setting the header directly bypasses that
-	rc.SetHeader("Authorization", fmt.Sprintf("Bearer %s", token))
+	// Keep the token outside Resty's shared header map because task API calls and heartbeats use this client concurrently.
+	var authorization atomic.Value
+	authorization.Store(fmt.Sprintf("Bearer %s", token))
+	rc.AddRequestMiddleware(func(_ *resty.Client, req *resty.Request) error {
+		req.Header.Set("Authorization", authorization.Load().(string))
+		return nil
+	})
+	rc.AddResponseMiddleware(func(_ *resty.Client, response *resty.Response) error {
+		if refreshedToken := response.Header().Get(refreshedAPITokenHeader); refreshedToken != "" {
+			authorization.Store(fmt.Sprintf("Bearer %s", refreshedToken))
+		}
+		return nil
+	})
 
 	opts := []ClientOption{
 		WithClient(rc),

go-sdk/pkg/api/client_test.go

@@ -0,0 +1,85 @@
+// 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 api
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWithBearerTokenUsesRefreshedAPIToken(t *testing.T) {
+	var requestCount atomic.Int32
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requestNumber := requestCount.Add(1)
+		expectedToken := "execution-token"
+		if requestNumber == 1 {
+			expectedToken = "workload-token"
+		}
+		if got := r.Header.Get("Authorization"); got != fmt.Sprintf("Bearer %s", expectedToken) {
+			http.Error(
+				w,
+				fmt.Sprintf("unexpected authorization header %q", got),
+				http.StatusUnauthorized,
+			)
+			return
+		}
+
+		switch requestNumber {
+		case 1:
+			w.Header().Set("Content-Type", "application/json")
+			w.Header().Set("Refreshed-API-Token", "execution-token")
+			_, _ = w.Write([]byte("{}"))
+		case 2:
+			w.WriteHeader(http.StatusNoContent)
+		default:
+			http.Error(w, "unexpected request", http.StatusInternalServerError)
+		}
+	}))
+	t.Cleanup(server.Close)
+
+	client, err := NewDefaultClient(server.URL)
+	require.NoError(t, err)
+	authenticatedClient, err := client.(*Client).WithBearerToken("workload-token")
+	require.NoError(t, err)
+
+	taskInstanceID := uuid.New()
+	_, err = authenticatedClient.TaskInstances().
+		Run(context.Background(), taskInstanceID, &TIEnterRunningPayload{
+			Hostname:  "worker",
+			Pid:       1,
+			StartDate: time.Now(),
+			State:     Running,
+			Unixname:  "airflow",
+		})
+	require.NoError(t, err)
+	err = authenticatedClient.TaskInstances().
+		Heartbeat(context.Background(), taskInstanceID, &TIHeartbeatInfo{
+			Hostname: "worker",
+			Pid:      1,
+		})
+	require.NoError(t, err)
+	require.EqualValues(t, 2, requestCount.Load())
+}