uranusjr commented on code in PR #75498: URL: https://github.com/apache/airflow/pull/75498#discussion_r4397885536
########## airflow-core/src/airflow/dag_processing/bundles/filesystem_artifact_store.py: ########## @@ -0,0 +1,473 @@ +# 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. +"""Shared-filesystem artifact store for manifest-backed Dag bundles, and its filesystem helpers.""" + +from __future__ import annotations + +import json +import logging +import os +import shutil +import stat +import tempfile +from contextlib import contextmanager, suppress +from fcntl import LOCK_EX, LOCK_UN, flock +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from airflow.dag_processing.bundles.artifact_store import ( + ArtifactStore, + BundleManifestNotFoundError, + BundleManifestRefDurabilityError, +) +from airflow.dag_processing.bundles.manifest import ( + MANIFEST_FILE_NAME, + SHA256_VERSION_PREFIX, + BundleManifestError, + BundleManifestSourceChangedError, + compute_file_sha256, + is_sha256_hex, + serialize_bundle_version_manifest, + validate_bundle_relative_path, +) + +if TYPE_CHECKING: + from collections.abc import Callable + +log = logging.getLogger(__name__) + + +class FilesystemArtifactStore(ArtifactStore): + """ + Artifact store for a ``published_root`` on a filesystem shared by publishers and Airflow. + + Layout under ``published_root``: ``refs/<bundle>/latest.json`` (release reference), + ``versions/<bundle>/<version>/`` (snapshots) and ``_locks/<bundle>.lock`` (publication lock). + + Publishers need write access to the publication directories and read/write access to + the lock file. For different OS users, provision these with a shared publisher group. + """ + + def __init__(self, *, bundle_name: str, published_root: Path, cache_root: Path) -> None: + self.bundle_name = bundle_name + if paths_overlap(published_root, cache_root): + raise ValueError( + "published_root must not overlap Airflow's bundle cache. Configure separate " + "authoritative publication and dag_bundle_storage_path locations." + ) + self._root = published_root + self._ref_path = published_root / "refs" / bundle_name / "latest.json" + self._snapshots_root = published_root / "versions" / bundle_name + self.lock_path = published_root / "_locks" / f"{bundle_name}.lock" + + @property + def root(self) -> Path: + return self._root + + @property + def ref_path(self) -> Path: + return self._ref_path + + @property + def snapshots_root(self) -> Path: + return self._snapshots_root + + def snapshot_path(self, version: str) -> Path: + return self._snapshots_root / version + + def read_ref(self, *, missing_message: str, invalid_message: str) -> dict[str, Any]: + return read_json_file( + self._ref_path, missing_message=missing_message, invalid_message=invalid_message + ) + + def write_ref(self, payload: dict[str, Any]) -> None: + _write_json_atomically(self._ref_path, payload) + + @contextmanager + def publication_guard(self): + """Serialize publishers through a lock stored with the authoritative snapshots.""" + _ensure_public_dir(self.lock_path.parent) + # Linux NFS implements flock with fcntl locks, which require a writable + # descriptor for an exclusive lock. + try: + fd = os.open(self.lock_path, os.O_CREAT | os.O_RDWR, 0o664) + except PermissionError as e: + raise BundleManifestError( + f"Cannot open publication lock {self.lock_path}: publishers need read/write " + "access. Set its group to a shared publisher group and grant group read/write " + "access (chmod 0664), and ensure the lock directory is writable." + ) from e + with os.fdopen(fd, "r+b") as lock_file: + # os.open masks the mode with the umask; repair it whenever this publisher + # owns the file so other members of the publisher group can acquire it. + with suppress(OSError): + os.fchmod(fd, stat.S_IMODE(os.fstat(fd).st_mode) | 0o664) + flock(lock_file, LOCK_EX) + try: + yield + finally: + flock(lock_file, LOCK_UN) + + def prepare_publish_areas(self) -> None: + _ensure_public_dir(self._snapshots_root) + _ensure_public_dir(self._ref_path.parent) + + def validate_source_paths(self, source_path: Path, *, cache_versions_dir: Path) -> None: + _validate_publish_paths( + source_path=source_path, + published_root=self._root, + versions_dir=cache_versions_dir, + ) + + def snapshot_exists(self, version: str) -> bool: + return (self._snapshots_root / version).is_dir() + + def fetch_snapshot( + self, + version: str, + destination: Path, + *, + structural_validator: Callable[[Path], None], + ) -> None: + published_version_path = self._snapshots_root / version + if not published_version_path.is_dir(): + raise BundleManifestNotFoundError( + f"Bundle '{self.bundle_name}' version '{version}' is not published at " + f"{published_version_path}. Publish or restore the immutable snapshot before " + "running pinned work." + ) + structural_validator(published_version_path) + shutil.copytree(published_version_path, destination, copy_function=shutil.copy2, dirs_exist_ok=True) + + def publish_snapshot( + self, + version: str, + *, + manifest: dict[str, Any], + source_root: Path, + validate_existing: Callable[[Path], None], + ) -> bool: + version_path = self._snapshots_root / version + if version_path.exists(): + if not version_path.is_dir(): + raise FileExistsError(f"Bundle snapshot path exists but is not a directory: {version_path}") + validate_existing(version_path) + return False + _materialize_manifest_snapshot( + source_path=source_root, + manifest=manifest, + versions_dir=self._snapshots_root, + version_path=version_path, + ) + return True + + def sweep_publish_temps(self) -> None: + remove_orphaned_temp_snapshots(self._snapshots_root) + _remove_orphaned_atomic_temp_files(self._ref_path) + + +def read_json_file(path: Path, *, missing_message: str, invalid_message: str) -> dict[str, Any]: + """Read a JSON object from a published document, raising ``missing_message`` if it is absent.""" + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError as e: + raise BundleManifestNotFoundError(missing_message) from e Review Comment: Other OSErrors escape raw. #72476 deliberately wraps every source OSError as a BundleManifestError. Follow the same approach here. -- 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]
