diff --git a/airflow-e2e-tests/docker/go.yml b/airflow-e2e-tests/docker/go.yml index c142696376586..38fde7634d2c9 100644 --- a/airflow-e2e-tests/docker/go.yml +++ b/airflow-e2e-tests/docker/go.yml @@ -22,8 +22,15 @@ # image can exec it directly -- no extra runtime needs to be installed. We only # bind-mount the bundle into the directory the ExecutableCoordinator scans and # point the worker at the "golang" queue where @task.stub tasks are routed. +# The native-Dag bundle sits in the Dags folder, so the Dag processor execs it +# to parse its Dags, with the same USER and HOME as the worker. --- services: + airflow-dag-processor: + environment: + USER: airflow + HOME: /home/airflow + airflow-worker: # The bundle is built with CGO_ENABLED=0, so the SDK's user.Current() call at # init falls back to Go's pure-Go resolver, which reads $USER / $HOME and diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py index 9814cb5147f6e..e57a8df764f98 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py @@ -44,6 +44,8 @@ GO_SDK_BUNDLE_NAME, GO_SDK_DAGS_PATH, GO_SDK_EXAMPLE_BUNDLE_PKG, + GO_SDK_NATIVE_BUNDLE_NAME, + GO_SDK_NATIVE_BUNDLE_PKG, GO_SDK_ROOT_PATH, GO_SDK_STATE_STORE_RETENTION_DAYS, JAVA_COMPOSE_PATH, @@ -536,10 +538,10 @@ def _setup_java_sdk_integration(dot_env_file, tmp_dir): os.environ["ENV_FILE_PATH"] = str(dot_env_file) -def _run_go_sdk_pack(output_path, *, capture_output=False, native=False): +def _run_go_sdk_pack(package, output_path, *, capture_output=False, native=False): """Run ``go tool airflow-go-pack`` natively or inside the pinned Go toolchain container. - ``go tool airflow-go-pack`` builds the bundle package, reads its + ``go tool airflow-go-pack`` builds the bundle *package*, reads its --airflow-metadata, and appends the source + airflow-metadata.yaml + the AFBNDL01 trailer, writing a single self-contained executable bundle. CGO_ENABLED=0 yields a fully static binary that runs on the stock worker. @@ -569,7 +571,7 @@ def _run_go_sdk_pack(output_path, *, capture_output=False, native=False): "airflow-go-pack", "--output", str(output_path), - GO_SDK_EXAMPLE_BUNDLE_PKG, + package, ] else: cwd = None @@ -600,20 +602,20 @@ def _run_go_sdk_pack(output_path, *, capture_output=False, native=False): "airflow-go-pack", "--output", f"{container_bin_dir}/{output_path.name}", - GO_SDK_EXAMPLE_BUNDLE_PKG, + package, ] return subprocess.run(argv, cwd=cwd, env=env, check=True, capture_output=capture_output, text=True) -def _pack_go_sdk_example_bundle(*, native=False): - """Build the Go SDK example bundle, capturing output so a failure prints the build log.""" - output_path = GO_SDK_BIN_PATH / GO_SDK_BUNDLE_NAME +def _pack_go_sdk_bundle(package, name, *, native=False): + """Build a Go SDK bundle, capturing output so a failure prints the build log.""" + output_path = GO_SDK_BIN_PATH / name mode_label = "host toolchain" if native else GO_BUILDER_IMAGE - console.print(f"[yellow]Building Go SDK example bundle ({mode_label})...") + console.print(f"[yellow]Building Go SDK bundle {package} ({mode_label})...") try: - completed = _run_go_sdk_pack(output_path, capture_output=True, native=native) + completed = _run_go_sdk_pack(package, output_path, capture_output=True, native=native) except subprocess.CalledProcessError as e: - console.print("[red]Go SDK example bundle build failed:") + console.print(f"[red]Go SDK bundle {package} build failed:") console.print(e.stdout, e.stderr, sep="\n", markup=False, soft_wrap=True) raise console.print(completed.stdout, completed.stderr, sep="\n", markup=False, soft_wrap=True) @@ -627,11 +629,16 @@ def _setup_go_sdk_integration(dot_env_file, tmp_dir): as the ``go-task-handlers`` Dag bundle, copies the Python stub Dag, and writes the coordinator configuration. + The native-Dag example is packed too, into the Dags folder, where the Dag + processor runs it to parse its Dags. Its tasks run on the ``golang`` queue + through the same coordinator. + The packed bundle is a statically linked native executable (built with ``CGO_ENABLED=0``), so the stock Airflow worker image can exec it directly without a Go toolchain or any extra runtime installed -- see ``go.yml``. """ - _pack_go_sdk_example_bundle(native=LANG_SDK_NATIVE_TOOLCHAIN) + _pack_go_sdk_bundle(GO_SDK_EXAMPLE_BUNDLE_PKG, GO_SDK_BUNDLE_NAME, native=LANG_SDK_NATIVE_TOOLCHAIN) + _pack_go_sdk_bundle(GO_SDK_NATIVE_BUNDLE_PKG, GO_SDK_NATIVE_BUNDLE_NAME, native=LANG_SDK_NATIVE_TOOLCHAIN) # Copy the compose override into the temp directory. copyfile(GO_COMPOSE_PATH, tmp_dir / "go.yml") @@ -644,13 +651,25 @@ def _setup_go_sdk_integration(dot_env_file, tmp_dir): packed_bundle = go_bundles_dir / GO_SDK_BUNDLE_NAME copyfile(GO_SDK_BIN_PATH / GO_SDK_BUNDLE_NAME, packed_bundle) os.chmod(packed_bundle, 0o755) + # This bundle only registers task handlers, so the Dag processor must not parse it if the + # directory is ever mounted there. Tasks do not read .airflowignore. + (go_bundles_dir / ".airflowignore").write_text("*\n") + + # The native-Dag bundle goes into the Dags folder, where the Dag processor runs it to parse. + native_dags_dir = tmp_dir / "dags" / "go" + native_dags_dir.mkdir() + packed_native_bundle = native_dags_dir / GO_SDK_NATIVE_BUNDLE_NAME + copyfile(GO_SDK_BIN_PATH / GO_SDK_NATIVE_BUNDLE_NAME, packed_native_bundle) + os.chmod(packed_native_bundle, 0o755) # Copy the Go SDK example stub Dag so Airflow can discover and serialize it. copyfile(GO_SDK_DAGS_PATH / "go_examples.py", tmp_dir / "dags" / "go_examples.py") # Coordinator registry: maps the logical name "go-sdk" to ExecutableCoordinator, # which scans the go-task-handlers Dag bundle for the packed bundle by dag_id. - # Queue mapping: routes tasks on the "golang" queue to "go-sdk". + # Queue mapping: routes tasks on the "golang" queue to "go-sdk". It is the only + # ExecutableCoordinator, so it also parses the native Dags in the Dags folder, and runs + # their tasks from that bundle, not from go-task-handlers. dag_bundle_config = _build_dag_bundle_config({"go-task-handlers": "/opt/airflow/go-bundles"}) coordinator_config = json.dumps( { diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py b/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py index d3ca78fe170b9..87e7fa5bdf6f0 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py @@ -92,6 +92,9 @@ GO_SDK_EXAMPLE_BUNDLE_PKG = "./example/bundle" # Name of the packed bundle binary (matches the example bundle's package dir name). GO_SDK_BUNDLE_NAME = "example_dags" +# The native-Dag example: Dags declared in Go and parsed from the binary by the Dag processor. +GO_SDK_NATIVE_BUNDLE_PKG = "./example/native" +GO_SDK_NATIVE_BUNDLE_NAME = "go_native_dags" # Where airflow-go-pack writes the packed bundle inside the repo (go-sdk/bin is gitignored). GO_SDK_BIN_PATH = GO_SDK_ROOT_PATH / "bin" GO_COMPOSE_PATH = AIRFLOW_ROOT_PATH / "airflow-e2e-tests" / "docker" / "go.yml" diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py b/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py index 87b9b816637ca..b48faed98139b 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/e2e_test_utils/clients.py @@ -19,6 +19,7 @@ import time from datetime import UTC, datetime from functools import cached_property +from http import HTTPStatus import boto3 import requests @@ -143,11 +144,7 @@ def trigger_dag(self, dag_id: str, json=None): def wait_for_dag_run(self, dag_id: str, run_id: str, timeout=300, check_interval=5): start_time = time.time() while time.time() - start_time < timeout: - response = self._make_request( - method="GET", - endpoint=f"dags/{dag_id}/dagRuns/{run_id}", - ) - state = response.get("state") + state = self.get_dag_run(dag_id=dag_id, run_id=run_id).get("state") if state in {"success", "failed"}: return state time.sleep(check_interval) @@ -170,6 +167,29 @@ def get_task_state_store(self, dag_id: str, run_id: str, task_id: str, key: str) endpoint=f"dags/{dag_id}/dagRuns/{run_id}/taskInstances/{task_id}/state-store/{key}", ) + def set_variable(self, key: str, value: str, description: str | None = None): + """Create or replace an Airflow Variable via API.""" + body = {"key": key, "value": value, "description": description} + try: + return self._make_request(method="POST", endpoint="variables", json=body) + except requests.HTTPError as exc: + # 409 == it already exists, from an earlier run of the same suite. + if exc.response is None or exc.response.status_code != HTTPStatus.CONFLICT: + raise + return self._make_request(method="PATCH", endpoint=f"variables/{key}", json=body) + + def get_tasks(self, dag_id: str): + """List a Dag's tasks, with the edges each one carries.""" + return self._make_request(method="GET", endpoint=f"dags/{dag_id}/tasks") + + def get_dag_source(self, dag_id: str): + """Get the source code stored for a Dag's latest version.""" + return self._make_request(method="GET", endpoint=f"dagSources/{dag_id}") + + def get_dag_run(self, dag_id: str, run_id: str): + """Get a Dag run, with its state, run type and conf.""" + return self._make_request(method="GET", endpoint=f"dags/{dag_id}/dagRuns/{run_id}") + def trigger_dag_and_wait(self, dag_id: str, json=None): """Trigger a DAG and wait for it to complete.""" self.un_pause_dag(dag_id) diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_native_dag.py b/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_native_dag.py new file mode 100644 index 0000000000000..43fcff7f4c8cf --- /dev/null +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_native_dag.py @@ -0,0 +1,220 @@ +# 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. +"""E2E tests for the Go SDK native Dags of ``go-sdk/example/native``. + +Run with:: + + E2E_TEST_MODE=go_sdk uv run --project airflow-e2e-tests pytest \\ + tests/airflow_e2e_tests/go_sdk_tests/test_go_sdk_native_dag.py -xvs + +No Python file declares these Dags. ``conftest._setup_go_sdk_integration`` packs the example with +``airflow-go-pack`` into the Dags folder, where the Dag processor claims it through the +``ExecutableCoordinator`` and runs it to parse. The worker runs each task from the same binary, on +the ``golang`` queue the Dag sets. + +``go_native_pipeline`` has a nested task group, a fan-in, order-only edges, an ``If``, a ``Switch`` +and a trigger-rule join. ``go_native_report`` is declared in a second file, so its Code view shows +that file. ``go_native_trigger`` triggers ``go_native_report``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import UTC, datetime + +import pytest + +from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient + +# The Dag processor execs the binary to parse it, so the Dags appear some time after the stack is up. +_PARSE_TIMEOUT = 600 +_GO_TASK_TIMEOUT = 600 + +_PIPELINE_DAG_ID = "go_native_pipeline" +_REPORT_DAG_ID = "go_native_report" +_TRIGGER_DAG_ID = "go_native_trigger" + +# Read by the example's tasks; see go-sdk/example/native/main.go. +_GREETING_VARIABLE = "go_native_greeting" +_CADENCE_VARIABLE = "go_native_cadence" +_GREETING = "hello from the e2e test" +# The test_http Connection that conftest._setup_go_sdk_integration defines. +_CONNECTION_HOST = "example.com" + + +@dataclass +class _CompletedRun: + client: AirflowClient + run_id: str + state: str + ti_states: dict[str, str] + + def xcom(self, task_id: str, key: str = "return_value"): + return self.client.get_xcom_value( + dag_id=_PIPELINE_DAG_ID, task_id=task_id, run_id=self.run_id, key=key + ).get("value") + + +@pytest.fixture(scope="module") +def parsed_dags() -> AirflowClient: + """A client that has waited for the Dag processor to parse and register all three Dags.""" + client = AirflowClient() + for dag_id in (_PIPELINE_DAG_ID, _REPORT_DAG_ID, _TRIGGER_DAG_ID): + client.wait_for_dag(dag_id, timeout=_PARSE_TIMEOUT) + return client + + +@pytest.fixture(scope="module") +def completed_run(parsed_dags: AirflowClient) -> _CompletedRun: + """Trigger the pipeline once, with the Variables its tasks read.""" + client = parsed_dags + client.set_variable(_GREETING_VARIABLE, _GREETING) + client.set_variable(_CADENCE_VARIABLE, "weekly") + + resp = client.trigger_dag(_PIPELINE_DAG_ID, json={"logical_date": datetime.now(UTC).isoformat()}) + run_id = resp["dag_run_id"] + state = client.wait_for_dag_run(dag_id=_PIPELINE_DAG_ID, run_id=run_id, timeout=_GO_TASK_TIMEOUT) + ti_resp = client.get_task_instances(dag_id=_PIPELINE_DAG_ID, run_id=run_id) + return _CompletedRun( + client=client, + run_id=run_id, + state=state, + ti_states={ti["task_id"]: ti.get("state") for ti in ti_resp.get("task_instances", [])}, + ) + + +def _edges(client: AirflowClient, dag_id: str) -> tuple[dict[str, set[str]], dict[str, set[str]]]: + """Return the downstream and the upstream task ids of every task of *dag_id*.""" + tasks = client.get_tasks(dag_id).get("tasks", []) + downstream = {task["task_id"]: set(task["downstream_task_ids"]) for task in tasks} + # The API reports downstream edges only, so upstream is read by inverting them. + upstream: dict[str, set[str]] = {task_id: set() for task_id in downstream} + for task_id, down in downstream.items(): + for child in down: + upstream.setdefault(child, set()).add(task_id) + return downstream, upstream + + +def test_the_dags_the_binary_parsed_are_registered(parsed_dags: AirflowClient): + """The Dags exist without any Python file declaring them.""" + for dag_id in (_PIPELINE_DAG_ID, _REPORT_DAG_ID, _TRIGGER_DAG_ID): + dag = parsed_dags.get_dag(dag_id) + assert dag["dag_id"] == dag_id + assert {tag["name"] for tag in dag.get("tags") or []} >= {"go", "native"} + + +def test_the_graph_carries_every_construct(parsed_dags: AirflowClient): + """Group prefixes, the group edge, the fan-in, the branches and the join survive parsing.""" + downstream, upstream = _edges(parsed_dags, _PIPELINE_DAG_ID) + + # The nested groups prefixed their tasks, and the task-to-group edge expanded onto the + # tasks that the group starts with. + assert {"extract.regions.north", "extract.regions.south", "extract.summarize"} <= set(downstream) + assert downstream["seed"] >= {"extract.regions.north", "extract.regions.south"} + # The fan-in. + assert upstream["extract.summarize"] >= {"extract.regions.north", "extract.regions.south"} + # The condition and the switch reach their candidates. + assert downstream["anyRows"] >= {"loadRows", "reportEmpty"} + assert downstream["pick"] >= {"publishDaily", "publishWeekly"} + # Cleanup sits behind every branch outcome, which is what its trigger rule is for. + assert upstream["cleanup"] >= {"loadRows", "reportEmpty", "publishDaily", "publishWeekly"} + + +def test_dag_run_succeeded(completed_run: _CompletedRun): + assert completed_run.state == "success", ( + f"expected the run to succeed; got {completed_run.state!r}. task states: {completed_run.ti_states}" + ) + + +def test_the_taken_branches_ran_and_the_others_skipped(completed_run: _CompletedRun): + """A branch is a run-time skip, so the states are what prove it worked.""" + states = completed_run.ti_states + skipped = {"reportEmpty", "publishDaily"} + ran = { + "seed", + "extract.regions.north", + "extract.regions.south", + "extract.summarize", + "loadRows", + "anyRows", + "pick", + "publishWeekly", + "cleanup", + } + assert set(states) == ran | skipped, states + for task_id in ran: + assert states[task_id] == "success", f"{task_id!r} did not succeed. task states: {states}" + for task_id in skipped: + assert states[task_id] == "skipped", f"{task_id!r} was not skipped. task states: {states}" + + +def test_xcoms_flow_between_go_tasks(completed_run: _CompletedRun): + """Each task got its literals and upstream results, and its own return came back.""" + assert completed_run.xcom("extract.regions.north") == {"region": "north", "rows": 3} + assert completed_run.xcom("extract.regions.south") == {"region": "south", "rows": 2} + assert completed_run.xcom("extract.summarize") == {"total": 5, "regions": 2} + assert completed_run.xcom("loadRows") == "s3://bucket/out" + + +def test_the_decisions_are_recorded(completed_run: _CompletedRun): + """A condition returns its boolean, and a switch the task id it chose.""" + assert completed_run.xcom("anyRows") is True + assert completed_run.xcom("pick") == "publishWeekly" + + +def test_the_variable_and_connection_reached_the_go_task(completed_run: _CompletedRun): + """``seed`` read the REST-set Variable and the Connection that conftest defines.""" + assert completed_run.xcom("seed") == {"greeting": _GREETING, "host": _CONNECTION_HOST} + + +def test_a_task_pulls_the_xcom_another_pushed_under_its_own_key(completed_run: _CompletedRun): + note = f"seeded {_GREETING} from {_CONNECTION_HOST}" + assert completed_run.xcom("seed", key="seed_note") == note + assert completed_run.xcom("cleanup") == note + + +def test_the_dag_source_is_the_file_that_declares_the_dag(parsed_dags: AirflowClient): + """The Code view shows the source file of each Dag, not always the entrypoint.""" + report = parsed_dags.get_dag_source(_REPORT_DAG_ID) + assert 'airflow.Dag("go_native_report"' in report["content"] + assert "func main(" not in report["content"] + assert report.get("language") == "go" + + pipeline = parsed_dags.get_dag_source(_PIPELINE_DAG_ID) + assert "func main(" in pipeline["content"] + + +def test_the_trigger_started_the_report_run(parsed_dags: AirflowClient): + client = parsed_dags + # Dags are paused at creation here, so the run the trigger starts would stay queued. + client.un_pause_dag(_REPORT_DAG_ID) + resp = client.trigger_dag(_TRIGGER_DAG_ID, json={"logical_date": datetime.now(UTC).isoformat()}) + trigger_run_id = resp["dag_run_id"] + state = client.wait_for_dag_run(dag_id=_TRIGGER_DAG_ID, run_id=trigger_run_id, timeout=_GO_TASK_TIMEOUT) + assert state == "success" + + report_run_id = client.get_xcom_value( + dag_id=_TRIGGER_DAG_ID, task_id="trigger_report", run_id=trigger_run_id, key="trigger_run_id" + )["value"] + report_state = client.wait_for_dag_run( + dag_id=_REPORT_DAG_ID, run_id=report_run_id, timeout=_GO_TASK_TIMEOUT + ) + assert report_state == "success" + + report_run = client.get_dag_run(dag_id=_REPORT_DAG_ID, run_id=report_run_id) + assert report_run["run_type"] == "operator_triggered" + assert report_run["conf"] == {"triggered_by": _TRIGGER_DAG_ID} diff --git a/go-sdk/example/native/main.go b/go-sdk/example/native/main.go new file mode 100644 index 0000000000000..fda85d96eb601 --- /dev/null +++ b/go-sdk/example/native/main.go @@ -0,0 +1,177 @@ +// 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 is an example of Dags written entirely in Go with airflow.Dag, with no Python stub +// Dag. The Dag processor runs the packed binary to parse them, and the worker runs it for each +// task. +// +// Pack it with `go tool airflow-go-pack ./example/native` and put the result in a Dag bundle. +package main + +import ( + "fmt" + "log" + + "github.com/apache/airflow/go-sdk/airflow" +) + +type regionRows struct { + Region string `json:"region"` + Rows int `json:"rows"` +} + +type summary struct { + Total int `json:"total"` + Regions int `json:"regions"` +} + +// pipeline holds the tasks that the switch chooses from, because a decider returns a *TaskRef. +type pipeline struct { + daily, weekly *airflow.TaskRef +} + +func main() { + bundle := airflow.Bundle() + bundle.Register( + newPipelineDag(), + newTriggerDag(), + newReportDag(), + ) + + if err := bundle.Serve(); err != nil { + log.Fatal(err) + } +} + +func newPipelineDag() *airflow.DagRef { + dag := airflow.Dag("go_native_pipeline", airflow.DagSpec{ + Queue: "golang", + Tags: []string{"go", "native"}, + }) + + seeded := dag.Task(seed) + + extract := dag.TaskGroup("extract") + regions := extract.TaskGroup("regions") + north := regions.Task( + countRows, + airflow.TaskSpec{TaskID: "north"}, + airflow.Inputs(airflow.Literal("north"), airflow.Literal(3)), + ) + south := regions.Task( + countRows, + airflow.TaskSpec{TaskID: "south"}, + airflow.Inputs(airflow.Literal("south"), airflow.Literal(2)), + ) + summed := extract.Task(summarize, airflow.Inputs(north, south)) + seeded.Before(airflow.Label(extract, "fan out")) + + loaded := dag.Task(loadRows, airflow.Inputs(summed, airflow.Literal("s3://bucket/out"))) + empty := dag.Task(reportEmpty) + dag.If(anyRows, airflow.Inputs(summed)).Then(loaded).Else(empty) + + p := &pipeline{ + daily: dag.Task(publishDaily), + weekly: dag.Task(publishWeekly), + } + dag.Switch(p.pick, airflow.Inputs(summed)).Case(p.daily).Case(p.weekly) + + // Each If and Switch skips one side, so cleanup has to run after a skipped upstream task. + dag.Task(cleanup, airflow.TaskSpec{TriggerRule: airflow.TriggerRuleNoneFailedMinOneSuccess}). + After(loaded, empty, p.daily, p.weekly) + + return dag +} + +func newTriggerDag() *airflow.DagRef { + dag := airflow.Dag( + "go_native_trigger", + airflow.DagSpec{Queue: "golang", Tags: []string{"go", "native"}}, + ) + dag.Task( + airflow.TriggerDagRun(airflow.TriggerDagRunSpec{ + DagID: "go_native_report", + Conf: map[string]any{"triggered_by": "go_native_trigger"}, + }), + airflow.TaskSpec{TaskID: "trigger_report"}, + ) + return dag +} + +// seed reads a Variable and a Connection, and leaves a note for cleanup under its own XCom key. +func seed(actx airflow.Context) (map[string]any, error) { + client := actx.Client() + greeting, err := client.GetVariable(actx, "go_native_greeting") + if err != nil { + return nil, err + } + conn, err := client.GetConnection(actx, "test_http") + if err != nil { + return nil, err + } + note := fmt.Sprintf("seeded %s from %s", greeting, conn.Host) + if err := client.PushXCom(actx, actx.TaskInstance(), "seed_note", note); err != nil { + return nil, err + } + return map[string]any{"greeting": greeting, "host": conn.Host}, nil +} + +func countRows(_ airflow.Context, region string, rows int) (regionRows, error) { + return regionRows{Region: region, Rows: rows}, nil +} + +func summarize(_ airflow.Context, north, south regionRows) (summary, error) { + return summary{Total: north.Rows + south.Rows, Regions: 2}, nil +} + +func anyRows(_ airflow.Context, summed summary) (bool, error) { + return summed.Total > 0, nil +} + +func loadRows(_ airflow.Context, _ summary, target string) (string, error) { + return target, nil +} + +func reportEmpty(airflow.Context) (string, error) { + return "no rows", nil +} + +// pick chooses the task that the Variable names. +func (p *pipeline) pick(actx airflow.Context, _ summary) (*airflow.TaskRef, error) { + cadence, err := actx.Client().GetVariable(actx, "go_native_cadence") + if err != nil { + return nil, err + } + if cadence == "weekly" { + return p.weekly, nil + } + return p.daily, nil +} + +func publishDaily(airflow.Context) (string, error) { + return "daily", nil +} + +func publishWeekly(airflow.Context) (string, error) { + return "weekly", nil +} + +// cleanup returns the note that seed pushed. +func cleanup(actx airflow.Context) (any, error) { + ti := actx.TaskInstance() + return actx.Client().GetXCom(actx, ti.DagID, ti.RunID, "seed", nil, "seed_note", nil) +} diff --git a/go-sdk/example/native/report.go b/go-sdk/example/native/report.go new file mode 100644 index 0000000000000..1a22c49c387af --- /dev/null +++ b/go-sdk/example/native/report.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 main + +import "github.com/apache/airflow/go-sdk/airflow" + +// newReportDag lives in its own file, so the Code view of go_native_report shows this file and +// not main.go. +func newReportDag() *airflow.DagRef { + dag := airflow.Dag("go_native_report", airflow.DagSpec{ + Queue: "golang", + Tags: []string{"go", "native"}, + }) + dag.Task(render) + return dag +} + +func render(airflow.Context) (string, error) { + return "report rendered", nil +}