изменения для нового сервиса #11
@ -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',
|
||||
]
|
||||
|
||||
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)
|
||||
|
||||
@ -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"'
|
||||
},)
|
||||
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()}')}"
|
||||
},)
|
||||
|
||||
138
api/src/app.py
138
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")
|
||||
app.mount("/", StaticFiles(directory="web", html=True), name="web")
|
||||
|
||||
@ -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, port=settings.OPENBAO.SMB.SMB_PORT):
|
||||
with smbclient.open_file(f'{new_dir}/{fname}', 'r', port=settings.OPENBAO.SMB.SMB_PORT) 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)
|
||||
@ -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()
|
||||
]
|
||||
@ -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,107 +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}/Файлы/{last_part}', 'wb', port=settings.OPENBAO.SMB.SMB_PORT) 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()
|
||||
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))
|
||||
|
||||
@ -12,11 +12,12 @@ 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_PORT: int = Field(default=445, ge=1, le=65535)
|
||||
@ -24,6 +25,7 @@ class SMBSettings(BaseSettings):
|
||||
SMB_PATH: str
|
||||
SMB_LOGIN: str
|
||||
SMB_PASSWORD: str
|
||||
SMB_PORT: int = 445
|
||||
|
||||
class OpenBAOSettings(BaseSettings):
|
||||
DATABASE: DBSettings
|
||||
@ -51,4 +53,4 @@ class Settings(BaseSettings):
|
||||
def get_settings() -> Settings:
|
||||
return Settings()
|
||||
|
||||
settings = get_settings()
|
||||
settings = get_settings()
|
||||
|
||||
80
api/src/core/pdfworker_nats.py
Normal file
80
api/src/core/pdfworker_nats.py
Normal file
@ -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()
|
||||
12
api/src/core/pdfworker_tasks.py
Normal file
12
api/src/core/pdfworker_tasks.py
Normal file
@ -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)
|
||||
@ -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"
|
||||
HEADER= "header"
|
||||
|
||||
61
api/src/schemas/pdfworker.py
Normal file
61
api/src/schemas/pdfworker.py
Normal file
@ -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
|
||||
@ -9,7 +9,7 @@ from datetime import datetime, timezone
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
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}
|
||||
@ -32,8 +32,14 @@ class FileService:
|
||||
'size_bytes': file.size,
|
||||
'uploaded_at': datetime.now(ZoneInfo("Europe/Moscow"))
|
||||
}
|
||||
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}'
|
||||
@ -41,16 +47,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)
|
||||
@ -74,15 +75,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
|
||||
}
|
||||
|
||||
|
||||
79
api/src/services/pdfworker_results.py
Normal file
79
api/src/services/pdfworker_results.py
Normal file
@ -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
|
||||
156
api/tests/unit/test_pdfworker_contract.py
Normal file
156
api/tests/unit/test_pdfworker_contract.py
Normal file
@ -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()
|
||||
223
api/tests/unit/test_pdfworker_persistence.py
Normal file
223
api/tests/unit/test_pdfworker_persistence.py
Normal file
@ -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
|
||||
Loading…
x
Reference in New Issue
Block a user