jason810496 commented on code in PR #73277:
URL: https://github.com/apache/airflow/pull/73277#discussion_r4048223642
##########
go-sdk/pkg/execution/integration_test.go:
##########
@@ -546,6 +547,90 @@ func TestRunTaskInjectsRuntimeContext(t *testing.T) {
assert.Equal(t, end, *dagRun.DataIntervalEnd)
}
+// A handler taking an airflow.Context gets on that one value everything
+// the runtime used to hand over as separate parameters.
+func TestRunTaskInjectsAirflowContext(t *testing.T) {
+ logical := time.Date(2026, 6, 9, 12, 0, 0, 0, time.UTC)
+
+ var got airflow.Context
+ bundle := buildBundle(t, func(r bundlev1.Registry) {
+ r.AddDag("test_dag").AddTaskWithName("ctxgrab",
+ func(actx airflow.Context) error {
+ got = actx
+ return nil
+ })
+ })
+
+ details := &genmodels.StartupDetails{
+ TI: genmodels.TaskInstance{
+ ID: "550e8400-e29b-41d4-a716-446655440000",
+ DagID: "test_dag",
+ TaskID: "ctxgrab",
+ RunID: "run1",
+ TryNumber: 2,
+ MapIndex: ptr(-1),
+ },
+ BundleInfo: genmodels.BundleInfo{Name: "test", Version: "1.0"},
+ TIContext: genmodels.TIRunContext{
+ DagRun: genmodels.DagRun{LogicalDate: logical},
+ },
+ }
+
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
+
+ result := RunTask(context.Background(), bundle, details, comm, logger)
+ assertSucceedTask(t, result)
+
+ assert.Same(t, logger, got.Logger(), "the task's logger must arrive on
the Context")
+ assert.NotNil(t, got.Client(), "the coordinator-backed client must
arrive on the Context")
+
+ ti := got.TaskInstance()
+ assert.Equal(t, "test_dag", ti.DagID)
+ assert.Equal(t, "run1", ti.RunID)
+ assert.Equal(t, "ctxgrab", ti.TaskID)
+ assert.Equal(t, 2, ti.TryNumber)
+ assert.Nil(t, ti.MapIndex, "an unmapped task (map_index -1) must
surface as nil")
+
+ dagRun := got.DagRun()
+ assert.Equal(t, "test_dag", dagRun.DagID)
+ assert.Equal(t, "run1", dagRun.RunID)
+ require.NotNil(t, dagRun.LogicalDate)
+ assert.Equal(t, logical, *dagRun.LogicalDate)
+
+ // A helper taking a plain context.Context recovers the same surface.
+ recovered, ok := airflow.FromContext(context.Context(got))
+ require.True(t, ok)
+ assert.Equal(t, ti, recovered.TaskInstance())
+}
+
+// Serve traps SIGINT/SIGTERM into the context it hands RunTask, so a
+// supervisor shutdown reaches the handler on actx.Done().
+func TestRunTaskAirflowContextHonorsShutdown(t *testing.T) {
+ bundle := buildBundle(t, func(r bundlev1.Registry) {
+ r.AddDag("test_dag").AddTaskWithName("ctxcheck",
+ func(actx airflow.Context) error {
+ select {
+ case <-actx.Done():
+ return actx.Err()
+ default:
+ return errors.New("actx.Done() did not
fire on a cancelled task context")
+ }
+ })
+ })
+
+ details := newStartupDetails("ctxcheck")
+
+ ctx, cancel := context.WithCancel(context.Background())
+ cancel()
+
+ logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+ comm := NewCoordinatorComm(bytes.NewReader(nil), io.Discard, logger)
+
+ result := RunTask(ctx, bundle, details, comm, logger)
+ assertTaskState(t, result, genmodels.TaskStateStateFailed)
Review Comment:
Good catch. Fixed in e7d05e1c71: the handler records which branch fired
based on the flag.
##########
go-sdk/airflow/context.go:
##########
@@ -0,0 +1,117 @@
+// 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 airflow
+
+import (
+ "context"
+ "log/slog"
+
+ "github.com/apache/airflow/go-sdk/sdk"
+)
+
+type (
+ // TaskInstance identifies the task instance a handler is running for.
+ TaskInstance = sdk.TaskInstance
+
+ // DagRun identifies the Dag run the task instance belongs to,
+ // and carries its scheduling timestamps.
+ DagRun = sdk.DagRun
+)
+
+// Context is the first parameter of every task handler.
+// It is a context.Context bound to the running task, so actx.Done() fires
when the supervisor
+// asks the task to stop, and it exposes what Airflow gives the task:
[Context.Logger],
+// [Context.Client], [Context.TaskInstance] and [Context.DagRun].
+//
+// [NewContext] builds one.
+type Context struct {
+ context.Context
+
+ values taskValues
+}
+
+// taskValues is what NewContext stores on the context chain so FromContext
can recover it.
+type taskValues struct {
+ logger *slog.Logger
+ client sdk.Client
+ ti TaskInstance
+ dagRun DagRun
+}
+
+type contextKey struct{}
+
+// NewContext returns a [Context] backed by ctx. It panics if ctx is nil.
+//
+// The runtime calls it when it binds a handler's first parameter.
+// Call it directly to unit-test a handler:
+//
+// actx := airflow.NewContext(
+// t.Context(), slog.Default(), fakeClient,
+// airflow.TaskInstance{DagID: "py_etl", TaskID: "transform",
TryNumber: 1},
+// airflow.DagRun{DagID: "py_etl", RunID: "run1"},
+// )
+// require.NoError(t, transform(actx, "US"))
+func NewContext(
+ ctx context.Context,
+ logger *slog.Logger,
+ client sdk.Client,
+ ti TaskInstance,
+ dagRun DagRun,
+) Context {
+ if ctx == nil {
+ panic("airflow.NewContext: cannot create airflow.Context from
nil context.Context")
+ }
+ values := taskValues{logger: logger, client: client, ti: ti, dagRun:
dagRun}
+ return Context{Context: context.WithValue(ctx, contextKey{}, values),
values: values}
+}
+
+// FromContext recovers the Airflow surface inside a helper typed as a plain
context.Context,
+// reporting whether ctx carries one.
+//
+// It succeeds for the [Context] a handler was given and for any context
derived from it.
+// The returned Context keeps ctx, so a deadline or cancellation added on the
way down applies.
+func FromContext(ctx context.Context) (Context, bool) {
+ if ctx == nil {
+ return Context{}, false
+ }
+ values, ok := ctx.Value(contextKey{}).(taskValues)
+ if !ok {
+ return Context{}, false
+ }
+ return Context{Context: ctx, values: values}, true
+}
+
+// Logger writes to the task's Airflow log. It never returns nil.
+//
+// The logger takes a context, so pass the same Context:
actx.Logger().InfoContext(actx, "msg").
+func (c Context) Logger() *slog.Logger {
+ if c.values.logger == nil {
+ return slog.Default()
+ }
+ return c.values.logger
+}
+
+// Client reads Airflow Variables, Connections and XCom.
+// Its calls take a context, so pass the same Context:
actx.Client().GetVariable(actx, "name").
+func (c Context) Client() sdk.Client { return c.values.client }
+
+// TaskInstance identifies the task instance that is executing.
+func (c Context) TaskInstance() TaskInstance { return c.values.ti }
Review Comment:
Nice catch as well. Fixed in 9d39b4d6d1 by having nil validation at
`NewContext` stage so that the getter will always be existed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
