amoghrajesh commented on code in PR #68283: URL: https://github.com/apache/airflow/pull/68283#discussion_r3497990329
########## 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: Good catch. Switched to `urllib.parse.quote(value, safe="")` which is reversible, so all distinct inputs produce distinct segments. Also added guards for empty, `.`, and `...`. Tests updated to assert collision freedom and reject the invalid inputs. -- 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]
