kaxil commented on code in PR #68283: URL: https://github.com/apache/airflow/pull/68283#discussion_r3497601548
########## providers/common/io/src/airflow/providers/common/io/state_store/backend.py: ########## @@ -0,0 +1,249 @@ +# 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 +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: + value = conf.getint(SECTION, "state_store_objectstorage_threshold", fallback=0) + if value < 0: + raise ValueError( + f"[{SECTION}] state_store_objectstorage_threshold must be non-negative, got {value}." + ) + return value + + +def _get_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 _sanitise_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("\\", "_") Review Comment: `_sanitise_segment` replacing `/`->`_` is lossy, so two distinct scopes can map to the same path. It's not reachable through `dag_id`/`task_id` (both go through `validate_key`, restricted to `[\w.-]`, no slash), but it is through the two free-form segments: - the state-store **key**: `TaskStateStoreAccessor.set()` (`task-sdk/src/airflow/sdk/execution_time/context.py:570`) passes `key` straight through with no validation, so `key="connector/offset"` and `key="connector_offset"` both build `.../connector_offset`; - the **asset identifier**: `scope.name or scope.uri` (L96 / L186), where free-form URIs like `s3://bucket/x` and `s3://bucket_x` both build `assets/s3:__bucket_x`. Colliding scopes silently overwrite and read each other's state, and since `serialize_*_to_ref` builds the offload path with the same helper, `delete`/`clear` then act on the wrong object. Separately, `..` / `.` / `""` pass through unchanged: `key=".."` builds `.../<map_index>/..`, which raises `IsADirectoryError` on the local filesystem and resolves to a sibling prefix on object stores (still under the base path, so `_is_storage_ref` won't flag it). Suggest a reversible per-segment encoding such as `urllib.parse.quote(value, safe="")`, plus rejecting empty/`.`/`..`. The XCom backend sidesteps this by using a `uuid4`-based path, so there's no parity to preserve here. (Minor: `test_build_task_path_sanitises_slashes` only asserts the slash is gone, not that two distinct inputs stay distinct, so it won't catch the collision.) -- 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]
