kaxil commented on code in PR #68283: URL: https://github.com/apache/airflow/pull/68283#discussion_r3484539884
########## providers/common/io/src/airflow/providers/common/io/state_store/backend.py: ########## @@ -0,0 +1,245 @@ +# 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 contextlib +import json +from functools import cache +from typing import TYPE_CHECKING +from urllib.parse import urlsplit + +import fsspec.utils + +from airflow.providers.common.compat.sdk import conf + +if TYPE_CHECKING: + from datetime import datetime + + from pydantic import JsonValue + from sqlalchemy.ext.asyncio import AsyncSession + from sqlalchemy.orm import Session + + +from airflow.sdk import ObjectStoragePath +from airflow.sdk._shared.state import AssetScope, BaseStoreBackend, StoreScope, TaskScope + +SECTION = "common.io" + + +@cache +def _get_base_path() -> ObjectStoragePath: + return ObjectStoragePath(conf.get_mandatory_value(SECTION, "state_store_objectstorage_path")) + + +@cache +def _get_compression() -> str | None: + value = conf.get(SECTION, "state_store_objectstorage_compression", fallback=None) + return value or None + + +@cache +def _get_threshold() -> int: + return conf.getint(SECTION, "state_store_objectstorage_threshold", fallback=0) + + +def _compression_suffix() -> str: + compression = _get_compression() + if not compression: + return "" + for suffix, c in fsspec.utils.compressions.items(): + if c == compression: + return f".{suffix}" + raise ValueError(f"Compression {compression!r} is not supported.") + + +def _safe_segment(value: str) -> str: + """ + Sanitise a string for use as a single path segment. + + This is a simple implementation that replaces slashes with underscores. + """ + return value.replace("/", "_").replace("\\", "_") + + +def _task_path(scope: TaskScope, key: str) -> ObjectStoragePath: + suffix = _compression_suffix() + return ( + _get_base_path() + / _safe_segment(scope.dag_id) + / _safe_segment(scope.run_id) + / _safe_segment(scope.task_id) + / str(scope.map_index) + / f"{_safe_segment(key)}{suffix}" + ) + + +def _asset_path(scope: AssetScope, key: str) -> ObjectStoragePath: + suffix = _compression_suffix() + asset_ref = _safe_segment(scope.name or scope.uri or str(scope.asset_id)) + return _get_base_path() / "assets" / asset_ref / f"{_safe_segment(key)}{suffix}" + + +def _write(path: ObjectStoragePath, value: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + compression = _get_compression() + with path.open(mode="wb", compression=compression) as f: + f.write(value.encode("utf-8")) + + +def _read(path: ObjectStoragePath) -> str | None: + try: + with path.open(mode="rb", compression="infer") as f: + return f.read().decode("utf-8") + except FileNotFoundError: + return None + + +def _is_storage_ref(value: str) -> bool: + try: + if not urlsplit(value).scheme: + return False + return ObjectStoragePath(value).is_relative_to(_get_base_path()) + except Exception: + return False + + +def _scope_path(scope: StoreScope, key: str) -> ObjectStoragePath: + match scope: + case TaskScope(): + return _task_path(scope, key) + case AssetScope(): + return _asset_path(scope, key) + case _: + raise TypeError(f"Unknown scope type: {type(scope)}") + + +class StateStoreObjectStorageBackend(BaseStoreBackend): + """ + Object-storage backend for task and asset store. + + Config keys (all under ``[common.io]``): + + - ``state_store_objectstorage_path``: base path, e.g. ``s3://conn_id@bucket/task-state/`` + - ``state_store_objectstorage_compression``: optional compression, e.g. ``gzip`` + """ + + def get(self, scope: StoreScope, key: str, *, session: Session | None = None) -> str | None: + return _read(_scope_path(scope, key)) + + def set( + self, + scope: StoreScope, + key: str, + value: str, + *, + expires_at: datetime | None = None, + session: Session | None = None, + ) -> None: + _write(_scope_path(scope, key), value) + + def delete(self, scope: StoreScope, key: str, *, session: Session | None = None) -> None: + path = _scope_path(scope, key) + with contextlib.suppress(FileNotFoundError): + path.unlink(missing_ok=True) + + def clear( + self, scope: StoreScope, *, all_map_indices: bool = False, session: Session | None = None + ) -> None: + match scope: + case TaskScope(): + if all_map_indices: + prefix = ( + _get_base_path() + / _safe_segment(scope.dag_id) + / _safe_segment(scope.run_id) + / _safe_segment(scope.task_id) + ) + for p in prefix.glob("*/*"): + p.unlink(missing_ok=True) + else: + prefix = ( + _get_base_path() + / _safe_segment(scope.dag_id) + / _safe_segment(scope.run_id) + / _safe_segment(scope.task_id) + / str(scope.map_index) + ) + for p in prefix.glob("*"): + p.unlink(missing_ok=True) + case AssetScope(): + asset_ref = _safe_segment(scope.name or scope.uri or str(scope.asset_id)) + prefix = _get_base_path() / "assets" / asset_ref + for p in prefix.glob("*"): + p.unlink(missing_ok=True) + + async def aget(self, scope: StoreScope, key: str, *, session: AsyncSession | None = None) -> str | None: + raise NotImplementedError + + async def aset( + self, + scope: StoreScope, + key: str, + value: str, + *, + expires_at: datetime | None = None, + session: AsyncSession | None = None, + ) -> None: + raise NotImplementedError + + async def adelete(self, scope: StoreScope, key: str, *, session: AsyncSession | None = None) -> None: + raise NotImplementedError + + async def aclear( + self, scope: StoreScope, *, all_map_indices: bool = False, session: AsyncSession | None = None + ) -> None: + raise NotImplementedError + + def serialize_task_store_to_ref(self, *, value: JsonValue, key: str, scope: TaskScope) -> str: Review Comment: After the rename in this PR, the base class and the worker SDK use these method names *with* the `state` segment, but the four methods here (this line plus L219/L229/L237) are named without it: - `BaseStoreBackend` defines `serialize_task_state_store_to_ref` / `deserialize_task_state_store_from_ref` / `serialize_asset_state_store_to_ref` / `deserialize_asset_state_store_from_ref` in `task-sdk/src/airflow/sdk/_shared/state/__init__.py:251,271,281,301`. - The worker calls those same `state` names: `task-sdk/src/airflow/sdk/execution_time/context.py:559,608,714,744`. Those base methods are concrete (not abstract) and default to `return json.dumps(value)` / `return json.loads(stored)`, so `StateStoreObjectStorageBackend` still instantiates and these overrides are never invoked. The threshold check and the offload to object storage never run; the worker stores the value inline through the base default, so the backend behaves like plain DB storage. The unit tests call these methods directly by the no-`state` names, so they pass without exercising the path the SDK actually uses. Renaming the four methods to include `state` (and routing the tests through the renamed names) wires it back up. ########## providers/common/io/docs/state_store_backend.rst: ########## @@ -0,0 +1,81 @@ +.. 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. + +Object Storage State Store Backend +=================================== + +The default state store backend is :class:`~airflow.state.metastore.MetastoreStateBackend`, which persists +task and asset state in the Airflow metadata database. For larger values, you may want to store state on +object storage instead. + +To enable object storage for task and asset state store, set ``backend`` in the ``[state_store]`` section to +``airflow.providers.common.io.state_store.backend.StateStateStoreObjectStorageBackend``, and set Review Comment: Typo: `StateStateStoreObjectStorageBackend` has a doubled `State`. The class is `StateStoreObjectStorageBackend` (`backend.py:130`). The same typo is on lines 26, 43, 52, and 61, so copy-pasting any of these `backend = ...` snippets into `airflow.cfg` raises `ImportError`. ########## providers/common/io/provider.yaml: ########## @@ -115,3 +115,31 @@ config: type: string example: "gz" default: "" + state_store_objectstorage_path: + description: | + Base path on object storage for the task/asset state store backend, in URL format. + When set, StateStoreObjectStorageBackend will persist task and asset state under this + prefix, organised as <dag_id>/<run_id>/<task_id>/<map_index>/<key> for tasks and + assets/<asset_identifier>/<key> for assets. + version_added: 1.8.0 + type: string + example: "s3://conn_id@bucket/task-state/" + default: "" + state_store_objectstorage_threshold: + description: | + Threshold in bytes for offloading serialized state store values to object storage. 0 means + always offload to object storage. Any positive number means values will be offloaded + only when their serialized size is at least that many bytes. Must be non-negative. + version_added: 1.8.0 + type: integer + example: "1000000" + default: "0" + state_store_objectstorage_compression: + description: | + Compression algorithm to use when writing task/asset state store values to object storage. + Supported algorithms are a.o.: snappy, zip, gzip, bz2, and lzma. If not specified, + no compression will be used. The same algorithm must be available on all workers. + version_added: 1.8.0 + type: string + example: "gz" Review Comment: The example `"gz"` isn't a usable value here. `_get_compression()` returns the string verbatim and `_compression_suffix()` matches it against the *values* of `fsspec.utils.compressions`, which are codec names (`gzip`, `bz2`, `lzma`, `xz`, `zip`, `lz4`, `zstd`), not extensions like `gz`. So `state_store_objectstorage_compression = gz` raises `ValueError: Compression 'gz' is not supported.` on the first read or write. Use `gzip` (the class docstring and the `.rst` already use `gzip`). Separately, `snappy` in the description above also isn't in `fsspec.utils.compressions` unless `python-snappy` is installed, so it fails the same way out of the box. -- 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]
