"""Lease-fenced persistence for the PARSE_DOCUMENT workflow.""" from __future__ import annotations import hashlib import json import re import uuid from collections.abc import Callable, Mapping from dataclasses import dataclass from typing import Any, Protocol, cast import psycopg from psycopg import Connection from psycopg.rows import dict_row from app.core.config import Settings from app.core.secrets import SecretFileError from app.persistence.job_queue import BackgroundJob, JobLease, LeaseLostError from app.services.document_ingestion import ( CloudTextPolicy, IngestionArtifact, IngestionStatus, ) type WorkflowRow = dict[str, Any] type ConnectionFactory = Callable[[str, int], Connection[WorkflowRow]] _DB_ID_NAMESPACE = uuid.UUID("07fb2c87-b4c9-5dd8-b965-fb5d8e53e350") _HASH = re.compile(r"^[0-9a-f]{64}$") _ERROR_CODE = re.compile(r"^[A-Z][A-Z0-9_]{0,127}$") _NORMALIZATION_PROFILE_HASH = hashlib.sha256(b"document-ingestion-normalized-text-v1").hexdigest() @dataclass(frozen=True, slots=True) class DocumentSource: upload_id: uuid.UUID document_id: uuid.UUID knowledge_base_id: uuid.UUID access_scope_id: uuid.UUID filename: str mime_type: str storage_key: uuid.UUID byte_size: int raw_sha256: str @dataclass(frozen=True, slots=True) class PlannedPage: id: uuid.UUID ordinal: int page_number: int | None text: str text_sha256: str line_start: int line_end: int @dataclass(frozen=True, slots=True) class PlannedBlock: id: uuid.UUID page_id: uuid.UUID ordinal: int kind: str text: str text_sha256: str section_path: tuple[str, ...] anchor_id: str normalized_text_sha256: str char_start: int char_end: int line_start: int line_end: int page_start: int | None page_end: int | None @dataclass(frozen=True, slots=True) class PlannedChunk: id: uuid.UUID ordinal: int display_text: str cloud_text: str cloud_text_sha256: str embedding_prefix: str embedding_text: str embedding_text_sha256: str token_count: int page_start: int | None page_end: int | None section_path: tuple[str, ...] metadata: Mapping[str, object] @dataclass(frozen=True, slots=True) class ArtifactPlan: version_id: uuid.UUID parser_profile_hash: str normalization_profile_hash: str chunk_profile_hash: str cloud_policy_id: str outbound_manifest_sha256: str | None review_state: str version_status: str version_error_code: str | None expected_chunk_count: int document_status: str job_stage: str pages: tuple[PlannedPage, ...] blocks: tuple[PlannedBlock, ...] chunks: tuple[PlannedChunk, ...] class InvalidDocumentJobError(RuntimeError): def __init__(self) -> None: super().__init__("document job contract is invalid") class DocumentWorkflowPersistenceError(RuntimeError): def __init__(self) -> None: super().__init__("document workflow persistence unavailable") class ArtifactConflictError(RuntimeError): def __init__(self) -> None: super().__init__("stored ingestion artifact conflicts with deterministic output") class DocumentWorkflowRepository(Protocol): def load_source(self, job: BackgroundJob) -> DocumentSource: ... def record_terminal_parse_failure( self, *, lease: JobLease, source: DocumentSource, error_code: str, ) -> None: ... def persist_artifact( self, *, lease: JobLease, source: DocumentSource, artifact: IngestionArtifact, cloud_policy: CloudTextPolicy, embedding_model: str, embedding_dimension: int, ) -> ArtifactPlan: ... LOAD_SOURCE_SQL = """ SELECT upload.id AS upload_id, upload.document_id, upload.knowledge_base_id, upload.access_scope_id, upload.original_filename, upload.declared_mime_type, upload.storage_key, upload.actual_size, upload.actual_sha256, document.filename AS document_filename, document.mime_type AS document_mime_type, document.raw_sha256 AS document_raw_sha256, document.storage_key AS document_storage_key FROM rag.background_jobs AS job JOIN rag.document_uploads AS upload ON upload.parse_job_id = job.id JOIN rag.documents AS document ON document.id = upload.document_id AND document.knowledge_base_id = upload.knowledge_base_id AND document.access_scope_id = upload.access_scope_id WHERE job.id = %s AND job.status = 'RUNNING' AND job.lease_owner = %s AND job.lease_token = %s AND job.lease_until >= now() AND job.job_type = 'PARSE_DOCUMENT' AND job.required_capability = 'document_parse' AND job.resource_type = 'document' AND job.resource_id = %s AND upload.id = %s AND upload.document_id = %s AND upload.status = 'COMPLETED' AND upload.actual_size = upload.expected_size AND upload.actual_sha256 = upload.expected_sha256 AND upload.actual_sha256 = document.raw_sha256 AND upload.storage_key::text = document.storage_key AND upload.original_filename = document.filename AND upload.declared_mime_type = document.mime_type AND document.deleted_at IS NULL LIMIT 1 """ FAIL_PARSE_SQL = """ WITH fenced AS ( SELECT job.id AS job_id, document.id AS document_id FROM rag.background_jobs AS job JOIN rag.document_uploads AS upload ON upload.parse_job_id = job.id JOIN rag.documents AS document ON document.id = upload.document_id AND document.knowledge_base_id = upload.knowledge_base_id AND document.access_scope_id = upload.access_scope_id WHERE job.id = %s AND job.status = 'RUNNING' AND job.lease_owner = %s AND job.lease_token = %s AND job.lease_until >= now() AND job.job_type = 'PARSE_DOCUMENT' AND job.required_capability = 'document_parse' AND job.resource_type = 'document' AND job.resource_id = %s AND upload.id = %s AND upload.document_id = %s AND upload.knowledge_base_id = %s AND upload.access_scope_id = %s AND upload.actual_size = %s AND upload.actual_sha256 = %s AND upload.storage_key = %s AND upload.status = 'COMPLETED' AND document.raw_sha256 = %s AND document.storage_key = %s AND document.deleted_at IS NULL FOR UPDATE OF job, document ), updated_document AS ( UPDATE rag.documents AS document SET status = 'FAILED', updated_at = now() FROM fenced WHERE document.id = fenced.document_id RETURNING document.id ) UPDATE rag.background_jobs AS job SET stage = 'PARSE_REJECTED', progress = 100, last_error_code = %s, last_error_message = 'Document parsing was rejected by a deterministic local policy.', updated_at = now() FROM fenced JOIN updated_document ON updated_document.id = fenced.document_id WHERE job.id = fenced.job_id RETURNING job.id """ INSERT_VERSION_SQL = """ INSERT INTO rag.document_versions ( id, document_id, parser_profile_hash, ocr_profile_hash, normalization_profile_hash, chunk_profile_hash, cloud_policy_id, outbound_manifest_sha256, review_state, embedding_profile_hash, status, expected_chunk_count, error_code ) VALUES (%s, %s, %s, NULL, %s, %s, %s, %s, %s, NULL, %s, %s, %s) ON CONFLICT DO NOTHING """ INSERT_PAGE_SQL = """ INSERT INTO rag.document_pages ( id, document_version_id, ordinal, page_number, display_text, text_sha256, line_start, line_end ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) ON CONFLICT DO NOTHING """ INSERT_BLOCK_SQL = """ INSERT INTO rag.document_blocks ( id, document_version_id, page_id, ordinal, block_kind, display_text, text_sha256, section_path, anchor_id, normalized_text_sha256, char_start, char_end, line_start, line_end, page_start, page_end ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb, %s, %s, %s, %s, %s, %s, %s, %s) ON CONFLICT DO NOTHING """ INSERT_MANIFEST_ITEM_SQL = """ INSERT INTO rag.outbound_manifest_items ( document_version_id, ordinal, outbound_manifest_sha256, cloud_text_sha256, embedding_text_sha256 ) VALUES (%s, %s, %s, %s, %s) ON CONFLICT DO NOTHING """ INSERT_CHUNK_SQL = """ INSERT INTO rag.chunks ( id, knowledge_base_id, document_id, document_version_id, access_scope_id, ordinal, display_text, cloud_text, cloud_text_sha256, embedding_prefix, embedding_text, embedding_text_sha256, embedded_text_sha256, embedding_profile_hash, outbound_manifest_sha256, token_count, page_start, page_end, section_path, metadata, embedding_model, embedding_dimension, embedding, approval_status, index_status, searchable ) VALUES ( %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, NULL, NULL, %s, %s, %s, %s, %s::jsonb, %s::jsonb, %s, %s, NULL, 'LOCAL_PARSED_PENDING_CLOUD_REVIEW', 'PENDING', false ) ON CONFLICT DO NOTHING """ SELECT_VERSION_SQL = """ SELECT id, document_id, parser_profile_hash, normalization_profile_hash, chunk_profile_hash, cloud_policy_id, outbound_manifest_sha256, expected_chunk_count FROM rag.document_versions WHERE id = %s AND document_id = %s """ SELECT_PAGES_SQL = """ SELECT id, ordinal, page_number, display_text, text_sha256, line_start, line_end FROM rag.document_pages WHERE document_version_id = %s ORDER BY ordinal """ SELECT_BLOCKS_SQL = """ SELECT id, page_id, ordinal, block_kind, display_text, text_sha256, section_path, anchor_id, normalized_text_sha256, char_start, char_end, line_start, line_end, page_start, page_end FROM rag.document_blocks WHERE document_version_id = %s ORDER BY ordinal """ SELECT_MANIFEST_ITEMS_SQL = """ SELECT ordinal, outbound_manifest_sha256, cloud_text_sha256, embedding_text_sha256 FROM rag.outbound_manifest_items WHERE document_version_id = %s ORDER BY ordinal """ SELECT_CHUNKS_SQL = """ SELECT id, ordinal, display_text, cloud_text, cloud_text_sha256, embedding_prefix, embedding_text, embedding_text_sha256, token_count, page_start, page_end, section_path, metadata FROM rag.chunks WHERE document_version_id = %s AND deleted_at IS NULL ORDER BY ordinal """ UPDATE_DOCUMENT_SQL = """ UPDATE rag.documents SET status = CASE WHEN active_version_id IS NULL THEN %s ELSE status END, updated_at = now() WHERE id = %s AND knowledge_base_id = %s AND access_scope_id = %s AND raw_sha256 = %s AND storage_key = %s AND deleted_at IS NULL RETURNING id """ FINALIZE_JOB_SQL = """ UPDATE rag.background_jobs AS job SET stage = %s, progress = 100, last_error_code = NULL, last_error_message = NULL, updated_at = now() WHERE job.id = %s AND job.status = 'RUNNING' AND job.lease_owner = %s AND job.lease_token = %s AND job.lease_until >= now() AND job.job_type = 'PARSE_DOCUMENT' AND job.required_capability = 'document_parse' AND job.resource_type = 'document' AND job.resource_id = %s AND EXISTS ( SELECT 1 FROM rag.document_uploads AS upload JOIN rag.documents AS document ON document.id = upload.document_id AND document.knowledge_base_id = upload.knowledge_base_id AND document.access_scope_id = upload.access_scope_id WHERE upload.parse_job_id = job.id AND upload.id = %s AND upload.document_id = %s AND upload.knowledge_base_id = %s AND upload.access_scope_id = %s AND upload.actual_size = %s AND upload.actual_sha256 = %s AND upload.storage_key = %s AND upload.status = 'COMPLETED' AND document.raw_sha256 = %s AND document.storage_key = %s AND document.deleted_at IS NULL ) RETURNING job.id """ class PostgresDocumentWorkflowRepository: def __init__( self, settings: Settings, *, connect_timeout: int = 5, connection_factory: ConnectionFactory | None = None, ) -> None: self._settings = settings self._connect_timeout = connect_timeout self._connection_factory = connection_factory or _default_connection_factory def _dsn(self) -> str: return ( self._settings.database_url() .set(drivername="postgresql") .render_as_string(hide_password=False) ) def load_source(self, job: BackgroundJob) -> DocumentSource: upload_id, document_id = _validate_job_contract(job) parameters = ( job.id, job.lease.worker_id, job.lease.lease_token, document_id, upload_id, document_id, ) try: with self._connection_factory(self._dsn(), self._connect_timeout) as connection: with connection.transaction(): row = connection.execute(LOAD_SOURCE_SQL, parameters).fetchone() except (OSError, SecretFileError, psycopg.Error): raise DocumentWorkflowPersistenceError from None if row is None: raise LeaseLostError("document source fence is no longer valid") source = _source_from_row(row) if source.upload_id != upload_id or source.document_id != document_id: raise ArtifactConflictError return source def record_terminal_parse_failure( self, *, lease: JobLease, source: DocumentSource, error_code: str, ) -> None: if _ERROR_CODE.fullmatch(error_code) is None: raise ValueError("error_code must be a stable uppercase identifier") parameters = _fence_parameters(lease, source) + (error_code,) try: with self._connection_factory(self._dsn(), self._connect_timeout) as connection: with connection.transaction(): row = connection.execute(FAIL_PARSE_SQL, parameters).fetchone() if row is None: raise LeaseLostError("document failure fence is no longer valid") except (OSError, SecretFileError, psycopg.Error): raise DocumentWorkflowPersistenceError from None def persist_artifact( self, *, lease: JobLease, source: DocumentSource, artifact: IngestionArtifact, cloud_policy: CloudTextPolicy, embedding_model: str, embedding_dimension: int, ) -> ArtifactPlan: plan = plan_artifact( source, artifact, cloud_policy=cloud_policy, ) if not embedding_model.strip() or embedding_dimension != 1024: raise ArtifactConflictError try: with self._connection_factory(self._dsn(), self._connect_timeout) as connection: with connection.transaction(): self._insert_plan( connection, source, plan, embedding_model=embedding_model, embedding_dimension=embedding_dimension, ) self._verify_plan(connection, source, plan) document = connection.execute( UPDATE_DOCUMENT_SQL, ( plan.document_status, source.document_id, source.knowledge_base_id, source.access_scope_id, source.raw_sha256, str(source.storage_key), ), ).fetchone() if document is None: raise ArtifactConflictError finalized = connection.execute( FINALIZE_JOB_SQL, (plan.job_stage, *_fence_parameters(lease, source)), ).fetchone() if finalized is None: raise LeaseLostError("document artifact fence is no longer valid") except (OSError, SecretFileError, psycopg.Error): raise DocumentWorkflowPersistenceError from None return plan @staticmethod def _insert_plan( connection: Connection[WorkflowRow], source: DocumentSource, plan: ArtifactPlan, *, embedding_model: str, embedding_dimension: int, ) -> None: connection.execute( INSERT_VERSION_SQL, ( plan.version_id, source.document_id, plan.parser_profile_hash, plan.normalization_profile_hash, plan.chunk_profile_hash, plan.cloud_policy_id, plan.outbound_manifest_sha256, plan.review_state, plan.version_status, plan.expected_chunk_count, plan.version_error_code, ), ) with connection.cursor() as cursor: pages = [ ( page.id, plan.version_id, page.ordinal, page.page_number, page.text, page.text_sha256, page.line_start, page.line_end, ) for page in plan.pages ] if pages: cursor.executemany(INSERT_PAGE_SQL, pages) blocks = [ ( block.id, plan.version_id, block.page_id, block.ordinal, block.kind, block.text, block.text_sha256, _json(block.section_path), block.anchor_id, block.normalized_text_sha256, block.char_start, block.char_end, block.line_start, block.line_end, block.page_start, block.page_end, ) for block in plan.blocks ] if blocks: cursor.executemany(INSERT_BLOCK_SQL, blocks) if plan.outbound_manifest_sha256 is not None and plan.chunks: cursor.executemany( INSERT_MANIFEST_ITEM_SQL, [ ( plan.version_id, chunk.ordinal, plan.outbound_manifest_sha256, chunk.cloud_text_sha256, chunk.embedding_text_sha256, ) for chunk in plan.chunks ], ) chunks = [ ( chunk.id, source.knowledge_base_id, source.document_id, plan.version_id, source.access_scope_id, chunk.ordinal, chunk.display_text, chunk.cloud_text, chunk.cloud_text_sha256, chunk.embedding_prefix, chunk.embedding_text, chunk.embedding_text_sha256, plan.outbound_manifest_sha256, chunk.token_count, chunk.page_start, chunk.page_end, _json(chunk.section_path), _json(chunk.metadata), embedding_model, embedding_dimension, ) for chunk in plan.chunks ] if chunks: cursor.executemany(INSERT_CHUNK_SQL, chunks) @staticmethod def _verify_plan( connection: Connection[WorkflowRow], source: DocumentSource, plan: ArtifactPlan, ) -> None: version = connection.execute( SELECT_VERSION_SQL, (plan.version_id, source.document_id) ).fetchone() expected_version = ( plan.version_id, source.document_id, plan.parser_profile_hash, plan.normalization_profile_hash, plan.chunk_profile_hash, plan.cloud_policy_id, plan.outbound_manifest_sha256, plan.expected_chunk_count, ) if version is None or _version_tuple(version) != expected_version: raise ArtifactConflictError pages = list(connection.execute(SELECT_PAGES_SQL, (plan.version_id,)).fetchall()) blocks = list(connection.execute(SELECT_BLOCKS_SQL, (plan.version_id,)).fetchall()) manifests = list( connection.execute(SELECT_MANIFEST_ITEMS_SQL, (plan.version_id,)).fetchall() ) chunks = list(connection.execute(SELECT_CHUNKS_SQL, (plan.version_id,)).fetchall()) if tuple(_page_tuple(row) for row in pages) != tuple( _planned_page_tuple(value) for value in plan.pages ): raise ArtifactConflictError if tuple(_block_tuple(row) for row in blocks) != tuple( _planned_block_tuple(value) for value in plan.blocks ): raise ArtifactConflictError expected_manifests = ( tuple( ( chunk.ordinal, plan.outbound_manifest_sha256, chunk.cloud_text_sha256, chunk.embedding_text_sha256, ) for chunk in plan.chunks ) if plan.outbound_manifest_sha256 is not None else () ) if tuple(_manifest_tuple(row) for row in manifests) != expected_manifests: raise ArtifactConflictError if tuple(_chunk_tuple(row) for row in chunks) != tuple( _planned_chunk_tuple(value) for value in plan.chunks ): raise ArtifactConflictError def plan_artifact( source: DocumentSource, artifact: IngestionArtifact, *, cloud_policy: CloudTextPolicy, ) -> ArtifactPlan: if artifact.raw_sha256 != source.raw_sha256: raise ArtifactConflictError try: source_version_id = uuid.UUID(artifact.document_version_id) except ValueError: raise ArtifactConflictError from None version_id = _namespaced_id( "version", source.knowledge_base_id, source.document_id, source_version_id, ) if artifact.status is IngestionStatus.OCR_REQUIRED: if artifact.pages or artifact.blocks or artifact.chunks or artifact.manifest is not None: raise ArtifactConflictError error_code = artifact.limitations[0] if artifact.limitations else "OCR_REQUIRED" if _ERROR_CODE.fullmatch(error_code) is None: raise ArtifactConflictError return ArtifactPlan( version_id=version_id, parser_profile_hash=_checked_hash(artifact.parser_profile_hash), normalization_profile_hash=_NORMALIZATION_PROFILE_HASH, chunk_profile_hash=_checked_hash(artifact.chunk_profile_hash), cloud_policy_id=cloud_policy.policy_id, outbound_manifest_sha256=None, review_state="QUARANTINED_LOCAL_REVIEW", version_status="PENDING", version_error_code=error_code, expected_chunk_count=0, document_status="LOCAL_OCR_REQUIRED", job_stage="OCR_REQUIRED", pages=(), blocks=(), chunks=(), ) if ( artifact.status is not IngestionStatus.READY_FOR_LOCAL_REVIEW or artifact.manifest is None or artifact.normalized_text_sha256 is None or not artifact.pages or not artifact.blocks or not artifact.chunks or artifact.manifest.cloud_policy_id != cloud_policy.policy_id or artifact.manifest.raw_sha256 != source.raw_sha256 or artifact.manifest.normalized_text_sha256 != artifact.normalized_text_sha256 ): raise ArtifactConflictError page_number_ids: dict[int | None, uuid.UUID] = {} pages: list[PlannedPage] = [] for page in artifact.pages: page_id = _namespaced_id("page", source.knowledge_base_id, version_id, page.page_id) if page.page_number in page_number_ids: raise ArtifactConflictError page_number_ids[page.page_number] = page_id pages.append( PlannedPage( id=page_id, ordinal=page.ordinal, page_number=page.page_number, text=page.text, text_sha256=_checked_hash(page.text_sha256), line_start=page.line_start, line_end=page.line_end, ) ) block_ids: dict[str, uuid.UUID] = {} blocks: list[PlannedBlock] = [] for block in artifact.blocks: block_id = _namespaced_id("block", source.knowledge_base_id, version_id, block.block_id) block_ids[block.block_id] = block_id block_page_id = page_number_ids.get(block.anchor.page_start) if block_page_id is None: raise ArtifactConflictError blocks.append( PlannedBlock( id=block_id, page_id=block_page_id, ordinal=block.ordinal, kind=block.kind.value, text=block.text, text_sha256=_checked_hash(block.text_sha256), section_path=block.section_path, anchor_id=_checked_hash(block.anchor.anchor_id), normalized_text_sha256=_checked_hash(block.anchor.normalized_text_sha256), char_start=block.anchor.char_start, char_end=block.anchor.char_end, line_start=block.anchor.line_start, line_end=block.anchor.line_end, page_start=block.anchor.page_start, page_end=block.anchor.page_end, ) ) chunks: list[PlannedChunk] = [] for chunk in artifact.chunks: try: remapped_blocks = tuple(str(block_ids[value]) for value in chunk.anchor.block_ids) except KeyError: raise ArtifactConflictError from None chunk_id = _namespaced_id("chunk", source.knowledge_base_id, version_id, chunk.chunk_id) chunks.append( PlannedChunk( id=chunk_id, ordinal=chunk.ordinal, display_text=chunk.display_text, cloud_text=chunk.cloud_text, cloud_text_sha256=_checked_hash(chunk.cloud_text_sha256), embedding_prefix=chunk.embedding_prefix, embedding_text=chunk.embedding_text, embedding_text_sha256=_checked_hash(chunk.embedding_text_sha256), token_count=chunk.token_count, page_start=chunk.anchor.page_start, page_end=chunk.anchor.page_end, section_path=chunk.section_path, metadata={ "source_anchor": { "anchor_id": _checked_hash(chunk.anchor.anchor_id), "normalized_text_sha256": _checked_hash( chunk.anchor.normalized_text_sha256 ), "char_start": chunk.anchor.char_start, "char_end": chunk.anchor.char_end, "line_start": chunk.anchor.line_start, "line_end": chunk.anchor.line_end, "page_start": chunk.anchor.page_start, "page_end": chunk.anchor.page_end, "block_ids": remapped_blocks, }, "ingestion_profile": "document-ingestion-v1", }, ) ) if [item.ordinal for item in artifact.manifest.items] != [chunk.ordinal for chunk in chunks]: raise ArtifactConflictError for item, source_chunk, planned_chunk in zip( artifact.manifest.items, artifact.chunks, chunks, strict=True ): if ( item.chunk_id != source_chunk.chunk_id or item.anchor_id != source_chunk.anchor.anchor_id or item.cloud_text_sha256 != planned_chunk.cloud_text_sha256 or item.embedding_text_sha256 != planned_chunk.embedding_text_sha256 ): raise ArtifactConflictError return ArtifactPlan( version_id=version_id, parser_profile_hash=_checked_hash(artifact.parser_profile_hash), normalization_profile_hash=_NORMALIZATION_PROFILE_HASH, chunk_profile_hash=_checked_hash(artifact.chunk_profile_hash), cloud_policy_id=cloud_policy.policy_id, outbound_manifest_sha256=_checked_hash(artifact.manifest.manifest_sha256), review_state="LOCAL_PARSED_PENDING_CLOUD_REVIEW", version_status="PROCESSING", version_error_code=None, expected_chunk_count=len(chunks), document_status="LOCAL_PARSED_PENDING_CLOUD_REVIEW", job_stage="LOCAL_PARSED_PENDING_CLOUD_REVIEW", pages=tuple(pages), blocks=tuple(blocks), chunks=tuple(chunks), ) def _validate_job_contract(job: BackgroundJob) -> tuple[uuid.UUID, uuid.UUID]: if ( job.lease.job_id != job.id or not job.lease.worker_id.strip() or job.job_type != "PARSE_DOCUMENT" or job.required_capability != "document_parse" or job.resource_type != "document" or set(job.payload) != {"upload_id", "document_id"} ): raise InvalidDocumentJobError try: upload_id = uuid.UUID(cast(str, job.payload["upload_id"])) document_id = uuid.UUID(cast(str, job.payload["document_id"])) except (ValueError, TypeError, KeyError): raise InvalidDocumentJobError from None if document_id != job.resource_id: raise InvalidDocumentJobError return upload_id, document_id def _source_from_row(row: Mapping[str, object]) -> DocumentSource: source = DocumentSource( upload_id=_uuid(row, "upload_id"), document_id=_uuid(row, "document_id"), knowledge_base_id=_uuid(row, "knowledge_base_id"), access_scope_id=_uuid(row, "access_scope_id"), filename=_text(row, "original_filename"), mime_type=_text(row, "declared_mime_type"), storage_key=_uuid(row, "storage_key"), byte_size=_integer(row, "actual_size"), raw_sha256=_checked_hash(_text(row, "actual_sha256")), ) if ( _text(row, "document_filename") != source.filename or _text(row, "document_mime_type") != source.mime_type or _checked_hash(_text(row, "document_raw_sha256")) != source.raw_sha256 or _text(row, "document_storage_key") != str(source.storage_key) ): raise ArtifactConflictError return source def _fence_parameters(lease: JobLease, source: DocumentSource) -> tuple[object, ...]: return ( lease.job_id, lease.worker_id, lease.lease_token, source.document_id, source.upload_id, source.document_id, source.knowledge_base_id, source.access_scope_id, source.byte_size, source.raw_sha256, source.storage_key, source.raw_sha256, str(source.storage_key), ) def _namespaced_id(kind: str, *parts: object) -> uuid.UUID: return uuid.uuid5(_DB_ID_NAMESPACE, "\x1f".join((kind, *(str(value) for value in parts)))) def _checked_hash(value: str) -> str: if _HASH.fullmatch(value) is None: raise ArtifactConflictError return value def _json(value: object) -> str: return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")) def _default_connection_factory(dsn: str, timeout: int) -> Connection[WorkflowRow]: return psycopg.connect( dsn, connect_timeout=timeout, row_factory=dict_row, application_name="geological-rag-document-workflow", ) def _version_tuple(row: Mapping[str, object]) -> tuple[object, ...]: return ( _uuid(row, "id"), _uuid(row, "document_id"), _checked_hash(_text(row, "parser_profile_hash")), _checked_hash(_text(row, "normalization_profile_hash")), _checked_hash(_text(row, "chunk_profile_hash")), _text(row, "cloud_policy_id"), _optional_hash(row, "outbound_manifest_sha256"), _integer(row, "expected_chunk_count"), ) def _planned_page_tuple(value: PlannedPage) -> tuple[object, ...]: return ( value.id, value.ordinal, value.page_number, value.text, value.text_sha256, value.line_start, value.line_end, ) def _page_tuple(row: Mapping[str, object]) -> tuple[object, ...]: return ( _uuid(row, "id"), _integer(row, "ordinal"), _optional_integer(row, "page_number"), _text(row, "display_text"), _checked_hash(_text(row, "text_sha256")), _integer(row, "line_start"), _integer(row, "line_end"), ) def _planned_block_tuple(value: PlannedBlock) -> tuple[object, ...]: return ( value.id, value.page_id, value.ordinal, value.kind, value.text, value.text_sha256, value.section_path, value.anchor_id, value.normalized_text_sha256, value.char_start, value.char_end, value.line_start, value.line_end, value.page_start, value.page_end, ) def _block_tuple(row: Mapping[str, object]) -> tuple[object, ...]: return ( _uuid(row, "id"), _uuid(row, "page_id"), _integer(row, "ordinal"), _text(row, "block_kind"), _text(row, "display_text"), _checked_hash(_text(row, "text_sha256")), _string_tuple(row, "section_path"), _checked_hash(_text(row, "anchor_id")), _checked_hash(_text(row, "normalized_text_sha256")), _integer(row, "char_start"), _integer(row, "char_end"), _integer(row, "line_start"), _integer(row, "line_end"), _optional_integer(row, "page_start"), _optional_integer(row, "page_end"), ) def _manifest_tuple(row: Mapping[str, object]) -> tuple[object, ...]: return ( _integer(row, "ordinal"), _checked_hash(_text(row, "outbound_manifest_sha256")), _checked_hash(_text(row, "cloud_text_sha256")), _checked_hash(_text(row, "embedding_text_sha256")), ) def _planned_chunk_tuple(value: PlannedChunk) -> tuple[object, ...]: return ( value.id, value.ordinal, value.display_text, value.cloud_text, value.cloud_text_sha256, value.embedding_prefix, value.embedding_text, value.embedding_text_sha256, value.token_count, value.page_start, value.page_end, value.section_path, _canonical_json(value.metadata), ) def _chunk_tuple(row: Mapping[str, object]) -> tuple[object, ...]: metadata = row.get("metadata") if not isinstance(metadata, dict): raise ArtifactConflictError return ( _uuid(row, "id"), _integer(row, "ordinal"), _text(row, "display_text"), _text(row, "cloud_text"), _checked_hash(_text(row, "cloud_text_sha256")), _text(row, "embedding_prefix"), _text(row, "embedding_text"), _checked_hash(_text(row, "embedding_text_sha256")), _integer(row, "token_count"), _optional_integer(row, "page_start"), _optional_integer(row, "page_end"), _string_tuple(row, "section_path"), metadata, ) def _canonical_json(value: object) -> object: """Match PostgreSQL jsonb's array representation for tuple-backed plans.""" return json.loads(_json(value)) def _uuid(row: Mapping[str, object], name: str) -> uuid.UUID: value = row.get(name) if isinstance(value, uuid.UUID): return value if isinstance(value, str): try: return uuid.UUID(value) except ValueError: pass raise ArtifactConflictError def _text(row: Mapping[str, object], name: str) -> str: value = row.get(name) if not isinstance(value, str) or not value.strip(): raise ArtifactConflictError return value def _integer(row: Mapping[str, object], name: str) -> int: value = row.get(name) if isinstance(value, bool) or not isinstance(value, int): raise ArtifactConflictError return value def _optional_integer(row: Mapping[str, object], name: str) -> int | None: return None if row.get(name) is None else _integer(row, name) def _optional_hash(row: Mapping[str, object], name: str) -> str | None: return None if row.get(name) is None else _checked_hash(_text(row, name)) def _string_tuple(row: Mapping[str, object], name: str) -> tuple[str, ...]: value = row.get(name) if not isinstance(value, list) or not all(isinstance(item, str) for item in value): raise ArtifactConflictError return tuple(value)