diff --git a/.github/workflows/airflow-e2e-tests.yml b/.github/workflows/airflow-e2e-tests.yml index e7f1a8c790ce1..dd83f67bb157f 100644 --- a/.github/workflows/airflow-e2e-tests.yml +++ b/.github/workflows/airflow-e2e-tests.yml @@ -49,7 +49,7 @@ on: # yamllint disable-line rule:truthy type: string required: true e2e_test_mode: - description: "Test mode - basic, remote_log, remote_log_elasticsearch, remote_log_opensearch, xcom_object_storage, or event_driven" # yamllint disable-line rule:line-length + description: "Test mode - basic, remote_log, remote_log_elasticsearch, remote_log_opensearch, xcom_object_storage, event_driven, or java_sdk" # yamllint disable-line rule:line-length type: string default: "basic" @@ -80,7 +80,7 @@ on: # yamllint disable-line rule:truthy type: string default: "" e2e_test_mode: - description: "Test mode - basic, remote_log, remote_log_elasticsearch, remote_log_opensearch, xcom_object_storage, or event_driven" # yamllint disable-line rule:line-length + description: "Test mode - basic, remote_log, remote_log_elasticsearch, remote_log_opensearch, xcom_object_storage, event_driven, or java_sdk" # yamllint disable-line rule:line-length type: string default: "basic" diff --git a/.github/workflows/ci-amd.yml b/.github/workflows/ci-amd.yml index 47b6429308719..7c4f563007222 100644 --- a/.github/workflows/ci-amd.yml +++ b/.github/workflows/ci-amd.yml @@ -126,6 +126,7 @@ jobs: run-api-tests: ${{ steps.selective-checks.outputs.run-api-tests }} run-coverage: ${{ steps.source-run-info.outputs.run-coverage }} run-go-sdk-tests: ${{ steps.selective-checks.outputs.run-go-sdk-tests }} + run-java-sdk-tests: ${{ steps.selective-checks.outputs.run-java-sdk-tests }} run-helm-tests: ${{ steps.selective-checks.outputs.run-helm-tests }} run-kubernetes-tests: ${{ steps.selective-checks.outputs.run-kubernetes-tests }} run-mypy-providers: ${{ steps.selective-checks.outputs.run-mypy-providers }} @@ -954,6 +955,37 @@ jobs: working-directory: ./go-sdk run: gotestsum --format github-actions ./... + tests-java-sdk: + name: "Java SDK tests" + needs: [build-info] + runs-on: ${{ fromJSON(needs.build-info.outputs.runner-type) }} + timeout-minutes: 15 + permissions: + contents: read + packages: read + if: needs.build-info.outputs.run-java-sdk-tests == 'true' + env: + GITHUB_REPOSITORY: ${{ github.repository }} + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GITHUB_USERNAME: ${{ github.actor }} + VERBOSE: "true" + steps: + - name: "Checkout ${{ github.ref }} ( ${{ github.sha }} )" + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + # keep this in sync with the jvmTarget in java-sdk/build.gradle.kts + - name: Setup Java + uses: actions/setup-java@be666c2fcd27ec809703dec50e508c2fdc7f6654 # v5.2.0 + with: + distribution: 'temurin' + java-version: '11' + - name: "Cleanup dist files" + run: rm -fv ./dist/* + - name: Run Java SDK tests + working-directory: ./java-sdk + run: ./gradlew test + tests-airflow-ctl: name: "Airflow CTL tests" uses: ./.github/workflows/airflow-distributions-tests.yml @@ -1007,6 +1039,7 @@ jobs: - tests-task-sdk - tests-airflow-ctl - tests-go-sdk + - tests-java-sdk - tests-with-lowest-direct-resolution-core - tests-with-lowest-direct-resolution-providers uses: ./.github/workflows/finalize-tests.yml diff --git a/.github/workflows/ci-arm.yml b/.github/workflows/ci-arm.yml index 62ac6a2d99c1e..b750c4a949c14 100644 --- a/.github/workflows/ci-arm.yml +++ b/.github/workflows/ci-arm.yml @@ -116,6 +116,7 @@ jobs: run-api-tests: ${{ steps.selective-checks.outputs.run-api-tests }} run-coverage: ${{ steps.source-run-info.outputs.run-coverage }} run-go-sdk-tests: ${{ steps.selective-checks.outputs.run-go-sdk-tests }} + run-java-sdk-tests: ${{ steps.selective-checks.outputs.run-java-sdk-tests }} run-helm-tests: ${{ steps.selective-checks.outputs.run-helm-tests }} run-kubernetes-tests: ${{ steps.selective-checks.outputs.run-kubernetes-tests }} run-mypy-providers: ${{ steps.selective-checks.outputs.run-mypy-providers }} @@ -944,6 +945,37 @@ jobs: working-directory: ./go-sdk run: gotestsum --format github-actions ./... + tests-java-sdk: + name: "Java SDK tests" + needs: [build-info] + runs-on: ${{ fromJSON(needs.build-info.outputs.runner-type) }} + timeout-minutes: 15 + permissions: + contents: read + packages: read + if: needs.build-info.outputs.run-java-sdk-tests == 'true' + env: + GITHUB_REPOSITORY: ${{ github.repository }} + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GITHUB_USERNAME: ${{ github.actor }} + VERBOSE: "true" + steps: + - name: "Checkout ${{ github.ref }} ( ${{ github.sha }} )" + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + # keep this in sync with the jvmTarget in java-sdk/build.gradle.kts + - name: Setup Java + uses: actions/setup-java@be666c2fcd27ec809703dec50e508c2fdc7f6654 # v5.2.0 + with: + distribution: 'temurin' + java-version: '11' + - name: "Cleanup dist files" + run: rm -fv ./dist/* + - name: Run Java SDK tests + working-directory: ./java-sdk + run: ./gradlew test + tests-airflow-ctl: name: "Airflow CTL tests" uses: ./.github/workflows/airflow-distributions-tests.yml @@ -997,6 +1029,7 @@ jobs: - tests-task-sdk - tests-airflow-ctl - tests-go-sdk + - tests-java-sdk - tests-with-lowest-direct-resolution-core - tests-with-lowest-direct-resolution-providers uses: ./.github/workflows/finalize-tests.yml diff --git a/.github/workflows/codeql-analysis.yml b/.github/workflows/codeql-analysis.yml index 1aabf1df58bac..9417c3b3773b4 100644 --- a/.github/workflows/codeql-analysis.yml +++ b/.github/workflows/codeql-analysis.yml @@ -39,7 +39,8 @@ jobs: strategy: fail-fast: false matrix: - language: ['python', 'javascript', 'actions', 'go'] + language: ['python', 'javascript', 'actions', 'go', 'java'] + permissions: actions: read contents: read @@ -51,14 +52,28 @@ jobs: with: persist-credentials: false + # keep this in sync with the jvmTarget in java-sdk/build.gradle.kts + - name: Setup Java + if: matrix.language == 'java' + uses: actions/setup-java@be666c2fcd27ec809703dec50e508c2fdc7f6654 # v5.2.0 + with: + distribution: 'temurin' + java-version: '11' + - name: Initialize CodeQL uses: github/codeql-action/init@9e0d7b8d25671d64c341c19c0152d693099fb5ba # v4.35.5 with: languages: ${{ matrix.language }} - name: Autobuild + if: matrix.language != 'java' uses: github/codeql-action/autobuild@9e0d7b8d25671d64c341c19c0152d693099fb5ba # v4.35.5 + - name: Build Java SDK + if: matrix.language == 'java' + working-directory: java-sdk + run: ./gradlew classes testClasses + - name: Perform CodeQL Analysis uses: github/codeql-action/analyze@9e0d7b8d25671d64c341c19c0152d693099fb5ba # v4.35.5 with: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c5d02ea0a8421..ad15f0b15d62a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -62,6 +62,9 @@ repos: exclude: | (?x) ^\.github/| + ^java-sdk/gradlew$| + ^java-sdk/gradlew\.bat$| + ^java-sdk/gradle/| ^scripts/ci/license-templates/ args: - --comment-style @@ -680,6 +683,9 @@ repos: ^airflow-core/src/airflow/utils/db\.py$| ^airflow-core/src/airflow/utils/trigger_rule\.py$| ^airflow-core/tests/| + ^java-sdk/gradlew$| + ^java-sdk/gradlew\.bat$| + ^java-sdk/gradle| ^task-sdk/tests/| ^.*changelog\.(rst|txt)$| ^.*CHANGELOG\.(rst|txt)$| @@ -1201,3 +1207,10 @@ repos: language: python files: .*test.*\.py$ pass_filenames: true + - id: ktlint + name: Run ktlint format + description: "Use ktlint (via Gradle) to format Kotlin and Java files" + entry: ./java-sdk/gradlew -p ./java-sdk ktlintFormat + language: system + pass_filenames: false + files: ^java-sdk/.*$ diff --git a/.rat-excludes b/.rat-excludes index afee79c2ed0f3..01f7a3acb8b43 100644 --- a/.rat-excludes +++ b/.rat-excludes @@ -327,6 +327,11 @@ www-hash.txt # Vendored-in code /src/airflow/providers/google/_vendor/* +# Java SDK build outputs +/java-sdk/build/* +/java-sdk/sdk/build/* +/java-sdk/example/build/* + # Git ignore file .gitignore diff --git a/airflow-e2e-tests/docker/Dockerfile.java b/airflow-e2e-tests/docker/Dockerfile.java new file mode 100644 index 0000000000000..7fc3363825ecb --- /dev/null +++ b/airflow-e2e-tests/docker/Dockerfile.java @@ -0,0 +1,27 @@ +# 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. + +# Extends the standard Airflow image with a headless JRE so JavaCoordinator +# can spawn JVM subprocesses for @task.stub tasks. +ARG DOCKER_IMAGE +FROM ${DOCKER_IMAGE} + +USER root +RUN apt-get update \ + && apt-get install -y --no-install-recommends default-jre-headless \ + && rm -rf /var/lib/apt/lists/* +USER airflow diff --git a/airflow-e2e-tests/docker/java.yml b/airflow-e2e-tests/docker/java.yml new file mode 100644 index 0000000000000..3a01c66dd181b --- /dev/null +++ b/airflow-e2e-tests/docker/java.yml @@ -0,0 +1,31 @@ +# 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. + +# Docker Compose override for java_sdk E2E test mode. +# +# Replaces the stock airflow-worker image with one that has a JRE installed +# (built by conftest._setup_java_sdk_integration via Dockerfile.java), mounts +# the pre-built example bundle JARs under /opt/airflow/jars, and configures +# the worker to consume the "java" Celery queue where @task.stub tasks are +# routed. +--- +services: + airflow-worker: + image: airflow-java-worker + volumes: + - ./jars:/opt/airflow/jars:ro + command: celery worker -q java,default diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py index 95952e4ae80f0..aa4e51d7f540c 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py @@ -18,6 +18,7 @@ import json import os +import subprocess from datetime import datetime from pathlib import Path from shutil import copyfile, copytree @@ -27,6 +28,7 @@ from testcontainers.compose import DockerCompose from airflow_e2e_tests.constants import ( + AIRFLOW_ROOT_PATH, AIRFLOW_SERVICES_FOR_PROVIDER_MOUNT, AWS_INIT_PATH, DOCKER_COMPOSE_HOST_PORT, @@ -35,6 +37,10 @@ E2E_DAGS_FOLDER, E2E_TEST_MODE, ELASTICSEARCH_PATH, + JAVA_COMPOSE_PATH, + JAVA_DOCKERFILE_PATH, + JAVA_SDK_DAGS_PATH, + JAVA_SDK_EXAMPLE_LIBS_PATH, KAFKA_DIR_PATH, LOCALSTACK_PATH, LOGS_FOLDER, @@ -234,6 +240,99 @@ def _setup_xcom_object_storage_integration(dot_env_file, tmp_dir): os.environ["ENV_FILE_PATH"] = str(dot_env_file) +def _setup_java_sdk_integration(dot_env_file, tmp_dir): + """Set up the java_sdk E2E test mode. + + Builds the Java example bundle via the Gradle wrapper, then builds a + Java-capable Airflow worker image, copies the JARs into the temp directory, + and writes the coordinator configuration. + """ + # Build the example bundle inside an ephemeral JDK container so the host + # does not need Java installed. + # + # --user keeps build outputs owned by the current user (not root). + # GRADLE_USER_HOME persists the Gradle distribution and dependency cache in + # java-sdk/.gradle/ (already gitignored) so the first run downloads once + # and subsequent runs skip straight to compilation. + # --no-daemon avoids a background JVM that would outlive the container. + console.print("[yellow]Building Java SDK example bundle (eclipse-temurin:17-jdk)...") + subprocess.run( + [ + "docker", + "run", + "--rm", + "--user", + f"{os.getuid()}:{os.getgid()}", + "-e", + "GRADLE_USER_HOME=/repo/java-sdk/.gradle", + # Mount java-sdk/ at /java-sdk (the Gradle root project). + "-v", + f"{AIRFLOW_ROOT_PATH}:/repo", + "-w", + "/repo/java-sdk", + "eclipse-temurin:17-jdk", + "./gradlew", + ":example:installDist", + "--no-daemon", + ], + check=True, + ) + + # Copy compose override and Dockerfile into the temp directory. + copyfile(JAVA_COMPOSE_PATH, tmp_dir / "java.yml") + copyfile(JAVA_DOCKERFILE_PATH, tmp_dir / "Dockerfile.java") + + # Copy all JARs from installDist output so the compose bind-mount ./jars + # gives the worker everything JavaCoordinator needs to build a classpath. + copytree(JAVA_SDK_EXAMPLE_LIBS_PATH, tmp_dir / "jars") + + # Copy the Java SDK example Dag file so Airflow can discover it. + copyfile(JAVA_SDK_DAGS_PATH / "java_examples.py", tmp_dir / "dags" / "java_examples.py") + + # Build a local Docker image that extends DOCKER_IMAGE with a JRE. + # We do this explicitly so testcontainers' DockerCompose.start() does not + # need to handle the build itself (which avoids --no-build vs --build flag + # uncertainty across testcontainers versions). + console.print(f"[yellow]Building airflow-java-worker image on top of {DOCKER_IMAGE}...") + subprocess.run( + [ + "docker", + "build", + "--build-arg", + f"DOCKER_IMAGE={DOCKER_IMAGE}", + "-t", + "airflow-java-worker", + "-f", + str(tmp_dir / "Dockerfile.java"), + str(tmp_dir), + ], + check=True, + ) + + # Coordinator registry: maps the logical name "java-jdk" to JavaCoordinator. + # Queue mapping: routes tasks on the "java" Celery queue to "java-jdk". + coordinator_config = json.dumps( + { + "java-jdk": { + "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", + "kwargs": {"jars_root": ["/opt/airflow/jars"]}, + } + } + ) + queue_to_coordinator = json.dumps({"java": "java-jdk"}) + + dot_env_file.write_text( + f"AIRFLOW_UID={os.getuid()}\n" + # Single-quote the JSON values so Docker Compose reads them literally. + f"AIRFLOW__SDK__COORDINATORS='{coordinator_config}'\n" + f"AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{queue_to_coordinator}'\n" + # Connection and variable expected by the Java example bundle tasks. + "AIRFLOW_CONN_TEST_HTTP=http://test:test@example.com/\n" + "AIRFLOW_VAR_MY_VARIABLE=test_value\n" + ) + os.environ["ENV_FILE_PATH"] = str(dot_env_file) + + def spin_up_airflow_environment(tmp_path_factory: pytest.TempPathFactory): tmp_dir = tmp_path_factory.mktemp("breeze-airflow-e2e-tests") @@ -275,6 +374,9 @@ def spin_up_airflow_environment(tmp_path_factory: pytest.TempPathFactory): elif E2E_TEST_MODE == "event_driven": compose_file_names.extend(["kafka.yml", "providers-mount.yml"]) _setup_event_driven_integration(dot_env_file, tmp_dir) + elif E2E_TEST_MODE == "java_sdk": + compose_file_names.append("java.yml") + _setup_java_sdk_integration(dot_env_file, tmp_dir) # # Please Do not use this Fernet key in any deployments! Please generate your own key. diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py b/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py index 7e27dd6edc6fe..a81eeb0d6a089 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/constants.py @@ -50,6 +50,13 @@ KAFKA_DIR_PATH = AIRFLOW_ROOT_PATH / "airflow-e2e-tests" / "docker" / "kafka" +# Java SDK E2E test paths +JAVA_SDK_ROOT_PATH = AIRFLOW_ROOT_PATH / "java-sdk" +JAVA_SDK_DAGS_PATH = JAVA_SDK_ROOT_PATH / "dags" +JAVA_SDK_EXAMPLE_LIBS_PATH = JAVA_SDK_ROOT_PATH / "example" / "build" / "install" / "example" / "lib" +JAVA_COMPOSE_PATH = AIRFLOW_ROOT_PATH / "airflow-e2e-tests" / "docker" / "java.yml" +JAVA_DOCKERFILE_PATH = AIRFLOW_ROOT_PATH / "airflow-e2e-tests" / "docker" / "Dockerfile.java" + # Local provider sources are mounted into the airflow containers under this directory so # ``_PIP_ADDITIONAL_REQUIREMENTS`` can install the in-tree (latest, possibly unreleased) # provider rather than the published one from PyPI. diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/__init__.py b/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py b/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py new file mode 100644 index 0000000000000..0e99a399f2d44 --- /dev/null +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_dag.py @@ -0,0 +1,128 @@ +# 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 Java SDK via the annotation-based example bundle. + +Run with:: + + E2E_TEST_MODE=java_sdk uv run --project airflow-e2e-tests pytest \\ + tests/airflow_e2e_tests/java_sdk_tests/ -xvs + +What is verified +---------------- +The test triggers the ``java_annotation_example`` Dag, which has this task +graph:: + + python_task_1 >> extract >> transform >> python_task_2 + +* ``extract`` and ``transform`` are ``@task.stub(queue="java")`` stubs whose + implementations live in ``AnnotationExample.java``. Both run via + ``JavaCoordinator``, which spawns a JVM subprocess for each. +* ``extract`` reads an XCom from ``python_task_1``, fetches the ``test_http`` + connection, and returns a timestamp (long). +* ``transform`` reads the XCom from ``extract``, fetches the ``my_variable`` + Airflow variable, and returns a timestamp (long). + +The test asserts that both Java task instances reach state ``success``, which +confirms: + +1. ``JavaCoordinator`` correctly discovers and launches the JVM JAR. +2. The wire protocol (supervisor → JVM → supervisor) round-trips + ``StartupDetails`` and the task result (``SucceedTask``/``TaskState``). +3. XCom reads and API calls (getXCom, getConnection, getVariable) work + end-to-end through the Task Execution API. +""" + +from __future__ import annotations + +from datetime import datetime, timezone + +import pytest + +from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient + +# The Java extract task sleeps 6 s + coordinator startup; allow plenty of room. +_JAVA_TASK_TIMEOUT = 600 + + +class TestJavaSDKAnnotationExample: + """Verify the annotation-based Java SDK example executes correctly.""" + + airflow_client = AirflowClient() + + @pytest.mark.parametrize("dag_id", ["java_annotation_example"]) + def test_java_tasks_execute_successfully(self, dag_id: str): + """Both Java stubs in the annotation example must succeed.""" + resp = self.airflow_client.trigger_dag( + dag_id, + json={"logical_date": datetime.now(timezone.utc).isoformat()}, + ) + run_id = resp["dag_run_id"] + + # Wait for all tasks to finish (success or failed). + dag_state = self.airflow_client.wait_for_dag_run( + dag_id=dag_id, + run_id=run_id, + timeout=_JAVA_TASK_TIMEOUT, + ) + + # Collect task-instance states for a detailed failure message. + ti_resp = self.airflow_client.get_task_instances(dag_id=dag_id, run_id=run_id) + ti_map = {ti["task_id"]: ti for ti in ti_resp.get("task_instances", [])} + + extract_ti = ti_map.get("extract", {}) + transform_ti = ti_map.get("transform", {}) + + assert extract_ti.get("state") == "success", ( + f"Java 'extract' task did not succeed.\n" + f" task state : {extract_ti.get('state')!r}\n" + f" dag state : {dag_state!r}\n" + f" all tasks : { {k: v.get('state') for k, v in ti_map.items()} }" + ) + assert transform_ti.get("state") == "success", ( + f"Java 'transform' task did not succeed.\n" + f" task state : {transform_ti.get('state')!r}\n" + f" dag state : {dag_state!r}\n" + f" all tasks : { {k: v.get('state') for k, v in ti_map.items()} }" + ) + + def test_transform_xcom_is_numeric_timestamp(self): + """The value returned by the Java 'transform' task must be a positive integer.""" + resp = self.airflow_client.trigger_dag( + "java_annotation_example", + json={"logical_date": datetime.now(timezone.utc).isoformat()}, + ) + run_id = resp["dag_run_id"] + + self.airflow_client.wait_for_dag_run( + dag_id="java_annotation_example", + run_id=run_id, + timeout=_JAVA_TASK_TIMEOUT, + ) + + xcom = self.airflow_client.get_xcom_value( + dag_id="java_annotation_example", + task_id="transform", + run_id=run_id, + key="return_value", + ) + value = xcom.get("value") + assert isinstance(value, int), ( + f"Expected 'transform' XCom to be an integer (millisecond timestamp), got {value!r}" + ) + assert value > 0, ( + f"Expected 'transform' XCom to be a positive integer (millisecond timestamp), got {value!r}" + ) diff --git a/dev/breeze/doc/images/output_testing_airflow-e2e-tests.svg b/dev/breeze/doc/images/output_testing_airflow-e2e-tests.svg index b6a293c2ec9b1..290e96c8a3251 100644 --- a/dev/breeze/doc/images/output_testing_airflow-e2e-tests.svg +++ b/dev/breeze/doc/images/output_testing_airflow-e2e-tests.svg @@ -135,7 +135,7 @@ (TEXT) --e2e-test-mode               Specify the mode to use for E2E tests. [default: basic] (basic|remote_log|remote_log_elasticsearch|remote_log_opensearch|xcom_object_sto -rage|event_driven) +rage|event_driven|java_sdk) ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ╭─ Common options ─────────────────────────────────────────────────────────────────────────────────────────────────────╮ --verbose-vPrint verbose information about performed steps. diff --git a/dev/breeze/doc/images/output_testing_airflow-e2e-tests.txt b/dev/breeze/doc/images/output_testing_airflow-e2e-tests.txt index 264b35b577b4c..e855c1591bd13 100644 --- a/dev/breeze/doc/images/output_testing_airflow-e2e-tests.txt +++ b/dev/breeze/doc/images/output_testing_airflow-e2e-tests.txt @@ -1 +1 @@ -19ec59301c48da362cdb7e2bcef8c89f +bf6ebaaf5870518bae6d39a46326a769 diff --git a/dev/breeze/src/airflow_breeze/commands/testing_commands.py b/dev/breeze/src/airflow_breeze/commands/testing_commands.py index ed78236eb4edd..eb06d8fd967ff 100644 --- a/dev/breeze/src/airflow_breeze/commands/testing_commands.py +++ b/dev/breeze/src/airflow_breeze/commands/testing_commands.py @@ -1420,6 +1420,7 @@ def python_api_client_tests( "remote_log_opensearch", "xcom_object_storage", "event_driven", + "java_sdk", ], case_sensitive=False, ), diff --git a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py index 91fcd6cd9afc1..76bc3f0e7a0f6 100644 --- a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py +++ b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py @@ -101,6 +101,7 @@ ("docs", "/opt/airflow/docs"), ("generated", "/opt/airflow/generated"), ("go-sdk", "/opt/airflow/go-sdk"), + ("java-sdk", "/opt/airflow/java-sdk"), ("kubernetes-tests", "/opt/airflow/kubernetes-tests"), ("logs", "/root/airflow/logs"), ("providers", "/opt/airflow/providers"), diff --git a/dev/breeze/src/airflow_breeze/utils/selective_checks.py b/dev/breeze/src/airflow_breeze/utils/selective_checks.py index 5e2a91699908e..4dd6854fb1617 100644 --- a/dev/breeze/src/airflow_breeze/utils/selective_checks.py +++ b/dev/breeze/src/airflow_breeze/utils/selective_checks.py @@ -116,6 +116,7 @@ class FileGroupForCi(Enum): TASK_SDK_FILES = auto() TASK_SDK_INTEGRATION_TEST_FILES = auto() GO_SDK_FILES = auto() + JAVA_SDK_FILES = auto() AIRFLOW_CTL_FILES = auto() AIRFLOW_CTL_INTEGRATION_TEST_FILES = auto() BREEZE_INTEGRATION_TEST_FILES = auto() @@ -366,6 +367,9 @@ def __hash__(self): FileGroupForCi.GO_SDK_FILES: [ r"^go-sdk/.*\.go$", ], + FileGroupForCi.JAVA_SDK_FILES: [ + r"^java-sdk/", + ], FileGroupForCi.ASSET_FILES: [ r"^airflow-core/src/airflow/assets/", r"^airflow-core/src/airflow/models/assets/", @@ -1009,6 +1013,10 @@ def run_task_sdk_integration_tests(self) -> bool: def run_go_sdk_tests(self) -> bool: return self._should_be_run(FileGroupForCi.GO_SDK_FILES) + @cached_property + def run_java_sdk_tests(self) -> bool: + return self._should_be_run(FileGroupForCi.JAVA_SDK_FILES) + @cached_property def run_airflow_ctl_tests(self) -> bool: return self._should_be_run(FileGroupForCi.AIRFLOW_CTL_FILES) diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index fdc795161bbe3..222d69aab8a23 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -1151,6 +1151,7 @@ opsgenie Optimise optimise optimizationObjective +OptIn optionality ora oracledb diff --git a/java-sdk/.editorconfig b/java-sdk/.editorconfig new file mode 100644 index 0000000000000..1b89a6e999824 --- /dev/null +++ b/java-sdk/.editorconfig @@ -0,0 +1,31 @@ +# +# 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. + +root = true + +[*] +end_of_line = lf +insert_final_newline = true +indent_style = space +indent_size = 4 + +[*.java] +indent_size = 2 + +[*.kt] +indent_size = 2 diff --git a/java-sdk/.gitattributes b/java-sdk/.gitattributes new file mode 100644 index 0000000000000..a87d264c425cf --- /dev/null +++ b/java-sdk/.gitattributes @@ -0,0 +1,11 @@ +# +# https://help.github.com/articles/dealing-with-line-endings/ +# +# Linux start script should use lf +/gradlew text eol=lf + +# These are Windows script files and should use crlf +*.bat text eol=crlf + +# Binary files should be left untouched +*.jar binary diff --git a/java-sdk/.gitignore b/java-sdk/.gitignore new file mode 100644 index 0000000000000..bf1f44332ebc8 --- /dev/null +++ b/java-sdk/.gitignore @@ -0,0 +1,56 @@ +.gradle +build/ +!gradle/wrapper/gradle-wrapper.jar +!**/src/main/**/build/ +!**/src/test/**/build/ +.kotlin + +### IntelliJ IDEA ### +.idea/modules.xml +.idea/jarRepositories.xml +.idea/compiler.xml +.idea/libraries/ +*.iws +*.iml +*.ipr +out/ +!**/src/main/**/out/ +!**/src/test/**/out/ + +### Eclipse ### +.apt_generated +.classpath +.factorypath +.project +.settings +.springBeans +.sts4-cache +bin/ +!**/src/main/**/bin/ +!**/src/test/**/bin/ + +### NetBeans ### +/nbproject/private/ +/nbbuild/ +/dist/ +/nbdist/ +/.nb-gradle/ + +### VS Code ### +.vscode/ + +### Mac OS ### +.DS_Store +# Ignore Gradle build output directory +build + +### Artifacts of airflow standalone command ### +airflow.cfg +airflow.db +simple_auth_manager_passwords.json.generated +logs/dag_id=* +logs/dag_processor + +### Compatibility Test Results ### +validation/serialization/serialized_java.json +validation/serialization/serialized_python.json diff --git a/java-sdk/README.md b/java-sdk/README.md new file mode 100644 index 0000000000000..495bfad3eede3 --- /dev/null +++ b/java-sdk/README.md @@ -0,0 +1,101 @@ + + +# Airflow Java SDK + +A **JVM** SDK for Apache Airflow. You can use any JVM-compatible language to write +workflow bundles, and have Airflow consume the result. + +The SDK and execution-time logic is implemented in Kotlin. +An example is bundled showing how the SDK can be used in Java. + +## Building the SDK + +```bash +./gradlew build +``` + +## Running the example + +* Put the [DAG with stub tasks](./dags) to somewhere Airflow can find. + +* Ensure the `java` command is available in the same environment the Airflow + task worker is in. + +* Package the example and its dependencies into JARs in + `./example/build/install/example/lib` + + ```bash + ./gradlew :example:installDist + ``` + +* Configure Airflow to route tasks in the *java* queue to be run with Java: + + ```bash + export AIRFLOW__SDK__COORDINATORS='{ + "java": { + "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", + "kwargs": {"jars_root": ["/opt/airflow/java-sdk/example/build/install/example/lib"]} + } + }' + export AIRFLOW__SDK__QUEUE_TO_COORDINATOR='{"java": "java"}' + ``` + +* Ensure the Connection and Variable needed by the example DAG are available: + + ```bash + export AIRFLOW_CONN_TEST_HTTP='{ + "conn_type": "http", + "login": "user", + "password": "pass", + "host": "example.com", + "port": 1234, + "extra": {"param1": "val1", "param2": "val2"} + }' + export AIRFLOW_VAR_MY_VARIABLE=123 + ``` + +## Technical Details + +The user uses the SDK to implement a Java application that implements task +methods, and metadata on which DAG and task each method should be used +for. + +When the Airflow Supervisor identifies a task should be run with Java, it +launches the Java application as a subprocess. The Java application accepts +flags `--comm` and `--logs` from the command line to identify TCP sockets it +should connect to, and communicates with the Supervisor through these channels +during execution. + +1. On connection, the Supervisor immediately sends a StartupDetails message + through the comm socket. +2. The Java application finds and executes the relevant method. +3. During execution, the Java application uses the comm socket to retrieve + information (e.g. Variable) from, and send data (e.g. XCom) to Airflow. +4. The Java application informs the comm socket to tell the Supervisor the + task's terminal state. +5. The Java application exits. + +During the Java application's lifetime, it also sends log messages generated by +the SDK (not user code) through the logs socket, so the Supervisor can append +them to Airflow logs. + +Communication uses the same formats as the Python-based processes. + +See [Architectural Design Records](./adr) in the `adr` directory to learn more. diff --git a/java-sdk/adr/0001-java-sdk-airflow-integration.md b/java-sdk/adr/0001-java-sdk-airflow-integration.md new file mode 100644 index 0000000000000..1c7b73f5689b2 --- /dev/null +++ b/java-sdk/adr/0001-java-sdk-airflow-integration.md @@ -0,0 +1,384 @@ + + +# ADR-0001: Java SDK Airflow Integration + +## Status + +Accepted + +## Context + +Airflow's current execution model is Python-only: DAGs are Python files, tasks are Python +callables, and the supervisor communicates with the forked task process via UNIX domain +socketpairs. To support tasks authored in other languages (starting with Java), we need an +architecture that: + +- Allows non-Python tasks to coexist with Python tasks in the same DAG via `@task.stub`. +- Reuses the existing task-runner infrastructure so Airflow extensions (XCom backends, + connections, variables) stay in Python. +- Is extensible to other languages (Go, TypeScript, etc.) without per-language changes to + Airflow Core. + +The only missing piece is a way for the task runner to hand off execution to a +foreign-language process and still drive the same API-call lifecycle. + +## Decision + +### Writing a Non-Python Task + +There are two ways to author a Java task, both producing a task class the SDK runtime can +discover and execute. + +#### Interface-Based + +Implement the `Task` interface with `execute(Context context, Client client)`. `Context` +provides static run-time data (logical date, run ID, etc.), and `Client` provides access to +Airflow services (connections, variables, XCom): + +```java +public static class Extract implements Task { + public void execute(Context context, Client client) throws Exception { + var connection = client.getConnection("test_http"); + client.setXCom(new Date().getTime()); + } +} + +public static class Transform implements Task { + public void execute(Context context, Client client) { + var extractXcom = client.getXCom("extract"); + client.setXCom(new Date().getTime()); + } +} + +public static Dag build() { + var dag = new Dag("java_interface_example"); + dag.addTask("extract", Extract.class); + dag.addTask("transform", Transform.class); + return dag; +} +``` + +#### Annotation-Based + +Use `@Builder.Dag`, `@Builder.Task`, and `@Builder.XCom` annotations on a class and its +methods. The SDK's annotation processor generates a `Builder` class with `Task` +implementations at build time. XCom inputs declared via `@Builder.XCom` are fetched +automatically; non-`void` return values are pushed as XCom. + +```java +@Builder.Dag(id = "java_annotation_example") +public class AnnotationExample { + + @Builder.Task(id = "extract") + public long extractValue(Client client) throws InterruptedException { + var connection = client.getConnection("test_http"); + Thread.sleep(6000); + return new Date().getTime(); // automatically pushed as XCom + } + + @Builder.Task(id = "transform") + public long transformValue(Client client, + @Builder.XCom(task = "extract") long extracted) { + // `extracted` is pulled from the "extract" task XCom automatically + return new Date().getTime(); + } + + @Builder.Task + public void load(@Builder.XCom(task = "transform") long transformed) { + throw new RuntimeException("I failed"); + } +} +``` + +The generated `AnnotationExampleBuilder.build()` returns a fully configured `Dag`. The +annotation-based interface is generally preferred for new code because it eliminates +boilerplate and makes XCom data-flow explicit in method signatures. + +Both approaches register tasks with `BundleBuilder.getDags()` and are served by the same +`Server` entry point: + +```java +public class ExampleBundleBuilder implements BundleBuilder { + @Override + public Iterable getDags() { + return List.of(InterfaceExampleBuilder.build(), AnnotationExampleBuilder.build()); + } + + public static void main(String[] args) { + var bundle = new ExampleBundleBuilder().build(); + Server.create(args).serve(bundle); + } +} +``` + +### Integrating Non-Python Tasks into a DAG: `@task.stub` + +DAG authors declare a non-Python task in a Python DAG file using `@task.stub` and specify a +queue. Python and Java tasks coexist in the same pipeline; the DAG remains defined in Python: + +```python +@task() +def python_task_1(): + return "value_from_python_task_1" + + +@task.stub(queue="java") +def extract(): ... + + +@task.stub(queue="java") +def transform(): ... + + +@task() +def python_task_2(transformed): + print(transformed) + + +@dag(dag_id="java_interface_example") +def simple_dag(): + python_task_1() >> extract() >> transform() >> python_task_2() +``` + +The `@task.stub` declarations carry no Python implementation — execution is delegated to +the coordinator identified by the task's `queue`. + +### Public API Surface: `Client` and `Context` + +The Java task interface is `void execute(Context context, Client client)`. Two design choices +warrant explanation. + +**Why both `Context` and `Client`?** The Java SDK exposes two objects, mirroring the Go SDK: + +| Object | Holds | Lifecycle | +|---|---|---| +| `Context` | Static run-time data (`ds`, `ti`, logical date, run-id, etc.) | Populated once from `StartupDetails`, read-only during execution | +| `Client` | Active accessors that perform Execution API calls (connections, variables, XCom) | Each method call is a synchronous request/response over the comm channel | + +In Python, magic objects on the context (e.g., `outlet_events`) can perform Execution API +calls transparently because of the language's flexibility. Java is more rigid; making +`Context` itself perform background API calls would require significantly more wiring without +much user-visible benefit. Splitting the two surfaces makes the API call boundary explicit at +the type level. + +**Why is `execute` `void`?** Returning a value from `execute` would imply an automatic XCom +push. Java's static type system does not have a clean equivalent of Python's "return any +object, get a default-keyed XCom" pattern, and explicit `client.setXCom(...)` calls keep the +wire-level behavior obvious. (The annotation-based interface infers XCom pushes from +non-`void` return types, providing the same convenience without losing type clarity.) + +### Coordinator Interface and Code Reuse + +`BaseCoordinator`, defined in +`task-sdk/src/airflow/sdk/execution_time/coordinator.py`, exposes a single `execute_task` +method. Subclasses implement this to start the language-specific subprocess and return when +it finishes (with exit code and final task state). + +The Task SDK provides two coordinator implementations alongside `BaseCoordinator`: + +- **`_PythonCoordinator`** (built-in, not user-configurable) — implements `execute_task` by + calling `ActivitySubprocess.start()`, which creates UNIX domain socketpairs and forks a + child Python process. The child inherits the request socket on fd 0 and uses the existing + task-runner main function. This is the path taken for all Python tasks today. + +- **`JavaCoordinator`** (in `airflow.sdk.coordinators.java`) — implements `execute_task` by + creating two TCP server sockets on `127.0.0.1`, spawning the JVM bundle process via + `subprocess.Popen`, and waiting for the Java process to connect back to those servers. + It uses `_JavaActivitySubprocess`, a subclass of `ActivitySubprocess`, so the request + handling, heartbeating, and state management logic is fully shared with the Python path. + +The key design benefit: `ActivitySubprocess` owns the supervisor-side event loop +(heartbeating, API request proxying, state management). Both `_PythonCoordinator` and +`JavaCoordinator` create a subprocess and hand it an `ActivitySubprocess` instance; only the +subprocess start-up and socket establishment differ. Adding a third language requires +implementing `execute_task` in a new `BaseCoordinator` subclass, with no changes to Airflow +Core. + +The `client` parameter passed to `execute_task` is the already-authenticated Execution API +client. It is passed through to `ActivitySubprocess`, which uses it to forward the subprocess's +API requests (getVariable, getConnection, setXCom, etc.) to the API server. + +### Supervisor–Subprocess Communication + +Python tasks and Java tasks use different channels between the supervisor and the task +subprocess: + +**Python tasks** — the supervisor creates UNIX domain socketpairs before forking. The child +process inherits the request socket on fd 0 (and stdout/stderr on fd 1/2 via separate +socketpairs). No network stack is involved. + +**Java tasks** — the coordinator creates two TCP server sockets on `127.0.0.1` (one for the +msgpack comm channel, one for structured logs), then spawns the JVM process via +`subprocess.Popen`. The Java process connects *back* to those servers. From that point on, +the supervisor drives the same msgpack-framed request/response exchange as with Python tasks, +over TCP instead of UNIX sockets. + +The Java SDK process is agnostic to transport — it sees a TCP socket carrying msgpack frames +and behaves identically regardless of whether the other end is a Python supervisor or any +other implementation of the same protocol. + +### The Coordinator Layer + +When a task is dispatched, `CoordinatorManager` (in the same module as `BaseCoordinator`) +resolves the task's `queue` to a registered coordinator instance and calls `execute_task`. + +The Java coordinator ships as part of the Task SDK and is importable as +`airflow.sdk.coordinators.java.JavaCoordinator`. The `airflow.sdk.coordinators` namespace +package is structured to allow future separation into standalone distributions without +changing import paths. For packaging and registration details, see +[ADR-0005](0005-coordinator-packaging.md). + +### Architecture Overview + +``` + Airflow Backend Language Runtime Subprocess (Java in this example) + ─────────────── ────────────────────────────────────────────────── + +┌──────────────────────────────┐ +│ DAG File (Python) │ +│ │ +│ @task.stub(queue="java") │ +│ def my_java_task(): │ +│ ... │ +└──────────────┬───────────────┘ + │ (standard Python parsing) +┌──────────────▼───────────────┐ +│ Metadata DB │ +│ │ +│ task_instance.queue = "java"│ +└──────────────┬───────────────┘ + │ +┌──────────────▼───────────────┐ +│ Scheduler │ +│ │ +│ Reads queue from TI │ +│ ──► ExecuteTask workload │ +│ (includes queue) │ +└──────────────┬───────────────┘ + │ +┌──────────────▼───────────────┐ ┌──────────────────────────────┐ +│ Execution API │ │ Runtime Subprocess (Java) │ +│ │ │ │ +│ TI.queue ──► Startup │ │ execute_task() starts JVM │ +│ Details │ │ process, accepts TCP conn │ +└──────────────┬───────────────┘ │ │ + │ └──────────────▲───────────────┘ +┌──────────────▼───────────────┐ │ TCP +│ Supervisor │ │ +│ │ │ +│ CoordinatorManager │ │ +│ resolves queue via │ │ +│ [sdk] queue_to_coordinator ┼───────────────────────────────────┘ +│ → JavaCoordinator │ +└──────────────────────────────┘ +``` + +### Java Coordinator Configuration + +`JavaCoordinator` (in `task-sdk/src/airflow/sdk/coordinators/java/coordinator.py`) accepts +three configuration parameters via `kwargs`: + +| Parameter | Default | Description | +|---|---|---| +| `java_executable` | `"java"` | Path to the `java` binary | +| `jvm_args` | `[]` | Extra JVM arguments (e.g. `["-Xmx1024m"]`) | +| `jars_root` | `[]` | Directories scanned for the bundle JAR (`Main-Class` manifest entry is the entry point) | + +### Integration Points — Required Changes + +**1. Decorator — DAG Author Interface** + +DAG authors declare a non-Python task using `@task.stub` and specify a queue: + +```python +@task.stub(queue="java") +def my_java_task(): ... +``` + +**2. Execution API — Task Queues Routed to the Worker** + +`[sdk] coordinators` is a JSON object keyed by coordinator name. Each entry supplies a +`classpath` (resolved via `import_string`) and free-form `kwargs` passed to the class +constructor. `[sdk] queue_to_coordinator` maps queue names to those keys: + +```ini +[sdk] +coordinators = { + "jdk-17": { + "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", + "kwargs": {"java_executable": "java", "jars_root": ["/opt/airflow/jars"]} + } +} +queue_to_coordinator = {"java": "jdk-17"} +``` + +Tasks on the `java` queue are routed to the entry named `jdk-17`. Multiple entries with +the same `classpath` (e.g. `jdk-11` and `jdk-17`) are independent instances with different +`kwargs`; there is no subclassing needed for per-runtime variants. + +For the full configuration schema and multi-JDK examples, see +[ADR-0005](0005-coordinator-packaging.md). + +### Implementation Language: Kotlin (with a Java-First Public API) + +The user-facing API surface (`Task`, `Client`, `Context`, `Dag`, `BundleBuilder`) is +published as Java types and is the contract bundle authors program against. The SDK +*implementation* — `CoordinatorComm`, `Server`, `Task.kt`, `Frame.kt` — is written in +Kotlin. + +Kotlin compiles to the same JVM bytecode as Java and is fully interoperable, so this choice +is invisible to bundle authors at runtime. The practical reasons for using Kotlin internally: + +- **Null safety** is part of the type system, removing a large class of latent NPEs in the + comm/serde paths. +- **Coroutines and structured I/O** simplify the synchronous-over-async pattern used by + `Client.getVariable()` and friends. +- **Less boilerplate** in serialization and frame encoding code, which is the bulk of the SDK. + +Because the user-facing API is Java, "Java SDK" remains the accurate name from a DAG-author +perspective. A future rename to "JVM SDK" has been floated but is not adopted here; it can be +revisited if/when Scala or other JVM-language bindings are proposed. + +## Consequences + +### New Interfaces + +| Component | New Interface | Change Type | +|-----------|--------------|-------------| +| `BaseCoordinator` | Abstract base with single `execute_task` hook, defined in Task SDK | New class | +| `airflow.sdk.coordinators` | Namespace package for language coordinator modules | New namespace | +| `@task.stub` decorator | `queue: str \| None` parameter | Additive | +| `[sdk] coordinators` | Airflow configuration: JSON object of named coordinator entries | New option | +| `[sdk] queue_to_coordinator` | Airflow configuration mapping queue name → coordinator entry key | New option | +| `CoordinatorManager.for_queue` | Resolves queue → coordinator, falls back to `_PythonCoordinator` | New code path | + +### What Becomes Easier + +- Adding a new language runtime requires only a `BaseCoordinator` subclass and a corresponding + entry in `[sdk] coordinators` — no changes to Airflow Core and no provider plumbing. +- DAG authors can mix Python and non-Python tasks in the same pipeline. +- The existing `ActivitySubprocess` infrastructure (heartbeating, state management, API + request proxying) is reused for all language runtimes. + +### What Becomes Harder + +- Each non-Python task involves an additional subprocess and a TCP connection to it. +- Debugging non-Python tasks requires understanding the communication between the supervisor + and the language runtime. diff --git a/java-sdk/adr/0002-workload-execution.md b/java-sdk/adr/0002-workload-execution.md new file mode 100644 index 0000000000000..acfd91418c06a --- /dev/null +++ b/java-sdk/adr/0002-workload-execution.md @@ -0,0 +1,404 @@ + + +# ADR-0002: Workload Execution — Language-Specific Task Execution + +## Status + +Accepted + +## Context + +Airflow's standard task runner executes Python callables. To support tasks written in other languages, the pipeline needs an extension point where a language-specific coordinator can intercept the execution, delegate to an external runtime process, and bridge the Task SDK protocol so the external process can access Airflow services (connections, variables, XCom) during execution. + +This ADR details the task execution side of the coordinator architecture described in +[ADR-0001](0001-java-sdk-airflow-integration.md). It starts with the generic model — the +abstract contracts and expected behavior that any language must implement — then walks through +Java as a concrete example. + +The Python-side `BaseCoordinator` interface, `CoordinatorManager`, and the supervisor changes +needed to support them are implemented in the Task SDK alongside the Java coordinator. + +## Decision + +### Extension Point: `BaseCoordinator` + +`BaseCoordinator` (in `task-sdk/src/airflow/sdk/execution_time/coordinator.py`) exposes a +single `execute_task` method. Subclasses implement this to start the language-specific +subprocess and block until it completes, returning an `ExecutionResult` (exit code + final +task state string). + +There is no lower-level `task_execution_cmd` hook; each coordinator implementation is free to +start its subprocess however it likes. For details on the built-in Python coordinator and how +the Java coordinator reuses `ActivitySubprocess` infrastructure, see +[ADR-0001 — Coordinator Interface and Code Reuse](0001-java-sdk-airflow-integration.md#coordinator-interface-and-code-reuse). + +Coordinators contribute to the `airflow.sdk.coordinators` namespace package and are activated +through `[sdk] coordinators` in `airflow.cfg` — there is no `provider.yaml` involvement. + +### Registration and Discovery + +Coordinators are registered in `[sdk] coordinators` and routed via `[sdk] queue_to_coordinator` +in `airflow.cfg`. See [ADR-0001 — Java Coordinator Configuration](0001-java-sdk-airflow-integration.md#java-coordinator-configuration) +for a configuration example, and [ADR-0005](0005-coordinator-packaging.md) for the full +configuration schema and rationale. + +`CoordinatorManager.for_queue(ti.queue)` resolves the queue to a coordinator instance (or falls +back to `_PythonCoordinator`) and returns it to the supervisor, which then calls +`coordinator.execute_task(...)`. See +[ADR-0001 — The Coordinator Layer](0001-java-sdk-airflow-integration.md#the-coordinator-layer) +for how `CoordinatorManager` loads and caches coordinator instances. + +### Expected E2E Flow + +``` +Airflow Executor (dispatches task) + │ + ▼ +supervise_task() ← supervisor process entry point + │ + ├─ coordinator = CoordinatorManager.for_queue(ti.queue) + │ └─ returns JavaCoordinator (or _PythonCoordinator as fallback) + │ + ▼ +coordinator.execute_task(what, dag_rel_path, bundle_info, client, ...) + │ + │ [Python path: _PythonCoordinator] + ├─ ActivitySubprocess.start() + │ ├─ create UNIX domain socketpairs (requests on fd 0, stdout/stderr on fd 1/2) + │ ├─ fork child Python process + │ ├─ child runs task_runner main function + │ └─ supervisor drives event loop (heartbeats, API proxying) + │ + │ [Java path: JavaCoordinator] + └─ _JavaActivitySubprocess.start() + ├─ create TCP servers on 127.0.0.1:random (comm + logs) + ├─ spawn Java bundle process via subprocess.Popen + ├─ accept TCP connections from Java process + ├─ send StartupDetails to Java process over comm socket + └─ supervisor drives the same event loop as the Python path +``` + +For the transport details (UNIX socketpairs for Python, TCP loopback for Java) and why they +differ, see [ADR-0001 — Supervisor–Subprocess Communication](0001-java-sdk-airflow-integration.md#supervisorsubprocess-communication). + +### Expected Message Sequence + +Task execution is a multi-round conversation. The supervisor and the language runtime exchange +msgpack-framed messages directly over their shared channel (a UNIX socket for Python, a TCP +socket for Java): + +``` +Airflow Supervisor Language Runtime + │ │ + ├── StartupDetails ────────────────────────────►│ + │ │ + │ ├── Look up task from bundle + │ │ + │ ┌────────────────────┤ + │◄── GetConnection(conn_id)┤ Task code runs │ + ├── ConnectionResult ─────►│ and may request: │ + │◄── GetVariable(key) ─────┤ │ + ├── VariableResult ───────►│ │ + │◄── GetXCom(key, ...) ────┤ │ + ├── XComResult ───────────►│ │ + │◄── SetXCom(key, value..) ┤ │ + ├── (empty response) ─────►│ │ + │ └────────────────────┤ + │ │ + │◄── SucceedTask / TaskState ───────────────────┤ + │ (terminal — no response) │ + │ └── exit(0) +``` + +### Task SDK Protocol Messages + +The language runtime exchanges these message types with the Airflow supervisor: + +**Runtime → Supervisor (requests):** + +| Message | Fields | Purpose | +|---|---|---| +| `GetConnection` | `conn_id` | Fetch an Airflow connection by ID | +| `GetVariable` | `key` | Fetch an Airflow variable by key | +| `GetXCom` | `key`, `dag_id`, `task_id`, `run_id`, `map_index?`, `include_prior_dates?` | Fetch an XCom value | +| `SetXCom` | `key`, `value`, `dag_id`, `task_id`, `run_id`, `map_index`, `mapped_length?` | Store an XCom value | +| `SucceedTask` | `end_date`, `task_outlets?`, `outlet_events?` | Terminal: task succeeded | +| `TaskState` | `state` (`"failed"`, `"removed"`, `"skipped"`), `end_date` | Terminal: task ended non-successfully | + +**Supervisor → Runtime (responses):** + +| Message | Fields | In response to | +|---|---|---| +| `ConnectionResult` | `conn_id`, `conn_type`, `host`, `schema`, `login`, `password`, `port`, `extra` | `GetConnection` | +| `VariableResult` | `key`, `value` | `GetVariable` | +| `XComResult` | `key`, `value` | `GetXCom` | +| (empty) | | `SetXCom` | +| `ErrorResponse` | `error`, `detail` | Any request that failed server-side | + +**Framing:** Every message is a length-prefixed msgpack frame. Requests are `[id, body]` (2-element array); responses are `[id, body, error]` (3-element array). The `id` field correlates request/response pairs. + +### Request/Response Semantics + +The task execution follows a synchronous request/response pattern from the runtime's perspective: + +1. The runtime sends a request frame (e.g., `GetVariable`) with an incrementing `id` +2. The supervisor reads the frame, fulfills the request (e.g., calls the Execution API), and sends back a response with the same `id` +3. The runtime blocks until it receives the response +4. This repeats for each Airflow service call the task code makes +5. When the task finishes, the runtime sends a terminal message (`SucceedTask` or `TaskState`) — no response is expected, and the process exits + +### IPC Forward-Compatibility Contract + +The supervisor-to-runtime IPC schema (the messages enumerated above plus `StartupDetails`) is shared between Airflow Core (Python) and every language SDK. A formal AIP for this protocol is expected as follow-up work; until then, this section pins down the rules that the Java SDK assumes and that any future SDK (Go, Rust, …) must follow. + +**Codec rule (load-bearing).** Every SDK MUST configure its decoder to ignore unknown fields: + +- Python side: `msgspec` / Pydantic models are forward-compatible by default. +- Java side: `TaskSdkFrames.kt` configures the Jackson `ObjectMapper` with `FAIL_ON_UNKNOWN_PROPERTIES = false`. A short comment at that call site documents that this is contract, not preference — flipping it back to the Jackson default would break forward compatibility with Core. +- Any new SDK: pick a codec configuration that mirrors this (silent drop of unknown fields). + +This rule is what makes additive Core changes safe to ship without bumping a version on every SDK. The analogous trap — generated clients that emit their *own* allowlist check before the configured mapper sees the bytes — has bitten downstream Java consumers in unrelated systems; flagging the contract here makes it visible to future SDK authors. + +**Change classification.** + +| Change to a message | Status | Required action | +|---|---|---| +| Add a new optional field | **Non-breaking.** Decoders ignore it; old SDKs unaffected. | None. Just ship it. | +| Add a new required field | Breaking. | Deprecation cycle: ship as optional first, populate from Core, wait for SDKs to consume it, then tighten. | +| Rename a field | Breaking. | Deprecation cycle: emit both names from Core during transition. | +| Change a field's type | Breaking. | Deprecation cycle, typically via a new field name + parallel emission. | +| Remove a required field | Breaking. **Especially dangerous in Java**: `lateinit var` properties on `StartupDetails` deserialize silently and only throw `UninitializedPropertyAccessException` on first access, so the failure surfaces inside user task code rather than at the protocol boundary. | Deprecation cycle. Prefer making the field optional first, then remove after a release in which all SDKs have absorbed the change. | + +**Recommended testing.** A small contract test on the SDK side should feed the decoder synthetic frames that exercise the rules above — an unknown field, a missing optional field, a `null` in an optional position — so that a future codec-config regression is caught before it reaches users. Such IPC-envelope tests are currently in the follow-up bucket. + +### Runtime Lifecycle and Worker Capability + +The language runtime is **ephemeral and one-process-per-task**: + +- Each task instance launches its own `java -classpath /* --comm=… --logs=…` (or the equivalent for another language). The lifetime of that process is the lifetime of the task. There is no pooling or warm-pool reuse. +- Parallelism on a single worker therefore equals the number of concurrently running task processes. Five concurrent Java tasks on one worker means five JVMs. +- DAG parsing has the same shape: each `DagFileProcessorProcess` child handles one parse request and exits. The language runtime spawned underneath it inherits that ephemerality. + +**Worker capability is opt-in.** A worker can run a non-Python task only if the Task SDK (which includes the language coordinator module) is installed, the matching coordinator instance is declared in `[sdk] coordinators`, and the language toolchain (e.g., a JRE) is on the host. There is no requirement that every worker support every language. Routing relies on: + +| Layer | Mechanism | +|---|---| +| Author intent | `@task.stub` declares `queue="java"` (or any custom queue) | +| Worker selection | The executor (Celery, Kubernetes, etc.) routes the task to a worker that consumes that queue, exactly as it does for Python tasks today | +| Runtime selection | Inside the task runner, `[sdk] queue_to_coordinator` maps the queue name to a coordinator instance name; that name is resolved against `[sdk] coordinators` to obtain the configured class and its `kwargs`; `CoordinatorManager.for_queue` instantiates the coordinator and `execute_task` is called | + +The deployment model is the same one that already applies to Python providers: install what your DAGs need, on the hosts they run on. Multi-language workers are possible (install both providers and both toolchains) but not required. + +**JAR / artifact version compatibility.** The Java SDK embeds its version in the bundle JAR via the `Airflow-Java-SDK-Version` manifest attribute. Validating that a bundle's SDK version matches the installed `JavaCoordinator` version at execution time is planned but not yet wired in; this is a follow-up to add before promoting the SDK out of preview. + +### StartupDetails + +The first message the runtime receives is `StartupDetails`, which provides full context for the task: + +| Field | Type | Description | +|---|---|---| +| `ti` | `TaskInstance` | id, task_id, dag_id, run_id, try_number, dag_version_id, map_index, context_carrier | +| `dag_rel_path` | string | Relative path to the DAG file / bundle | +| `bundle_info` | `BundleInfo` | name, version | +| `start_date` | datetime | When this task attempt started | +| `ti_context` | `TIRunContext` | DAG run context (logical date, data interval, etc.) | +| `sentry_integration` | string | Sentry DSN for error reporting (optional) | + +### What a Language SDK Must Implement + +For task execution, a new language SDK needs: + +1. **A `BaseCoordinator` subclass** with: + - An `__init__` that accepts the kwargs declared in `[sdk] coordinators` (e.g., interpreter path, language-specific runtime flags) + - `execute_task(...)` — starts the language-specific subprocess, drives the `ActivitySubprocess` event loop, and returns when the task finishes + +2. **A runtime process** that: + - Accepts `--comm=host:port` and `--logs=host:port` CLI arguments + - Connects to both TCP addresses + - Reads a `StartupDetails` msgpack frame from the comm channel + - Looks up the task to execute from its bundle using `ti.dag_id` and `ti.task_id` + - Executes the task, making `GetConnection`/`GetVariable`/`GetXCom`/`SetXCom` requests as needed + - Sends `SucceedTask` on success or `TaskState("failed")` on failure + - Exits + +3. **A task interface** that user code implements (analogous to Python's `@task` decorator or `BaseOperator`) + +4. **A client API** that wraps the socket protocol behind a simple interface (get_connection, get_variable, get_xcom, set_xcom) so task authors don't deal with framing + +5. **Distribution** under `airflow.sdk.coordinators.` — currently shipped as part of the Task SDK; a standalone distribution is possible in the future without changing the import path + +### Java as a Concrete Example + +**JavaCoordinator (Python side):** + +See [ADR-0001 — Coordinator Interface and Code Reuse](0001-java-sdk-airflow-integration.md#coordinator-interface-and-code-reuse) +for how `JavaCoordinator` implements `execute_task`, why it uses `_JavaActivitySubprocess`, +and how the Python and Java paths share `ActivitySubprocess` infrastructure. For configuration +parameters and an `airflow.cfg` example, see +[ADR-0001 — Java Coordinator Configuration](0001-java-sdk-airflow-integration.md#java-coordinator-configuration). + +**Java SDK Task Interface:** + +User task code implements a single-method interface: + +```java +// sdk: org.apache.airflow.sdk.Task +public interface Task { + void execute(Client client) throws Exception; +} +``` + +The `Client` provides access to Airflow services: + +```java +// sdk: org.apache.airflow.sdk.Client +public class Client { + // Access task metadata + public StartupDetails getDetails(); + + // Airflow services + public Connection getConnection(String id); + public Object getVariable(String key); + public Object getXCom(String key, String dagId, String taskId, String runId, ...); + public void setXCom(String key, Object value); // defaults: key="return_value", dagId/taskId/runId from current task +} +``` + +**Java SDK Task Execution Flow:** + +When the bundle process receives `StartupDetails`: + +``` +CoordinatorComm.handleIncoming(frame) + │ + ├── frame.body is StartupDetails + │ ti: TaskInstance (id, dagId, taskId, runId, tryNumber, ...) + │ dagRelPath, bundleInfo, startDate, tiContext + │ + ▼ +TaskRunner.run(bundle, request, comm) + │ + ├── Create Client(request, CoordinatorClient(comm)) + │ CoordinatorClient wraps the comm channel behind the Client interface + │ + ├── Look up task class: + │ bundle.dags[request.ti.dagId]?.tasks[request.ti.taskId] + │ └── if not found → return TaskState("removed") + │ + ├── Instantiate task: + │ task.getDeclaredConstructor().newInstance() + │ + ├── Execute: + │ try { + │ instance.execute(client) ← USER TASK CODE RUNS HERE + │ return SucceedTask() + │ } catch (Exception e) { + │ return TaskState("failed") + │ } + │ + ▼ +sendMessage(frame.id, result) ← sends SucceedTask or TaskState back +shutDownRequested = true ← one-shot, process will exit +``` + +**Java SDK Airflow Service Access:** + +When user task code calls `client.getVariable("my_key")`, the call chain is: + +``` +client.getVariable("my_key") // Client.kt (public SDK) + │ + └── impl.getVariable("my_key") // CoordinatorClient (execution) + │ + └── runBlocking { // blocks the calling thread + comm.communicate( // CoordinatorComm + GetVariable(key = "my_key") + ) + } + │ + ├── sendMessage(nextId++, GetVariable) // encode + write to comm socket + │ ├── encode: [id, {"type": "GetVariable", "key": "my_key"}] + │ └── write: [4-byte len][msgpack] + │ + ├── processOnce(::handle) // block until response arrives + │ ├── read 4-byte length prefix + │ ├── read payload + │ └── decode: [id, {"type": "VariableResult", ...}, null] + │ + └── return response.value // unwrap VariableResponse +``` + +This is fully synchronous from the task code's perspective — `getVariable()` blocks until the supervisor responds. + +**Java SDK Example Task Implementation:** + +```java +public static class Extract implements Task { + public void execute(Client client) throws Exception { + // Read XCom from a Python task in the same DAG + var pythonXcom = client.getXCom("python_task_1"); + + // Access Airflow connections + var connection = client.getConnection("test_http"); + + // Do work... + Thread.sleep(6000); + + // Push XCom for downstream tasks (Java or Python) + client.setXCom(new Date().getTime()); + } +} + +public static class Transform implements Task { + public void execute(Client client) { + // Read XCom from upstream Java task + var extractXcom = client.getXCom("extract"); + + // Access Airflow variables + var variable = client.getVariable("my_variable"); + + // Push XCom (readable by downstream Python tasks) + client.setXCom(new Date().getTime()); + } +} + +public static class Load implements Task { + public void execute(Client client) { + var xcom = client.getXCom("transform"); + throw new RuntimeException("I failed"); + // Exception → TaskRunner catches → sends TaskState("failed") + } +} +``` + +**Java SDK Complete Bundle Entry Point:** + +See [ADR-0001 — Writing a Non-Python Task](0001-java-sdk-airflow-integration.md#writing-a-non-python-task) +for the full `BundleBuilder` / `Server.create(args).serve(bundle)` pattern. From the task +execution perspective, `main()` is the JVM entry point the coordinator launches; `StartupDetails` +is the first message received, which triggers `runTask()`, and the process exits after the +terminal `SucceedTask`/`TaskState` response. + +## Consequences + +- Task execution for any language reuses the same coordinator pattern, keeping the extension surface small. +- The multi-round protocol (GetConnection, GetVariable, etc.) means the language runtime has full access to Airflow services without reimplementing them — they stay in Python. +- The synchronous request/response model is simple for language SDK authors but adds a round-trip per service call. +- Task authors interact with a simple `Client` interface, completely abstracted from the underlying socket protocol. diff --git a/java-sdk/adr/0003-pure-java-dags.md b/java-sdk/adr/0003-pure-java-dags.md new file mode 100644 index 0000000000000..e02740a50a80b --- /dev/null +++ b/java-sdk/adr/0003-pure-java-dags.md @@ -0,0 +1,258 @@ + + +# ADR-0003: Pure Java DAGs — Build-Time Packaging and Code Visibility + +## Status + +Proposed + +> **Note:** This ADR describes pure Java DAG authoring (entire DAGs written in Java without a +> Python file), which was removed from the scope of +> [AIP-108](https://cwiki.apache.org/confluence/x/pY4mGQ). Per AIP-108, Java tasks are +> declared as `@task.stub` in ordinary Python DAG files; pure Java DAG authoring is left to a +> future proposal, likely after [AIP-85](https://cwiki.apache.org/confluence/x/_Q7OEg) +> stabilises. The `BundleBuilder` interface and `Bundle`/`BundleBuilder.getDags()` mechanism +> remain in the SDK as the internal registry that the task execution runtime uses to locate task +> classes — they are not a public DAG-authoring surface in the current release. + +## Context + +[ADR-0001](0001-java-sdk-airflow-integration.md) originally introduced two ways to integrate non-Python tasks: `@task.stub` (mixed Python+Java DAGs) and pure Java DAGs (entire DAG in Java via `BundleBuilder`). [ADR-0004](0004-dag-parsing.md) and [ADR-0002](0002-workload-execution.md) describe the coordinator infrastructure for DAG parsing and task execution respectively. + +This ADR focuses on the Java-SDK-specific concerns that would make pure Java DAGs work end-to-end — build-time metadata generation, source code packaging for UI visibility, and JAR manifest conventions — rather than the shared coordinator infrastructure already covered in those ADRs. + +The central challenge is that Airflow Core expects to read DAG metadata and source code from files on disk or from the metadata DB. A JAR is an opaque binary — Airflow cannot `open()` it and read Python source. The Java SDK would need to bridge this gap at build time by embedding machine-readable metadata and human-readable source into the JAR itself. + +## Decision + +### JAR Manifest Conventions + +The JAR manifest (`META-INF/MANIFEST.MF`) carries three SDK-specific attributes that Airflow and the Java SDK use to bootstrap a bundle: + +| Attribute | Example Value | Purpose | +|---|---------------------------------------------------|---| +| `Main-Class` | `org.apache.airflow.example.ExampleBundleBuilder` | Standard Java attribute; the coordinator uses it to launch the JVM | +| `Airflow-Java-SDK-Metadata` | `airflow-metadata.yaml` | Points to the embedded metadata file (dag IDs, task IDs) | +| `Airflow-Java-SDK-Dag-Code` | `JavaExampleBuilder.java` | Points to the embedded source file for Airflow UI display | + +These attributes are set in the Gradle build (see [Build-Time Packaging](#build-time-packaging-gradle) below). The Python-side coordinator reads `Main-Class` to construct the launch command; `BundleScanner` reads `Airflow-Java-SDK-Metadata` to discover DAG IDs without launching the JVM. + +### Build-Time Metadata: `airflow-metadata.yaml` + +At build time, the SDK runs `BundleInspector` — a build-time utility that reflectively instantiates the user's `BundleBuilder` class, calls `getDags()`, and writes a YAML file listing every DAG ID and its task IDs: + +```yaml +dags: + java_example: + tasks: + - extract + - transform + - load +``` + +This file is embedded in the JAR root and referenced by the `Airflow-Java-SDK-Metadata` manifest attribute. + +**Why build-time, not runtime?** The metadata must be available before the JVM starts. `BundleScanner` reads it from the JAR to discover which DAG IDs a bundle contains — this is used for `@task.stub` routing (mapping a `dag_id` to the correct bundle's classpath) without paying JVM startup cost. For pure Java DAGs, the coordinator already knows the bundle path, but the metadata is still useful for validation and tooling. + +**`BundleInspector`:** + +```kotlin +object BundleInspector { + @JvmStatic + fun main(args: Array) { + val className = args[0] + val outputPath = args[1] + val clazz = Class.forName(className) + val instance = clazz.getDeclaredConstructor().newInstance() as? BundleBuilder + ?: error("$className does not implement BundleBuilder") + val dags = instance.getDags() + File(outputPath).apply { parentFile.mkdirs() }.writeText(toYaml(dags)) + } + + internal fun toYaml(dags: List): String = buildString { + appendLine("dags:") + for (dag in dags) { + appendLine(" ${dag.dagId}:") + appendLine(" tasks:") + for (taskId in dag.tasks.keys) { + appendLine(" - $taskId") + } + } + } +} +``` + +### Source Code Packaging for UI Visibility + +Airflow stores DAG source code in the `dag_code` table and displays it in the web UI. For Python DAGs this is trivial — `DagCode.write_code()` reads the `.py` file from disk. For a JAR, the raw bytecode is not human-readable. + +The solution: pack the original `.java` source file into the JAR at build time. The `Airflow-Java-SDK-Dag-Code` manifest attribute tells the coordinator which file to extract. + +On the Python side, `get_code_from_file()` on the coordinator: + +1. Opens the JAR as a ZIP +2. Reads the `Airflow-Java-SDK-Dag-Code` attribute from the manifest +3. Extracts and returns the raw `.java` source + +This lets Airflow's existing `DagCode` infrastructure store and display Java source code with no changes to Airflow Core. + +### Build-Time Packaging (Gradle) + +The `example/build.gradle.kts` shows the complete packaging pattern: + +```kotlin +val bundleMainClass = application.mainClass.get() +val metadataFileName = "airflow-metadata.yaml" +val metadataOutputDir = layout.buildDirectory.dir("airflow-metadata") +val dagCodeSourcePath = bundleMainClass.replace('.', '/') + ".java" +val dagCodeFileName = bundleMainClass.substringAfterLast('.') + ".java" + +// 1. Run BundleInspector at compile time to generate metadata +val inspectBundle = tasks.register("inspectBundle") { + dependsOn("classes") + classpath = sourceSets.main.get().runtimeClasspath + mainClass.set("org.apache.airflow.sdk.BundleInspector") + args = listOf(bundleMainClass, metadataOutputDir.get().file(metadataFileName).asFile.absolutePath) +} + +// 2. Pack metadata + source into the JAR +tasks.withType { + dependsOn(inspectBundle) + from(metadataOutputDir) // airflow-metadata.yaml + from("src/java/$dagCodeSourcePath") // raw .java source file + manifest { + attributes( + "Main-Class" to bundleMainClass, + "Airflow-Java-SDK-Version" to project.version, + "Airflow-Java-SDK-Metadata" to metadataFileName, + "Airflow-Java-SDK-Dag-Code" to dagCodeFileName, + ) + } +} +``` + +The resulting JAR contains: + +``` +example.jar +├── META-INF/MANIFEST.MF (Main-Class, SDK attributes) +├── airflow-metadata.yaml (dag IDs + task IDs) +├── JavaExampleBuilder.java (raw source for UI display) +├── org/apache/airflow/example/ +│ ├── JavaExampleBuildser.class (compiled bundle entry point) +│ ├── JavaExampleBuilder$Extract.class +│ ├── JavaExampleBuilder$Transform.class +│ └── JavaExampleBuilder$Load.class +└── ... (SDK + dependency classes) +``` + +### `BundleScanner` — Runtime Bundle Discovery + +`BundleScanner` reads JAR manifests at runtime to discover bundles without launching the JVM. This is used by the `@task.stub` path to resolve which bundle contains a given `dag_id`. + +```kotlin +data class ResolvedBundle( + val mainClass: String, // From Main-Class manifest attribute + val classpath: String, // All JARs in bundle directory, colon-separated +) + +fun scanBundles(bundlesDir: Path): Map +``` + +It supports two directory layouts: + +- **Nested**: each subdirectory of `bundlesDir` is a bundle home (e.g., `bundles/my-app/lib/*.jar`) +- **Flat**: `bundlesDir` itself contains the JARs (e.g., `bundles/*.jar`) + +For each JAR, it reads the `Airflow-Java-SDK-Metadata` manifest attribute, extracts the referenced YAML, parses DAG IDs, and returns a mapping from `dag_id` to `ResolvedBundle`. + +### The BundleBuilder Authoring API + +Bundle authors implement builder classes to define their DAGs: + +```java +public class JavaExampleBuilder { + + public static class Extract implements Task { + public void execute(Client client) throws Exception { + var connection = client.getConnection("test_http"); + client.setXCom(new Date().getTime()); + } + } + + public static class Transform implements Task { + public void execute(Client client) { + var extract_xcom = client.getXCom("extract"); + client.setXCom(new Date().getTime()); + } + } + + public static Dag build() { + var dag = new Dag("java_example", null, "@daily"); + dag.addTask("extract", Extract.class, List.of()); + dag.addTask("transform", Transform.class, List.of("extract")); + return dag; + } +} +``` + +and then collect DAGs with a BundleBuilder: + +```java +public class ExampleBundleBuilder implements BundleBuilder { + public Iterable getDags() { + return List.of(JavaExampleBuilder.build()) + } + + public static void main(String[] args) { + var bundle = new ExampleBundleBuilder().build(); + Server.create(args).serve(bundle); + } +} +``` + +The `main()` method is the JVM entry point that the coordinator launches. It wires the `BundleBuilder` to the SDK's TCP communication layer (`Server` → `CoordinatorComm`), which handles DAG parsing requests and task execution commands as described in [ADR-0004](0004-dag-parsing.md) and [ADR-0002](0002-workload-execution.md). + +> **Note:** The current `BundleBuilder` interface is subject to review before the SDK reaches 1.0. Subclassing `Dag` directly may be a more natural fit and is being considered for post-OSS-integration. + +### Deployment and Updates + +A reasonable concern about JAR-based DAGs is whether updating a bundle requires draining or restarting the DAG processor / workers — Python source files are flexible because everything is read fresh on each parse, but a long-lived JVM holding a JAR open could pin an old version. + +The design avoids this by leaning on the same ephemerality that Python uses: + +- **DAG processor.** `DagFileProcessorManager` is long-lived, but each `DagFileProcessorProcess` child is one-shot and exits after returning a `DagFileParseRequest`. The Java runtime spawned underneath it (`java -classpath /* …`) shares that lifetime — it loads the JAR fresh on every parse, then exits. Replacing the JAR on disk takes effect on the next scheduled parse with no manager restart. +- **Workers.** Each task instance launches its own JVM ([ADR-0002 — Runtime Lifecycle and Worker Capability](0002-workload-execution.md#runtime-lifecycle-and-worker-capability)). The classloader is process-scoped; a swapped JAR is picked up the next time a task starts. There is no warm JVM pool to invalidate. + +In practice, "updating a Java DAG bundle" is the same shape as "updating a Python DAG file": drop the new file (or directory of JARs) into the bundle location and let normal scheduling pick it up. The version that runs a given task instance is determined at task start, not at worker start. + +Two operational details worth flagging: + +- **Atomic swap.** Writing a JAR in place while a task happens to be loading it can yield a corrupted read. Operators should prefer the standard "write to a temp name, rename into place" pattern, which the file system handles atomically on POSIX. This is the same guidance that already applies to Python file-system bundles. +- **Mid-run version skew.** Because the version is resolved per task launch, a long-running DAG run can in principle observe one bundle version for an upstream task and a different version for a downstream task if a swap happens between them. Bundle-version validation against `Airflow-Java-SDK-Bundle-Version` (planned — distinct from `Airflow-Java-SDK-Version`, which identifies the SDK toolkit; see [ADR-0002](0002-workload-execution.md#runtime-lifecycle-and-worker-capability)) gives operators a way to detect skew if it matters; the data-plane consequences (XCom shape changes, etc.) are the bundle author's responsibility, exactly as with Python. + +## Consequences + +- JAR bundles are self-contained: metadata, source, and compiled code are all in one artifact, simplifying deployment (copy one directory of JARs). +- Build-time metadata generation means DAG IDs can be discovered without JVM startup — important for `BundleScanner` and tooling. +- Source code packaging enables Airflow UI display with no changes to Airflow Core's `DagCode` infrastructure. +- The manifest convention (`Airflow-Java-SDK-*` attributes) is extensible — future attributes can carry additional metadata without breaking existing tooling. +- The build-time `BundleInspector` step adds a compile-time dependency on the SDK and requires the `BundleBuilder` class to be instantiable without side effects (no I/O, no connections in the constructor). +- Bundle authors must follow the Gradle packaging pattern (or replicate it in Maven/other build tools) — this is SDK-specific boilerplate that doesn't exist for Python DAGs. diff --git a/java-sdk/adr/0004-dag-parsing.md b/java-sdk/adr/0004-dag-parsing.md new file mode 100644 index 0000000000000..564e237ea29a7 --- /dev/null +++ b/java-sdk/adr/0004-dag-parsing.md @@ -0,0 +1,414 @@ + + +# ADR-0004: DAG Parsing — Language-Specific DAG File Processing + +## Status + +Proposed + +> **Note:** This ADR describes language-native DAG file parsing, which was removed from the +> scope of [AIP-108](https://cwiki.apache.org/confluence/x/pY4mGQ). Per AIP-108, Java tasks are +> declared as `@task.stub` in ordinary Python DAG files; the Java SDK and coordinator are +> responsible only for *task execution*, not for DAG parsing. This document is retained for +> future reference if language-native DAG parsing is revisited in a follow-up AIP (likely after +> [AIP-85](https://cwiki.apache.org/confluence/x/_Q7OEg) stabilises). + +## Context + +Airflow's standard DAG file processor only understands Python files. To support DAGs defined in other languages (Java, Go, Rust, etc.), the pipeline would need an extension point where a language-specific processor can intercept the parsing request, delegate to an external runtime, and return a result in the same format the Airflow scheduler expects. + +This ADR details the DAG parsing side of the coordinator architecture described in [ADR-0001](0001-java-sdk-airflow-integration.md). It starts with the generic model — the abstract contracts and expected behavior that any language must implement — then walks through Java as a concrete example. + +## Decision + +### Extension Point: `BaseCoordinator` + +A single abstract base class — `BaseCoordinator` — handles both DAG parsing and task execution. Concrete subclasses ship as standalone distributions (`apache-airflow-coordinators-`) under the shared `airflow.sdk.coordinators` namespace package; they are **not** Airflow providers and are not registered through `provider.yaml`. For DAG parsing, a subclass must implement two methods: + +| Method | Signature | Responsibility | +|---|---|---| +| `can_handle_dag_file` | `(bundle_name, path) -> bool` | Return `True` if this coordinator should handle the given file. Default returns `False`; subclasses add language-specific checks (e.g., "is this a JAR with a Main-Class?"). | +| `dag_parsing_cmd` | `(dag_file_path, bundle_name, bundle_path, comm_addr, logs_addr) -> list[str]` | Return the full command to launch the language runtime. `comm_addr` and `logs_addr` are `host:port` strings the process must connect to. | + +### Registration + +Coordinators are configured in `airflow.cfg`. See +[ADR-0001 — Java Coordinator Configuration](0001-java-sdk-airflow-integration.md#java-coordinator-configuration) +for a configuration example, and [ADR-0005](0005-coordinator-packaging.md) for the full +schema. A single instance entry covers both DAG parsing and task execution — there are no +separate registries for the two roles. + +**Per-host opt-in.** A coordinator becomes active on a given DAG processor host only when +its module is installed there *and* its instance appears in `[sdk] coordinators`. A deployment +can run a Python-only DAG processor pool and a separate Java-capable pool by omitting the +coordinator config on the Python-only hosts. There is no requirement that every parser carry +a JDK. + +### Discovery: `_resolve_processor_target()` + +When `DagFileProcessorProcess.start()` needs to parse a file: + +``` +_resolve_processor_target(path, bundle_name, bundle_path) + for entry in conf.get("sdk", "coordinators"): + coordinator_cls = import_string(entry["classpath"]) + coordinator = coordinator_cls(name=entry["name"], **entry.get("kwargs", {})) + if coordinator.can_handle_dag_file(bundle_name, path): + return functools.partial(coordinator.run_dag_parsing, path=..., bundle_name=..., bundle_path=...) + return None # fall back to default Python parser +``` + +The first coordinator instance whose `can_handle_dag_file()` returns `True` wins. If none match, the default Python `_parse_file_entrypoint` runs. Instances are constructed lazily from `[sdk] coordinators` and cached for the lifetime of the host process. + +### Transport: Why msgpack over TCP Loopback + +A natural reviewer question is "why a custom-looking framed-msgpack protocol over `127.0.0.1:`, and not Unix sockets / gRPC / HTTP REST?" Two clarifications are important: + +1. **The protocol is not new for the Java SDK.** Length-prefixed msgpack frames are the existing transport between the Airflow supervisor and the Python task runner (see `task-sdk/src/airflow/sdk/execution_time/supervisor.py` and `comms.py`). The coordinator bridge wires the language-runtime sockets onto that same byte stream — it does not define a new wire format. Migrating it would be a separate, pan-SDK change. +2. **Forward-compat for IPC messages is treated as a contract**, not as a transport choice. The decoder rules that all SDKs must follow are stated in [ADR-0002 — IPC Forward-Compatibility Contract](0002-workload-execution.md#ipc-forward-compatibility-contract). + +#### Alternatives considered + +| Option | Why not (today) | +|---|---| +| **Unix domain sockets** instead of TCP loopback | Avoids the IPv6/dual-stack concern with `127.0.0.1`, and matches conventions like Docker's `/var/run/docker.sock`. Worth revisiting once a formal IPC AIP lands; not adopted now because it would diverge from the existing Python supervisor transport, which is also TCP loopback. | +| **gRPC / Protocol Buffers** | Would require defining an intermediate IDL for `DagFileParseRequest`, `StartupDetails`, etc. The internal serialization that the language runtime returns (DagSerialization v3) is *not* expressible as a flat ProtoBuf without losing information — see "Cross-SDK serialization compatibility" below. gRPC would replace one custom-looking layer with two: ProtoBuf for transport plus a separate JSON-shaped DAG payload nested inside it. | +| **HTTP REST** | Adds an HTTP server in every language runtime and an HTTP client in the supervisor for a strictly local, single-peer connection. None of HTTP's value (intermediaries, caching, content negotiation) applies. The Java SDK's `Supervisor.kt` already does HTTP for the *Execution API* (Edge-worker path); the comm channel between supervisor and language runtime is intentionally lower-level. | +| **Keep msgpack-over-TCP** (chosen) | Reuses the existing supervisor transport unchanged; the bridge is a pure byte forwarder. New language SDKs only need a length-prefixed-msgpack codec, which exists in every target language. | + +A formal AIP for the supervisor-to-runtime comm protocol is expected as a follow-up once two or more language SDKs (Java, Go) are in tree; that AIP is the natural place to revisit transport and framing. + +### Cross-SDK Serialization Compatibility + +The `DagFileParsingResult` payload that a language runtime returns is the *Airflow internal* serialized DAG format, not an SDK-defined schema. The authoritative reference is `airflow-core/src/airflow/serialization/schema.json`, which describes `LazyDeserializedDAG` (see `airflow-core/src/airflow/dag_processing/processor.py` and `airflow-core/src/airflow/serialization/serialized_objects.py`). The scheduler reads this format directly into its internal model — any divergence is a parsing failure. + +**Why a per-language reimplementation rather than codegen?** The first attempt was to generate POJOs from `schema.json` (similar to how Pydantic models are generated from OpenAPI specs). That approach was abandoned because the generated types miss the wrapping/unwrapping rules that distinguish "decorated" fields (kept as `{"__type", "__var"}`) from "non-decorated" fields (unwrapped to the bare value), as well as the timetable/task encoding rules listed below. Wiring an extra translation layer on top of generated types added more code than implementing the serializer directly per language. + +**Compatibility strategy.** Each language SDK ships its own serializer plus a cross-SDK validator: + +- A shared `test_dags.yaml` defines logical fixtures. +- Python emits `serialized_python.json` via `DagSerialization.serialize_dag()`. +- Each language SDK emits `serialized_.json` via its own serializer. +- `compare.py` does a field-by-field comparison and fails on divergence. + +This validator is planned to run as a CI gate (PR #65959). A complementary direction (suggested by reviewers, deferred): publish JSON schemas for the IPC envelope types themselves (`DagFileParsingResult`, `StartupDetails`, `TaskInstance`), which are currently undocumented because they were Python-to-Python only. That work is out of scope for the Java SDK PR but is a sensible next step once a second language SDK is in tree. + +### What the Base Class Handles Automatically + +The matched coordinator's `run_dag_parsing()` (a concrete method on `BaseCoordinator`) delegates to `_runtime_subprocess_entrypoint()`, which handles all the TCP/process plumbing: + +1. Creates two TCP servers on `127.0.0.1` with random ports (comm + logs) +2. Creates a stderr socketpair +3. Calls `dag_parsing_cmd()` to get the command +4. Spawns the subprocess with `stdin=DEVNULL` (does NOT inherit fd 0) +5. Accepts TCP connections from the subprocess +6. Wraps fd 0 as `supervisor_comm` via `os.dup(0)` +7. Runs `_bridge()` — a raw byte forwarder between fd 0 and the TCP comm socket + +### Expected E2E Flow + +``` +Airflow Dag-Processor + │ + ▼ +DagFileProcessorProcess.start(path, bundle_name, bundle_path) + │ + ├─ _resolve_processor_target() + │ └─ iterates instances from [sdk] coordinators (airflow.cfg) + │ └─ first can_handle_dag_file() == True wins + │ + ▼ +WatchedSubprocess.start(target=coordinator.run_dag_parsing) + │ + [fork — child process gets fd 0 as Unix domain socket to supervisor] + │ + ▼ (in child) +Coordinator.run_dag_parsing(path, bundle_name, bundle_path) + │ + ▼ +BaseCoordinator._runtime_subprocess_entrypoint(DagParsingInfo) + │ + ├─ 1. Create TCP comm_server + logs_server on 127.0.0.1:random + ├─ 2. Create stderr socketpair + ├─ 3. Call dag_parsing_cmd() → get launch command + ├─ 4. Popen(cmd, stdin=DEVNULL, stderr=child_stderr) + ├─ 5. Accept TCP connections from the language runtime + ├─ 6. supervisor_comm = socket(fileno=os.dup(0)) + └─ 7. _bridge() — raw byte forwarding until process exits +``` + +### Expected Message Sequence + +Once the bridge is running, the Airflow supervisor and the language runtime communicate directly through the bridge (raw bytes, no re-encoding): + +``` +Airflow Supervisor Bridge Language Runtime + │ │ │ + ├── DagFileParseRequest ──────────┼──────────────────────►│ + │ [4-byte len][msgpack frame] │ raw byte forward │ + │ │ │ + │ │ ├── parse DAGs from + │ │ │ bundle/file + │ │ │ + │◄── DagFileParsingResult ────────┼───────────────────────┤ + │ [4-byte len][msgpack frame] │ raw byte forward │ + │ │ │ + │ │ └── exit(0) + │ │ + │ └── drain remaining bytes (5s deadline) + │ close all sockets +``` + +### DagFileParsingResult Format + +The language runtime must produce a `DagFileParsingResult` that matches Python Airflow's DagSerialization format exactly. The Airflow scheduler deserializes this into its internal model — any divergence causes parsing failures. + +**Envelope:** + +``` +{ + "type": "DagFileParsingResult", + "fileloc": "", + "serialized_dags": [ + { + "data": { + "__version": 3, + "dag": { } + } + }, + ... + ] +} +``` + +**Serialized DAG structure** (version 3): + +| Field | Type | Required | Description | +|---|---|---|---| +| `dag_id` | string | yes | Unique identifier | +| `fileloc` | string | yes | Source file path (can be empty) | +| `relative_fileloc` | string | yes | Relative source path (can be empty) | +| `timezone` | string | yes | Always `"UTC"` | +| `timetable` | `{__type, __var}` | yes | Schedule timetable (see below) | +| `tasks` | list | yes | Serialized task list | +| `dag_dependencies` | list | yes | Empty list for non-Python DAGs | +| `task_group` | map | yes | Flat root task group | +| `edge_info` | map | yes | Empty map | +| `params` | list | yes | DAG-level parameters | +| `description` | string | if set | | +| `start_date` | float (epoch) | if set | Unwrapped from `__type`/`__var` | +| `end_date` | float (epoch) | if set | Unwrapped from `__type`/`__var` | +| `tags` | list | if non-empty | Unwrapped from `__type`/`__var` | +| `catchup` | bool | if `true` | | +| `max_active_tasks` | int | if non-default | | +| `max_active_runs` | int | if non-default | | + +**Timetable encoding:** + +| Schedule | `__type` | `__var` | +|---|---|---| +| `null` | `airflow.timetables.simple.NullTimetable` | `{}` | +| `@once` | `airflow.timetables.simple.OnceTimetable` | `{}` | +| `@continuous` | `airflow.timetables.simple.ContinuousTimetable` | `{}` | +| cron expr | `airflow.timetables.trigger.CronTriggerTimetable` | `{expression, timezone, interval, run_immediately}` | + +**Task encoding:** + +``` +{ + "__type": "operator", + "__var": { + "task_id": "", + "task_type": "", + "_task_module": "", + "downstream_task_ids": [""] // only if non-empty + } +} +``` + +**Value type encoding** (for complex fields): + +| Type | Encoding | +|---|---| +| datetime | `{"__type": "datetime", "__var": }` | +| timedelta | `{"__type": "timedelta", "__var": }` | +| dict | `{"__type": "dict", "__var": {k: serialize(v), ...}}` | +| set | `{"__type": "set", "__var": [sorted_items]}` | +| list | `[serialize(item), ...]` (no wrapper) | +| primitives | pass through unchanged | + +**Non-decorated vs decorated fields:** Some fields (like `start_date`, `end_date`, `tags`) are "non-decorated" — they are serialized with `__type`/`__var` wrapping but then unwrapped to just the `__var` value before inclusion in the DAG dict. Other fields (like `default_args`, `access_control`) are "decorated" — they keep the `__type`/`__var` wrapper. This matches Python's `serialize_to_json` behavior. + +### What a Language Provider Must Implement + +For DAG parsing, a new language provider needs: + +1. **A `BaseCoordinator` subclass** with: + - `can_handle_dag_file()` — language-specific file detection (e.g., "is this a JAR?", "is this a .go file?") + - `dag_parsing_cmd()` — returns the command to launch the runtime + +2. **A runtime process** that: + - Accepts `--comm=host:port` and `--logs=host:port` CLI arguments + - Connects to both TCP addresses + - Reads a `DagFileParseRequest` msgpack frame from the comm channel + - Parses the DAGs from the bundle + - Serializes the result to DagSerialization v3 format + - Sends back a `DagFileParsingResult` msgpack frame + - Exits + +3. **Registration** as an entry in `[sdk] coordinators` in `airflow.cfg`, pointing `classpath` at the importable subclass under `airflow.sdk.coordinators.` + +### Java as a Concrete Example + +**JavaCoordinator:** + +The Java SDK implements all DAG-parsing contracts in a single `BaseCoordinator` subclass shipped as `apache-airflow-coordinators-java`: + +```python +# Distribution: apache-airflow-coordinators-java +# Module: airflow.sdk.coordinators.java.coordinator +class JavaCoordinator(BaseCoordinator): + def __init__(self, *, name, java_executable="java", jvm_args=None, jdk_home=None): + self.name = name + self.java_executable = java_executable + self.jvm_args = list(jvm_args or []) + self.jdk_home = jdk_home + + def can_handle_dag_file(self, bundle_name, path) -> bool: + # Returns True when path is a JAR with a Main-Class manifest entry + with contextlib.suppress(FileNotFoundError): + return find_main_class(Path(path)) is not None + return False + + def dag_parsing_cmd(self, *, dag_file_path, bundle_name, bundle_path, comm_addr, logs_addr): + main_class = find_main_class(Path(dag_file_path)) + return [ + self.java_executable, + *self.jvm_args, + "-classpath", + f"{bundle_path}/*", + main_class, + f"--comm={comm_addr}", + f"--logs={logs_addr}", + ] +``` + +`can_handle_dag_file()` checks that the file is a JAR with a `Main-Class` in its manifest. This ensures the coordinator only claims files it can actually handle. + +The classpath is `/*` — a wildcard that includes all JARs in the directory (the application JAR plus its dependencies). The `java_executable` and `jvm_args` come from the per-instance `kwargs` declared in `[sdk] coordinators`, so multiple instances (e.g., `jdk-11`, `jdk-17`) can launch different JVMs with different flags from the same class. + +No separate `JavaDagFileProcessor` class is needed — `BaseCoordinator` consolidates file detection, DAG parsing, and task execution into a single extension point. + +**Java SDK Bundle Process:** + +The Java bundle process (`Server.kt`) starts, connects to both TCP servers, and enters `CoordinatorComm.startProcessing()`. When it receives a `DagFileParseRequest`: + +``` +CoordinatorComm.handleIncoming(frame) + │ + ├── frame.body is DagFileParseRequest + │ file: String ← the path from the request + │ + ▼ +DagParser(request.file).parse(bundle) + │ + ├── Returns DagParsingResult(fileloc=file, dags=bundle.dags) + │ The DAGs were already loaded into the Bundle at startup + │ via BundleBuilder.getDags() + │ + ▼ +sendMessage(frame.id, result) + │ + ├── CoordinatorComm.encode(OutgoingFrame(id, result)) + │ ├── detects DagParsingResult type + │ └── calls result.serialize() ← Serde.kt + │ + ├── DagParsingResult.serialize() + │ ├── Wraps each DAG: {"data": {"__version": 3, "dag": dag.serialize(id)}} + │ ├── Dag.serialize() produces the full v3 format: + │ │ timetable, tasks, task_group, params, optional fields... + │ ├── Task.serialize() wraps as {"__type": "operator", "__var": {...}} + │ └── serializeValue() handles datetime/timedelta/dict/set encoding + │ + ├── TaskSdkFrames.encodeRequest(id, serializedMap) + │ ├── Converts map to msgpack: [id, body] + │ └── Returns byte array + │ + └── Writes [4-byte length prefix][msgpack payload] to comm channel + +shutDownRequested = true ← one-shot, process will exit +``` + +**Java SDK BundleBuilder Interface:** + +Bundle authors implement `BundleBuilder` to define their DAGs: + +```java +public class ExampleBundleBuilder implements BundleBuilder { + @Override + public List getDags() { + var dag = new Dag("java_example", null, "@daily"); + dag.addTask("extract", Extract.class, List.of()); + dag.addTask("transform", Transform.class, List.of("extract")); + dag.addTask("load", Load.class, List.of("transform")); + return List.of(dag); + } + + public static void main(String[] args) { + var bundle = new ExampleBundleBuilder().build(); + Server.create(args).serve(bundle); + } +} +``` + +The `Dag` class provides a fluent API: + +- `dagId`, `description`, `schedule` (cron or preset), `startDate`, `endDate`, and all standard Airflow DAG parameters +- `addTask(id, taskClass, dependsOn)` — registers a task and its upstream dependencies +- Dependencies are stored as a `dependants` map (parent → set of children), serialized as `downstream_task_ids` + +**Java SDK Serialization Compatibility:** + +The serialization in `Serde.kt` is validated against Python's output: + +```bash +# 1. Java generates serialized output +./gradlew sdk:test +# → writes validation/serialization/serialized_java.json + +# 2. Python generates the same DAGs +uv run validation/serialization/serialize_python.py \ + validation/serialization/test_dags.yaml \ + validation/serialization/serialized_python.json + +# 3. Field-by-field comparison +uv run validation/serialization/compare.py \ + validation/serialization/serialized_python.json \ + validation/serialization/serialized_java.json +``` + +Both share test cases defined in `test_dags.yaml`, ensuring the Java SDK produces byte-identical output to Python's `DagSerialization.serialize_dag()` for the same inputs. + +## Consequences + +- The DAG file processor can be extended to any language without modifying Airflow Core — only a `BaseCoordinator` subclass distributed as `apache-airflow-coordinators-` plus an entry in `[sdk] coordinators` is needed. +- The language runtime must produce exact DagSerialization v3 JSON, requiring cross-language validation infrastructure (e.g., `test_dags.yaml` + `compare.py`). +- The base class absorbs all TCP/process plumbing, so language providers only implement two methods for DAG parsing. +- The subprocess bridge adds latency and a process boundary; DAG parsing for non-Python files is inherently slower than in-process Python parsing. diff --git a/java-sdk/adr/0005-coordinator-packaging.md b/java-sdk/adr/0005-coordinator-packaging.md new file mode 100644 index 0000000000000..263b57a4fdee4 --- /dev/null +++ b/java-sdk/adr/0005-coordinator-packaging.md @@ -0,0 +1,160 @@ + + +# ADR-0005: Coordinator Packaging, Module Layout, and Registration + +## Status + +Proposed + +> **Note:** This ADR describes coordinator packaging as a standalone distribution separate +> from the Task SDK. The current plan is to ship coordinators as part of the Task SDK +> (`apache-airflow-task-sdk`); a separate distribution may be introduced later once the +> coordinator interface exits experimental status, but is not committed. This document is +> retained for reference if that split is revisited. Tracked operationally in +> [apache/airflow#66451](https://github.com/apache/airflow/issues/66451). + +## Context + +[ADR-0001](0001-java-sdk-airflow-integration.md) introduces a coordinator extension point. +Reviewers on PR #65958 raised three related but separable questions: + +1. **PyPI package name.** Should the Java coordinator ship as + `apache-airflow-providers-sdk-java` (consistent with every other provider) or under + `airflow.sdk.coordinators` as part of the Task SDK (recognizing that "language + coordinator" is a structurally new kind of distribution that does not behave like + operators/hooks/sensors)? +2. **Source-tree module layout.** Should it live under `providers/sdk/java/` alongside other + providers, or as a subpackage of the Task SDK? +3. **Discovery / registration mechanism.** Should coordinator classes be discovered through + the existing `ProvidersManager`, or through some other mechanism? + +A second concern, raised separately, is **runtime configuration**: a single `JavaCoordinator` +class is not enough to express "use JDK 11 for the legacy queue and JDK 17 for the modern +queue, with different `-Xmx` values." Class-only registration forces operators to subclass for +every variant or hardcode environment lookups, which the issue calls out explicitly: + +> How can I use different JDK version? How can I use different JVM arguments? We hardcoded the +> subprocess cmd … so users have to subclass another Coordinator to override the Java config. +> — [apache/airflow#66451](https://github.com/apache/airflow/issues/66451) + +## Decision + +### A. Distribution: included in the Task SDK + +The Java coordinator ships as part of the **Task SDK** (`apache-airflow-task-sdk`) and is +importable as `airflow.sdk.coordinators.java.JavaCoordinator`. This avoids extra packaging +infrastructure (separate release cadence, testing matrix, constraints files) while the +coordinator interface is still stabilising. New language coordinators (`go`, `typescript`, …) +follow the same model. + +A coordinator distribution exposes: + +- A `BaseCoordinator` subclass under `airflow.sdk.coordinators.`. +- No operators, hooks, sensors, triggers, or `provider.yaml`. + +### B. Module layout: namespace package under `airflow.sdk.coordinators` + +Each coordinator contributes a subpackage to the **namespace package** `airflow.sdk.coordinators`. +The Task SDK owns the namespace; individual language coordinators add +`airflow.sdk.coordinators.`. + +The Java coordinator therefore resolves as: + +```python +from airflow.utils.module_loading import import_string + +JavaCoordinator = import_string("airflow.sdk.coordinators.java.JavaCoordinator") +``` + +Both Airflow Core (DAG processor) and the Task SDK (task runner) import coordinators by this +path. The namespace package layout means the physical distribution can change in the future +without altering import paths or user configuration. + +### C. Discovery via `[sdk] coordinators` (Airflow configuration) + +Coordinators are **not** discovered through `ProvidersManager` / +`ProvidersManagerTaskRuntime`, and there is no `coordinators` key in `provider.yaml`. They are +registered as named instances in `airflow.cfg`: + +```ini +[sdk] +coordinators = { + "jdk-11": { + "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", + "kwargs": { + "java_executable": "/usr/lib/jvm/java-11-openjdk-amd64/bin/java", + "jvm_args": ["-Xmx512m"], + "jars_root": ["/files/legacy/lib"] + } + }, + "jdk-17": { + "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", + "kwargs": { + "java_executable": "/usr/lib/jvm/java-17-openjdk-amd64/bin/java", + "jvm_args": ["-Xmx1024m", "-Xms256m"], + "jars_root": ["/files/new/lib"] + } + } +} + +queue_to_coordinator = {"legacy-java-queue": "jdk-11", "modern-java-queue": "jdk-17"} +``` + +`[sdk] coordinators` is a JSON object: the key is the coordinator's name (used as the routing +target in `[sdk] queue_to_coordinator`), and the value supplies `classpath` and free-form +`kwargs` passed to the constructor. Two entries with the same `classpath` (e.g., both +`JavaCoordinator`) but different keys and `kwargs` are independent instances — this is how +JDK 11 and JDK 17 tasks run on the same worker without subclassing. + +### Why not `provider.yaml` / `ProvidersManager`? + +Coordinators are not providers in the Airflow sense: + +- They expose no operators / hooks / sensors / triggers. +- They are consumed by both Airflow Core (in the DAG processor) **and** the Task SDK (in the + task runner). The provider system is not designed to be loaded from inside a worker subprocess + that intentionally has no Airflow-Core import. +- They need **per-instance** runtime configuration (interpreter path, JVM flags, …). + `provider.yaml` registers classes, not instances, and bolting kwargs onto provider entries + would distort the provider data model. +- A coordinator is the only thing in this distribution; there is no benefit to sharing the + provider's discoverability surface (`airflow providers list`, etc.). On the contrary, listing + `apache-airflow-providers-sdk-java` next to AWS/GCP providers is misleading for users. + +Putting the registry in `airflow.cfg` keeps the data model honest (instances, with their kwargs) +and makes the per-host opt-in (install + config-edit) explicit rather than implicit +(install-implies-active). + +## Consequences + +- The Java coordinator ships **inside the Task SDK**; the namespace package layout ensures that + if packaging arrangements change in the future, no changes to user configuration or Airflow + Core are required. +- **`airflow.sdk.coordinators`** is a namespace package owned by the Task SDK; language + coordinator modules contribute subpackages to it. Multiple coordinator modules can be + installed side by side without colliding. +- **`[sdk] coordinators`** carries instance-level configuration; **`[sdk] queue_to_coordinator`** + carries queue → instance routing. +- Multiple instances of the same coordinator class (e.g., `jdk-11` and `jdk-17`) can be + registered with different `kwargs` and bound to different queues — solving the multi-JDK and + JVM-flag use cases raised in + [apache/airflow#66451](https://github.com/apache/airflow/issues/66451) without subclassing. +- The provider registry no longer shows coordinators, removing the "Java appears, Go does not" + asymmetry that earlier drafts of this ADR flagged as a transitional UX wart. diff --git a/java-sdk/build.gradle.kts b/java-sdk/build.gradle.kts new file mode 100644 index 0000000000000..950f09db5422f --- /dev/null +++ b/java-sdk/build.gradle.kts @@ -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. + */ + +import com.diffplug.gradle.spotless.SpotlessExtension +import org.jetbrains.kotlin.gradle.dsl.JvmTarget + +plugins { + kotlin("jvm") version "2.3.0" + id("com.diffplug.spotless") version "7.2.1" // Last version supporting JDK 11. + id("org.jlleitschuh.gradle.ktlint") version "14.0.1" +} + +allprojects { + apply(plugin = "com.diffplug.spotless") + apply(plugin = "org.jetbrains.kotlin.jvm") + apply(plugin = "org.jlleitschuh.gradle.ktlint") + + repositories { mavenCentral() } + + java { + toolchain { + languageVersion.set(JavaLanguageVersion.of(11)) + } + sourceCompatibility = JavaVersion.VERSION_11 + } + kotlin { + compilerOptions { + // If this is changed, also change "Setup Java" in codeql-analysis.yml. + jvmTarget = JvmTarget.JVM_11 + } + } + + configure { + java { + target("**/*.java") + googleJavaFormat().formatJavadoc(false) + trimTrailingWhitespace() + endWithNewline() + } + } +} diff --git a/java-sdk/dags/java_examples.py b/java-sdk/dags/java_examples.py new file mode 100644 index 0000000000000..a85bbce305609 --- /dev/null +++ b/java-sdk/dags/java_examples.py @@ -0,0 +1,60 @@ +# 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. + +from __future__ import annotations + +from airflow.sdk import dag, task + + +@task() +def python_task_1(): + print("python_task_1") + print("Push Python Task 'python_task_1' XCom:") + return "value_from_python_task_1" + + +@task.stub(queue="java") +def extract(): ... + + +@task.stub(queue="java") +def transform(): ... + + +@task() +def python_task_2(transformed): + print("python_task_2") + print("Pull Java Task 'transform' XCom:") + print(transformed) + + +@dag(dag_id="java_interface_example") +def java_interface_example(): + transformed = transform() + python_task_1() >> extract() >> transformed + python_task_2(transformed) + + +@dag(dag_id="java_annotation_example") +def java_annotation_example(): + transformed = transform() + python_task_1() >> extract() >> transformed + python_task_2(transformed) + + +java_interface_example() +java_annotation_example() diff --git a/java-sdk/example/build.gradle.kts b/java-sdk/example/build.gradle.kts new file mode 100644 index 0000000000000..5fd1a8f3b82ed --- /dev/null +++ b/java-sdk/example/build.gradle.kts @@ -0,0 +1,48 @@ +/* + * 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. + */ + +plugins { + application +} + +dependencies { + annotationProcessor(project(":sdk")) + implementation(project(":sdk")) + implementation("org.slf4j:slf4j-simple:2.0.17") +} + +sourceSets { + main { + java.srcDir("src/java") + } +} + +application { + mainClass = "org.apache.airflow.example.ExampleBundleBuilder" +} + +tasks.withType { + manifest { + attributes( + "Main-Class" to application.mainClass.get(), + "Implementation-Title" to "Example Java bundle", + "Implementation-Version" to "1", + ) + } +} diff --git a/java-sdk/example/src/java/org/apache/airflow/example/AnnotationExample.java b/java-sdk/example/src/java/org/apache/airflow/example/AnnotationExample.java new file mode 100644 index 0000000000000..010bb90d9440a --- /dev/null +++ b/java-sdk/example/src/java/org/apache/airflow/example/AnnotationExample.java @@ -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. + */ + +package org.apache.airflow.example; + +import org.apache.airflow.sdk.*; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Date; + +@SuppressWarnings("DuplicatedCode") +@Builder.Dag(id = "java_annotation_example") +public class AnnotationExample { + private static final Logger logger = LoggerFactory.getLogger(AnnotationExample.class); + + @Builder.Task(id = "extract") + public long extractValue(Client client) throws InterruptedException { + logger.info("Hello from task"); + + var pythonXcom = client.getXCom("python_task_1"); + logger.info("Got XCom from Python Task 'python_task_1' {}", pythonXcom); + + var connection = client.getConnection("test_http"); + logger.info("Got con {}", connection); + + for (var i = 0; i < 3; i++) { + logger.info("Beep {}, next time will be {}", i, new Date()); + Thread.sleep(2 * 1000); + } + + logger.info("Goodbye from task"); + return new Date().getTime(); + } + + @Builder.Task(id = "transform") + public long transformValue(Client client, @Builder.XCom(task = "extract") long extracted) { + logger.info("Got XCom from 'extract' {}", extracted); + + var variable = client.getVariable("my_variable"); + logger.info("Got variable {}", variable); + + logger.info("Push XCom to python task 2"); + return new Date().getTime(); + } + + @Builder.Task + public void load(@Builder.XCom(task = "transform") long transformed) { + logger.info("Got XCom from 'transform' {}", transformed); + throw new RuntimeException("I failed"); + } +} diff --git a/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java new file mode 100644 index 0000000000000..0aa729d00030e --- /dev/null +++ b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java @@ -0,0 +1,37 @@ +/* + * 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 org.apache.airflow.example; + +import org.apache.airflow.sdk.*; +import org.jetbrains.annotations.NotNull; +import java.util.List; + +public class ExampleBundleBuilder implements BundleBuilder { + @NotNull + @Override + public Iterable getDags() { + return List.of(InterfaceExampleBuilder.build(), AnnotationExampleBuilder.build()); + } + + public static void main(String[] args) { + var bundle = new ExampleBundleBuilder().build(); + Server.create(args).serve(bundle); + } +} diff --git a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java new file mode 100644 index 0000000000000..7853927b66a50 --- /dev/null +++ b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java @@ -0,0 +1,80 @@ +/* + * 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 org.apache.airflow.example; + +import java.util.Date; +import org.apache.airflow.sdk.*; +import org.jetbrains.annotations.NotNull; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +@SuppressWarnings("DuplicatedCode") +public class InterfaceExampleBuilder { + private static final Logger logger = LoggerFactory.getLogger(InterfaceExampleBuilder.class); + + public static class Extract implements Task { + public void execute(@NotNull Context context, Client client) throws Exception { + logger.info("Hello from task"); + + var pythonInput = client.getXCom("python_task_1"); + logger.info("Got XCom from Python Task 'python_task_1' {}", pythonInput); + + var connection = client.getConnection("test_http"); + logger.info("Got con {}", connection); + + for (var i = 0; i < 3; i++) { + logger.info("Beep {}, next time will be {}", i, new Date()); + Thread.sleep(2 * 1000); + } + + client.setXCom(new Date().getTime()); + logger.info("Goodbye from task"); + } + } + + public static class Transform implements Task { + public void execute(@NotNull Context context, Client client) { + var extracted = client.getXCom("extract"); + logger.info("Got XCom from 'extract' {}", extracted); + + var variable = client.getVariable("my_variable"); + logger.info("Got variable {}", variable); + + logger.info("Push XCom to python task 2"); + client.setXCom(new Date().getTime()); + } + } + + public static class Load implements Task { + public void execute(@NotNull Context context, Client client) { + var transformed = client.getXCom("transform"); + logger.info("Got XCom from 'transform' {}", transformed); + throw new RuntimeException("I failed"); + } + } + + public static Dag build() { + var dag = new Dag("java_interface_example"); + dag.addTask("extract", Extract.class); + dag.addTask("transform", Transform.class); + dag.addTask("load", Load.class); + return dag; + } +} diff --git a/java-sdk/gradle.properties b/java-sdk/gradle.properties new file mode 100644 index 0000000000000..cab8c837412af --- /dev/null +++ b/java-sdk/gradle.properties @@ -0,0 +1,20 @@ +# 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. + +org.gradle.configuration-cache=true + +airflowSupervisorSchemaVersion=2026-06-16 diff --git a/java-sdk/gradle/libs.versions.toml b/java-sdk/gradle/libs.versions.toml new file mode 100644 index 0000000000000..a0ac505d527fe --- /dev/null +++ b/java-sdk/gradle/libs.versions.toml @@ -0,0 +1,19 @@ +# 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. + +# This file was generated by the Gradle 'init' task. +# https://docs.gradle.org/current/userguide/platforms.html#sub::toml-dependencies-format diff --git a/java-sdk/gradle/wrapper/gradle-wrapper.jar b/java-sdk/gradle/wrapper/gradle-wrapper.jar new file mode 100644 index 0000000000000..1b33c55baabb5 Binary files /dev/null and b/java-sdk/gradle/wrapper/gradle-wrapper.jar differ diff --git a/java-sdk/gradle/wrapper/gradle-wrapper.properties b/java-sdk/gradle/wrapper/gradle-wrapper.properties new file mode 100644 index 0000000000000..ed9937b14b9fe --- /dev/null +++ b/java-sdk/gradle/wrapper/gradle-wrapper.properties @@ -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. + +distributionBase=GRADLE_USER_HOME +distributionPath=wrapper/dists +distributionUrl=https\://services.gradle.org/distributions/gradle-8.14.4-bin.zip +networkTimeout=10000 +validateDistributionUrl=true +zipStoreBase=GRADLE_USER_HOME +zipStorePath=wrapper/dists diff --git a/java-sdk/gradlew b/java-sdk/gradlew new file mode 100755 index 0000000000000..342f59614db8e --- /dev/null +++ b/java-sdk/gradlew @@ -0,0 +1,233 @@ +#!/bin/sh + +############################################################################## +# +# Gradle start up script for POSIX generated by Gradle. +# +# Important for running: +# +# (1) You need a POSIX-compliant shell to run this script. If your /bin/sh is +# noncompliant, but you have some other compliant shell such as ksh or +# bash, then to run this script, type that shell name before the whole +# command line, like: +# +# ksh Gradle +# +# Busybox and similar reduced shells will NOT work, because this script +# requires all of these POSIX shell features: +# * functions; +# * expansions «$var», «${var}», «${var:-default}», «${var+SET}», +# «${var#prefix}», «${var%suffix}», and «$( cmd )»; +# * compound commands having a testable exit status, especially «case»; +# * various built-in commands including «command», «set», and «ulimit». +# +# Important for patching: +# +# (2) This script targets any POSIX shell, so it avoids extensions provided +# by Bash, Ksh, etc; in particular arrays are avoided. +# +# The "traditional" practice of packing multiple parameters into a +# space-separated string is a well documented source of bugs and security +# problems, so this is (mostly) avoided, by progressively accumulating +# options in "$@", and eventually passing that to Java. +# +# Where the inherited environment variables (DEFAULT_JVM_OPTS, JAVA_OPTS, +# and GRADLE_OPTS) rely on word-splitting, this is performed explicitly; +# see the in-line comments for details. +# +# There are tweaks for specific operating systems such as AIX, CygWin, +# Darwin, MinGW, and NonStop. +# +# (3) This script is generated from the Groovy template +# https://github.com/gradle/gradle/blob/HEAD/platforms/jvm/plugins-application/src/main/resources/org/gradle/api/internal/plugins/unixStartScript.txt +# within the Gradle project. +# +# You can find Gradle at https://github.com/gradle/gradle/. +# +############################################################################## + +# Attempt to set APP_HOME + +# Resolve links: $0 may be a link +app_path=$0 + +# Need this for daisy-chained symlinks. +while + APP_HOME=${app_path%"${app_path##*/}"} # leaves a trailing /; empty if no leading path + [ -h "$app_path" ] +do + ls=$( ls -ld "$app_path" ) + link=${ls#*' -> '} + case $link in #( + /*) app_path=$link ;; #( + *) app_path=$APP_HOME$link ;; + esac +done + +# This is normally unused +# shellcheck disable=SC2034 +APP_BASE_NAME=${0##*/} +# Discard cd standard output in case $CDPATH is set (https://github.com/gradle/gradle/issues/25036) +APP_HOME=$( cd -P "${APP_HOME:-./}" > /dev/null && printf '%s\n' "$PWD" ) || exit + +# Use the maximum available, or set MAX_FD != -1 to use that value. +MAX_FD=maximum + +warn () { + echo "$*" +} >&2 + +die () { + echo + echo "$*" + echo + exit 1 +} >&2 + +# OS specific support (must be 'true' or 'false'). +cygwin=false +msys=false +darwin=false +nonstop=false +case "$( uname )" in #( + CYGWIN* ) cygwin=true ;; #( + Darwin* ) darwin=true ;; #( + MSYS* | MINGW* ) msys=true ;; #( + NONSTOP* ) nonstop=true ;; +esac + +CLASSPATH="\\\"\\\"" + + +# Determine the Java command to use to start the JVM. +if [ -n "$JAVA_HOME" ] ; then + if [ -x "$JAVA_HOME/jre/sh/java" ] ; then + # IBM's JDK on AIX uses strange locations for the executables + JAVACMD=$JAVA_HOME/jre/sh/java + else + JAVACMD=$JAVA_HOME/bin/java + fi + if [ ! -x "$JAVACMD" ] ; then + die "ERROR: JAVA_HOME is set to an invalid directory: $JAVA_HOME + +Please set the JAVA_HOME variable in your environment to match the +location of your Java installation." + fi +else + JAVACMD=java + if ! command -v java >/dev/null 2>&1 + then + die "ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. + +Please set the JAVA_HOME variable in your environment to match the +location of your Java installation." + fi +fi + +# Increase the maximum file descriptors if we can. +if ! "$cygwin" && ! "$darwin" && ! "$nonstop" ; then + case $MAX_FD in #( + max*) + # In POSIX sh, ulimit -H is undefined. That's why the result is checked to see if it worked. + # shellcheck disable=SC2039,SC3045 + MAX_FD=$( ulimit -H -n ) || + warn "Could not query maximum file descriptor limit" + esac + case $MAX_FD in #( + '' | soft) :;; #( + *) + # In POSIX sh, ulimit -n is undefined. That's why the result is checked to see if it worked. + # shellcheck disable=SC2039,SC3045 + ulimit -n "$MAX_FD" || + warn "Could not set maximum file descriptor limit to $MAX_FD" + esac +fi + +# Collect all arguments for the java command, stacking in reverse order: +# * args from the command line +# * the main class name +# * -classpath +# * -D...appname settings +# * --module-path (only if needed) +# * DEFAULT_JVM_OPTS, JAVA_OPTS, and GRADLE_OPTS environment variables. + +# For Cygwin or MSYS, switch paths to Windows format before running java +if "$cygwin" || "$msys" ; then + APP_HOME=$( cygpath --path --mixed "$APP_HOME" ) + CLASSPATH=$( cygpath --path --mixed "$CLASSPATH" ) + + JAVACMD=$( cygpath --unix "$JAVACMD" ) + + # Now convert the arguments - kludge to limit ourselves to /bin/sh + for arg do + if + case $arg in #( + -*) false ;; # don't mess with options #( + /?*) t=${arg#/} t=/${t%%/*} # looks like a POSIX filepath + [ -e "$t" ] ;; #( + *) false ;; + esac + then + arg=$( cygpath --path --ignore --mixed "$arg" ) + fi + # Roll the args list around exactly as many times as the number of + # args, so each arg winds up back in the position where it started, but + # possibly modified. + # + # NB: a `for` loop captures its iteration list before it begins, so + # changing the positional parameters here affects neither the number of + # iterations, nor the values presented in `arg`. + shift # remove old arg + set -- "$@" "$arg" # push replacement arg + done +fi + + +# Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script. +DEFAULT_JVM_OPTS='"-Xmx64m" "-Xms64m"' + +# Collect all arguments for the java command: +# * DEFAULT_JVM_OPTS, JAVA_OPTS, and optsEnvironmentVar are not allowed to contain shell fragments, +# and any embedded shellness will be escaped. +# * For example: A user cannot expect ${Hostname} to be expanded, as it is an environment variable and will be +# treated as '${Hostname}' itself on the command line. + +set -- \ + "-Dorg.gradle.appname=$APP_BASE_NAME" \ + -classpath "$CLASSPATH" \ + -jar "$APP_HOME/gradle/wrapper/gradle-wrapper.jar" \ + "$@" + +# Stop when "xargs" is not available. +if ! command -v xargs >/dev/null 2>&1 +then + die "xargs is not available" +fi + +# Use "xargs" to parse quoted args. +# +# With -n1 it outputs one arg per line, with the quotes and backslashes removed. +# +# In Bash we could simply go: +# +# readarray ARGS < <( xargs -n1 <<<"$var" ) && +# set -- "${ARGS[@]}" "$@" +# +# but POSIX shell has neither arrays nor command substitution, so instead we +# post-process each arg (as a line of input to sed) to backslash-escape any +# character that might be a shell metacharacter, then use eval to reverse +# that process (while maintaining the separation between arguments), and wrap +# the whole thing up as a single "set" statement. +# +# This will of course break if any of these variables contains a newline or +# an unmatched quote. +# + +eval "set -- $( + printf '%s\n' "$DEFAULT_JVM_OPTS $JAVA_OPTS $GRADLE_OPTS" | + xargs -n1 | + sed ' s~[^-[:alnum:]+,./:=@_]~\\&~g; ' | + tr '\n' ' ' + )" '"$@"' + +exec "$JAVACMD" "$@" diff --git a/java-sdk/gradlew.bat b/java-sdk/gradlew.bat new file mode 100644 index 0000000000000..b656b6ae26efd --- /dev/null +++ b/java-sdk/gradlew.bat @@ -0,0 +1,76 @@ +@if "%DEBUG%"=="" @echo off +@rem ########################################################################## +@rem +@rem Gradle startup script for Windows +@rem +@rem ########################################################################## + +@rem Set local scope for the variables with windows NT shell +if "%OS%"=="Windows_NT" setlocal + +set DIRNAME=%~dp0 +if "%DIRNAME%"=="" set DIRNAME=. +@rem This is normally unused +set APP_BASE_NAME=%~n0 +set APP_HOME=%DIRNAME% + +@rem Resolve any "." and ".." in APP_HOME to make it shorter. +for %%i in ("%APP_HOME%") do set APP_HOME=%%~fi + +@rem Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script. +set DEFAULT_JVM_OPTS="-Xmx64m" "-Xms64m" + +@rem Find java.exe +if defined JAVA_HOME goto findJavaFromJavaHome + +set JAVA_EXE=java.exe +%JAVA_EXE% -version >NUL 2>&1 +if %ERRORLEVEL% equ 0 goto execute + +echo. 1>&2 +echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. 1>&2 +echo. 1>&2 +echo Please set the JAVA_HOME variable in your environment to match the 1>&2 +echo location of your Java installation. 1>&2 + +goto fail + +:findJavaFromJavaHome +set JAVA_HOME=%JAVA_HOME:"=% +set JAVA_EXE=%JAVA_HOME%/bin/java.exe + +if exist "%JAVA_EXE%" goto execute + +echo. 1>&2 +echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME% 1>&2 +echo. 1>&2 +echo Please set the JAVA_HOME variable in your environment to match the 1>&2 +echo location of your Java installation. 1>&2 + +goto fail + +:execute +@rem Setup the command line + +set CLASSPATH= + + +@rem Execute Gradle +"%JAVA_EXE%" %DEFAULT_JVM_OPTS% %JAVA_OPTS% %GRADLE_OPTS% "-Dorg.gradle.appname=%APP_BASE_NAME%" -classpath "%CLASSPATH%" -jar "%APP_HOME%\gradle\wrapper\gradle-wrapper.jar" %* + +:end +@rem End local scope for the variables with windows NT shell +if %ERRORLEVEL% equ 0 goto mainEnd + +:fail +rem Set variable GRADLE_EXIT_CONSOLE if you need the _script_ return code instead of +rem the _cmd.exe /c_ return code! +set EXIT_CODE=%ERRORLEVEL% +if %EXIT_CODE% equ 0 set EXIT_CODE=1 +if not ""=="%GRADLE_EXIT_CONSOLE%" exit %EXIT_CODE% +exit /b %EXIT_CODE% + +:mainEnd +if "%OS%"=="Windows_NT" endlocal + +:omega diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts new file mode 100644 index 0000000000000..ce56e699fe7b0 --- /dev/null +++ b/java-sdk/sdk/build.gradle.kts @@ -0,0 +1,218 @@ +/* + * 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. + */ + +buildscript { + repositories { + mavenCentral() + } +} + +val airflowSupervisorSchemaVersion: String by project + +plugins { + kotlin("plugin.serialization") version "2.3.0" + id("org.jsonschema2pojo") version "1.2.2" +} + +// TODO: Use a hosted file instead. +val schemaInput = rootProject.file("../task-sdk/src/airflow/sdk/execution_time/schema/schema.json") +val pointersDir = layout.buildDirectory.dir("schema-pointers/main") +val jsonSchemaPackage = "org.apache.airflow.sdk.execution.comm" +val discriminatorDir = layout.buildDirectory.dir("generated-resources/main/src/main/kotlin") + +dependencies { + compileOnly("com.github.spotbugs:spotbugs-annotations:4.9.8") + compileOnly("javax.annotation:javax.annotation-api:1.3.2") + + implementation("com.fasterxml.jackson.core:jackson-annotations:2.21") + implementation("com.fasterxml.jackson.core:jackson-core:2.21.1") + implementation("com.fasterxml.jackson.core:jackson-databind:2.21.0") + implementation("com.fasterxml.jackson.dataformat:jackson-dataformat-yaml:2.21.0") + implementation("com.fasterxml.jackson.datatype:jackson-datatype-jsr310:2.21.0") + implementation("com.squareup:javapoet:1.13.0") + implementation("com.xenomachina:kotlin-argparser:2.0.7") + implementation("io.ktor:ktor-network:3.3.3") + implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.10.2") + implementation("org.jetbrains.kotlinx:kotlinx-datetime:0.7.1") + implementation("org.jetbrains.kotlinx:kotlinx-serialization-json:1.10.0") + implementation("org.msgpack:msgpack-core:0.9.11") + implementation("org.msgpack:jackson-dataformat-msgpack:0.9.11") + + testImplementation(kotlin("test")) + testImplementation("com.google.testing.compile:compile-testing:0.23.0") + testImplementation("com.squareup.okhttp3:mockwebserver:4.12.0") +} + +// jsonSchema2Pojo does not accept the single JSON Schema file directly. +// It needs a list of schema files, each containing a "$ref" pointer to +// a $def. This task walks over all $ref items in the Supervisor Schema +// file and generates one JSON file with $ref for each one. +abstract class GeneratePointersTask : DefaultTask() { + @get:InputFile + abstract val schemaFile: RegularFileProperty + + @get:OutputDirectory + abstract val targetDirectory: DirectoryProperty + + @TaskAction + fun generate() { + val srcFile = schemaFile.get().asFile + val outDir = targetDirectory.get().asFile.also { it.mkdirs() } + + srcFile.copyTo(outDir.resolve(srcFile.name), overwrite = true) + + com.fasterxml.jackson.databind + .ObjectMapper() + .readTree(srcFile) + .path("\$defs") + .fieldNames() + .forEach { type -> + outDir + .resolve("$type.json") + .writeText("""{"${"$"}ref": "${srcFile.name}#/${"$"}defs/$type"}""" + "\n") + } + } +} + +// Generate a name->class mapping of known jsonSchema2Pojo models. +// This is needed for type discrimination in the MessagePack decoder. +abstract class GenerateDiscriminatorTask : DefaultTask() { + @get:Input + abstract val modelPackage: Property + + @get:InputFile + abstract val schemaFile: RegularFileProperty + + @get:OutputDirectory + abstract val targetDirectory: DirectoryProperty + + @TaskAction + fun generate() { + data class Entry( + val wireType: String, + val className: String, + ) + + val entries = + buildList { + com.fasterxml.jackson.databind + .ObjectMapper() + .readTree(schemaFile.get().asFile) + .path("\$defs") + .fields() + .forEach { (className, def) -> + val constNode = def.path("properties").path("type").path("const") + if (!constNode.isMissingNode && !constNode.isNull) { + add(Entry(constNode.asText(), className)) + } + } + }.sortedBy { it.className } + + val outDir = + targetDirectory + .get() + .asFile + .resolve("org/apache/airflow/sdk/execution/comm") + .also { it.mkdirs() } + + outDir.resolve("Discriminator.kt").writeText( + buildString { + appendLine("package ${modelPackage.get()}") + appendLine() + appendLine("// Maps every wire `type` discriminator string to its generated model class.") + appendLine("// Generated from the Supervisor Schema; do not edit by hand.") + appendLine("internal object Discriminator {") + appendLine(" val types: Map> = mapOf(") + entries.forEach { appendLine(" \"${it.wireType}\" to ${it.className}::class.java,") } + appendLine(" )") + appendLine("}") + }, + ) + } +} + +tasks.register("generateDiscriminator") { + description = "Generate Discriminator to wire type strings to model classes" + schemaFile = layout.file(provider { schemaInput }) + modelPackage = jsonSchemaPackage + targetDirectory = discriminatorDir +} + +tasks.register("generatePointers") { + description = "Generate pointer files for jsonSchema2Pojo" + schemaFile = layout.file(provider { schemaInput }) + targetDirectory = pointersDir +} + +jsonSchema2Pojo { + setSource(listOf(pointersDir.get().asFile)) + targetPackage = jsonSchemaPackage + targetDirectory = + layout.buildDirectory + .dir("generate-resources/main/src/main/java") + .get() + .asFile + setAnnotationStyle("jackson") + dateTimeType = "java.time.OffsetDateTime" + generateBuilders = false + includeAdditionalProperties = false + includeConstructors = false + includeHashcodeAndEquals = true + includeJsr305Annotations = true + includeToString = true + initializeCollections = true + removeOldOutput = true + useTitleAsClassname = true +} + +sourceSets { + main { + java.srcDir(layout.buildDirectory.dir("generate-resources/main/src/main/java")) + kotlin.srcDir(discriminatorDir) + } +} + +tasks.named("generateJsonSchema2Pojo") { + dependsOn("generatePointers") +} + +tasks.named("compileJava") { + dependsOn("generateJsonSchema2Pojo") +} + +tasks.named("compileKotlin") { + dependsOn("generateJsonSchema2Pojo") + dependsOn("generateDiscriminator") +} + +tasks.named("runKtlintCheckOverMainSourceSet") { + dependsOn("generateJsonSchema2Pojo") +} + +tasks.withType { + manifest { + attributes( + "Airflow-Supervisor-Schema-Version" to airflowSupervisorSchemaVersion, + ) + } +} + +tasks.withType { + useJUnitPlatform() +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt new file mode 100644 index 0000000000000..19d977d4aba76 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt @@ -0,0 +1,273 @@ +/* + * 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. + */ + +@file:Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN") + +package org.apache.airflow.sdk + +import com.squareup.javapoet.ClassName +import com.squareup.javapoet.JavaFile +import com.squareup.javapoet.MethodSpec +import com.squareup.javapoet.TypeName +import com.squareup.javapoet.TypeSpec +import javax.annotation.processing.AbstractProcessor +import javax.annotation.processing.ProcessingEnvironment +import javax.annotation.processing.RoundEnvironment +import javax.annotation.processing.SupportedAnnotationTypes +import javax.annotation.processing.SupportedSourceVersion +import javax.lang.model.SourceVersion +import javax.lang.model.element.ExecutableElement +import javax.lang.model.element.Modifier +import javax.lang.model.element.TypeElement +import javax.lang.model.type.TypeKind +import javax.lang.model.type.TypeMirror +import javax.tools.Diagnostic + +/** + * Container for the annotation-based Dag-authoring API. + * + *

This class is not instantiated directly. Its nested annotations drive the + * [BuilderProcessor] annotation processor, which generates a {@code *Builder} class + * for each class annotated with [Dag]. + * + *

Example: + *

{@code
+ * @Builder.Dag(id = "my_pipeline")
+ * public class MyPipeline {
+ *
+ *     @Builder.Task(id = "extract")
+ *     public long extract(Client client) { ... }
+ *
+ *     @Builder.Task(id = "transform")
+ *     public long transform(Client client,
+ *                           @Builder.XCom(task = "extract") long extracted) { ... }
+ * }
+ * }
+ * + *

The processor generates {@code MyPipelineBuilder.build()}, which returns a + * fully-wired [Dag] ready to add to a [Bundle]. + */ +class Builder internal constructor() { + /** + * Annotation to automate a Dag-builder pattern. + * + * When applied on a class Foo, this generates a FooBuilder class with a static build method + * to create the Dag structure automatically. + * + * @param id Override the Dag ID. If empty or not provided, the annotated class's name is used by default. + * @param to Name of the Dag-builder class. If empty or not provided, use the annotated class name + "Builder". + */ + @Target(AnnotationTarget.CLASS) + @MustBeDocumented + annotation class Dag( + val id: String = "", + val to: String = "", + ) + + /** + * Annotation to automate task definition in a Dag-builder pattern. + * + * @param id Override the task ID. If empty or not provided, the annotated function's name is used by default. + */ + @Target(AnnotationTarget.FUNCTION) + @MustBeDocumented + annotation class Task( + val id: String = "", + ) + + /** + * Annotation to mark a task definition's method parameter as an XCom input. + * + * @param task The task ID to pull. If empty or not given, the annotated parameter's name is used by default. + * @param key The XCom key to pull. Defaults to the task's return value. + */ + @Target(AnnotationTarget.VALUE_PARAMETER) + @MustBeDocumented + annotation class XCom( + val task: String = "", + val key: String = Client.XCOM_RETURN_KEY, + ) +} + +/** + * @suppress + * + * Annotation processor for [Builder.Dag]. Registered as a standard javac processor via + * {@code META-INF/services/javax.annotation.processing.Processor}; not intended to be + * instantiated or referenced directly. + * + *

For each class annotated with [Builder.Dag], generates a {@code *Builder} class + * containing: + *

    + *
  • One inner class per [Builder.Task]-annotated method, implementing [Task].
  • + *
  • A static {@code build()} method that constructs the [Dag] and registers those + * inner classes as tasks.
  • + *
+ * + *

[Builder.XCom]-annotated parameters are resolved via {@code client.getXCom} in the + * generated {@code execute} body, with the result cast to the parameter's declared type. + * Non-{@code void} return values are forwarded to {@code client.setXCom}. + */ +@SupportedAnnotationTypes("org.apache.airflow.sdk.Builder.Dag") +@SupportedSourceVersion(SourceVersion.RELEASE_11) +class BuilderProcessor : AbstractProcessor() { + override fun process( + annotations: Set, + roundEnv: RoundEnvironment, + ): Boolean { + if (annotations.isEmpty()) return false + roundEnv.getElementsAnnotatedWith(Builder.Dag::class.java).filterIsInstance().forEach { el -> + with(processingEnv) { + runCatching { + JavaFile + .builder( + elementUtils.getPackageOf(el).qualifiedName.toString(), + buildDag(el), + ).build() + .writeTo(filer) + }.onFailure { e -> + messager.printMessage( + Diagnostic.Kind.ERROR, + e.message ?: "Unknown error", + el, + ) + } + } + } + return true + } + + private fun buildDag(el: TypeElement): TypeSpec { + val ann = el.getAnnotation(Builder.Dag::class.java)!! + + val builderClass = + TypeSpec + .classBuilder(ann.to.ifBlank { "${el.simpleName}Builder" }) + .addModifiers(Modifier.PUBLIC, Modifier.FINAL) + + val buildMethod = + MethodSpec + .methodBuilder("build") + .addModifiers(Modifier.PUBLIC, Modifier.STATIC) + .returns(ClassName.get(Dag::class.java)) + .addStatement($$"var dag = new $T($S)", ClassName.get(Dag::class.java), ann.id.ifBlank { el.simpleName }) + + for (inner in el.enclosedElements) { + if (inner !is ExecutableElement) continue + if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}") + + val ann = inner.getAnnotation(Builder.Task::class.java) ?: continue + val innerName = inner.simpleName.toString().replaceFirstChar(Char::uppercase) + + val task = buildTask(innerName, inner, el) + builderClass.addType(task.spec) + + buildMethod.addStatement( + $$"dag.addTask($S, $L.class)", + ann.id.ifBlank { inner.simpleName }, + innerName, + ) + } + + buildMethod.addStatement("return dag") + builderClass.addMethod(buildMethod.build()) + return builderClass.build() + } + + private fun buildTask( + name: String, + inner: ExecutableElement, + parent: TypeElement, + ): BuildTaskResult { + val clientType = ClassName.get(Client::class.java) + val contextType = ClassName.get(Context::class.java) + + val executeSpec = + MethodSpec + .methodBuilder("execute") + .addAnnotation(Override::class.java) + .addModifiers(Modifier.PUBLIC) + .returns(TypeName.VOID) + .addParameter(contextType, "context") + .addParameter(clientType, "client") + .addException(Exception::class.java) + + val required = mutableListOf() + val innerArgs = + with(processingEnv) { + inner.parameters.joinToString { param -> + val anno = param.getAnnotation(Builder.XCom::class.java) + val type = param.asType() + when { + anno != null -> + param.simpleName.toString().also { + required += RequiredXCom(type, it, anno.task.ifBlank { it }) + } + isType(type, clientType) -> "client" + isType(type, contextType) -> "context" + else -> throw IllegalArgumentException("Unsupported task parameter '${param.simpleName}' with type: $type") + } + } + } + required.forEach { + executeSpec.addStatement( + $$"var $L = ($T) client.getXCom($S)", + it.paramName, + with(TypeName.get(it.paramType)) { if (isPrimitive) box() else this }, + it.taskId, + ) + } + if (inner.returnType.kind == TypeKind.VOID) { + $$"new $T().$L($L)" + } else { + $$"client.setXCom(new $T().$L($L))" + }.also { + executeSpec.addStatement( + it, + ClassName.get(parent), + inner.simpleName, + innerArgs, + ) + } + + val spec = + TypeSpec + .classBuilder(name) + .addSuperinterface(Task::class.java) + .addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC) + .addMethod(executeSpec.build()) + .build() + return BuildTaskResult(spec) + } +} + +private fun ProcessingEnvironment.isType( + t: TypeMirror, + c: ClassName, +): Boolean = typeUtils.isSameType(t, elementUtils.getTypeElement(c.canonicalName()).asType()) + +private data class RequiredXCom( + val paramType: TypeMirror, + val paramName: String, + val taskId: String, +) + +private data class BuildTaskResult( + val spec: TypeSpec, +) diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt new file mode 100644 index 0000000000000..5ceb0ffac8ae8 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt @@ -0,0 +1,86 @@ +/* + * 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 org.apache.airflow.sdk + +/** + * An immutable snapshot of all [Dag]s that this JVM process can execute. + * + *

Build a [Bundle] by implementing [BundleBuilder], then pass it to [Server.serve] + * to start accepting task-execution requests. + * + * @property version Implementation version read from the JAR manifest + * ({@code Implementation-Version}); falls back to {@code "0"} if absent. + * @property dags All registered Dags keyed by [Dag.id]. Insertion order is preserved. + */ +class Bundle( + val version: String, + dags: Iterable, +) { + val dags: Map = dags.associateByDagId() +} + +private fun Iterable.associateByDagId(): Map { + val dagMap = linkedMapOf() + for (dag in this) { + require(dagMap.putIfAbsent(dag.id, dag) == null) { + "Dags in bundle have duplicate ID: ${dag.id}" + } + } + return dagMap +} + +/** + * Entry point for declaring the [Dag]s that this bundle contains. + * + *

Implement this interface in the class named as {@code Main-Class} in your JAR + * manifest. The build tooling instantiates it at compile time to record Dag and task + * IDs in the manifest, enabling inspection without running the full process. + * + *

{@code
+ * public class MyBundleBuilder implements BundleBuilder {
+ *     @Override
+ *     public Iterable getDags() {
+ *         return List.of(MyDagBuilder.build());
+ *     }
+ *
+ *     public static void main(String[] args) {
+ *         Server.create(args).serve(new MyBundleBuilder().build());
+ *     }
+ * }
+ * }
+ */ +interface BundleBuilder { + /** + * Returns all [Dag]s that belong to this bundle. + * + *

Called once during [build]; Dag IDs must be unique across the returned collection. + */ + fun getDags(): Iterable + + /** + * Constructs a [Bundle] from the Dags returned by [getDags]. + * + *

The bundle version is taken from the JAR's {@code Implementation-Version} manifest + * attribute, or {@code "0"} if that attribute is absent. + * + * @throws IllegalArgumentException if any two Dags share the same ID. + */ + fun build(): Bundle = Bundle(this::class.java.`package`.implementationVersion ?: "0", getDags()) +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt new file mode 100644 index 0000000000000..f4dee2d2f6e1d --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt @@ -0,0 +1,149 @@ +/* + * 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 org.apache.airflow.sdk + +import org.apache.airflow.sdk.execution.Client +import org.apache.airflow.sdk.execution.comm.StartupDetails + +/** + * A connection registered in Airflow's connection store. + * + * @property id Connection ID as configured in Airflow. + * @property type Connection type (e.g. {@code "http"}, {@code "postgres"}), if configured. + * @property host Hostname, if configured. + * @property schema Schema or database name, if configured. + * @property login Username, if configured. + * @property password Password, if configured. + * @property port Port number, if configured. + * @property extra JSON blob of extra connection parameters, if configured. + */ +data class Connection( + @JvmField val id: String, + @JvmField val type: String?, + @JvmField val host: String?, + @JvmField val schema: String?, + @JvmField val login: String?, + @JvmField val password: String?, + @JvmField val port: Int?, + @JvmField val extra: Any?, +) + +/** + * Client for Airflow API calls scoped to the current task instance. + * + *

An instance is provided when a task is being executed. All reads and writes are + * automatically scoped to the current Dag run and task instance unless you pass + * explicit IDs. + */ +class Client internal constructor( + internal val details: StartupDetails, + internal val impl: Client, +) { + internal companion object { + /** + * Default XCom key used for a task's return value ({@value}). + */ + const val XCOM_RETURN_KEY = "return_value" + } + + /** + * Retrieves a connection from the Airflow connection store. + * + * @param id Connection ID as configured in Airflow. + * @return The connection. + * @throws ApiError if the connection does not exist or the API call fails. + */ + fun getConnection(id: String): Connection = + with(impl.getConnection(id)) { + Connection( + id = connId, + type = connType, + host = host as String?, + schema = schema as String?, + login = login as String?, + password = password as String?, + port = port as Int?, + extra = extra, + ) + } + + /** + * Retrieves an Airflow variable. + * + * @param key Variable key. + * @return The variable value, or {@code null} if the variable is not set. + * @throws ApiError if the API call fails. + */ + fun getVariable(key: String): Any? = impl.getVariable(key).value + + /** + * Reads an XCom value pushed by another task. + * + *

The current Dag run's [dagId][TaskInstance.dagId] and [runId][TaskInstance.runId] + * are used by default; override them only when reading across Dags or runs. + * + * @param key XCom key to read; defaults to [XCOM_RETURN_KEY]. + * @param dagId Dag that owns the XCom; defaults to the current Dag. + * @param taskId Task that pushed the XCom. + * @param runId Run that produced the XCom; defaults to the current run. + * @param mapIndex Map index of the source task instance, or {@code null} for non-mapped tasks. + * @param includePriorDates If {@code true}, also search earlier Dag-run dates. + * @return The XCom value, or {@code null} if none was pushed. + * @throws ApiError if the API call fails. + */ + @JvmOverloads fun getXCom( + key: String = XCOM_RETURN_KEY, + dagId: String = details.ti.dagId, + taskId: String, + runId: String = details.ti.runId, + mapIndex: Int? = null, + includePriorDates: Boolean = false, + ): Any? = + impl + .getXCom( + key = key, + dagId = dagId, + taskId = taskId, + runId = runId, + mapIndex = mapIndex, + includePriorDates = includePriorDates, + ).value + + /** + * Pushes an XCom value for downstream tasks to read. + * + *

The current task instance's identifiers are used automatically. + * + * @param key XCom key; defaults to [XCOM_RETURN_KEY]. + * @param value Value to push. Must be JSON-serializable. + * @throws ApiError if the API call fails. + */ + @JvmOverloads fun setXCom( + key: String = XCOM_RETURN_KEY, + value: Any, + ) = impl.setXCom( + key = key, + value = value, + dagId = details.ti.dagId, + taskId = details.ti.taskId, + runId = details.ti.runId, + mapIndex = details.ti.mapIndex ?: -1, + ) +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt new file mode 100644 index 0000000000000..4842c393154e6 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt @@ -0,0 +1,72 @@ +/* + * 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 org.apache.airflow.sdk + +import org.apache.airflow.sdk.execution.comm.StartupDetails + +/** + * Identifies the Dag run that the current task instance belongs to. + * + * @property dagId ID of the Dag being run. + * @property runId Unique identifier for this Dag run. + */ +data class DagRun( + @JvmField val dagId: String, + @JvmField val runId: String, +) + +/** + * Identifies the task instance that is currently executing. + * + * @property dagId ID of the parent Dag. + * @property runId ID of the Dag run that triggered this instance. + * @property taskId ID of the task within the Dag. + * @property mapIndex Index within a mapped task group, if this is a mapped task instance. + * @property tryNumber How many times this task instance has been attempted (1-based). + */ +data class TaskInstance( + @JvmField val dagId: String, + @JvmField val runId: String, + @JvmField val taskId: String, + @JvmField val mapIndex: Int?, + @JvmField val tryNumber: Int, +) + +/** + * Runtime context passed to the task execution. + * + *

Provides metadata about the current Dag run and task instance. + * Use [Client] to interact with Airflow at runtime (connections, variables, XComs). + * + * @property dagRun Dag run the currently executing task instance belongs to. + * @property ti Currently executing task instance. + */ +data class Context( + @JvmField val dagRun: DagRun, + @JvmField val ti: TaskInstance, +) { + internal companion object { + fun from(request: StartupDetails): Context = + Context( + dagRun = with(request.tiContext.dagRun) { DagRun(dagId, runId) }, + ti = with(request.ti) { TaskInstance(dagId, runId, taskId, mapIndex, tryNumber) }, + ) + } +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt new file mode 100644 index 0000000000000..c7bb2277b55c6 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt @@ -0,0 +1,92 @@ +/* + * 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 org.apache.airflow.sdk + +import kotlin.Throws + +/** + * A collection of tasks with directional dependencies. + * + *

Create a [Dag] directly and register tasks with [addTask]. + * + * The [Builder.Dag] annotation should generally be preferred in user code, where the + * annotation processor generates the wiring for you. Only use this class directly if + * you need to do low-level plumbing. + * + * @param id Dag identifier. Must contain only ASCII alphanumeric characters, dashes, + * dots, or underscores; must be unique within a [Bundle]. + * + * @see Builder.Dag + */ +class Dag( + val id: String, // TODO: charset check? +) { + internal var tasks = mutableMapOf>() + + /** + * Registers a task with this Dag. + * + *

The class must have a public no-argument constructor and implement [Task]. + * Task IDs must be unique within a Dag. + * + * @param id Task identifier, unique within this Dag. + * @param definition Class that implements [Task]. Must have a public no-arg constructor. + * @return This Dag, for chaining. + */ + fun addTask( + id: String, + definition: Class, + ): Dag { + // TODO: Check duplicate key. + tasks[id] = definition + return this + } +} + +/** + * A single unit of work executed by Airflow. + * + *

Prefer using the [Builder.Task] annotation with [Builder.Dag] to have the + * annotation processor generate an implementation for you. Only use this interface if + * you need to do low-level plumbing. + * + *

Implement this interface to define task logic. Airflow instantiates the class via + * its no-argument constructor, then calls [execute] once per task-instance run. + * + * @see Builder.Dag + * @see Builder.Task + */ +interface Task { + /** + * Executes this task. + * + *

Any exception thrown marks the task instance as failed. Use [client] to read + * connections, variables, pull XComs, or to push an XCom for downstream tasks. + * + * @param context Runtime metadata for the current Dag run and task instance. + * @param client Client for Airflow API calls scoped to this task instance. + * @throws Exception on failure; the task instance is marked failed. + */ + @Throws(Exception::class) + fun execute( + context: Context, + client: Client, + ) +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt new file mode 100644 index 0000000000000..e963863533213 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt @@ -0,0 +1,159 @@ +/* + * 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 org.apache.airflow.sdk + +import com.xenomachina.argparser.ArgParser +import io.ktor.network.selector.SelectorManager +import io.ktor.network.sockets.InetSocketAddress +import io.ktor.network.sockets.aSocket +import io.ktor.network.sockets.openReadChannel +import io.ktor.network.sockets.openWriteChannel +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.coroutineScope +import kotlinx.coroutines.launch +import kotlinx.coroutines.runBlocking +import org.apache.airflow.sdk.execution.CoordinatorComm +import org.apache.airflow.sdk.execution.LogSender +import org.apache.airflow.sdk.execution.Logger +import kotlin.text.substringAfterLast +import kotlin.text.substringBeforeLast + +private class Args( + parser: ArgParser, +) { + private fun parseAddress(address: String): InetSocketAddress = + InetSocketAddress( + address.substringBeforeLast(':'), + address.substringAfterLast(':').toInt(), + ) + + val comm by parser.storing("--comm", help = "Address (host:port) to communicate with parent") { + parseAddress(this) + } + val logs by parser.storing("--logs", help = "Address (host:port) to send Airflow logs to") { + parseAddress(this) + } +} + +/** + * Thrown when an Airflow API call returns an error response. + * + *

Extends [IllegalStateException] so callers can handle it without checked-exception + * machinery. + */ +class ApiError( + message: String, +) : IllegalStateException(message) + +/** + * Connects this JVM process to the Airflow coordinator and dispatches task-execution + * requests to the registered [Bundle]. + * + *

The typical entry point is: + *

{@code
+ * public static void main(String[] args) {
+ *     Server.create(args).serve(new MyBundleBuilder().build());
+ * }
+ * }
+ * + *

The process exits when the coordinator closes the connection (normally after + * one task-instance execution). + */ +class Server( + private val comm: InetSocketAddress, + private val logs: InetSocketAddress, +) { + companion object { + /** + * Parses coordinator addresses from command-line arguments and returns a + * ready-to-use [Server]. + * + *

The arguments are supplied automatically by Airflow and are not intended + * to be constructed by hand: + *

    + *
  • {@code --comm host:port} — address for task-execution messages.
  • + *
  • {@code --logs host:port} — address for log forwarding.
  • + *
+ * + * @param args Command-line arguments as received by {@code main}. + * @return A configured [Server] ready to call [serve]. + */ + @JvmStatic + fun create(args: Array): Server { + val args = ArgParser(args).parseInto(::Args) + return Server(args.comm, args.logs) + } + } + + private val logger = Logger(Server::class) + + /** + * Blocking entry point: connects to the coordinator and serves task-execution + * requests from the given [bundle]. + * + *

This is a convenience wrapper around [serveAsync] for use from a plain + * {@code main} method. Prefer [serveAsync] when calling from an existing coroutine. + * The call returns when the coordinator closes the connection (normally after + * one task-instance execution). + * + * @param bundle Bundle containing all Dags this process can execute. + * + * @see [serveAsync] + */ + fun serve(bundle: Bundle) { + runBlocking { launch { serveAsync(bundle) } } + } + + /** + * Suspending entry point: connects to the coordinator and serves task-execution + * requests from the given [bundle]. + * + *

Opens both the task-execution channel ({@code --comm}) and the log-forwarding + * channel ({@code --logs}) concurrently, then processes incoming requests until the + * coordinator closes the connection (normally after one task-instance execution). + * The coroutine returns once both channels have been closed. + * + *

Use this variant when calling from an existing coroutine scope; use the + * blocking [serve] from a plain {@code main} method. + * + * @param bundle Bundle containing all Dags this process can execute. + * + * @see [serve] + */ + suspend fun serveAsync(bundle: Bundle) = + coroutineScope { + launch { + aSocket(SelectorManager(Dispatchers.IO)).tcp().connect(comm).use { socket -> + logger.debug("Connected comm", mapOf("addr" to comm)) + CoordinatorComm( + bundle, + socket.openReadChannel(), + socket.openWriteChannel(autoFlush = true), + ).startProcessing() + } + } + launch { + aSocket(SelectorManager(Dispatchers.IO)).tcp().connect(logs).use { socket -> + logger.debug("Connected logs", mapOf("addr" to logs)) + LogSender.configure(socket.openWriteChannel(autoFlush = true)) + } + } + } +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt new file mode 100644 index 0000000000000..306e23640df1a --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt @@ -0,0 +1,131 @@ +/* + * 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 org.apache.airflow.sdk.execution + +import kotlinx.coroutines.runBlocking +import org.apache.airflow.sdk.execution.comm.ConnectionResult +import org.apache.airflow.sdk.execution.comm.GetConnection +import org.apache.airflow.sdk.execution.comm.GetVariable +import org.apache.airflow.sdk.execution.comm.GetXCom +import org.apache.airflow.sdk.execution.comm.SetXCom +import org.apache.airflow.sdk.execution.comm.VariableResult +import org.apache.airflow.sdk.execution.comm.XComResult + +/** + * @suppress + * + * Transport contract between [org.apache.airflow.sdk.Client] and the coordinator. + * + * Implementations translate each SDK method call into the appropriate + * message and unwrap the raw response model into the value expected by the public + * SDK layer. + * + * Currently, the only production implementation is [CoordinatorClient]. A test + * double can be supplied via the internal [org.apache.airflow.sdk.Client] + * constructor to exercise task logic without a live coordinator. + */ +interface Client { + fun getConnection(id: String): ConnectionResult + + fun getVariable(key: String): VariableResult + + fun getXCom( + key: String, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int? = null, + includePriorDates: Boolean = false, + ): XComResult + + fun setXCom( + key: String, + value: Any, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int, + ) +} + +/** + * @suppress + * + * Production [Client] implementation backed by a live comm. + * + * Each method serializes the request into the appropriate message type (e.g. + * [GetConnection], [GetXCom]), sends it over the comm, and returns the + * unwrapped response model. All calls block the calling thread because task + * [execute][org.apache.airflow.sdk.Task.execute] runs on a plain thread, not + * inside a coroutine. + */ +class CoordinatorClient( + val exec: CoordinatorComm, +) : Client { + override fun getConnection(id: String) = + runBlocking { + exec.communicate(GetConnection().apply { connId = id }) + } + + override fun getVariable(key: String) = + runBlocking { + exec.communicate(GetVariable().also { it.key = key }) + } + + override fun setXCom( + key: String, + value: Any, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int, + ) { + val message = + SetXCom().also { + it.key = key + it.value = value + it.dagId = dagId + it.taskId = taskId + it.runId = runId + it.mapIndex = mapIndex + } + runBlocking { exec.communicate(message) } + } + + override fun getXCom( + key: String, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int?, + includePriorDates: Boolean, + ): XComResult { + val message = + GetXCom().also { + it.key = key + it.dagId = dagId + it.taskId = taskId + it.runId = runId + it.mapIndex = mapIndex + it.includePriorDates = includePriorDates + } + return runBlocking { exec.communicate(message) } + } +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt new file mode 100644 index 0000000000000..86bf3ef7e0457 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Comm.kt @@ -0,0 +1,138 @@ +/* + * 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 org.apache.airflow.sdk.execution + +import io.ktor.utils.io.ByteReadChannel +import io.ktor.utils.io.ByteWriteChannel +import io.ktor.utils.io.readByteArray +import io.ktor.utils.io.writeByteArray +import org.apache.airflow.sdk.ApiError +import org.apache.airflow.sdk.Bundle +import org.apache.airflow.sdk.execution.comm.ErrorResponse +import org.apache.airflow.sdk.execution.comm.StartupDetails +import kotlin.concurrent.atomics.AtomicInt +import kotlin.concurrent.atomics.ExperimentalAtomicApi +import kotlin.system.exitProcess + +data class IncomingFrame( + val id: Int, + val body: Any?, +) + +data class OutgoingFrame( + val id: Int, + val body: Any, +) + +@OptIn(ExperimentalAtomicApi::class) +class CoordinatorComm( + private val bundle: Bundle, + private val reader: ByteReadChannel, + private val writer: ByteWriteChannel, +) { + internal companion object { + private val logger = Logger(CoordinatorComm::class) + + fun encode(outgoing: OutgoingFrame) = Frame.encodeRequest(outgoing.id, outgoing.body) + + fun decode(bytes: ByteArray) = Frame.decode(bytes) + } + + private val nextId = AtomicInt(0) + private var shutDownRequested = false + + suspend fun startProcessing() { + while (!shutDownRequested) { + processOnce(::handleIncoming) + } + logger.debug("Goodbye") + } + + private suspend fun processOnce(handle: suspend (IncomingFrame) -> Unit) { + val prefix = reader.readByteArray(4) // First 4 bytes as length. + if (prefix.size != 4) { // Something is terribly wrong. Let's bail. + logger.error("Need 4 prefix bytes", mapOf("actual" to prefix.size)) + shutDownRequested = true + return + } + + val payloadLength = Frame.parseLengthPrefix(prefix) + val payload = reader.readByteArray(payloadLength) + if (payload.size != payloadLength) { // Something is terribly wrong. Let's bail. + logger.error( + "Payload length not right", + mapOf("expect" to payloadLength, "receive" to payload.size), + ) + shutDownRequested = true + return + } + val frame = decode(payload) + logger.debug("Handling", mapOf("id" to frame.id)) + handle(frame) + } + + private suspend fun sendMessage( + id: Int, + body: Any, + ) { + val data = encode(OutgoingFrame(id, body)) + logger.debug("Sending", mapOf("id" to id, "body" to body)) + writer.writeByteArray(Frame.lengthPrefix(data.size)) + writer.writeByteArray(data) + } + + suspend fun handleIncoming(frame: IncomingFrame) { + when (val request = frame.body) { + null -> {} + is ErrorResponse -> { + println("Error!! id=${frame.id} [${request.error}] ${request.detail}") // TODO: Handle error. + exitProcess(1) + } + is StartupDetails -> { + sendMessage(frame.id, runTask(bundle, request, this)) + shutDownRequested = true + } + } + } + + @Throws(ApiError::class) + suspend fun communicateImpl(body: Any): Any { + var frame: IncomingFrame? = null + + suspend fun handle(f: IncomingFrame) { + frame = f + } + sendMessage(nextId.fetchAndAdd(1), body) + processOnce(::handle) + if (frame == null) { + throw ApiError("No response received") + } + return frame.body ?: Unit + } + + @Throws(ApiError::class) + suspend inline fun communicate(request: Any): T { + when (val response = communicateImpl(request)) { + is ErrorResponse -> throw ApiError("[${response.error}] ${response.detail}") + is T -> return response + else -> throw ApiError("Unexpected response type ${response::class.java}") + } + } +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Frame.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Frame.kt new file mode 100644 index 0000000000000..a3815d6140141 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Frame.kt @@ -0,0 +1,102 @@ +/* + * 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 org.apache.airflow.sdk.execution + +import com.fasterxml.jackson.databind.DeserializationFeature +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.databind.SerializationFeature +import com.fasterxml.jackson.databind.util.StdDateFormat +import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule +import org.apache.airflow.sdk.execution.comm.Discriminator +import org.msgpack.core.MessagePack +import java.io.ByteArrayOutputStream + +object Frame { + private val mapper = + ObjectMapper().apply { + configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false) + disable(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS) + registerModule(JavaTimeModule()) + registerModule(TimestampToJavaOffsetDateTimeModule()) + setDateFormat(StdDateFormat().withColonInTimeZone(true)) + } + + fun encodeRequest( + id: Int, + body: Any, + ): ByteArray = encodeFrame(id, body) + + fun decode(bytes: ByteArray): IncomingFrame { + val unpacker = MessagePack.newDefaultUnpacker(bytes) + val headerSize = unpacker.unpackArrayHeader() + check(headerSize >= 1) { "Unexpected Task SDK frame arity $headerSize" } + + val id = unpacker.unpackInt() + val rawBody = if (headerSize >= 2) unpacker.unpackAny() else null + val rawError = if (headerSize >= 3) unpacker.unpackAny() else null + unpacker.close() + + val body = decodeMessage(rawError) ?: decodeMessage(rawBody) + return IncomingFrame(id, body) + } + + fun lengthPrefix(length: Int) = + byteArrayOf( + (length shr 24).toByte(), + (length shr 16).toByte(), + (length shr 8).toByte(), + length.toByte(), + ) + + fun parseLengthPrefix(prefix: ByteArray): Int { + check(prefix.size == 4) { "Need 4 prefix bytes" } + return prefix.fold(0) { acc, byte -> (acc shl 8) or (byte.toInt() and 0xff) } + } + + private fun encodeFrame( + id: Int, + body: Any?, + ): ByteArray { + val payload = ByteArrayOutputStream() + val packer = MessagePack.newDefaultPacker(payload) + packer.packArrayHeader(2) + packer.packInt(id) + packer.packAny(body?.let(::toBody)) + packer.close() + return payload.toByteArray() + } + + private fun decodeMessage(raw: Any?): Any? { + val body = raw as? Map<*, *> ?: return raw + return mapper.convertValue(body, Discriminator.discriminate(body) ?: return raw) + } + + @Suppress("UNCHECKED_CAST") + private fun toBody(value: Any): Map = + when (value) { + is Map<*, *> -> value as Map + else -> mapper.convertValue(value, MutableMap::class.java) as MutableMap + } +} + +private fun Discriminator.discriminate(body: Map<*, *>): Class<*>? { + val name = body["type"] as? String ?: return null + return types[name] ?: error("Unsupported supervisor message type: $name") +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt new file mode 100644 index 0000000000000..27005f92082cb --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt @@ -0,0 +1,128 @@ +/* + * 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 org.apache.airflow.sdk.execution + +import io.ktor.utils.io.ByteWriteChannel +import io.ktor.utils.io.writeString +import kotlinx.coroutines.runBlocking +import kotlinx.datetime.LocalDateTime +import kotlinx.datetime.TimeZone +import kotlinx.datetime.toLocalDateTime +import kotlinx.serialization.json.JsonElement +import kotlinx.serialization.json.JsonNull +import kotlinx.serialization.json.JsonPrimitive +import kotlinx.serialization.json.buildJsonArray +import kotlinx.serialization.json.buildJsonObject +import java.util.concurrent.ConcurrentLinkedDeque +import kotlin.reflect.KClass +import kotlin.time.Clock + +enum class Level { ERROR, DEBUG, } + +internal data class LogMessage( + val event: String, + val arguments: Map, + val logger: Logger, + val level: Level, + val timestamp: LocalDateTime = Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()), +) + +internal class Logger( + cls: KClass<*>, +) { + val name: String? = cls.java.typeName + + // TODO: Actually implement level filtering. + @Suppress("UNUSED_PARAMETER") + fun isEnabledForLevel(level: Level): Boolean = true + + fun debug( + message: String, + arguments: Map = emptyMap(), + ) { + log(Level.DEBUG, message, arguments) + } + + fun error( + message: String, + arguments: Map = emptyMap(), + ) { + log(Level.ERROR, message, arguments) + } + + private fun log( + level: Level, + event: String, + arguments: Map, + ) { + if (!isEnabledForLevel(level)) return + LogSender.send(LogMessage(event, arguments, this, level)) + } +} + +internal object LogSender { + private var writer: ByteWriteChannel? = null + private val messages: ConcurrentLinkedDeque = ConcurrentLinkedDeque() + + fun configure(channel: ByteWriteChannel) { + writer = channel + if (!channel.isClosedForWrite) { + while (messages.isNotEmpty()) { + sendTo(channel, messages.removeFirst()) + } + } + } + + fun send(message: LogMessage) { + val channel = writer + if (channel == null || channel.isClosedForWrite) { + messages.addLast(message) + } else { + sendTo(channel, message) + } + } + + private fun sendTo( + writer: ByteWriteChannel, + message: LogMessage, + ) { + val map = message.arguments.toMutableMap() + map["event"] = message.event + map["level"] = message.level.name.lowercase() + map["logger"] = message.logger.name ?: "(java)" + map["timestamp"] = message.timestamp + // TODO: Can this be done asynchronously instead? + runBlocking { writer.writeString("${map.toJsonElement()}\n") } + } +} + +private fun Any?.toJsonElement(): JsonElement = + when (this) { + is JsonElement -> this + is Map<*, *> -> + buildJsonObject { + forEach { (k, v) -> put(k.toString(), v.toJsonElement()) } + } + is Iterable<*> -> buildJsonArray { forEach { add(it.toJsonElement()) } } + is Number -> JsonPrimitive(this) + is String -> JsonPrimitive(this) + null -> JsonNull + else -> JsonPrimitive(toString()) // Also correctly handles Kotlinx DateTime. + } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/MsgPack.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/MsgPack.kt new file mode 100644 index 0000000000000..88fa3f632d959 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/MsgPack.kt @@ -0,0 +1,155 @@ +/* + * 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 org.apache.airflow.sdk.execution + +import com.fasterxml.jackson.core.JsonParser +import com.fasterxml.jackson.core.JsonToken +import com.fasterxml.jackson.core.util.JacksonFeatureSet +import com.fasterxml.jackson.databind.DeserializationContext +import com.fasterxml.jackson.databind.deser.std.StdDeserializer +import com.fasterxml.jackson.databind.module.SimpleModule +import com.fasterxml.jackson.datatype.jsr310.JavaTimeFeature +import com.fasterxml.jackson.datatype.jsr310.deser.InstantDeserializer +import org.msgpack.core.ExtensionTypeHeader +import org.msgpack.core.MessagePack +import org.msgpack.core.MessagePacker +import org.msgpack.core.MessageUnpacker +import org.msgpack.jackson.dataformat.MessagePackExtensionType +import org.msgpack.value.ArrayValue +import org.msgpack.value.ExtensionValue +import org.msgpack.value.MapValue +import org.msgpack.value.Value +import org.msgpack.value.ValueType +import java.math.BigInteger +import java.time.OffsetDateTime +import java.time.ZoneOffset + +private fun MessagePacker.packByteArray(data: ByteArray) { + packBinaryHeader(data.size) + data.forEach { packByte(it) } +} + +private fun MessagePacker.packMap(data: Map<*, *>) { + packMapHeader(data.size) + data.forEach { (k, v) -> + check(k is String) + packString(k) + packAny(v) + } +} + +private fun MessagePacker.packCollection(data: Collection<*>) { + packArrayHeader(data.size) + data.forEach { packAny(it) } +} + +fun MessagePacker.packAny(data: Any?) { + when (data) { + null -> packNil() + is Boolean -> packBoolean(data) + is Byte -> packByte(data) + is Short -> packShort(data) + is Int -> packInt(data) + is Long -> packLong(data) + is BigInteger -> packBigInteger(data) + is Float -> packFloat(data) + is Double -> packDouble(data) + is ByteArray -> packByteArray(data) + is String -> packString(data) + is Map<*, *> -> packMap(data) + is Collection<*> -> packCollection(data) + else -> throw IllegalArgumentException("Unsupported data type: $data") + } +} + +private fun ArrayValue.decodeArray(): List<*> = + mutableListOf().also { + iterator().forEach { v -> it.add(v.decode()) } + } + +private fun MapValue.decodeMap(): Map<*, *> = + mutableMapOf().also { + entrySet().forEach { (k, v) -> it[k.asStringValue().asString()] = v.decode() } + } + +private fun ExtensionValue.decodeExtensionValue(): Any? = + when (type) { + TimestampToJavaOffsetDateTimeModule.EXT_TYPE -> { + MessagePack + .newDefaultUnpacker(data) + .unpackTimestamp(ExtensionTypeHeader(type, data.size)) + .atOffset(ZoneOffset.UTC) + } + else -> throw IllegalArgumentException("Unsupported extension type: $this") + } + +private fun Value.decode(): Any? = + when (valueType) { + ValueType.NIL -> null + ValueType.BOOLEAN -> asBooleanValue().boolean + ValueType.INTEGER -> with(asIntegerValue()) { if (isInLongRange) asLong() else asBigInteger() } + ValueType.FLOAT -> asFloatValue().toDouble() + ValueType.STRING -> asStringValue().asString() + ValueType.BINARY -> asBinaryValue().asByteArray() + ValueType.ARRAY -> asArrayValue().decodeArray() + ValueType.MAP -> asMapValue().decodeMap() + ValueType.EXTENSION -> asExtensionValue().decodeExtensionValue() + + else -> throw IllegalArgumentException("Unsupported data type: $this") + } + +fun MessageUnpacker.unpackAny(): Any? = unpackValue().decode() + +class TimestampToJavaOffsetDateTimeModule : SimpleModule() { + companion object { + const val EXT_TYPE: Byte = -1 + } + + class OffsetDateTimeDeserializer : StdDeserializer(OffsetDateTime::class.java) { + val instantDeserializer: InstantDeserializer = + InstantDeserializer.OFFSET_DATE_TIME.withFeatures( + JacksonFeatureSet.fromDefaults(JavaTimeFeature.entries.toTypedArray()), + ) + + override fun deserialize( + p: JsonParser, + ctxt: DeserializationContext, + ): OffsetDateTime { + if (p.currentToken == JsonToken.VALUE_EMBEDDED_OBJECT) { + deserializeMsgPackTimestamp(p)?.let { return it } + } + return instantDeserializer.deserialize(p, ctxt) + } + + private fun deserializeMsgPackTimestamp(p: JsonParser): OffsetDateTime? { + val ext = p.readValueAs(MessagePackExtensionType::class.java) + if (ext.type != EXT_TYPE) { + return null + } + val unpacker = MessagePack.newDefaultUnpacker(ext.data) + val instant = unpacker.unpackTimestamp(ExtensionTypeHeader(EXT_TYPE, ext.data.size)) + return instant.atOffset(ZoneOffset.UTC) + } + } + + init { + addDeserializer(OffsetDateTime::class.java, OffsetDateTimeDeserializer()) + } +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt new file mode 100644 index 0000000000000..70e093a288fa1 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt @@ -0,0 +1,86 @@ +/* + * 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 org.apache.airflow.sdk.execution + +import org.apache.airflow.sdk.Bundle +import org.apache.airflow.sdk.Client +import org.apache.airflow.sdk.Context +import org.apache.airflow.sdk.execution.comm.AssetProfile +import org.apache.airflow.sdk.execution.comm.StartupDetails +import org.apache.airflow.sdk.execution.comm.SucceedTask +import org.apache.airflow.sdk.execution.comm.TaskState +import java.time.OffsetDateTime + +internal object TaskResult { + fun success( + endDate: OffsetDateTime = OffsetDateTime.now(), + taskOutlets: List = emptyList(), + outletEvents: List> = emptyList(), + renderedMapIndex: String? = null, + ) = SucceedTask().also { + it.state = "success" + it.endDate = endDate + it.taskOutlets = taskOutlets + it.outletEvents = outletEvents + it.renderedMapIndex = renderedMapIndex + } + + fun of( + state: TaskState.State, + endDate: OffsetDateTime = OffsetDateTime.now(), + renderedMapIndex: String? = null, + ) = TaskState().also { + it.state = state + it.endDate = endDate + it.renderedMapIndex = renderedMapIndex + } +} + +internal object TaskRunner { + val logger = Logger(TaskRunner::class) + + internal fun runTask( + bundle: Bundle, + request: StartupDetails, + client: Client, + ): Any { + val task = bundle.dags[request.ti.dagId]?.tasks[request.ti.taskId] ?: return TaskResult.of(TaskState.State.REMOVED) + return try { + task.getDeclaredConstructor().newInstance().execute(Context.from(request), client) + TaskResult.success() + } catch (e: Exception) { + logger.error("Error executing task", mapOf("ti" to request.ti, "error" to e, "trace" to e.stackTraceToString())) + e.printStackTrace() + TaskResult.of(TaskState.State.FAILED) + } + } +} + +internal fun runTask( + bundle: Bundle, + request: StartupDetails, + comm: CoordinatorComm, +): Any = TaskRunner.runTask(bundle, request, Client(request, CoordinatorClient(comm))) + +internal fun runTask( + bundle: Bundle, + request: StartupDetails, + client: Client, +) = TaskRunner.runTask(bundle, request, client) diff --git a/java-sdk/sdk/src/main/resources/META-INF/services/javax.annotation.processing.Processor b/java-sdk/sdk/src/main/resources/META-INF/services/javax.annotation.processing.Processor new file mode 100644 index 0000000000000..cc7c3d49ef78a --- /dev/null +++ b/java-sdk/sdk/src/main/resources/META-INF/services/javax.annotation.processing.Processor @@ -0,0 +1,18 @@ +# +# 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. +org.apache.airflow.sdk.BuilderProcessor diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt new file mode 100644 index 0000000000000..bd3aae81d13f9 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt @@ -0,0 +1,239 @@ +/* + * 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 org.apache.airflow.sdk + +import com.google.testing.compile.CompilationSubject.assertThat +import com.google.testing.compile.Compiler +import com.google.testing.compile.JavaFileObjectSubject +import com.google.testing.compile.JavaFileObjects +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test + +private fun compile(source: String) = + Compiler.javac().withProcessors(BuilderProcessor()).compile( + JavaFileObjects.forSourceString("org.apache.airflow.example.TestExample", source), + ) + +private fun JavaFileObjectSubject.hasSourceEquivalentTo( + qual: String, + source: String, +) = hasSourceEquivalentTo( + JavaFileObjects.forSourceString(qual, source), +) + +class BuilderTest { + @Test + @DisplayName("generate builder for dag class") + fun generateBuilderForDagClass() { + val compilation = + compile( + """ + package org.apache.airflow.example; + + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.Context; + + @Builder.Dag + public class TestExample { + @Builder.Task + public void t1() {} + + @Builder.Task + public int t2(Client client) { + return (Integer) client.getXCom("t0"); + } + + @Builder.Task + public void t3(Context ctx, @Builder.XCom(task = "t2") int value) { + System.out.println(String.format("%s %s", ctx.ti, value)); + } + } + """, + ) + + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleBuilder") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleBuilder", + """ + package org.apache.airflow.example; + + import java.lang.Exception; + import java.lang.Integer; + import java.lang.Override; + import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.Context; + import org.apache.airflow.sdk.Dag; + import org.apache.airflow.sdk.Task; + + public final class TestExampleBuilder { + public static Dag build() { + var dag = new Dag("TestExample"); + dag.addTask("t1", T1.class); + dag.addTask("t2", T2.class); + dag.addTask("t3", T3.class); + return dag; + } + public static final class T1 implements Task { + @Override + public void execute(Context context, Client client) throws Exception { + new TestExample().t1(); + } + } + public static final class T2 implements Task { + @Override + public void execute(Context context, Client client) throws Exception { + client.setXCom(new TestExample().t2(client)); + } + } + public static final class T3 implements Task { + @Override + public void execute(Context context, Client client) throws Exception { + var value = (Integer) client.getXCom("t2"); + new TestExample().t3(context, value); + } + } + } + """, + ) + } + + @Test + @DisplayName("generate builder for dag class with custom dag id") + fun generateBuilderWithCustomDagId() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag(id = "foo") public class TestExample {} + """, + ) + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleBuilder") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleBuilder", + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Dag; + public final class TestExampleBuilder { public static Dag build() { var dag = new Dag("foo"); return dag; } } + """, + ) + } + + @Test + @DisplayName("generate builder for dag class with custom class name") + fun generateBuilderWithCustomClassName() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag(to = "Foo") public class TestExample {} + """, + ) + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.Foo") + .hasSourceEquivalentTo( + "org.apache.airflow.example.Foo", + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Dag; + public final class Foo { public static Dag build() { var dag = new Dag("TestExample"); return dag; } } + """, + ) + } + + @Test + @DisplayName("generate builder for dag class with custom task name") + fun generateBuilderForDagClassWithCustomTaskName() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { @Builder.Task(id = "foo") public void t1() {} } + """, + ) + + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleBuilder") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleBuilder", + """ + package org.apache.airflow.example; + import java.lang.Exception; + import java.lang.Override; + import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.Context; + import org.apache.airflow.sdk.Dag; + import org.apache.airflow.sdk.Task; + public final class TestExampleBuilder { + public static Dag build() { + var dag = new Dag("TestExample"); + dag.addTask("foo", T1.class); + return dag; + } + public static final class T1 implements Task { + @Override public void execute(Context context, Client client) throws Exception { new TestExample().t1(); } + } + } + """, + ) + } + + @Test + @DisplayName("generate builder for dag class with invalid task parameter") + fun generateBuilderForDagClassWithInvalidTaskParameter() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { @Builder.Task(id = "foo") public void t1(String client) {} } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Unsupported task parameter 'client' with type: java.lang.String", + ) + } + + @Test + @DisplayName("generate builder for dag class with varargs task parameter") + fun generateBuilderForDagClassWithVarArgsTaskParameter() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { @Builder.Task(id = "foo") public void t1(String... client) {} } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Cannot create task from vararg function t1", + ) + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt new file mode 100644 index 0000000000000..48bae797cbe20 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt @@ -0,0 +1,47 @@ +/* + * 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 org.apache.airflow.sdk + +import org.junit.jupiter.api.Assertions +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test + +internal class BundleTest { + @Test + @DisplayName("Should index dags by dagId") + fun shouldIndexDagsByDagId() { + val dag = Dag("dag") + + val bundle = Bundle("0", listOf(dag)) + + Assertions.assertEquals(mapOf("dag" to dag), bundle.dags) + } + + @Test + @DisplayName("Should reject duplicate dag ids") + fun shouldRejectDuplicateDagIds() { + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + Bundle("0", listOf(Dag("dag"), Dag("dag"))) + } + + Assertions.assertEquals("Dags in bundle have duplicate ID: dag", error.message) + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt new file mode 100644 index 0000000000000..30a5db40a3389 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt @@ -0,0 +1,86 @@ +/* + * 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 org.apache.airflow.sdk.execution + +import org.apache.airflow.sdk.execution.comm.StartupDetails +import org.junit.jupiter.api.Assertions +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test +import java.time.OffsetDateTime +import java.time.ZoneOffset + +fun byteArrayFromHexString(hexString: String): ByteArray = + hexString + .split(' ', '\r', '\n') + .filter { it.isNotEmpty() } + .map { it.toUByte(16).toByte() } + .toByteArray() + +class CommsTest { + @Test + @DisplayName("Should decode startup details") + fun shouldDecodeStartupDetails() { + // [2, msg, null] with msg coming from + // https://github.com/astronomer/airflow/blob/f39c8da8/task-sdk/tests/task_sdk/execution_time/test_comms.py#L73-L108 + val data = + """ + 92 02 88 a4 74 79 70 65 ae 53 74 61 72 74 75 70 44 65 74 61 69 6c 73 a2 74 69 86 a2 69 64 d9 24 + 34 64 38 32 38 61 36 32 2d 61 34 31 37 2d 34 39 33 36 2d 61 37 61 36 2d 32 62 33 66 61 62 61 63 + 65 63 61 62 a7 74 61 73 6b 5f 69 64 a1 61 aa 74 72 79 5f 6e 75 6d 62 65 72 01 a6 72 75 6e 5f 69 + 64 a1 62 a6 64 61 67 5f 69 64 a1 63 ae 64 61 67 5f 76 65 72 73 69 6f 6e 5f 69 64 d9 24 34 64 38 + 32 38 61 36 32 2d 61 34 31 37 2d 34 39 33 36 2d 61 37 61 36 2d 32 62 33 66 61 62 61 63 65 63 61 + 62 aa 74 69 5f 63 6f 6e 74 65 78 74 85 a7 64 61 67 5f 72 75 6e 8c a6 64 61 67 5f 69 64 a1 63 a6 + 72 75 6e 5f 69 64 a1 62 ac 6c 6f 67 69 63 61 6c 5f 64 61 74 65 b4 32 30 32 34 2d 31 32 2d 30 31 + 54 30 31 3a 30 30 3a 30 30 5a b3 64 61 74 61 5f 69 6e 74 65 72 76 61 6c 5f 73 74 61 72 74 b4 32 + 30 32 34 2d 31 32 2d 30 31 54 30 30 3a 30 30 3a 30 30 5a b1 64 61 74 61 5f 69 6e 74 65 72 76 61 + 6c 5f 65 6e 64 b4 32 30 32 34 2d 31 32 2d 30 31 54 30 31 3a 30 30 3a 30 30 5a aa 73 74 61 72 74 + 5f 64 61 74 65 b4 32 30 32 34 2d 31 32 2d 30 31 54 30 31 3a 30 30 3a 30 30 5a a9 72 75 6e 5f 61 + 66 74 65 72 b4 32 30 32 34 2d 31 32 2d 30 31 54 30 31 3a 30 30 3a 30 30 5a a8 65 6e 64 5f 64 61 + 74 65 c0 a8 72 75 6e 5f 74 79 70 65 a6 6d 61 6e 75 61 6c a5 73 74 61 74 65 a7 73 75 63 63 65 73 + 73 a4 63 6f 6e 66 c0 b5 63 6f 6e 73 75 6d 65 64 5f 61 73 73 65 74 5f 65 76 65 6e 74 73 90 a9 6d + 61 78 5f 74 72 69 65 73 00 ac 73 68 6f 75 6c 64 5f 72 65 74 72 79 c2 a9 76 61 72 69 61 62 6c 65 + 73 c0 ab 63 6f 6e 6e 65 63 74 69 6f 6e 73 c0 a4 66 69 6c 65 a9 2f 64 65 76 2f 6e 75 6c 6c aa 73 + 74 61 72 74 5f 64 61 74 65 b4 32 30 32 34 2d 31 32 2d 30 31 54 30 31 3a 30 30 3a 30 30 5a ac 64 + 61 67 5f 72 65 6c 5f 70 61 74 68 a9 2f 64 65 76 2f 6e 75 6c 6c ab 62 75 6e 64 6c 65 5f 69 6e 66 + 6f 82 a4 6e 61 6d 65 a8 61 6e 79 2d 6e 61 6d 65 a7 76 65 72 73 69 6f 6e ab 61 6e 79 2d 76 65 72 + 73 69 6f 6e b2 73 65 6e 74 72 79 5f 69 6e 74 65 67 72 61 74 69 6f 6e a0 c0 + """.trimIndent() + val result = CoordinatorComm.decode(byteArrayFromHexString(data)) + Assertions.assertInstanceOf(IncomingFrame::class.java, result) + Assertions.assertInstanceOf(StartupDetails::class.java, result.body) + } + + @Test + @DisplayName("Should serialize all fields") + fun shouldEncodeSucceedTask() { + val endDate = OffsetDateTime.of(2024, 12, 1, 1, 0, 0, 0, ZoneOffset.UTC) + val bytes = CoordinatorComm.encode(OutgoingFrame(3, TaskResult.success(endDate = endDate))) + val actual = bytes.toHexString(HexFormat { bytes { byteSeparator = " " } }) + + val expected = + """ + 92 03 85 a5 73 74 61 74 65 a7 73 75 63 63 65 73 73 a8 65 6e 64 5f 64 61 74 65 b4 32 30 32 34 2d + 31 32 2d 30 31 54 30 31 3a 30 30 3a 30 30 5a ac 74 61 73 6b 5f 6f 75 74 6c 65 74 73 90 ad 6f 75 + 74 6c 65 74 5f 65 76 65 6e 74 73 90 a4 74 79 70 65 ab 53 75 63 63 65 65 64 54 61 73 6b + """.trimIndent().replace('\n', ' ') + + Assertions.assertEquals(expected, actual) + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt new file mode 100644 index 0000000000000..19796a384d9c0 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt @@ -0,0 +1,147 @@ +/* + * 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 org.apache.airflow.sdk.execution + +import org.apache.airflow.sdk.Bundle +import org.apache.airflow.sdk.Client +import org.apache.airflow.sdk.Context +import org.apache.airflow.sdk.Dag +import org.apache.airflow.sdk.Task +import org.apache.airflow.sdk.execution.comm.BundleInfo +import org.apache.airflow.sdk.execution.comm.DagRun +import org.apache.airflow.sdk.execution.comm.StartupDetails +import org.apache.airflow.sdk.execution.comm.SucceedTask +import org.apache.airflow.sdk.execution.comm.TIRunContext +import org.apache.airflow.sdk.execution.comm.TaskInstanceDTO +import org.apache.airflow.sdk.execution.comm.TaskState +import org.junit.jupiter.api.Assertions +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test +import java.time.OffsetDateTime +import java.util.UUID + +class TaskTest { + @Test + @DisplayName("Should execute task and return success") + fun shouldExecuteTaskAndReturnSuccess() { + val result = runTask(bundleWith("success", SuccessTask::class.java), startupDetails(taskId = "success"), noOpClient()) + + Assertions.assertInstanceOf(SucceedTask::class.java, result) + } + + @Test + @DisplayName("Should return removed when task is missing") + fun shouldReturnRemovedWhenTaskIsMissing() { + val result = runTask(bundleWith("other", SuccessTask::class.java), startupDetails(taskId = "missing"), noOpClient()) + + Assertions.assertInstanceOf(TaskState::class.java, result) + Assertions.assertEquals(TaskState.State.REMOVED, (result as TaskState).state) + } + + @Test + @DisplayName("Should return failed when task throws") + fun shouldReturnFailedWhenTaskThrows() { + val result = runTask(bundleWith("failing", FailingTask::class.java), startupDetails(taskId = "failing"), noOpClient()) + + Assertions.assertInstanceOf(TaskState::class.java, result) + Assertions.assertEquals(TaskState.State.FAILED, (result as TaskState).state) + } + + private fun bundleWith( + taskId: String, + taskClass: Class, + ): Bundle { + val dag = Dag("test_dag") + dag.addTask(taskId, taskClass) + return Bundle("1", listOf(dag)) + } + + private fun startupDetails(taskId: String): StartupDetails = + StartupDetails().also { + it.ti = + TaskInstanceDTO().also { o -> + o.id = UUID.randomUUID() + o.taskId = taskId + o.dagId = "test_dag" + o.runId = "manual__2026-03-31T00:00:00+00:00" + o.tryNumber = 1 + o.dagVersionId = UUID.randomUUID() + } + it.dagRelPath = "/dev/null" + it.bundleInfo = + BundleInfo().also { info -> + info.name = "bundle" + info.version = "1" + } + it.startDate = OffsetDateTime.parse("2026-03-31T00:00:00Z") + it.tiContext = + TIRunContext().apply { + dagRun = + DagRun().apply { + dagId = "test_dag" + runId = "manual__2026-03-31T00:00:00+00:00" + } + } + it.sentryIntegration = "" + } + + private fun noOpClient() = + Client( + startupDetails(taskId = "unused"), + object : org.apache.airflow.sdk.execution.Client { + override fun getConnection(id: String) = throw UnsupportedOperationException("not used in test") + + override fun getVariable(key: String) = throw UnsupportedOperationException("not used in test") + + override fun getXCom( + key: String, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int?, + includePriorDates: Boolean, + ) = throw UnsupportedOperationException("not used in test") + + override fun setXCom( + key: String, + value: Any, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int, + ): Unit = throw UnsupportedOperationException("not used in test") + }, + ) + + class SuccessTask : Task { + override fun execute( + context: Context, + client: Client, + ) { + } + } + + class FailingTask : Task { + override fun execute( + context: Context, + client: Client, + ): Unit = throw IllegalStateException("boom") + } +} diff --git a/java-sdk/settings.gradle.kts b/java-sdk/settings.gradle.kts new file mode 100644 index 0000000000000..0892437c13d3e --- /dev/null +++ b/java-sdk/settings.gradle.kts @@ -0,0 +1,13 @@ +/* + * This file was generated by the Gradle 'init' task. + * + * The settings file is used to specify which projects to include in your build. + * For more detailed information on multi-project builds, please refer to https://docs.gradle.org/9.2.1/userguide/multi_project_builds.html in the Gradle documentation. + */ + +plugins { + id("org.gradle.toolchains.foojay-resolver-convention").version("0.10.0") +} + +rootProject.name = "airflow-java-sdk" +include("example", "sdk") diff --git a/scripts/ci/docker-compose/local.yml b/scripts/ci/docker-compose/local.yml index f91198189f6db..0e5073d251272 100644 --- a/scripts/ci/docker-compose/local.yml +++ b/scripts/ci/docker-compose/local.yml @@ -96,6 +96,9 @@ services: - type: bind source: ../../../go-sdk target: /opt/airflow/go-sdk + - type: bind + source: ../../../java-sdk + target: /opt/airflow/java-sdk - type: bind source: ../../../kubernetes-tests target: /opt/airflow/kubernetes-tests