jason810496 commented on code in PR #75184: URL: https://github.com/apache/airflow/pull/75184#discussion_r4440580486
########## go-sdk/capabilities.yaml: ########## Review Comment: Kept 3.4.0; added a `custom-xcom-backend` row, supported since Airflow 3.5.0. [`54db61b`](https://github.com/apache/airflow/pull/75184/commits/54db61bffa508ce77632e858f15c0e1793bd054a) ########## task-sdk/src/airflow/sdk/coordinators/_subprocess.py: ########## @@ -604,6 +658,220 @@ def _handle_set_task_state_store( ) return super()._handle_set_task_state_store(msg, log, req_id) + def _monitor_subprocess(self): + try: + super()._monitor_subprocess() + finally: + self._stop_xcom_backend() + + def _start_xcom_backend_call( + self, call: XComBackendCall, finish: Callable[[XComBackendResult], None] + ) -> None: + """Send ``call`` to the XCom backend process, which starts on the first call it handles.""" + if self._xcom_backend is not None and not self._xcom_backend.can_run_a_call: + # The process has exited, so a call sent to it would wait for an answer that never comes. + self._stop_xcom_backend() Review Comment: `stop()` now reaps a process known to be dead without servicing the shared selector. [`54db61b`](https://github.com/apache/airflow/pull/75184/commits/54db61bffa508ce77632e858f15c0e1793bd054a) ########## task-sdk/src/airflow/sdk/exceptions.py: ########## @@ -101,6 +101,9 @@ class ErrorType(enum.Enum): PERMISSION_DENIED = "PERMISSION_DENIED" GENERIC_ERROR = "GENERIC_ERROR" API_SERVER_ERROR = "API_SERVER_ERROR" + # The configured XCom backend could not run, or refused, an XCom of a language SDK task. Distinct + # from API_SERVER_ERROR, which would point a user at the API server rather than their own backend. Review Comment: Moved to the PR description. [`54db61b`](https://github.com/apache/airflow/pull/75184/commits/54db61bffa508ce77632e858f15c0e1793bd054a) ########## task-sdk/src/airflow/sdk/coordinators/_xcom_backend.py: ########## @@ -0,0 +1,337 @@ +# +# 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. +""" +The XCom backend for the tasks of a language SDK runtime. + +A language SDK runtime sends its XCom requests to the supervisor, which runs the configured XCom backend +on them before it forwards them to the API server. The backend runs in a helper process, started per task +the way the supervisor starts a Python task process, so the long-lived supervisor process keeps none of +the backend's state, such as a cached filesystem that a task process forked from it later could not use. + +The supervisor services the helper process from the selector of the task it serves, so a backend call +never blocks that task's request loop, which goes on heartbeating and enforcing ``execution_timeout``. +""" + +from __future__ import annotations + +import contextlib +import selectors +import signal +import time +from collections import deque +from collections.abc import Callable +from typing import TYPE_CHECKING, Annotated, ClassVar, Literal +from uuid import UUID + +import attrs +import structlog +from pydantic import BaseModel, Field, JsonValue, TypeAdapter + +from airflow.sdk.bases.xcom import BaseXCom, XComCoordinate +from airflow.sdk.execution_time.comms import ( + ConnectionResult, + ErrorResponse, + GetConnection, + GetVariable, + MaskSecret, + VariableResult, + XComResult, +) +from airflow.sdk.execution_time.supervisor import ( + ResponseSent, + WatchedSubprocess, + register_request_method, + task_process_uses_exec, +) + +if TYPE_CHECKING: + from socket import socket + from typing import Self + + from structlog.typing import FilteringBoundLogger + + from airflow.sdk.api.client import Client + +log: FilteringBoundLogger = structlog.get_logger(logger_name="xcom_backend") + +_EXIT_TIMEOUT = 10.0 + + +class XComBackendError(RuntimeError): + """ + An XCom backend call failed. + + For an exception raised in the helper process, the message holds its type and message, since the + exception itself does not cross the process boundary. + """ + + +class _Coordinate(BaseModel): + dag_id: str + run_id: str + task_id: str + ti_id: UUID | None = None + region_id: UUID | None = None + region_index: int | None = None + + +class XComBackendCall(BaseModel): + """ + A call of an XCom backend hook, sent to the helper process. + + ``value`` is a value as the runtime sends it for ``store``, the XCom row for ``load`` and ``purge``, + and a list of rows for ``load_many``. + """ + + hook: Literal["store", "load", "load_many", "purge", "exit"] + key: str = "" + value: JsonValue = None + coordinate: _Coordinate | None = None + map_index: int = -1 + type: Literal["XComBackendCall"] = "XComBackendCall" + + @classmethod + def store(cls, coordinate: XComCoordinate, map_index: int, key: str, value: JsonValue) -> Self: + return cls( + hook="store", + key=key, + value=value, + coordinate=_Coordinate.model_validate(attrs.asdict(coordinate)), + map_index=map_index, + ) + + +class XComBackendResult(BaseModel): + """The outcome of an XCom backend call, sent by the helper process.""" + + value: JsonValue = None + error: str | None = None + type: Literal["XComBackendResult"] = "XComBackendResult" + + +class AwaitXComBackendCall(BaseModel): + """Sent by the helper process when it is ready for its first call.""" + + type: Literal["AwaitXComBackendCall"] = "AwaitXComBackendCall" + + +_ToXComBackendProcess = Annotated[ + XComBackendCall | ConnectionResult | VariableResult | ErrorResponse, + Field(discriminator="type"), +] +_FromXComBackendProcess = Annotated[ + AwaitXComBackendCall | XComBackendResult | GetConnection | GetVariable | MaskSecret, + Field(discriminator="type"), +] + + +def runs_in_helper(xcom: type[BaseXCom]) -> bool: + """Whether task XComs go through the helper process, which runs every backend but the default one.""" + return xcom is not BaseXCom + + +_PendingCall = tuple[XComBackendCall, Callable[[XComBackendResult], None]] + + [email protected](kw_only=True) +class XComBackendProcess(WatchedSubprocess): + """ + The helper process that runs XCom backend calls for the task of a language SDK runtime. + + The process asks for a call, runs it, and answers with its result, which is also how it asks for the + next call. Each call carries the callback that takes its result, so a caller never has to match an + answer to its own call. Connection and Variable lookups of the backend reach the API server through + ``client``. + """ + + client: Client + decoder: ClassVar[TypeAdapter[_FromXComBackendProcess]] = TypeAdapter(_FromXComBackendProcess) + + _call_request_id: int | None = attrs.field(default=None, init=False) + _queued_calls: deque[_PendingCall] = attrs.field(factory=deque, init=False) + _running_call: Callable[[XComBackendResult], None] | None = attrs.field(default=None, init=False) + + @classmethod + def start( # type: ignore[override] + cls, + *, + id: UUID, + client: Client, + logger: FilteringBoundLogger, + selector: selectors.BaseSelector, + ) -> Self: + return super().start( + id=id, + client=client, + # The selector of the task this process serves, so that the task's request loop services both. + selector=selector, + logger=logger, + target=run_xcom_backend, + use_exec=task_process_uses_exec(), + ) + + @property + def can_run_a_call(self) -> bool: + """Whether the process can still be sent a call, which a process that has exited cannot.""" + return self.stdin in self._open_sockets + + def send_call(self, call: XComBackendCall, finish: Callable[[XComBackendResult], None]) -> None: + """ + Send ``call`` to the process, which answers it by calling ``finish`` with its result. + + The process runs one call at a time, so a call waits its turn. ``finish`` also takes the error of + a call the process never answers, such as when it exits while the call is waiting. + """ + self._queued_calls.append((call, finish)) + self._send_next_call() + + def _send_next_call(self) -> None: + """Hand the oldest waiting call to the process, if it has asked for one and runs none.""" + if self._call_request_id is None or self._running_call is not None or not self._queued_calls: + return + call, self._running_call = self._queued_calls.popleft() + request_id, self._call_request_id = self._call_request_id, None + self.send_msg(call, request_id=request_id) + + def _fail_waiting_calls(self, error: str) -> None: + """Answer every call the process has not answered with ``error``.""" + running, self._running_call = self._running_call, None + waiting = [running] if running is not None else [] + waiting += [finish for _, finish in self._queued_calls] + self._queued_calls.clear() + for finish in waiting: + finish(XComBackendResult(error=error)) + + def stop(self) -> None: + """Let the helper process exit, and kill it if it does not exit in time.""" + # The caller gives up on the calls it has not answered, so they are dropped rather than failed. Review Comment: `stop()` now fails them, so no caller can be left waiting. [`54db61b`](https://github.com/apache/airflow/pull/75184/commits/54db61bffa508ce77632e858f15c0e1793bd054a) ########## scripts/tests/ci/prek/test_sync_java_sdk_supervisor_schema.py: ########## @@ -0,0 +1,132 @@ +# 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 + +import json + +import pytest +import sync_java_sdk_supervisor_schema as syncer + +VERSION = "2027-02-28" + + +def write_schema(path, api_version=VERSION, **extra): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps({"api_version": api_version, **extra})) + return path + + [email protected] +def repo(tmp_path): + """A checkout whose declared version, owned schema and vendored copy are all in step.""" + (tmp_path / syncer.GRADLE_PROPERTIES).parent.mkdir(parents=True, exist_ok=True) + (tmp_path / syncer.GRADLE_PROPERTIES).write_text( + f"org.gradle.caching=true\nairflowSupervisorSchemaVersion={VERSION}\n" + ) + write_schema(tmp_path / syncer.SOURCE_SCHEMA, title="supervisor") + write_schema(tmp_path / syncer.VENDORED_SCHEMA, title="supervisor") + return tmp_path + + +def test_the_copy_is_vendored_inside_the_java_sdk(): + # A copy outside java-sdk/ would leave the Gradle build reaching across the monorepo, + # which is the coupling vendoring removes. + assert str(syncer.VENDORED_SCHEMA).startswith("java-sdk/") + assert not str(syncer.SOURCE_SCHEMA).startswith("java-sdk/") Review Comment: Dropped it — it only asserted constants. [`7af6b46`](https://github.com/apache/airflow/pull/75942/commits/7af6b460abbea522f8d2702da43f11f736f7edd2) ########## task-sdk/src/airflow/sdk/coordinators/_xcom_backend.py: ########## @@ -0,0 +1,337 @@ +# +# 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. +""" +The XCom backend for the tasks of a language SDK runtime. + +A language SDK runtime sends its XCom requests to the supervisor, which runs the configured XCom backend +on them before it forwards them to the API server. The backend runs in a helper process, started per task +the way the supervisor starts a Python task process, so the long-lived supervisor process keeps none of +the backend's state, such as a cached filesystem that a task process forked from it later could not use. + +The supervisor services the helper process from the selector of the task it serves, so a backend call +never blocks that task's request loop, which goes on heartbeating and enforcing ``execution_timeout``. +""" + +from __future__ import annotations + +import contextlib +import selectors +import signal +import time +from collections import deque +from collections.abc import Callable +from typing import TYPE_CHECKING, Annotated, ClassVar, Literal +from uuid import UUID + +import attrs +import structlog +from pydantic import BaseModel, Field, JsonValue, TypeAdapter + +from airflow.sdk.bases.xcom import BaseXCom, XComCoordinate +from airflow.sdk.execution_time.comms import ( + ConnectionResult, + ErrorResponse, + GetConnection, + GetVariable, + MaskSecret, + VariableResult, + XComResult, +) +from airflow.sdk.execution_time.supervisor import ( + ResponseSent, + WatchedSubprocess, + register_request_method, + task_process_uses_exec, +) + +if TYPE_CHECKING: + from socket import socket + from typing import Self + + from structlog.typing import FilteringBoundLogger + + from airflow.sdk.api.client import Client + +log: FilteringBoundLogger = structlog.get_logger(logger_name="xcom_backend") + +_EXIT_TIMEOUT = 10.0 + + +class XComBackendError(RuntimeError): + """ + An XCom backend call failed. + + For an exception raised in the helper process, the message holds its type and message, since the + exception itself does not cross the process boundary. + """ + + +class _Coordinate(BaseModel): + dag_id: str + run_id: str + task_id: str + ti_id: UUID | None = None + region_id: UUID | None = None + region_index: int | None = None + + +class XComBackendCall(BaseModel): + """ + A call of an XCom backend hook, sent to the helper process. + + ``value`` is a value as the runtime sends it for ``store``, the XCom row for ``load`` and ``purge``, + and a list of rows for ``load_many``. + """ + + hook: Literal["store", "load", "load_many", "purge", "exit"] + key: str = "" + value: JsonValue = None + coordinate: _Coordinate | None = None + map_index: int = -1 + type: Literal["XComBackendCall"] = "XComBackendCall" + + @classmethod + def store(cls, coordinate: XComCoordinate, map_index: int, key: str, value: JsonValue) -> Self: + return cls( + hook="store", + key=key, + value=value, + coordinate=_Coordinate.model_validate(attrs.asdict(coordinate)), + map_index=map_index, + ) + + +class XComBackendResult(BaseModel): + """The outcome of an XCom backend call, sent by the helper process.""" + + value: JsonValue = None + error: str | None = None + type: Literal["XComBackendResult"] = "XComBackendResult" + + +class AwaitXComBackendCall(BaseModel): + """Sent by the helper process when it is ready for its first call.""" + + type: Literal["AwaitXComBackendCall"] = "AwaitXComBackendCall" + + +_ToXComBackendProcess = Annotated[ + XComBackendCall | ConnectionResult | VariableResult | ErrorResponse, + Field(discriminator="type"), +] +_FromXComBackendProcess = Annotated[ + AwaitXComBackendCall | XComBackendResult | GetConnection | GetVariable | MaskSecret, + Field(discriminator="type"), +] + + +def runs_in_helper(xcom: type[BaseXCom]) -> bool: + """Whether task XComs go through the helper process, which runs every backend but the default one.""" + return xcom is not BaseXCom + + +_PendingCall = tuple[XComBackendCall, Callable[[XComBackendResult], None]] + + [email protected](kw_only=True) +class XComBackendProcess(WatchedSubprocess): + """ + The helper process that runs XCom backend calls for the task of a language SDK runtime. + + The process asks for a call, runs it, and answers with its result, which is also how it asks for the + next call. Each call carries the callback that takes its result, so a caller never has to match an + answer to its own call. Connection and Variable lookups of the backend reach the API server through + ``client``. + """ + + client: Client + decoder: ClassVar[TypeAdapter[_FromXComBackendProcess]] = TypeAdapter(_FromXComBackendProcess) + + _call_request_id: int | None = attrs.field(default=None, init=False) + _queued_calls: deque[_PendingCall] = attrs.field(factory=deque, init=False) + _running_call: Callable[[XComBackendResult], None] | None = attrs.field(default=None, init=False) + + @classmethod + def start( # type: ignore[override] + cls, + *, + id: UUID, + client: Client, + logger: FilteringBoundLogger, + selector: selectors.BaseSelector, + ) -> Self: + return super().start( + id=id, + client=client, + # The selector of the task this process serves, so that the task's request loop services both. + selector=selector, + logger=logger, + target=run_xcom_backend, + use_exec=task_process_uses_exec(), + ) + + @property + def can_run_a_call(self) -> bool: + """Whether the process can still be sent a call, which a process that has exited cannot.""" + return self.stdin in self._open_sockets + + def send_call(self, call: XComBackendCall, finish: Callable[[XComBackendResult], None]) -> None: + """ + Send ``call`` to the process, which answers it by calling ``finish`` with its result. + + The process runs one call at a time, so a call waits its turn. ``finish`` also takes the error of + a call the process never answers, such as when it exits while the call is waiting. + """ + self._queued_calls.append((call, finish)) + self._send_next_call() + + def _send_next_call(self) -> None: + """Hand the oldest waiting call to the process, if it has asked for one and runs none.""" + if self._call_request_id is None or self._running_call is not None or not self._queued_calls: + return + call, self._running_call = self._queued_calls.popleft() + request_id, self._call_request_id = self._call_request_id, None + self.send_msg(call, request_id=request_id) + + def _fail_waiting_calls(self, error: str) -> None: + """Answer every call the process has not answered with ``error``.""" + running, self._running_call = self._running_call, None + waiting = [running] if running is not None else [] + waiting += [finish for _, finish in self._queued_calls] + self._queued_calls.clear() + for finish in waiting: + finish(XComBackendResult(error=error)) + + def stop(self) -> None: + """Let the helper process exit, and kill it if it does not exit in time.""" + # The caller gives up on the calls it has not answered, so they are dropped rather than failed. + self._queued_calls.clear() + self._running_call = None + if self._exit_code is None and self._call_request_id is not None: + # The process may be dead already, such as when the task it serves was killed. + with contextlib.suppress(OSError): + self.send_msg(XComBackendCall(hook="exit"), request_id=self._call_request_id) + self._call_request_id = None + deadline = time.monotonic() + _EXIT_TIMEOUT + while self._open_sockets and (now := time.monotonic()) < deadline: + self._service_subprocess(max_wait_time=deadline - now) + if self._exit_code is None: + with contextlib.suppress(self._process.TimeoutExpired): + self._exit_code = self._process.wait(timeout=max(0.0, deadline - time.monotonic())) + if self._exit_code is None: + self.kill(signal.SIGTERM, force=True) + # Only this process's own sockets, since the selector belongs to the task it serves. + for sock in list(self._open_sockets): + self._on_socket_closed(sock) + sock.close() + + def _on_socket_closed(self, sock: socket) -> None: + super()._on_socket_closed(sock) + if sock is self.stdin: + # The process has exited, or is exiting, so it answers none of the calls it was still to run. + self._fail_waiting_calls("The XCom backend process exited before it answered") + + def _handle_await_call( + self, msg: AwaitXComBackendCall, log: FilteringBoundLogger, req_id: int + ) -> ResponseSent: + self._call_request_id = req_id + self._send_next_call() + return ResponseSent.ALREADY_SENT + + def _handle_result(self, msg: XComBackendResult, log: FilteringBoundLogger, req_id: int) -> ResponseSent: + self._call_request_id = req_id + finish, self._running_call = self._running_call, None + self._send_next_call() + if finish is not None: + finish(msg) + return ResponseSent.ALREADY_SENT + + _request_handlers = { + **WatchedSubprocess._get_shared_request_handlers(GetConnection, GetVariable, MaskSecret), + **dict( + [ + register_request_method(AwaitXComBackendCall, _handle_await_call), + register_request_method(XComBackendResult, _handle_result), + ] + ), + } + + +def run_xcom_backend() -> None: + """Run the XCom backend calls that the supervisor sends, until it sends ``exit``.""" + # Lazy: these load the task runtime and the configured backend, which only the helper process needs. Review Comment: Removed. [`54db61b`](https://github.com/apache/airflow/pull/75184/commits/54db61bffa508ce77632e858f15c0e1793bd054a) ########## dev/breeze/src/airflow_breeze/utils/selective_checks.py: ########## @@ -228,6 +228,17 @@ def __hash__(self): r"^scripts/ci/lang_sdk_serialization/.*", ) +# Sources of the `lang_sdk_xcom_object_storage` e2e mode, which runs the configured XCom backend for +# TypeScript, Go and Java task XComs: the coordinator-side backend, the base XCom class it calls, the +# object storage backend, and the mode's own compose override and tests. +LANG_SDK_XCOM_BACKEND_FILES = ( + r"^task-sdk/src/airflow/sdk/coordinators/_xcom_backend\.py$", + r"^task-sdk/src/airflow/sdk/bases/xcom\.py$", + r"^providers/common/io/src/airflow/providers/common/io/xcom/.*", + r"^airflow-e2e-tests/docker/lang-sdk-xcom-object-storage\.yml$", + r"^airflow-e2e-tests/tests/airflow_e2e_tests/lang_sdk_xcom_object_storage_tests/.*", +) Review Comment: Now its own group and `run-lang-sdk-xcom-e2e-tests` output; the job's `if:` follows it. [`54db61b`](https://github.com/apache/airflow/pull/75184/commits/54db61bffa508ce77632e858f15c0e1793bd054a) ########## task-sdk/src/airflow/sdk/coordinators/_xcom_backend.py: ########## @@ -0,0 +1,337 @@ +# +# 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. +""" +The XCom backend for the tasks of a language SDK runtime. + +A language SDK runtime sends its XCom requests to the supervisor, which runs the configured XCom backend +on them before it forwards them to the API server. The backend runs in a helper process, started per task +the way the supervisor starts a Python task process, so the long-lived supervisor process keeps none of +the backend's state, such as a cached filesystem that a task process forked from it later could not use. + +The supervisor services the helper process from the selector of the task it serves, so a backend call +never blocks that task's request loop, which goes on heartbeating and enforcing ``execution_timeout``. +""" + +from __future__ import annotations + +import contextlib +import selectors +import signal +import time +from collections import deque +from collections.abc import Callable +from typing import TYPE_CHECKING, Annotated, ClassVar, Literal +from uuid import UUID + +import attrs +import structlog +from pydantic import BaseModel, Field, JsonValue, TypeAdapter + +from airflow.sdk.bases.xcom import BaseXCom, XComCoordinate +from airflow.sdk.execution_time.comms import ( + ConnectionResult, + ErrorResponse, + GetConnection, + GetVariable, + MaskSecret, + VariableResult, + XComResult, +) +from airflow.sdk.execution_time.supervisor import ( + ResponseSent, + WatchedSubprocess, + register_request_method, + task_process_uses_exec, +) + +if TYPE_CHECKING: + from socket import socket + from typing import Self + + from structlog.typing import FilteringBoundLogger + + from airflow.sdk.api.client import Client + +log: FilteringBoundLogger = structlog.get_logger(logger_name="xcom_backend") + +_EXIT_TIMEOUT = 10.0 + + +class XComBackendError(RuntimeError): + """ + An XCom backend call failed. + + For an exception raised in the helper process, the message holds its type and message, since the + exception itself does not cross the process boundary. + """ + + +class _Coordinate(BaseModel): + dag_id: str + run_id: str + task_id: str + ti_id: UUID | None = None + region_id: UUID | None = None + region_index: int | None = None + + +class XComBackendCall(BaseModel): + """ + A call of an XCom backend hook, sent to the helper process. + + ``value`` is a value as the runtime sends it for ``store``, the XCom row for ``load`` and ``purge``, + and a list of rows for ``load_many``. + """ + + hook: Literal["store", "load", "load_many", "purge", "exit"] + key: str = "" + value: JsonValue = None + coordinate: _Coordinate | None = None + map_index: int = -1 + type: Literal["XComBackendCall"] = "XComBackendCall" + + @classmethod + def store(cls, coordinate: XComCoordinate, map_index: int, key: str, value: JsonValue) -> Self: + return cls( + hook="store", + key=key, + value=value, + coordinate=_Coordinate.model_validate(attrs.asdict(coordinate)), + map_index=map_index, + ) + + +class XComBackendResult(BaseModel): + """The outcome of an XCom backend call, sent by the helper process.""" + + value: JsonValue = None + error: str | None = None + type: Literal["XComBackendResult"] = "XComBackendResult" + + +class AwaitXComBackendCall(BaseModel): + """Sent by the helper process when it is ready for its first call.""" + + type: Literal["AwaitXComBackendCall"] = "AwaitXComBackendCall" + + +_ToXComBackendProcess = Annotated[ + XComBackendCall | ConnectionResult | VariableResult | ErrorResponse, + Field(discriminator="type"), +] +_FromXComBackendProcess = Annotated[ + AwaitXComBackendCall | XComBackendResult | GetConnection | GetVariable | MaskSecret, + Field(discriminator="type"), +] + + +def runs_in_helper(xcom: type[BaseXCom]) -> bool: + """Whether task XComs go through the helper process, which runs every backend but the default one.""" + return xcom is not BaseXCom + + +_PendingCall = tuple[XComBackendCall, Callable[[XComBackendResult], None]] + + [email protected](kw_only=True) +class XComBackendProcess(WatchedSubprocess): + """ + The helper process that runs XCom backend calls for the task of a language SDK runtime. + + The process asks for a call, runs it, and answers with its result, which is also how it asks for the + next call. Each call carries the callback that takes its result, so a caller never has to match an + answer to its own call. Connection and Variable lookups of the backend reach the API server through + ``client``. + """ + + client: Client + decoder: ClassVar[TypeAdapter[_FromXComBackendProcess]] = TypeAdapter(_FromXComBackendProcess) + + _call_request_id: int | None = attrs.field(default=None, init=False) + _queued_calls: deque[_PendingCall] = attrs.field(factory=deque, init=False) + _running_call: Callable[[XComBackendResult], None] | None = attrs.field(default=None, init=False) + + @classmethod + def start( # type: ignore[override] + cls, + *, + id: UUID, + client: Client, + logger: FilteringBoundLogger, + selector: selectors.BaseSelector, + ) -> Self: + return super().start( + id=id, + client=client, + # The selector of the task this process serves, so that the task's request loop services both. + selector=selector, + logger=logger, + target=run_xcom_backend, + use_exec=task_process_uses_exec(), + ) + + @property + def can_run_a_call(self) -> bool: + """Whether the process can still be sent a call, which a process that has exited cannot.""" + return self.stdin in self._open_sockets + + def send_call(self, call: XComBackendCall, finish: Callable[[XComBackendResult], None]) -> None: + """ + Send ``call`` to the process, which answers it by calling ``finish`` with its result. + + The process runs one call at a time, so a call waits its turn. ``finish`` also takes the error of + a call the process never answers, such as when it exits while the call is waiting. + """ + self._queued_calls.append((call, finish)) + self._send_next_call() + + def _send_next_call(self) -> None: + """Hand the oldest waiting call to the process, if it has asked for one and runs none.""" + if self._call_request_id is None or self._running_call is not None or not self._queued_calls: + return + call, self._running_call = self._queued_calls.popleft() + request_id, self._call_request_id = self._call_request_id, None + self.send_msg(call, request_id=request_id) Review Comment: Caught: the call fails instead of raising out of the request handler. [`54db61b`](https://github.com/apache/airflow/pull/75184/commits/54db61bffa508ce77632e858f15c0e1793bd054a) ########## task-sdk/src/airflow/sdk/coordinators/_xcom_backend.py: ########## @@ -0,0 +1,337 @@ +# +# 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. +""" +The XCom backend for the tasks of a language SDK runtime. + +A language SDK runtime sends its XCom requests to the supervisor, which runs the configured XCom backend +on them before it forwards them to the API server. The backend runs in a helper process, started per task +the way the supervisor starts a Python task process, so the long-lived supervisor process keeps none of +the backend's state, such as a cached filesystem that a task process forked from it later could not use. + +The supervisor services the helper process from the selector of the task it serves, so a backend call +never blocks that task's request loop, which goes on heartbeating and enforcing ``execution_timeout``. +""" + +from __future__ import annotations + +import contextlib +import selectors +import signal +import time +from collections import deque +from collections.abc import Callable +from typing import TYPE_CHECKING, Annotated, ClassVar, Literal +from uuid import UUID + +import attrs +import structlog +from pydantic import BaseModel, Field, JsonValue, TypeAdapter + +from airflow.sdk.bases.xcom import BaseXCom, XComCoordinate +from airflow.sdk.execution_time.comms import ( + ConnectionResult, + ErrorResponse, + GetConnection, + GetVariable, + MaskSecret, + VariableResult, + XComResult, +) +from airflow.sdk.execution_time.supervisor import ( + ResponseSent, + WatchedSubprocess, + register_request_method, + task_process_uses_exec, +) + +if TYPE_CHECKING: + from socket import socket + from typing import Self + + from structlog.typing import FilteringBoundLogger + + from airflow.sdk.api.client import Client + +log: FilteringBoundLogger = structlog.get_logger(logger_name="xcom_backend") + +_EXIT_TIMEOUT = 10.0 + + +class XComBackendError(RuntimeError): + """ + An XCom backend call failed. + + For an exception raised in the helper process, the message holds its type and message, since the + exception itself does not cross the process boundary. + """ + + +class _Coordinate(BaseModel): + dag_id: str + run_id: str + task_id: str + ti_id: UUID | None = None + region_id: UUID | None = None + region_index: int | None = None + + +class XComBackendCall(BaseModel): + """ + A call of an XCom backend hook, sent to the helper process. + + ``value`` is a value as the runtime sends it for ``store``, the XCom row for ``load`` and ``purge``, + and a list of rows for ``load_many``. + """ + + hook: Literal["store", "load", "load_many", "purge", "exit"] + key: str = "" + value: JsonValue = None + coordinate: _Coordinate | None = None + map_index: int = -1 + type: Literal["XComBackendCall"] = "XComBackendCall" + + @classmethod + def store(cls, coordinate: XComCoordinate, map_index: int, key: str, value: JsonValue) -> Self: + return cls( + hook="store", + key=key, + value=value, + coordinate=_Coordinate.model_validate(attrs.asdict(coordinate)), + map_index=map_index, + ) + + +class XComBackendResult(BaseModel): + """The outcome of an XCom backend call, sent by the helper process.""" + + value: JsonValue = None + error: str | None = None + type: Literal["XComBackendResult"] = "XComBackendResult" + + +class AwaitXComBackendCall(BaseModel): + """Sent by the helper process when it is ready for its first call.""" + + type: Literal["AwaitXComBackendCall"] = "AwaitXComBackendCall" + + +_ToXComBackendProcess = Annotated[ + XComBackendCall | ConnectionResult | VariableResult | ErrorResponse, + Field(discriminator="type"), +] +_FromXComBackendProcess = Annotated[ + AwaitXComBackendCall | XComBackendResult | GetConnection | GetVariable | MaskSecret, + Field(discriminator="type"), +] + + +def runs_in_helper(xcom: type[BaseXCom]) -> bool: + """Whether task XComs go through the helper process, which runs every backend but the default one.""" + return xcom is not BaseXCom + + +_PendingCall = tuple[XComBackendCall, Callable[[XComBackendResult], None]] + + [email protected](kw_only=True) +class XComBackendProcess(WatchedSubprocess): + """ + The helper process that runs XCom backend calls for the task of a language SDK runtime. + + The process asks for a call, runs it, and answers with its result, which is also how it asks for the + next call. Each call carries the callback that takes its result, so a caller never has to match an + answer to its own call. Connection and Variable lookups of the backend reach the API server through + ``client``. + """ + + client: Client + decoder: ClassVar[TypeAdapter[_FromXComBackendProcess]] = TypeAdapter(_FromXComBackendProcess) + + _call_request_id: int | None = attrs.field(default=None, init=False) + _queued_calls: deque[_PendingCall] = attrs.field(factory=deque, init=False) + _running_call: Callable[[XComBackendResult], None] | None = attrs.field(default=None, init=False) + + @classmethod + def start( # type: ignore[override] + cls, + *, + id: UUID, + client: Client, + logger: FilteringBoundLogger, + selector: selectors.BaseSelector, + ) -> Self: + return super().start( + id=id, + client=client, + # The selector of the task this process serves, so that the task's request loop services both. + selector=selector, + logger=logger, + target=run_xcom_backend, + use_exec=task_process_uses_exec(), + ) + + @property + def can_run_a_call(self) -> bool: + """Whether the process can still be sent a call, which a process that has exited cannot.""" + return self.stdin in self._open_sockets + + def send_call(self, call: XComBackendCall, finish: Callable[[XComBackendResult], None]) -> None: + """ + Send ``call`` to the process, which answers it by calling ``finish`` with its result. + + The process runs one call at a time, so a call waits its turn. ``finish`` also takes the error of + a call the process never answers, such as when it exits while the call is waiting. + """ + self._queued_calls.append((call, finish)) + self._send_next_call() + + def _send_next_call(self) -> None: + """Hand the oldest waiting call to the process, if it has asked for one and runs none.""" + if self._call_request_id is None or self._running_call is not None or not self._queued_calls: + return + call, self._running_call = self._queued_calls.popleft() + request_id, self._call_request_id = self._call_request_id, None + self.send_msg(call, request_id=request_id) + + def _fail_waiting_calls(self, error: str) -> None: + """Answer every call the process has not answered with ``error``.""" + running, self._running_call = self._running_call, None + waiting = [running] if running is not None else [] + waiting += [finish for _, finish in self._queued_calls] + self._queued_calls.clear() + for finish in waiting: + finish(XComBackendResult(error=error)) + + def stop(self) -> None: + """Let the helper process exit, and kill it if it does not exit in time.""" + # The caller gives up on the calls it has not answered, so they are dropped rather than failed. + self._queued_calls.clear() + self._running_call = None + if self._exit_code is None and self._call_request_id is not None: + # The process may be dead already, such as when the task it serves was killed. + with contextlib.suppress(OSError): + self.send_msg(XComBackendCall(hook="exit"), request_id=self._call_request_id) + self._call_request_id = None + deadline = time.monotonic() + _EXIT_TIMEOUT + while self._open_sockets and (now := time.monotonic()) < deadline: + self._service_subprocess(max_wait_time=deadline - now) + if self._exit_code is None: + with contextlib.suppress(self._process.TimeoutExpired): + self._exit_code = self._process.wait(timeout=max(0.0, deadline - time.monotonic())) + if self._exit_code is None: + self.kill(signal.SIGTERM, force=True) + # Only this process's own sockets, since the selector belongs to the task it serves. + for sock in list(self._open_sockets): + self._on_socket_closed(sock) + sock.close() + + def _on_socket_closed(self, sock: socket) -> None: + super()._on_socket_closed(sock) + if sock is self.stdin: + # The process has exited, or is exiting, so it answers none of the calls it was still to run. + self._fail_waiting_calls("The XCom backend process exited before it answered") + + def _handle_await_call( + self, msg: AwaitXComBackendCall, log: FilteringBoundLogger, req_id: int + ) -> ResponseSent: + self._call_request_id = req_id + self._send_next_call() + return ResponseSent.ALREADY_SENT + + def _handle_result(self, msg: XComBackendResult, log: FilteringBoundLogger, req_id: int) -> ResponseSent: + self._call_request_id = req_id + finish, self._running_call = self._running_call, None + self._send_next_call() + if finish is not None: + finish(msg) + return ResponseSent.ALREADY_SENT + + _request_handlers = { + **WatchedSubprocess._get_shared_request_handlers(GetConnection, GetVariable, MaskSecret), + **dict( + [ + register_request_method(AwaitXComBackendCall, _handle_await_call), + register_request_method(XComBackendResult, _handle_result), + ] + ), + } + + +def run_xcom_backend() -> None: + """Run the XCom backend calls that the supervisor sends, until it sends ``exit``.""" + # Lazy: these load the task runtime and the configured backend, which only the helper process needs. + from airflow.sdk.execution_time import task_runner + from airflow.sdk.execution_time.comms import CommsDecoder + from airflow.sdk.execution_time.xcom import XCom + + comms = CommsDecoder[_ToXComBackendProcess, _FromXComBackendProcess]( + log=log, body_decoder=TypeAdapter(_ToXComBackendProcess) + ) + # The secrets backends of a task process look up Connections and Variables through SUPERVISOR_COMMS. + task_runner.SUPERVISOR_COMMS = comms # type: ignore[assignment] + runner = _HookRunner(XCom) + call = comms.send(AwaitXComBackendCall()) + while isinstance(call, XComBackendCall) and call.hook != "exit": + call = comms.send(runner.run(call)) + + [email protected] +class _HookRunner: + """Runs the hooks of an XCom backend that the supervisor's calls name, one method per hook.""" + + xcom: type[BaseXCom] + + def run(self, call: XComBackendCall) -> XComBackendResult: + try: + if (hook := getattr(self, call.hook, None)) is None: + raise ValueError(f"Unknown XCom backend hook {call.hook!r}") + return hook(call) + except Exception as e: + log.exception("XCom backend call failed", hook=call.hook, key=call.key) + return XComBackendResult(error=f"{type(e).__name__}: {e}") + + def store(self, call: XComBackendCall) -> XComBackendResult: + # Lazy: only the helper process decodes and encodes XCom values. Review Comment: Removed. [`54db61b`](https://github.com/apache/airflow/pull/75184/commits/54db61bffa508ce77632e858f15c0e1793bd054a) ########## java-sdk/sdk/build.gradle.kts: ########## @@ -185,6 +193,20 @@ abstract class SyncSupervisorSchemaTask : DefaultTask() { fun sync() { val file = schemaFile.get().asFile val version = schemaVersion.get() + // The schema in this repository is the source of truth, and is compared by content: the + // remote only serves released versions, so an in-progress one is only ever found here, and + // a change that leaves api_version alone still has to reach the copy. + val source = sourceFile.asFile.orNull + if (source != null && apiVersionOf(source) == version) { + if (file.exists() && file.readBytes().contentEquals(source.readBytes())) { + logger.lifecycle("Supervisor Schema is up-to-date (api_version=$version).") + return + } + logger.lifecycle("Refreshing Supervisor Schema from $source") + file.parentFile.mkdirs() + Files.copy(source.toPath(), file.toPath(), StandardCopyOption.REPLACE_EXISTING) + return + } Review Comment: Gradle now copies the in-repo source whenever it exists, so it agrees with the hook. [`7af6b46`](https://github.com/apache/airflow/pull/75942/commits/7af6b460abbea522f8d2702da43f11f736f7edd2) ########## task-sdk/src/airflow/sdk/coordinators/_subprocess.py: ########## @@ -604,6 +658,220 @@ def _handle_set_task_state_store( ) return super()._handle_set_task_state_store(msg, log, req_id) + def _monitor_subprocess(self): + try: + super()._monitor_subprocess() + finally: + self._stop_xcom_backend() + + def _start_xcom_backend_call( + self, call: XComBackendCall, finish: Callable[[XComBackendResult], None] + ) -> None: + """Send ``call`` to the XCom backend process, which starts on the first call it handles.""" + if self._xcom_backend is not None and not self._xcom_backend.can_run_a_call: + # The process has exited, so a call sent to it would wait for an answer that never comes. + self._stop_xcom_backend() + if self._xcom_backend is None: + self._xcom_backend = XComBackendProcess.start( + id=self.id, + client=self.client, + logger=self.process_log, + selector=self.selector, + ) + self._xcom_backend.send_call(call, finish) + + def _stop_xcom_backend(self) -> None: + """Stop the XCom backend process, so that the next call starts one that can answer it.""" + if (backend := self._xcom_backend) is not None: + self._xcom_backend = None + backend.stop() + + def _call_xcom_backend( + self, + call: XComBackendCall, + then: Callable[[XComBackendResult], RequestResult], + req_id: int, + ) -> ResponseSent: + """ + Answer request ``req_id`` with ``then(result)`` once the XCom backend has run ``call``. + + The request loop goes on servicing the task and the backend process while the call runs, so it + keeps heartbeating and stops the task when it reaches its ``execution_timeout``. A call that + fails fails the request, so the runtime fails the task. + """ + + def finish(result: XComBackendResult) -> None: + try: + if result.error is not None: + raise XComBackendError(result.error) + resp, dump_opts = then(result) + except Exception as e: + self.process_log.exception("Failed to handle the XCom backend result", request_id=req_id) + self.send_msg(None, request_id=req_id, error=_build_request_error(e)) + else: + self.send_msg(resp, request_id=req_id, error=None, **dump_opts) + + self._start_xcom_backend_call(call, finish) + return ResponseSent.ALREADY_SENT + + def _run_xcom_backend_call(self, call: XComBackendCall) -> None: + """ + Run ``call`` and wait for it, for a caller with no request loop to return to. + + :raises XComBackendError: when the call fails or is not answered in time. + """ + answer: list[XComBackendResult] = [] + self._start_xcom_backend_call(call, answer.append) + deadline = time.monotonic() + _XCOM_BACKEND_CALL_TIMEOUT + # The task process has not been told to start, so servicing it here reads no request of its own. + while not answer and (now := time.monotonic()) < deadline: + self._service_subprocess(max_wait_time=deadline - now) + if not answer: + # The process is still running the call, and what it does next is unknown, so it is replaced. + self._stop_xcom_backend() + raise XComBackendError("The XCom backend did not answer in time") + if (error := answer[0].error) is not None: + raise XComBackendError(error) + + def _handle_set_xcom( + self, msg: SetXCom, log: FilteringBoundLogger, req_id: int + ) -> RequestResult | ResponseSent: + xcom = _get_xcom_backend() + if not (is_stored_through_backend(msg.key) and not _is_extra_link(msg.key) and runs_in_helper(xcom)): + return handle_set_xcom(self.client, msg) + ti = cast("_OwnTaskInstance", self.ti) + if is_own_xcom(ti, msg.dag_id, msg.run_id, msg.task_id, msg.resolved_map_index): + coordinate = build_own_coordinate(ti) + else: + coordinate = XComCoordinate(msg.dag_id, msg.run_id, msg.task_id) + # A loop iteration has map index -1, unlike a mapped task. A mapped task inside a loop is not + # caught here because the supervisor does not parse the Dag, while Python rejects such a Dag + # when it is parsed. + if coordinate.region_id is not None and msg.resolved_map_index == -1 and is_legacy_xcom_backend(xcom): + error = _build_request_error(build_loop_error(xcom, f"Task {msg.task_id!r}")) + self.send_msg(None, request_id=req_id, error=error) + return ResponseSent.ALREADY_SENT + return self._call_xcom_backend( + XComBackendCall.store(coordinate, msg.resolved_map_index, msg.key, msg.value), + lambda stored: handle_set_xcom(self.client, msg.model_copy(update={"value": stored.value})), + req_id, + ) + + def _handle_get_xcom( + self, msg: GetXCom, log: FilteringBoundLogger, req_id: int + ) -> RequestResult | ResponseSent: + result, dump_opts = handle_get_xcom(self.client, msg) + if not (isinstance(result, XComResult) and result.value is not None and _is_loaded_through_backend()): + return result, dump_opts + return self._call_xcom_backend( + XComBackendCall(hook="load", key=msg.key, value=result.value), + lambda loaded: (result.model_copy(update={"value": loaded.value}), dump_opts), + req_id, + ) + + def _handle_get_xcom_sequence_item( + self, msg: GetXComSequenceItem, log: FilteringBoundLogger, req_id: int + ) -> RequestResult | ResponseSent: + result, dump_opts = handle_get_xcom_sequence_item(self.client, msg) + if not (isinstance(result, XComSequenceIndexResult) and _is_loaded_through_backend()): + return result, dump_opts + return self._call_xcom_backend( + XComBackendCall(hook="load", key=msg.key, value=result.root), + lambda loaded: (XComSequenceIndexResult(root=loaded.value), dump_opts), + req_id, + ) + + def _handle_get_xcom_sequence_slice( + self, msg: GetXComSequenceSlice, log: FilteringBoundLogger, req_id: int + ) -> RequestResult | ResponseSent: + result, dump_opts = handle_get_xcom_sequence_slice(self.client, msg) + if not (isinstance(result, XComSequenceSliceResult) and result.root and _is_loaded_through_backend()): + return result, dump_opts + return self._call_xcom_backend( + XComBackendCall(hook="load_many", key=msg.key, value=result.root), + lambda loaded: ( + XComSequenceSliceResult(root=cast("list[JsonValue]", loaded.value)), + dump_opts, + ), + req_id, + ) + + def _handle_delete_xcom( + self, msg: DeleteXCom, log: FilteringBoundLogger, req_id: int + ) -> RequestResult | ResponseSent: + result, purge = self._delete_xcom(msg) + if purge is None: + return result + # A failed purge fails the request, as it fails a Python task that calls XCom.delete. + return self._call_xcom_backend(purge, lambda _: result, req_id) + + def _delete_xcom(self, msg: DeleteXCom) -> tuple[RequestResult, XComBackendCall | None]: + """ + Delete an XCom, and return the call that purges it from the XCom backend as ``XCom.delete`` does. + + The read is not guarded, for the same reason ``XCom.delete`` does not guard it: a row the + supervisor cannot read is a row it cannot purge, and silently skipping the purge would leave + the stored data behind without anything saying so. + """ + if not runs_in_helper(_get_xcom_backend()): + return handle_delete_xcom(self.client, msg), None + row, _ = handle_get_xcom( + self.client, + GetXCom( + key=msg.key, + dag_id=msg.dag_id, + run_id=msg.run_id, + task_id=msg.task_id, + map_index=msg.map_index, + ), + ) + result = handle_delete_xcom(self.client, msg) + if isinstance(row, XComResult) and row.value is not None: + return result, XComBackendCall(hook="purge", key=msg.key, value=row.value) + return result, None + + _request_handlers: ClassVar[dict[type[BaseModel], RequestHandler[Any]]] = { + **ActivitySubprocess._request_handlers, + **dict( + [ + register_request_method(DeleteXCom, _handle_delete_xcom), + register_request_method(GetXCom, _handle_get_xcom), + register_request_method(GetXComSequenceItem, _handle_get_xcom_sequence_item), + register_request_method(GetXComSequenceSlice, _handle_get_xcom_sequence_slice), + register_request_method(SetXCom, _handle_set_xcom), + ] + ), + } + + +def _build_request_error(exc: Exception) -> ErrorResponse: + """Report ``exc`` to the runtime as an XCom backend failure rather than an API server one.""" + return ErrorResponse( + error=ErrorType.XCOM_BACKEND_ERROR, + detail={"status_code": None, "message": str(exc), "exception_type": type(exc).__name__}, + ) + + +def _get_xcom_backend() -> type[BaseXCom]: + # Lazy: resolving the configured backend imports it, which only a worker that runs a task needs. Review Comment: Removed. [`54db61b`](https://github.com/apache/airflow/pull/75184/commits/54db61bffa508ce77632e858f15c0e1793bd054a) ########## scripts/ci/prek/sync_java_sdk_supervisor_schema.py: ########## @@ -16,65 +16,110 @@ # specific language governing permissions and limitations # under the License. """ -Keep the Java SDK's bundled Supervisor Schema in sync with ``airflowSupervisorSchemaVersion``. +Refresh the Java SDK's vendored copy of the supervisor schema that the Python Task SDK owns. -The Gradle task ``:sdk:syncSupervisorSchema`` downloads a fresh ``schema.json`` when the -``api_version`` inside it differs from the version declared in ``java-sdk/gradle.properties``. -Starting Gradle for that comparison costs over a minute on a cold CI runner (wrapper download, -JVM start, plugin resolution), so this hook does the same comparison in Python first and only -hands over to Gradle when the two versions actually differ. +``java-sdk`` vendors the schema so it builds from a source distribution with no network, the +way ``go-sdk`` and ``ts-sdk`` vendor theirs. It used to refresh that copy by downloading the +version named in ``java-sdk/gradle.properties``, which has two problems: the remote only +publishes released versions, so an in-progress one cannot be fetched at all, and a schema +change that leaves ``api_version`` alone moved nothing, so the copy went stale unnoticed. + +The copy is now refreshed from the schema in this repository and compared by content, so any +change to the Python source moves it in the same commit. Copying is mechanical, so this copies +for you and fails; ``git add`` the refreshed copy and recommit. ``airflowSupervisorSchemaVersion`` +stays the declared version and is asserted against the source rather than used to fetch it. + +Outside the monorepo the source is absent and the copy is already what the distribution +shipped, so the download is left as the fallback. + +Run it from the repo root, through prek: + + prek run sync-java-sdk-supervisor-schema + +Exits 0 when the copy already matched its source, 1 when it refreshed it. """ from __future__ import annotations import json import re +import shutil import subprocess import sys from pathlib import Path REPO_ROOT = Path(__file__).resolve().parents[3] -JAVA_SDK_DIR = REPO_ROOT / "java-sdk" -GRADLE_PROPERTIES = JAVA_SDK_DIR / "gradle.properties" -SCHEMA_FILE = JAVA_SDK_DIR / "sdk" / "schema" / "schema.json" + +GRADLE_PROPERTIES = Path("java-sdk/gradle.properties") +VENDORED_SCHEMA = Path("java-sdk/sdk/schema/schema.json") +SOURCE_SCHEMA = Path("task-sdk/src/airflow/sdk/execution_time/schema/schema.json") VERSION_PATTERN = re.compile(r"^airflowSupervisorSchemaVersion\s*=\s*(\S+)\s*$", re.MULTILINE) -def configured_version() -> str: - match = VERSION_PATTERN.search(GRADLE_PROPERTIES.read_text()) +def configured_version(repo_root: Path = REPO_ROOT) -> str: + """Return the version ``gradle.properties`` declares, which the schema has to agree with.""" + properties = repo_root / GRADLE_PROPERTIES + match = VERSION_PATTERN.search(properties.read_text()) if not match: - print(f"airflowSupervisorSchemaVersion is not set in {GRADLE_PROPERTIES}", file=sys.stderr) - sys.exit(1) + raise SystemExit(f"airflowSupervisorSchemaVersion is not set in {GRADLE_PROPERTIES}") return match.group(1) -def bundled_version() -> str | None: - if not SCHEMA_FILE.exists(): +def api_version_of(path: Path) -> str | None: + """Return the ``api_version`` of a schema file, or None when it is missing or unreadable.""" + if not path.is_file(): return None try: - with SCHEMA_FILE.open() as schema: + with path.open() as schema: return json.load(schema).get("api_version") except json.JSONDecodeError: - # A truncated or conflict-marked file counts as drift; Gradle overwrites it. + # A truncated or conflict-marked file counts as drift; it is overwritten. return None -def main() -> int: - expected = configured_version() - actual = bundled_version() - if actual == expected: - print(f"Supervisor Schema is up-to-date (api_version={expected}).") - return 0 - print( - f"Supervisor Schema api_version={actual!r} differs from configured {expected!r}, syncing via Gradle." +def refresh_from_source(expected: str, repo_root: Path = REPO_ROOT) -> tuple[int, str]: Review Comment: Agreed. The hook now copies by content and only checks the pin is a version the bundle defines. [`7af6b46`](https://github.com/apache/airflow/pull/75942/commits/7af6b460abbea522f8d2702da43f11f736f7edd2) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
