From 0db107c44cd237e28052860471ad6d9c0af97a3a Mon Sep 17 00:00:00 2001 From: drewrukin Date: Wed, 10 Jun 2026 18:26:42 +0000 Subject: [PATCH 1/2] Go SDK: stop logging fetched edge workloads verbatim --- go-sdk/edge/worker.go | 32 +++++++++++- go-sdk/edge/worker_test.go | 100 +++++++++++++++++++++++++++++++++++++ 2 files changed, 130 insertions(+), 2 deletions(-) create mode 100644 go-sdk/edge/worker_test.go diff --git a/go-sdk/edge/worker.go b/go-sdk/edge/worker.go index 3930806e99b1e..c5a326f0124c1 100644 --- a/go-sdk/edge/worker.go +++ b/go-sdk/edge/worker.go @@ -64,6 +64,27 @@ type worker struct { sysInfo map[string]edgeapi.WorkerStateBody_Sysinfo_AdditionalProperties } +func buildFetchedJobLogAttrs(job edgeapi.EdgeJobFetched) []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 +} + var ( HeartbeatInterval = 30 * time.Second DeregisterTimeout = 5 * time.Second @@ -361,7 +382,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.LogAttrs(ctx, slog.LevelInfo, "Fetched job", buildFetchedJobLogAttrs(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 +405,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 0000000000000..4cf6547b01312 --- /dev/null +++ b/go-sdk/edge/worker_test.go @@ -0,0 +1,100 @@ +// 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, nil)), + 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) + + logOutput := logBuffer.String() + require.Contains(t, logOutput, "Fetched job") + require.Contains(t, logOutput, "example_dag") + require.Contains(t, logOutput, "example_task") + require.NotContains(t, logOutput, secretToken) + require.NotContains(t, logOutput, "\"token\"") +} From ee4d12a825572e0f4eda7471c758c19ab6fd027b Mon Sep 17 00:00:00 2001 From: drewrukin Date: Mon, 15 Jun 2026 10:36:42 +0000 Subject: [PATCH 2/2] Use LogValuer for edge workload logging --- go-sdk/edge/worker.go | 35 ++++++++++++++++++++++++++++++++--- go-sdk/edge/worker_test.go | 15 +++++++++++---- 2 files changed, 43 insertions(+), 7 deletions(-) diff --git a/go-sdk/edge/worker.go b/go-sdk/edge/worker.go index c5a326f0124c1..5a06f152e890c 100644 --- a/go-sdk/edge/worker.go +++ b/go-sdk/edge/worker.go @@ -64,7 +64,9 @@ type worker struct { sysInfo map[string]edgeapi.WorkerStateBody_Sysinfo_AdditionalProperties } -func buildFetchedJobLogAttrs(job edgeapi.EdgeJobFetched) []slog.Attr { +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), @@ -85,6 +87,10 @@ func buildFetchedJobLogAttrs(job edgeapi.EdgeJobFetched) []slog.Attr { return attrs } +func (job fetchedJobForLog) LogValue() slog.Value { + return slog.GroupValue(buildFetchedJobLogAttrs(job)...) +} + var ( HeartbeatInterval = 30 * time.Second DeregisterTimeout = 5 * time.Second @@ -329,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 { @@ -382,7 +411,7 @@ func (w *worker) fetchJob(ctx context.Context) (*bundlev1.ExecuteTaskWorkload, i return nil, 0, nil } - w.logger.LogAttrs(ctx, slog.LevelInfo, "Fetched job", buildFetchedJobLogAttrs(edgeapi.EdgeJobFetched{ + w.logger.Info("Fetched job", "job", fetchedJobForLog(edgeapi.EdgeJobFetched{ Command: resp.Command, ConcurrencySlots: resp.ConcurrencySlots, DagId: resp.DagId, @@ -390,7 +419,7 @@ func (w *worker) fetchJob(ctx context.Context) (*bundlev1.ExecuteTaskWorkload, i RunId: resp.RunId, TaskId: resp.TaskId, TryNumber: resp.TryNumber, - })...) + })) // Round trip via json. Inefficient, but easy to code asJSON, err := json.Marshal(resp.Command) diff --git a/go-sdk/edge/worker_test.go b/go-sdk/edge/worker_test.go index 4cf6547b01312..78d119f0ec9fe 100644 --- a/go-sdk/edge/worker_test.go +++ b/go-sdk/edge/worker_test.go @@ -77,10 +77,12 @@ func TestFetchJobDoesNotLogToken(t *testing.T) { require.NoError(t, err) w := &worker{ - hostname: "test-worker", - client: client, - queues: []string{"default"}, - logger: slog.New(slog.NewJSONHandler(&logBuffer, nil)), + 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) @@ -90,9 +92,14 @@ func TestFetchJobDoesNotLogToken(t *testing.T) { 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)