feat: update le docs logic

src/infrastructure/database/repositories/legal_entity_repository.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 from typing import Any
 
-from fastapi import status
 from sqlalchemy import func, select, update
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.exc import IntegrityError, SQLAlchemyError
@@ -11,6 +10,7 @@ from src.application.abstractions.repositories import ILegalEntityRepository
 from src.application.contracts import ILogger
 from src.application.domain.entities.organization import LegalEntityEntity
 from src.application.domain.exceptions import ApplicationException
+from src.application.domain.organization_documents import DOCUMENT_TYPE_TO_COLUMN, ORGANIZATION_DOCUMENT_TYPES
 from src.infrastructure.database.models import LegalEntityModel
 
 
@@ -19,6 +19,12 @@ class LegalEntityRepository(ILegalEntityRepository):
         self._session = session
         self._logger = logger
 
+    def _document_s3_keys_from_model(self, m: LegalEntityModel) -> dict[str, str | None]:
+        return {
+            document_type: getattr(m, DOCUMENT_TYPE_TO_COLUMN[document_type])
+            for document_type in ORGANIZATION_DOCUMENT_TYPES
+        }
+
     def _to_entity(self, m: LegalEntityModel) -> LegalEntityEntity:
         return LegalEntityEntity(
             id=m.id,
@@ -37,6 +43,7 @@ class LegalEntityRepository(ILegalEntityRepository):
             kyc_verified=m.kyc_verified,
             kyc_verified_at=m.kyc_verified_at,
             encrypted_mnemonic=m.encrypted_mnemonic,
+            document_s3_keys=self._document_s3_keys_from_model(m),
             created_by=m.created_by,
             created_at=m.created_at,
             updated_at=m.updated_at,
@@ -121,3 +128,19 @@ class LegalEntityRepository(ILegalEntityRepository):
             .values(encrypted_mnemonic=encrypted_mnemonic)
         )
         await self._session.flush()
+
+    async def get_document_s3_key(self, organization_id: str, document_type: str) -> str | None:
+        org = await self.get_by_id(organization_id)
+        return org.document_s3_keys.get(document_type)
+
+    async def set_document_s3_key(self, organization_id: str, document_type: str, s3_key: str) -> LegalEntityEntity:
+        column = DOCUMENT_TYPE_TO_COLUMN.get(document_type)
+        if column is None:
+            raise ApplicationException(status_code=400, message='Invalid document type')
+        await self._session.execute(
+            update(LegalEntityModel)
+            .where(LegalEntityModel.id == organization_id)
+            .values(**{column: s3_key})
+        )
+        await self._session.flush()
+        return await self.get_by_id(organization_id)