airflow/go-sdk/example/bundle/taskflowbinding/taskflowbinding.go

// 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 taskflowbinding contains TaskFlow argument-binding examples.
package taskflowbinding

import (
	"fmt"
	"reflect"

	"github.com/apache/airflow/go-sdk/airflow"
)

// Config is an object passed through XCom.
type Config struct {
	Environment string `json:"environment"`
	Region      string `json:"region"`
	Debug       bool   `json:"debug"`
}

// MakeConfig returns a Config XCom.
func MakeConfig(actx airflow.Context) (any, error) {
	cfg := Config{Environment: "production", Region: "eu-west-1", Debug: true}
	actx.Logger().InfoContext(
		actx,
		"Pushing config",
		"environment",
		cfg.Environment,
		"region",
		cfg.Region,
		"debug",
		cfg.Debug,
	)
	return cfg, nil
}

// MakeNumbers returns an integer-slice XCom.
func MakeNumbers(actx airflow.Context) (any, error) {
	numbers := []int{1, 1, 2, 3, 5, 8}
	actx.Logger().InfoContext(actx, "Pushing numbers", "numbers", fmt.Sprint(numbers))
	return numbers, nil
}

// MakeRegion returns a region XCom.
func MakeRegion(actx airflow.Context) (any, error) {
	region := "eu-west-1"
	actx.Logger().InfoContext(actx, "Pushing region", "region", region)
	return region, nil
}

// ViaFlatArgs exercises positional binding for literals, defaults, and XComs.
func ViaFlatArgs(
	actx airflow.Context,
	name string,
	count int,
	ratio float64,
	enabled bool,
	tags []string,
	config Config,
	numbers []int,
	note *string,
) (any, error) {
	if name != "summary" || count != 3 || ratio != 2.5 || !enabled {
		return nil, fmt.Errorf(
			"scalar literals bound incorrectly: name=%q count=%d ratio=%v enabled=%v",
			name, count, ratio, enabled,
		)
	}
	if want := []string{"metrics", "hourly"}; !reflect.DeepEqual(tags, want) {
		return nil, fmt.Errorf("array literal bound incorrectly: tags=%v, want %v", tags, want)
	}
	if want := (Config{Environment: "production", Region: "eu-west-1", Debug: true}); config != want {
		return nil, fmt.Errorf("object XCom bound incorrectly: config=%+v, want %+v", config, want)
	}
	if want := []int{1, 1, 2, 3, 5, 8}; !reflect.DeepEqual(numbers, want) {
		return nil, fmt.Errorf("array XCom bound incorrectly: numbers=%v, want %v", numbers, want)
	}
	if note != nil {
		return nil, fmt.Errorf("defaulted None bound incorrectly: note=%q, want nil", *note)
	}

	sum := 0
	for _, n := range numbers {
		sum += n
	}
	actx.Logger().InfoContext(actx, "Bound TaskFlow arguments",
		"name", name,
		"count", count,
		"ratio", ratio,
		"enabled", enabled,
		"tags", fmt.Sprint(tags),
		"environment", config.Environment,
		"sum", sum,
	)
	return map[string]any{
		"name":          name,
		"count":         count,
		"ratio":         ratio,
		"enabled":       enabled,
		"tags":          tags,
		"environment":   config.Environment,
		"debug":         config.Debug,
		"sum":           sum,
		"note_was_null": note == nil,
	}, nil
}

// ViaStructNoTagsInput binds fields by folded Go name.
type ViaStructNoTagsInput struct {
	RegionCode string
	Threshold  float64
}

// ViaStructNoTags exercises binding without `arg:` tags.
func ViaStructNoTags(actx airflow.Context, input ViaStructNoTagsInput) (any, error) {
	if input.RegionCode != "eu-west-1" || input.Threshold != 0.75 {
		return nil, fmt.Errorf(
			"struct fields bound incorrectly: region_code=%q threshold=%v",
			input.RegionCode,
			input.Threshold,
		)
	}

	actx.Logger().InfoContext(actx, "Bound struct (no tags)",
		"region_code", input.RegionCode,
		"threshold", input.Threshold,
	)
	return map[string]any{
		"region_code": input.RegionCode,
		"threshold":   input.Threshold,
	}, nil
}

// ViaStructArgTagInput binds fields with explicit `arg:` tags.
type ViaStructArgTagInput struct {
	Region    string  `arg:"region_code"`
	Threshold float64 `arg:"threshold"`
}

// ViaStructArgTag exercises explicit field-name binding.
func ViaStructArgTag(actx airflow.Context, input ViaStructArgTagInput) (any, error) {
	if input.Region != "eu-west-1" || input.Threshold != 0.75 {
		return nil, fmt.Errorf(
			"struct fields bound incorrectly: region=%q threshold=%v",
			input.Region,
			input.Threshold,
		)
	}

	actx.Logger().InfoContext(actx, "Bound struct (arg: tag)",
		"region", input.Region,
		"threshold", input.Threshold,
	)
	return map[string]any{
		"region":    input.Region,
		"threshold": input.Threshold,
	}, nil
}

// ViaStructDefaultArgInput claims only the explicitly passed argument.
type ViaStructDefaultArgInput struct {
	Region string `arg:"region_code"`
}

// ViaStructDefaultArg exercises a captured stub default no struct field claims.
func ViaStructDefaultArg(
	actx airflow.Context,
	input ViaStructDefaultArgInput,
) (any, error) {
	if input.Region != "eu-west-1" {
		return nil, fmt.Errorf("struct field bound incorrectly: region=%q", input.Region)
	}

	actx.Logger().
		InfoContext(actx, "Bound struct (defaulted arg unclaimed)", "region", input.Region)
	return map[string]any{"region": input.Region}, nil
}

// ViaStructMoreArgsInput declares fewer fields than its call passes arguments.
type ViaStructMoreArgsInput struct {
	Region string `arg:"region_code"`
}

// ViaStructMoreArgs exercises the Dag passing an argument the struct does not
// declare: warned about, and the task still runs.
func ViaStructMoreArgs(actx airflow.Context, input ViaStructMoreArgsInput) (any, error) {
	if input.Region != "eu-west-1" {
		return nil, fmt.Errorf("struct field bound incorrectly: region=%q", input.Region)
	}

	actx.Logger().InfoContext(actx, "Bound struct (call passed more)", "region", input.Region)
	return map[string]any{"region": input.Region}, nil
}

// ViaStructFewerArgsInput declares more fields than its call passes arguments.
type ViaStructFewerArgsInput struct {
	Region   string `arg:"region_code"`
	NotInDag string `arg:"not_in_dag"`
}

// ViaStructFewerArgs exercises the struct declaring an argument the Dag's call
// does not pass: warned about, and the field keeps its Go zero value.
func ViaStructFewerArgs(actx airflow.Context, input ViaStructFewerArgsInput) (any, error) {
	if input.Region != "eu-west-1" {
		return nil, fmt.Errorf("struct field bound incorrectly: region=%q", input.Region)
	}
	if input.NotInDag != "" {
		return nil, fmt.Errorf(
			"expected the undeclared field to keep its zero value, got not_in_dag=%q",
			input.NotInDag,
		)
	}

	actx.Logger().InfoContext(actx, "Bound struct (call passed fewer)",
		"region", input.Region,
		"not_in_dag_was_empty", input.NotInDag == "",
	)
	return map[string]any{
		"region":               input.Region,
		"not_in_dag_was_empty": input.NotInDag == "",
	}, nil
}

// FlatMapConfig receives one dict as a whole value.
type FlatMapConfig struct {
	Region string `json:"region"`
	Count  int    `json:"count"`
}

// ViaFlatMap exercises whole-value dict decoding into a struct.
func ViaFlatMap(actx airflow.Context, config FlatMapConfig) (any, error) {
	if config.Region != "eu-west-1" || config.Count != 3 {
		return nil, fmt.Errorf(
			"whole-value map bound incorrectly: region=%q count=%d",
			config.Region,
			config.Count,
		)
	}

	actx.Logger().InfoContext(actx, "Bound whole map into struct",
		"region", config.Region,
		"count", config.Count,
	)
	return map[string]any{"region": config.Region, "count": config.Count}, nil
}

// StructMapInput binds a dict to a map field.
type StructMapInput struct {
	Payload map[string]any `arg:"payload"`
}

// ViaStructMap exercises dict binding to a struct field.
func ViaStructMap(actx airflow.Context, input StructMapInput) (any, error) {
	region, _ := input.Payload["region"].(string)
	if region != "eu-west-1" {
		return nil, fmt.Errorf("map field bound incorrectly: payload=%v", input.Payload)
	}

	actx.Logger().
		InfoContext(actx, "Bound map onto struct field", "payload", fmt.Sprint(input.Payload))
	return map[string]any{"payload": input.Payload}, nil
}

// ViaPlainMap exercises dict decoding into a typed map.
func ViaPlainMap(actx airflow.Context, labels map[string]string) (any, error) {
	if labels["team"] != "data" || labels["tier"] != "gold" {
		return nil, fmt.Errorf("plain map bound incorrectly: labels=%v", labels)
	}

	actx.Logger().InfoContext(actx, "Bound dict into a plain map", "labels", fmt.Sprint(labels))
	return map[string]any{"team": labels["team"], "tier": labels["tier"]}, nil
}