YoVinchen
ecdb10c37a
Some checks failed
verify / verify (push) Has been cancelled
Separate local parsing from model indexing, bind review decisions to immutable manifests, persist vectors behind active profiles, and expose retrieval, chat, evaluation, and document workflows through the React workbench. Constraint: Live Bailian authentication currently fails for all three configured capabilities Rejected: Direct upload-to-embedding flow | bypasses local review and manifest binding Confidence: high Scope-risk: broad Directive: Keep private-data deployment blocked until authentication, RBAC, and separate database roles land Tested: make verify; fresh and replay Docker document smoke; worker recovery smoke; frozen synthetic evaluation; migration 0003-0004 roundtrip Not-tested: Successful live Bailian calls, OCR, real multi-user authorization
1022 lines
30 KiB
Python
1022 lines
30 KiB
Python
"""Short-transaction PostgreSQL persistence for governed document ingestion."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import re
|
|
import uuid
|
|
from collections.abc import Callable, Mapping, Sequence
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
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
|
|
|
|
type DocumentRow = dict[str, Any]
|
|
type ConnectionFactory = Callable[[str, int], Connection[DocumentRow]]
|
|
|
|
_HASH = re.compile(r"^[0-9a-f]{64}$")
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class DocumentActor:
|
|
subject: str
|
|
knowledge_base_id: uuid.UUID
|
|
access_scope_id: uuid.UUID
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class DocumentUpload:
|
|
id: uuid.UUID
|
|
filename: str
|
|
declared_mime_type: str
|
|
expected_size: int
|
|
expected_sha256: str
|
|
storage_key: uuid.UUID
|
|
actual_size: int | None
|
|
actual_sha256: str | None
|
|
status: str
|
|
document_id: uuid.UUID | None
|
|
parse_job_id: uuid.UUID | None
|
|
created_at: datetime
|
|
updated_at: datetime
|
|
completed_at: datetime | None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class SafeJob:
|
|
id: uuid.UUID
|
|
job_type: str
|
|
stage: str
|
|
status: str
|
|
progress: int
|
|
attempt: int
|
|
max_attempts: int
|
|
last_error_code: str | None
|
|
created_at: datetime
|
|
updated_at: datetime
|
|
finished_at: datetime | None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class DocumentSummary:
|
|
id: uuid.UUID
|
|
filename: str
|
|
mime_type: str
|
|
raw_sha256: str
|
|
status: str
|
|
active_version_id: uuid.UUID | None
|
|
created_at: datetime
|
|
updated_at: datetime
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class DocumentDetail:
|
|
document: DocumentSummary
|
|
version_count: int
|
|
page_count: int
|
|
block_count: int
|
|
chunk_count: int
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ReviewVersion:
|
|
id: uuid.UUID
|
|
review_state: str
|
|
review_revision: int
|
|
status: str
|
|
parser_profile_hash: str
|
|
chunk_profile_hash: str
|
|
cloud_policy_id: str
|
|
outbound_manifest_sha256: str | None
|
|
expected_chunk_count: int | None
|
|
error_code: str | None
|
|
created_at: datetime
|
|
completed_at: datetime | None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ReviewPage:
|
|
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 ReviewBlock:
|
|
id: uuid.UUID
|
|
ordinal: int
|
|
kind: str
|
|
text: str
|
|
text_sha256: str
|
|
section_path: tuple[str, ...]
|
|
anchor_id: 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 ReviewChunk:
|
|
ordinal: int
|
|
display_text: str
|
|
cloud_text: str
|
|
cloud_text_sha256: str
|
|
embedding_text_sha256: str
|
|
token_count: int
|
|
page_start: int | None
|
|
page_end: int | None
|
|
section_path: tuple[str, ...]
|
|
approval_status: str
|
|
index_status: str
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ReviewBundle:
|
|
document: DocumentSummary
|
|
version: ReviewVersion | None
|
|
pages: tuple[ReviewPage, ...]
|
|
blocks: tuple[ReviewBlock, ...]
|
|
chunks: tuple[ReviewChunk, ...]
|
|
next_ordinal: int | None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class DocumentListPage:
|
|
items: tuple[DocumentSummary, ...]
|
|
next_cursor: uuid.UUID | None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class CompletedUpload:
|
|
upload: DocumentUpload
|
|
document: DocumentSummary
|
|
job: SafeJob
|
|
|
|
|
|
class DocumentPersistenceError(RuntimeError):
|
|
def __init__(self) -> None:
|
|
super().__init__("document persistence unavailable")
|
|
|
|
|
|
class IdempotencyConflictError(RuntimeError):
|
|
def __init__(self) -> None:
|
|
super().__init__("idempotency key conflicts with an earlier request")
|
|
|
|
|
|
class UploadStateConflictError(RuntimeError):
|
|
def __init__(self) -> None:
|
|
super().__init__("document upload state does not permit this operation")
|
|
|
|
|
|
class DocumentsRepository(Protocol):
|
|
def create_upload(
|
|
self,
|
|
*,
|
|
actor: DocumentActor,
|
|
idempotency_key_hash: str,
|
|
request_fingerprint: str,
|
|
filename: str,
|
|
declared_mime_type: str,
|
|
expected_size: int,
|
|
expected_sha256: str,
|
|
storage_key: uuid.UUID,
|
|
trace_id: uuid.UUID,
|
|
) -> tuple[DocumentUpload, bool]: ...
|
|
|
|
def get_upload(self, actor: DocumentActor, upload_id: uuid.UUID) -> DocumentUpload | None: ...
|
|
|
|
def mark_upload_stored(
|
|
self,
|
|
*,
|
|
actor: DocumentActor,
|
|
upload_id: uuid.UUID,
|
|
actual_size: int,
|
|
actual_sha256: str,
|
|
trace_id: uuid.UUID,
|
|
) -> DocumentUpload: ...
|
|
|
|
def complete_upload(
|
|
self,
|
|
*,
|
|
actor: DocumentActor,
|
|
upload_id: uuid.UUID,
|
|
trace_id: uuid.UUID,
|
|
) -> CompletedUpload: ...
|
|
|
|
def get_job(self, actor: DocumentActor, job_id: uuid.UUID) -> SafeJob | None: ...
|
|
|
|
def list_documents(
|
|
self,
|
|
actor: DocumentActor,
|
|
*,
|
|
cursor: uuid.UUID | None,
|
|
limit: int,
|
|
) -> DocumentListPage: ...
|
|
|
|
def get_document(
|
|
self, actor: DocumentActor, document_id: uuid.UUID
|
|
) -> DocumentDetail | None: ...
|
|
|
|
def get_review_bundle(
|
|
self,
|
|
actor: DocumentActor,
|
|
document_id: uuid.UUID,
|
|
*,
|
|
after_ordinal: int,
|
|
limit: int,
|
|
) -> ReviewBundle | None: ...
|
|
|
|
|
|
CREATE_UPLOAD_SQL = """
|
|
WITH inserted AS (
|
|
INSERT INTO rag.document_uploads (
|
|
actor_subject,
|
|
knowledge_base_id,
|
|
access_scope_id,
|
|
idempotency_key_hash,
|
|
request_fingerprint,
|
|
original_filename,
|
|
declared_mime_type,
|
|
expected_size,
|
|
expected_sha256,
|
|
storage_key
|
|
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT (actor_subject, idempotency_key_hash) DO NOTHING
|
|
RETURNING *, true AS created
|
|
), audit AS (
|
|
INSERT INTO rag.document_upload_events (
|
|
upload_id, actor_subject, event_type, trace_id
|
|
)
|
|
SELECT id, actor_subject, 'CREATED', %s
|
|
FROM inserted
|
|
), selected AS (
|
|
SELECT *, false AS created
|
|
FROM rag.document_uploads
|
|
WHERE actor_subject = %s
|
|
AND knowledge_base_id = %s
|
|
AND access_scope_id = %s
|
|
AND idempotency_key_hash = %s
|
|
)
|
|
SELECT * FROM inserted
|
|
UNION ALL
|
|
SELECT * FROM selected
|
|
LIMIT 1
|
|
"""
|
|
|
|
GET_UPLOAD_SQL = """
|
|
SELECT *
|
|
FROM rag.document_uploads
|
|
WHERE id = %s
|
|
AND actor_subject = %s
|
|
AND knowledge_base_id = %s
|
|
AND access_scope_id = %s
|
|
LIMIT 1
|
|
"""
|
|
|
|
MARK_STORED_SQL = """
|
|
WITH locked AS (
|
|
SELECT *
|
|
FROM rag.document_uploads
|
|
WHERE id = %s
|
|
AND actor_subject = %s
|
|
AND knowledge_base_id = %s
|
|
AND access_scope_id = %s
|
|
AND status IN ('CREATED', 'STORED')
|
|
FOR UPDATE
|
|
), updated AS (
|
|
UPDATE rag.document_uploads AS upload
|
|
SET status = 'STORED',
|
|
actual_size = %s,
|
|
actual_sha256 = %s,
|
|
updated_at = now()
|
|
FROM locked
|
|
WHERE upload.id = locked.id
|
|
AND locked.expected_size = %s
|
|
AND locked.expected_sha256 = %s
|
|
RETURNING upload.*, locked.status AS previous_status
|
|
), audit AS (
|
|
INSERT INTO rag.document_upload_events (
|
|
upload_id, actor_subject, event_type, trace_id
|
|
)
|
|
SELECT id, actor_subject, 'STORED', %s
|
|
FROM updated
|
|
WHERE previous_status = 'CREATED'
|
|
)
|
|
SELECT * FROM updated
|
|
"""
|
|
|
|
COMPLETE_UPLOAD_SQL = """
|
|
WITH locked AS (
|
|
SELECT *
|
|
FROM rag.document_uploads
|
|
WHERE id = %s
|
|
AND actor_subject = %s
|
|
AND knowledge_base_id = %s
|
|
AND access_scope_id = %s
|
|
AND status IN ('STORED', 'COMPLETED')
|
|
FOR UPDATE
|
|
), upserted_document AS (
|
|
INSERT INTO rag.documents (
|
|
knowledge_base_id,
|
|
access_scope_id,
|
|
raw_sha256,
|
|
filename,
|
|
storage_key,
|
|
mime_type,
|
|
status
|
|
)
|
|
SELECT
|
|
knowledge_base_id,
|
|
access_scope_id,
|
|
actual_sha256,
|
|
original_filename,
|
|
storage_key::text,
|
|
declared_mime_type,
|
|
'QUARANTINED_LOCAL_REVIEW'
|
|
FROM locked
|
|
ON CONFLICT (knowledge_base_id, raw_sha256)
|
|
DO UPDATE SET updated_at = rag.documents.updated_at
|
|
WHERE rag.documents.access_scope_id = EXCLUDED.access_scope_id
|
|
RETURNING rag.documents.*
|
|
), upserted_job AS (
|
|
INSERT INTO rag.background_jobs (
|
|
job_type,
|
|
required_capability,
|
|
resource_type,
|
|
resource_id,
|
|
idempotency_key,
|
|
payload,
|
|
stage,
|
|
status,
|
|
max_attempts
|
|
)
|
|
SELECT
|
|
'PARSE_DOCUMENT',
|
|
'document_parse',
|
|
'document',
|
|
document.id,
|
|
'parse-document:' || document.id::text || ':' || document.raw_sha256,
|
|
jsonb_build_object(
|
|
'document_id', document.id::text,
|
|
'upload_id', locked.id::text
|
|
),
|
|
'PENDING',
|
|
'QUEUED',
|
|
3
|
|
FROM upserted_document AS document
|
|
CROSS JOIN locked
|
|
ON CONFLICT (job_type, idempotency_key)
|
|
DO UPDATE SET updated_at = rag.background_jobs.updated_at
|
|
RETURNING rag.background_jobs.*
|
|
), updated_upload AS (
|
|
UPDATE rag.document_uploads AS upload
|
|
SET status = 'COMPLETED',
|
|
document_id = document.id,
|
|
parse_job_id = job.id,
|
|
completed_at = COALESCE(upload.completed_at, now()),
|
|
updated_at = now()
|
|
FROM locked
|
|
CROSS JOIN upserted_document AS document
|
|
CROSS JOIN upserted_job AS job
|
|
WHERE upload.id = locked.id
|
|
RETURNING upload.*, locked.status AS previous_status
|
|
), audit AS (
|
|
INSERT INTO rag.document_upload_events (
|
|
upload_id, actor_subject, event_type, trace_id
|
|
)
|
|
SELECT id, actor_subject, 'COMPLETED', %s
|
|
FROM updated_upload
|
|
WHERE previous_status = 'STORED'
|
|
)
|
|
SELECT
|
|
upload.*,
|
|
document.filename AS document_filename,
|
|
document.mime_type AS document_mime_type,
|
|
document.raw_sha256 AS document_raw_sha256,
|
|
document.status AS document_status,
|
|
document.active_version_id AS document_active_version_id,
|
|
document.created_at AS document_created_at,
|
|
document.updated_at AS document_updated_at,
|
|
job.id AS job_id,
|
|
job.job_type AS job_job_type,
|
|
job.stage AS job_stage,
|
|
job.status AS job_status,
|
|
job.progress AS job_progress,
|
|
job.attempt AS job_attempt,
|
|
job.max_attempts AS job_max_attempts,
|
|
job.last_error_code AS job_last_error_code,
|
|
job.created_at AS job_created_at,
|
|
job.updated_at AS job_updated_at,
|
|
job.finished_at AS job_finished_at
|
|
FROM updated_upload AS upload
|
|
JOIN upserted_document AS document ON document.id = upload.document_id
|
|
JOIN upserted_job AS job ON job.id = upload.parse_job_id
|
|
"""
|
|
|
|
GET_JOB_SQL = """
|
|
SELECT
|
|
job.id,
|
|
job.job_type,
|
|
job.stage,
|
|
job.status,
|
|
job.progress,
|
|
job.attempt,
|
|
job.max_attempts,
|
|
job.last_error_code,
|
|
job.created_at,
|
|
job.updated_at,
|
|
job.finished_at
|
|
FROM rag.background_jobs AS job
|
|
WHERE job.id = %s
|
|
AND EXISTS (
|
|
SELECT 1
|
|
FROM rag.document_uploads AS upload
|
|
LEFT JOIN rag.document_versions AS version
|
|
ON job.job_type = 'EMBED_DOCUMENT'
|
|
AND job.resource_type = 'document_version'
|
|
AND job.resource_id = version.id
|
|
AND version.document_id = upload.document_id
|
|
WHERE upload.actor_subject = %s
|
|
AND upload.knowledge_base_id = %s
|
|
AND upload.access_scope_id = %s
|
|
AND (
|
|
upload.parse_job_id = job.id
|
|
OR version.id IS NOT NULL
|
|
)
|
|
)
|
|
LIMIT 1
|
|
"""
|
|
|
|
LIST_DOCUMENTS_SQL = """
|
|
SELECT
|
|
document.id,
|
|
document.filename,
|
|
document.mime_type,
|
|
document.raw_sha256,
|
|
document.status,
|
|
document.active_version_id,
|
|
document.created_at,
|
|
document.updated_at
|
|
FROM rag.documents AS document
|
|
WHERE document.knowledge_base_id = %s
|
|
AND document.access_scope_id = %s
|
|
AND document.deleted_at IS NULL
|
|
AND (%s::uuid IS NULL OR document.id < %s)
|
|
ORDER BY document.id DESC
|
|
LIMIT %s
|
|
"""
|
|
|
|
GET_DOCUMENT_SQL = """
|
|
SELECT
|
|
document.id,
|
|
document.filename,
|
|
document.mime_type,
|
|
document.raw_sha256,
|
|
document.status,
|
|
document.active_version_id,
|
|
document.created_at,
|
|
document.updated_at,
|
|
(SELECT count(*) FROM rag.document_versions AS version
|
|
WHERE version.document_id = document.id)::integer AS version_count,
|
|
(SELECT count(*) FROM rag.document_pages AS page
|
|
JOIN rag.document_versions AS version ON version.id = page.document_version_id
|
|
WHERE version.document_id = document.id)::integer AS page_count,
|
|
(SELECT count(*) FROM rag.document_blocks AS block
|
|
JOIN rag.document_versions AS version ON version.id = block.document_version_id
|
|
WHERE version.document_id = document.id)::integer AS block_count,
|
|
(SELECT count(*) FROM rag.chunks AS chunk
|
|
WHERE chunk.document_id = document.id)::integer AS chunk_count
|
|
FROM rag.documents AS document
|
|
WHERE document.id = %s
|
|
AND document.knowledge_base_id = %s
|
|
AND document.access_scope_id = %s
|
|
AND document.deleted_at IS NULL
|
|
LIMIT 1
|
|
"""
|
|
|
|
GET_LATEST_VERSION_SQL = """
|
|
SELECT
|
|
version.id,
|
|
version.review_state,
|
|
version.review_revision,
|
|
version.status,
|
|
version.parser_profile_hash,
|
|
version.chunk_profile_hash,
|
|
version.cloud_policy_id,
|
|
version.outbound_manifest_sha256,
|
|
version.expected_chunk_count,
|
|
version.error_code,
|
|
version.created_at,
|
|
version.completed_at
|
|
FROM rag.document_versions AS version
|
|
WHERE version.document_id = %s
|
|
ORDER BY version.created_at DESC, version.id DESC
|
|
LIMIT 1
|
|
"""
|
|
|
|
GET_REVIEW_PAGES_SQL = """
|
|
SELECT id, ordinal, page_number, display_text, text_sha256, line_start, line_end
|
|
FROM rag.document_pages
|
|
WHERE document_version_id = %s
|
|
AND ordinal > %s
|
|
ORDER BY ordinal
|
|
LIMIT %s
|
|
"""
|
|
|
|
GET_REVIEW_BLOCKS_SQL = """
|
|
SELECT
|
|
id, ordinal, block_kind, display_text, text_sha256, section_path,
|
|
anchor_id, char_start, char_end, line_start, line_end, page_start, page_end
|
|
FROM rag.document_blocks
|
|
WHERE document_version_id = %s
|
|
AND ordinal > %s
|
|
ORDER BY ordinal
|
|
LIMIT %s
|
|
"""
|
|
|
|
GET_REVIEW_CHUNKS_SQL = """
|
|
SELECT
|
|
ordinal, display_text, cloud_text, cloud_text_sha256, embedding_text_sha256,
|
|
token_count, page_start, page_end, section_path, approval_status, index_status
|
|
FROM rag.chunks
|
|
WHERE document_version_id = %s
|
|
AND ordinal > %s
|
|
AND deleted_at IS NULL
|
|
ORDER BY ordinal
|
|
LIMIT %s
|
|
"""
|
|
|
|
|
|
def idempotency_key_hash(actor: DocumentActor, key: uuid.UUID) -> str:
|
|
return _sha256(f"{actor.subject}\x1f{key}")
|
|
|
|
|
|
def upload_request_fingerprint(
|
|
*, filename: str, declared_mime_type: str, expected_size: int, expected_sha256: str
|
|
) -> str:
|
|
serialized = json.dumps(
|
|
{
|
|
"filename": filename,
|
|
"declared_mime_type": declared_mime_type,
|
|
"expected_size": expected_size,
|
|
"expected_sha256": expected_sha256,
|
|
},
|
|
ensure_ascii=False,
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
)
|
|
return _sha256(serialized)
|
|
|
|
|
|
def _sha256(value: str) -> str:
|
|
return hashlib.sha256(value.encode()).hexdigest()
|
|
|
|
|
|
def _default_connection_factory(dsn: str, timeout: int) -> Connection[DocumentRow]:
|
|
return psycopg.connect(
|
|
dsn,
|
|
connect_timeout=timeout,
|
|
row_factory=dict_row,
|
|
application_name="geological-rag-documents",
|
|
)
|
|
|
|
|
|
class PostgresDocumentsRepository:
|
|
def __init__(
|
|
self,
|
|
settings: Settings,
|
|
*,
|
|
connect_timeout: int = 5,
|
|
connection_factory: ConnectionFactory = _default_connection_factory,
|
|
) -> None:
|
|
self._settings = settings
|
|
self._connect_timeout = connect_timeout
|
|
self._connection_factory = connection_factory
|
|
|
|
def _dsn(self) -> str:
|
|
return (
|
|
self._settings.database_url()
|
|
.set(drivername="postgresql")
|
|
.render_as_string(hide_password=False)
|
|
)
|
|
|
|
def _one(self, statement: str, parameters: Sequence[object]) -> DocumentRow | None:
|
|
try:
|
|
with self._connection_factory(self._dsn(), self._connect_timeout) as connection:
|
|
with connection.transaction():
|
|
return connection.execute(statement, parameters).fetchone()
|
|
except (OSError, SecretFileError, psycopg.Error):
|
|
raise DocumentPersistenceError from None
|
|
|
|
def create_upload(
|
|
self,
|
|
*,
|
|
actor: DocumentActor,
|
|
idempotency_key_hash: str,
|
|
request_fingerprint: str,
|
|
filename: str,
|
|
declared_mime_type: str,
|
|
expected_size: int,
|
|
expected_sha256: str,
|
|
storage_key: uuid.UUID,
|
|
trace_id: uuid.UUID,
|
|
) -> tuple[DocumentUpload, bool]:
|
|
row = self._one(
|
|
CREATE_UPLOAD_SQL,
|
|
(
|
|
actor.subject,
|
|
actor.knowledge_base_id,
|
|
actor.access_scope_id,
|
|
idempotency_key_hash,
|
|
request_fingerprint,
|
|
filename,
|
|
declared_mime_type,
|
|
expected_size,
|
|
expected_sha256,
|
|
storage_key,
|
|
trace_id,
|
|
actor.subject,
|
|
actor.knowledge_base_id,
|
|
actor.access_scope_id,
|
|
idempotency_key_hash,
|
|
),
|
|
)
|
|
if row is None:
|
|
raise DocumentPersistenceError
|
|
if row.get("request_fingerprint") != request_fingerprint:
|
|
raise IdempotencyConflictError
|
|
return _upload(row), bool(row.get("created"))
|
|
|
|
def get_upload(self, actor: DocumentActor, upload_id: uuid.UUID) -> DocumentUpload | None:
|
|
row = self._one(GET_UPLOAD_SQL, _actor_resource(actor, upload_id))
|
|
return None if row is None else _upload(row)
|
|
|
|
def mark_upload_stored(
|
|
self,
|
|
*,
|
|
actor: DocumentActor,
|
|
upload_id: uuid.UUID,
|
|
actual_size: int,
|
|
actual_sha256: str,
|
|
trace_id: uuid.UUID,
|
|
) -> DocumentUpload:
|
|
row = self._one(
|
|
MARK_STORED_SQL,
|
|
(
|
|
upload_id,
|
|
actor.subject,
|
|
actor.knowledge_base_id,
|
|
actor.access_scope_id,
|
|
actual_size,
|
|
actual_sha256,
|
|
actual_size,
|
|
actual_sha256,
|
|
trace_id,
|
|
),
|
|
)
|
|
if row is None:
|
|
raise UploadStateConflictError
|
|
return _upload(row)
|
|
|
|
def complete_upload(
|
|
self,
|
|
*,
|
|
actor: DocumentActor,
|
|
upload_id: uuid.UUID,
|
|
trace_id: uuid.UUID,
|
|
) -> CompletedUpload:
|
|
row = self._one(
|
|
COMPLETE_UPLOAD_SQL,
|
|
(
|
|
upload_id,
|
|
actor.subject,
|
|
actor.knowledge_base_id,
|
|
actor.access_scope_id,
|
|
trace_id,
|
|
),
|
|
)
|
|
if row is None:
|
|
raise UploadStateConflictError
|
|
upload = _upload(row)
|
|
document = DocumentSummary(
|
|
id=_uuid(row, "document_id"),
|
|
filename=_text(row, "document_filename"),
|
|
mime_type=_text(row, "document_mime_type"),
|
|
raw_sha256=_hash(row, "document_raw_sha256"),
|
|
status=_text(row, "document_status"),
|
|
active_version_id=_optional_uuid(row, "document_active_version_id"),
|
|
created_at=_datetime(row, "document_created_at"),
|
|
updated_at=_datetime(row, "document_updated_at"),
|
|
)
|
|
job = _job(row, prefix="job_")
|
|
return CompletedUpload(upload, document, job)
|
|
|
|
def get_job(self, actor: DocumentActor, job_id: uuid.UUID) -> SafeJob | None:
|
|
row = self._one(GET_JOB_SQL, _actor_resource(actor, job_id))
|
|
return None if row is None else _job(row)
|
|
|
|
def list_documents(
|
|
self,
|
|
actor: DocumentActor,
|
|
*,
|
|
cursor: uuid.UUID | None,
|
|
limit: int,
|
|
) -> DocumentListPage:
|
|
rows = self._many(
|
|
LIST_DOCUMENTS_SQL,
|
|
(
|
|
actor.knowledge_base_id,
|
|
actor.access_scope_id,
|
|
cursor,
|
|
cursor,
|
|
limit + 1,
|
|
),
|
|
)
|
|
has_more = len(rows) > limit
|
|
items = tuple(_document(row) for row in rows[:limit])
|
|
next_cursor = items[-1].id if has_more and items else None
|
|
return DocumentListPage(items, next_cursor)
|
|
|
|
def get_document(self, actor: DocumentActor, document_id: uuid.UUID) -> DocumentDetail | None:
|
|
row = self._one(
|
|
GET_DOCUMENT_SQL,
|
|
(document_id, actor.knowledge_base_id, actor.access_scope_id),
|
|
)
|
|
if row is None:
|
|
return None
|
|
return DocumentDetail(
|
|
document=_document(row),
|
|
version_count=_integer(row, "version_count"),
|
|
page_count=_integer(row, "page_count"),
|
|
block_count=_integer(row, "block_count"),
|
|
chunk_count=_integer(row, "chunk_count"),
|
|
)
|
|
|
|
def get_review_bundle(
|
|
self,
|
|
actor: DocumentActor,
|
|
document_id: uuid.UUID,
|
|
*,
|
|
after_ordinal: int,
|
|
limit: int,
|
|
) -> ReviewBundle | None:
|
|
detail = self.get_document(actor, document_id)
|
|
if detail is None:
|
|
return None
|
|
try:
|
|
with self._connection_factory(self._dsn(), self._connect_timeout) as connection:
|
|
with connection.transaction():
|
|
version_row = connection.execute(
|
|
GET_LATEST_VERSION_SQL, (document_id,)
|
|
).fetchone()
|
|
if version_row is None:
|
|
return ReviewBundle(detail.document, None, (), (), (), None)
|
|
version = _review_version(version_row)
|
|
fetch_limit = limit + 1
|
|
pages_rows = list(
|
|
connection.execute(
|
|
GET_REVIEW_PAGES_SQL,
|
|
(version.id, after_ordinal, fetch_limit),
|
|
).fetchall()
|
|
)
|
|
block_rows = list(
|
|
connection.execute(
|
|
GET_REVIEW_BLOCKS_SQL,
|
|
(version.id, after_ordinal, fetch_limit),
|
|
).fetchall()
|
|
)
|
|
chunk_rows = list(
|
|
connection.execute(
|
|
GET_REVIEW_CHUNKS_SQL,
|
|
(version.id, after_ordinal, fetch_limit),
|
|
).fetchall()
|
|
)
|
|
except (OSError, SecretFileError, psycopg.Error):
|
|
raise DocumentPersistenceError from None
|
|
all_rows = (pages_rows, block_rows, chunk_rows)
|
|
has_more = any(len(rows) > limit for rows in all_rows)
|
|
maximum_ordinal = max(
|
|
(cast(int, row["ordinal"]) for rows in all_rows for row in rows[:limit]),
|
|
default=after_ordinal,
|
|
)
|
|
return ReviewBundle(
|
|
document=detail.document,
|
|
version=version,
|
|
pages=tuple(_review_page(row) for row in pages_rows[:limit]),
|
|
blocks=tuple(_review_block(row) for row in block_rows[:limit]),
|
|
chunks=tuple(_review_chunk(row) for row in chunk_rows[:limit]),
|
|
next_ordinal=maximum_ordinal if has_more else None,
|
|
)
|
|
|
|
def _many(self, statement: str, parameters: Sequence[object]) -> list[DocumentRow]:
|
|
try:
|
|
with self._connection_factory(self._dsn(), self._connect_timeout) as connection:
|
|
with connection.transaction():
|
|
return list(connection.execute(statement, parameters).fetchall())
|
|
except (OSError, SecretFileError, psycopg.Error):
|
|
raise DocumentPersistenceError from None
|
|
|
|
|
|
def _actor_resource(actor: DocumentActor, resource_id: uuid.UUID) -> tuple[object, ...]:
|
|
return (
|
|
resource_id,
|
|
actor.subject,
|
|
actor.knowledge_base_id,
|
|
actor.access_scope_id,
|
|
)
|
|
|
|
|
|
def _upload(row: Mapping[str, object]) -> DocumentUpload:
|
|
return DocumentUpload(
|
|
id=_uuid(row, "id"),
|
|
filename=_text(row, "original_filename"),
|
|
declared_mime_type=_text(row, "declared_mime_type"),
|
|
expected_size=_integer(row, "expected_size"),
|
|
expected_sha256=_hash(row, "expected_sha256"),
|
|
storage_key=_uuid(row, "storage_key"),
|
|
actual_size=_optional_integer(row, "actual_size"),
|
|
actual_sha256=_optional_hash(row, "actual_sha256"),
|
|
status=_text(row, "status"),
|
|
document_id=_optional_uuid(row, "document_id"),
|
|
parse_job_id=_optional_uuid(row, "parse_job_id"),
|
|
created_at=_datetime(row, "created_at"),
|
|
updated_at=_datetime(row, "updated_at"),
|
|
completed_at=_optional_datetime(row, "completed_at"),
|
|
)
|
|
|
|
|
|
def _document(row: Mapping[str, object]) -> DocumentSummary:
|
|
return DocumentSummary(
|
|
id=_uuid(row, "id"),
|
|
filename=_text(row, "filename"),
|
|
mime_type=_text(row, "mime_type"),
|
|
raw_sha256=_hash(row, "raw_sha256"),
|
|
status=_text(row, "status"),
|
|
active_version_id=_optional_uuid(row, "active_version_id"),
|
|
created_at=_datetime(row, "created_at"),
|
|
updated_at=_datetime(row, "updated_at"),
|
|
)
|
|
|
|
|
|
def _job(row: Mapping[str, object], *, prefix: str = "") -> SafeJob:
|
|
return SafeJob(
|
|
id=_uuid(row, f"{prefix}id" if prefix else "id"),
|
|
job_type=_text(row, f"{prefix}job_type"),
|
|
stage=_text(row, f"{prefix}stage"),
|
|
status=_text(row, f"{prefix}status"),
|
|
progress=_integer(row, f"{prefix}progress"),
|
|
attempt=_integer(row, f"{prefix}attempt"),
|
|
max_attempts=_integer(row, f"{prefix}max_attempts"),
|
|
last_error_code=_optional_text(row, f"{prefix}last_error_code"),
|
|
created_at=_datetime(row, f"{prefix}created_at"),
|
|
updated_at=_datetime(row, f"{prefix}updated_at"),
|
|
finished_at=_optional_datetime(row, f"{prefix}finished_at"),
|
|
)
|
|
|
|
|
|
def _review_version(row: Mapping[str, object]) -> ReviewVersion:
|
|
return ReviewVersion(
|
|
id=_uuid(row, "id"),
|
|
review_state=_text(row, "review_state"),
|
|
review_revision=_integer(row, "review_revision"),
|
|
status=_text(row, "status"),
|
|
parser_profile_hash=_hash(row, "parser_profile_hash"),
|
|
chunk_profile_hash=_hash(row, "chunk_profile_hash"),
|
|
cloud_policy_id=_text(row, "cloud_policy_id"),
|
|
outbound_manifest_sha256=_optional_hash(row, "outbound_manifest_sha256"),
|
|
expected_chunk_count=_optional_integer(row, "expected_chunk_count"),
|
|
error_code=_optional_text(row, "error_code"),
|
|
created_at=_datetime(row, "created_at"),
|
|
completed_at=_optional_datetime(row, "completed_at"),
|
|
)
|
|
|
|
|
|
def _review_page(row: Mapping[str, object]) -> ReviewPage:
|
|
return ReviewPage(
|
|
id=_uuid(row, "id"),
|
|
ordinal=_integer(row, "ordinal"),
|
|
page_number=_optional_integer(row, "page_number"),
|
|
text=_text(row, "display_text"),
|
|
text_sha256=_hash(row, "text_sha256"),
|
|
line_start=_integer(row, "line_start"),
|
|
line_end=_integer(row, "line_end"),
|
|
)
|
|
|
|
|
|
def _review_block(row: Mapping[str, object]) -> ReviewBlock:
|
|
section_path = row.get("section_path")
|
|
if not isinstance(section_path, list) or not all(
|
|
isinstance(value, str) for value in section_path
|
|
):
|
|
raise DocumentPersistenceError
|
|
return ReviewBlock(
|
|
id=_uuid(row, "id"),
|
|
ordinal=_integer(row, "ordinal"),
|
|
kind=_text(row, "block_kind"),
|
|
text=_text(row, "display_text"),
|
|
text_sha256=_hash(row, "text_sha256"),
|
|
section_path=tuple(section_path),
|
|
anchor_id=_hash(row, "anchor_id"),
|
|
char_start=_integer(row, "char_start"),
|
|
char_end=_integer(row, "char_end"),
|
|
line_start=_integer(row, "line_start"),
|
|
line_end=_integer(row, "line_end"),
|
|
page_start=_optional_integer(row, "page_start"),
|
|
page_end=_optional_integer(row, "page_end"),
|
|
)
|
|
|
|
|
|
def _review_chunk(row: Mapping[str, object]) -> ReviewChunk:
|
|
section_path = row.get("section_path")
|
|
if not isinstance(section_path, list) or not all(
|
|
isinstance(value, str) for value in section_path
|
|
):
|
|
raise DocumentPersistenceError
|
|
return ReviewChunk(
|
|
ordinal=_integer(row, "ordinal"),
|
|
display_text=_text(row, "display_text"),
|
|
cloud_text=_text(row, "cloud_text"),
|
|
cloud_text_sha256=_hash(row, "cloud_text_sha256"),
|
|
embedding_text_sha256=_hash(row, "embedding_text_sha256"),
|
|
token_count=_integer(row, "token_count"),
|
|
page_start=_optional_integer(row, "page_start"),
|
|
page_end=_optional_integer(row, "page_end"),
|
|
section_path=tuple(section_path),
|
|
approval_status=_text(row, "approval_status"),
|
|
index_status=_text(row, "index_status"),
|
|
)
|
|
|
|
|
|
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 DocumentPersistenceError
|
|
|
|
|
|
def _optional_uuid(row: Mapping[str, object], name: str) -> uuid.UUID | None:
|
|
return None if row.get(name) is None else _uuid(row, name)
|
|
|
|
|
|
def _text(row: Mapping[str, object], name: str) -> str:
|
|
value = row.get(name)
|
|
if not isinstance(value, str) or not value.strip():
|
|
raise DocumentPersistenceError
|
|
return value
|
|
|
|
|
|
def _optional_text(row: Mapping[str, object], name: str) -> str | None:
|
|
return None if row.get(name) is None else _text(row, name)
|
|
|
|
|
|
def _hash(row: Mapping[str, object], name: str) -> str:
|
|
value = _text(row, name)
|
|
if _HASH.fullmatch(value) is None:
|
|
raise DocumentPersistenceError
|
|
return value
|
|
|
|
|
|
def _optional_hash(row: Mapping[str, object], name: str) -> str | None:
|
|
return None if row.get(name) is None else _hash(row, name)
|
|
|
|
|
|
def _integer(row: Mapping[str, object], name: str) -> int:
|
|
value = row.get(name)
|
|
if isinstance(value, bool) or not isinstance(value, int):
|
|
raise DocumentPersistenceError
|
|
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 _datetime(row: Mapping[str, object], name: str) -> datetime:
|
|
value = row.get(name)
|
|
if not isinstance(value, datetime):
|
|
raise DocumentPersistenceError
|
|
return value
|
|
|
|
|
|
def _optional_datetime(row: Mapping[str, object], name: str) -> datetime | None:
|
|
return None if row.get(name) is None else _datetime(row, name)
|