diff --git a/.editorconfig b/.editorconfig index b3084d3f3aa..bdaedb2630b 100644 --- a/.editorconfig +++ b/.editorconfig @@ -58,3 +58,7 @@ indent_size = 2 [*.json] indent_size = 4 + +[*.go] +indent_style = tab +max_line_length = 110 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8795357cbdf..56718804648 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -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 diff --git a/.rat-excludes b/.rat-excludes index 75070872095..dd845ce28ff 100644 --- a/.rat-excludes +++ b/.rat-excludes @@ -174,3 +174,7 @@ auth_generated.py # hash files www-hash.txt + +# go setup files +go.mod +go.sum diff --git a/contributing-docs/08_static_code_checks.rst b/contributing-docs/08_static_code_checks.rst index 0741dabd0e8..bc343673245 100644 --- a/contributing-docs/08_static_code_checks.rst +++ b/contributing-docs/08_static_code_checks.rst @@ -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 | | +-----------------------------------------------------------+--------------------------------------------------------+---------+ diff --git a/dev/breeze/doc/images/output_static-checks.svg b/dev/breeze/doc/images/output_static-checks.svg index af541784efd..39e1fe79bde 100644 --- a/dev/breeze/doc/images/output_static-checks.svg +++ b/dev/breeze/doc/images/output_static-checks.svg @@ -362,11 +362,11 @@ │check-urlparse-usage-in-code | check-xml | check-zip-file-is-not-committed |     │ │codespell | compile-fab-assets | compile-ui-assets | compile-ui-assets-dev |     │ │create-missing-init-py-files-tests | debug-statements | detect-private-key |     │ -│doctoc | end-of-file-fixer | fix-encoding-pragma | flynt |                       │ +│doctoc | end-of-file-fixer | fix-encoding-pragma | flynt | gci |                 │ │generate-airflow-diagrams | generate-airflowctl-datamodels |                     │ │generate-airflowctl-help-images | generate-openapi-spec |                        │ │generate-openapi-spec-fab | generate-pypi-readme | generate-tasksdk-datamodels | │ -│generate-volumes-for-sources | identity | insert-license | kubeconform |         │ +│generate-volumes-for-sources | gofmt | identity | insert-license | kubeconform | │ │lint-chart-schema | lint-dockerfile | lint-helm-chart | lint-json-schema |       │ │lint-markdown | mixed-line-ending | mypy-airflow-core | mypy-airflow-ctl |       │ │mypy-dev | mypy-devel-common | mypy-providers | mypy-task-sdk |                  │ diff --git a/dev/breeze/doc/images/output_static-checks.txt b/dev/breeze/doc/images/output_static-checks.txt index 19fbb879951..555ae8e5508 100644 --- a/dev/breeze/doc/images/output_static-checks.txt +++ b/dev/breeze/doc/images/output_static-checks.txt @@ -1 +1 @@ -44ddf285d8b0b2f0369d4750ace7dde9 +06e827be83fd345223ec466cb843ccc0 diff --git a/dev/breeze/src/airflow_breeze/pre_commit_ids.py b/dev/breeze/src/airflow_breeze/pre_commit_ids.py index a4fc570dcec..fbef1499867 100644 --- a/dev/breeze/src/airflow_breeze/pre_commit_ids.py +++ b/dev/breeze/src/airflow_breeze/pre_commit_ids.py @@ -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", diff --git a/go-sdk/README.md b/go-sdk/README.md new file mode 100644 index 00000000000..a3f55246fa0 --- /dev/null +++ b/go-sdk/README.md @@ -0,0 +1,62 @@ + + +# 🚧 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) diff --git a/go-sdk/celery/app.go b/go-sdk/celery/app.go new file mode 100644 index 00000000000..9499ae46964 --- /dev/null +++ b/go-sdk/celery/app.go @@ -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 +} diff --git a/go-sdk/celery/cmd/root.go b/go-sdk/celery/cmd/root.go new file mode 100644 index 00000000000..7a2c77c3dbc --- /dev/null +++ b/go-sdk/celery/cmd/root.go @@ -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)) + } + }) +} diff --git a/go-sdk/celery/cmd/run.go b/go-sdk/celery/cmd/run.go new file mode 100644 index 00000000000..dd48b9ce290 --- /dev/null +++ b/go-sdk/celery/cmd/run.go @@ -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") +} diff --git a/go-sdk/celery/config.go b/go-sdk/celery/config.go new file mode 100644 index 00000000000..7f756b00d1a --- /dev/null +++ b/go-sdk/celery/config.go @@ -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"` +} diff --git a/go-sdk/example/main.go b/go-sdk/example/main.go new file mode 100644 index 00000000000..34022ce0061 --- /dev/null +++ b/go-sdk/example/main.go @@ -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 +} diff --git a/go-sdk/go.mod b/go-sdk/go.mod new file mode 100644 index 00000000000..0d61aa9fb61 --- /dev/null +++ b/go-sdk/go.mod @@ -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 diff --git a/go-sdk/go.sum b/go-sdk/go.sum new file mode 100644 index 00000000000..7555eecbe58 --- /dev/null +++ b/go-sdk/go.sum @@ -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= diff --git a/go-sdk/pkg/api/client.gen.go b/go-sdk/pkg/api/client.gen.go new file mode 100644 index 00000000000..ae57c39ba31 --- /dev/null +++ b/go-sdk/pkg/api/client.gen.go @@ -0,0 +1,2613 @@ +// Package api provides primitives to interact with the openapi HTTP API. +// +// Code generated by github.com/ashb/oapi-resty-codegen version v0.0.0-20250508101518-564f8c686e08 DO NOT EDIT. +// 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 ( + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + "time" + + "github.com/oapi-codegen/runtime" + openapi_types "github.com/oapi-codegen/runtime/types" + "resty.dev/v3" +) + +const ( + JWTBearerScopes = "JWTBearer.Scopes" +) + +// Defines values for DagRunState. +const ( + DagRunStateFailed DagRunState = "failed" + DagRunStateQueued DagRunState = "queued" + DagRunStateRunning DagRunState = "running" + DagRunStateSuccess DagRunState = "success" +) + +// Defines values for DagRunType. +const ( + AssetTriggered DagRunType = "asset_triggered" + Backfill DagRunType = "backfill" + Manual DagRunType = "manual" + Scheduled DagRunType = "scheduled" +) + +// Defines values for IntermediateTIState. +const ( + IntermediateTIStateDeferred IntermediateTIState = "deferred" + IntermediateTIStateQueued IntermediateTIState = "queued" + IntermediateTIStateRestarting IntermediateTIState = "restarting" + IntermediateTIStateScheduled IntermediateTIState = "scheduled" + IntermediateTIStateUpForReschedule IntermediateTIState = "up_for_reschedule" + IntermediateTIStateUpForRetry IntermediateTIState = "up_for_retry" + IntermediateTIStateUpstreamFailed IntermediateTIState = "upstream_failed" +) + +// Defines values for TIDeferredStatePayloadState. +const ( + Deferred TIDeferredStatePayloadState = "deferred" +) + +// Defines values for TIEnterRunningPayloadState. +const ( + Running TIEnterRunningPayloadState = "running" +) + +// Defines values for TIRescheduleStatePayloadState. +const ( + UpForReschedule TIRescheduleStatePayloadState = "up_for_reschedule" +) + +// Defines values for TIRetryStatePayloadState. +const ( + UpForRetry TIRetryStatePayloadState = "up_for_retry" +) + +// Defines values for TISuccessStatePayloadState. +const ( + TISuccessStatePayloadStateSuccess TISuccessStatePayloadState = "success" +) + +// Defines values for TerminalStateNonSuccess. +const ( + Failed TerminalStateNonSuccess = "failed" + Removed TerminalStateNonSuccess = "removed" + Skipped TerminalStateNonSuccess = "skipped" +) + +// Defines values for TerminalTIState. +const ( + TerminalTIStateFailed TerminalTIState = "failed" + TerminalTIStateRemoved TerminalTIState = "removed" + TerminalTIStateSkipped TerminalTIState = "skipped" + TerminalTIStateSuccess TerminalTIState = "success" +) + +// AssetAliasReferenceAssetEventDagRun Schema for AssetAliasModel used in AssetEventDagRunReference. +type AssetAliasReferenceAssetEventDagRun struct { + Name string `json:"name"` +} + +// AssetEventDagRunReference Schema for AssetEvent model used in DagRun. +type AssetEventDagRunReference struct { + // Asset Schema for AssetModel used in AssetEventDagRunReference. + Asset AssetReferenceAssetEventDagRun `json:"asset"` + Extra map[string]interface{} `json:"extra"` + SourceAliases []AssetAliasReferenceAssetEventDagRun `json:"source_aliases"` + SourceDagId *string `json:"source_dag_id"` + SourceMapIndex *int `json:"source_map_index"` + SourceRunId *string `json:"source_run_id"` + SourceTaskId *string `json:"source_task_id"` + Timestamp time.Time `json:"timestamp"` +} + +// AssetEventResponse Asset event schema with fields that are needed for Runtime. +type AssetEventResponse struct { + // Asset Asset schema for responses with fields that are needed for Runtime. + Asset AssetResponse `json:"asset"` + CreatedDagruns []DagRunAssetReference `json:"created_dagruns"` + Extra *map[string]interface{} `json:"extra"` + Id int `json:"id"` + SourceDagId *string `json:"source_dag_id"` + SourceMapIndex *int `json:"source_map_index,omitempty"` + SourceRunId *string `json:"source_run_id"` + SourceTaskId *string `json:"source_task_id"` + Timestamp time.Time `json:"timestamp"` +} + +// AssetEventsResponse Collection of AssetEventResponse. +type AssetEventsResponse struct { + AssetEvents []AssetEventResponse `json:"asset_events"` +} + +// AssetProfile Profile of an asset-like object. +// +// Asset will have name, uri defined, with type set to 'Asset'. +// AssetNameRef will have name defined, type set to 'AssetNameRef'. +// AssetUriRef will have uri defined, type set to 'AssetUriRef'. +// AssetAlias will have name defined, type set to 'AssetAlias'. +// +// Note that 'type' here is distinct from 'asset_type' the user declares on an +// Asset (or subclass). This field is for distinguishing between different +// asset-related types (Asset, AssetRef, or AssetAlias). +type AssetProfile struct { + Name *string `json:"name"` + Type string `json:"type"` + Uri *string `json:"uri"` +} + +// AssetReferenceAssetEventDagRun Schema for AssetModel used in AssetEventDagRunReference. +type AssetReferenceAssetEventDagRun struct { + Extra map[string]interface{} `json:"extra"` + Name string `json:"name"` + Uri string `json:"uri"` +} + +// AssetResponse Asset schema for responses with fields that are needed for Runtime. +type AssetResponse struct { + Extra *map[string]interface{} `json:"extra"` + Group string `json:"group"` + Name string `json:"name"` + Uri string `json:"uri"` +} + +// BundleInfo Schema for telling task which bundle to run with. +type BundleInfo struct { + Name string `json:"name"` + Version *string `json:"version"` +} + +// ConnectionResponse Connection schema for responses with fields that are needed for Runtime. +type ConnectionResponse struct { + ConnId string `json:"conn_id"` + ConnType string `json:"conn_type"` + Extra *string `json:"extra"` + Host *string `json:"host"` + Login *string `json:"login"` + Password *string `json:"password"` + Port *int `json:"port"` + Schema *string `json:"schema"` +} + +// DagRun Schema for DagRun model with minimal required fields needed for Runtime. +type DagRun struct { + ClearNumber *int `json:"clear_number,omitempty"` + Conf *map[string]interface{} `json:"conf,omitempty"` + ConsumedAssetEvents []AssetEventDagRunReference `json:"consumed_asset_events"` + DagId string `json:"dag_id"` + DataIntervalEnd *time.Time `json:"data_interval_end"` + DataIntervalStart *time.Time `json:"data_interval_start"` + EndDate *time.Time `json:"end_date"` + LogicalDate *time.Time `json:"logical_date"` + RunAfter time.Time `json:"run_after"` + RunId string `json:"run_id"` + + // RunType Class with DagRun types. + RunType DagRunType `json:"run_type"` + StartDate time.Time `json:"start_date"` +} + +// DagRunAssetReference DagRun serializer for asset responses. +type DagRunAssetReference struct { + DagId string `json:"dag_id"` + DataIntervalEnd *time.Time `json:"data_interval_end"` + DataIntervalStart *time.Time `json:"data_interval_start"` + EndDate *time.Time `json:"end_date"` + LogicalDate *time.Time `json:"logical_date"` + RunId string `json:"run_id"` + StartDate time.Time `json:"start_date"` + State string `json:"state"` +} + +// DagRunState All possible states that a DagRun can be in. +// +// These are "shared" with TaskInstanceState in some parts of the code, +// so please ensure that their values always match the ones with the +// same name in TaskInstanceState. +type DagRunState string + +// DagRunStateResponse Schema for DAG Run State response. +type DagRunStateResponse struct { + // State All possible states that a DagRun can be in. + // + // These are "shared" with TaskInstanceState in some parts of the code, + // so please ensure that their values always match the ones with the + // same name in TaskInstanceState. + State DagRunState `json:"state"` +} + +// DagRunType Class with DagRun types. +type DagRunType string + +// HTTPValidationError defines model for HTTPValidationError. +type HTTPValidationError struct { + Detail *[]ValidationError `json:"detail,omitempty"` +} + +// IntermediateTIState States that a Task Instance can be in that indicate it is not yet in a terminal or running state. +type IntermediateTIState string + +// JsonValue defines model for JsonValue. +type JsonValue struct { + union json.RawMessage +} + +// JsonValue0 defines model for . +type JsonValue0 = string + +// JsonValue1 defines model for . +type JsonValue1 = float32 + +// JsonValue2 defines model for . +type JsonValue2 = int + +// JsonValue3 defines model for . +type JsonValue3 = map[string]interface{} + +// JsonValue4 defines model for . +type JsonValue4 = []interface{} + +// JsonValue5 defines model for . +type JsonValue5 = bool + +// PrevSuccessfulDagRunResponse Schema for response with previous successful DagRun information for Task Template Context. +type PrevSuccessfulDagRunResponse struct { + DataIntervalEnd *time.Time `json:"data_interval_end"` + DataIntervalStart *time.Time `json:"data_interval_start"` + EndDate *time.Time `json:"end_date"` + StartDate *time.Time `json:"start_date"` +} + +// TIDeferredStatePayload Schema for updating TaskInstance to a deferred state. +type TIDeferredStatePayload struct { + Classpath string `json:"classpath"` + NextKwargs *TIDeferredStatePayload_NextKwargs `json:"next_kwargs,omitempty"` + NextMethod string `json:"next_method"` + State TIDeferredStatePayloadState `json:"state"` + TriggerKwargs *TIDeferredStatePayload_TriggerKwargs `json:"trigger_kwargs,omitempty"` + TriggerTimeout *string `json:"trigger_timeout"` +} + +// TIDeferredStatePayloadNextKwargs0 defines model for . +type TIDeferredStatePayloadNextKwargs0 map[string]interface{} + +// TIDeferredStatePayloadNextKwargs1 defines model for . +type TIDeferredStatePayloadNextKwargs1 = string + +// TIDeferredStatePayload_NextKwargs defines model for TIDeferredStatePayload.NextKwargs. +type TIDeferredStatePayload_NextKwargs struct { + union json.RawMessage +} + +// TIDeferredStatePayloadState defines model for TIDeferredStatePayload.State. +type TIDeferredStatePayloadState string + +// TIDeferredStatePayloadTriggerKwargs0 defines model for . +type TIDeferredStatePayloadTriggerKwargs0 map[string]interface{} + +// TIDeferredStatePayloadTriggerKwargs1 defines model for . +type TIDeferredStatePayloadTriggerKwargs1 = string + +// TIDeferredStatePayload_TriggerKwargs defines model for TIDeferredStatePayload.TriggerKwargs. +type TIDeferredStatePayload_TriggerKwargs struct { + union json.RawMessage +} + +// TIEnterRunningPayload Schema for updating TaskInstance to 'RUNNING' state with minimal required fields. +type TIEnterRunningPayload struct { + Hostname string `json:"hostname"` + Pid int `json:"pid"` + StartDate time.Time `json:"start_date"` + State TIEnterRunningPayloadState `json:"state"` + Unixname string `json:"unixname"` +} + +// TIEnterRunningPayloadState defines model for TIEnterRunningPayload.State. +type TIEnterRunningPayloadState string + +// TIHeartbeatInfo Schema for TaskInstance heartbeat endpoint. +type TIHeartbeatInfo struct { + Hostname string `json:"hostname"` + Pid int `json:"pid"` +} + +// TIRescheduleStatePayload Schema for updating TaskInstance to a up_for_reschedule state. +type TIRescheduleStatePayload struct { + EndDate time.Time `json:"end_date"` + RescheduleDate time.Time `json:"reschedule_date"` + State TIRescheduleStatePayloadState `json:"state"` +} + +// TIRescheduleStatePayloadState defines model for TIRescheduleStatePayload.State. +type TIRescheduleStatePayloadState string + +// TIRetryStatePayload Schema for updating TaskInstance to up_for_retry. +type TIRetryStatePayload struct { + EndDate time.Time `json:"end_date"` + State TIRetryStatePayloadState `json:"state"` +} + +// TIRetryStatePayloadState defines model for TIRetryStatePayload.State. +type TIRetryStatePayloadState string + +// TIRunContext Response schema for TaskInstance run context. +type TIRunContext struct { + Connections *[]ConnectionResponse `json:"connections,omitempty"` + + // DagRun Schema for DagRun model with minimal required fields needed for Runtime. + DagRun DagRun `json:"dag_run"` + MaxTries int `json:"max_tries"` + NextKwargs *TIRunContext_NextKwargs `json:"next_kwargs"` + NextMethod *string `json:"next_method"` + ShouldRetry bool `json:"should_retry"` + TaskRescheduleCount *int `json:"task_reschedule_count,omitempty"` + UpstreamMapIndexes *map[string]int `json:"upstream_map_indexes"` + Variables *[]VariableResponse `json:"variables,omitempty"` + XcomKeysToClear *[]string `json:"xcom_keys_to_clear,omitempty"` +} + +// TIRunContextNextKwargs0 defines model for . +type TIRunContextNextKwargs0 map[string]interface{} + +// TIRunContextNextKwargs1 defines model for . +type TIRunContextNextKwargs1 = string + +// TIRunContext_NextKwargs defines model for TIRunContext.NextKwargs. +type TIRunContext_NextKwargs struct { + union json.RawMessage +} + +// TISkippedDownstreamTasksStatePayload Schema for updating downstream tasks to a skipped state. +type TISkippedDownstreamTasksStatePayload struct { + Tasks []TISkippedDownstreamTasksStatePayload_Tasks_Item `json:"tasks"` +} + +// TISkippedDownstreamTasksStatePayloadTasks0 defines model for . +type TISkippedDownstreamTasksStatePayloadTasks0 = string + +// TISkippedDownstreamTasksStatePayloadTasks1 defines model for . +type TISkippedDownstreamTasksStatePayloadTasks1 = []interface{} + +// TISkippedDownstreamTasksStatePayload_Tasks_Item defines model for TISkippedDownstreamTasksStatePayload.tasks.Item. +type TISkippedDownstreamTasksStatePayload_Tasks_Item struct { + union json.RawMessage +} + +// TISuccessStatePayload Schema for updating TaskInstance to success state. +type TISuccessStatePayload struct { + EndDate time.Time `json:"end_date"` + OutletEvents *[]map[string]interface{} `json:"outlet_events,omitempty"` + State TISuccessStatePayloadState `json:"state"` + TaskOutlets *[]AssetProfile `json:"task_outlets,omitempty"` +} + +// TISuccessStatePayloadState defines model for TISuccessStatePayload.State. +type TISuccessStatePayloadState string + +// TITargetStatePayload Schema for updating TaskInstance to a target state, excluding terminal and running states. +type TITargetStatePayload struct { + // State States that a Task Instance can be in that indicate it is not yet in a terminal or running state. + State IntermediateTIState `json:"state"` +} + +// TITerminalStatePayload Schema for updating TaskInstance to a terminal state except SUCCESS state. +type TITerminalStatePayload struct { + EndDate time.Time `json:"end_date"` + + // State TaskInstance states that can be reported without extra information. + State TerminalStateNonSuccess `json:"state"` +} + +// TIUpdateStatePayload defines model for TIUpdateStatePayload. +type TIUpdateStatePayload struct { + union json.RawMessage +} + +// TaskInstance Schema for TaskInstance model with minimal required fields needed for Runtime. +type TaskInstance struct { + ContextCarrier *map[string]interface{} `json:"context_carrier"` + DagId string `json:"dag_id"` + Hostname *string `json:"hostname"` + Id openapi_types.UUID `json:"id"` + MapIndex *int `json:"map_index,omitempty"` + RunId string `json:"run_id"` + TaskId string `json:"task_id"` + TryNumber int `json:"try_number"` +} + +// TaskStatesResponse Response for task states with run_id, task and state. +type TaskStatesResponse struct { + TaskStates map[string]interface{} `json:"task_states"` +} + +// TerminalStateNonSuccess TaskInstance states that can be reported without extra information. +type TerminalStateNonSuccess string + +// TerminalTIState defines model for TerminalTIState. +type TerminalTIState string + +// TriggerDAGRunPayload Schema for Trigger DAG Run API request. +type TriggerDAGRunPayload struct { + Conf *map[string]interface{} `json:"conf,omitempty"` + LogicalDate *time.Time `json:"logical_date"` + ResetDagRun *bool `json:"reset_dag_run,omitempty"` +} + +// ValidationError defines model for ValidationError. +type ValidationError struct { + Loc []ValidationError_Loc_Item `json:"loc"` + Msg string `json:"msg"` + Type string `json:"type"` +} + +// ValidationErrorLoc0 defines model for . +type ValidationErrorLoc0 = string + +// ValidationErrorLoc1 defines model for . +type ValidationErrorLoc1 = int + +// ValidationError_Loc_Item defines model for ValidationError.loc.Item. +type ValidationError_Loc_Item struct { + union json.RawMessage +} + +// VariablePostBody Request body schema for creating variables. +type VariablePostBody struct { + Description *string `json:"description"` + Val *string `json:"val"` +} + +// VariableResponse Variable schema for responses with fields that are needed for Runtime. +type VariableResponse struct { + Key string `json:"key"` + Value *string `json:"value"` +} + +// XComResponse XCom schema for responses with fields that are needed for Runtime. +type XComResponse struct { + Key string `json:"key"` + Value *JsonValue `json:"value"` +} + +// GetAssetEventByAssetNameUriParams defines parameters for GetAssetEventByAssetNameUri. +type GetAssetEventByAssetNameUriParams struct { + // Name The name of the Asset + Name string `form:"name" json:"name"` + + // Uri The URI of the Asset + Uri string `form:"uri" json:"uri"` +} + +// GetAssetEventByAssetAliasParams defines parameters for GetAssetEventByAssetAlias. +type GetAssetEventByAssetAliasParams struct { + // Name The name of the Asset Alias + Name string `form:"name" json:"name"` +} + +// GetAssetByNameParams defines parameters for GetAssetByName. +type GetAssetByNameParams struct { + // Name The name of the Asset + Name string `form:"name" json:"name"` +} + +// GetAssetByUriParams defines parameters for GetAssetByUri. +type GetAssetByUriParams struct { + // Uri The URI of the Asset + Uri string `form:"uri" json:"uri"` +} + +// GetDrCountParams defines parameters for GetDrCount. +type GetDrCountParams struct { + DagId string `form:"dag_id" json:"dag_id"` + LogicalDates *[]time.Time `form:"logical_dates,omitempty" json:"logical_dates,omitempty"` + RunIds *[]string `form:"run_ids,omitempty" json:"run_ids,omitempty"` + States *[]string `form:"states,omitempty" json:"states,omitempty"` +} + +// GetTaskInstanceCountParams defines parameters for GetTaskInstanceCount. +type GetTaskInstanceCountParams struct { + DagId string `form:"dag_id" json:"dag_id"` + TaskIds *[]string `form:"task_ids,omitempty" json:"task_ids,omitempty"` + TaskGroupId *string `form:"task_group_id,omitempty" json:"task_group_id,omitempty"` + LogicalDates *[]time.Time `form:"logical_dates,omitempty" json:"logical_dates,omitempty"` + RunIds *[]string `form:"run_ids,omitempty" json:"run_ids,omitempty"` + States *[]string `form:"states,omitempty" json:"states,omitempty"` +} + +// GetTaskInstanceStatesParams defines parameters for GetTaskInstanceStates. +type GetTaskInstanceStatesParams struct { + DagId string `form:"dag_id" json:"dag_id"` + TaskIds *[]string `form:"task_ids,omitempty" json:"task_ids,omitempty"` + TaskGroupId *string `form:"task_group_id,omitempty" json:"task_group_id,omitempty"` + LogicalDates *[]time.Time `form:"logical_dates,omitempty" json:"logical_dates,omitempty"` + RunIds *[]string `form:"run_ids,omitempty" json:"run_ids,omitempty"` +} + +// TaskInstancePutRenderedFieldsJSONBody defines parameters for TaskInstancePutRenderedFields. +type TaskInstancePutRenderedFieldsJSONBody map[string]*JsonValue + +// TaskInstanceUpdateStateJSONBody defines parameters for TaskInstanceUpdateState. +type TaskInstanceUpdateStateJSONBody = TIUpdateStatePayload + +// DeleteXcomParams defines parameters for DeleteXcom. +type DeleteXcomParams struct { + MapIndex *int `form:"map_index,omitempty" json:"map_index,omitempty"` +} + +// GetXcomParams defines parameters for GetXcom. +type GetXcomParams struct { + MapIndex *int `form:"map_index,omitempty" json:"map_index,omitempty"` + IncludePriorDates *bool `form:"include_prior_dates,omitempty" json:"include_prior_dates,omitempty"` +} + +// HeadXcomParams defines parameters for HeadXcom. +type HeadXcomParams struct { + MapIndex *int `form:"map_index,omitempty" json:"map_index,omitempty"` +} + +// SetXcomParams defines parameters for SetXcom. +type SetXcomParams struct { + MapIndex *int `form:"map_index,omitempty" json:"map_index,omitempty"` + + // MappedLength Number of mapped tasks this value expands into + MappedLength *int `form:"mapped_length,omitempty" json:"mapped_length,omitempty"` +} + +// TriggerDagRunJSONRequestBody defines body for TriggerDagRun for application/json ContentType. +type TriggerDagRunJSONRequestBody = TriggerDAGRunPayload + +// TaskInstanceHeartbeatJSONRequestBody defines body for TaskInstanceHeartbeat for application/json ContentType. +type TaskInstanceHeartbeatJSONRequestBody = TIHeartbeatInfo + +// TaskInstancePutRenderedFieldsJSONRequestBody defines body for TaskInstancePutRenderedFields for application/json ContentType. +type TaskInstancePutRenderedFieldsJSONRequestBody TaskInstancePutRenderedFieldsJSONBody + +// TaskInstanceRunJSONRequestBody defines body for TaskInstanceRun for application/json ContentType. +type TaskInstanceRunJSONRequestBody = TIEnterRunningPayload + +// TaskInstanceSkipDownstreamJSONRequestBody defines body for TaskInstanceSkipDownstream for application/json ContentType. +type TaskInstanceSkipDownstreamJSONRequestBody = TISkippedDownstreamTasksStatePayload + +// TaskInstanceUpdateStateJSONRequestBody defines body for TaskInstanceUpdateState for application/json ContentType. +type TaskInstanceUpdateStateJSONRequestBody = TaskInstanceUpdateStateJSONBody + +// PutVariableJSONRequestBody defines body for PutVariable for application/json ContentType. +type PutVariableJSONRequestBody = VariablePostBody + +// SetXcomJSONRequestBody defines body for SetXcom for application/json ContentType. +type SetXcomJSONRequestBody = JsonValue + +// AsJsonValue0 returns the union data inside the JsonValue as a JsonValue0 +func (t JsonValue) AsJsonValue0() (JsonValue0, error) { + var body JsonValue0 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromJsonValue0 overwrites any union data inside the JsonValue as the provided JsonValue0 +func (t *JsonValue) FromJsonValue0(v JsonValue0) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeJsonValue0 performs a merge with any union data inside the JsonValue, using the provided JsonValue0 +func (t *JsonValue) MergeJsonValue0(v JsonValue0) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsJsonValue1 returns the union data inside the JsonValue as a JsonValue1 +func (t JsonValue) AsJsonValue1() (JsonValue1, error) { + var body JsonValue1 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromJsonValue1 overwrites any union data inside the JsonValue as the provided JsonValue1 +func (t *JsonValue) FromJsonValue1(v JsonValue1) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeJsonValue1 performs a merge with any union data inside the JsonValue, using the provided JsonValue1 +func (t *JsonValue) MergeJsonValue1(v JsonValue1) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsJsonValue2 returns the union data inside the JsonValue as a JsonValue2 +func (t JsonValue) AsJsonValue2() (JsonValue2, error) { + var body JsonValue2 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromJsonValue2 overwrites any union data inside the JsonValue as the provided JsonValue2 +func (t *JsonValue) FromJsonValue2(v JsonValue2) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeJsonValue2 performs a merge with any union data inside the JsonValue, using the provided JsonValue2 +func (t *JsonValue) MergeJsonValue2(v JsonValue2) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsJsonValue3 returns the union data inside the JsonValue as a JsonValue3 +func (t JsonValue) AsJsonValue3() (JsonValue3, error) { + var body JsonValue3 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromJsonValue3 overwrites any union data inside the JsonValue as the provided JsonValue3 +func (t *JsonValue) FromJsonValue3(v JsonValue3) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeJsonValue3 performs a merge with any union data inside the JsonValue, using the provided JsonValue3 +func (t *JsonValue) MergeJsonValue3(v JsonValue3) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsJsonValue4 returns the union data inside the JsonValue as a JsonValue4 +func (t JsonValue) AsJsonValue4() (JsonValue4, error) { + var body JsonValue4 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromJsonValue4 overwrites any union data inside the JsonValue as the provided JsonValue4 +func (t *JsonValue) FromJsonValue4(v JsonValue4) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeJsonValue4 performs a merge with any union data inside the JsonValue, using the provided JsonValue4 +func (t *JsonValue) MergeJsonValue4(v JsonValue4) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsJsonValue5 returns the union data inside the JsonValue as a JsonValue5 +func (t JsonValue) AsJsonValue5() (JsonValue5, error) { + var body JsonValue5 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromJsonValue5 overwrites any union data inside the JsonValue as the provided JsonValue5 +func (t *JsonValue) FromJsonValue5(v JsonValue5) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeJsonValue5 performs a merge with any union data inside the JsonValue, using the provided JsonValue5 +func (t *JsonValue) MergeJsonValue5(v JsonValue5) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +func (t JsonValue) MarshalJSON() ([]byte, error) { + b, err := t.union.MarshalJSON() + return b, err +} + +func (t *JsonValue) UnmarshalJSON(b []byte) error { + err := t.union.UnmarshalJSON(b) + return err +} + +// AsTIDeferredStatePayloadNextKwargs0 returns the union data inside the TIDeferredStatePayload_NextKwargs as a TIDeferredStatePayloadNextKwargs0 +func (t TIDeferredStatePayload_NextKwargs) AsTIDeferredStatePayloadNextKwargs0() (TIDeferredStatePayloadNextKwargs0, error) { + var body TIDeferredStatePayloadNextKwargs0 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromTIDeferredStatePayloadNextKwargs0 overwrites any union data inside the TIDeferredStatePayload_NextKwargs as the provided TIDeferredStatePayloadNextKwargs0 +func (t *TIDeferredStatePayload_NextKwargs) FromTIDeferredStatePayloadNextKwargs0(v TIDeferredStatePayloadNextKwargs0) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeTIDeferredStatePayloadNextKwargs0 performs a merge with any union data inside the TIDeferredStatePayload_NextKwargs, using the provided TIDeferredStatePayloadNextKwargs0 +func (t *TIDeferredStatePayload_NextKwargs) MergeTIDeferredStatePayloadNextKwargs0(v TIDeferredStatePayloadNextKwargs0) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsTIDeferredStatePayloadNextKwargs1 returns the union data inside the TIDeferredStatePayload_NextKwargs as a TIDeferredStatePayloadNextKwargs1 +func (t TIDeferredStatePayload_NextKwargs) AsTIDeferredStatePayloadNextKwargs1() (TIDeferredStatePayloadNextKwargs1, error) { + var body TIDeferredStatePayloadNextKwargs1 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromTIDeferredStatePayloadNextKwargs1 overwrites any union data inside the TIDeferredStatePayload_NextKwargs as the provided TIDeferredStatePayloadNextKwargs1 +func (t *TIDeferredStatePayload_NextKwargs) FromTIDeferredStatePayloadNextKwargs1(v TIDeferredStatePayloadNextKwargs1) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeTIDeferredStatePayloadNextKwargs1 performs a merge with any union data inside the TIDeferredStatePayload_NextKwargs, using the provided TIDeferredStatePayloadNextKwargs1 +func (t *TIDeferredStatePayload_NextKwargs) MergeTIDeferredStatePayloadNextKwargs1(v TIDeferredStatePayloadNextKwargs1) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +func (t TIDeferredStatePayload_NextKwargs) MarshalJSON() ([]byte, error) { + b, err := t.union.MarshalJSON() + return b, err +} + +func (t *TIDeferredStatePayload_NextKwargs) UnmarshalJSON(b []byte) error { + err := t.union.UnmarshalJSON(b) + return err +} + +// AsTIDeferredStatePayloadTriggerKwargs0 returns the union data inside the TIDeferredStatePayload_TriggerKwargs as a TIDeferredStatePayloadTriggerKwargs0 +func (t TIDeferredStatePayload_TriggerKwargs) AsTIDeferredStatePayloadTriggerKwargs0() (TIDeferredStatePayloadTriggerKwargs0, error) { + var body TIDeferredStatePayloadTriggerKwargs0 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromTIDeferredStatePayloadTriggerKwargs0 overwrites any union data inside the TIDeferredStatePayload_TriggerKwargs as the provided TIDeferredStatePayloadTriggerKwargs0 +func (t *TIDeferredStatePayload_TriggerKwargs) FromTIDeferredStatePayloadTriggerKwargs0(v TIDeferredStatePayloadTriggerKwargs0) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeTIDeferredStatePayloadTriggerKwargs0 performs a merge with any union data inside the TIDeferredStatePayload_TriggerKwargs, using the provided TIDeferredStatePayloadTriggerKwargs0 +func (t *TIDeferredStatePayload_TriggerKwargs) MergeTIDeferredStatePayloadTriggerKwargs0(v TIDeferredStatePayloadTriggerKwargs0) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsTIDeferredStatePayloadTriggerKwargs1 returns the union data inside the TIDeferredStatePayload_TriggerKwargs as a TIDeferredStatePayloadTriggerKwargs1 +func (t TIDeferredStatePayload_TriggerKwargs) AsTIDeferredStatePayloadTriggerKwargs1() (TIDeferredStatePayloadTriggerKwargs1, error) { + var body TIDeferredStatePayloadTriggerKwargs1 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromTIDeferredStatePayloadTriggerKwargs1 overwrites any union data inside the TIDeferredStatePayload_TriggerKwargs as the provided TIDeferredStatePayloadTriggerKwargs1 +func (t *TIDeferredStatePayload_TriggerKwargs) FromTIDeferredStatePayloadTriggerKwargs1(v TIDeferredStatePayloadTriggerKwargs1) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeTIDeferredStatePayloadTriggerKwargs1 performs a merge with any union data inside the TIDeferredStatePayload_TriggerKwargs, using the provided TIDeferredStatePayloadTriggerKwargs1 +func (t *TIDeferredStatePayload_TriggerKwargs) MergeTIDeferredStatePayloadTriggerKwargs1(v TIDeferredStatePayloadTriggerKwargs1) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +func (t TIDeferredStatePayload_TriggerKwargs) MarshalJSON() ([]byte, error) { + b, err := t.union.MarshalJSON() + return b, err +} + +func (t *TIDeferredStatePayload_TriggerKwargs) UnmarshalJSON(b []byte) error { + err := t.union.UnmarshalJSON(b) + return err +} + +// AsTIRunContextNextKwargs0 returns the union data inside the TIRunContext_NextKwargs as a TIRunContextNextKwargs0 +func (t TIRunContext_NextKwargs) AsTIRunContextNextKwargs0() (TIRunContextNextKwargs0, error) { + var body TIRunContextNextKwargs0 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromTIRunContextNextKwargs0 overwrites any union data inside the TIRunContext_NextKwargs as the provided TIRunContextNextKwargs0 +func (t *TIRunContext_NextKwargs) FromTIRunContextNextKwargs0(v TIRunContextNextKwargs0) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeTIRunContextNextKwargs0 performs a merge with any union data inside the TIRunContext_NextKwargs, using the provided TIRunContextNextKwargs0 +func (t *TIRunContext_NextKwargs) MergeTIRunContextNextKwargs0(v TIRunContextNextKwargs0) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsTIRunContextNextKwargs1 returns the union data inside the TIRunContext_NextKwargs as a TIRunContextNextKwargs1 +func (t TIRunContext_NextKwargs) AsTIRunContextNextKwargs1() (TIRunContextNextKwargs1, error) { + var body TIRunContextNextKwargs1 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromTIRunContextNextKwargs1 overwrites any union data inside the TIRunContext_NextKwargs as the provided TIRunContextNextKwargs1 +func (t *TIRunContext_NextKwargs) FromTIRunContextNextKwargs1(v TIRunContextNextKwargs1) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeTIRunContextNextKwargs1 performs a merge with any union data inside the TIRunContext_NextKwargs, using the provided TIRunContextNextKwargs1 +func (t *TIRunContext_NextKwargs) MergeTIRunContextNextKwargs1(v TIRunContextNextKwargs1) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +func (t TIRunContext_NextKwargs) MarshalJSON() ([]byte, error) { + b, err := t.union.MarshalJSON() + return b, err +} + +func (t *TIRunContext_NextKwargs) UnmarshalJSON(b []byte) error { + err := t.union.UnmarshalJSON(b) + return err +} + +// AsTISkippedDownstreamTasksStatePayloadTasks0 returns the union data inside the TISkippedDownstreamTasksStatePayload_Tasks_Item as a TISkippedDownstreamTasksStatePayloadTasks0 +func (t TISkippedDownstreamTasksStatePayload_Tasks_Item) AsTISkippedDownstreamTasksStatePayloadTasks0() (TISkippedDownstreamTasksStatePayloadTasks0, error) { + var body TISkippedDownstreamTasksStatePayloadTasks0 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromTISkippedDownstreamTasksStatePayloadTasks0 overwrites any union data inside the TISkippedDownstreamTasksStatePayload_Tasks_Item as the provided TISkippedDownstreamTasksStatePayloadTasks0 +func (t *TISkippedDownstreamTasksStatePayload_Tasks_Item) FromTISkippedDownstreamTasksStatePayloadTasks0(v TISkippedDownstreamTasksStatePayloadTasks0) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeTISkippedDownstreamTasksStatePayloadTasks0 performs a merge with any union data inside the TISkippedDownstreamTasksStatePayload_Tasks_Item, using the provided TISkippedDownstreamTasksStatePayloadTasks0 +func (t *TISkippedDownstreamTasksStatePayload_Tasks_Item) MergeTISkippedDownstreamTasksStatePayloadTasks0(v TISkippedDownstreamTasksStatePayloadTasks0) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsTISkippedDownstreamTasksStatePayloadTasks1 returns the union data inside the TISkippedDownstreamTasksStatePayload_Tasks_Item as a TISkippedDownstreamTasksStatePayloadTasks1 +func (t TISkippedDownstreamTasksStatePayload_Tasks_Item) AsTISkippedDownstreamTasksStatePayloadTasks1() (TISkippedDownstreamTasksStatePayloadTasks1, error) { + var body TISkippedDownstreamTasksStatePayloadTasks1 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromTISkippedDownstreamTasksStatePayloadTasks1 overwrites any union data inside the TISkippedDownstreamTasksStatePayload_Tasks_Item as the provided TISkippedDownstreamTasksStatePayloadTasks1 +func (t *TISkippedDownstreamTasksStatePayload_Tasks_Item) FromTISkippedDownstreamTasksStatePayloadTasks1(v TISkippedDownstreamTasksStatePayloadTasks1) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeTISkippedDownstreamTasksStatePayloadTasks1 performs a merge with any union data inside the TISkippedDownstreamTasksStatePayload_Tasks_Item, using the provided TISkippedDownstreamTasksStatePayloadTasks1 +func (t *TISkippedDownstreamTasksStatePayload_Tasks_Item) MergeTISkippedDownstreamTasksStatePayloadTasks1(v TISkippedDownstreamTasksStatePayloadTasks1) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +func (t TISkippedDownstreamTasksStatePayload_Tasks_Item) MarshalJSON() ([]byte, error) { + b, err := t.union.MarshalJSON() + return b, err +} + +func (t *TISkippedDownstreamTasksStatePayload_Tasks_Item) UnmarshalJSON(b []byte) error { + err := t.union.UnmarshalJSON(b) + return err +} + +// AsTITerminalStatePayload returns the union data inside the TIUpdateStatePayload as a TITerminalStatePayload +func (t TIUpdateStatePayload) AsTITerminalStatePayload() (TITerminalStatePayload, error) { + var body TITerminalStatePayload + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromTITerminalStatePayload overwrites any union data inside the TIUpdateStatePayload as the provided TITerminalStatePayload +func (t *TIUpdateStatePayload) FromTITerminalStatePayload(v TITerminalStatePayload) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeTITerminalStatePayload performs a merge with any union data inside the TIUpdateStatePayload, using the provided TITerminalStatePayload +func (t *TIUpdateStatePayload) MergeTITerminalStatePayload(v TITerminalStatePayload) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsTISuccessStatePayload returns the union data inside the TIUpdateStatePayload as a TISuccessStatePayload +func (t TIUpdateStatePayload) AsTISuccessStatePayload() (TISuccessStatePayload, error) { + var body TISuccessStatePayload + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromTISuccessStatePayload overwrites any union data inside the TIUpdateStatePayload as the provided TISuccessStatePayload +func (t *TIUpdateStatePayload) FromTISuccessStatePayload(v TISuccessStatePayload) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeTISuccessStatePayload performs a merge with any union data inside the TIUpdateStatePayload, using the provided TISuccessStatePayload +func (t *TIUpdateStatePayload) MergeTISuccessStatePayload(v TISuccessStatePayload) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsTITargetStatePayload returns the union data inside the TIUpdateStatePayload as a TITargetStatePayload +func (t TIUpdateStatePayload) AsTITargetStatePayload() (TITargetStatePayload, error) { + var body TITargetStatePayload + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromTITargetStatePayload overwrites any union data inside the TIUpdateStatePayload as the provided TITargetStatePayload +func (t *TIUpdateStatePayload) FromTITargetStatePayload(v TITargetStatePayload) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeTITargetStatePayload performs a merge with any union data inside the TIUpdateStatePayload, using the provided TITargetStatePayload +func (t *TIUpdateStatePayload) MergeTITargetStatePayload(v TITargetStatePayload) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsTIDeferredStatePayload returns the union data inside the TIUpdateStatePayload as a TIDeferredStatePayload +func (t TIUpdateStatePayload) AsTIDeferredStatePayload() (TIDeferredStatePayload, error) { + var body TIDeferredStatePayload + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromTIDeferredStatePayload overwrites any union data inside the TIUpdateStatePayload as the provided TIDeferredStatePayload +func (t *TIUpdateStatePayload) FromTIDeferredStatePayload(v TIDeferredStatePayload) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeTIDeferredStatePayload performs a merge with any union data inside the TIUpdateStatePayload, using the provided TIDeferredStatePayload +func (t *TIUpdateStatePayload) MergeTIDeferredStatePayload(v TIDeferredStatePayload) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsTIRescheduleStatePayload returns the union data inside the TIUpdateStatePayload as a TIRescheduleStatePayload +func (t TIUpdateStatePayload) AsTIRescheduleStatePayload() (TIRescheduleStatePayload, error) { + var body TIRescheduleStatePayload + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromTIRescheduleStatePayload overwrites any union data inside the TIUpdateStatePayload as the provided TIRescheduleStatePayload +func (t *TIUpdateStatePayload) FromTIRescheduleStatePayload(v TIRescheduleStatePayload) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeTIRescheduleStatePayload performs a merge with any union data inside the TIUpdateStatePayload, using the provided TIRescheduleStatePayload +func (t *TIUpdateStatePayload) MergeTIRescheduleStatePayload(v TIRescheduleStatePayload) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsTIRetryStatePayload returns the union data inside the TIUpdateStatePayload as a TIRetryStatePayload +func (t TIUpdateStatePayload) AsTIRetryStatePayload() (TIRetryStatePayload, error) { + var body TIRetryStatePayload + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromTIRetryStatePayload overwrites any union data inside the TIUpdateStatePayload as the provided TIRetryStatePayload +func (t *TIUpdateStatePayload) FromTIRetryStatePayload(v TIRetryStatePayload) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeTIRetryStatePayload performs a merge with any union data inside the TIUpdateStatePayload, using the provided TIRetryStatePayload +func (t *TIUpdateStatePayload) MergeTIRetryStatePayload(v TIRetryStatePayload) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +func (t TIUpdateStatePayload) MarshalJSON() ([]byte, error) { + b, err := t.union.MarshalJSON() + return b, err +} + +func (t *TIUpdateStatePayload) UnmarshalJSON(b []byte) error { + err := t.union.UnmarshalJSON(b) + return err +} + +// AsValidationErrorLoc0 returns the union data inside the ValidationError_Loc_Item as a ValidationErrorLoc0 +func (t ValidationError_Loc_Item) AsValidationErrorLoc0() (ValidationErrorLoc0, error) { + var body ValidationErrorLoc0 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromValidationErrorLoc0 overwrites any union data inside the ValidationError_Loc_Item as the provided ValidationErrorLoc0 +func (t *ValidationError_Loc_Item) FromValidationErrorLoc0(v ValidationErrorLoc0) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeValidationErrorLoc0 performs a merge with any union data inside the ValidationError_Loc_Item, using the provided ValidationErrorLoc0 +func (t *ValidationError_Loc_Item) MergeValidationErrorLoc0(v ValidationErrorLoc0) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +// AsValidationErrorLoc1 returns the union data inside the ValidationError_Loc_Item as a ValidationErrorLoc1 +func (t ValidationError_Loc_Item) AsValidationErrorLoc1() (ValidationErrorLoc1, error) { + var body ValidationErrorLoc1 + err := json.Unmarshal(t.union, &body) + return body, err +} + +// FromValidationErrorLoc1 overwrites any union data inside the ValidationError_Loc_Item as the provided ValidationErrorLoc1 +func (t *ValidationError_Loc_Item) FromValidationErrorLoc1(v ValidationErrorLoc1) error { + b, err := json.Marshal(v) + t.union = b + return err +} + +// MergeValidationErrorLoc1 performs a merge with any union data inside the ValidationError_Loc_Item, using the provided ValidationErrorLoc1 +func (t *ValidationError_Loc_Item) MergeValidationErrorLoc1(v ValidationErrorLoc1) error { + b, err := json.Marshal(v) + if err != nil { + return err + } + + merged, err := runtime.JSONMerge(t.union, b) + t.union = merged + return err +} + +func (t ValidationError_Loc_Item) MarshalJSON() ([]byte, error) { + b, err := t.union.MarshalJSON() + return b, err +} + +func (t *ValidationError_Loc_Item) UnmarshalJSON(b []byte) error { + err := t.union.UnmarshalJSON(b) + return err +} + +// Client which conforms to the OpenAPI3 specification for this service. +type Client struct { + // The endpoint of the server conforming to this interface, with scheme, + // https://api.deepmap.com for example. This can contain a path relative + // to the server, such as https://api.deepmap.com/dev-test, and all the + // paths in the swagger spec will be appended to the server. + Server string + + *resty.Client + + RequestMiddleware []resty.RequestMiddleware +} + +// ClientOption allows setting custom parameters during construction +type ClientOption func(*Client) error + +func NewClient(server string, opts ...ClientOption) (ClientInterface, error) { + // create a client with sane default values + client := Client{ + Server: server, + Client: resty.New(), + } + client.Client.SetBaseURL(client.Server) + // mutate client and add all optional params + for _, o := range opts { + if err := o(&client); err != nil { + return nil, err + } + } + // ensure the server URL always has a trailing slash + if !strings.HasSuffix(client.Server, "/") { + client.Server += "/" + } + + return &client, nil +} + +// WithClient allows overriding the default [resty.Client], which is +// automatically created using http.Client. +// +// If this is used the `server` base URL argument passed in will not be respected anymore +func WithClient(r *resty.Client) ClientOption { + return func(c *Client) error { + c.Client = r + return nil + } +} + +// WithRoundTripper method sets custom http.Transport or any http.RoundTripper +// compatible interface implementation in the Resty client +func WithRoundTripper(transport http.RoundTripper) ClientOption { + return func(c *Client) error { + c.Client.SetTransport(transport) + return nil + } +} + +// WithRequestMiddleware allows setting up a callback function, which will be +// called right before sending the request. This can be used to mutate the request. +func WithRequestMiddleware(mw resty.RequestMiddleware) ClientOption { + return func(c *Client) error { + c.RequestMiddleware = append(c.RequestMiddleware, mw) + c.Client = c.Client.AddRequestMiddleware(mw) + return nil + } +} + +func (c *Client) AssetEvents() AssetEventsClient { + return &assetEventsClient{c.Client} +} +func (c *Client) Assets() AssetsClient { + return &assetsClient{c.Client} +} +func (c *Client) Connections() ConnectionsClient { + return &connectionsClient{c.Client} +} +func (c *Client) DagRuns() DagRunsClient { + return &dagRunsClient{c.Client} +} +func (c *Client) TaskInstances() TaskInstancesClient { + return &taskInstancesClient{c.Client} +} +func (c *Client) TaskReschedules() TaskReschedulesClient { + return &taskReschedulesClient{c.Client} +} +func (c *Client) Variables() VariablesClient { + return &variablesClient{c.Client} +} +func (c *Client) Xcoms() XcomsClient { + return &xcomsClient{c.Client} +} + +type ClientInterface interface { + // AssetEvents deals with all the Asset Events endpoints + AssetEvents() AssetEventsClient + // Assets deals with all the Assets endpoints + Assets() AssetsClient + // Connections deals with all the Connections endpoints + Connections() ConnectionsClient + // DagRuns deals with all the Dag Runs endpoints + DagRuns() DagRunsClient + // TaskInstances deals with all the Task Instances endpoints + TaskInstances() TaskInstancesClient + // TaskReschedules deals with all the Task Reschedules endpoints + TaskReschedules() TaskReschedulesClient + // Variables deals with all the Variables endpoints + Variables() VariablesClient + // Xcoms deals with all the XComs endpoints + Xcoms() XcomsClient +} + +type GeneralHTTPError struct { + Response *resty.Response + JSON map[string]any + Text string +} + +var errorTypes = map[int]string{ + 1: "informational response", + 3: "redirect response", + 4: "client error", + 5: "server error", +} + +func (e GeneralHTTPError) Error() string { + var b strings.Builder + kind, ok := errorTypes[e.Response.StatusCode()/100] + if !ok { + kind = "unknown HTTP error" + } + fmt.Fprintf(&b, "%s '%s'", kind, e.Response.Status()) + if e.JSON != nil { + fmt.Fprintf(&b, " %v", e.JSON) + } else { + fmt.Fprintf(&b, " content=%q", e.Text) + } + return b.String() +} + +func HandleError(client *resty.Client, resp *resty.Response) error { + if !resp.IsError() { + return nil + } + + e := GeneralHTTPError{Response: resp} + + // fmt.Printf("HandleError: IsRead: %v Error: %T\n", resp.IsRead, resp.Error()) + e.Text = resp.String() + if resp.Header().Get("content-type") == "application/json" { + if json.Unmarshal([]byte(e.Text), &e.JSON) == nil { + e.Text = "" + } + } + + // Set the parsed error back into the object so `resp.Error()` returns the populated one! + resp.Request.SetError(&e) + + return &e +} + +type assetEventsClient struct { + *resty.Client +} + +// GetByAssetNameUriResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *assetEventsClient) GetByAssetNameUriResponse(ctx context.Context, params *GetAssetEventByAssetNameUriParams) (resp *resty.Response, err error) { + var res AssetEventsResponse + + var param0 string + + param0, err = runtime.StyleParamWithLocation("form", true, "name", runtime.ParamLocationQuery, params.Name) + if err != nil { + return nil, err + } + var param1 string + + param1, err = runtime.StyleParamWithLocation("form", true, "uri", runtime.ParamLocationQuery, params.Uri) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetQueryParam("name", param0). + SetQueryParam("uri", param1). + SetResult(&res). + Get("asset-events/by-asset") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *assetEventsClient) GetByAssetNameUri(ctx context.Context, params *GetAssetEventByAssetNameUriParams) (*AssetEventsResponse, error) { + res, err := c.GetByAssetNameUriResponse(ctx, params) + if err != nil { + return nil, err + } + + return res.Result().(*AssetEventsResponse), nil +} + +// GetByAssetAliasResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *assetEventsClient) GetByAssetAliasResponse(ctx context.Context, params *GetAssetEventByAssetAliasParams) (resp *resty.Response, err error) { + var res AssetEventsResponse + + var param0 string + + param0, err = runtime.StyleParamWithLocation("form", true, "name", runtime.ParamLocationQuery, params.Name) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetQueryParam("name", param0). + SetResult(&res). + Get("asset-events/by-asset-alias") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *assetEventsClient) GetByAssetAlias(ctx context.Context, params *GetAssetEventByAssetAliasParams) (*AssetEventsResponse, error) { + res, err := c.GetByAssetAliasResponse(ctx, params) + if err != nil { + return nil, err + } + + return res.Result().(*AssetEventsResponse), nil +} + +type assetsClient struct { + *resty.Client +} + +// GetByNameResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *assetsClient) GetByNameResponse(ctx context.Context, params *GetAssetByNameParams) (resp *resty.Response, err error) { + var res AssetResponse + + var param0 string + + param0, err = runtime.StyleParamWithLocation("form", true, "name", runtime.ParamLocationQuery, params.Name) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetQueryParam("name", param0). + SetResult(&res). + Get("assets/by-name") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *assetsClient) GetByName(ctx context.Context, params *GetAssetByNameParams) (*AssetResponse, error) { + res, err := c.GetByNameResponse(ctx, params) + if err != nil { + return nil, err + } + + return res.Result().(*AssetResponse), nil +} + +// GetByUriResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *assetsClient) GetByUriResponse(ctx context.Context, params *GetAssetByUriParams) (resp *resty.Response, err error) { + var res AssetResponse + + var param0 string + + param0, err = runtime.StyleParamWithLocation("form", true, "uri", runtime.ParamLocationQuery, params.Uri) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetQueryParam("uri", param0). + SetResult(&res). + Get("assets/by-uri") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *assetsClient) GetByUri(ctx context.Context, params *GetAssetByUriParams) (*AssetResponse, error) { + res, err := c.GetByUriResponse(ctx, params) + if err != nil { + return nil, err + } + + return res.Result().(*AssetResponse), nil +} + +type connectionsClient struct { + *resty.Client +} + +// GetResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *connectionsClient) GetResponse(ctx context.Context, connectionId string) (resp *resty.Response, err error) { + var res ConnectionResponse + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "connection_id", runtime.ParamLocationPath, connectionId) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetPathParam("connection_id", pathParam0). + SetResult(&res). + Get("connections/{connection_id}") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *connectionsClient) Get(ctx context.Context, connectionId string) (*ConnectionResponse, error) { + res, err := c.GetResponse(ctx, connectionId) + if err != nil { + return nil, err + } + + return res.Result().(*ConnectionResponse), nil +} + +type dagRunsClient struct { + *resty.Client +} + +// GetDrCountResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *dagRunsClient) GetDrCountResponse(ctx context.Context, params *GetDrCountParams) (resp *resty.Response, err error) { + var res int + + var param0 string + + param0, err = runtime.StyleParamWithLocation("form", true, "dag_id", runtime.ParamLocationQuery, params.DagId) + if err != nil { + return nil, err + } + var param1 string + + param1, err = runtime.StyleParamWithLocation("form", true, "logical_dates", runtime.ParamLocationQuery, params.LogicalDates) + if err != nil { + return nil, err + } + var param2 string + + param2, err = runtime.StyleParamWithLocation("form", true, "run_ids", runtime.ParamLocationQuery, params.RunIds) + if err != nil { + return nil, err + } + var param3 string + + param3, err = runtime.StyleParamWithLocation("form", true, "states", runtime.ParamLocationQuery, params.States) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetQueryParam("dag_id", param0). + SetQueryParam("logical_dates", param1). + SetQueryParam("run_ids", param2). + SetQueryParam("states", param3). + SetResult(&res). + Get("dag-runs/count") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *dagRunsClient) GetDrCount(ctx context.Context, params *GetDrCountParams) (*int, error) { + res, err := c.GetDrCountResponse(ctx, params) + if err != nil { + return nil, err + } + + return res.Result().(*int), nil +} + +// TriggerResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *dagRunsClient) TriggerResponse(ctx context.Context, dagId string, runId string, body *TriggerDAGRunPayload) (resp *resty.Response, err error) { + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "dag_id", runtime.ParamLocationPath, dagId) + if err != nil { + return nil, err + } + var pathParam1 string + + pathParam1, err = runtime.StyleParamWithLocation("simple", false, "run_id", runtime.ParamLocationPath, runId) + if err != nil { + return nil, err + } + + if body == nil { + return nil, fmt.Errorf("Trigger requires a non-nil body argument") + } + resp, err = c.R(). + SetContext(ctx). + SetPathParam("dag_id", pathParam0). + SetPathParam("run_id", pathParam1). + SetBody(body). + Post("dag-runs/{dag_id}/{run_id}") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *dagRunsClient) Trigger(ctx context.Context, dagId string, runId string, body *TriggerDAGRunPayload) error { + _, err := c.TriggerResponse(ctx, dagId, runId, body) + return err +} + +// ClearResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *dagRunsClient) ClearResponse(ctx context.Context, dagId string, runId string) (resp *resty.Response, err error) { + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "dag_id", runtime.ParamLocationPath, dagId) + if err != nil { + return nil, err + } + var pathParam1 string + + pathParam1, err = runtime.StyleParamWithLocation("simple", false, "run_id", runtime.ParamLocationPath, runId) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetPathParam("dag_id", pathParam0). + SetPathParam("run_id", pathParam1). + Post("dag-runs/{dag_id}/{run_id}/clear") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *dagRunsClient) Clear(ctx context.Context, dagId string, runId string) error { + _, err := c.ClearResponse(ctx, dagId, runId) + return err +} + +// GetStateResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *dagRunsClient) GetStateResponse(ctx context.Context, dagId string, runId string) (resp *resty.Response, err error) { + var res DagRunStateResponse + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "dag_id", runtime.ParamLocationPath, dagId) + if err != nil { + return nil, err + } + var pathParam1 string + + pathParam1, err = runtime.StyleParamWithLocation("simple", false, "run_id", runtime.ParamLocationPath, runId) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetPathParam("dag_id", pathParam0). + SetPathParam("run_id", pathParam1). + SetResult(&res). + Get("dag-runs/{dag_id}/{run_id}/state") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *dagRunsClient) GetState(ctx context.Context, dagId string, runId string) (*DagRunStateResponse, error) { + res, err := c.GetStateResponse(ctx, dagId, runId) + if err != nil { + return nil, err + } + + return res.Result().(*DagRunStateResponse), nil +} + +type taskInstancesClient struct { + *resty.Client +} + +// GetCountResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *taskInstancesClient) GetCountResponse(ctx context.Context, params *GetTaskInstanceCountParams) (resp *resty.Response, err error) { + var res int + + var param0 string + + param0, err = runtime.StyleParamWithLocation("form", true, "dag_id", runtime.ParamLocationQuery, params.DagId) + if err != nil { + return nil, err + } + var param1 string + + param1, err = runtime.StyleParamWithLocation("form", true, "task_ids", runtime.ParamLocationQuery, params.TaskIds) + if err != nil { + return nil, err + } + var param2 string + + param2, err = runtime.StyleParamWithLocation("form", true, "task_group_id", runtime.ParamLocationQuery, params.TaskGroupId) + if err != nil { + return nil, err + } + var param3 string + + param3, err = runtime.StyleParamWithLocation("form", true, "logical_dates", runtime.ParamLocationQuery, params.LogicalDates) + if err != nil { + return nil, err + } + var param4 string + + param4, err = runtime.StyleParamWithLocation("form", true, "run_ids", runtime.ParamLocationQuery, params.RunIds) + if err != nil { + return nil, err + } + var param5 string + + param5, err = runtime.StyleParamWithLocation("form", true, "states", runtime.ParamLocationQuery, params.States) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetQueryParam("dag_id", param0). + SetQueryParam("task_ids", param1). + SetQueryParam("task_group_id", param2). + SetQueryParam("logical_dates", param3). + SetQueryParam("run_ids", param4). + SetQueryParam("states", param5). + SetResult(&res). + Get("task-instances/count") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *taskInstancesClient) GetCount(ctx context.Context, params *GetTaskInstanceCountParams) (*int, error) { + res, err := c.GetCountResponse(ctx, params) + if err != nil { + return nil, err + } + + return res.Result().(*int), nil +} + +// GetStatesResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *taskInstancesClient) GetStatesResponse(ctx context.Context, params *GetTaskInstanceStatesParams) (resp *resty.Response, err error) { + var res TaskStatesResponse + + var param0 string + + param0, err = runtime.StyleParamWithLocation("form", true, "dag_id", runtime.ParamLocationQuery, params.DagId) + if err != nil { + return nil, err + } + var param1 string + + param1, err = runtime.StyleParamWithLocation("form", true, "task_ids", runtime.ParamLocationQuery, params.TaskIds) + if err != nil { + return nil, err + } + var param2 string + + param2, err = runtime.StyleParamWithLocation("form", true, "task_group_id", runtime.ParamLocationQuery, params.TaskGroupId) + if err != nil { + return nil, err + } + var param3 string + + param3, err = runtime.StyleParamWithLocation("form", true, "logical_dates", runtime.ParamLocationQuery, params.LogicalDates) + if err != nil { + return nil, err + } + var param4 string + + param4, err = runtime.StyleParamWithLocation("form", true, "run_ids", runtime.ParamLocationQuery, params.RunIds) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetQueryParam("dag_id", param0). + SetQueryParam("task_ids", param1). + SetQueryParam("task_group_id", param2). + SetQueryParam("logical_dates", param3). + SetQueryParam("run_ids", param4). + SetResult(&res). + Get("task-instances/states") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *taskInstancesClient) GetStates(ctx context.Context, params *GetTaskInstanceStatesParams) (*TaskStatesResponse, error) { + res, err := c.GetStatesResponse(ctx, params) + if err != nil { + return nil, err + } + + return res.Result().(*TaskStatesResponse), nil +} + +// HeartbeatResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *taskInstancesClient) HeartbeatResponse(ctx context.Context, taskInstanceId openapi_types.UUID, body *TIHeartbeatInfo) (resp *resty.Response, err error) { + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "task_instance_id", runtime.ParamLocationPath, taskInstanceId) + if err != nil { + return nil, err + } + + if body == nil { + return nil, fmt.Errorf("Heartbeat requires a non-nil body argument") + } + resp, err = c.R(). + SetContext(ctx). + SetPathParam("task_instance_id", pathParam0). + SetBody(body). + Put("task-instances/{task_instance_id}/heartbeat") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *taskInstancesClient) Heartbeat(ctx context.Context, taskInstanceId openapi_types.UUID, body *TIHeartbeatInfo) error { + _, err := c.HeartbeatResponse(ctx, taskInstanceId, body) + return err +} + +// GetPreviousSuccessfulDagrunResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *taskInstancesClient) GetPreviousSuccessfulDagrunResponse(ctx context.Context, taskInstanceId openapi_types.UUID) (resp *resty.Response, err error) { + var res PrevSuccessfulDagRunResponse + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "task_instance_id", runtime.ParamLocationPath, taskInstanceId) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetPathParam("task_instance_id", pathParam0). + SetResult(&res). + Get("task-instances/{task_instance_id}/previous-successful-dagrun") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *taskInstancesClient) GetPreviousSuccessfulDagrun(ctx context.Context, taskInstanceId openapi_types.UUID) (*PrevSuccessfulDagRunResponse, error) { + res, err := c.GetPreviousSuccessfulDagrunResponse(ctx, taskInstanceId) + if err != nil { + return nil, err + } + + return res.Result().(*PrevSuccessfulDagRunResponse), nil +} + +// PutRenderedFieldsResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *taskInstancesClient) PutRenderedFieldsResponse(ctx context.Context, taskInstanceId openapi_types.UUID, body *map[string]*JsonValue) (resp *resty.Response, err error) { + var res interface{} + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "task_instance_id", runtime.ParamLocationPath, taskInstanceId) + if err != nil { + return nil, err + } + + if body == nil { + return nil, fmt.Errorf("PutRenderedFields requires a non-nil body argument") + } + resp, err = c.R(). + SetContext(ctx). + SetPathParam("task_instance_id", pathParam0). + SetBody(body). + SetResult(&res). + Put("task-instances/{task_instance_id}/rtif") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *taskInstancesClient) PutRenderedFields(ctx context.Context, taskInstanceId openapi_types.UUID, body *map[string]*JsonValue) (*interface{}, error) { + res, err := c.PutRenderedFieldsResponse(ctx, taskInstanceId, body) + if err != nil { + return nil, err + } + + return res.Result().(*interface{}), nil +} + +// RunResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *taskInstancesClient) RunResponse(ctx context.Context, taskInstanceId openapi_types.UUID, body *TIEnterRunningPayload) (resp *resty.Response, err error) { + var res TIRunContext + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "task_instance_id", runtime.ParamLocationPath, taskInstanceId) + if err != nil { + return nil, err + } + + if body == nil { + return nil, fmt.Errorf("Run requires a non-nil body argument") + } + resp, err = c.R(). + SetContext(ctx). + SetPathParam("task_instance_id", pathParam0). + SetBody(body). + SetResult(&res). + Patch("task-instances/{task_instance_id}/run") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *taskInstancesClient) Run(ctx context.Context, taskInstanceId openapi_types.UUID, body *TIEnterRunningPayload) (*TIRunContext, error) { + res, err := c.RunResponse(ctx, taskInstanceId, body) + if err != nil { + return nil, err + } + + return res.Result().(*TIRunContext), nil +} + +// SkipDownstreamResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *taskInstancesClient) SkipDownstreamResponse(ctx context.Context, taskInstanceId openapi_types.UUID, body *TISkippedDownstreamTasksStatePayload) (resp *resty.Response, err error) { + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "task_instance_id", runtime.ParamLocationPath, taskInstanceId) + if err != nil { + return nil, err + } + + if body == nil { + return nil, fmt.Errorf("SkipDownstream requires a non-nil body argument") + } + resp, err = c.R(). + SetContext(ctx). + SetPathParam("task_instance_id", pathParam0). + SetBody(body). + Patch("task-instances/{task_instance_id}/skip-downstream") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *taskInstancesClient) SkipDownstream(ctx context.Context, taskInstanceId openapi_types.UUID, body *TISkippedDownstreamTasksStatePayload) error { + _, err := c.SkipDownstreamResponse(ctx, taskInstanceId, body) + return err +} + +// UpdateStateResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *taskInstancesClient) UpdateStateResponse(ctx context.Context, taskInstanceId openapi_types.UUID, body *TIUpdateStatePayload) (resp *resty.Response, err error) { + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "task_instance_id", runtime.ParamLocationPath, taskInstanceId) + if err != nil { + return nil, err + } + + if body == nil { + return nil, fmt.Errorf("UpdateState requires a non-nil body argument") + } + resp, err = c.R(). + SetContext(ctx). + SetPathParam("task_instance_id", pathParam0). + SetBody(body). + Patch("task-instances/{task_instance_id}/state") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *taskInstancesClient) UpdateState(ctx context.Context, taskInstanceId openapi_types.UUID, body *TIUpdateStatePayload) error { + _, err := c.UpdateStateResponse(ctx, taskInstanceId, body) + return err +} + +type taskReschedulesClient struct { + *resty.Client +} + +// GetStartDateResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *taskReschedulesClient) GetStartDateResponse(ctx context.Context, taskInstanceId openapi_types.UUID) (resp *resty.Response, err error) { + var res time.Time + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "task_instance_id", runtime.ParamLocationPath, taskInstanceId) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetPathParam("task_instance_id", pathParam0). + SetResult(&res). + Get("task-reschedules/{task_instance_id}/start_date") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *taskReschedulesClient) GetStartDate(ctx context.Context, taskInstanceId openapi_types.UUID) (*time.Time, error) { + res, err := c.GetStartDateResponse(ctx, taskInstanceId) + if err != nil { + return nil, err + } + + return res.Result().(*time.Time), nil +} + +type variablesClient struct { + *resty.Client +} + +// DeleteResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *variablesClient) DeleteResponse(ctx context.Context, variableKey string) (resp *resty.Response, err error) { + var res HTTPValidationError + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "variable_key", runtime.ParamLocationPath, variableKey) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetPathParam("variable_key", pathParam0). + SetResult(&res). + Delete("variables/{variable_key}") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *variablesClient) Delete(ctx context.Context, variableKey string) (*HTTPValidationError, error) { + res, err := c.DeleteResponse(ctx, variableKey) + if err != nil { + return nil, err + } + + return res.Result().(*HTTPValidationError), nil +} + +// GetResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *variablesClient) GetResponse(ctx context.Context, variableKey string) (resp *resty.Response, err error) { + var res VariableResponse + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "variable_key", runtime.ParamLocationPath, variableKey) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetPathParam("variable_key", pathParam0). + SetResult(&res). + Get("variables/{variable_key}") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *variablesClient) Get(ctx context.Context, variableKey string) (*VariableResponse, error) { + res, err := c.GetResponse(ctx, variableKey) + if err != nil { + return nil, err + } + + return res.Result().(*VariableResponse), nil +} + +// PutResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *variablesClient) PutResponse(ctx context.Context, variableKey string, body *VariablePostBody) (resp *resty.Response, err error) { + var res interface{} + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "variable_key", runtime.ParamLocationPath, variableKey) + if err != nil { + return nil, err + } + + if body == nil { + return nil, fmt.Errorf("Put requires a non-nil body argument") + } + resp, err = c.R(). + SetContext(ctx). + SetPathParam("variable_key", pathParam0). + SetBody(body). + SetResult(&res). + Put("variables/{variable_key}") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *variablesClient) Put(ctx context.Context, variableKey string, body *VariablePostBody) (*interface{}, error) { + res, err := c.PutResponse(ctx, variableKey, body) + if err != nil { + return nil, err + } + + return res.Result().(*interface{}), nil +} + +type xcomsClient struct { + *resty.Client +} + +// DeleteResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *xcomsClient) DeleteResponse(ctx context.Context, dagId string, runId string, taskId string, key string, params *DeleteXcomParams) (resp *resty.Response, err error) { + var res interface{} + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "dag_id", runtime.ParamLocationPath, dagId) + if err != nil { + return nil, err + } + var pathParam1 string + + pathParam1, err = runtime.StyleParamWithLocation("simple", false, "run_id", runtime.ParamLocationPath, runId) + if err != nil { + return nil, err + } + var pathParam2 string + + pathParam2, err = runtime.StyleParamWithLocation("simple", false, "task_id", runtime.ParamLocationPath, taskId) + if err != nil { + return nil, err + } + var pathParam3 string + + pathParam3, err = runtime.StyleParamWithLocation("simple", false, "key", runtime.ParamLocationPath, key) + if err != nil { + return nil, err + } + + var param0 string + + param0, err = runtime.StyleParamWithLocation("form", true, "map_index", runtime.ParamLocationQuery, params.MapIndex) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetPathParam("dag_id", pathParam0). + SetPathParam("run_id", pathParam1). + SetPathParam("task_id", pathParam2). + SetPathParam("key", pathParam3). + SetQueryParam("map_index", param0). + SetResult(&res). + Delete("xcoms/{dag_id}/{run_id}/{task_id}/{key}") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *xcomsClient) Delete(ctx context.Context, dagId string, runId string, taskId string, key string, params *DeleteXcomParams) (*interface{}, error) { + res, err := c.DeleteResponse(ctx, dagId, runId, taskId, key, params) + if err != nil { + return nil, err + } + + return res.Result().(*interface{}), nil +} + +// GetResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *xcomsClient) GetResponse(ctx context.Context, dagId string, runId string, taskId string, key string, params *GetXcomParams) (resp *resty.Response, err error) { + var res XComResponse + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "dag_id", runtime.ParamLocationPath, dagId) + if err != nil { + return nil, err + } + var pathParam1 string + + pathParam1, err = runtime.StyleParamWithLocation("simple", false, "run_id", runtime.ParamLocationPath, runId) + if err != nil { + return nil, err + } + var pathParam2 string + + pathParam2, err = runtime.StyleParamWithLocation("simple", false, "task_id", runtime.ParamLocationPath, taskId) + if err != nil { + return nil, err + } + var pathParam3 string + + pathParam3, err = runtime.StyleParamWithLocation("simple", false, "key", runtime.ParamLocationPath, key) + if err != nil { + return nil, err + } + + var param0 string + + param0, err = runtime.StyleParamWithLocation("form", true, "map_index", runtime.ParamLocationQuery, params.MapIndex) + if err != nil { + return nil, err + } + var param1 string + + param1, err = runtime.StyleParamWithLocation("form", true, "include_prior_dates", runtime.ParamLocationQuery, params.IncludePriorDates) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetPathParam("dag_id", pathParam0). + SetPathParam("run_id", pathParam1). + SetPathParam("task_id", pathParam2). + SetPathParam("key", pathParam3). + SetQueryParam("map_index", param0). + SetQueryParam("include_prior_dates", param1). + SetResult(&res). + Get("xcoms/{dag_id}/{run_id}/{task_id}/{key}") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *xcomsClient) Get(ctx context.Context, dagId string, runId string, taskId string, key string, params *GetXcomParams) (*XComResponse, error) { + res, err := c.GetResponse(ctx, dagId, runId, taskId, key, params) + if err != nil { + return nil, err + } + + return res.Result().(*XComResponse), nil +} + +// HeadResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *xcomsClient) HeadResponse(ctx context.Context, dagId string, runId string, taskId string, key string, params *HeadXcomParams) (resp *resty.Response, err error) { + var res interface{} + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "dag_id", runtime.ParamLocationPath, dagId) + if err != nil { + return nil, err + } + var pathParam1 string + + pathParam1, err = runtime.StyleParamWithLocation("simple", false, "run_id", runtime.ParamLocationPath, runId) + if err != nil { + return nil, err + } + var pathParam2 string + + pathParam2, err = runtime.StyleParamWithLocation("simple", false, "task_id", runtime.ParamLocationPath, taskId) + if err != nil { + return nil, err + } + var pathParam3 string + + pathParam3, err = runtime.StyleParamWithLocation("simple", false, "key", runtime.ParamLocationPath, key) + if err != nil { + return nil, err + } + + var param0 string + + param0, err = runtime.StyleParamWithLocation("form", true, "map_index", runtime.ParamLocationQuery, params.MapIndex) + if err != nil { + return nil, err + } + + resp, err = c.R(). + SetContext(ctx). + SetPathParam("dag_id", pathParam0). + SetPathParam("run_id", pathParam1). + SetPathParam("task_id", pathParam2). + SetPathParam("key", pathParam3). + SetQueryParam("map_index", param0). + SetResult(&res). + Head("xcoms/{dag_id}/{run_id}/{task_id}/{key}") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *xcomsClient) Head(ctx context.Context, dagId string, runId string, taskId string, key string, params *HeadXcomParams) (*interface{}, error) { + res, err := c.HeadResponse(ctx, dagId, runId, taskId, key, params) + if err != nil { + return nil, err + } + + return res.Result().(*interface{}), nil +} + +// SetResponse performs the HTTP request and returns the lower level [resty.Response] +func (c *xcomsClient) SetResponse(ctx context.Context, dagId string, runId string, taskId string, key string, params *SetXcomParams, body *JsonValue) (resp *resty.Response, err error) { + var res interface{} + + var pathParam0 string + + pathParam0, err = runtime.StyleParamWithLocation("simple", false, "dag_id", runtime.ParamLocationPath, dagId) + if err != nil { + return nil, err + } + var pathParam1 string + + pathParam1, err = runtime.StyleParamWithLocation("simple", false, "run_id", runtime.ParamLocationPath, runId) + if err != nil { + return nil, err + } + var pathParam2 string + + pathParam2, err = runtime.StyleParamWithLocation("simple", false, "task_id", runtime.ParamLocationPath, taskId) + if err != nil { + return nil, err + } + var pathParam3 string + + pathParam3, err = runtime.StyleParamWithLocation("simple", false, "key", runtime.ParamLocationPath, key) + if err != nil { + return nil, err + } + + var param0 string + + param0, err = runtime.StyleParamWithLocation("form", true, "map_index", runtime.ParamLocationQuery, params.MapIndex) + if err != nil { + return nil, err + } + var param1 string + + param1, err = runtime.StyleParamWithLocation("form", true, "mapped_length", runtime.ParamLocationQuery, params.MappedLength) + if err != nil { + return nil, err + } + + if body == nil { + return nil, fmt.Errorf("Set requires a non-nil body argument") + } + resp, err = c.R(). + SetContext(ctx). + SetPathParam("dag_id", pathParam0). + SetPathParam("run_id", pathParam1). + SetPathParam("task_id", pathParam2). + SetPathParam("key", pathParam3). + SetQueryParam("map_index", param0). + SetQueryParam("mapped_length", param1). + SetBody(body). + SetResult(&res). + Post("xcoms/{dag_id}/{run_id}/{task_id}/{key}") + if err != nil { + return resp, err + } + return resp, HandleError(c.Client, resp) +} + +func (c *xcomsClient) Set(ctx context.Context, dagId string, runId string, taskId string, key string, params *SetXcomParams, body *JsonValue) (*interface{}, error) { + res, err := c.SetResponse(ctx, dagId, runId, taskId, key, params, body) + if err != nil { + return nil, err + } + + return res.Result().(*interface{}), nil +} + +type AssetEventsClient interface { + GetByAssetNameUri(ctx context.Context, params *GetAssetEventByAssetNameUriParams) (*AssetEventsResponse, error) + // GetByAssetNameUriResponse is a lower level version of [GetByAssetNameUri] and provides access to the raw [resty.Response] + GetByAssetNameUriResponse(ctx context.Context, params *GetAssetEventByAssetNameUriParams) (*resty.Response, error) + + GetByAssetAlias(ctx context.Context, params *GetAssetEventByAssetAliasParams) (*AssetEventsResponse, error) + // GetByAssetAliasResponse is a lower level version of [GetByAssetAlias] and provides access to the raw [resty.Response] + GetByAssetAliasResponse(ctx context.Context, params *GetAssetEventByAssetAliasParams) (*resty.Response, error) +} + +var _ AssetEventsClient = (*assetEventsClient)(nil) + +type AssetsClient interface { + // Get an Airflow Asset by `name`. + GetByName(ctx context.Context, params *GetAssetByNameParams) (*AssetResponse, error) + // GetByNameResponse is a lower level version of [GetByName] and provides access to the raw [resty.Response] + GetByNameResponse(ctx context.Context, params *GetAssetByNameParams) (*resty.Response, error) + + // Get an Airflow Asset by `uri`. + GetByUri(ctx context.Context, params *GetAssetByUriParams) (*AssetResponse, error) + // GetByUriResponse is a lower level version of [GetByUri] and provides access to the raw [resty.Response] + GetByUriResponse(ctx context.Context, params *GetAssetByUriParams) (*resty.Response, error) +} + +var _ AssetsClient = (*assetsClient)(nil) + +type ConnectionsClient interface { + // Get an Airflow connection. + Get(ctx context.Context, connectionId string) (*ConnectionResponse, error) + // GetResponse is a lower level version of [Get] and provides access to the raw [resty.Response] + GetResponse(ctx context.Context, connectionId string) (*resty.Response, error) +} + +var _ ConnectionsClient = (*connectionsClient)(nil) + +type DagRunsClient interface { + // Get the count of DAG runs matching the given criteria. + GetDrCount(ctx context.Context, params *GetDrCountParams) (*int, error) + // GetDrCountResponse is a lower level version of [GetDrCount] and provides access to the raw [resty.Response] + GetDrCountResponse(ctx context.Context, params *GetDrCountParams) (*resty.Response, error) + + // Trigger a DAG Run. + Trigger(ctx context.Context, dagId string, runId string, body *TriggerDAGRunPayload) error + // TriggerResponse is a lower level version of [Trigger] and provides access to the raw [resty.Response] + TriggerResponse(ctx context.Context, dagId string, runId string, body *TriggerDAGRunPayload) (*resty.Response, error) + + // Clear a DAG Run. + Clear(ctx context.Context, dagId string, runId string) error + // ClearResponse is a lower level version of [Clear] and provides access to the raw [resty.Response] + ClearResponse(ctx context.Context, dagId string, runId string) (*resty.Response, error) + + // Get a DAG Run State. + GetState(ctx context.Context, dagId string, runId string) (*DagRunStateResponse, error) + // GetStateResponse is a lower level version of [GetState] and provides access to the raw [resty.Response] + GetStateResponse(ctx context.Context, dagId string, runId string) (*resty.Response, error) +} + +var _ DagRunsClient = (*dagRunsClient)(nil) + +type TaskInstancesClient interface { + // Get the count of task instances matching the given criteria. + GetCount(ctx context.Context, params *GetTaskInstanceCountParams) (*int, error) + // GetCountResponse is a lower level version of [GetCount] and provides access to the raw [resty.Response] + GetCountResponse(ctx context.Context, params *GetTaskInstanceCountParams) (*resty.Response, error) + + // Get the states for Task Instances with the given criteria. + GetStates(ctx context.Context, params *GetTaskInstanceStatesParams) (*TaskStatesResponse, error) + // GetStatesResponse is a lower level version of [GetStates] and provides access to the raw [resty.Response] + GetStatesResponse(ctx context.Context, params *GetTaskInstanceStatesParams) (*resty.Response, error) + + // Update the heartbeat of a TaskInstance to mark it as alive & still running. + Heartbeat(ctx context.Context, taskInstanceId openapi_types.UUID, body *TIHeartbeatInfo) error + // HeartbeatResponse is a lower level version of [Heartbeat] and provides access to the raw [resty.Response] + HeartbeatResponse(ctx context.Context, taskInstanceId openapi_types.UUID, body *TIHeartbeatInfo) (*resty.Response, error) + + // Get the previous successful DagRun for a TaskInstance. + // + // The data from this endpoint is used to get values for Task Context. + GetPreviousSuccessfulDagrun(ctx context.Context, taskInstanceId openapi_types.UUID) (*PrevSuccessfulDagRunResponse, error) + // GetPreviousSuccessfulDagrunResponse is a lower level version of [GetPreviousSuccessfulDagrun] and provides access to the raw [resty.Response] + GetPreviousSuccessfulDagrunResponse(ctx context.Context, taskInstanceId openapi_types.UUID) (*resty.Response, error) + + // Add an RTIF entry for a task instance, sent by the worker. + PutRenderedFields(ctx context.Context, taskInstanceId openapi_types.UUID, body *map[string]*JsonValue) (*interface{}, error) + // PutRenderedFieldsResponse is a lower level version of [PutRenderedFields] and provides access to the raw [resty.Response] + PutRenderedFieldsResponse(ctx context.Context, taskInstanceId openapi_types.UUID, body *map[string]*JsonValue) (*resty.Response, error) + + // Run a TaskInstance. + // + // This endpoint is used to start a TaskInstance that is in the QUEUED state. + Run(ctx context.Context, taskInstanceId openapi_types.UUID, body *TIEnterRunningPayload) (*TIRunContext, error) + // RunResponse is a lower level version of [Run] and provides access to the raw [resty.Response] + RunResponse(ctx context.Context, taskInstanceId openapi_types.UUID, body *TIEnterRunningPayload) (*resty.Response, error) + + SkipDownstream(ctx context.Context, taskInstanceId openapi_types.UUID, body *TISkippedDownstreamTasksStatePayload) error + // SkipDownstreamResponse is a lower level version of [SkipDownstream] and provides access to the raw [resty.Response] + SkipDownstreamResponse(ctx context.Context, taskInstanceId openapi_types.UUID, body *TISkippedDownstreamTasksStatePayload) (*resty.Response, error) + + // Update the state of a TaskInstance. + // + // Not all state transitions are valid, and transitioning to some states requires extra information to be + // passed along. (Check out the datamodels for details, the rendered docs might not reflect this accurately) + UpdateState(ctx context.Context, taskInstanceId openapi_types.UUID, body *TIUpdateStatePayload) error + // UpdateStateResponse is a lower level version of [UpdateState] and provides access to the raw [resty.Response] + UpdateStateResponse(ctx context.Context, taskInstanceId openapi_types.UUID, body *TIUpdateStatePayload) (*resty.Response, error) +} + +var _ TaskInstancesClient = (*taskInstancesClient)(nil) + +type TaskReschedulesClient interface { + // Get the first reschedule date if found, None if no records exist. + GetStartDate(ctx context.Context, taskInstanceId openapi_types.UUID) (*time.Time, error) + // GetStartDateResponse is a lower level version of [GetStartDate] and provides access to the raw [resty.Response] + GetStartDateResponse(ctx context.Context, taskInstanceId openapi_types.UUID) (*resty.Response, error) +} + +var _ TaskReschedulesClient = (*taskReschedulesClient)(nil) + +type VariablesClient interface { + // Delete an Airflow Variable. + Delete(ctx context.Context, variableKey string) (*HTTPValidationError, error) + // DeleteResponse is a lower level version of [Delete] and provides access to the raw [resty.Response] + DeleteResponse(ctx context.Context, variableKey string) (*resty.Response, error) + + // Get an Airflow Variable. + Get(ctx context.Context, variableKey string) (*VariableResponse, error) + // GetResponse is a lower level version of [Get] and provides access to the raw [resty.Response] + GetResponse(ctx context.Context, variableKey string) (*resty.Response, error) + + // Set an Airflow Variable. + Put(ctx context.Context, variableKey string, body *VariablePostBody) (*interface{}, error) + // PutResponse is a lower level version of [Put] and provides access to the raw [resty.Response] + PutResponse(ctx context.Context, variableKey string, body *VariablePostBody) (*resty.Response, error) +} + +var _ VariablesClient = (*variablesClient)(nil) + +type XcomsClient interface { + // Delete a single XCom Value + Delete(ctx context.Context, dagId string, runId string, taskId string, key string, params *DeleteXcomParams) (*interface{}, error) + // DeleteResponse is a lower level version of [Delete] and provides access to the raw [resty.Response] + DeleteResponse(ctx context.Context, dagId string, runId string, taskId string, key string, params *DeleteXcomParams) (*resty.Response, error) + + // Get a single XCom Value + Get(ctx context.Context, dagId string, runId string, taskId string, key string, params *GetXcomParams) (*XComResponse, error) + // GetResponse is a lower level version of [Get] and provides access to the raw [resty.Response] + GetResponse(ctx context.Context, dagId string, runId string, taskId string, key string, params *GetXcomParams) (*resty.Response, error) + + // Returns the count of mapped XCom values found in the `Content-Range` response header + Head(ctx context.Context, dagId string, runId string, taskId string, key string, params *HeadXcomParams) (*interface{}, error) + // HeadResponse is a lower level version of [Head] and provides access to the raw [resty.Response] + HeadResponse(ctx context.Context, dagId string, runId string, taskId string, key string, params *HeadXcomParams) (*resty.Response, error) + + // Set an Airflow XCom. + Set(ctx context.Context, dagId string, runId string, taskId string, key string, params *SetXcomParams, body *JsonValue) (*interface{}, error) + // SetResponse is a lower level version of [Set] and provides access to the raw [resty.Response] + SetResponse(ctx context.Context, dagId string, runId string, taskId string, key string, params *SetXcomParams, body *JsonValue) (*resty.Response, error) +} + +var _ XcomsClient = (*xcomsClient)(nil) diff --git a/go-sdk/pkg/api/client.go b/go-sdk/pkg/api/client.go new file mode 100644 index 00000000000..d40760fc58d --- /dev/null +++ b/go-sdk/pkg/api/client.go @@ -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...) +} diff --git a/go-sdk/pkg/api/init.go b/go-sdk/pkg/api/init.go new file mode 100644 index 00000000000..0b494c1a1cd --- /dev/null +++ b/go-sdk/pkg/api/init.go @@ -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()) + } +} diff --git a/go-sdk/pkg/api/models.go b/go-sdk/pkg/api/models.go new file mode 100644 index 00000000000..1dba6b07d41 --- /dev/null +++ b/go-sdk/pkg/api/models.go @@ -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"` +} diff --git a/go-sdk/pkg/api/oapi-codegen.yml b/go-sdk/pkg/api/oapi-codegen.yml new file mode 100644 index 00000000000..6f5422b091d --- /dev/null +++ b/go-sdk/pkg/api/oapi-codegen.yml @@ -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 diff --git a/go-sdk/pkg/api/overlay.yml b/go-sdk/pkg/api/overlay.yml new file mode 100644 index 00000000000..cfd7ae4d8ce --- /dev/null +++ b/go-sdk/pkg/api/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" diff --git a/go-sdk/pkg/logging/attrs.go b/go-sdk/pkg/logging/attrs.go new file mode 100644 index 00000000000..dde80914f7f --- /dev/null +++ b/go-sdk/pkg/logging/attrs.go @@ -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) +} diff --git a/go-sdk/pkg/logging/stdbridge.go b/go-sdk/pkg/logging/stdbridge.go new file mode 100644 index 00000000000..8733d718a45 --- /dev/null +++ b/go-sdk/pkg/logging/stdbridge.go @@ -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...) +} diff --git a/go-sdk/pkg/logging/tee.go b/go-sdk/pkg/logging/tee.go new file mode 100644 index 00000000000..eaf28d5d763 --- /dev/null +++ b/go-sdk/pkg/logging/tee.go @@ -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 +} diff --git a/go-sdk/pkg/sdkcontext/keys.go b/go-sdk/pkg/sdkcontext/keys.go new file mode 100644 index 00000000000..0dbc2c60194 --- /dev/null +++ b/go-sdk/pkg/sdkcontext/keys.go @@ -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{} +) diff --git a/go-sdk/sdk/doc.go b/go-sdk/sdk/doc.go new file mode 100644 index 00000000000..73e540fcd19 --- /dev/null +++ b/go-sdk/sdk/doc.go @@ -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 diff --git a/go-sdk/sdk/variable.go b/go-sdk/sdk/variable.go new file mode 100644 index 00000000000..715778a5f08 --- /dev/null +++ b/go-sdk/sdk/variable.go @@ -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 +} diff --git a/go-sdk/worker/init.go b/go-sdk/worker/init.go new file mode 100644 index 00000000000..0e28ab404e0 --- /dev/null +++ b/go-sdk/worker/init.go @@ -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() +} diff --git a/go-sdk/worker/registry.go b/go-sdk/worker/registry.go new file mode 100644 index 00000000000..04fb8712905 --- /dev/null +++ b/go-sdk/worker/registry.go @@ -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 ®istry{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 +} diff --git a/go-sdk/worker/registry_test.go b/go-sdk/worker/registry_test.go new file mode 100644 index 00000000000..5415325bad9 --- /dev/null +++ b/go-sdk/worker/registry_test.go @@ -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) + }) +} diff --git a/go-sdk/worker/runner.go b/go-sdk/worker/runner.go new file mode 100644 index 00000000000..f01cde3fdaf --- /dev/null +++ b/go-sdk/worker/runner.go @@ -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 +} diff --git a/go-sdk/worker/runner_test.go b/go-sdk/worker/runner_test.go new file mode 100644 index 00000000000..0787bb3adbe --- /dev/null +++ b/go-sdk/worker/runner_test.go @@ -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 "], + 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") +} diff --git a/go-sdk/worker/task.go b/go-sdk/worker/task.go new file mode 100644 index 00000000000..7e0703195f9 --- /dev/null +++ b/go-sdk/worker/task.go @@ -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 + // `, error`, or just `error` + if fnType.NumOut() < 1 || fnType.NumOut() > 2 { + return fmt.Errorf( + "task function %s has %d return values, must be `, 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) +}