86 lines
3.2 KiB
Python
86 lines
3.2 KiB
Python
import json
|
|
import uuid
|
|
|
|
from fastapi import UploadFile
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from sqlalchemy.inspection import inspect
|
|
from loguru import logger
|
|
|
|
from src.core.clients import S3Storage
|
|
from src.schemas.enums import FileStatus, FileType, RecognitionStatus, RecognitionPipeline
|
|
|
|
def get_dict(obj):
|
|
return {c.key: getattr(obj, c.key) for c in inspect(obj).mapper.column_attrs}
|
|
|
|
class FileService:
|
|
def __init__(self, storage: S3Storage, session: AsyncSession, file_repository, recognition_job_repo, nats_conn):
|
|
self.storage = storage
|
|
self.file_repo = file_repository
|
|
self.rec_job_repo = recognition_job_repo
|
|
self.nats_conn = nats_conn
|
|
self.session = session
|
|
|
|
async def save_file(self, file: UploadFile, legal_entity_id: uuid.UUID, quarter: str, year: int):
|
|
try:
|
|
db_payload = {
|
|
'original_name': file.filename,
|
|
'legal_entity_id': legal_entity_id,
|
|
'period_quarter': quarter,
|
|
'period_year': year,
|
|
'size_bytes': file.size,
|
|
}
|
|
file_ext = file.filename.lower().split('.')[-1]
|
|
db_payload['file_type'] = FileType[file_ext.upper()]
|
|
new_file = await self.file_repo.insert_one(db_payload)
|
|
|
|
s3_key = f'{str(new_file.id)}.{file_ext}'
|
|
await self.storage.put_object(s3_key, await file.read())
|
|
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)
|
|
new_file.current_recognition_job_id = rec_job.id
|
|
|
|
await self.session.commit()
|
|
|
|
|
|
return {'file': get_dict(new_file), 'rec_job': get_dict(rec_job)}
|
|
|
|
|
|
except Exception as e:
|
|
await self.session.rollback()
|
|
logger.exception("Ошибка при сохранении файла в БД")
|
|
|
|
raise
|
|
|
|
async def get_files_by_legal_entity(self, legal_entity_id: uuid.UUID):
|
|
result = await self.file_repo.get_by_legal_entity(legal_entity_id)
|
|
return [self._get_transformed_file_job_info(new_file) for new_file in result]
|
|
|
|
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))
|
|
|
|
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,
|
|
'original_name': stored_file_db_obj.original_name,
|
|
'document_type': stored_file_db_obj.document_type
|
|
}
|
|
|