From c024976bf77ddd573942fb28c933d6a9d77fb3c6 Mon Sep 17 00:00:00 2001 From: PotapovaA Date: Tue, 6 Oct 2026 13:13:36 +0300 Subject: [PATCH] =?UTF-8?q?=D0=B8=D0=B7=D0=BC=D0=B5=D0=BD=D0=B5=D0=BD?= =?UTF-8?q?=D0=B8=D1=8F=20=D0=B4=D0=BB=D1=8F=20=D0=BD=D0=BE=D0=B2=D0=BE?= =?UTF-8?q?=D0=B3=D0=BE=20=D1=81=D0=B5=D1=80=D0=B2=D0=B8=D1=81=D0=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api/src/__init__.py | 18 +- api/src/api/endpoints/files.py | 12 +- api/src/app.py | 138 +++-------- api/src/core/PDFWorker/PDFTempPoller.py | 168 ------------- api/src/core/async_cache.py | 52 ---- api/src/core/clients.py | 196 +++++++-------- api/src/core/config.py | 8 +- api/src/core/pdfworker_nats.py | 80 +++++++ api/src/core/pdfworker_tasks.py | 12 + api/src/schemas/enums.py | 5 +- api/src/schemas/pdfworker.py | 61 +++++ api/src/services/file_service.py | 29 ++- .../ofr_result_conversion}/__init__.py | 0 .../ofr_result_conversion}/consts.py | 0 .../ofr_result_conversion}/utils.py | 0 api/src/services/pdfworker_results.py | 79 +++++++ api/tests/unit/test_pdfworker_contract.py | 156 ++++++++++++ api/tests/unit/test_pdfworker_persistence.py | 223 ++++++++++++++++++ 18 files changed, 776 insertions(+), 461 deletions(-) delete mode 100644 api/src/core/PDFWorker/PDFTempPoller.py delete mode 100644 api/src/core/async_cache.py create mode 100644 api/src/core/pdfworker_nats.py create mode 100644 api/src/core/pdfworker_tasks.py create mode 100644 api/src/schemas/pdfworker.py rename api/src/{core/PDFWorker => services/ofr_result_conversion}/__init__.py (100%) rename api/src/{core/PDFWorker => services/ofr_result_conversion}/consts.py (100%) rename api/src/{core/PDFWorker => services/ofr_result_conversion}/utils.py (100%) create mode 100644 api/src/services/pdfworker_results.py create mode 100644 api/tests/unit/test_pdfworker_contract.py create mode 100644 api/tests/unit/test_pdfworker_persistence.py diff --git a/api/src/__init__.py b/api/src/__init__.py index 61c3faa..0ddffd0 100644 --- a/api/src/__init__.py +++ b/api/src/__init__.py @@ -1,7 +1,13 @@ -from .app import app -from .core.config import get_settings +"""Importing schemas or pure functions must not initialize the application.""" +__all__ = ["app", "get_settings"] -__all__ = [ - 'app', - 'get_settings', -] \ No newline at end of file + +def __getattr__(name): + if name == "app": + from .app import app + globals()[name] = app + return app + if name == "get_settings": + from .core.config import get_settings + return get_settings + raise AttributeError(name) diff --git a/api/src/api/endpoints/files.py b/api/src/api/endpoints/files.py index 17ee5a0..7810746 100644 --- a/api/src/api/endpoints/files.py +++ b/api/src/api/endpoints/files.py @@ -2,7 +2,9 @@ from fastapi import APIRouter, File, HTTPException, UploadFile, Form, Depends from fastapi.responses import FileResponse, StreamingResponse from src.api.dependencies.services import get_file_service, FileService from src.schemas.enums import PeriodQuarter +import mimetypes import uuid +from urllib.parse import quote router = APIRouter(prefix='/files', tags=['files']) @@ -25,7 +27,9 @@ async def get_files_by_legal_entity(legal_entity_id: uuid.UUID, service: FileSer @router.get('/{file_id}') async def get_file(file_id: uuid.UUID, service: FileService = Depends(get_file_service)): - file_bytes = await service.get_file_bytes(file_id) - return StreamingResponse(file_bytes.content, media_type="application/pdf", headers={ - "Content-Disposition": f'inline; filename="file.pdf"' - },) \ No newline at end of file + file_response, file_type = await service.get_file_bytes(file_id) + media_type = mimetypes.guess_type(f"file.{file_type.lower()}")[0] or "application/octet-stream" + disposition = "inline" if file_type.upper() == "PDF" else "attachment" + return StreamingResponse(file_response.content, media_type=media_type, headers={ + "Content-Disposition": f"{disposition}; filename*=UTF-8''{quote(f'{file_id}.{file_type.lower()}')}" + },) diff --git a/api/src/app.py b/api/src/app.py index 4f1edbb..2020f27 100644 --- a/api/src/app.py +++ b/api/src/app.py @@ -1,13 +1,11 @@ import os -import smbclient from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles from datetime import datetime, timezone -from src.core.PDFWorker.PDFTempPoller import PDFPollWorker from .core.config import get_settings # from .api.router import files_router @@ -15,17 +13,15 @@ from contextlib import asynccontextmanager from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware -from src.core.nats_callback import build_file_handler from fastapi.exceptions import HTTPException, StarletteHTTPException from src.api.router import le_router, files_router, cg_router, rf_office_router, osv_router, fsl_router from miniopy_async import Minio from src.core.config import settings from src.core.clients import S3Storage, PollWorker -from src.core.async_cache import AsyncInMemoryCache -# from src.core.logger import setup_logger -from nats.aio.client import Client as NATS -import nkeys -import base64 +from src.core.logger import setup_logger +from src.core.pdfworker_nats import pdfworker_connection +from src.services.pdfworker_results import apply_pdfworker_result +from src.db.db import async_session_maker from loguru import logger import os @@ -36,30 +32,13 @@ from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware from src.core.exception_handlers import register_exception_handlers +setup_logger(intercept_std_logging=False) settings = get_settings() VERSION = "1.0.0" UP_TIME = datetime.now(timezone.utc) -kp = nkeys.from_seed(settings.OPENBAO.NATS.NATS_SEED.encode('utf-8')) - - -def jwt_cb(): - return settings.OPENBAO.NATS.NATS_JWT.encode('utf-8') - - -def sig_cb(nonce): - if (isinstance(nonce, str)): - nonce = nonce.encode('utf-8') - signed = kp.sign(nonce) - return base64.b64encode(signed) - - -async def sub_cb(msg): - return None - - protected_middleware_available = False if not settings.DEBUG: try: @@ -91,83 +70,42 @@ _APP_ROOT_PATH = _normalize_root_path(settings.ROOT_PATH) @asynccontextmanager async def lifespan(app): - app.state.nats_subs = [] - app.state.storage = None app.state.nats = None + app.state.poller = None + app.state.storage = None + s3_client = Minio(settings.OPENBAO.MINIO.MINIO_URL, + access_key=settings.OPENBAO.MINIO.MINIO_ADMIN_LOGIN, + secret_key=settings.OPENBAO.MINIO.MINIO_ADMIN_PASSWORD, + secure=settings.OPENBAO.MINIO.MINIO_SECURE, cert_check=False) try: - smbclient.register_session(settings.OPENBAO.SMB.SMB_HOST, settings.OPENBAO.SMB.SMB_LOGIN, settings.OPENBAO.SMB.SMB_PASSWORD) - except: - logger.exception("Ошибка при регистрации сессии SMB") - s3_client = Minio(settings.OPENBAO.MINIO.MINIO_URL, access_key=settings.OPENBAO.MINIO.MINIO_ADMIN_LOGIN, secret_key=settings.OPENBAO.MINIO.MINIO_ADMIN_PASSWORD, cert_check=False) - try: - bucket_exists = await s3_client.bucket_exists(settings.OPENBAO.MINIO.MINIO_BUCKET) - logger.info("Бакет проверен. bucket_exists={bucket_exists}", bucket_exists=bucket_exists) - except Exception as exc: - logger.exception( - "Не удалось проверить наличие бакетa. error_category=s3_error, error_type={error_type}, error={error}", - error_type=type(exc).__name__, - error=repr(exc), - ) - return - storage = S3Storage(client=s3_client, bucket_name=settings.OPENBAO.MINIO.MINIO_BUCKET) - app.state.storage = storage - cache = AsyncInMemoryCache() - nc = NATS() - if not settings.DEBUG: - try: - await nc.connect(settings.OPENBAO.NATS.NATS_URL, user_jwt_cb=jwt_cb, signature_cb=sig_cb, - max_reconnect_attempts=2) - logger.info("Подключение к NATS произошло") - except Exception as exc: - logger.exception( - "Не удалось проверить наличие бакетa. error_category=nats_error, error_type={error_type}, error={error}", - error_type=type(exc).__name__, - error=repr(exc), - ) - return - poller = PollWorker(nc, cache, storage) - await poller.start() - pdf_poller = PDFPollWorker(cache, storage) - await pdf_poller.start() - app.state.nats = nc - app.state.poller = poller - app.state.pdf_poller = pdf_poller - if not settings.DEBUG: - handler_cb = build_file_handler(nc, storage) + if not await s3_client.bucket_exists(settings.OPENBAO.MINIO.MINIO_BUCKET): + raise RuntimeError("Configured MinIO bucket does not exist") + storage = S3Storage(s3_client, settings.OPENBAO.MINIO.MINIO_BUCKET) + app.state.storage = storage - try: - # sub_xls = await nc.subscribe('files.tabular_parsed', cb=handler_cb) - # sub_img = await nc.subscribe('files.img_parsed', cb=handler_cb) - # - # app.state.nats_subs = [sub_xls, sub_img] - logger.info("Подписка на топики завершена") - except Exception as exc: - logger.warning( - "Не удалось подписаться на тему. error_category=nats_error, error_type={error_type}, error={error}", - error_type=type(exc).__name__, - error=repr(exc), - ) - try: - logger.info("Приложение запущено") - yield - except Exception as exc: - logger.exception("ОШИБКА ПРИ РАБОТЕ С СЕКРЕТАМИ error={error}", error=repr(exc)) + async def on_result(event): + await apply_pdfworker_result(event, async_session_maker, storage) + + # NATS is required for PDFWorker in local DEBUG mode as well. + nats_settings = settings.OPENBAO.NATS + async with pdfworker_connection(nats_settings.NATS_URL, on_result, + jwt=nats_settings.NATS_JWT, + seed=nats_settings.NATS_SEED) as nc: + app.state.nats = nc + poller = PollWorker(nc) + app.state.poller = poller + try: + await poller.start() + logger.info("ARMKA PDFWorker publisher and result consumer started") + yield + finally: + await poller.stop() finally: - logger.info("Приложение закрывается") - smbclient.reset_connection_cache() - if not settings.DEBUG: - for sub in app.state.nats_subs: - logger.info("Отписываемся от топика") - await sub.unsubscribe() - if app.state.nats: - await app.state.nats.close() - if app.state.poller: - await app.state.poller.stop() - if app.state.storage: - await app.state.storage.client.close_session() - if app.state.pdf_poller: - await app.state.pdf_poller.stop() - logger.info("CLOSING APP AT ALL") + await s3_client.close_session() + app.state.nats = None + app.state.poller = None + app.state.storage = None + app = FastAPI( title="Fast API", @@ -247,4 +185,4 @@ class SPAStaticFiles(StaticFiles): app.mount("/", SPAStaticFiles(directory="web", html=True), name="web") -app.mount("/", StaticFiles(directory="web", html=True), name="web") \ No newline at end of file +app.mount("/", StaticFiles(directory="web", html=True), name="web") diff --git a/api/src/core/PDFWorker/PDFTempPoller.py b/api/src/core/PDFWorker/PDFTempPoller.py deleted file mode 100644 index 4ad0738..0000000 --- a/api/src/core/PDFWorker/PDFTempPoller.py +++ /dev/null @@ -1,168 +0,0 @@ -import asyncio -from datetime import datetime, timedelta -import json - -from alembic.util import status -from loguru import logger -from sqlalchemy import select, update, insert -import smbclient - -from src.core.PDFWorker.utils import json_to_db_rows -from src.core.async_cache import AsyncInMemoryCache, FileDictItem -from src.core.clients import S3Storage -from src.db.db import async_session_maker -from src.models import RecognitionJob, StoredFile, FinancialStatementMeta, FinancialStatementLine, FinancialStatementValue -from src.schemas.enums import FileStatus, DocumentType, RecognitionStatus -from src.core.config import settings - - -class PDFPollWorker: - def __init__(self, cache: AsyncInMemoryCache, s3: S3Storage): - self.task: asyncio.Task | None = None - self.stop_event = asyncio.Event() - self.cache = cache - self.interval = 10 - self.storage = s3 - self.smb_path = settings.OPENBAO.SMB.SMB_PATH - - async def start(self): - self.stop_event.clear() - self.task = asyncio.create_task(self.run()) - - async def stop(self): - self.stop_event.set() - if self.task: - self.task.cancel() - try: - await self.task - except asyncio.CancelledError: - pass - - async def run(self): - logger.info("Начат процесс опроса SMB") - while not self.stop_event.is_set(): - try: - await self.poll_once() - except Exception: - # logger.exception("Poll iteration failed") - pass - - try: - await asyncio.wait_for( - self.stop_event.wait(), - timeout=self.interval, - ) - except asyncio.TimeoutError: - pass - - logger.info("PDFPollWorker остановлен") - - async def poll_once(self): - try: - jobs: list[FileDictItem] = await self.cache.get_active_items() - logger.info("Получил jobs") - except Exception as e: - logger.exception("Не получилось достать данные из кэша") - jobs = [] - - all_files = smbclient.listdir(f'{self.smb_path}/Output') - - for job in jobs: - search_fname = f"{job.file_id}" - now = datetime.now() - try: - if (job.expires_at + timedelta(seconds=2 * self.interval)) < now: - try: - # await self.publish_nats_expired(job['file_id']) - await self.mark_timeout(job.job_id) - await self.cache.remove(job.job_id) - except Exception as exc: - logger.exception("Произошла ошибка с пометкой timeout. Начало: {started_at}. Ожидаемый конец: {expires_at}. job_id={job_id}", job_id=job.job_id, started_at=job.started_at, expires_at=job.expires_at) - - else: - if search_fname in all_files: - new_dir = f'{self.smb_path}/Output/{search_fname}' - - logger.info("Начинаю обработку файла по пути smb_dir={smb_dir}", smb_dir=new_dir) - files = [] - for fname in smbclient.listdir(new_dir): - with smbclient.open_file(f'{new_dir}/{fname}', 'r') as f: - files.append(json.loads(f.read())) - logger.info("В рамках job_id={job_id} файлов в SMB найдено: {files_count}", job_id=job.job_id, files_count=len(files)) - s3_json_dict = { - 'file_id': job.file_id, - 'job_id': job.job_id, - 'legal_entity_id': job.legal_entity_id, - 'files': files - } - - async with async_session_maker() as session: - try: - await self.storage.put_object(f'{job.file_id}.json', json.dumps(s3_json_dict, ensure_ascii=False).encode('utf-8')) - db_data = json_to_db_rows(s3_json_dict) - db_meta = db_data['meta'] - db_rows = db_data['rows'] - db_values = db_data['values'] - skipped = db_data['skipped_files'] - if len(skipped) > 0: - logger.warning("При считывании в job_id={job_id} некоторые из файлов вызвали ошибку skipped={skipped}", job_id=job.job_id, skipped=skipped) - if len(db_meta) == 0: - raise ValueError("Пустой db_meta") - await session.execute(insert(FinancialStatementMeta).values(db_meta)) - await session.execute(insert(FinancialStatementLine).values(db_rows)) - await session.execute(insert(FinancialStatementValue).values(db_values)) - update_job_stmt = update(RecognitionJob).where(RecognitionJob.id == job.job_id).values( - status=RecognitionStatus.DONE, finished_at=datetime.now()) - await session.execute(update_job_stmt) - update_file_stmt = ( - update(StoredFile) - .where(StoredFile.id == job.file_id) - .values(document_type=DocumentType.BUH_OT, status=FileStatus.CHECK) - ) - await session.execute(update_file_stmt) - await session.commit() - logger.info("Коммит завершен по файлу file_id={file_id}", file_id=job.file_id) - except: - await session.rollback() - await self.mark_errored(session, job.job_id) - await session.commit() - - raise - finally: - await self.cache.remove(job.job_id) - except Exception as exc: - logger.exception("Ошибка при обработке результата в pdfpollworker. job_id={job_id}, error={error}", job_id=str(job.job_id), error=repr(exc)) - - - async def mark_timeout(self, job_id): - async with async_session_maker() as session: - stmt = ( - update(RecognitionJob) - .where(RecognitionJob.id == job_id) - .values(status = RecognitionStatus.ERROR, error_message="Истекло время ожидания", finished_at=datetime.now()) - .returning(RecognitionJob) - ) - file_stmt = ( - update(StoredFile) - .where(StoredFile.current_recognition_job_id == job_id) - .values(status = FileStatus.ERROR) - ) - res = await session.execute(stmt) - file_res = await session.execute(file_stmt) - - await session.commit() - - async def mark_errored(self, session, job_id): - stmt = ( - update(RecognitionJob) - .where(RecognitionJob.id == job_id) - .values(status=RecognitionStatus.ERROR, error_message="Ошибка в интеграции с ИС", finished_at=datetime.now()) - .returning(RecognitionJob) - ) - file_stmt = ( - update(StoredFile) - .where(StoredFile.current_recognition_job_id == job_id) - .values(status=FileStatus.ERROR) - ) - res = await session.execute(stmt) - file_res = await session.execute(file_stmt) diff --git a/api/src/core/async_cache.py b/api/src/core/async_cache.py deleted file mode 100644 index 32e864c..0000000 --- a/api/src/core/async_cache.py +++ /dev/null @@ -1,52 +0,0 @@ -import asyncio -from datetime import datetime, timedelta -import uuid - -from dataclasses import dataclass -from datetime import datetime - -@dataclass -class FileDictItem: - started_at: datetime - expires_at: datetime - file_id: str - job_id: str - s3_key: str - legal_entity_id: str - -class AsyncInMemoryCache: - def __init__(self): - self.mapping: dict[str, FileDictItem] = {} - self._lock = asyncio.Lock() - - async def add(self, job_id: str, file_id: str, s3_key: str, legal_entity_id: str) -> FileDictItem: - now = datetime.now() - expiration_date = datetime.now() + timedelta(minutes=10) - - new_item = FileDictItem( - started_at= now, - expires_at= expiration_date, - file_id= str(file_id), - s3_key=s3_key, - job_id= str(job_id), - legal_entity_id = legal_entity_id - ) - async with self._lock: - self.mapping[str(job_id)] = new_item - return new_item - - async def get(self, job_id: uuid.UUID): - async with self._lock: - return self.mapping.get(str(job_id)) - - async def remove(self, source_link: str) -> None: - async with self._lock: - self.mapping.pop(source_link, None) - - - async def get_active_items(self) -> list[FileDictItem]: - async with self._lock: - return [ - item - for item in self.mapping.values() - ] \ No newline at end of file diff --git a/api/src/core/clients.py b/api/src/core/clients.py index 666a6db..d14e916 100644 --- a/api/src/core/clients.py +++ b/api/src/core/clients.py @@ -1,29 +1,21 @@ -from datetime import datetime -import json - -import smbclient -from miniopy_async import Minio +from datetime import datetime, timedelta, timezone from io import BytesIO import asyncio -from nats.aio.client import Client as NATS -from sqlalchemy import or_, select, update, and_ +from loguru import logger +from pydantic import ValidationError +from miniopy_async import Minio +from sqlalchemy import and_, or_, select, update -from src.core.async_cache import AsyncInMemoryCache from src.models import RecognitionJob, StoredFile from src.schemas.enums import RecognitionStatus, FileStatus from src.db.db import async_session_maker -from src.core.config import settings +from src.core.pdfworker_tasks import create_pdfworker_task +from src.core.pdfworker_nats import publish_pdfworker_task -from loguru import logger +PDFWORKER_TIMEOUT = timedelta(minutes=10) +PUBLISH_LEASE = timedelta(seconds=60) -ext_to_subject = { - "XLS": "file.tabular_uploaded", - "XLSX": "file.tabular_uploaded", - "PNG": "file.img_uploaded", - "PDF": "file.img_uploaded", - "JPG": "file.img_uploaded", -} class S3Storage: def __init__(self, client: Minio, bucket_name: str): @@ -46,15 +38,11 @@ class S3Storage: raise - class PollWorker: - def __init__(self, nats_client: NATS, cache: AsyncInMemoryCache, s3_client: S3Storage): - self.task: asyncio.Task | None = None + def __init__(self, nats_client, cache=None, s3_client=None): + self.task = None self.stop_event = asyncio.Event() - self.nats_client = nats_client - self.cache = cache - self.storage = s3_client - + self.nats = nats_client async def start(self): self.stop_event.clear() @@ -75,106 +63,82 @@ class PollWorker: try: await self.poll_once() except Exception: - # logger.exception("Poll iteration failed") - pass - + logger.exception("Ошибка отправки задач PDFWorker") try: - await asyncio.wait_for( - self.stop_event.wait(), - timeout=7, - ) + await asyncio.wait_for(self.stop_event.wait(), timeout=7) except asyncio.TimeoutError: pass - logger.info("PollWorker остановлен") - - async def poll_once(self): - try: - rows = await self.claim_pending_rows() - logger.info("Получил rows. Количество: {rlen}", rlen=len(rows)) - except Exception as exc: - logger.exception( - "Проблема при получении строк из БД по RecognitionJob. error_category=pollworker, error_type={error_type}, error={error}", - error_type=type(exc).__name__, - error=repr(exc), - ) - rows=[] - - for row in rows: - row_id = row["id"] + for command in await self.claim_pending_rows(): try: - # await self.publish_to_nats(row["subject"], row["row_payload"]) - - if row['subject'] == 'file.img_uploaded': - last_part = row['row_payload']['s3_key'].split('/')[-1] - with smbclient.open_file(f'{settings.OPENBAO.SMB.SMB_PATH}/Input/{last_part}', 'wb') as f: - data = await self.storage.get_object(row['row_payload']['s3_key']) - content = await data.read() - f.write(content) - await self.cache.add(row['row_payload']['job_id'], row['row_payload']['file_id'], - row['row_payload']['s3_key'], row['row_payload']['legal_entity_id']) - - - except Exception as exc: - # logger.exception( - # "Failed to publish row id={row_id}. error_category=pollworker, error_type={error_type}, error={error}", - # row_id=row_id, - # error_type=type(exc).__name__, - # error=repr(exc), - # ) - logger.exception( - "Ошибка при попытке положить сообщение в кэш. error_category=pdfpollworker, error_type={error_type}, error={error}", - row_id=row_id, - error_type=type(exc).__name__, - error=repr(exc), - ) - - await self.mark_queued(row_id) + await publish_pdfworker_task(self.nats, command) + await self.mark_published(command.task_id) + except Exception: + logger.exception("Не удалось опубликовать задачу PDFWorker task_id={}", command.task_id) + await self.mark_queued(command.task_id) async def claim_pending_rows(self): + now = datetime.now(timezone.utc) async with async_session_maker() as session: - stmt = ( - select(RecognitionJob, StoredFile) - .join(StoredFile, StoredFile.id == RecognitionJob.file_id) - .where(and_(RecognitionJob.status == RecognitionStatus.QUEUED, StoredFile.status.is_(None))) - .with_for_update(skip_locked=True, read=True) - ) - result = await session.execute(stmt) - result_scalars = result.all() + async with session.begin(): + stmt = ( + select(RecognitionJob, StoredFile) + .join(StoredFile, StoredFile.id == RecognitionJob.file_id) + .where( + StoredFile.current_recognition_job_id == RecognitionJob.id, + or_(RecognitionJob.status == RecognitionStatus.QUEUED, + and_(RecognitionJob.status == RecognitionStatus.PROCESSING, + RecognitionJob.queue_msg_id.is_(None), + or_(RecognitionJob.started_at.is_(None), + RecognitionJob.started_at < now - PUBLISH_LEASE))), + ) + .with_for_update(skip_locked=True) + ) + rows = (await session.execute(stmt)).all() + commands = [] + for job, file in rows: + requested = job.requested_at or now + if requested.tzinfo is None: + requested = requested.replace(tzinfo=timezone.utc) + requested = requested.astimezone(timezone.utc) + try: + command = create_pdfworker_task( + file.s3_key, file.id, task_id=job.id, + expires_at=requested + PDFWORKER_TIMEOUT, + ) + except ValidationError as exc: + job.status = RecognitionStatus.ERROR + job.finished_at = now + job.error_message = "Invalid PDFWorker command: " + ", ".join( + ".".join(map(str, error["loc"])) for error in exc.errors()) + file.status = FileStatus.ERROR + logger.error("Некорректная задача PDFWorker task_id={}: {}", job.id, job.error_message) + continue + job.requested_at = requested + job.started_at = now + job.status = RecognitionStatus.PROCESSING + file.status = FileStatus.PROCESSING + commands.append(command) + return commands - for (job, file) in result_scalars: - job.status = RecognitionStatus.PROCESSING - file.status = FileStatus.PROCESSING - job.requested_at = datetime.now() - - - await session.commit() - - return [{ - "id": job.id, - "subject": ext_to_subject[file.file_type], - "row_payload": { - "legal_entity_id": str(job.legal_entity_id), - "file_id": str(file.id), - "s3_key": str(file.s3_key), - "job_id": str(job.id) - } - - } for (job, file) in result_scalars] - - async def publish_to_nats(self, subject, payload): - if not settings.DEBUG: - # logger.info("Попытка отправить сообщение в subject={subject}", subject=subject) - await self.nats_client.publish(subject, json.dumps(payload).encode("utf-8")) - else: - print("PUBLISHING TO NATS", subject, payload) #print достаточен, поскеольку работает локально - print() - - async def mark_queued(self, row_id): + async def mark_published(self, task_id): async with async_session_maker() as session: - stmt = update(RecognitionJob).where(RecognitionJob.id == row_id).values(status=RecognitionStatus.QUEUED) - file_stmt = update(StoredFile).where(StoredFile.current_recognition_job_id == row_id).values(status=None) - await session.execute(stmt) - await session.execute(file_stmt) - await session.commit() \ No newline at end of file + async with session.begin(): + await session.execute(update(RecognitionJob).where( + RecognitionJob.id == task_id, + RecognitionJob.status == RecognitionStatus.PROCESSING, + ).values(queue_msg_id=f"pdfworker:prepare:{task_id}")) + + async def mark_queued(self, task_id): + async with async_session_maker() as session: + async with session.begin(): + job = (await session.execute(select(RecognitionJob).where( + RecognitionJob.id == task_id).with_for_update())).scalar_one_or_none() + # A result may have committed before a publish timeout/retry. + if job is None or job.status not in (RecognitionStatus.PROCESSING, RecognitionStatus.QUEUED): + return + job.status = RecognitionStatus.QUEUED + await session.execute(update(StoredFile).where( + StoredFile.current_recognition_job_id == task_id, + ).values(status=None)) diff --git a/api/src/core/config.py b/api/src/core/config.py index c87a231..355ab28 100644 --- a/api/src/core/config.py +++ b/api/src/core/config.py @@ -10,17 +10,19 @@ class MinioSettings(BaseSettings): MINIO_ADMIN_LOGIN: str MINIO_ADMIN_PASSWORD: str MINIO_BUCKET: str + MINIO_SECURE: bool = True class NATSSettings(BaseSettings): NATS_URL: str - NATS_SEED: str - NATS_JWT: str + NATS_SEED: str = "" + NATS_JWT: str = "" class SMBSettings(BaseSettings): SMB_HOST: str SMB_PATH: str SMB_LOGIN: str SMB_PASSWORD: str + SMB_PORT: int = 445 class OpenBAOSettings(BaseSettings): DATABASE: DBSettings @@ -48,4 +50,4 @@ class Settings(BaseSettings): def get_settings() -> Settings: return Settings() -settings = get_settings() \ No newline at end of file +settings = get_settings() diff --git a/api/src/core/pdfworker_nats.py b/api/src/core/pdfworker_nats.py new file mode 100644 index 0000000..7ad7fc6 --- /dev/null +++ b/api/src/core/pdfworker_nats.py @@ -0,0 +1,80 @@ +"""NATS transport; application supplies a validated-result handler.""" +import base64 +from contextlib import asynccontextmanager +import logging + +from nats.aio.client import Client as NATS +import nkeys +from pydantic import ValidationError + +from src.schemas.pdfworker import (COMMAND_SUBJECT, RESULT_SUBJECT, + PDFWorkerCommand, PDFWorkerResult) + +logger = logging.getLogger(__name__) + + +def nats_auth(jwt: str = "", seed: str = "") -> dict: + if bool(jwt) != bool(seed): + raise ValueError("NATS JWT and 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: + return base64.b64encode(key.sign(nonce.encode() if isinstance(nonce, str) else nonce)) + finally: + key.wipe() + return {"user_jwt_cb": jwt_cb, "signature_cb": signature_cb} + + +async def publish_pdfworker_task(nc, command: PDFWorkerCommand): + command = PDFWorkerCommand.model_validate(command) + await nc.publish(COMMAND_SUBJECT, command.model_dump_json().encode()) + await nc.flush() + + +async def handle_result_message(msg, handle_result): + try: + event = PDFWorkerResult.model_validate_json(msg.data) + except ValidationError as exc: + # No payload/secrets in the log; field locations are enough to diagnose. + logger.error("Invalid PDFWorker result: %s", exc.errors(include_input=False, include_context=False)) + return + try: + await handle_result(event) + except Exception: + logger.exception("PDFWorker result processing failed: task_id=%s", event.task_id) + + +@asynccontextmanager +async def pdfworker_connection(url, handle_result, *, jwt="", seed=""): + async def error_cb(error): + logger.error("NATS connection error: %s", type(error).__name__) + + nc = NATS() + subscription = None + try: + await nc.connect(url, **nats_auth(jwt, seed), connect_timeout=5, + max_reconnect_attempts=2, drain_timeout=60, error_cb=error_cb) + async def callback(msg): + await handle_result_message(msg, handle_result) + subscription = await nc.subscribe(RESULT_SUBJECT, cb=callback) + await nc.flush() + yield nc + finally: + # Drain queued callbacks while DB/MinIO dependencies are still alive. + try: + if subscription is not None: + await subscription.drain() + finally: + if nc is not None and not nc.is_closed: + try: + if nc.is_connected: + await nc.drain() + finally: + if not nc.is_closed: + await nc.close() diff --git a/api/src/core/pdfworker_tasks.py b/api/src/core/pdfworker_tasks.py new file mode 100644 index 0000000..0558bd5 --- /dev/null +++ b/api/src/core/pdfworker_tasks.py @@ -0,0 +1,12 @@ +"""Pure command construction: no settings, network connections or database writes.""" +from datetime import datetime +from uuid import UUID + +from src.schemas.pdfworker import PDFWorkerCommand + + +def create_pdfworker_task(s3_key: str, file_id: UUID | str, *, task_id: UUID | str, + expires_at: datetime | str) -> PDFWorkerCommand: + """Use persisted task_id/expires_at on retries; preserve the complete object key.""" + return PDFWorkerCommand(task_id=task_id, file_id=file_id, source_s3_key=s3_key, + expires_at=expires_at) diff --git a/api/src/schemas/enums.py b/api/src/schemas/enums.py index dee2046..6b26efe 100644 --- a/api/src/schemas/enums.py +++ b/api/src/schemas/enums.py @@ -20,12 +20,15 @@ class DocumentType(StrEnum): OFR = "ОФР" OTHER = "Прочий документ" +PDFWORKER_DOCUMENT_TYPE = DocumentType.OFR + class FileType(StrEnum): XLSX = "XLSX" XLS = "XLS" PDF = "PDF" PNG = "PNG" JPG = "JPG" + FILE = "FILE" class RecognitionPipeline(StrEnum): TABULAR = "tabular" @@ -62,4 +65,4 @@ class FslRowKind(StrEnum): ROW = "row" TOTAL = "total" BALANCE = "balance" - HEADER= "header" \ No newline at end of file + HEADER= "header" diff --git a/api/src/schemas/pdfworker.py b/api/src/schemas/pdfworker.py new file mode 100644 index 0000000..d4aad20 --- /dev/null +++ b/api/src/schemas/pdfworker.py @@ -0,0 +1,61 @@ +"""PDFWorker wire protocol. Kept compatible with PDFWorker schemas/messages.py.""" +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/services/file_service.py b/api/src/services/file_service.py index e5674fd..d5cd7a9 100644 --- a/api/src/services/file_service.py +++ b/api/src/services/file_service.py @@ -7,7 +7,7 @@ from sqlalchemy.inspection import inspect from loguru import logger from src.core.clients import S3Storage -from src.schemas.enums import FileStatus, FileType, RecognitionStatus, RecognitionPipeline +from src.schemas.enums import PDFWORKER_DOCUMENT_TYPE, FileType, RecognitionStatus def get_dict(obj): return {c.key: getattr(obj, c.key) for c in inspect(obj).mapper.column_attrs} @@ -29,8 +29,14 @@ class FileService: 'period_year': year, 'size_bytes': file.size, } - file_ext = file.filename.lower().split('.')[-1] - db_payload['file_type'] = FileType[file_ext.upper()] + if not file.filename or '.' not in file.filename: + raise ValueError("Имя файла должно содержать расширение") + file_ext = file.filename.rsplit('.', 1)[-1].lower() + if (not file_ext or file_ext.endswith(' ') or + any(ord(char) < 32 or char in '\\/:*?"<>|' for char in file_ext)): + raise ValueError("Некорректное расширение файла") + db_payload['file_type'] = FileType.__members__.get(file_ext.upper(), FileType.FILE) + db_payload['document_type'] = PDFWORKER_DOCUMENT_TYPE new_file = await self.file_repo.insert_one(db_payload) s3_key = f'{str(new_file.id)}.{file_ext}' @@ -38,16 +44,11 @@ class FileService: new_file.s3_key = s3_key # new_file.status = FileStatus.UPLOADED - job_pipeline = RecognitionPipeline.IMAGE - if (file_ext in ['xlsx', 'xls']): - job_pipeline = RecognitionPipeline.TABULAR - rec_job_db_payload = { 'file_id': new_file.id, 'legal_entity_id': legal_entity_id, 'status': RecognitionStatus.QUEUED, 'attempt': 1, - 'pipeline': job_pipeline, } rec_job = await self.rec_job_repo.insert_one(rec_job_db_payload) @@ -71,15 +72,21 @@ class FileService: async def get_file_bytes(self, file_id): result = await self.file_repo.get_by_id(file_id) - return (await self.storage.get_object(result.s3_key)) + return await self.storage.get_object(result.s3_key), self._source_extension(result) + + @staticmethod + def _source_extension(stored_file_db_obj): + name = stored_file_db_obj.s3_key or stored_file_db_obj.original_name + if '.' in name: + return name.rsplit('.', 1)[-1].upper() + return stored_file_db_obj.file_type.value def _get_transformed_file_job_info(self, stored_file_db_obj): return { 'id': stored_file_db_obj.id, 'legal_entity_id': stored_file_db_obj.legal_entity_id, 'status': stored_file_db_obj.status, - 'file_type': stored_file_db_obj.file_type, + 'file_type': self._source_extension(stored_file_db_obj), 'original_name': stored_file_db_obj.original_name, 'document_type': stored_file_db_obj.document_type } - diff --git a/api/src/core/PDFWorker/__init__.py b/api/src/services/ofr_result_conversion/__init__.py similarity index 100% rename from api/src/core/PDFWorker/__init__.py rename to api/src/services/ofr_result_conversion/__init__.py diff --git a/api/src/core/PDFWorker/consts.py b/api/src/services/ofr_result_conversion/consts.py similarity index 100% rename from api/src/core/PDFWorker/consts.py rename to api/src/services/ofr_result_conversion/consts.py diff --git a/api/src/core/PDFWorker/utils.py b/api/src/services/ofr_result_conversion/utils.py similarity index 100% rename from api/src/core/PDFWorker/utils.py rename to api/src/services/ofr_result_conversion/utils.py diff --git a/api/src/services/pdfworker_results.py b/api/src/services/pdfworker_results.py new file mode 100644 index 0000000..a80cf3f --- /dev/null +++ b/api/src/services/pdfworker_results.py @@ -0,0 +1,79 @@ +"""Receive PDFWorker results and persist them through the selected handler.""" +import json +import logging + +from sqlalchemy import select, insert +from src.models import (RecognitionJob, StoredFile, FinancialStatementMeta, + FinancialStatementLine, FinancialStatementValue) +from src.schemas.enums import RecognitionStatus, FileStatus, DocumentType, PDFWORKER_DOCUMENT_TYPE +from src.services.ofr_result_conversion.utils import json_to_db_rows + +logger = logging.getLogger(__name__) + + +def result_payload(data, event, legal_entity_id): + """Validate linkage and use ARMKA's authoritative IDs for DB transformation.""" + if not isinstance(data, dict) or not isinstance(data.get("files"), list): + raise ValueError("Expected a merged JSON object with files") + for name, expected in (("task_id", event.task_id), ("job_id", event.task_id), ("file_id", event.file_id)): + if str(data.get(name)) != str(expected): + raise ValueError(f"Result JSON has a mismatched {name}") + return {**data, "legal_entity_id": str(legal_entity_id)} + + +async def process_ofr_result(session, payload, event, job): + """Persist the OFR payload; the caller owns the transaction and job status.""" + rows = json_to_db_rows(result_payload(payload, event, job.legal_entity_id)) + if not rows['meta']: + raise ValueError("Result contains no financial statement metadata") + for model, key in ((FinancialStatementMeta, 'meta'), + (FinancialStatementLine, 'rows'), + (FinancialStatementValue, 'values')): + if rows[key]: + await session.execute(insert(model).values(rows[key])) + if rows.get('skipped_files'): + logger.warning("Task %s: skipped %s result parts", event.task_id, len(rows['skipped_files'])) + + +RESULT_HANDLERS = {DocumentType.OFR: process_ofr_result} + + +async def apply_pdfworker_result(event, session_factory, storage): + async with session_factory() as session: + async with session.begin(): + job = (await session.execute(select(RecognitionJob).where( + RecognitionJob.id == event.task_id).with_for_update())).scalar_one_or_none() + if job is None: + # e.g. an independent smoke test uses its own UUIDs on the same subject. + logger.warning("Ignoring PDFWorker result for unknown task %s", event.task_id) + return + if job.file_id != event.file_id: + raise ValueError("Result file_id does not match the task") + if job.status in (RecognitionStatus.DONE, RecognitionStatus.ERROR): + logger.info("Ignoring duplicate terminal result for %s", event.task_id) + return + file = (await session.execute(select(StoredFile).where( + StoredFile.id == job.file_id).with_for_update())).scalar_one() + if file.current_recognition_job_id != job.id: + logger.warning("Ignoring result for superseded task %s", event.task_id) + return + if event.status == "ready": + handler = RESULT_HANDLERS.get(PDFWORKER_DOCUMENT_TYPE) + if handler is None: + raise ValueError(f"No PDFWorker result handler for {PDFWORKER_DOCUMENT_TYPE}") + response = await storage.get_object(event.result_s3_key) + try: + payload = json.loads(await response.read()) + finally: + response.release() + await handler(session, payload, event, job) + job.raw_result_s3_key = event.result_s3_key + job.status = RecognitionStatus.DONE + job.error_message = None + file.status = FileStatus.CHECK + else: + job.status = RecognitionStatus.ERROR + job.error_message = f"PDFWorker {event.status}: {event.error_code}" + file.status = FileStatus.ERROR + file.document_type = PDFWORKER_DOCUMENT_TYPE + job.finished_at = event.finished_at diff --git a/api/tests/unit/test_pdfworker_contract.py b/api/tests/unit/test_pdfworker_contract.py new file mode 100644 index 0000000..101e115 --- /dev/null +++ b/api/tests/unit/test_pdfworker_contract.py @@ -0,0 +1,156 @@ +from datetime import datetime, timedelta, timezone +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock +from uuid import uuid4 +import pytest +from pydantic import ValidationError +from src.schemas.pdfworker import PDFWorkerCommand, PDFWorkerResult +from src.core.pdfworker_tasks import create_pdfworker_task +from src.core.pdfworker_nats import handle_result_message, publish_pdfworker_task, pdfworker_connection, nats_auth + + +def command(): + return create_pdfworker_task('uploads/2026/report.pdf', uuid4(), task_id=uuid4(), + expires_at=datetime.now(timezone.utc) + timedelta(minutes=10)) + + +def event(status='ready'): + task = command() + return PDFWorkerResult(task_id=task.task_id, file_id=task.file_id, status=status, + result_s3_key='pdfworker/results/result.json' if status == 'ready' else None, + error_code=None if status == 'ready' else 'storage_unavailable', + finished_at=datetime.now(timezone.utc)) + + +def message(data): + return SimpleNamespace(data=data) + + +def test_full_key_and_persisted_identity(): + a = command() + b = create_pdfworker_task(a.source_s3_key, a.file_id, task_id=a.task_id, expires_at=a.expires_at) + assert a == b + assert a.source_s3_key == 'uploads/2026/report.pdf' + assert 'version' not in a.model_dump() + assert 'legal_entity_id' not in a.model_dump() + assert PDFWorkerCommand.model_validate_json(a.model_dump_json()) == a + + +@pytest.mark.parametrize('extension', ['docx', 'zip', 'jpeg', 'csv', 'c++', 'данные']) +def test_arbitrary_safe_extension(extension): + task = command() + data = {**task.model_dump(), 'source_s3_key': f'uploads/report.{extension}'} + assert PDFWorkerCommand.model_validate(data).source_s3_key == data['source_s3_key'] + + +@pytest.mark.parametrize('field,value', [('version', 2), ('task_id', 'bad'), ('file_id', 'bad'), + ('expires_at', 'test'), ('expires_at', '2026-10-01T12:00:00'), + ('source_s3_key', 's3://bucket/doc.pdf'), ('source_s3_key', ''), + ('source_s3_key', 'uploads/report.bad:name'), ('source_files', ['x.json'])]) +def test_invalid_command(field, value): + with pytest.raises(ValidationError): + PDFWorkerCommand.model_validate({**command().model_dump(), field: value}) + + +@pytest.mark.parametrize('updates', [dict(version=2), dict(status='success'), dict(error_code='unknown'), + dict(result_s3_key=None), dict(error_code='storage_unavailable'), + dict(status='timeout', result_s3_key=None, error_code='source_not_found'), + dict(status='error', result_s3_key=None, error_code='deadline_exceeded')]) +def test_invalid_result(updates): + with pytest.raises(ValidationError): + PDFWorkerResult.model_validate({**event().model_dump(), **updates}) + + +def test_expired_command_still_valid_for_timeout_response(): + data = command().model_dump() + data['expires_at'] = datetime.now(timezone.utc) - timedelta(days=1) + assert PDFWorkerCommand.model_validate(data) + + +@pytest.mark.asyncio +async def test_publish_command_and_flush(): + nc = SimpleNamespace(publish=AsyncMock(), flush=AsyncMock()) + task = command() + await publish_pdfworker_task(nc, task) + assert nc.publish.call_args.args[0] == 'files.tabular_uploaded' + assert PDFWorkerCommand.model_validate_json(nc.publish.call_args.args[1]) == task + nc.flush.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize('status', ['ready', 'error']) +async def test_result_handler(status): + calls = [] + msg = message(event(status).model_dump_json().encode()) + async def handle(result): + assert result.status == status + calls.append('commit') + await handle_result_message(msg, handle) + assert calls == ['commit'] + + +@pytest.mark.asyncio +async def test_invalid_message_is_ignored(): + msg, handler = message(b'{}'), AsyncMock() + await handle_result_message(msg, handler) + handler.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_handler_failure_is_logged(): + msg = message(event().model_dump_json().encode()) + await handle_result_message(msg, AsyncMock(side_effect=RuntimeError('db down'))) + + +@pytest.mark.asyncio +@pytest.mark.parametrize('fail_subscribe', [False, True]) +async def test_lifecycle_cleanup(monkeypatch, fail_subscribe): + from src.core import pdfworker_nats + sub = SimpleNamespace(drain=AsyncMock()) + subscribe = AsyncMock(return_value=sub) + if fail_subscribe: + subscribe.side_effect = RuntimeError('subscription denied') + nc = SimpleNamespace(is_closed=False, is_connected=True, connect=AsyncMock(), + subscribe=subscribe, flush=AsyncMock(), drain=AsyncMock(), close=AsyncMock()) + monkeypatch.setattr(pdfworker_nats, 'NATS', lambda: nc) + if fail_subscribe: + with pytest.raises(RuntimeError): + async with pdfworker_connection('nats://localhost', AsyncMock()): + pytest.fail('should not yield') + else: + async with pdfworker_connection('nats://localhost', AsyncMock()) as connected: + assert connected is nc + pass + subscribe.assert_awaited_once() + sub.drain.assert_awaited_once() + nc.drain.assert_awaited_once() + nc.close.assert_awaited_once() + + +def test_optional_auth(): + assert nats_auth() == {} + with pytest.raises(ValueError): + nats_auth('jwt', '') + + +def test_models_match_pdfworker_checkout(): + counterpart = Path(__file__).resolve().parents[4] / 'DKB_ARMKA_FILE_FLOW/api/src/schemas/messages.py' + if not counterpart.exists(): + pytest.skip('PDFWorker checkout not available') + from runpy import run_path + models = run_path(str(counterpart)) + assert PDFWorkerCommand.model_json_schema() == models['PDFWorkerCommand'].model_json_schema() + assert PDFWorkerResult.model_json_schema() == models['PDFWorkerResult'].model_json_schema() + + +@pytest.mark.asyncio +async def test_failed_connect_closes_allocated_client(monkeypatch): + from src.core import pdfworker_nats + nc = SimpleNamespace(is_closed=False, is_connected=False, + connect=AsyncMock(side_effect=RuntimeError('offline')), close=AsyncMock()) + monkeypatch.setattr(pdfworker_nats, 'NATS', lambda: nc) + with pytest.raises(RuntimeError, match='offline'): + async with pdfworker_connection('nats://localhost', AsyncMock()): + pytest.fail('should not yield') + nc.close.assert_awaited_once() diff --git a/api/tests/unit/test_pdfworker_persistence.py b/api/tests/unit/test_pdfworker_persistence.py new file mode 100644 index 0000000..3aafe75 --- /dev/null +++ b/api/tests/unit/test_pdfworker_persistence.py @@ -0,0 +1,223 @@ +import json +from io import BytesIO +from datetime import date, datetime, timedelta, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock +from uuid import uuid4 +import pytest +import pytest_asyncio +from fastapi import UploadFile +from sqlalchemy import select, func +from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker +from src.models import (RecognitionJob, StoredFile, FinancialStatementMeta, + FinancialStatementLine, FinancialStatementValue) +from src.models.base import Base +from src.schemas.enums import RecognitionStatus, RecognitionPipeline, FileStatus, DocumentType, FileType, StatementType, FslRowKind +from src.schemas.pdfworker import PDFWorkerResult +from src.services import pdfworker_results +from src.services.file_service import FileService +from src.repositories.file_repository import FileRepository +from src.repositories.recognition_job_repo import RecognitionJobRepository +from src.core import clients + + +@pytest_asyncio.fixture +async def db(): + engine = create_async_engine('sqlite+aiosqlite:///:memory:') + tables = [StoredFile.__table__, RecognitionJob.__table__, FinancialStatementMeta.__table__, + FinancialStatementLine.__table__, FinancialStatementValue.__table__] + async with engine.begin() as connection: + await connection.run_sync(lambda conn: Base.metadata.create_all(conn, tables=tables)) + factory = async_sessionmaker(engine, expire_on_commit=False) + ids = SimpleNamespace(task=uuid4(), file=uuid4(), legal=uuid4()) + async with factory() as session: + session.add(StoredFile(id=ids.file, legal_entity_id=ids.legal, file_type=FileType.PDF, + document_type=DocumentType.OFR, + original_name='report.pdf', s3_key='folder/nested/report.pdf', + size_bytes=123, period_year=2026, current_recognition_job_id=ids.task)) + session.add(RecognitionJob(id=ids.task, file_id=ids.file, legal_entity_id=ids.legal, + pipeline=RecognitionPipeline.IMAGE, status=RecognitionStatus.QUEUED, attempt=1)) + await session.commit() + yield factory, ids + await engine.dispose() + + +def ready_event(ids, status='ready'): + return PDFWorkerResult(task_id=ids.task, file_id=ids.file, status=status, + result_s3_key='results/full/path.json' if status == 'ready' else None, + error_code=None if status == 'ready' else 'storage_unavailable', finished_at=datetime.now(timezone.utc)) + + +def setup_storage(ids, monkeypatch): + payload = dict(task_id=str(ids.task), job_id=str(ids.task), file_id=str(ids.file), files=[{}]) + response = SimpleNamespace(read=AsyncMock(return_value=json.dumps(payload).encode()), release=Mock()) + storage = SimpleNamespace(get_object=AsyncMock(return_value=response)) + meta_id, line_id = uuid4(), uuid4() + converter = Mock(return_value=dict( + meta=[dict(id=meta_id, stored_file_id=ids.file, statement_type=StatementType.BALANCE, + inn='123', inn_verified=True, units='руб', units_verified=True, + report_date=date(2026,1,1), report_date_verified=True)], + rows=[dict(id=line_id, recognition_job_id=ids.task, meta_id=meta_id, + legal_entity_id=ids.legal, line_code='1600', line_code_recognized=True, + line_name='Итого', line_name_recognized=True, row_kind=FslRowKind.ROW, ordinal=1)], + values=[dict(id=uuid4(), line_id=line_id, amount=100, string_value='100', column_index=1, is_recognized=True)], + skipped_files=[])) + monkeypatch.setattr(pdfworker_results, 'json_to_db_rows', converter) + return storage, response, converter + + +@pytest.mark.asyncio +async def test_ready_commits_once(db, monkeypatch): + factory, ids = db + storage, response, converter = setup_storage(ids, monkeypatch) + event = ready_event(ids) + await pdfworker_results.apply_pdfworker_result(event, factory, storage) + await pdfworker_results.apply_pdfworker_result(event, factory, storage) + storage.get_object.assert_awaited_once_with('results/full/path.json') + response.release.assert_called_once() + assert converter.call_args.args[0]['legal_entity_id'] == str(ids.legal) + async with factory() as session: + assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.DONE + stored_file = await session.get(StoredFile, ids.file) + assert stored_file.status == FileStatus.CHECK + assert stored_file.document_type == DocumentType.OFR + for model in (FinancialStatementMeta, FinancialStatementLine, FinancialStatementValue): + assert await session.scalar(select(func.count()).select_from(model)) == 1 + + +@pytest.mark.asyncio +async def test_upload_accepts_new_extension_and_saves_ofr_type(db): + factory, ids = db + storage = Mock(put_object=AsyncMock()) + async with factory() as session: + service = FileService(storage, session, FileRepository(session), + RecognitionJobRepository(session), None) + uploaded = UploadFile(filename='report.docx', file=BytesIO(b'data'), size=4) + created = await service.save_file(uploaded, ids.legal, 'Q1', 2026) + stored = await session.get(StoredFile, created['file']['id']) + job = await session.get(RecognitionJob, created['rec_job']['id']) + assert stored.file_type == FileType.FILE + assert stored.document_type == DocumentType.OFR + assert job.file_id == stored.id + assert service._get_transformed_file_job_info(stored)['file_type'] == 'DOCX' + storage.put_object.assert_awaited_once_with(f'{stored.id}.docx', b'data') + + +@pytest.mark.asyncio +@pytest.mark.parametrize('status', ['error', 'timeout']) +async def test_error_without_minio(db, status): + factory, ids = db + data = ready_event(ids, 'error').model_dump() + if status == 'timeout': + data.update(status='timeout', error_code='deadline_exceeded') + event = PDFWorkerResult(**data) + storage = Mock(get_object=AsyncMock()) + await pdfworker_results.apply_pdfworker_result(event, factory, storage) + await pdfworker_results.apply_pdfworker_result(event, factory, storage) + storage.get_object.assert_not_awaited() + async with factory() as session: + job = await session.get(RecognitionJob, ids.task) + assert job.status == RecognitionStatus.ERROR + assert event.error_code in job.error_message + assert (await session.get(StoredFile, ids.file)).status == FileStatus.ERROR + + +@pytest.mark.asyncio +async def test_invalid_json_rolls_back(db, monkeypatch): + factory, ids = db + storage, response, _ = setup_storage(ids, monkeypatch) + response.read.return_value = b'{"files": []}' + with pytest.raises(ValueError): + await pdfworker_results.apply_pdfworker_result(ready_event(ids), factory, storage) + response.release.assert_called_once() + async with factory() as session: + assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.QUEUED + assert await session.scalar(select(func.count()).select_from(FinancialStatementMeta)) == 0 + + +@pytest.mark.asyncio +async def test_pollworker_recovers_same_command(db, monkeypatch): + factory, ids = db + monkeypatch.setattr(clients, 'async_session_maker', factory) + worker = clients.PollWorker(SimpleNamespace(jetstream=lambda: SimpleNamespace())) + first, = await worker.claim_pending_rows() + assert first.source_s3_key == 'folder/nested/report.pdf' + assert first.task_id == ids.task + assert not await worker.claim_pending_rows() + async with factory() as session: + job = await session.get(RecognitionJob, ids.task) + job.started_at = datetime.now(timezone.utc) - timedelta(minutes=2) + await session.commit() + retry, = await worker.claim_pending_rows() + assert first.model_dump_json() == retry.model_dump_json() + await worker.mark_published(first.task_id) + assert not await worker.claim_pending_rows() + + +@pytest.mark.asyncio +async def test_cannot_requeue_finished_job(db, monkeypatch): + factory, ids = db + monkeypatch.setattr(clients, 'async_session_maker', factory) + await pdfworker_results.apply_pdfworker_result(ready_event(ids, 'error'), factory, Mock()) + worker = clients.PollWorker(SimpleNamespace(jetstream=lambda: None)) + await worker.mark_queued(ids.task) + async with factory() as session: + assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.ERROR + assert (await session.get(StoredFile, ids.file)).status == FileStatus.ERROR + + +@pytest.mark.asyncio +async def test_invalid_stored_key_does_not_block_queue(db, monkeypatch): + factory, ids = db + monkeypatch.setattr(clients, 'async_session_maker', factory) + async with factory() as session: + file = await session.get(StoredFile, ids.file) + file.s3_key = None + await session.commit() + worker = clients.PollWorker(SimpleNamespace(jetstream=lambda: None)) + assert await worker.claim_pending_rows() == [] + async with factory() as session: + job = await session.get(RecognitionJob, ids.task) + assert job.status == RecognitionStatus.ERROR + assert 'source_s3_key' in job.error_message + + +@pytest.mark.asyncio +async def test_wrong_file_id_does_not_change_task(db): + factory, ids = db + data = ready_event(ids, 'error').model_dump() + data['file_id'] = uuid4() + with pytest.raises(ValueError, match='file_id'): + await pdfworker_results.apply_pdfworker_result(PDFWorkerResult(**data), factory, Mock()) + async with factory() as session: + assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.QUEUED + + +@pytest.mark.asyncio +async def test_dispatches_flagged_document_type(db, monkeypatch): + factory, ids = db + storage, _, _ = setup_storage(ids, monkeypatch) + async with factory() as session: + stored = await session.get(StoredFile, ids.file) + stored.document_type = DocumentType.OSV + await session.commit() + handler = AsyncMock() + monkeypatch.setitem(pdfworker_results.RESULT_HANDLERS, DocumentType.OFR, handler) + await pdfworker_results.apply_pdfworker_result(ready_event(ids), factory, storage) + handler.assert_awaited_once() + async with factory() as session: + job = await session.get(RecognitionJob, ids.task) + assert job.status == RecognitionStatus.DONE + assert (await session.get(StoredFile, ids.file)).document_type == DocumentType.OFR + + +@pytest.mark.asyncio +async def test_missing_handler_does_not_complete_job(db, monkeypatch): + factory, ids = db + monkeypatch.delitem(pdfworker_results.RESULT_HANDLERS, DocumentType.OFR) + storage = Mock(get_object=AsyncMock()) + with pytest.raises(ValueError, match='handler'): + await pdfworker_results.apply_pdfworker_result(ready_event(ids), factory, storage) + storage.get_object.assert_not_awaited() + async with factory() as session: + assert (await session.get(RecognitionJob, ids.task)).status == RecognitionStatus.QUEUED -- 2.47.2