"""数据处理原始源文件的受控本地对象存储。""" from __future__ import annotations import os import re import stat import unicodedata import uuid from dataclasses import dataclass from functools import lru_cache from pathlib import Path, PurePosixPath from typing import Iterable, Iterator from urllib.parse import quote, unquote, urlsplit class DataProcessStorageError(ValueError): """本地对象引用或文件系统状态不安全。""" @dataclass(frozen=True, slots=True) class StagedSourceObject: """尚未发布的原始文件；绝对路径仅在存储模块内部流转。""" reference: str _temporary_path: Path _relative_path: PurePosixPath def _default_storage_root() -> Path: return Path(__file__).resolve().parents[3] / "storage" / "data-process" def _configured_storage_root() -> Path: configured = os.getenv("DATA_PROCESS_STORAGE_DIR", "").strip() if not configured: return _default_storage_root() path = Path(configured).expanduser() # 相对配置固定以 backend 目录为基准， # 避免从不同 cwd 启动时写入不同位置。 return path if path.is_absolute() else Path(__file__).resolve().parents[3] / path def _safe_component(value: str, label: str) -> str: if not value or value in {".", ".."} or len(value) > 128: raise DataProcessStorageError(f"invalid {label}") if not value[0].isalnum() or any( not (character.isalnum() or character in {"-", "_", "."}) for character in value ): raise DataProcessStorageError(f"invalid {label}") return value def _safe_basename(value: str) -> str: if not value or len(value.encode("utf-8")) > 255: raise DataProcessStorageError("invalid source file name") if value != Path(value).name or "/" in value or "\\" in value or "\x00" in value: raise DataProcessStorageError("invalid source file name") if value in {".", ".."} or any( unicodedata.category(character).startswith("C") for character in value ): raise DataProcessStorageError("invalid source file name") return value class LocalDataProcessStorage: """只允许访问配置根目录下的版本化原始文件。""" def __init__(self, root: str | os.PathLike[str] | Path | None = None) -> None: configured = Path(root) if root is not None else _configured_storage_root() configured = configured.expanduser() if configured.exists() and configured.is_symlink(): raise DataProcessStorageError("data process storage root must not be a symlink") configured.mkdir(parents=True, exist_ok=True, mode=0o700) self._root = configured.resolve(strict=True) # StagedSourceObject 本身是普通 dataclass，不能只依赖其中的路径字段判断 # 来源；只接受由当前存储实例实际签发的对象， # 避免调用方伪造暂存路径。 self._issued_staged_objects: dict[Path, StagedSourceObject] = {} self._ensure_directory(self._root / ".staging") @property def root(self) -> Path: """仅供运维和测试检查；API 响应不得序列化该属性。""" return self._root def new_batch_id(self) -> str: return f"batch-{uuid.uuid4().hex}" def stage_bytes( self, *, batch_id: str, task_id: str, source_file_id: str, version: int, name: str, content: bytes, ) -> StagedSourceObject: batch_id = _safe_component(batch_id, "batch id") task_id = _safe_component(task_id, "task id") source_file_id = _safe_component(source_file_id, "source file id") if isinstance(version, bool) or not isinstance(version, int) or version < 1: raise DataProcessStorageError("invalid source file version") basename = _safe_basename(name) if not isinstance(content, bytes): raise TypeError("content must be bytes") batch_directory = self._ensure_directory(self._root / ".staging" / batch_id) temporary_path = batch_directory / f"{source_file_id}-{uuid.uuid4().hex}.tmp" flags = os.O_CREAT | os.O_EXCL | os.O_WRONLY if hasattr(os, "O_NOFOLLOW"): flags |= os.O_NOFOLLOW descriptor = os.open(temporary_path, flags, 0o600) try: with os.fdopen(descriptor, "wb", closefd=True) as stream: stream.write(content) stream.flush() os.fsync(stream.fileno()) except Exception: temporary_path.unlink(missing_ok=True) raise relative_path = PurePosixPath( task_id, source_file_id, f"v{version}", basename, ) reference = ( "local://data-process/" f"{task_id}/{source_file_id}/v{version}/{quote(basename, safe='')}" ) staged = StagedSourceObject(reference, temporary_path, relative_path) self._issued_staged_objects[temporary_path] = staged return staged def stage_copy( self, *, batch_id: str, source_reference: str, expected_source_task_id: str, expected_source_file_id: str, task_id: str, source_file_id: str, version: int, name: str, ) -> StagedSourceObject: """为不可变源对象创建独立目录项，不把大文件重新读入内存。""" batch_id = _safe_component(batch_id, "batch id") task_id = _safe_component(task_id, "task id") source_file_id = _safe_component(source_file_id, "source file id") if isinstance(version, bool) or not isinstance(version, int) or version < 1: raise DataProcessStorageError("invalid source file version") basename = _safe_basename(name) source_relative = self._relative_from_reference(source_reference) if source_relative is None: raise DataProcessStorageError("original source object is not available") self._assert_expected_owner( source_relative, expected_task_id=expected_source_task_id, expected_source_file_id=expected_source_file_id, ) descriptor, source_info = self._open_read_descriptor(source_relative) os.close(descriptor) batch_directory = self._ensure_directory(self._root / ".staging" / batch_id) temporary_path = batch_directory / f"{source_file_id}-{uuid.uuid4().hex}.tmp" source_path = self._path_for_relative(source_relative) try: os.link(source_path, temporary_path, follow_symlinks=False) copy_info = temporary_path.lstat() if ( not stat.S_ISREG(copy_info.st_mode) or source_info.st_dev != copy_info.st_dev or source_info.st_ino != copy_info.st_ino ): raise DataProcessStorageError("source storage object changed while copying") except Exception: temporary_path.unlink(missing_ok=True) raise relative_path = PurePosixPath( task_id, source_file_id, f"v{version}", basename, ) reference = ( "local://data-process/" f"{task_id}/{source_file_id}/v{version}/{quote(basename, safe='')}" ) staged = StagedSourceObject(reference, temporary_path, relative_path) self._issued_staged_objects[temporary_path] = staged return staged def publish(self, objects: Iterable[StagedSourceObject]) -> None: staged = list(objects) published: list[StagedSourceObject] = [] try: seen_temporary_paths: set[Path] = set() for item in staged: self._validate_staged_object(item, require_file=True) if item._temporary_path in seen_temporary_paths: raise DataProcessStorageError("duplicate staged source object") seen_temporary_paths.add(item._temporary_path) for item in staged: final_path = self._path_for_relative(item._relative_path) self._ensure_directory(final_path.parent) if final_path.exists() or final_path.is_symlink(): raise DataProcessStorageError("source storage object already exists") os.link(item._temporary_path, final_path, follow_symlinks=False) published.append(item) item._temporary_path.unlink() self._fsync_directory(final_path.parent) except Exception: for item in reversed(published): try: self.delete(item.reference) except Exception: # 回滚必须尽量处理其余对象，并保留真正的发布异常。 pass for item in staged: try: self.discard([item]) except Exception: pass raise self.discard(staged) def discard(self, objects: Iterable[StagedSourceObject]) -> None: staged = list(objects) for item in staged: self._validate_staged_object(item, require_file=False) batch_directories: set[Path] = set() first_error: Exception | None = None for item in staged: temporary_path = item._temporary_path try: temporary_path.unlink(missing_ok=True) except Exception as exc: if first_error is None: first_error = exc else: self._issued_staged_objects.pop(temporary_path, None) batch_directories.add(temporary_path.parent) for directory in batch_directories: self._remove_empty_directory(directory) if first_error is not None: raise first_error def read(self, reference: str) -> bytes | None: """读取 local 引用；旧 ``db://`` 对象返回 ``None`` 由数据库正文兜底。""" relative_path = self._relative_from_reference(reference) if relative_path is None: return None descriptor, _ = self._open_read_descriptor(relative_path) with os.fdopen(descriptor, "rb", closefd=True) as stream: return stream.read() def file_size( self, reference: str, *, expected_task_id: str, expected_source_file_id: str, ) -> int | None: """返回受控 local 对象大小；旧 ``db://`` 对象没有原始文件。""" relative_path = self._relative_from_reference(reference) if relative_path is None: return None self._assert_expected_owner( relative_path, expected_task_id=expected_task_id, expected_source_file_id=expected_source_file_id, ) descriptor, info = self._open_read_descriptor(relative_path) os.close(descriptor) return info.st_size def iter_bytes( self, reference: str, *, expected_task_id: str, expected_source_file_id: str, expected_size: int, start: int = 0, length: int | None = None, chunk_size: int = 256 * 1024, ) -> Iterator[bytes]: """按范围流式读取原始文件，避免 PDF 预览把大文件整体载入内存。""" relative_path = self._relative_from_reference(reference) if relative_path is None: raise DataProcessStorageError("original source object is not available") self._assert_expected_owner( relative_path, expected_task_id=expected_task_id, expected_source_file_id=expected_source_file_id, ) if start < 0 or expected_size < 0 or chunk_size < 1: raise DataProcessStorageError("invalid source byte range") descriptor, info = self._open_read_descriptor(relative_path) if info.st_size != expected_size: os.close(descriptor) raise DataProcessStorageError("source object size does not match metadata") remaining = expected_size - start if length is None else length if remaining < 0 or start + remaining > expected_size: os.close(descriptor) raise DataProcessStorageError("invalid source byte range") with os.fdopen(descriptor, "rb", closefd=True) as stream: stream.seek(start) while remaining: chunk = stream.read(min(chunk_size, remaining)) if not chunk: raise DataProcessStorageError("source object ended unexpectedly") remaining -= len(chunk) yield chunk def validate_owner( self, reference: str, *, expected_task_id: str, expected_source_file_id: str, ) -> bool: """校验 local 引用归属；旧 ``db://`` 引用无需文件系统处理。""" relative_path = self._relative_from_reference(reference) if relative_path is None: return False self._assert_expected_owner( relative_path, expected_task_id=expected_task_id, expected_source_file_id=expected_source_file_id, ) return True def _open_read_descriptor( self, relative_path: PurePosixPath, ) -> tuple[int, os.stat_result]: path = self._path_for_relative(relative_path) self._assert_controlled_parent(path) try: before_open = path.lstat() except FileNotFoundError as exc: raise DataProcessStorageError("source storage object does not exist") from exc if stat.S_ISLNK(before_open.st_mode) or not stat.S_ISREG(before_open.st_mode): raise DataProcessStorageError("source storage object is not a regular file") flags = os.O_RDONLY if hasattr(os, "O_NOFOLLOW"): flags |= os.O_NOFOLLOW descriptor = os.open(path, flags) after_open = os.fstat(descriptor) if ( not stat.S_ISREG(after_open.st_mode) or before_open.st_dev != after_open.st_dev or before_open.st_ino != after_open.st_ino ): os.close(descriptor) raise DataProcessStorageError("source storage object changed while opening") return descriptor, after_open def delete( self, reference: str, *, expected_task_id: str | None = None, expected_source_file_id: str | None = None, ) -> bool: """删除受控 local 对象；旧 ``db://`` 引用保持不变。""" relative_path = self._relative_from_reference(reference) if relative_path is None: return False if (expected_task_id is None) != (expected_source_file_id is None): raise DataProcessStorageError("both expected storage owner fields are required") if expected_task_id is not None and expected_source_file_id is not None: self._assert_expected_owner( relative_path, expected_task_id=expected_task_id, expected_source_file_id=expected_source_file_id, ) path = self._path_for_relative(relative_path) self._assert_controlled_parent(path) try: info = path.lstat() except FileNotFoundError: return False if stat.S_ISLNK(info.st_mode) or not stat.S_ISREG(info.st_mode): raise DataProcessStorageError("refusing to delete a non-regular storage object") path.unlink() self._fsync_directory(path.parent) for directory in (path.parent, path.parent.parent, path.parent.parent.parent): self._remove_empty_directory(directory) return True @staticmethod def _assert_expected_owner( relative_path: PurePosixPath, *, expected_task_id: str, expected_source_file_id: str, ) -> None: task_id = _safe_component(expected_task_id, "expected task id") source_file_id = _safe_component( expected_source_file_id, "expected source file id", ) if relative_path.parts[:2] != (task_id, source_file_id): raise DataProcessStorageError("source storage object owner mismatch") def _relative_from_reference(self, reference: str) -> PurePosixPath | None: if reference.startswith("db://"): return None parsed = urlsplit(reference) if parsed.scheme != "local" or parsed.netloc != "data-process": raise DataProcessStorageError("unsupported source storage reference") if parsed.query or parsed.fragment or "\\" in parsed.path: raise DataProcessStorageError("unsafe source storage reference") raw_parts = parsed.path.lstrip("/").split("/") if len(raw_parts) != 4: raise DataProcessStorageError("unsafe source storage reference") if any(re.search(r"%(?![0-9A-Fa-f]{2})", part) for part in raw_parts): raise DataProcessStorageError("unsafe source storage reference") try: decoded = [unquote(part, encoding="utf-8", errors="strict") for part in raw_parts] except UnicodeDecodeError as exc: raise DataProcessStorageError("unsafe source storage reference") from exc if any("/" in part or "\\" in part for part in decoded): raise DataProcessStorageError("unsafe source storage reference") canonical_parts = [ quote(decoded[0], safe="-_."), quote(decoded[1], safe="-_."), quote(decoded[2], safe="-_."), quote(decoded[3], safe=""), ] if canonical_parts != raw_parts: raise DataProcessStorageError("source storage reference is not canonical") task_id = _safe_component(decoded[0], "task id") source_file_id = _safe_component(decoded[1], "source file id") version_component = decoded[2] if not version_component.startswith("v") or not version_component[1:].isdigit(): raise DataProcessStorageError("invalid source file version") version = int(version_component[1:]) if version < 1: raise DataProcessStorageError("invalid source file version") basename = _safe_basename(decoded[3]) return PurePosixPath(task_id, source_file_id, f"v{version}", basename) def _path_for_relative(self, relative_path: PurePosixPath) -> Path: if relative_path.is_absolute() or any( part in {"", ".", ".."} for part in relative_path.parts ): raise DataProcessStorageError("storage path escapes the configured root") path = self._root.joinpath(*relative_path.parts) self._assert_controlled_parent(path) return path def _validate_staged_object( self, item: StagedSourceObject, *, require_file: bool, ) -> None: if not isinstance(item, StagedSourceObject): raise DataProcessStorageError("invalid staged source object") if self._issued_staged_objects.get(item._temporary_path) is not item: raise DataProcessStorageError("staged source object was not issued by this storage") expected_relative = self._relative_from_reference(item.reference) if expected_relative is None or expected_relative != item._relative_path: raise DataProcessStorageError("staged source object reference mismatch") staging_root = self._root / ".staging" try: relative_temporary = item._temporary_path.relative_to(staging_root) except ValueError as exc: raise DataProcessStorageError("staged source object escapes staging") from exc if len(relative_temporary.parts) != 2: raise DataProcessStorageError("invalid staged source object path") _safe_component(relative_temporary.parts[0], "batch id") _safe_basename(relative_temporary.parts[1]) self._assert_controlled_parent(item._temporary_path) try: info = item._temporary_path.lstat() except FileNotFoundError: if require_file: raise DataProcessStorageError("staged source object does not exist") from None return if stat.S_ISLNK(info.st_mode) or not stat.S_ISREG(info.st_mode): raise DataProcessStorageError("staged source object is not a regular file") def _ensure_directory(self, directory: Path) -> Path: try: relative = directory.relative_to(self._root) except ValueError as exc: raise DataProcessStorageError("storage path escapes the configured root") from exc current = self._root for component in relative.parts: current = current / component try: current.mkdir(mode=0o700) except FileExistsError: pass info = current.lstat() if stat.S_ISLNK(info.st_mode) or not stat.S_ISDIR(info.st_mode): raise DataProcessStorageError("storage path contains a symlink or non-directory") return directory def _assert_controlled_parent(self, path: Path) -> None: try: relative_parent = path.parent.relative_to(self._root) except ValueError as exc: raise DataProcessStorageError("storage path escapes the configured root") from exc current = self._root for component in relative_parent.parts: current = current / component if not current.exists(): continue info = current.lstat() if stat.S_ISLNK(info.st_mode) or not stat.S_ISDIR(info.st_mode): raise DataProcessStorageError("storage path contains a symlink or non-directory") @staticmethod def _fsync_directory(directory: Path) -> None: descriptor = os.open(directory, os.O_RDONLY) try: os.fsync(descriptor) finally: os.close(descriptor) def _remove_empty_directory(self, directory: Path) -> None: if directory in {self._root, self._root / ".staging"}: return self._assert_controlled_parent(directory / "placeholder") try: directory.rmdir() except (FileNotFoundError, OSError): return @lru_cache def get_data_process_storage() -> LocalDataProcessStorage: return LocalDataProcessStorage() __all__ = [ "DataProcessStorageError", "LocalDataProcessStorage", "StagedSourceObject", "get_data_process_storage", ]