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