Lee-W commented on code in PR #68283: URL: https://github.com/apache/airflow/pull/68283#discussion_r3496625142
########## providers/common/io/src/airflow/providers/common/io/state_store/backend.py: ########## @@ -0,0 +1,247 @@ +# 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 _compression_suffix() -> str: Review Comment: ```suggestion def _get_compression_suffix() -> str: ``` nit ########## providers/common/io/src/airflow/providers/common/io/state_store/backend.py: ########## @@ -0,0 +1,247 @@ +# 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 _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: Review Comment: ```suggestion def _build_task_path(scope: TaskScope, key: str) -> ObjectStoragePath: ``` nit ########## providers/common/io/src/airflow/providers/common/io/state_store/backend.py: ########## @@ -0,0 +1,247 @@ +# 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 _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: Review Comment: ```suggestion def _build_asset_path(scope: AssetScope, key: str) -> ObjectStoragePath: ``` nit ########## providers/common/io/src/airflow/providers/common/io/state_store/backend.py: ########## @@ -0,0 +1,247 @@ +# 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 _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: + _scope_path(scope, key).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)) Review Comment: ```suggestion asset_identifier = _safe_segment(scope.name or scope.uri or str(scope.asset_id)) ``` ########## providers/common/io/src/airflow/providers/common/io/state_store/backend.py: ########## @@ -0,0 +1,247 @@ +# 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 _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: Review Comment: ```suggestion def _sanitise_segment(value: str) -> str: ``` nit ########## providers/common/io/src/airflow/providers/common/io/state_store/backend.py: ########## @@ -0,0 +1,247 @@ +# 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 _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: Review Comment: ```suggestion def _write_to_object_storage(path: ObjectStoragePath, value: str) -> None: ``` ########## providers/common/io/tests/unit/common/io/state_store/test_backend.py: ########## @@ -0,0 +1,266 @@ +# 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 pytest + +from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS + +if not AIRFLOW_V_3_3_PLUS: + pytest.skip("Store backend requires Airflow >= 3.3", allow_module_level=True) + +from airflow.providers.common.io.state_store import backend +from airflow.providers.common.io.state_store.backend import ( + StateStoreObjectStorageBackend, + _asset_path, + _read, + _task_path, + _write, +) +from airflow.sdk import ObjectStoragePath + +from tests_common.test_utils.config import conf_vars + +if AIRFLOW_V_3_3_PLUS: + from airflow.sdk.state import AssetScope, TaskScope + + [email protected](autouse=True) +def clear_caches(): + backend._get_base_path.cache_clear() + backend._get_compression.cache_clear() + backend._get_threshold.cache_clear() + yield + backend._get_base_path.cache_clear() + backend._get_compression.cache_clear() + backend._get_threshold.cache_clear() + + [email protected] +def base_path(tmp_path): + store_path = tmp_path / "store" + store_path.mkdir() + return f"file://{store_path}" + + [email protected] +def conf_overrides(base_path): + from tests_common.test_utils.config import conf_vars + + with conf_vars( + { + ("common.io", "state_store_objectstorage_path"): base_path, + ("common.io", "state_store_objectstorage_compression"): "", + } + ): + yield base_path + + [email protected](not AIRFLOW_V_3_3_PLUS, reason="task state store requires Airflow >= 3.3") +class TestPathBuilders: + def test_task_path_segments(self, conf_overrides): + scope = TaskScope(dag_id="my_dag", run_id="run_1", task_id="my_task", map_index=-1) + path = _task_path(scope, "job_id") + assert str(path).endswith("my_dag/run_1/my_task/-1/job_id") + + def test_task_path_sanitises_slashes(self, conf_overrides): + scope = TaskScope(dag_id="a/b", run_id="r/1", task_id="t/x", map_index=0) + path = _task_path(scope, "key/name") + # a/b is sanitised to a_b + assert "a/b" not in str(path) + assert "a_b" in str(path) + assert "key_name" in str(path) + assert "key/name" not in str(path) Review Comment: Instead of using `in` and `not in`, it would be better for us to just compare with the whole path. ########## providers/common/io/src/airflow/providers/common/io/state_store/backend.py: ########## @@ -0,0 +1,247 @@ +# 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 _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: Review Comment: ```suggestion def _read_from_object_storage(path: ObjectStoragePath) -> str | None: ``` ########## providers/common/io/tests/unit/common/io/state_store/test_backend.py: ########## @@ -0,0 +1,266 @@ +# 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 pytest + +from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS + +if not AIRFLOW_V_3_3_PLUS: + pytest.skip("Store backend requires Airflow >= 3.3", allow_module_level=True) + +from airflow.providers.common.io.state_store import backend +from airflow.providers.common.io.state_store.backend import ( + StateStoreObjectStorageBackend, + _asset_path, + _read, + _task_path, + _write, +) +from airflow.sdk import ObjectStoragePath + +from tests_common.test_utils.config import conf_vars + +if AIRFLOW_V_3_3_PLUS: + from airflow.sdk.state import AssetScope, TaskScope + + [email protected](autouse=True) +def clear_caches(): + backend._get_base_path.cache_clear() + backend._get_compression.cache_clear() + backend._get_threshold.cache_clear() + yield + backend._get_base_path.cache_clear() + backend._get_compression.cache_clear() + backend._get_threshold.cache_clear() + + [email protected] +def base_path(tmp_path): + store_path = tmp_path / "store" + store_path.mkdir() + return f"file://{store_path}" + + [email protected] +def conf_overrides(base_path): + from tests_common.test_utils.config import conf_vars + Review Comment: ```suggestion ``` ########## providers/common/io/src/airflow/providers/common/io/state_store/backend.py: ########## @@ -0,0 +1,247 @@ +# 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 _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: + _scope_path(scope, key).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) Review Comment: Do we want to add a safeguard here so that we won't silently miss new store (if we add it in the future) ```suggestion p.unlink(missing_ok=True) case _: raise TypeError(f"Unknown scope type: {type(scope)}") ``` ########## providers/common/io/tests/unit/common/io/state_store/test_backend.py: ########## @@ -0,0 +1,266 @@ +# 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 pytest + +from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS + +if not AIRFLOW_V_3_3_PLUS: + pytest.skip("Store backend requires Airflow >= 3.3", allow_module_level=True) + +from airflow.providers.common.io.state_store import backend +from airflow.providers.common.io.state_store.backend import ( + StateStoreObjectStorageBackend, + _asset_path, + _read, + _task_path, + _write, +) +from airflow.sdk import ObjectStoragePath + +from tests_common.test_utils.config import conf_vars + +if AIRFLOW_V_3_3_PLUS: + from airflow.sdk.state import AssetScope, TaskScope + + [email protected](autouse=True) +def clear_caches(): + backend._get_base_path.cache_clear() + backend._get_compression.cache_clear() + backend._get_threshold.cache_clear() + yield + backend._get_base_path.cache_clear() + backend._get_compression.cache_clear() + backend._get_threshold.cache_clear() + + [email protected] +def base_path(tmp_path): + store_path = tmp_path / "store" + store_path.mkdir() + return f"file://{store_path}" + + [email protected] +def conf_overrides(base_path): + from tests_common.test_utils.config import conf_vars + + with conf_vars( + { + ("common.io", "state_store_objectstorage_path"): base_path, + ("common.io", "state_store_objectstorage_compression"): "", + } + ): + yield base_path + + [email protected](not AIRFLOW_V_3_3_PLUS, reason="task state store requires Airflow >= 3.3") Review Comment: Do we still need it here given that we already have line 23-24 -- 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]
