Initial version of the Golang Task SDK for Apache Airflow (#50964)

This introduces the Golang Task SDK, to allow Airflow tasks to be written in Go. At this stage the code is working well enough to be able to run tasks written in Go (via a Celery broker of all things, because there was a ready made Go client for it.) but it is not close to being GA yet. Some of the known features gaps are documented in the README, we should convert those into tickets The code in this PR is a good start, but we need more tests (there are some failing now as I don't think Go has TODO tests like pytest does) and the code layout (pkg/ folder vs others etc) need some work (but in a future PR) My goal here with this PR at this stage is two-fold: 1. Get this up for more eyes to check I haven't made any huge mistakes, and 2. Allow multiple people to work on this concurrently without making the PR even bigger. As a result unless there are an glaring mistakes/things that need to be fixed now I would like to merge this and we can improve on it over time.
2026-11-04 07:12:48 +00:00 · 2025-05-23 16:04:23 +01:00
1 parent 341ab096c6
commit fb37f959ed
@@ -58,3 +58,7 @@ indent_size = 2

 [*.json]
 indent_size = 4
+
+[*.go]
+indent_style = tab
+max_line_length = 110
@@ -170,6 +170,21 @@ repos:
          - --fuzzy-match-generates-todo
        files: >
          \.cfg$|\.conf$|\.ini$|\.ldif$|\.properties$|\.service$|\.tf$|Dockerfile.*$
+  # https://github.com/Lucas-C/pre-commit-hooks/pull/103
+  - repo: https://github.com/ashb/pre-commit-hooks
+    rev: 36ddb15cd7a607003b06bf11c28045c91ea97dc0
+    hooks:
+      - id: insert-license
+        name: Add license for all Go files
+        types: [go]
+        args:
+          - --comment-style
+          - "|//|"
+          - --license-filepath
+          - scripts/ci/license-templates/LICENSE.txt
+          - --insert-license-after-regex
+          # We need this 'generated by' line at the top for `golines` to not format it
+          - '// Code generated by .*'
  - repo: local
    hooks:
      - id: check-min-python-version
@@ -406,7 +421,7 @@ repos:
        entry: ./scripts/ci/pre_commit/replace_bad_characters.py
        language: python
        types: [file, text]
-        exclude: ^clients/gen/go\.sh$|^\.gitmodules$|.*/dist/.*
+        exclude: ^clients/gen/go\.sh$|^\.gitmodules$|.*/dist/.*|\.go$|/go\.(mod|sum)$
        additional_dependencies: ['rich>=12.4.4']
      - id: lint-dockerfile
        name: Lint Dockerfile
@@ -1231,6 +1246,29 @@ repos:
        files: ^airflow-core/src/airflow/migrations/versions/.*\.py$
        exclude:
          airflow-core/src/airflow/migrations/versions/0028_3_0_0_drop_ab_user_id_foreign_key.py
+      - id: generate-airflowctl-help-images
+        name: Generate SVG from Airflow CTL Commands
+        entry: ./scripts/ci/pre_commit/capture_airflowctl_help.py
+        language: python
+        pass_filenames: false
+        files:
+          ^airflow-ctl/src/airflowctl/api/operations.py|airflow-ctl/src/airflowctl/cli/.*\.py$
+        additional_dependencies: ['rich>=12.4.4']
+      - id: gofmt
+        name: Format go code
+        entry: golines --base-formatter=gofumpt --write-output --max-len=100 --chain-split-dots
+        additional_dependencies: [github.com/segmentio/golines@latest, mvdan.cc/gofumpt@v0.8.0]
+        files: ^go-sdk/
+        types: [go]
+        language: golang
+      - id: gci
+        name: Consistent import ordering for Go files
+        # Since this is invoked from the root folder, not go-sdk/, gci can't auto-detect the prefix
+        entry: gci write -s standard -s default -s "prefix(github.com/apache/airflow)"
+        additional_dependencies: [github.com/daixiang0/gci@v0.13.6]
+        files: ^go-sdk/
+        types: [go]
+        language: golang
        ## ADD MOST PRE-COMMITS ABOVE THAT LINE
        # The below pre-commits are those requiring CI image to be built
      - id: mypy-dev
@@ -174,3 +174,7 @@ auth_generated.py

 # hash files
 www-hash.txt
+
+# go setup files
+go.mod
+go.sum
@@ -271,11 +271,14 @@ require Breeze Docker image to be built locally.
 +-----------------------------------------------------------+--------------------------------------------------------+---------+
 | flynt                                                     | Run flynt string format converter for Python           |         |
 +-----------------------------------------------------------+--------------------------------------------------------+---------+
+| gci                                                       | Consistent import ordering for Go files                |         |
+-----------------------------------------------------------+--------------------------------------------------------+---------+
 | generate-airflow-diagrams                                 | Generate airflow diagrams                              |         |
 +-----------------------------------------------------------+--------------------------------------------------------+---------+
 | generate-airflowctl-datamodels                            | Generate Datamodels for AirflowCTL                     | *       |
 +-----------------------------------------------------------+--------------------------------------------------------+---------+
-| generate-airflowctl-help-images                           | Generate SVG from Airflow CTL Commands                 | *       |
+| generate-airflowctl-help-images                           | * Generate SVG from Airflow CTL Commands               | *       |
+|                                                           | * Generate SVG from Airflow CTL Commands               |         |
 +-----------------------------------------------------------+--------------------------------------------------------+---------+
 | generate-openapi-spec                                     | Generate the FastAPI API spec                          | *       |
 +-----------------------------------------------------------+--------------------------------------------------------+---------+
@@ -287,6 +290,8 @@ require Breeze Docker image to be built locally.
 +-----------------------------------------------------------+--------------------------------------------------------+---------+
 | generate-volumes-for-sources                              | Generate volumes for docker compose                    |         |
 +-----------------------------------------------------------+--------------------------------------------------------+---------+
+| gofmt                                                     | Format go code                                         |         |
+-----------------------------------------------------------+--------------------------------------------------------+---------+
 | identity                                                  | Print checked files                                    |         |
 +-----------------------------------------------------------+--------------------------------------------------------+---------+
 | insert-license                                            | * Add license for all SQL files                        |         |
@@ -300,6 +305,7 @@ require Breeze Docker image to be built locally.
 |                                                           | * Add license for all YAML files except Helm templates |         |
 |                                                           | * Add license for all Markdown files                   |         |
 |                                                           | * Add license for all other files                      |         |
+|                                                           | * Add license for all Go files                         |         |
 +-----------------------------------------------------------+--------------------------------------------------------+---------+
 | kubeconform                                               | Kubeconform check on our helm chart                    |         |
 +-----------------------------------------------------------+--------------------------------------------------------+---------+
@@ -362,11 +362,11 @@
 </text><text class="breeze-static-checks-r5" x="0" y="947.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-38)">│</text><text class="breeze-static-checks-r7" x="451.4" y="947.2" textLength="988.2" clip-path="url(#breeze-static-checks-line-38)">check-urlparse-usage-in-code&#160;|&#160;check-xml&#160;|&#160;check-zip-file-is-not-committed&#160;|&#160;&#160;&#160;&#160;&#160;</text><text class="breeze-static-checks-r5" x="1451.8" y="947.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-38)">│</text><text class="breeze-static-checks-r1" x="1464" y="947.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-38)">
 </text><text class="breeze-static-checks-r5" x="0" y="971.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-39)">│</text><text class="breeze-static-checks-r7" x="451.4" y="971.6" textLength="988.2" clip-path="url(#breeze-static-checks-line-39)">codespell&#160;|&#160;compile-fab-assets&#160;|&#160;compile-ui-assets&#160;|&#160;compile-ui-assets-dev&#160;|&#160;&#160;&#160;&#160;&#160;</text><text class="breeze-static-checks-r5" x="1451.8" y="971.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-39)">│</text><text class="breeze-static-checks-r1" x="1464" y="971.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-39)">
 </text><text class="breeze-static-checks-r5" x="0" y="996" textLength="12.2" clip-path="url(#breeze-static-checks-line-40)">│</text><text class="breeze-static-checks-r7" x="451.4" y="996" textLength="988.2" clip-path="url(#breeze-static-checks-line-40)">create-missing-init-py-files-tests&#160;|&#160;debug-statements&#160;|&#160;detect-private-key&#160;|&#160;&#160;&#160;&#160;&#160;</text><text class="breeze-static-checks-r5" x="1451.8" y="996" textLength="12.2" clip-path="url(#breeze-static-checks-line-40)">│</text><text class="breeze-static-checks-r1" x="1464" y="996" textLength="12.2" clip-path="url(#breeze-static-checks-line-40)">
-</text><text class="breeze-static-checks-r5" x="0" y="1020.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-41)">│</text><text class="breeze-static-checks-r7" x="451.4" y="1020.4" textLength="988.2" clip-path="url(#breeze-static-checks-line-41)">doctoc&#160;|&#160;end-of-file-fixer&#160;|&#160;fix-encoding-pragma&#160;|&#160;flynt&#160;|&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;</text><text class="breeze-static-checks-r5" x="1451.8" y="1020.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-41)">│</text><text class="breeze-static-checks-r1" x="1464" y="1020.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-41)">
+</text><text class="breeze-static-checks-r5" x="0" y="1020.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-41)">│</text><text class="breeze-static-checks-r7" x="451.4" y="1020.4" textLength="988.2" clip-path="url(#breeze-static-checks-line-41)">doctoc&#160;|&#160;end-of-file-fixer&#160;|&#160;fix-encoding-pragma&#160;|&#160;flynt&#160;|&#160;gci&#160;|&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;</text><text class="breeze-static-checks-r5" x="1451.8" y="1020.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-41)">│</text><text class="breeze-static-checks-r1" x="1464" y="1020.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-41)">
 </text><text class="breeze-static-checks-r5" x="0" y="1044.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-42)">│</text><text class="breeze-static-checks-r7" x="451.4" y="1044.8" textLength="988.2" clip-path="url(#breeze-static-checks-line-42)">generate-airflow-diagrams&#160;|&#160;generate-airflowctl-datamodels&#160;|&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;</text><text class="breeze-static-checks-r5" x="1451.8" y="1044.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-42)">│</text><text class="breeze-static-checks-r1" x="1464" y="1044.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-42)">
 </text><text class="breeze-static-checks-r5" x="0" y="1069.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-43)">│</text><text class="breeze-static-checks-r7" x="451.4" y="1069.2" textLength="988.2" clip-path="url(#breeze-static-checks-line-43)">generate-airflowctl-help-images&#160;|&#160;generate-openapi-spec&#160;|&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;</text><text class="breeze-static-checks-r5" x="1451.8" y="1069.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-43)">│</text><text class="breeze-static-checks-r1" x="1464" y="1069.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-43)">
 </text><text class="breeze-static-checks-r5" x="0" y="1093.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-44)">│</text><text class="breeze-static-checks-r7" x="451.4" y="1093.6" textLength="988.2" clip-path="url(#breeze-static-checks-line-44)">generate-openapi-spec-fab&#160;|&#160;generate-pypi-readme&#160;|&#160;generate-tasksdk-datamodels&#160;|&#160;</text><text class="breeze-static-checks-r5" x="1451.8" y="1093.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-44)">│</text><text class="breeze-static-checks-r1" x="1464" y="1093.6" textLength="12.2" clip-path="url(#breeze-static-checks-line-44)">
-</text><text class="breeze-static-checks-r5" x="0" y="1118" textLength="12.2" clip-path="url(#breeze-static-checks-line-45)">│</text><text class="breeze-static-checks-r7" x="451.4" y="1118" textLength="988.2" clip-path="url(#breeze-static-checks-line-45)">generate-volumes-for-sources&#160;|&#160;identity&#160;|&#160;insert-license&#160;|&#160;kubeconform&#160;|&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;</text><text class="breeze-static-checks-r5" x="1451.8" y="1118" textLength="12.2" clip-path="url(#breeze-static-checks-line-45)">│</text><text class="breeze-static-checks-r1" x="1464" y="1118" textLength="12.2" clip-path="url(#breeze-static-checks-line-45)">
+</text><text class="breeze-static-checks-r5" x="0" y="1118" textLength="12.2" clip-path="url(#breeze-static-checks-line-45)">│</text><text class="breeze-static-checks-r7" x="451.4" y="1118" textLength="988.2" clip-path="url(#breeze-static-checks-line-45)">generate-volumes-for-sources&#160;|&#160;gofmt&#160;|&#160;identity&#160;|&#160;insert-license&#160;|&#160;kubeconform&#160;|&#160;</text><text class="breeze-static-checks-r5" x="1451.8" y="1118" textLength="12.2" clip-path="url(#breeze-static-checks-line-45)">│</text><text class="breeze-static-checks-r1" x="1464" y="1118" textLength="12.2" clip-path="url(#breeze-static-checks-line-45)">
 </text><text class="breeze-static-checks-r5" x="0" y="1142.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-46)">│</text><text class="breeze-static-checks-r7" x="451.4" y="1142.4" textLength="988.2" clip-path="url(#breeze-static-checks-line-46)">lint-chart-schema&#160;|&#160;lint-dockerfile&#160;|&#160;lint-helm-chart&#160;|&#160;lint-json-schema&#160;|&#160;&#160;&#160;&#160;&#160;&#160;&#160;</text><text class="breeze-static-checks-r5" x="1451.8" y="1142.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-46)">│</text><text class="breeze-static-checks-r1" x="1464" y="1142.4" textLength="12.2" clip-path="url(#breeze-static-checks-line-46)">
 </text><text class="breeze-static-checks-r5" x="0" y="1166.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-47)">│</text><text class="breeze-static-checks-r7" x="451.4" y="1166.8" textLength="988.2" clip-path="url(#breeze-static-checks-line-47)">lint-markdown&#160;|&#160;mixed-line-ending&#160;|&#160;mypy-airflow-core&#160;|&#160;mypy-airflow-ctl&#160;|&#160;&#160;&#160;&#160;&#160;&#160;&#160;</text><text class="breeze-static-checks-r5" x="1451.8" y="1166.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-47)">│</text><text class="breeze-static-checks-r1" x="1464" y="1166.8" textLength="12.2" clip-path="url(#breeze-static-checks-line-47)">
 </text><text class="breeze-static-checks-r5" x="0" y="1191.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-48)">│</text><text class="breeze-static-checks-r7" x="451.4" y="1191.2" textLength="988.2" clip-path="url(#breeze-static-checks-line-48)">mypy-dev&#160;|&#160;mypy-devel-common&#160;|&#160;mypy-providers&#160;|&#160;mypy-task-sdk&#160;|&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;&#160;</text><text class="breeze-static-checks-r5" x="1451.8" y="1191.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-48)">│</text><text class="breeze-static-checks-r1" x="1464" y="1191.2" textLength="12.2" clip-path="url(#breeze-static-checks-line-48)">
@@ -1 +1 @@
-44ddf285d8b0b2f0369d4750ace7dde9
+06e827be83fd345223ec466cb843ccc0
@@ -99,6 +99,7 @@ PRE_COMMIT_LIST = [
    "end-of-file-fixer",
    "fix-encoding-pragma",
    "flynt",
+    "gci",
    "generate-airflow-diagrams",
    "generate-airflowctl-datamodels",
    "generate-airflowctl-help-images",
@@ -107,6 +108,7 @@ PRE_COMMIT_LIST = [
    "generate-pypi-readme",
    "generate-tasksdk-datamodels",
    "generate-volumes-for-sources",
+    "gofmt",
    "identity",
    "insert-license",
    "kubeconform",
@@ -0,0 +1,62 @@
+<!--
+ 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.
+ -->
+
+# 🚧 Apache Airflow Go Task SDK 🚧
+
+> [!NOTE]
+> This Golang SDK is under active development and is not ready for prime-time yet.
+
+This README is primarily aimed at developers working _on_ the Go-SDK itself. Users wishing to write Airflow tasks in Go should look at the reference docs, but those don't exist yet.
+
+## How It Works
+
+The Go SDK uses the Task Execution Interface (TEI or Task API) introduced in AIP-72 with Airflow 3.0.0.
+
+The Task API however does not provide a means to get the `ExecuteTaskWorkload` to the go worker itself. For the short term, we make use of [gopher-celery](github.com/marselester/gopher-celery) to get tasks from a Redis broker. Longer term we will likely need to stabilize the Edge Executor API and write a go client for that.
+
+Since Go is a compiled language (putting aside projects such as [YAEGI](https://github.com/traefik/yaegi) that allow go to be interpreted) all tasks must be a) compiled in to the binary, and b) "registered" inside the worker process in order to be executed.
+
+## Current state
+
+This SDK currently will:
+
+- Get tasks from Celery queue(s)
+- Run registered tasks (no support for dag versioning or loading of multiple different "bundles")
+- Heartbeat and report the final state of the final TI
+- Allow access to Variables
+
+## Known missing features
+
+A non-exhaustive list of features we have yet to implement
+
+- Reading of Airflow Connections
+- Support for putting tasks into state other than success or failed/up-for-retry (deferred, failed-without-retries etc.)
+- HTTP Log server to view logs from in-progress tasks
+- Remote task logs (i.e. S3/GCS etc)
+- XCom reading/writing from API server
+- XCom reading/writing from other XCom backends
+
+
+## Future Direction
+
+This is more of an "it would be nice to have" than any plan or commitment, and a place to record ideas.
+
+- Support multiple versions by compiling tasks/bundles into plugins and make use of [go-plugin](https://github.com/hashicorp/go-plugin) (This is how Terraform providers work)
+
+  This would enable use to have executor code and task code in separate processes, and to be able to have a single worker execute different bundles/versions of tasks (i.e. we'd have a go executor process that launches versioned plugin bundles to actually execute the task)
@@ -0,0 +1,106 @@
+// 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 celery
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"log/slog"
+	"os"
+	"os/signal"
+
+	celery "github.com/marselester/gopher-celery"
+	celeryredis "github.com/marselester/gopher-celery/goredis"
+	"github.com/redis/go-redis/v9"
+	"github.com/spf13/viper"
+
+	"github.com/apache/airflow/go-sdk/pkg/api"
+	"github.com/apache/airflow/go-sdk/pkg/sdkcontext"
+	"github.com/apache/airflow/go-sdk/worker"
+)
+
+func Run(ctx context.Context, config Config) error {
+	worker, ok := ctx.Value(sdkcontext.WorkerContextKey).(worker.Worker)
+	if !ok {
+		slog.ErrorContext(ctx, "Worker missing in context")
+		return fmt.Errorf("worker value missing in context")
+	}
+
+	worker, err := worker.WithServer(viper.GetString("execution-api-url"))
+	if err != nil {
+		slog.ErrorContext(ctx, "Error setting ExecutionAPI sxerver for worker", "err", err)
+		return err
+	}
+
+	// Store the updated worker in context
+	ctx = context.WithValue(ctx, sdkcontext.WorkerContextKey, worker)
+
+	ctx, stop := signal.NotifyContext(ctx, os.Interrupt)
+	defer stop()
+
+	c := redis.NewClient(&redis.Options{
+		Addr: config.BrokerAddr,
+	})
+	defer func() {
+		if err := c.Close(); err != nil {
+			slog.ErrorContext(ctx, "failed to close Redis client", "err", err)
+		}
+	}()
+
+	if _, err := c.Ping(ctx).Result(); err != nil {
+		slog.ErrorContext(ctx, "Redis connection failed", "err", err)
+		return err
+	}
+	broker := celeryredis.NewBroker(
+		celeryredis.WithClient(c),
+	)
+
+	if len(config.Queues) == 0 {
+		return fmt.Errorf("no queues defined")
+	}
+	broker.Observe(config.Queues)
+
+	app := celery.NewApp(
+		celery.WithBroker(broker),
+	)
+
+	for _, queue := range config.Queues {
+		app.Register(
+			"execute_workload",
+			queue,
+			func(ctx context.Context, p *celery.TaskParam) error {
+				p.NameArgs("payload")
+				payload := p.MustString("payload")
+
+				var workload api.ExecuteTaskWorkload
+				if err := json.Unmarshal([]byte(payload), &workload); err != nil {
+					return err
+				}
+				return worker.ExecuteTaskWorkload(ctx, workload)
+			},
+		)
+	}
+
+	slog.Info("waiting for tasks", "queues", config.Queues)
+	err = app.Run(ctx)
+	if err != nil {
+		slog.Error("program stopped", "error", err)
+	}
+	return err
+}
@@ -0,0 +1,153 @@
+// 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 cmd
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"path"
+	"strings"
+
+	cc "github.com/ivanpirog/coloredcobra"
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+	"github.com/spf13/viper"
+
+	"github.com/apache/airflow/go-sdk/pkg/sdkcontext"
+	"github.com/apache/airflow/go-sdk/worker"
+)
+
+var cfgFile string
+
+// rootCmd represents the base command when called without any subcommands
+var rootCmd = &cobra.Command{
+	Use:   "airflow-go-worker",
+	Short: "A brief description of your application",
+	Long: `Airflow worker for running Go tasks.
+
+All options (other than ` + "`--config`" + `) can be specified in the config file using
+the same name as the CLI argument but without the ` + "`--`" + ` prefix.`,
+	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
+		return initializeConfig(cmd)
+	},
+}
+
+// Execute is the main entrypoint, and runs the Celery broker for the given worker
+func Execute(worker worker.Worker) {
+	// TODO: This should possibly just take a task Registry, not a worker object
+	cc.Init(&cc.Config{
+		RootCmd:       rootCmd,
+		Headings:      cc.Bold,
+		Commands:      cc.Yellow + cc.Bold,
+		Example:       cc.Italic,
+		ExecName:      cc.HiMagenta + cc.Bold,
+		Flags:         cc.Green,
+		FlagsDataType: cc.Italic + cc.White,
+	})
+	// Store the worker in the context so we can pull it out later
+	ctx := context.WithValue(context.Background(), sdkcontext.WorkerContextKey, worker)
+	err := rootCmd.ExecuteContext(ctx)
+	if err != nil {
+		os.Exit(1)
+	}
+}
+
+func init() {
+	rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "",
+		"config file (default is $HOME/airflow/go-sdk.yaml)")
+	rootCmd.AddCommand(runCmd)
+}
+
+var envKeyReplacer *strings.Replacer = strings.NewReplacer(".", "__", "-", "_")
+
+func setupViper() error {
+	if cfgFile != "" {
+		// Use config file from the flag.
+		viper.SetConfigFile(cfgFile)
+	} else {
+		airflowHome := os.Getenv("AIRFLOW_HOME")
+		if airflowHome == "" {
+			home, err := os.UserHomeDir()
+			cobra.CheckErr(err)
+			airflowHome = path.Join(home, "airflow")
+		}
+
+		// Search config in home directory with name ".go-sdk" (without extension).
+		viper.AddConfigPath(airflowHome)
+		viper.SetConfigType("yaml")
+		viper.SetConfigName("go-sdk.yaml")
+	}
+
+	viper.SetOptions(viper.ExperimentalBindStruct())
+
+	// Attempt to read the config file, gracefully ignoring errors
+	// caused by a config file not being found. Return an error
+	// if we cannot parse the config file.
+	if err := viper.ReadInConfig(); err != nil {
+		// It's okay if there isn't a config file
+		if _, ok := err.(viper.ConfigFileNotFoundError); !ok {
+			return err
+		}
+	}
+
+	// Make the prefix be AIRFLOW__ -- viper adds an extra `_` automatically
+	viper.SetEnvPrefix("AIRFLOW_")
+	// Set the key replacer to replace "__" with "."
+	viper.SetEnvKeyReplacer(envKeyReplacer)
+
+	viper.AutomaticEnv() // read in environment variables that match
+	return nil
+}
+
+// initConfig reads in config file and ENV variables if set.
+func initializeConfig(cmd *cobra.Command) error {
+	if err := setupViper(); err != nil {
+		return err
+	}
+
+	// Bind the current command's flags to viper
+	bindFlags(cmd, envKeyReplacer)
+
+	return nil
+}
+
+// Bind each cobra flag to its associated viper configuration (config file and environment variable)
+// This approach cribbed from https://github.com/carolynvs/stingoftheviper/blob/19bd73117f0285436505ca17616cbc394d22e63d/main.go
+func bindFlags(cmd *cobra.Command, replacer *strings.Replacer) {
+	cmd.Flags().VisitAll(func(f *pflag.Flag) {
+		// Determine the naming convention of the flags when represented in the config file
+
+		// Since viper does case-insensitive comparisons, we don't need to bother fixing the case, and only need to remove the hyphens.
+		configName := replacer.Replace(f.Name)
+
+		// Use cli flags for preference over env or config file!
+		if f.Changed {
+			// Apply the viper config value to the flag when the flag is not set and viper has a value
+			if slice, ok := f.Value.(pflag.SliceValue); ok {
+				viper.Set(configName, slice.GetSlice())
+			} else {
+				viper.Set(configName, f.Value.String())
+			}
+		} else if viper.IsSet(configName) {
+			// If we have a viper config but no flag, set the flag value. This lets `MarkRequiredFlag` work
+			val := viper.Get(configName)
+			cmd.Flags().Set(f.Name, fmt.Sprintf("%v", val))
+		}
+	})
+}
@@ -0,0 +1,49 @@
+// 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 cmd
+
+import (
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+
+	"github.com/apache/airflow/go-sdk/celery"
+)
+
+var runCmd = &cobra.Command{
+	Use:   "run",
+	Short: "Connect to Celery broker and run Airflow workloads",
+
+	RunE: func(cmd *cobra.Command, args []string) error {
+		var config celery.Config
+
+		viper.BindPFlags(cmd.Flags())
+
+		err := viper.Unmarshal(&config)
+		cobra.CheckErr(err)
+
+		return celery.Run(cmd.Context(), config)
+	},
+}
+
+func init() {
+	runCmd.Flags().String("broker-address", "", "Celery Broker host:port to connect to")
+	runCmd.Flags().
+		String("execution-api-url", "http://localhost:8080/execution/", "Execution API to connect to")
+	runCmd.Flags().StringSlice("queues", []string{"default"}, "Celery queues to listen on")
+	runCmd.MarkFlagRequired("broker-address")
+}
@@ -0,0 +1,24 @@
+// 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 celery
+
+type Config struct {
+	BrokerAddr string   `mapstructure:"broker_address"`
+	Port       int      `mapstructure:"port"`
+	Queues     []string `mapstructure:"queues"`
+}
@@ -0,0 +1,90 @@
+// 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 main
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"os"
+	"time"
+
+	"github.com/MatusOllah/slogcolor"
+
+	airflow "github.com/apache/airflow/go-sdk/celery/cmd"
+	"github.com/apache/airflow/go-sdk/sdk"
+	"github.com/apache/airflow/go-sdk/worker"
+)
+
+// Ensure the worker knows how to execute the tasks
+func registerTasks(worker worker.Worker) {
+	worker.RegisterTask("tutorial_dag", extract)
+	worker.RegisterTask("tutorial_dag", transform)
+	worker.RegisterTask("tutorial_dag", load)
+}
+
+func extract(ctx context.Context, log *slog.Logger) error {
+	log.Info("Hello from task")
+	for range 10 {
+
+		// Once per loop,.check if we've been asked to cancel!
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+		log.Info("After the beep the time will be", "time", time.Now())
+		time.Sleep(2 * time.Second)
+	}
+	log.Info("Goodbye from task")
+
+	return nil
+}
+
+func transform(ctx context.Context, log *slog.Logger) error {
+	key := "my_variable"
+	val, err := sdk.VariableGet(ctx, key)
+	if err != nil {
+		return err
+	}
+	log.Info("Obtained variable", key, val)
+	return nil
+}
+
+func load() error {
+	return fmt.Errorf("Please fail")
+}
+
+func main() {
+	logger := makeLogger()
+	slog.SetDefault(logger)
+	logger.Debug("Starting up")
+	worker := worker.New(logger)
+	registerTasks(worker)
+
+	airflow.Execute(worker)
+}
+
+func makeLogger() *slog.Logger {
+	opts := *slogcolor.DefaultOptions
+	leveler := &slog.LevelVar{}
+	leveler.Set(slog.LevelDebug)
+	opts.Level = leveler
+	log := slog.New(slogcolor.NewHandler(os.Stderr, &opts))
+	return log
+}
@@ -0,0 +1,62 @@
+module github.com/apache/airflow/go-sdk
+
+go 1.24.1
+
+require (
+	github.com/ivanpirog/coloredcobra v1.0.1
+	github.com/oapi-codegen/runtime v1.1.1
+	github.com/spf13/cobra v1.9.1
+	github.com/spf13/pflag v1.0.6
+	github.com/spf13/viper v1.20.1
+	github.com/stretchr/testify v1.10.0
+	github.com/swaggest/assertjson v1.9.0
+	resty.dev/v3 v3.0.0-beta.2
+)
+
+require (
+	github.com/bool64/shared v0.1.5 // indirect
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
+	github.com/fsnotify/fsnotify v1.8.0 // indirect
+	github.com/go-kit/log v0.2.1 // indirect
+	github.com/go-logfmt/logfmt v0.6.0 // indirect
+	github.com/go-viper/mapstructure/v2 v2.2.1 // indirect
+	github.com/gomodule/redigo v1.8.9 // indirect
+	github.com/iancoleman/orderedmap v0.3.0 // indirect
+	github.com/inconshreveable/mousetrap v1.1.0 // indirect
+	github.com/nxadm/tail v1.4.11 // indirect
+	github.com/pelletier/go-toml/v2 v2.2.3 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/rogpeppe/go-internal v1.12.0 // indirect
+	github.com/sagikazarmark/locafero v0.7.0 // indirect
+	github.com/sergi/go-diff v1.1.0 // indirect
+	github.com/sourcegraph/conc v0.3.0 // indirect
+	github.com/spf13/afero v1.12.0 // indirect
+	github.com/spf13/cast v1.7.1 // indirect
+	github.com/subosito/gotenv v1.6.0 // indirect
+	github.com/yudai/gojsondiff v1.0.0 // indirect
+	github.com/yudai/golcs v0.0.0-20170316035057-ecda9a501e82 // indirect
+	go.uber.org/multierr v1.10.0 // indirect
+	golang.org/x/net v0.33.0 // indirect
+	golang.org/x/sync v0.10.0 // indirect
+	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
+)
+
+require (
+	github.com/MatusOllah/slogcolor v1.6.0
+	github.com/apapsch/go-jsonmerge/v2 v2.0.0 // indirect
+	github.com/fatih/color v1.18.0 // indirect
+	github.com/google/uuid v1.6.0
+	github.com/jarcoal/httpmock v1.4.0
+	github.com/marselester/gopher-celery v1.0.0
+	github.com/mattn/go-colorable v0.1.14 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/redis/go-redis/v9 v9.7.3
+	golang.org/x/sys v0.32.0 // indirect
+	golang.org/x/text v0.21.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)
+
+// https://github.com/ivanpirog/coloredcobra/pull/7
+replace github.com/ivanpirog/coloredcobra => github.com/gabe565/coloredcobra v0.0.0-20240807081640-cadd85f208af
@@ -0,0 +1,154 @@
+github.com/MatusOllah/slogcolor v1.6.0 h1:JAKer0xj5l1jYTXyQvs5ggqmJqYDuLnxgR9jfMAd+sI=
+github.com/MatusOllah/slogcolor v1.6.0/go.mod h1:5y1H50XuQIBvuYTJlmokWi+4FuPiJN5L7Z0jM4K4bYA=
+github.com/RaveNoX/go-jsoncommentstrip v1.0.0/go.mod h1:78ihd09MekBnJnxpICcwzCMzGrKSKYe4AqU6PDYYpjk=
+github.com/apapsch/go-jsonmerge/v2 v2.0.0 h1:axGnT1gRIfimI7gJifB699GoE/oq+F2MU7Dml6nw9rQ=
+github.com/apapsch/go-jsonmerge/v2 v2.0.0/go.mod h1:lvDnEdqiQrp0O42VQGgmlKpxL1AP2+08jFMw88y4klk=
+github.com/bmatcuk/doublestar v1.1.1/go.mod h1:UD6OnuiIn0yFxxA2le/rnRU1G4RaI4UvFv1sNto9p6w=
+github.com/bool64/dev v0.2.29 h1:x+syGyh+0eWtOzQ1ItvLzOGIWyNWnyjXpHIcpF2HvL4=
+github.com/bool64/dev v0.2.29/go.mod h1:iJbh1y/HkunEPhgebWRNcs8wfGq7sjvJ6W5iabL8ACg=
+github.com/bool64/shared v0.1.5 h1:fp3eUhBsrSjNCQPcSdQqZxxh9bBwrYiZ+zOKFkM0/2E=
+github.com/bool64/shared v0.1.5/go.mod h1:081yz68YC9jeFB3+Bbmno2RFWvGKv1lPKkMP6MHJlPs=
+github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs=
+github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c=
+github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA=
+github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0=
+github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
+github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/cpuguy83/go-md2man/v2 v2.0.1/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
+github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78=
+github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc=
+github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
+github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM=
+github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU=
+github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
+github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
+github.com/fsnotify/fsnotify v1.6.0/go.mod h1:sl3t1tCWJFWoRz9R8WJCbQihKKwmorjAbSClcnxKAGw=
+github.com/fsnotify/fsnotify v1.8.0 h1:dAwr6QBTBZIkG8roQaJjGof0pp0EeF+tNV7YBP3F/8M=
+github.com/fsnotify/fsnotify v1.8.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0=
+github.com/gabe565/coloredcobra v0.0.0-20240807081640-cadd85f208af h1:0f/Um4MeHOMJyTdDytEpirTKQHAeBbZCIeukspwCDhQ=
+github.com/gabe565/coloredcobra v0.0.0-20240807081640-cadd85f208af/go.mod h1:iho4nEKcnwZFiniGSdcgdvRgZNjxm+h20acv8vqmN6Q=
+github.com/go-kit/log v0.2.1 h1:MRVx0/zhvdseW+Gza6N9rVzU/IVzaeE1SFI4raAhmBU=
+github.com/go-kit/log v0.2.1/go.mod h1:NwTd00d/i8cPZ3xOwwiv2PO5MOcx78fFErGNcVmBjv0=
+github.com/go-logfmt/logfmt v0.6.0 h1:wGYYu3uicYdqXVgoYbvnkrPVXkuLM1p1ifugDMEdRi4=
+github.com/go-logfmt/logfmt v0.6.0/go.mod h1:WYhtIu8zTZfxdn5+rREduYbwxfcBr/Vr6KEVveWlfTs=
+github.com/go-viper/mapstructure/v2 v2.2.1 h1:ZAaOCxANMuZx5RCeg0mBdEZk7DZasvvZIxtHqx8aGss=
+github.com/go-viper/mapstructure/v2 v2.2.1/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
+github.com/gomodule/redigo v1.8.9 h1:Sl3u+2BI/kk+VEatbj0scLdrFhjPmbxOc1myhDP41ws=
+github.com/gomodule/redigo v1.8.9/go.mod h1:7ArFNvsTjH8GMMzB4uy1snslv2BwmginuMs06a1uzZE=
+github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/iancoleman/orderedmap v0.3.0 h1:5cbR2grmZR/DiVt+VJopEhtVs9YGInGIxAoMJn+Ichc=
+github.com/iancoleman/orderedmap v0.3.0/go.mod h1:XuLcCUkdL5owUCQeF2Ue9uuw1EptkJDkXXS7VoV7XGE=
+github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
+github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
+github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
+github.com/jarcoal/httpmock v1.4.0 h1:BvhqnH0JAYbNudL2GMJKgOHe2CtKlzJ/5rWKyp+hc2k=
+github.com/jarcoal/httpmock v1.4.0/go.mod h1:ftW1xULwo+j0R0JJkJIIi7UKigZUXCLLanykgjwBXL0=
+github.com/juju/gnuflag v0.0.0-20171113085948-2ce1bb71843d/go.mod h1:2PavIy+JPciBPrBUjwbNvtwB6RQlve+hkpll6QSNmOE=
+github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
+github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
+github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
+github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
+github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
+github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
+github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
+github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
+github.com/marselester/gopher-celery v1.0.0 h1:Dlu83eROAhUsJk+SCOJDsjAAlPh/fsN5KVySZb5lHYo=
+github.com/marselester/gopher-celery v1.0.0/go.mod h1:nrZHmlN19mpqs12bVTzdDoDo20GEKx0dYazWzrO7Y5Y=
+github.com/mattn/go-colorable v0.1.9/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
+github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
+github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
+github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU=
+github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94=
+github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
+github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
+github.com/maxatome/go-testdeep v1.14.0 h1:rRlLv1+kI8eOI3OaBXZwb3O7xY3exRzdW5QyX48g9wI=
+github.com/maxatome/go-testdeep v1.14.0/go.mod h1:lPZc/HAcJMP92l7yI6TRz1aZN5URwUBUAfUNvrclaNM=
+github.com/nxadm/tail v1.4.11 h1:8feyoE3OzPrcshW5/MJ4sGESc5cqmGkGCWlco4l0bqY=
+github.com/nxadm/tail v1.4.11/go.mod h1:OTaG3NK980DZzxbRq6lEuzgU+mug70nY11sMd4JXXHc=
+github.com/oapi-codegen/runtime v1.1.1 h1:EXLHh0DXIJnWhdRPN2w4MXAzFyE4CskzhNLUmtpMYro=
+github.com/oapi-codegen/runtime v1.1.1/go.mod h1:SK9X900oXmPWilYR5/WKPzt3Kqxn/uS/+lbpREv+eCg=
+github.com/onsi/ginkgo v1.15.2 h1:l77YT15o814C2qVL47NOyjV/6RbaP7kKdrvZnxQ3Org=
+github.com/onsi/ginkgo v1.15.2/go.mod h1:Dd6YFfwBW84ETqqtL0CPyPXillHgY6XhQH3uuCCTr/o=
+github.com/onsi/gomega v1.11.0 h1:+CqWgvj0OZycCaqclBD1pxKHAU+tOkHmQIWvDHq2aug=
+github.com/onsi/gomega v1.11.0/go.mod h1:azGKhqFUon9Vuj0YmTfLSmx0FUwqXYSTl5re8lQLTUg=
+github.com/pelletier/go-toml/v2 v2.2.3 h1:YmeHyLY8mFWbdkNWwpr+qIL2bEqT0o95WSdkNHvL12M=
+github.com/pelletier/go-toml/v2 v2.2.3/go.mod h1:MfCQTFTvCcUyyvvwm1+G6H/jORL20Xlb6rzQu9GuUkc=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/redis/go-redis/v9 v9.7.3 h1:YpPyAayJV+XErNsatSElgRZZVCwXX9QzkKYNvO7x0wM=
+github.com/redis/go-redis/v9 v9.7.3/go.mod h1:bGUrSggJ9X9GUmZpZNEOQKaANxSGgOEBRltRTZHSvrA=
+github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8=
+github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4=
+github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
+github.com/sagikazarmark/locafero v0.7.0 h1:5MqpDsTGNDhY8sGp0Aowyf0qKsPrhewaLSsFaodPcyo=
+github.com/sagikazarmark/locafero v0.7.0/go.mod h1:2za3Cg5rMaTMoG/2Ulr9AwtFaIppKXTRYnozin4aB5k=
+github.com/sergi/go-diff v1.1.0 h1:we8PVUC3FE2uYfodKH/nBHMSetSfHDR6scGdBi+erh0=
+github.com/sergi/go-diff v1.1.0/go.mod h1:STckp+ISIX8hZLjrqAeVduY0gWCT9IjLuqbuNXdaHfM=
+github.com/sourcegraph/conc v0.3.0 h1:OQTbbt6P72L20UqAkXXuLOj79LfEanQ+YQFNpLA9ySo=
+github.com/sourcegraph/conc v0.3.0/go.mod h1:Sdozi7LEKbFPqYX2/J+iBAM6HpqSLTASQIKqDmF7Mt0=
+github.com/spf13/afero v1.12.0 h1:UcOPyRBYczmFn6yvphxkn9ZEOY65cpwGKb5mL36mrqs=
+github.com/spf13/afero v1.12.0/go.mod h1:ZTlWwG4/ahT8W7T0WQ5uYmjI9duaLQGy3Q2OAl4sk/4=
+github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y=
+github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
+github.com/spf13/cobra v1.4.0/go.mod h1:Wo4iy3BUC+X2Fybo0PDqwJIv3dNRiZLHQymsfxlB84g=
+github.com/spf13/cobra v1.9.1 h1:CXSaggrXdbHK9CF+8ywj8Amf7PBRmPCOJugH954Nnlo=
+github.com/spf13/cobra v1.9.1/go.mod h1:nDyEzZ8ogv936Cinf6g1RU9MRY64Ir93oCnqb9wxYW0=
+github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
+github.com/spf13/pflag v1.0.6 h1:jFzHGLGAlb3ruxLB8MhbI6A8+AQX/2eW4qeyNZXNp2o=
+github.com/spf13/pflag v1.0.6/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
+github.com/spf13/viper v1.20.1 h1:ZMi+z/lvLyPSCoNtFCpqjy0S4kPbirhpTMwl8BkW9X4=
+github.com/spf13/viper v1.20.1/go.mod h1:P9Mdzt1zoHIG8m2eZQinpiBjo6kCmZSKBClNNqjJvu4=
+github.com/spkg/bom v0.0.0-20160624110644-59b7046e48ad/go.mod h1:qLr4V1qq6nMqFKkMo8ZTx3f+BZEkzsRUY10Xsm2mwU0=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
+github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
+github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8=
+github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU=
+github.com/swaggest/assertjson v1.9.0 h1:dKu0BfJkIxv/xe//mkCrK5yZbs79jL7OVf9Ija7o2xQ=
+github.com/swaggest/assertjson v1.9.0/go.mod h1:b+ZKX2VRiUjxfUIal0HDN85W0nHPAYUbYH5WkkSsFsU=
+github.com/yudai/gojsondiff v1.0.0 h1:27cbfqXLVEJ1o8I6v3y9lg8Ydm53EKqHXAOMxEGlCOA=
+github.com/yudai/gojsondiff v1.0.0/go.mod h1:AY32+k2cwILAkW1fbgxQ5mUmMiZFgLIV+FBNExI05xg=
+github.com/yudai/golcs v0.0.0-20170316035057-ecda9a501e82 h1:BHyfKlQyqbsFN5p3IfnEUduWvb9is428/nNb5L3U01M=
+github.com/yudai/golcs v0.0.0-20170316035057-ecda9a501e82/go.mod h1:lgjkn3NuSvDfVJdfcVVdX+jpBxNmX4rDAzaS45IcYoM=
+github.com/yudai/pp v2.0.1+incompatible h1:Q4//iY4pNF6yPLZIigmvcl7k/bPgrcTPIFIcmawg5bI=
+github.com/yudai/pp v2.0.1+incompatible/go.mod h1:PuxR/8QJ7cyCkFp/aUDS+JY727OFEZkTdatxwunjIkc=
+go.uber.org/multierr v1.10.0 h1:S0h4aNzvfcFsC3dRF1jLoaov7oRaKqRGC/pUEJ2yvPQ=
+go.uber.org/multierr v1.10.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
+golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I=
+golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4=
+golang.org/x/sync v0.10.0 h1:3NQrjDixjgGwUOCaF8w2+VYHv0Ve/vGYSbdkTa98gmQ=
+golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
+golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20200223170610-d5e6a3e2c0ae/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220908164124-27713097b956/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.32.0 h1:s77OFDvIQeibCmezSnk/q6iAfkdiQaJi4VzroCFrN20=
+golang.org/x/sys v0.32.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
+golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
+golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
+gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ=
+gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw=
+gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
+gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+resty.dev/v3 v3.0.0-beta.2 h1:xu4mGAdbCLuc3kbk7eddWfWm4JfhwDtdapwss5nCjnQ=
+resty.dev/v3 v3.0.0-beta.2/go.mod h1:OgkqiPvTDtOuV4MGZuUDhwOpkY8enjOsjjMzeOHefy4=
@@ -0,0 +1,73 @@
+// 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 (
+	"fmt"
+	"maps"
+
+	"github.com/google/uuid"
+	"resty.dev/v3"
+)
+
+const API_VERSION = "2025-04-11"
+
+//go:generate go run -modfile=../../tools.mod github.com/ashb/oapi-resty-codegen --config oapi-codegen.yml ../../openapi-fmt.json
+
+func correlationIdInjector(_ *resty.Client, req *resty.Request) error {
+	if uuid, err := uuid.NewV7(); err != nil {
+		return err
+	} else {
+		req.Header.Set("Correlation-Id", uuid.String())
+	}
+	return nil
+}
+
+func NewDefaultClient(server string, opts ...ClientOption) (ClientInterface, error) {
+	rc := resty.New()
+	rc.SetBaseURL(server)
+	rc.SetHeader("Airflow-API-Version", API_VERSION)
+	return NewClient(
+		server,
+		WithClient(rc),
+		WithRequestMiddleware(correlationIdInjector),
+	)
+}
+
+// WithBearerToken creates a copy of the client (reusing the underlying http.Client) adding in a Bearer token auth to all requests
+func (c *Client) WithBearerToken(token string) (ClientInterface, error) {
+	rc := resty.NewWithClient(c.Client.Client())
+	maps.Copy(rc.Header(), c.Client.Header())
+	rc.SetBaseURL(c.Server)
+	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 herej and setting the header directly by-passes that
+	rc.SetHeader("Authorization", fmt.Sprintf("Bearer %s", token))
+
+	opts := []ClientOption{
+		WithClient(rc),
+	}
+	for _, mw := range c.RequestMiddleware {
+		opts = append(opts, WithRequestMiddleware(mw))
+	}
+
+	return NewClient(c.Server, opts...)
+}
@@ -0,0 +1,42 @@
+// 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 (
+	"os"
+	"os/user"
+)
+
+var (
+	Hostname string
+	Username string
+)
+
+func init() {
+	user, err := user.Current()
+	if err != nil {
+		panic(err.Error())
+	}
+
+	Username = user.Username
+
+	Hostname, err = os.Hostname()
+	if err != nil {
+		panic(err.Error())
+	}
+}
@@ -0,0 +1,25 @@
+// 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
+
+type ExecuteTaskWorkload struct {
+	Token      string       `json:"token"`
+	TI         TaskInstance `json:"ti"`
+	BundleInfo BundleInfo   `json:"bundle_info"`
+	LogPath    *string      `json:"log_path,omitempty"`
+}
@@ -0,0 +1,32 @@
+# 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.
+
+# yamllint disable rule:line-length
+# yaml-language-server: $schema=https://raw.githubusercontent.com/oapi-codegen/oapi-codegen/HEAD/configuration-schema.json
+# yamllint enable
+---
+package: api
+output: client.gen.go
+generate:
+  models: true
+  client: true
+output-options:
+  exclude-tags: [Health]
+  # Include the unreferenced schemas/models, such as `TaskInstance`
+  skip-prune: true
+  overlay:
+    path: overlay.yml
@@ -0,0 +1,68 @@
+# 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.
+---
+# yamllint disable rule:line-length
+overlay: 1.0.0
+info:
+  title: OpenAPI Overlay
+  version: 0.0.0
+actions:
+  - description: |-
+      Remove Airflow-API-Version header
+
+      Since the version is consistent across the entire client, we don't want it in the request objects
+    target: $.paths.*.*.parameters[?(@.name == "Airflow-API-Version")]
+    remove: true
+
+  # Customize the operation IDs so it produced nice function names
+  - target: $.paths['/dag-runs/count']
+    update:
+      operationId: get_dag_run_count
+  - target: $.paths['/dag-runs/{dag_id}/{run_id}/state'].get
+    update:
+      operationId: get_dag_run_state
+  - target: $.paths['/task-instances/{task_instance_id}/heartbeat'].put
+    update:
+      operationId: task_instance_heartbeat
+  - target: $.paths['/task-instances/{task_instance_id}/rtif'].put
+    update:
+      operationId: task_instance_put_rendered_fields
+  - target: $.paths['/task-instances/{task_instance_id}/run'].patch
+    update:
+      operationId: task_instance_run
+  - target: $.paths['/task-instances/{task_instance_id}/state'].patch
+    update:
+      operationId: task_instance_update_state
+  - target: $.paths['/task-instances/{task_instance_id}/skip-downstream'].patch
+    update:
+      operationId: task_instance_skip_downstream
+
+  - target: $.paths['/task-instances/{task_instance_id}/state'].patch.requestBody.content['application/json'].schema
+    description: Override UpdateState body to not be an embedded type
+    update:
+      x-go-type: TIUpdateStatePayload
+  - target: $.components.schemas
+    update:
+      TIUpdateStatePayload:
+        oneOf:
+          - $ref: "#/components/schemas/TITerminalStatePayload"
+          - $ref: "#/components/schemas/TISuccessStatePayload"
+          - $ref: "#/components/schemas/TITargetStatePayload"
+          - $ref: "#/components/schemas/TIDeferredStatePayload"
+          - $ref: "#/components/schemas/TIRescheduleStatePayload"
+          - $ref: "#/components/schemas/TIRetryStatePayload"
+        title: "TI UpdateState Payload"
@@ -0,0 +1,29 @@
+// 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 logging
+
+import "log/slog"
+
+// AttrErr returns an [slog.Attr] that will only include the given error if it is non nill
+func AttrErr(err error) slog.Attr {
+	if err == nil {
+		// Return an empty attr
+		return slog.Attr{}
+	}
+	return slog.Any("err", err)
+}
@@ -0,0 +1,57 @@
+// 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 logging
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"runtime"
+	"time"
+)
+
+type RestyLoggerBridge struct {
+	Handler slog.Handler
+	Context context.Context
+}
+
+func (b *RestyLoggerBridge) log(level slog.Level, format string, args ...any) {
+	if !b.Handler.Enabled(b.Context, level) {
+		return
+	}
+	var pcs [1]uintptr
+	runtime.Callers(3, pcs[:]) // skip [Callers, log, Infof etc]
+	r := slog.NewRecord(time.Now(), level, fmt.Sprintf(format, args...), pcs[0])
+	_ = b.Handler.Handle(b.Context, r)
+}
+
+func (b *RestyLoggerBridge) Debugf(format string, v ...any) {
+	b.log(slog.LevelDebug, format, v...)
+}
+
+func (b *RestyLoggerBridge) Infof(format string, v ...any) {
+	b.log(slog.LevelInfo, format, v...)
+}
+
+func (b *RestyLoggerBridge) Warnf(format string, v ...any) {
+	b.log(slog.LevelWarn, format, v...)
+}
+
+func (b *RestyLoggerBridge) Errorf(format string, v ...any) {
+	b.log(slog.LevelError, format, v...)
+}
@@ -0,0 +1,77 @@
+// 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 logging
+
+import (
+	"context"
+	"log/slog"
+)
+
+var _ slog.Handler = (*TeeLogger)(nil)
+
+type TeeLogger struct {
+	handlers []slog.Handler
+}
+
+func NewTeeLogger(handlers ...slog.Handler) *TeeLogger {
+	return &TeeLogger{
+		handlers: handlers,
+	}
+}
+
+func (t *TeeLogger) Enabled(ctx context.Context, level slog.Level) bool {
+	for _, handler := range t.handlers {
+		if handler.Enabled(ctx, level) {
+			return true
+		}
+	}
+
+	return false
+}
+
+func (t *TeeLogger) Handle(ctx context.Context, record slog.Record) error {
+	for _, handler := range t.handlers {
+		if err := handler.Handle(ctx, record); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (t *TeeLogger) WithAttrs(attrs []slog.Attr) slog.Handler {
+	ret := &TeeLogger{
+		handlers: make([]slog.Handler, len(t.handlers)),
+	}
+	for k, handler := range t.handlers {
+		ret.handlers[k] = handler.WithAttrs(attrs)
+	}
+
+	return ret
+}
+
+func (t *TeeLogger) WithGroup(name string) slog.Handler {
+	ret := &TeeLogger{
+		handlers: make([]slog.Handler, len(t.handlers)),
+	}
+	for k, handler := range t.handlers {
+		ret.handlers[k] = handler.WithGroup(name)
+	}
+
+	return ret
+}
@@ -0,0 +1,35 @@
+// 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 sdkcontext
+
+type (
+	workoadContextKey   struct{}
+	apiClientContextKey struct{}
+	workerContextKey    struct{}
+	runtimeTIContextKey struct{}
+)
+
+var (
+	// WorkloadContextKey stores the current workload
+	//
+	// workload := ctx.Value(sdkcontext.WorkloadContextKey).(api.ExecuteTaskWorkload)
+	WorkloadContextKey  = workoadContextKey{}
+	RuntimeTIContextKey = runtimeTIContextKey{}
+	ApiClientContextKey = apiClientContextKey{}
+	WorkerContextKey    = workerContextKey{}
+)
@@ -0,0 +1,21 @@
+// 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.
+
+/*
+sdk implements the interface for Go Tasks to access the Airflow state
+*/
+package sdk
@@ -0,0 +1,43 @@
+// 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 sdk
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/apache/airflow/go-sdk/pkg/api"
+	"github.com/apache/airflow/go-sdk/pkg/sdkcontext"
+)
+
+func VariableGet(ctx context.Context, key string) (string, error) {
+	client := ctx.Value(sdkcontext.ApiClientContextKey).(api.ClientInterface)
+
+	resp, err := client.Variables().Get(ctx, key)
+	if err != nil {
+		var httpError *api.GeneralHTTPError
+		if errors.As(err, &httpError) && httpError.Response.StatusCode() == 404 {
+			// TODO: return a custom error message!
+			return "", fmt.Errorf("variable %q not found: %d", key, httpError.Response.StatusCode())
+		}
+		return "", err
+	}
+	// TODO: Handle deserialization etc!
+	return *resp.Value, nil
+}
@@ -0,0 +1,45 @@
+// 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 worker
+
+import (
+	"os"
+	"os/user"
+)
+
+var (
+	Hostname string
+	Username string
+	PID      int
+)
+
+func init() {
+	user, err := user.Current()
+	if err != nil {
+		panic(err.Error())
+	}
+
+	Username = user.Username
+
+	Hostname, err = os.Hostname()
+	if err != nil {
+		panic(err.Error())
+	}
+
+	PID = os.Getpid()
+}
@@ -0,0 +1,98 @@
+// 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 worker
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"reflect"
+	"runtime"
+	"strings"
+	"sync"
+)
+
+type (
+	// task is an interface of an task implementation.
+	Task interface {
+		Execute(ctx context.Context, logger *slog.Logger) error
+	}
+	registry struct {
+		sync.RWMutex
+		taskFuncMap map[string]map[string]Task
+	}
+)
+
+func newRegistry() *registry {
+	return &registry{taskFuncMap: make(map[string]map[string]Task)}
+}
+
+func getFnName(fn reflect.Value) string {
+	fullName := runtime.FuncForPC(fn.Pointer()).Name()
+	parts := strings.Split(fullName, ".")
+	fnName := parts[len(parts)-1]
+	// Go adds `-fm` suffix to a method names
+	return strings.TrimSuffix(fnName, "-fm")
+}
+
+func (r *registry) RegisterTask(dagId string, fn any) {
+	val := reflect.ValueOf(fn)
+
+	if val.Kind() != reflect.Func {
+		panic(fmt.Errorf("task fn was a %s, not a func", val.Kind()))
+	}
+
+	fnName := getFnName(val)
+
+	r.RegisterTaskWithName(dagId, fnName, fn)
+}
+
+func (r *registry) RegisterTaskWithName(dagId, taskId string, fn any) {
+	task, err := NewTaskFunction(fn)
+	if err != nil {
+		panic(fmt.Errorf("error registering task %q for DAG %q: %w", taskId, dagId, err))
+	}
+
+	r.RWMutex.Lock()
+	defer r.RWMutex.Unlock()
+
+	dagTasks, exists := r.taskFuncMap[dagId]
+
+	if !exists {
+		dagTasks = make(map[string]Task)
+		r.taskFuncMap[dagId] = dagTasks
+	}
+
+	_, exists = dagTasks[taskId]
+	if exists {
+		panic(fmt.Errorf("taskId %q is already registered for DAG %q", taskId, dagId))
+	}
+	dagTasks[taskId] = task
+}
+
+func (r *registry) LookupTask(dagId, taskId string) (task Task, exists bool) {
+	r.RLock()
+	defer r.RUnlock()
+
+	dagTasks, exists := r.taskFuncMap[dagId]
+	if !exists {
+		return nil, false
+	}
+	task, exists = dagTasks[taskId]
+	return task, exists
+}
@@ -0,0 +1,67 @@
+// 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 worker
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func A_B_c() error {
+	return nil
+}
+
+func NotErrorRet() int {
+	return 0
+}
+
+func TestRegisterName(t *testing.T) {
+	reg := newRegistry()
+	reg.RegisterTask("mydag", A_B_c)
+
+	task, found := reg.LookupTask("mydag", "A_B_c")
+	assert.True(t, found)
+	assert.NotNil(t, task)
+}
+
+func TestNotAFunction(t *testing.T) {
+	reg := newRegistry()
+	assert.PanicsWithError(t, "task fn was a string, not a func", func() {
+		reg.RegisterTask("mydag", "not a fn")
+	})
+}
+
+func TestTaskReturnNotError(t *testing.T) {
+	reg := newRegistry()
+	assert.PanicsWithError(
+		t,
+		"error registering task \"NotErrorRet\" for DAG \"mydag\": expected task function github.com/apache/airflow/go-sdk/worker.NotErrorRet last return value to return error but found int",
+		func() {
+			reg.RegisterTask("mydag", NotErrorRet)
+		},
+	)
+}
+
+func TestRegisterDuplicateTask(t *testing.T) {
+	reg := newRegistry()
+	reg.RegisterTask("mydag", A_B_c)
+	assert.PanicsWithError(t, `taskId "A_B_c" is already registered for DAG "mydag"`, func() {
+		reg.RegisterTask("mydag", A_B_c)
+	})
+}
@@ -0,0 +1,402 @@
+// 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 worker
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"log/slog"
+	"os"
+	"path"
+	"runtime/debug"
+	"time"
+
+	"github.com/spf13/viper"
+
+	"github.com/apache/airflow/go-sdk/pkg/api"
+	"github.com/apache/airflow/go-sdk/pkg/logging"
+	"github.com/apache/airflow/go-sdk/pkg/sdkcontext"
+)
+
+type (
+	Registry interface {
+		// RegisterTask registers a new task with this task registry
+		// A task takes a [context.Context] and input and returns a (result, error) or just error.
+		//
+		// This method panics if taskFunc doesn't comply with the expected format or tries to register a duplicate task
+		RegisterTask(dagid string, fn any)
+
+		RegisterTaskWithName(dagId, taskId string, fn any)
+
+		LookupTask(dagId, taskId string) (task Task, exists bool)
+	}
+
+	Worker interface {
+		Registry
+
+		ExecuteTaskWorkload(ctx context.Context, workload api.ExecuteTaskWorkload) error
+
+		WithServer(server string) (Worker, error)
+	}
+
+	worker struct {
+		Registry
+		client            api.ClientInterface
+		logger            *slog.Logger
+		heartbeatInterval time.Duration
+		reportTimeout     time.Duration
+	}
+)
+
+var ErrTaskCancelledAfterFailedHeartbeat = errors.New("task cancelled after failed heartbeat")
+
+func New(logger *slog.Logger) Worker {
+	return &worker{
+		logger:            logger,
+		heartbeatInterval: HeartbeatInterval,
+		reportTimeout:     ReportTimeout,
+		Registry:          newRegistry(),
+	}
+}
+
+const (
+	HeartbeatInterval = 5 * time.Second
+	ReportTimeout     = 10 * time.Second
+)
+
+func (w *worker) WithServer(server string) (Worker, error) {
+	client, err := api.NewDefaultClient(server)
+	if err != nil {
+		return nil, err
+	}
+	return &worker{
+		Registry:          w.Registry,
+		client:            client,
+		logger:            w.logger,
+		heartbeatInterval: w.heartbeatInterval,
+		reportTimeout:     w.reportTimeout,
+	}, nil
+}
+
+type heartbeater struct {
+	heartbeatInterval time.Duration
+	logger            *slog.Logger
+	taskCanceller     context.CancelCauseFunc
+	taskLogger        *slog.Logger
+}
+
+// Run a heartbeat loop until our ctx is Done
+//
+// If there is an error reporting too many heartbeats then the task context will be cancelled to stop the task.
+func (h *heartbeater) Run(
+	ctx context.Context,
+) {
+	h.logger.DebugContext(ctx, "Starting heartbeater", "heartbeat", h.heartbeatInterval)
+
+	workload := ctx.Value(sdkcontext.WorkloadContextKey).(api.ExecuteTaskWorkload)
+	client := ctx.Value(sdkcontext.ApiClientContextKey).(api.ClientInterface)
+
+	ticker := time.NewTicker(h.heartbeatInterval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			// TODO: Record when last successful heartbeat was, and fail+abort if we haven't managed to record
+			// one recently enough
+
+			err := client.TaskInstances().
+				Heartbeat(ctx, workload.TI.Id, &api.TIHeartbeatInfo{
+					Hostname: Hostname,
+					Pid:      os.Getpid(),
+				})
+			if err != nil {
+				h.logger.InfoContext(ctx, "heartbeating", slog.Any("error", err))
+			}
+
+			var httpError *api.GeneralHTTPError
+			if errors.As(err, &httpError) {
+				resp := httpError.Response
+				if resp != nil && resp.StatusCode() == 404 || resp.StatusCode() == 409 {
+					h.logger.ErrorContext(
+						ctx,
+						"Server indicated the task shouldn't be running anymore",
+						"status_code", resp.StatusCode(),
+						"details", httpError.JSON,
+					)
+					// Log something in the task log file too
+					h.taskLogger.ErrorContext(
+						ctx,
+						"Server indicated the task shouldn't be running anymore. Terminating workload",
+						"status_code",
+						resp.StatusCode(),
+						"details",
+						httpError.JSON,
+					)
+
+					h.taskCanceller(ErrTaskCancelledAfterFailedHeartbeat)
+					return
+				}
+
+			}
+		case <-ctx.Done():
+			// Something signaled us to stop.
+			return
+		}
+	}
+}
+
+func (w *worker) ExecuteTaskWorkload(ctx context.Context, workload api.ExecuteTaskWorkload) error {
+	// Store the workload in the context so we can get at task id, etc, variables
+	taskContext, cancelTaskCtx := context.WithCancelCause(
+		context.WithValue(ctx, sdkcontext.WorkloadContextKey, workload),
+	)
+	defer cancelTaskCtx(context.Canceled)
+
+	// Task Logger is for the task's _own_ logs. We want a logger for our heartbeat etc.
+	logger := w.logger.With(
+		slog.String("dag_id", workload.TI.DagId),
+		slog.String("task_id", workload.TI.TaskId),
+		slog.String("ti_id", workload.TI.Id.String()),
+	)
+
+	workloadClient := w.client
+	if c, ok := workloadClient.(*api.Client); ok {
+		var err error
+		workloadClient, err = c.WithBearerToken(workload.Token)
+		if err != nil {
+			logger.ErrorContext(ctx, "Could not create client", slog.Any("error", err))
+			return err
+		}
+
+		c = workloadClient.(*api.Client)
+		c.Client.SetLogger(&logging.RestyLoggerBridge{Handler: logger.Handler(), Context: ctx})
+		c.Client.SetDebug(viper.GetBool("api_client.debug"))
+	}
+
+	reportStateFailed := func() error {
+		body := &api.TIUpdateStatePayload{}
+		body.FromTITerminalStatePayload(api.TITerminalStatePayload{
+			State:   api.TerminalStateNonSuccess(api.TerminalTIStateFailed),
+			EndDate: time.Now().UTC(),
+		})
+		return workloadClient.TaskInstances().UpdateState(
+			ctx,
+			workload.TI.Id,
+			body,
+		)
+	}
+
+	// Store the configured API client in the context so we can get it out for accessing Variables etc.
+	taskContext = context.WithValue(taskContext, sdkcontext.ApiClientContextKey, workloadClient)
+
+	taskLogger, err := w.setupTaskLogger(ctx, workload)
+	if err != nil {
+		logger.ErrorContext(taskContext, "Could not create logger", slog.Any("error", err))
+		_ = reportStateFailed()
+		return err
+	}
+
+	task, exists := w.LookupTask(workload.TI.DagId, workload.TI.TaskId)
+	if !exists {
+		taskLogger.ErrorContext(
+			taskContext,
+			"Task not registered",
+			"dag_id",
+			workload.TI.DagId,
+			"task_id",
+			workload.TI.TaskId,
+		)
+		return reportStateFailed()
+	}
+
+	// TODO: Timeout etc on the context
+	// TODO: Add in retries on the api client
+
+	runtimeContext, err := workloadClient.TaskInstances().
+		Run(ctx, workload.TI.Id, &api.TIEnterRunningPayload{
+			Hostname:  Hostname,
+			Unixname:  Username,
+			Pid:       PID,
+			State:     api.Running,
+			StartDate: time.Now().UTC(),
+		})
+	logger.LogAttrs(
+		ctx,
+		slog.LevelInfo,
+		"Start context",
+		slog.Any("resp", runtimeContext),
+		logging.AttrErr(err),
+	)
+	if err != nil {
+		var httpError *api.GeneralHTTPError
+		if errors.As(err, &httpError) {
+			resp := httpError.Response
+			taskLogger.ErrorContext(
+				taskContext,
+				"Server reported error when attempting to start task",
+				slog.Any("error", httpError),
+				slog.Int("status_code", resp.StatusCode()),
+				slog.Group(
+					"request",
+					"url",
+					resp.Request.URL,
+					"method",
+					resp.Request.Method,
+					"headers",
+					resp.Request.Header,
+				),
+			)
+		}
+		return err
+	}
+
+	taskContext = context.WithValue(taskContext, sdkcontext.RuntimeTIContextKey, runtimeContext)
+
+	// Make the heartbreat Context a child of the task context, so it finishes automatically when the task
+	// context does.
+	heartbeatCtx, stopHeartbeating := context.WithCancel(taskContext)
+	defer stopHeartbeating()
+
+	defer func() {
+		if r := recover(); r != nil {
+			stopHeartbeating()
+
+			// Create a new context that isn't cancelled to give us time to report
+			ctx, cancel := context.WithTimeout(context.WithoutCancel(taskContext), w.reportTimeout)
+			defer cancel()
+			taskLogger.ErrorContext(
+				ctx,
+				"Recovered in task",
+				slog.Any("error", r),
+				"stack",
+				string(debug.Stack()),
+			)
+			_ = reportStateFailed()
+		}
+	}()
+
+	heartbeater := &heartbeater{
+		heartbeatInterval: w.heartbeatInterval,
+		logger:            logger,
+		taskCanceller:     cancelTaskCtx,
+		taskLogger:        taskLogger,
+	}
+
+	go heartbeater.Run(heartbeatCtx)
+
+	err = task.Execute(taskContext, taskLogger)
+	endTime := time.Now().UTC()
+
+	stopHeartbeating()
+
+	var finalState api.TerminalTIState
+	body := &api.TIUpdateStatePayload{}
+
+	if taskContext.Err() == ErrTaskCancelledAfterFailedHeartbeat {
+		// We've already logged when we failed to heartbeat, don't do it again
+		finalState = api.TerminalTIStateFailed
+		body.FromTITerminalStatePayload(api.TITerminalStatePayload{
+			State:   api.TerminalStateNonSuccess(finalState),
+			EndDate: time.Now().UTC(),
+		})
+	} else if err != nil {
+		logger.InfoContext(ctx, "Task returned an error, execution complete")
+		taskLogger.InfoContext(
+			taskContext,
+			"Task failed with error, marking task as failed",
+			"error",
+			err,
+			"type",
+			fmt.Sprintf("%T", err),
+		)
+		finalState = api.TerminalTIStateFailed
+		body.FromTITerminalStatePayload(api.TITerminalStatePayload{
+			State:   api.TerminalStateNonSuccess(finalState),
+			EndDate: time.Now().UTC(),
+		})
+	} else {
+		logger.InfoContext(ctx, "Task succeeded")
+
+		taskLogger.InfoContext(
+			taskContext,
+			"Task succeeded",
+		)
+		finalState = api.TerminalTIStateSuccess
+		body.FromTISuccessStatePayload(api.TISuccessStatePayload{
+			EndDate: endTime,
+			State:   api.TISuccessStatePayloadState(finalState),
+		})
+	}
+
+	err = workloadClient.TaskInstances().UpdateState(
+		ctx,
+		workload.TI.Id,
+		body,
+	)
+	if err != nil {
+		taskLogger.Error(
+			"Error reporting final state to server",
+			"error",
+			err,
+			"final_state",
+			finalState,
+		)
+		return err
+	}
+
+	return err
+}
+
+func (w *worker) setupTaskLogger(
+	_ context.Context,
+	workload api.ExecuteTaskWorkload,
+) (*slog.Logger, error) {
+	// Create a logger that:
+	// - only exits the go-routine on panic, not the whole program
+	// - Writes JSON to a file
+	// - And streams output to stdout in a nice format too
+
+	base := viper.GetString("logging.base_log_path")
+	filename := path.Join(base, *workload.LogPath)
+	dir := path.Dir(filename)
+
+	// TODO: umask?
+	err := os.MkdirAll(dir, 0o750)
+	if err != nil {
+		return nil, err
+	}
+
+	fh, err := os.Create(filename)
+	if err != nil {
+		return nil, err
+	}
+	taskLogger := slog.New(
+		logging.NewTeeLogger(
+			slog.NewJSONHandler(
+				fh,
+				&slog.HandlerOptions{
+					AddSource:   false,
+					ReplaceAttr: nil,
+				},
+			),
+			slog.Default().Handler(),
+		),
+	)
+	return taskLogger, nil
+}
@@ -0,0 +1,211 @@
+// 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 worker
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log/slog"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/jarcoal/httpmock"
+	"github.com/stretchr/testify/suite"
+	"github.com/swaggest/assertjson"
+
+	"github.com/apache/airflow/go-sdk/pkg/api"
+)
+
+const ExecutionAPIServer = "http://localhost:9999/execution"
+
+var TestWorkload = api.ExecuteTaskWorkload{
+	Token: "",
+	// {"context_carrier":{},"dag_id":"tutorial_dag","hostname":null,"id":"0196ab8a-5c97-7d4f-b431-e3f49ce20b7f","map_index":-1,"run_id":"manual__2025-05-07T15:48:39.420678+00:00","task_id":"extract","try_number":5}
+	TI: api.TaskInstance{
+		ContextCarrier: new(map[string]any),
+		DagId:          "tutorial_dag",
+		RunId:          "manual__2025-05-07T15:48:39.420678+00:00",
+		TaskId:         "extract",
+		Id:             uuid.MustParse("0196ab8a-5c97-7d4f-b431-e3f49ce20b7f"),
+		// MapIndex:       new(int),
+		TryNumber: 1,
+	},
+	BundleInfo: api.BundleInfo{
+		Name:    "example_dags",
+		Version: nil,
+	},
+	// LogPath: new(string),
+}
+
+func init() {
+	idx := -1
+	log := "dag_id=tutorial_dag/run_id=manual__2025-05-07T15:48:39.420678+00:00/task_id=extract/attempt=5.log"
+	TestWorkload.TI.MapIndex = &idx
+	TestWorkload.LogPath = &log
+}
+
+type WorkerSuite struct {
+	suite.Suite
+	worker    Worker
+	transport *httpmock.MockTransport
+}
+
+func TestWorkerSuite(t *testing.T) {
+	suite.Run(t, &WorkerSuite{})
+}
+
+func (s *WorkerSuite) SetupTest() {
+	s.worker = New(slog.Default())
+
+	s.transport = httpmock.NewMockTransport()
+	client, err := api.NewClient(ExecutionAPIServer, api.WithRoundTripper(s.transport))
+	s.Require().NoError(err)
+
+	s.worker.(*worker).client = client
+}
+
+func (s *WorkerSuite) TestWithServer() {
+	s.worker.(*worker).heartbeatInterval = 100 * time.Millisecond
+	iface, err := s.worker.WithServer("http://abc.com")
+
+	s.Require().NoError(err)
+	w := iface.(*worker)
+	s.Equal(100*time.Millisecond, w.heartbeatInterval)
+	s.Equal(w.client.(*api.Client).BaseURL(), "http://abc.com")
+}
+
+// ExpectTaskRun sets up  a matcher for the "/task-instances/{id}/run" end point and adds a finalize check
+// that it has been called
+func (s *WorkerSuite) ExpectTaskRun() {
+	s.transport.RegisterResponder(
+		http.MethodPatch,
+		"=~^"+ExecutionAPIServer+"/task-instances/[0-9a-f-]{36}/run",
+		httpmock.NewJsonResponderOrPanic(200, map[string]any{}),
+	)
+	s.T().Cleanup(func() {
+		callCounts := s.transport.GetCallCountInfo()
+
+		s.Equal(callCounts["PATCH =~^"+ExecutionAPIServer+"/task-instances/[0-9a-f-]{36}/run"], 1)
+	})
+}
+
+// ExpecteTaskSta sets up  a matcher for the "/task-instances/{id}/state" with the given state end point and adds a finalize check
+// that it has been called
+func (s *WorkerSuite) ExpectTaskState(state any) {
+	s.transport.RegisterMatcherResponder(
+		http.MethodPatch,
+		"=~^"+ExecutionAPIServer+"/task-instances/[0-9a-f-]{36}/state",
+		s.BodyJSONMatches(fmt.Appendf(nil, `{"state": %q}`, state)),
+		httpmock.NewJsonResponderOrPanic(200, map[string]any{}),
+	)
+	s.T().Cleanup(func() {
+		callCounts := s.transport.GetCallCountInfo()
+
+		s.Equal(
+			callCounts["PATCH =~^"+ExecutionAPIServer+"/task-instances/[0-9a-f-]{36}/state <BodyJSONMatches>"],
+			1,
+			"Actual call counts: %#v",
+			callCounts,
+		)
+	})
+}
+
+// BodyContainsJSON creates an httpmock Matcher that will check that the http.Request body contains the given
+// JSON fields
+//
+// The request can contain extra JSON fields. See [github.com/swaggest/assertjson.Matches] for more info
+func (s *WorkerSuite) BodyJSONMatches(expected []byte) httpmock.Matcher {
+	matcher := httpmock.NewMatcher("BodyJSONMatches", func(req *http.Request) bool {
+		b, err := io.ReadAll(req.Body)
+		return err == nil && assertjson.Matches(s.T(), expected, b)
+	})
+	return matcher
+}
+
+// TestTaskNotRegisteredErrors checks that when a task cannot be found we report "success" on the Workload but
+// report the task as failed to the Execution API server
+func (s *WorkerSuite) TestTaskNotRegisteredErrors() {
+	s.transport.RegisterMatcherResponder(
+		http.MethodPatch,
+		"=~^"+ExecutionAPIServer+"/task-instances/[0-9a-f-]{36}/state",
+		s.BodyJSONMatches([]byte(`{"state": "failed"}`)),
+		httpmock.NewJsonResponderOrPanic(200, map[string]any{}),
+	)
+	err := s.worker.ExecuteTaskWorkload(context.Background(), TestWorkload)
+
+	s.NoError(err, "ExecuteTaskWorkload should not report an error")
+
+	s.Equal(1, s.transport.GetTotalCallCount(), "State was reported to server")
+}
+
+// TestStartContextErrorTaskDoesntStart checks that if the /run endpoint returns an error that task doesn't
+// start, but that it is logged
+func (s *WorkerSuite) TestStartContextErrorTaskDoesntStart() {
+	s.T().Fatal("TODO")
+}
+
+// TestTaskPanicReportsFailedState tests that when the task/user code panics that we catch it and report thr
+// error upstream
+func (s *WorkerSuite) TestTaskPanicReportsFailedState() {
+	s.T().Fatal("TODO")
+}
+
+func (s *WorkerSuite) TestTaskReturnErrorReportsFailedState() {
+	s.T().Fatal("TODO")
+}
+
+func (s *WorkerSuite) TestTaskHeartbeatsWhlieRunning() {
+	s.worker.RegisterTaskWithName("tutorial_dag", "extract", func() error {
+		time.Sleep(time.Second)
+		return nil
+	})
+
+	s.ExpectTaskRun()
+	s.ExpectTaskState(api.TerminalTIStateSuccess)
+	s.transport.RegisterResponder(
+		http.MethodPut,
+		"=~^"+ExecutionAPIServer+"/task-instances/[0-9a-f-]{36}/heartbeat",
+		httpmock.NewJsonResponderOrPanic(200, map[string]any{}),
+	)
+
+	s.worker.(*worker).heartbeatInterval = 100 * time.Millisecond
+	err := s.worker.ExecuteTaskWorkload(context.Background(), TestWorkload)
+	s.NoError(err, "ExecuteTaskWorkload should not report an error")
+
+	callCounts := s.transport.GetCallCountInfo()
+
+	// Since we heartbeat every 100ms and run for 1 second, we should expect 10 heartbeat calls. But allow +/-
+	// 1 due to timing imprecision
+	s.InDelta(
+		10,
+		callCounts["PUT =~^"+ExecutionAPIServer+"/task-instances/[0-9a-f-]{36}/heartbeat"],
+		1,
+		"Actual call counts: %#v", callCounts,
+	)
+}
+
+func (s *WorkerSuite) TestTaskHeatbeatErrorStopsTaskAndLogs() {
+	s.T().Fatal("TODO")
+}
+
+func (s *WorkerSuite) TestTokenRefreshHeaderRespected() {
+	s.T().Fatal("TODO")
+}
@@ -0,0 +1,139 @@
+// 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 worker
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"reflect"
+	"runtime"
+)
+
+type taskFunction struct {
+	fn       reflect.Value
+	fullName string
+}
+
+var _ Task = (*taskFunction)(nil)
+
+func NewTaskFunction(fn any) (Task, error) {
+	v := reflect.ValueOf(fn)
+	fullName := runtime.FuncForPC(v.Pointer()).Name()
+	f := &taskFunction{v, fullName}
+	return f, f.validateFn(v.Type())
+}
+
+func (f *taskFunction) Execute(ctx context.Context, logger *slog.Logger) error {
+	fnType := f.fn.Type()
+
+	reflectArgs := make([]reflect.Value, fnType.NumIn())
+	for i := range reflectArgs {
+		in := fnType.In(i)
+
+		switch {
+		case isContext(in):
+			reflectArgs[i] = reflect.ValueOf(ctx)
+		case isLogger(in):
+			reflectArgs[i] = reflect.ValueOf(logger)
+		default:
+			// TODO: deal with other value types. For now they will all be Zero values unless it's a context
+			reflectArgs[i] = reflect.Zero(in)
+		}
+	}
+	slog.Debug("Attempting to call fn", "fn", f.fn, "args", reflectArgs)
+	retValues := f.fn.Call(reflectArgs)
+
+	var err error
+	if errResult := retValues[len(retValues)-1].Interface(); errResult != nil {
+		var ok bool
+		if err, ok = errResult.(error); !ok {
+			return fmt.Errorf(
+				"failed to extract task error result as it is not of error interface: %v",
+				errResult,
+			)
+		}
+	}
+	// If there are two results, convert the first only if it's not a nil pointer
+	var res any
+	if len(retValues) > 1 && (retValues[0].Kind() != reflect.Ptr || !retValues[0].IsNil()) {
+		res = retValues[0].Interface()
+	}
+	// TODO: send the result to XCom
+	_ = res
+	return err
+}
+
+func (f *taskFunction) validateFn(fnType reflect.Type) error {
+	if fnType.Kind() != reflect.Func {
+		return fmt.Errorf("expected a func as input but was %s", fnType.Kind())
+	}
+
+	// Return values
+	//     `<result>, error`,  or just `error`
+	if fnType.NumOut() < 1 || fnType.NumOut() > 2 {
+		return fmt.Errorf(
+			"task function %s has %d return values, must be `<result>, error` or just `error`",
+			f.fullName,
+			fnType.NumOut(),
+		)
+	}
+	if fnType.NumOut() > 1 && !isValidResultType(fnType.Out(0)) {
+		return fmt.Errorf(
+			"expected task function %s first return value to return valid type but found: %v",
+			f.fullName,
+			fnType.Out(0).Kind(),
+		)
+	}
+	if !isError(fnType.Out(fnType.NumOut() - 1)) {
+		return fmt.Errorf(
+			"expected task function %s last return value to return error but found %v",
+			f.fullName,
+			fnType.Out(fnType.NumOut()-1).Kind(),
+		)
+	}
+	return nil
+}
+
+func isValidResultType(inType reflect.Type) bool {
+	// https://golang.org/pkg/reflect/#Kind
+	switch inType.Kind() {
+	case reflect.Func, reflect.Chan, reflect.UnsafePointer:
+		return false
+	}
+
+	return true
+}
+
+var (
+	errorType      = reflect.TypeOf((*error)(nil)).Elem()
+	contextType    = reflect.TypeOf((*context.Context)(nil)).Elem()
+	slogLoggerType = reflect.TypeOf((*slog.Logger)(nil))
+)
+
+func isError(inType reflect.Type) bool {
+	return inType != nil && inType.Implements(errorType)
+}
+
+func isContext(inType reflect.Type) bool {
+	return inType != nil && inType.Implements(contextType)
+}
+
+func isLogger(inType reflect.Type) bool {
+	return inType != nil && inType.AssignableTo(slogLoggerType)
+}