commit 596f547086a2378e02d56cb0ac9a866005effdf1 Author: PotapovaA Date: Tue Oct 6 15:01:30 2026 +0300 добавления сервиса PDFWorker для обработки файлов через NATS, MinIO и SMB diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bdeb8a5 --- /dev/null +++ b/.gitignore @@ -0,0 +1,10 @@ +.env +.venv/ +__pycache__/ +.pytest_cache/ +*.pyc +.DS_Store + + +.local/ +graphify-out diff --git a/README.md b/README.md new file mode 100644 index 0000000..b48bd21 --- /dev/null +++ b/README.md @@ -0,0 +1,92 @@ +# Контракты FileWorker + +## NATS: ARMKA → FileWorker + +Subject: `files.tabular_uploaded`. Формат: JSON в UTF-8. + +```json +{ + "task_id": "c609e0d2-f5c4-4a78-b6dc-686a280a2c21", + "file_id": "5e92bc4f-d919-487c-aed4-c293de5c8c96", + "source_s3_key": "uploads/report.pdf", + "expires_at": "2026-10-06T15:00:00Z" +} +``` + +Все поля обязательны. `task_id` — UUID задачи (`RecognitionJob.id`), +`file_id` — UUID исходного файла. `source_s3_key` — ключ объекта внутри +настроенного бакета MinIO, с расширением, без URL и начального `/`. +`expires_at` — срок задачи в ISO 8601 с часовым поясом. +Все расширения используют этот subject; `document_type` не передаётся. + +## NATS: FileWorker → ARMKA + +Subject: `files.tabular_parsed`. Формат: JSON в UTF-8. + +```json +{ + "task_id": "c609e0d2-f5c4-4a78-b6dc-686a280a2c21", + "file_id": "5e92bc4f-d919-487c-aed4-c293de5c8c96", + "status": "ready", + "result_s3_key": "FileWorker/results/c609e0d2-f5c4-4a78-b6dc-686a280a2c21.json", + "error_code": null, + "finished_at": "2026-10-06T14:30:00Z" +} +``` + +`task_id` и `file_id` совпадают с командой. Они, `status` и `finished_at` +обязательны. `finished_at` — ISO 8601 с часовым поясом. + +| status | result_s3_key | error_code | +|---|---|---| +| `ready` | Ключ результата в MinIO | `null` или отсутствует | +| `timeout` | `null` или отсутствует | `deadline_exceeded` | +| `error` | `null` или отсутствует | `source_not_found`, `invalid_source_json` или `storage_unavailable` | + +`storage_unavailable` допустим по схеме, но сейчас ошибки хранилища повторяются +внутри FileWorker и не публикуются с этим кодом. +В обоих сообщениях дополнительные поля запрещены, ключи MinIO имеют длину 1–1024 символа. +Core NATS не сохраняет сообщения для отсутствующих подписчиков; +`flush()` не подтверждает обработку сообщения ARMKA. + +## Файлы + +Все SMB-пути ниже относительны к `SMB_PATH`. + +| Направление | Путь и формат | +|---|---| +| FileWorker → распознаватель | `Input/.<расширение>` — исходный файл | +| Распознаватель → FileWorker | `Output//*.json` — JSON-результаты | +| FileWorker → MinIO | `/.json`, по умолчанию `FileWorker/results/.json` | + +Распознаватель игнорирует `.tmp`. JSON-результаты одной задачи появляются +полным батчем и больше не изменяются. Пустая или отсутствующая папка означает ожидание. +В итоговом JSON каждый исходный JSON становится элементом `files`, +в порядке имён без учёта регистра: + +```json +{ + "task_id": "c609e0d2-f5c4-4a78-b6dc-686a280a2c21", + "job_id": "c609e0d2-f5c4-4a78-b6dc-686a280a2c21", + "file_id": "5e92bc4f-d919-487c-aed4-c293de5c8c96", + "files": [{"example": "содержимое JSON распознавателя"}] +} +``` + +## HTTP + +Префикс `ROOT_PATH` по умолчанию: `/aurora/apps/fastapi-dkb-armka-pdf-worker`. + +| Метод и путь | Ответ | +|---|---| +| `GET /healthcheck` | `200 {"status":"ok"}`; `503`, если NATS или обработчик недоступен | +| `GET /tasks/active` | `200 {"tasks":[...]}`; `503`, если экземпляр FileService недоступен | + +Элемент `tasks`: `task_id`, `file_id`, `source_s3_key`, `status`. +Статусы: `received`, `checking_result`, `checking_smb`, `reading_source`, +`sending_to_smb`, `discovering_result`, `waiting_for_result`, `merging_result`, +`reading_result`, `uploading_result`, `cleaning_smb`, `publishing_result`, +`retrying_storage`, `retrying_publish`. +При сбое добавляется `last_error`: `stage`, `type`, `message`, `location`, `occurred_at`. +Последняя ошибка сохраняется между попытками. Список относится к текущему процессу, +очищается при перезапуске; завершённые задачи из него удаляются. diff --git a/api/requirements-dev.txt b/api/requirements-dev.txt new file mode 100644 index 0000000..d83b683 --- /dev/null +++ b/api/requirements-dev.txt @@ -0,0 +1,4 @@ +-r requirements.txt +pytest==9.1.1 + +pytest-asyncio==1.4.0 diff --git a/api/requirements.txt b/api/requirements.txt new file mode 100644 index 0000000..a3afb8b --- /dev/null +++ b/api/requirements.txt @@ -0,0 +1,7 @@ +nats-py[nkeys]==2.14.0 +fastapi==0.115.7 +uvicorn==0.34.0 +smbprotocol==1.15.0 +pydantic==2.13.4 +pydantic-settings==2.14.2 +miniopy-async==1.23.4 diff --git a/api/run_server.py b/api/run_server.py new file mode 100644 index 0000000..be4e40f --- /dev/null +++ b/api/run_server.py @@ -0,0 +1,14 @@ +"""Run the pdfworker FastAPI application with its JetStream lifespan.""" +import uvicorn + +from src.core.config import get_settings + + +def main(): + settings = get_settings() + uvicorn.run("src.app:app", host=settings.server.HOST, + port=settings.server.PORT, lifespan="on") + + +if __name__ == "__main__": + main() diff --git a/api/run_worker.py b/api/run_worker.py new file mode 100644 index 0000000..6b4080e --- /dev/null +++ b/api/run_worker.py @@ -0,0 +1,4 @@ +from run_server import main + +if __name__ == "__main__": + main() diff --git a/api/src/__init__.py b/api/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/api/src/app.py b/api/src/app.py new file mode 100644 index 0000000..49e7422 --- /dev/null +++ b/api/src/app.py @@ -0,0 +1,78 @@ +"""HTTP application and lifetime of the pdfworker NATS subscriber.""" +import asyncio +from contextlib import asynccontextmanager + +import nats +from fastapi import FastAPI, HTTPException + +from src.core.clients import SMBStorage +from src.core.config import get_settings +from src.core.minio_client import MinioStorage +from src.core.nats_auth import nats_auth +from src.core.nats_worker import consume +from src.services.file_service import FileService + +@asynccontextmanager +async def lifespan(app: FastAPI): + settings = get_settings() + app.root_path = settings.ROOT_PATH + nats_settings = settings.OPENBAO.NATS + storage = SMBStorage(settings) + objects = None + nc = None + worker = None + stop = asyncio.Event() + try: + objects = MinioStorage(settings) + await asyncio.to_thread(storage.connect) + nc = await nats.connect(nats_settings.NATS_URL, **nats_auth(nats_settings), + max_reconnect_attempts=-1, drain_timeout=60) + subscription = await nc.subscribe(nats_settings.NATS_COMMAND_SUBJECT) + await nc.flush() + service = FileService(storage, objects) + worker = asyncio.create_task( + consume(subscription, nc, service, settings, stop), + name="pdfworker-consumer", + ) + app.state.nats = nc + app.state.worker = worker + app.state.service = service + yield + finally: + stop.set() + try: + if worker is not None: + await worker + finally: + try: + if nc is not None and not nc.is_closed: + await nc.close() + finally: + try: + await asyncio.to_thread(storage.close) + finally: + if objects is not None: + await objects.close() + app.state.nats = None + app.state.worker = None + app.state.service = None + + +app = FastAPI(lifespan=lifespan) + + +@app.get("/healthcheck") +async def healthcheck(): + nc = getattr(app.state, "nats", None) + worker = getattr(app.state, "worker", None) + if nc is None or not nc.is_connected or worker is None or worker.done(): + raise HTTPException(status_code=503, detail="pdfworker consumer unavailable") + return {"status": "ok"} + + +@app.get("/tasks/active") +async def active_tasks(): + service = getattr(app.state, "service", None) + if service is None: + raise HTTPException(status_code=503, detail="pdfworker consumer unavailable") + return {"tasks": list(service.active.values())} diff --git a/api/src/core/__init__.py b/api/src/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/api/src/core/clients.py b/api/src/core/clients.py new file mode 100644 index 0000000..4518b38 --- /dev/null +++ b/api/src/core/clients.py @@ -0,0 +1,226 @@ +import asyncio +import errno +import json +import ntpath +import stat +from contextlib import asynccontextmanager +from uuid import uuid4 + +import smbclient +from smbclient._io import SMBFileTransaction, set_info +from smbprotocol.header import NtStatus +from smbprotocol.exceptions import SMBOSError +from smbprotocol.file_info import FileRenameInformation +from smbprotocol.open import FilePipePrinterAccessMask + +from src.core.config import Settings +from src.schemas.tasks import PrepareTask, relative_path + + +class SMBStorage: + def __init__(self, settings: Settings): + self.max_source_bytes = settings.MAX_SOURCE_BYTES + self.settings = settings.OPENBAO.SMB + settings = self.settings + self._connection_cache = {} + self.dispatch_root = self._root(settings.SMB_DISPATCH_ROOT) + self.input_root = self._root(settings.SMB_INPUT_ROOT) + self.output_root = self._root(settings.SMB_OUTPUT_ROOT) + for root_path in (self.input_root, self.output_root): + a, b = self.dispatch_root.casefold(), root_path.casefold() + if a == b or a.startswith(b + "\\") or b.startswith(a + "\\"): + raise ValueError("Dispatch, results and staging roots must be separate") + if (self.input_root.casefold() == self.output_root.casefold() or + self.output_root.casefold().startswith(self.input_root.casefold() + "\\") or + self.input_root.casefold().startswith(self.output_root.casefold() + "\\")): + raise ValueError("Input and output roots must be separate") + + def _remove_file(self, path, task): + try: + smbclient.remove(path, **self._connection_options()) + except FileNotFoundError: + pass + except (SMBOSError, OSError, ValueError): + pass + + def _root(self, path: str) -> str: + path = path.replace("/", "\\").rstrip("\\") + drive, tail = ntpath.splitdrive(path) + if (not drive.startswith("\\\\") or len(drive.split("\\")) != 4 or + drive.split("\\")[2].casefold() != self.settings.SMB_HOST.casefold()): + raise ValueError("SMB root must be a UNC path on SMB_HOST") + if any(part in (".", "..") for part in tail.split("\\")): + raise ValueError("SMB root must not contain traversal") + return path + + def _connection_options(self) -> dict: + return { + "username": self.settings.SMB_LOGIN, + "password": self.settings.SMB_PASSWORD.get_secret_value(), + "port": self.settings.SMB_PORT, + "connection_timeout": self.settings.SMB_TIMEOUT_SECONDS, + "connection_cache": self._connection_cache, + } + + def connect(self): + smbclient.register_session(self.settings.SMB_HOST, **self._connection_options()) + for path in (self.output_root, self.dispatch_root): + smbclient.makedirs(path, exist_ok=True, **self._connection_options()) + + def close(self): + smbclient.reset_connection_cache(connection_cache=self._connection_cache) + + def source_path(self, path: str) -> str: + return ntpath.join(self.input_root, relative_path(path).replace("/", "\\")) + + def result_path(self, task: PrepareTask) -> str: + return ntpath.join(self.output_root, f"{task.task_id}.json") + + def exists(self, path: str) -> bool: + try: + return stat.S_ISREG(smbclient.stat(path, **self._connection_options()).st_mode) + except FileNotFoundError: + return False + except SMBOSError as exc: + if exc.errno == errno.ENOENT or exc.ntstatus == NtStatus.STATUS_NO_SUCH_FILE: + return False + raise + + def discover_results(self, task: PrepareTask) -> list[str]: + directory = self.source_path(str(task.task_id)) + try: + names = smbclient.listdir(directory, **self._connection_options()) + except FileNotFoundError: + return [] + except SMBOSError as exc: + if exc.errno == errno.ENOENT or exc.ntstatus == NtStatus.STATUS_NO_SUCH_FILE: + return [] + raise + paths = [] + for name in names: + if not name.lower().endswith(".json"): + continue + if "/" in name or "\\" in name: + continue + relative = f"{task.task_id}/{name}" + if self.exists(self.source_path(relative)): + paths.append(relative) + return paths + + def _publish(self, temporary: str, destination: str): + with smbclient.open_file( + temporary, mode="rb", buffering=0, share_access="rwd", + desired_access=FilePipePrinterAccessMask.DELETE, + **self._connection_options(), + ) as source: + info = FileRenameInformation() + info["replace_if_exists"] = False + info["file_name"] = ntpath.join(ntpath.dirname(source.fd.file_name), + ntpath.basename(destination)) + try: + with SMBFileTransaction(source) as transaction: + set_info(transaction, info) + except FileExistsError: + if not self.exists(destination): + raise + + def merge(self, task: PrepareTask, source_files: list[str]) -> str: + if not source_files: + raise ValueError("No result files to merge") + destination = self.result_path(task) + temporary = destination + f".{uuid4()}.tmp" + metadata = { + "task_id": str(task.task_id), "job_id": str(task.task_id), + "file_id": str(task.file_id), + } + try: + with smbclient.open_file(temporary, mode="x", encoding="utf-8", **self._connection_options()) as output: + output.write(json.dumps(metadata, ensure_ascii=False)[:-1] + ', "files":[') + for index, path in enumerate(source_files): + with smbclient.open_file(self.source_path(path), mode="rb", **self._connection_options()) as source: + raw = source.read(self.max_source_bytes + 1) + if len(raw) > self.max_source_bytes: + raise ValueError("Source exceeds MAX_SOURCE_BYTES") + data = json.loads(raw) + if index: + output.write(",") + json.dump(data, output, ensure_ascii=False, allow_nan=False) + del raw, data + output.write("]}") + self._publish(temporary, destination) + return destination + finally: + self._remove_file(temporary, task) + + def _receipt_path(self, task): + return ntpath.join(self.output_root, f"{task.task_id}.submitted.json") + + def submitted(self, task): + return self.exists(self._receipt_path(task)) + + async def dispatch(self, task, stream): + destination = ntpath.join(self.dispatch_root, task.smb_input_name) + temporary = destination + f".{uuid4()}.tmp" + output = None + try: + output = await asyncio.to_thread(smbclient.open_file, temporary, mode="xb", + **self._connection_options()) + while chunk := await stream.read(1024 * 1024): + await asyncio.to_thread(output.write, chunk) + await asyncio.to_thread(output.close) + output = None + await asyncio.to_thread(self._publish, temporary, destination) + finally: + try: + if output is not None: + await asyncio.to_thread(output.close) + finally: + await asyncio.to_thread(self._remove_file, temporary, task) + def record_receipt(): + receipt = self._receipt_path(task) + temporary = receipt + f".{uuid4()}.tmp" + try: + with smbclient.open_file(temporary, mode="x", encoding="utf-8", **self._connection_options()) as output: + output.write("{}") + self._publish(temporary, receipt) + finally: + self._remove_file(temporary, task) + await asyncio.to_thread(record_receipt) + + def cleanup_result(self, task): + directory = self.source_path(str(task.task_id)) + try: + source_files = self.discover_results(task) + except (OSError, ValueError): + source_files = [] + for path in source_files: + self._remove_file(self.source_path(path), task) + for path in (self.result_path(task), self._receipt_path(task), + ntpath.join(self.dispatch_root, task.smb_input_name)): + self._remove_file(path, task) + try: + smbclient.rmdir(directory, **self._connection_options()) + except FileNotFoundError: + pass + except (OSError, ValueError): + pass + + @asynccontextmanager + async def read_result(self, task): + path = self.result_path(task) + size = (await asyncio.to_thread(smbclient.stat, path, + **self._connection_options())).st_size + source = await asyncio.to_thread(smbclient.open_file, path, mode="rb", + **self._connection_options()) + try: + yield AsyncFileReader(source), size + finally: + await asyncio.to_thread(source.close) + + +class AsyncFileReader: + def __init__(self, source): + self.source = source + + async def read(self, size): + return await asyncio.to_thread(self.source.read, size) diff --git a/api/src/core/config.py b/api/src/core/config.py new file mode 100644 index 0000000..0d0dcc0 --- /dev/null +++ b/api/src/core/config.py @@ -0,0 +1,81 @@ +import ntpath +from pathlib import Path + +from pydantic import BaseModel, Field, SecretStr, field_validator +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class MinioSettings(BaseModel): + MINIO_URL: str + MINIO_ADMIN_LOGIN: str + MINIO_ADMIN_PASSWORD: SecretStr + MINIO_BUCKET: str + MINIO_SECURE: bool = True + MINIO_RESULT_PREFIX: str = "pdfworker/results" + + +class NATSSettings(BaseModel): + NATS_URL: str + NATS_JWT: SecretStr | None = None + NATS_SEED: SecretStr | None = None + NATS_COMMAND_SUBJECT: str = "files.tabular_uploaded" + NATS_EVENT_SUBJECT: str = "files.tabular_parsed" + + +class SMBSettings(BaseModel): + SMB_HOST: str + SMB_PATH: str + SMB_PORT: int = Field(default=445, ge=1, le=65535) + SMB_LOGIN: str + SMB_PASSWORD: SecretStr + SMB_TIMEOUT_SECONDS: int = Field(default=30, ge=1, le=120) + + @field_validator("SMB_PATH") + @classmethod + def normalize_path(cls, value: str) -> str: + path = value.replace("/", "\\") + return "\\" + path if path.startswith("\\") and not path.startswith("\\\\") else path + + @property + def SMB_DISPATCH_ROOT(self) -> str: + return ntpath.join(self.SMB_PATH, "Input") + + @property + def SMB_INPUT_ROOT(self) -> str: + return ntpath.join(self.SMB_PATH, "Output") + + @property + def SMB_OUTPUT_ROOT(self) -> str: + return ntpath.join(self.SMB_PATH, "Merged") + + +class OpenBAOSettings(BaseModel): + MINIO: MinioSettings + NATS: NATSSettings = Field(default_factory=NATSSettings) + SMB: SMBSettings + + +class ServerSettings(BaseSettings): + model_config = SettingsConfigDict(env_prefix="SERVER_") + + HOST: str = "0.0.0.0" + PORT: int = 8000 + + +class Settings(BaseSettings): + model_config = SettingsConfigDict( + env_file=Path(__file__).resolve().parents[2] / ".env", + env_nested_delimiter="__", + extra="ignore", + hide_input_in_errors=True, + ) + + OPENBAO: OpenBAOSettings + ROOT_PATH: str = "/aurora/apps/fastapi-dkb-armka-pdf-worker" + server: ServerSettings = Field(default_factory=ServerSettings) + POLL_INTERVAL_SECONDS: int = Field(default=10, ge=1, le=300) + MAX_SOURCE_BYTES: int = Field(default=100_000_000, ge=1) + + +def get_settings() -> Settings: + return Settings() diff --git a/api/src/core/errors.py b/api/src/core/errors.py new file mode 100644 index 0000000..cd10fa9 --- /dev/null +++ b/api/src/core/errors.py @@ -0,0 +1,2 @@ +class SourceNotFound(FileNotFoundError): + """The MinIO source does not exist.""" diff --git a/api/src/core/minio_client.py b/api/src/core/minio_client.py new file mode 100644 index 0000000..01b1788 --- /dev/null +++ b/api/src/core/minio_client.py @@ -0,0 +1,53 @@ +import asyncio +from contextlib import asynccontextmanager + +from aiohttp import ClientError +from miniopy_async import Minio +from miniopy_async.error import MinioException, S3Error + +from src.core.errors import SourceNotFound + + +class MinioStorage: + def __init__(self, settings): + settings = settings.OPENBAO.MINIO + self.bucket = settings.MINIO_BUCKET + self.prefix = settings.MINIO_RESULT_PREFIX.strip("/") + self.client = Minio(settings.MINIO_URL, access_key=settings.MINIO_ADMIN_LOGIN, + secret_key=settings.MINIO_ADMIN_PASSWORD.get_secret_value(), + secure=settings.MINIO_SECURE) + + async def close(self): + await self.client.close_session() + + def result_key(self, task): + return f"{self.prefix}/{task.task_id}.json" + + async def result_exists(self, task): + try: + await self.client.stat_object(self.bucket, self.result_key(task)) + except S3Error as exc: + if exc.code in {"NoSuchKey", "NoSuchObject"}: + return False + raise + return True + + @asynccontextmanager + async def read_source(self, key): + try: + response = await self.client.get_object(self.bucket, key) + except S3Error as exc: + if exc.code in {"NoSuchKey", "NoSuchObject"}: + raise SourceNotFound(key) from exc + raise + try: + yield response.content + finally: + response.release() + + async def upload_result(self, task, stream, size): + await self.client.put_object(self.bucket, self.result_key(task), stream, length=size, + content_type="application/json") + + +MINIO_ERRORS = (MinioException, ClientError, asyncio.TimeoutError) diff --git a/api/src/core/nats_auth.py b/api/src/core/nats_auth.py new file mode 100644 index 0000000..4f652d1 --- /dev/null +++ b/api/src/core/nats_auth.py @@ -0,0 +1,28 @@ +"""NATS JWT authentication using credentials supplied through environment variables.""" +import base64 + +import nkeys + +from src.core.config import NATSSettings + + +def nats_auth(settings: NATSSettings) -> dict: + jwt = settings.NATS_JWT.get_secret_value() if settings.NATS_JWT else "" + seed = settings.NATS_SEED.get_secret_value() if settings.NATS_SEED else "" + if bool(jwt) != bool(seed): + raise ValueError("NATS_JWT and NATS_SEED must be provided together") + if not jwt: + return {} + + def jwt_cb(): + return jwt.encode() + + def signature_cb(nonce): + key = nkeys.from_seed(seed.encode()) + try: + payload = nonce.encode() if isinstance(nonce, str) else nonce + return base64.b64encode(key.sign(payload)) + finally: + key.wipe() + + return {"user_jwt_cb": jwt_cb, "signature_cb": signature_cb} diff --git a/api/src/core/nats_worker.py b/api/src/core/nats_worker.py new file mode 100644 index 0000000..79de5b7 --- /dev/null +++ b/api/src/core/nats_worker.py @@ -0,0 +1,54 @@ +import asyncio + +from nats.errors import TimeoutError as NATSTimeoutError +from pydantic import ValidationError + +from smbprotocol.exceptions import SMBException +from src.core.minio_client import MINIO_ERRORS +from src.schemas.tasks import PrepareTask + +async def process_message(msg, nc, service, settings): + try: + task = PrepareTask.model_validate_json(msg.data) + except ValidationError: + return + service.set_status(task, "received") + try: + while True: + try: + result = await service.process(task) + except (OSError, SMBException, *MINIO_ERRORS) as exc: + service.set_status(task, "retrying_storage", error=exc) + result = None + if result is None: + await asyncio.sleep(settings.POLL_INTERVAL_SECONDS) + continue + try: + service.set_status(task, "publishing_result") + await nc.publish(settings.OPENBAO.NATS.NATS_EVENT_SUBJECT, result.model_dump_json().encode()) + await nc.flush() + except Exception as exc: + service.set_status(task, "retrying_publish", error=exc) + await asyncio.sleep(settings.POLL_INTERVAL_SECONDS) + continue + service.acknowledged(task) + return + finally: + service.forget() + + +async def consume(subscription, nc, service, settings, stop): + pending = set() + try: + while not stop.is_set(): + try: + msg = await subscription.next_msg(timeout=1) + except NATSTimeoutError: + continue + task = asyncio.create_task(process_message(msg, nc, service, settings)) + pending.add(task) + task.add_done_callback(pending.discard) + finally: + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) diff --git a/api/src/schemas/__init__.py b/api/src/schemas/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/api/src/schemas/messages.py b/api/src/schemas/messages.py new file mode 100644 index 0000000..e593de5 --- /dev/null +++ b/api/src/schemas/messages.py @@ -0,0 +1,60 @@ +from pathlib import PurePosixPath +from typing import Literal +from uuid import UUID + +from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, field_validator, model_validator + +COMMAND_SUBJECT = "files.tabular_uploaded" +RESULT_SUBJECT = "files.tabular_parsed" +ErrorCode = Literal["deadline_exceeded", "source_not_found", "invalid_source_json", + "storage_unavailable"] + + +def validate_object_key(key: str) -> str: + if not key.strip() or key.startswith(("/", "http://", "https://", "s3://")): + raise ValueError("Expected the full object key inside the configured MinIO bucket, not a URL") + return key + + +class PDFWorkerCommand(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + task_id: UUID + file_id: UUID + source_s3_key: str = Field(min_length=1, max_length=1024) + expires_at: AwareDatetime + + @field_validator("source_s3_key") + @classmethod + def validate_source(cls, value): + validate_object_key(value) + suffix = PurePosixPath(value).suffix[1:] + if (not suffix or suffix.endswith(' ') or + any(ord(char) < 32 or char in '\\/:*?"<>|' for char in suffix)): + raise ValueError("Expected a safe source extension") + return value + + +class PDFWorkerResult(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + task_id: UUID + file_id: UUID + status: Literal["ready", "timeout", "error"] + result_s3_key: str | None = Field(default=None, min_length=1, max_length=1024) + error_code: ErrorCode | None = None + finished_at: AwareDatetime + + @field_validator("result_s3_key") + @classmethod + def validate_result_key(cls, value): + return validate_object_key(value) if value is not None else value + + @model_validator(mode="after") + def validate_status(self): + if self.status == "ready": + if self.result_s3_key is None or self.error_code is not None: + raise ValueError("ready requires result_s3_key and no error_code") + elif self.result_s3_key is not None or self.error_code is None: + raise ValueError("Failure requires error_code and no result_s3_key") + elif (self.status == "timeout") != (self.error_code == "deadline_exceeded"): + raise ValueError("deadline_exceeded is only valid with timeout") + return self diff --git a/api/src/schemas/tasks.py b/api/src/schemas/tasks.py new file mode 100644 index 0000000..2d0bfa0 --- /dev/null +++ b/api/src/schemas/tasks.py @@ -0,0 +1,24 @@ +from pathlib import PurePosixPath + +from src.schemas.messages import PDFWorkerCommand, PDFWorkerResult + + +def relative_path(value: str) -> str: + value = value.replace("\\", "/") + if (not value or value.startswith("/") or + any(part in ("", ".", "..") for part in value.split("/")) or + any(char in value for char in ':*?"<>|\x00')): + raise ValueError("Expected a relative SMB path without traversal or wildcards") + if any(part.endswith((".", " ")) for part in value.split("/")): + raise ValueError("Trailing dots/spaces are not allowed") + return value + + +class PrepareTask(PDFWorkerCommand): + @property + def smb_input_name(self) -> str: + return f"{self.task_id}{PurePosixPath(self.source_s3_key).suffix.lower()}" + + +class TaskResult(PDFWorkerResult): + pass diff --git a/api/src/services/__init__.py b/api/src/services/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/api/src/services/file_service.py b/api/src/services/file_service.py new file mode 100644 index 0000000..39b39a6 --- /dev/null +++ b/api/src/services/file_service.py @@ -0,0 +1,86 @@ +import asyncio +import traceback +from datetime import datetime, timezone + +from src.core.errors import SourceNotFound +from src.schemas.tasks import PrepareTask, TaskResult + + +class FileService: + def __init__(self, storage, objects): + self.storage = storage + self.objects = objects + self._ready = set() + self.active = {} + + def set_status(self, task, status, error=None): + previous = self.active.get(asyncio.current_task(), {}) + entry = { + **previous, + "task_id": str(task.task_id), "file_id": str(task.file_id), + "source_s3_key": task.source_s3_key, "status": status, + } + if error is not None: + frames = traceback.extract_tb(error.__traceback__) + frame = next((frame for frame in reversed(frames) + if "/src/" in frame.filename), None) + entry["last_error"] = { + "stage": previous.get("status"), + "type": type(error).__name__, + "message": str(error)[:2000], + "location": (f"{frame.filename.split('/src/')[-1]}:{frame.lineno} " + f"({frame.name})") if frame else None, + "occurred_at": datetime.now(timezone.utc).isoformat(), + } + self.active[asyncio.current_task()] = entry + + def forget(self): + self.active.pop(asyncio.current_task(), None) + + def acknowledged(self, task): + self._ready.discard(task.task_id) + + async def process(self, task: PrepareTask) -> TaskResult | None: + def result(status, **kwargs): + return TaskResult(task_id=task.task_id, file_id=task.file_id, + status=status, finished_at=datetime.now(timezone.utc), **kwargs) + + try: + self.set_status(task, "checking_result") + if task.task_id in self._ready: + return result("ready", result_s3_key=self.objects.result_key(task)) + if await self.objects.result_exists(task): + self._ready.add(task.task_id) + self.set_status(task, "cleaning_smb") + await asyncio.to_thread(self.storage.cleanup_result, task) + return result("ready", result_s3_key=self.objects.result_key(task)) + self.set_status(task, "checking_smb") + submitted = await asyncio.to_thread(self.storage.submitted, task) + if not (submitted and await asyncio.to_thread(self.storage.exists, self.storage.result_path(task))): + if datetime.now(timezone.utc) >= task.expires_at: + return result("timeout", error_code="deadline_exceeded") + if not submitted: + self.set_status(task, "reading_source") + async with self.objects.read_source(task.source_s3_key) as source: + self.set_status(task, "sending_to_smb") + await self.storage.dispatch(task, source) + self.set_status(task, "discovering_result") + source_files = await asyncio.to_thread(self.storage.discover_results, task) + if not source_files: + self.set_status(task, "waiting_for_result") + return None + try: + self.set_status(task, "merging_result") + await asyncio.to_thread(self.storage.merge, task, source_files) + except (ValueError, UnicodeError): + return result("error", error_code="invalid_source_json") + self.set_status(task, "reading_result") + async with self.storage.read_result(task) as (stream, size): + self.set_status(task, "uploading_result") + await self.objects.upload_result(task, stream, size) + self._ready.add(task.task_id) + self.set_status(task, "cleaning_smb") + await asyncio.to_thread(self.storage.cleanup_result, task) + return result("ready", result_s3_key=self.objects.result_key(task)) + except SourceNotFound: + return result("error", error_code="source_not_found")