amoghrajesh commented on code in PR #68349:
URL: https://github.com/apache/airflow/pull/68349#discussion_r3393706304
##########
go-sdk/sdk/context.go:
##########
@@ -20,29 +20,48 @@ package sdk
import (
"context"
"time"
-
- "github.com/apache/airflow/go-sdk/pkg/sdkcontext"
)
-// RuntimeContext carries the identifiers and scheduling timestamps of the task
-// instance that is currently executing, along with the Dag run it belongs to.
-// It is the Go equivalent of the execution context the Python and Java SDKs
-// expose to task authors.
+// TIRunContext is the execution context handed to a task. It embeds the
+// standard context.Context (cancellation, deadline, request-scoped values) and
+// additionally carries the identifiers and scheduling timestamps of the task
+// instance that is executing, along with the Dag run it belongs to. It is the
+// Go equivalent of the execution context the Python and Java SDKs expose to
+// task authors.
//
-// Retrieve it inside a task function with CurrentContext:
+// The runtime injects it into a task function by parameter type, so declare it
+// as the task's context argument and read the fields directly:
//
-// func myTask(ctx context.Context, log *slog.Logger) error {
-// rc := sdk.CurrentContext(ctx)
-// log.Info("running", "task_id", rc.TI.TaskID, "run_id",
rc.DagRun.RunID)
+// func myTask(ctx sdk.TIRunContext, log *slog.Logger) error {
+// log.Info("running", "task_id", ctx.TI.TaskID, "run_id",
ctx.DagRun.RunID)
// return nil
// }
-type RuntimeContext struct {
+//
+// Because it embeds context.Context it is itself a context.Context: pass it
+// straight to client calls, select on ctx.Done(), or hand it to downstream
+// helpers that take a context.Context.
+//
+// Embedding a context.Context in a struct is normally discouraged, but it is a
+// deliberate exception here: the runtime needs a concrete type to bind by
+// parameter, and the value lives only for the duration of a single task
+// execution. The runtime always sets the embedded Context before invoking the
+// task. When constructing one yourself (for example in a unit test) set the
+// Context field, otherwise the context.Context methods dereference a nil
+// interface and panic:
+//
+// ctx := sdk.TIRunContext{Context: context.Background()}
+type TIRunContext struct {
+ context.Context
Review Comment:
One thing worth discussing: embedding `context.Context` in a struct is a Go
anti pattern the docs explicitly advises against
(https://pkg.go.dev/context#hdr-Contexts_and_structs). This is cos because
structs tend to outlive the request they were created for, so the stored
context's cancellation/deadline silently carries over into unintended reuse.
The alternative that avoids this entirely: keep `context.Context` as
`context.Context` and inject the `TI/DagRun` metadata as a separate plain
struct maybe:
```go
func myTask(ctx context.Context, run sdk.RunInfo, log *slog.Logger) error {
log.Info("running", "task_id", run.TI.TaskID)
client.GetVariable(ctx, "my-var")
}
```
worth considering before the API stabilises.
--
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]
