YG_FT/backend/app/modules/data_process/storage.py

"""数据处理原始源文件的受控本地对象存储。"""

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",
]