"""Authenticated durable import lifecycle; filesystem work remains behind private adapters.""" import asyncio import hashlib import os from contextlib import asynccontextmanager from datetime import timedelta from uuid import UUID from sqlalchemy import func, select, text from kf_oj.contracts.admin_problems import ( IMPORT_LIFETIME_SECONDS, MAX_OWNER_CANDIDATES, MAX_TOTAL_CANDIDATES, MAX_UPLOAD_CHUNK_BYTES, MAX_UPLOAD_CHUNKS, STAGING_MAX_BYTES, STAGING_RESERVATION_BYTES, UPLOAD_DEADLINE_SECONDS, CandidatePreview, CreateProblemImport, ProblemImportCapabilities, ProblemImportDetail, ProblemImportPage, ProblemImportSummary, ) from kf_oj.infra.artifacts.staging import ImportStaging from kf_oj.infra.artifacts.store import ArtifactError from kf_oj.modules.identity.public import MaintenanceDenied, ProblemMaintenancePort from kf_oj.modules.operations.public import MaintenanceAuditPort from kf_oj.modules.problems.import_models import ImportPool, ProblemCandidate from kf_oj.modules.problems.import_models import ProblemImportOperation as Operation from kf_oj.modules.problems.import_policy import ImportUnavailable as ImportUnavailable from kf_oj.modules.problems.import_policy import policy_digest as policy_digest ACTIVE = ("uploading", "uploaded", "validating", "ready") class ImportConflict(ValueError): pass class ImportBadDigest(ImportConflict): pass class ImportMissing(ValueError): pass def summary(operation) -> ProblemImportSummary: return ProblemImportSummary( operation_id=operation.id, state=operation.state, expected_archive_digest=operation.archive_digest, archive_bytes=operation.archive_bytes, created_at=operation.created_at, expires_at=operation.expires_at, updated_at=operation.updated_at, failure_code=operation.failure_code, candidate_id=operation.candidate_id, published_version_id=operation.published_version_id, ) class ProblemImports: def __init__( self, sessions, staging: ImportStaging, identity: ProblemMaintenancePort, audit: MaintenanceAuditPort, policy, ): self.sessions, self.staging, self.identity, self.audit, self.policy = ( sessions, staging, identity, audit, policy, ) @asynccontextmanager async def write(self, principal): async with self.sessions.begin() as session: await session.execute(text("SET LOCAL statement_timeout='5s'")) await session.execute(text("SET LOCAL lock_timeout='2s'")) await self.identity.require_problem_maintenance(session, principal) yield session @staticmethod async def pool(session): pool = await session.get(ImportPool, 1, with_for_update=True) if pool is None: raise ImportUnavailable("import capacity unavailable") actual = await session.scalar( select(func.coalesce(func.sum(Operation.reserved_bytes), 0)).where( Operation.reserved_bytes > 0 ) ) if pool.reserved_bytes != actual: raise ImportUnavailable("import capacity inconsistent") return pool async def capabilities(self): async with self.sessions() as session: pool = await session.get(ImportPool, 1) now = await session.scalar(select(func.clock_timestamp())) available = bool( pool and pool.worker_ready_until and pool.worker_ready_until > now and pool.worker_policy_digest == policy_digest(self.policy) ) return ProblemImportCapabilities( uploads_available=available, validation_available=available ) async def create(self, principal, request: CreateProblemImport): # Immutable receipt retries need no filesystem slot, including while a parser/GC # owns the root lock. Missing receipts still enter the locked creation path below. async with self.write(principal) as session: existing = await session.get(Operation, request.operation_id) if existing: if ( existing.owner_id != principal.account_id or existing.archive_digest != request.archive_digest or existing.archive_bytes != request.archive_bytes ): raise ImportConflict("operation key already bound") return summary(existing) try: with self.staging.exclusive(): await asyncio.to_thread(self.staging.usage) async with self.write(principal) as session: pool = await self.pool(session) existing = await session.get( Operation, request.operation_id, with_for_update=True ) if existing: if ( existing.owner_id != principal.account_id or existing.archive_digest != request.archive_digest or existing.archive_bytes != request.archive_bytes ): raise ImportConflict("operation key already bound") return summary(existing) now = await session.scalar(select(func.clock_timestamp())) if ( not pool.worker_ready_until or pool.worker_ready_until <= now or pool.worker_policy_digest != policy_digest(self.policy) ): raise ImportUnavailable("import validation unavailable") global_count = await session.scalar( select(func.count()) .select_from(Operation) .where(Operation.state.in_(ACTIVE)) ) owner_count = await session.scalar( select(func.count()) .select_from(Operation) .where( Operation.state.in_(ACTIVE), Operation.owner_id == principal.account_id ) ) busy = await session.scalar( select(func.count()) .select_from(Operation) .where( Operation.state.in_(("uploading", "validating")), Operation.lease_until > now, ) ) if ( global_count >= MAX_TOTAL_CANDIDATES or owner_count >= MAX_OWNER_CANDIDATES or busy or pool.reserved_bytes + STAGING_RESERVATION_BYTES > STAGING_MAX_BYTES ): raise ImportUnavailable("import capacity exhausted") operation = Operation( id=request.operation_id, owner_id=principal.account_id, authorization_version=principal.authorization_version, archive_digest=request.archive_digest, archive_bytes=request.archive_bytes, state="uploading", created_at=now, updated_at=now, expires_at=now + timedelta(seconds=IMPORT_LIFETIME_SECONDS), lease_until=now + timedelta(seconds=UPLOAD_DEADLINE_SECONDS + 15), attempts=0, fence=0, reserved_bytes=STAGING_RESERVATION_BYTES, ) session.add(operation) pool.reserved_bytes += STAGING_RESERVATION_BYTES await self.audit.record( session, principal.account_id, principal.authorization_version, operation.id, "upload_created", ) await session.flush() return summary(operation) except ArtifactError: raise ImportUnavailable("import storage unavailable") from None async def detail(self, principal, operation_id): # Recheck live authority even when a principal was injected by an internal caller. async with self.write(principal) as session: op = await session.get(Operation, operation_id) if op is None or op.owner_id != principal.account_id: raise ImportMissing("import unavailable") candidate = ( await session.get(ProblemCandidate, op.candidate_id) if op.candidate_id else None ) return ProblemImportDetail( **summary(op).model_dump(), candidate=CandidatePreview.model_validate(candidate.preview) if candidate else None, ) async def list(self, principal, after=None, limit=20): if not 1 <= limit <= 100: raise ValueError("invalid page limit") async with self.write(principal) as session: query = select(Operation).where(Operation.owner_id == principal.account_id) if after: query = query.where(Operation.id > after) rows = list(await session.scalars(query.order_by(Operation.id).limit(limit + 1))) return ProblemImportPage( items=tuple(summary(o) for o in rows[:limit]), next_after=rows[limit - 1].id if len(rows) > limit else None, ) async def upload(self, principal, operation_id: UUID, chunks, content_length: int | None): # Reject missing/foreign receipts before consulting shared filesystem capacity. # The locked transaction below still rechecks authority and ownership before writes. async with self.write(principal) as session: op = await session.get(Operation, operation_id) if op is None or op.owner_id != principal.account_id: raise ImportMissing("import unavailable") if content_length is not None and content_length != op.archive_bytes: raise ImportConflict("archive length mismatch") descriptor = operation_fd = None fence = None received = 0 chunk_count = 0 digest = hashlib.sha256() try: with self.staging.exclusive(): async with self.write(principal) as session: await self.pool(session) op = await session.get(Operation, operation_id, with_for_update=True) if not op or op.owner_id != principal.account_id: raise ImportMissing("import unavailable") accepted = op.state in ("uploaded", "validating", "ready", "published") now = await session.scalar(select(func.clock_timestamp())) if ( not accepted and ( op.state != "uploading" or op.lease_until is None or op.lease_until <= now ) ) or op.expires_at <= now: raise ImportConflict("upload no longer available") if content_length is not None and content_length != op.archive_bytes: raise ImportConflict("archive length mismatch") if not accepted: op.fence += 1 fence = op.fence op.lease_until = now + timedelta(seconds=UPLOAD_DEADLINE_SECONDS + 15) expected_bytes, expected_digest = op.archive_bytes, op.archive_digest if not accepted: operation_fd = self.staging.operation(operation_id, create=True) descriptor = os.open( "archive", os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW | os.O_CLOEXEC, 0o600, dir_fd=operation_fd, ) async with asyncio.timeout(UPLOAD_DEADLINE_SECONDS): async for chunk in chunks: chunk_count += 1 received += len(chunk) if ( type(chunk) is not bytes or len(chunk) > MAX_UPLOAD_CHUNK_BYTES or chunk_count > MAX_UPLOAD_CHUNKS or received > expected_bytes ): raise ImportConflict("archive bytes outside policy") digest.update(chunk) if accepted: continue # Cancellation waits for the finite writer outside the event loop. task = asyncio.create_task(asyncio.to_thread(write_all, descriptor, chunk)) try: await asyncio.shield(task) except BaseException: await task raise if received != expected_bytes or digest.hexdigest() != expected_digest: raise ImportBadDigest("archive digest or length mismatch") if accepted: async with self.write(principal) as session: op = await session.get(Operation, operation_id) if not op or op.state not in ( "uploaded", "validating", "ready", "published", ): raise ImportConflict("import receipt no longer available") return summary(op) await asyncio.to_thread(os.fsync, descriptor) os.close(descriptor) descriptor = None os.fsync(operation_fd) async with self.write(principal) as session: await self.pool(session) op = await session.get(Operation, operation_id, with_for_update=True) now = await session.scalar(select(func.clock_timestamp())) if op.state != "uploading" or op.fence != fence or op.lease_until <= now: raise ImportConflict("upload receipt expired") op.state = "uploaded" op.lease_until = None op.updated_at = now await self.audit.record( session, op.owner_id, op.authorization_version, op.id, "upload_received" ) result = summary(op) return result except BaseException as error: # Failure bookkeeping cannot authorize a new business write. It only fences this # accepted upload and releases capacity after physical bytes have been removed. if descriptor is not None: os.close(descriptor) descriptor = None if operation_fd is not None: os.close(operation_fd) operation_fd = None if fence is not None: try: with self.staging.exclusive(): await asyncio.to_thread(self.staging.remove, operation_id) async with self.sessions.begin() as session: pool = await self.pool(session) op = await session.get(Operation, operation_id, with_for_update=True) if op and op.state == "uploading" and op.fence == fence: op.state = "failed" op.failure_code = ( "upload_deadline" if isinstance(error, TimeoutError) else "digest_mismatch" if isinstance(error, ImportBadDigest) else "authority_revoked" if isinstance(error, MaintenanceDenied) else "storage_unavailable" if isinstance(error, OSError) else "upload_interrupted" ) op.updated_at = await session.scalar(select(func.clock_timestamp())) op.lease_until = None pool.reserved_bytes -= op.reserved_bytes op.reserved_bytes = 0 await self.audit.record( session, op.owner_id, op.authorization_version, op.id, "import_failed", ) except (ArtifactError, OSError): pass # Durable reservation retained for fenced GC. raise finally: if descriptor is not None: os.close(descriptor) if operation_fd is not None: os.close(operation_fd) async def discard(self, principal, operation_id): async with self.write(principal) as session: op = await session.get(Operation, operation_id) if op is None or op.owner_id != principal.account_id: raise ImportMissing("import unavailable") if op.state == "discarded": return summary(op) try: with self.staging.exclusive(): async with self.write(principal) as session: pool = await self.pool(session) op = await session.get(Operation, operation_id, with_for_update=True) if not op or op.owner_id != principal.account_id: raise ImportMissing("import unavailable") if op.state == "published": raise ImportConflict("published import cannot be discarded") if op.state == "discarded": return summary(op) await asyncio.to_thread(self.staging.remove, operation_id) op.state = "discarded" op.candidate_id = None op.failure_code = None op.lease_until = None op.fence += 1 op.updated_at = await session.scalar(select(func.clock_timestamp())) pool.reserved_bytes -= op.reserved_bytes op.reserved_bytes = 0 await self.audit.record( session, principal.account_id, principal.authorization_version, op.id, "import_discarded", ) return summary(op) except ArtifactError: raise ImportUnavailable("import storage unavailable") from None def write_all(descriptor, data): remaining = memoryview(data) while remaining: size = os.write(descriptor, remaining) if size <= 0: raise OSError("short staging write") remaining = remaining[size:]