zcbot/core/artifact_lifecycle.py

"""Database-backed lifecycle operations for published workspace artifacts."""
from __future__ import annotations

import hashlib
import mimetypes
import os
import shutil
from datetime import datetime, timezone
from pathlib import Path
from uuid import UUID, uuid4

from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert as pg_insert

from .artifacts import ARTIFACT_TRASH_DIR, ArtifactRef
from .storage import session_scope
from .storage.models import Artifact


def _rel(root: Path, path: Path) -> str:
    return Path(path).resolve().relative_to(Path(root).resolve()).as_posix()


def _hash_file(path: Path) -> str:
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def register_published_artifacts(
    *,
    user_id: UUID,
    task_id: UUID,
    user_root: Path,
    working_dir: Path,
    refs: tuple[dict, ...],
) -> tuple[dict, ...]:
    """Upsert active artifact identities and return version-2 message refs."""
    root = Path(user_root).resolve()
    wd = Path(working_dir).resolve()
    output: list[dict] = []
    with session_scope() as session:
        for ref in refs:
            task_path = str(ref.get("path") or "")
            path = (wd / Path(task_path)).resolve()
            path.relative_to(wd)
            if not path.is_file():
                continue
            current_path = _rel(root, path)
            label = str(ref.get("label") or "")
            media_type = mimetypes.guess_type(path.name)[0]
            size_bytes = path.stat().st_size
            content_sha256 = _hash_file(path)
            statement = pg_insert(Artifact).values(
                user_id=user_id,
                origin_task_id=task_id,
                current_path=current_path,
                label=label,
                media_type=media_type,
                size_bytes=size_bytes,
                content_sha256=content_sha256,
            ).on_conflict_do_update(
                index_elements=[Artifact.user_id, Artifact.current_path],
                index_where=Artifact.status == "active",
                set_={
                    "label": label,
                    "media_type": media_type,
                    "size_bytes": size_bytes,
                    "content_sha256": content_sha256,
                    "updated_at": datetime.now(timezone.utc),
                },
            ).returning(Artifact.artifact_id)
            artifact_id = session.execute(statement).scalar_one()
            output.append(ArtifactRef(
                path=task_path,
                label=label,
                artifact_id=artifact_id,
                version=2,
            ).as_dict())
    return tuple(output)


def rename_active_artifacts(
    *,
    user_id: UUID,
    user_root: Path,
    old_path: Path,
    new_path: Path,
) -> int:
    """Rewrite active artifact paths for one file or directory subtree."""
    root = Path(user_root).resolve()
    old_rel = _rel(root, old_path)
    new_rel = _rel(root, new_path)
    with session_scope() as session:
        rows = session.execute(
            select(Artifact).where(
                Artifact.user_id == user_id,
                Artifact.status == "active",
            )
        ).scalars().all()
        changed = 0
        for row in rows:
            if row.current_path == old_rel:
                suffix = ""
            elif row.current_path.startswith(old_rel + "/"):
                suffix = row.current_path[len(old_rel):]
            else:
                continue
            row.current_path = new_rel + suffix
            changed += 1
        return changed


def copy_active_artifacts(
    *,
    user_id: UUID,
    user_root: Path,
    source: Path,
    target: Path,
) -> int:
    """Create independent artifact identities for copied files."""
    root = Path(user_root).resolve()
    source_rel = _rel(root, source)
    target_rel = _rel(root, target)
    with session_scope() as session:
        sources = session.execute(
            select(Artifact).where(
                Artifact.user_id == user_id,
                Artifact.status == "active",
            )
        ).scalars().all()
        created = 0
        for original in sources:
            if original.current_path == source_rel:
                suffix = ""
            elif original.current_path.startswith(source_rel + "/"):
                suffix = original.current_path[len(source_rel):]
            else:
                continue
            copied_path = target_rel + suffix
            session.add(Artifact(
                user_id=user_id,
                origin_task_id=original.origin_task_id,
                copied_from_artifact_id=original.artifact_id,
                current_path=copied_path,
                label=(original.label + " 副本").strip(),
                media_type=original.media_type,
                size_bytes=original.size_bytes,
                content_sha256=original.content_sha256,
            ))
            created += 1
        return created


def trash_active_artifacts(
    *,
    user_id: UUID,
    user_root: Path,
    target: Path,
) -> int:
    """Move matching active artifacts to hidden trash and mark them deleted."""
    root = Path(user_root).resolve()
    source = Path(target).resolve()
    source_rel = _rel(root, source)
    entry = (
        root / ARTIFACT_TRASH_DIR
        / datetime.now(timezone.utc).strftime("%Y/%m/%d")
        / uuid4().hex
    )
    moved: list[tuple[Path, Path]] = []
    try:
        with session_scope() as session:
            rows = session.execute(
                select(Artifact).where(
                    Artifact.user_id == user_id,
                    Artifact.status == "active",
                ).with_for_update()
            ).scalars().all()
            matches = [
                row for row in rows
                if row.current_path == source_rel
                or row.current_path.startswith(source_rel + "/")
            ]
            if not matches:
                return 0
            for row in matches:
                original = root / Path(row.current_path)
                if not original.is_file():
                    continue
                destination = entry / "files" / Path(row.current_path)
                destination.parent.mkdir(parents=True, exist_ok=True)
                os.replace(original, destination)
                moved.append((original, destination))
                row.status = "deleted"
                row.deleted_at = datetime.now(timezone.utc)
                row.trash_path = _rel(root, destination)
        return len(moved)
    except Exception:
        for original, destination in reversed(moved):
            try:
                original.parent.mkdir(parents=True, exist_ok=True)
                os.replace(destination, original)
            except OSError:
                pass
        shutil.rmtree(entry, ignore_errors=True)
        raise