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 63b8bedca66 Go SDK: stop logging fetched edge workloads verbatim
(#68355)
63b8bedca66 is described below
commit 63b8bedca666252cc81f4a8d11616bb1be346e78
Author: drewrukin <[email protected]>
AuthorDate: Tue Jul 7 12:34:16 2026 +0300
Go SDK: stop logging fetched edge workloads verbatim (#68355)
---
go-sdk/edge/worker.go | 61 +++++++++++++++++++++++++-
go-sdk/edge/worker_test.go | 107 +++++++++++++++++++++++++++++++++++++++++++++
2 files changed, 166 insertions(+), 2 deletions(-)
diff --git a/go-sdk/edge/worker.go b/go-sdk/edge/worker.go
index 3930806e99b..5a06f152e89 100644
--- a/go-sdk/edge/worker.go
+++ b/go-sdk/edge/worker.go
@@ -64,6 +64,33 @@ type worker struct {
sysInfo map[string]edgeapi.WorkerStateBody_Sysinfo_AdditionalProperties
}
+type fetchedJobForLog edgeapi.EdgeJobFetched
+
+func buildFetchedJobLogAttrs(job fetchedJobForLog) []slog.Attr {
+ attrs := []slog.Attr{
+ slog.String("dag_id", job.DagId),
+ slog.String("task_id", job.TaskId),
+ slog.String("run_id", job.RunId),
+ slog.Int("try_number", job.TryNumber),
+ slog.Int("map_index", job.MapIndex),
+ slog.Int("concurrency_slots", job.ConcurrencySlots),
+ slog.String("bundle_name", job.Command.BundleInfo.Name),
+ }
+
+ if job.Command.BundleInfo.Version != nil {
+ attrs = append(attrs, slog.String("bundle_version",
*job.Command.BundleInfo.Version))
+ }
+ if job.Command.Ti.Queue != "" {
+ attrs = append(attrs, slog.String("queue",
job.Command.Ti.Queue))
+ }
+
+ return attrs
+}
+
+func (job fetchedJobForLog) LogValue() slog.Value {
+ return slog.GroupValue(buildFetchedJobLogAttrs(job)...)
+}
+
var (
HeartbeatInterval = 30 * time.Second
DeregisterTimeout = 5 * time.Second
@@ -308,6 +335,29 @@ type jobInfo struct {
ConcurrencySlots int32
}
+func (job jobInfo) LogValue() slog.Value {
+ mapIndex := -1
+ if job.TI.MapIndex != nil {
+ mapIndex = *job.TI.MapIndex
+ }
+
+ attrs := []slog.Attr{
+ slog.String("dag_id", job.TI.DagId),
+ slog.String("task_id", job.TI.TaskId),
+ slog.String("run_id", job.TI.RunId),
+ slog.Int("try_number", job.TI.TryNumber),
+ slog.Int("map_index", mapIndex),
+ slog.Int("concurrency_slots", int(job.ConcurrencySlots)),
+ slog.String("bundle_name", job.BundleInfo.Name),
+ }
+
+ if job.BundleInfo.Version != nil {
+ attrs = append(attrs, slog.String("bundle_version",
*job.BundleInfo.Version))
+ }
+
+ return slog.GroupValue(attrs...)
+}
+
// fetchJobsForever will fetch jobs from the API server every second (if there
is capacity), sending the
// resulting workloads out on the channel.
func (w *worker) fetchJobsForever(ctx context.Context) <-chan jobInfo {
@@ -361,7 +411,15 @@ func (w *worker) fetchJob(ctx context.Context)
(*bundlev1.ExecuteTaskWorkload, i
return nil, 0, nil
}
- w.logger.Info("fetchJob", "resp", fmt.Sprintf("%#v\n", resp))
+ w.logger.Info("Fetched job", "job",
fetchedJobForLog(edgeapi.EdgeJobFetched{
+ Command: resp.Command,
+ ConcurrencySlots: resp.ConcurrencySlots,
+ DagId: resp.DagId,
+ MapIndex: resp.MapIndex,
+ RunId: resp.RunId,
+ TaskId: resp.TaskId,
+ TryNumber: resp.TryNumber,
+ }))
// Round trip via json. Inefficient, but easy to code
asJSON, err := json.Marshal(resp.Command)
@@ -376,7 +434,6 @@ func (w *worker) fetchJob(ctx context.Context)
(*bundlev1.ExecuteTaskWorkload, i
// TODO: Report this to API server
return nil, 0, fmt.Errorf("unable to unmarshal into workload
%w", err)
}
- w.logger.Info("fetchJob", "out", fmt.Sprintf("%#v\n", out))
return &out, int32(resp.ConcurrencySlots), nil
}
diff --git a/go-sdk/edge/worker_test.go b/go-sdk/edge/worker_test.go
new file mode 100644
index 00000000000..78d119f0ec9
--- /dev/null
+++ b/go-sdk/edge/worker_test.go
@@ -0,0 +1,107 @@
+// 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 edge
+
+import (
+ "bytes"
+ "context"
+ "encoding/json"
+ "log/slog"
+ "net/http"
+ "net/http/httptest"
+ "testing"
+
+ "github.com/stretchr/testify/require"
+
+ "github.com/apache/airflow/go-sdk/pkg/edgeapi"
+)
+
+func TestFetchJobDoesNotLogToken(t *testing.T) {
+ var logBuffer bytes.Buffer
+
+ secretToken := "super-secret-edge-token"
+ version := "1.0.0"
+ logPath := "/tmp/example.log"
+
+ server := httptest.NewServer(http.HandlerFunc(func(w
http.ResponseWriter, r *http.Request) {
+ require.Equal(t, http.MethodPost, r.Method)
+ require.Equal(t, "/edge_worker/v1/jobs/fetch/test-worker",
r.URL.Path)
+
+ w.Header().Set("Content-Type", "application/json")
+ err := json.NewEncoder(w).Encode(edgeapi.EdgeJobFetched{
+ Command: edgeapi.ExecuteTask{
+ BundleInfo: edgeapi.BundleInfo{
+ Name: "example-bundle",
+ Version: &version,
+ },
+ DagRelPath: "dags/example.py",
+ LogPath: &logPath,
+ Ti: edgeapi.TaskInstance{
+ DagId: "example_dag",
+ Queue: "default",
+ RunId:
"manual__2026-06-10T00:00:00+00:00",
+ TaskId: "example_task",
+ TryNumber: 2,
+ PoolSlots: 1,
+ PriorityWeight: 1,
+ },
+ Token: secretToken,
+ },
+ ConcurrencySlots: 3,
+ DagId: "example_dag",
+ MapIndex: -1,
+ RunId: "manual__2026-06-10T00:00:00+00:00",
+ TaskId: "example_task",
+ TryNumber: 2,
+ })
+ require.NoError(t, err)
+ }))
+ t.Cleanup(server.Close)
+
+ client, err := edgeapi.NewClient(server.URL + "/")
+ require.NoError(t, err)
+
+ w := &worker{
+ hostname: "test-worker",
+ client: client,
+ queues: []string{"default"},
+ logger: slog.New(slog.NewJSONHandler(&logBuffer,
&slog.HandlerOptions{
+ Level: slog.LevelDebug,
+ })),
+ maxConcurrency: 16,
+ }
+ w.freeConcurrency.Store(4)
+
+ workload, slots, err := w.fetchJob(context.Background())
+ require.NoError(t, err)
+ require.NotNil(t, workload)
+ require.Equal(t, int32(3), slots)
+ require.Equal(t, secretToken, workload.Token)
+ w.logger.Debug("Got allocation", "workload", jobInfo{
+ ExecuteTaskWorkload: *workload,
+ ConcurrencySlots: slots,
+ })
+
+ logOutput := logBuffer.String()
+ require.Contains(t, logOutput, "Fetched job")
+ require.Contains(t, logOutput, "Got allocation")
+ require.Contains(t, logOutput, "example_dag")
+ require.Contains(t, logOutput, "example_task")
+ require.NotContains(t, logOutput, secretToken)
+ require.NotContains(t, logOutput, "\"token\"")
+}
