Release govoplan-files v0.1.26: speed archive workflows and unify file tools

src/govoplan_files/backend/routes/uploads.py

@@ -3,9 +3,11 @@ from __future__ import annotations
 import hashlib
 import hmac
 import json
+from contextlib import ExitStack, contextmanager
 from datetime import UTC, datetime, timedelta
+from pathlib import Path
 from typing import Literal
-from fastapi import APIRouter, Depends, File as FastAPIFile, Form, UploadFile
+from fastapi import APIRouter, Depends, File as FastAPIFile, Form, HTTPException, UploadFile
 from sqlalchemy.orm import Session
 
 from govoplan_core.auth import ApiPrincipal, require_scope
@@ -24,6 +26,10 @@ from govoplan_files.backend.schemas import (
 from govoplan_files.backend.db.models import FileAsset
 from govoplan_core.db.session import get_session
 from govoplan_files.backend.runtime import settings
+from govoplan_files.backend.archive_work import (
+    ArchiveProgress, ArchiveWorkExpired, discard_staged_upload, read_progress,
+    stage_upload, use_staged_upload,
+)
 from govoplan_files.backend.storage.paths import (
     UnsafeFilePathError,
     normalize_folder,
@@ -44,7 +50,7 @@ from govoplan_files.backend.storage.files import (
 
 
 from govoplan_files.backend.route_support import (
-    _asset_response,
+    _asset_list_response,
     _audit_connector_imports,
     _cleanup_temp_file,
     _connector_policy_error,
@@ -60,6 +66,57 @@ router = APIRouter(prefix="/files", tags=["files"])
 _ARCHIVE_PREVIEW_PURPOSE = "files.archive-preview.v1"
 
 
+@contextmanager
+def _archive_source(file, staged_upload_id, preview_token, principal):
+    if isinstance(staged_upload_id, str):
+        if file is not None and hasattr(file, "file"):
+            raise FileStorageError("Choose either the staged archive or a new upload")
+        if not isinstance(preview_token, str) or not preview_token:
+            raise FileStorageError("A valid preview token is required for a staged archive")
+        payload = open_transient_payload(preview_token, ttl_seconds=settings.file_archive_preview_ttl_seconds)
+        expected = {"purpose": _ARCHIVE_PREVIEW_PURPOSE, "tenant_id": principal.tenant_id,
+                    "user_id": principal.user.id, "staged_upload_id": staged_upload_id}
+        if any(payload.get(key) != value for key, value in expected.items()) or not isinstance(payload.get("filename"), str):
+            raise FileStorageError("Archive preview does not match this staged upload")
+        with use_staged_upload(settings, staged_upload_id, tenant_id=principal.tenant_id, user_id=principal.user.id) as path:
+            token_digest = payload.get("archive_sha256")
+            if not isinstance(token_digest, str) or not hmac.compare_digest(_archive_sha256(path), token_digest):
+                raise FileStorageError("Archive contents changed after preview; preview it again")
+            yield path, payload["filename"]
+        return
+    if file is None or not hasattr(file, "file"):
+        raise FileStorageError("Select an archive to upload")
+    filename = file.filename or "archive"
+    path = _spool_limited_upload_to_temp(file, max_bytes=settings.file_upload_zip_max_bytes, suffix=_archive_suffix(filename))
+    try:
+        yield path, filename
+    finally:
+        try:
+            _cleanup_temp_file(path)
+        except OSError:
+            pass
+
+
+@router.get("/archive-progress/{operation_id}")
+async def archive_operation_progress(operation_id: str, principal: ApiPrincipal = Depends(require_scope("files:file:upload"))):
+    # Small local receipt read stays available even while sync upload workers
+    # are busy. The request owns the transaction; this is not a job retry API.
+    try:
+        return read_progress(settings, operation_id, tenant_id=principal.tenant_id, user_id=principal.user.id)
+    except ArchiveWorkExpired as exc:
+        raise HTTPException(status_code=404, detail=str(exc)) from exc
+    except FileStorageError as exc:
+        raise _http_error(exc) from exc
+
+
+@router.delete("/archive-staging/{stage_id}", status_code=204)
+def release_staged_archive(stage_id: str, principal: ApiPrincipal = Depends(require_scope("files:file:upload"))):
+    try:
+        discard_staged_upload(settings, stage_id, tenant_id=principal.tenant_id, user_id=principal.user.id)
+    except FileStorageError as exc:
+        raise _http_error(exc) from exc
+
+
 def _archive_suffix(filename: str) -> str:
     lowered = filename.casefold()
     for suffix in (".tar.bz2", ".tar.gz", ".tar.xz", ".tbz2", ".tgz", ".txz", ".tar", ".zip"):
@@ -135,24 +192,22 @@ def _validate_archive_preview_token(
 
 @router.post("/archive-preview", response_model=ArchivePreviewResponse)
 def preview_archive_upload(
-    file: UploadFile = FastAPIFile(...),
+    file: UploadFile | None = FastAPIFile(default=None),
     owner_type: Literal["user", "group"] = Form(default="user"),
     owner_id: str | None = Form(default=None),
     path: str = Form(default=""),
     campaign_id: str | None = Form(default=None),
     password: str | None = Form(default=None),
+    retain_upload: bool = Form(default=False),
+    staged_upload_id: str | None = Form(default=None),
+    preview_token: str | None = Form(default=None),
     principal: ApiPrincipal = Depends(require_scope("files:file:upload")),
 ):
     target_owner = owner_id or principal.user.id
-    archive_path: str | None = None
-    filename = file.filename or "archive"
+    sources = ExitStack()
     try:
+        archive_path, filename = sources.enter_context(_archive_source(file, staged_upload_id, preview_token, principal))
         archive_format = archive_format_for_filename(filename)
-        archive_path = _spool_limited_upload_to_temp(
-            file,
-            max_bytes=settings.file_upload_zip_max_bytes,
-            suffix=_archive_suffix(filename),
-        )
         inspection = inspect_archive(
             archive_path,
             filename=filename,
@@ -163,6 +218,9 @@ def preview_archive_upload(
         )
         digest = _archive_sha256(archive_path)
         normalized_path = normalize_folder(path)
+        stage_id = staged_upload_id if isinstance(staged_upload_id, str) else None
+        if retain_upload is True and stage_id is None:
+            stage_id = stage_upload(settings, archive_path, tenant_id=principal.tenant_id, user_id=principal.user.id)
         preview_token = seal_transient_payload(
             {
                 "purpose": _ARCHIVE_PREVIEW_PURPOSE,
@@ -174,13 +232,21 @@ def preview_archive_upload(
                 "campaign_id": campaign_id or "",
                 "archive_format": archive_format,
                 "archive_sha256": digest,
+                **({"staged_upload_id": stage_id, "filename": filename} if stage_id else {}),
             }
         )
         expires_at = datetime.now(UTC) + timedelta(
             seconds=settings.file_archive_preview_ttl_seconds
         )
+        if isinstance(staged_upload_id, str):
+            # Repreviewing (for example with a password) does not renew the
+            # underlying private stage's lifetime.
+            expires_at = min(expires_at, datetime.fromtimestamp(
+                Path(archive_path).stat().st_mtime, UTC
+            ) + timedelta(seconds=settings.file_archive_preview_ttl_seconds))
         return ArchivePreviewResponse(
             preview_token=preview_token,
+            staged_upload_id=stage_id,
             archive_format=inspection.archive_format,
             entries=[
                 ArchiveEntryResponse(
@@ -200,16 +266,17 @@ def preview_archive_upload(
             password_verified=inspection.password_verified,
             expires_at=expires_at.isoformat(),
         )
-    except (FileStorageError, UnsafeFilePathError, ValueError) as exc:
+    except ArchiveWorkExpired as exc:
+        raise HTTPException(status_code=410, detail=str(exc)) from exc
+    except (FileStorageError, TransientPayloadError, UnsafeFilePathError, ValueError) as exc:
         raise _http_error(exc) from exc
     finally:
-        if archive_path:
-            _cleanup_temp_file(archive_path)
+        sources.close()
 
 
 @router.post("/archive-confirm", response_model=FileUploadResponse)
 def confirm_archive_upload(
-    file: UploadFile = FastAPIFile(...),
+    file: UploadFile | None = FastAPIFile(default=None),
     preview_token: str = Form(...),
     selected_paths_json: str = Form(...),
     owner_type: Literal["user", "group"] = Form(default="user"),
@@ -225,13 +292,17 @@ def confirm_archive_upload(
     source_revision: str | None = Form(default=None),
     connector_policy_json: str | None = Form(default=None),
     encryption_vault_id: str | None = Form(default=None),
+    staged_upload_id: str | None = Form(default=None),
+    operation_id: str | None = Form(default=None),
     session: Session = Depends(get_session),
     principal: ApiPrincipal = Depends(require_scope("files:file:upload")),
 ):
     target_owner = owner_id or principal.user.id
-    archive_path: str | None = None
-    filename = file.filename or "archive"
+    sources = ExitStack()
+    progress = None
+    committed = False
     try:
+        archive_path, filename = sources.enter_context(_archive_source(file, staged_upload_id, preview_token, principal))
         raw_resolutions = (
             json.loads(conflict_resolutions_json)
             if conflict_resolutions_json
@@ -251,11 +322,6 @@ def confirm_archive_upload(
         )
         archive_format = archive_format_for_filename(filename)
         normalized_path = normalize_folder(path)
-        archive_path = _spool_limited_upload_to_temp(
-            file,
-            max_bytes=settings.file_upload_zip_max_bytes,
-            suffix=_archive_suffix(filename),
-        )
         digest = _archive_sha256(archive_path)
         _validate_archive_preview_token(
             preview_token,
@@ -268,6 +334,7 @@ def confirm_archive_upload(
             archive_format=archive_format,
             archive_sha256=digest,
         )
+        progress = ArchiveProgress(settings, operation_id, tenant_id=principal.tenant_id, user_id=principal.user.id)
         extracted = extract_archive_upload(
             session,
             tenant_id=principal.tenant_id,
@@ -289,16 +356,20 @@ def confirm_archive_upload(
             max_file_bytes=settings.file_upload_max_bytes,
             max_expanded_bytes=settings.file_archive_max_expanded_bytes,
             max_expansion_ratio=settings.file_archive_max_expansion_ratio,
+            progress=progress,
         )
         uploaded_assets = [item.asset for item in extracted]
         _audit_connector_imports(session, principal, uploaded_assets)
+        progress("finalizing", len(extracted), len(extracted), progress.value["completed_bytes"], progress.value["total_bytes"])
+        session.flush()
+        response = FileUploadResponse(files=_asset_list_response(session, uploaded_assets, include_shares=True))
         session.commit()
-        return FileUploadResponse(
-            files=[
-                _asset_response(session, asset, include_shares=True)
-                for asset in uploaded_assets
-            ]
-        )
+        committed = True
+        progress.finish(True)
+        return response
+    except ArchiveWorkExpired as exc:
+        session.rollback()
+        raise HTTPException(status_code=410, detail=str(exc)) from exc
     except ConnectorPolicyDenied as exc:
         session.rollback()
         raise _connector_policy_error(exc) from exc
@@ -311,9 +382,21 @@ def confirm_archive_upload(
     ) as exc:
         session.rollback()
         raise _http_error(exc) from exc
+    except Exception:
+        session.rollback()
+        raise
     finally:
-        if archive_path:
-            _cleanup_temp_file(archive_path)
+        if progress and not committed:
+            progress.finish(False)
+        try:
+            sources.close()
+        except OSError:
+            pass
+        if committed and isinstance(staged_upload_id, str):
+            try:
+                discard_staged_upload(settings, staged_upload_id, tenant_id=principal.tenant_id, user_id=principal.user.id)
+            except (OSError, FileStorageError):
+                pass
 
 
 @router.post("/upload", response_model=FileUploadResponse)
@@ -402,6 +485,8 @@ def upload_files(
             )
             uploaded_assets.append(stored.asset)
         _audit_connector_imports(session, principal, uploaded_assets)
+        session.flush()
+        response = FileUploadResponse(files=_asset_list_response(session, uploaded_assets, include_shares=True))
         session.commit()
     except ConnectorPolicyDenied as exc:
         session.rollback()
@@ -409,12 +494,7 @@ def upload_files(
     except (FileStorageError, UnsafeFilePathError, ValueError) as exc:
         session.rollback()
         raise _http_error(exc) from exc
-    return FileUploadResponse(
-        files=[
-            _asset_response(session, asset, include_shares=True)
-            for asset in uploaded_assets
-        ]
-    )
+    return response
 
 
 @router.post("/upload-zip", response_model=FileUploadResponse)
@@ -469,6 +549,8 @@ def upload_zip(
             max_total_bytes=settings.file_upload_zip_max_bytes,
         )
         _audit_connector_imports(session, principal, [item.asset for item in extracted])
+        session.flush()
+        response = FileUploadResponse(files=_asset_list_response(session, [item.asset for item in extracted], include_shares=True))
         session.commit()
     except ConnectorPolicyDenied as exc:
         session.rollback()
@@ -479,9 +561,4 @@ def upload_zip(
     finally:
         if zip_path:
             _cleanup_temp_file(zip_path)
-    return FileUploadResponse(
-        files=[
-            _asset_response(session, item.asset, include_shares=True)
-            for item in extracted
-        ]
-    )
+    return response