Merge pull request #135 from dataforgoodfr/109_connect_dispatcher_to_db

feat: [DISPATCHER-109] - Update magnet link and status of rescues in database

Author: broija

dispatcher/api/models/logic.py

@@ -6,6 +6,9 @@
 import logging
 import uuid
 
+from sqlalchemy.orm import Session
+
+from rescue_api.models import Asset, Rescue, Rescuer
 from .payload import AssetModel
 
 # Configuration du logging
@@ -157,7 +160,95 @@ def allocate_assets(self, free_space_mb: float, node_id: str = None) -> Dict:
             "allocation_id": str(uuid.uuid4())
         }
 
-    def upsert_rescues(self, rescuer_id: int, assets: List[AssetModel]) -> Dict:
+
+    def upsert_rescues_to_db(self, rescuer_id: int, assets: List[AssetModel], db: Session) -> Dict:
+        if not self._rescuer_exists(rescuer_id=rescuer_id, db=db):
+            return {}
+
+        if not self._are_assets_data_consistent(assets=assets, db=db):
+            return {}
+
+        updated_rescues = []
+        inserted_rescues = []
+        not_committed_rescues = []
+
+        for asset in assets:
+            asset_id = int(asset.asset_id)
+            rescue = db.query(Rescue).filter(
+                (Rescue.rescuer_id == rescuer_id) & (Rescue.asset_id == asset_id)
+            ).first()
+
+            is_insertion_operation = False
+            if not rescue:
+                is_insertion_operation = True
+                rescue = Rescue(
+                    asset_id=asset_id,
+                    rescuer_id=rescuer_id,
+                    magnet_link=asset.magnet_link,
+                    status=asset.status.value.lower(),
+                )
+                db.add(rescue)
+            else:
+                rescue.magnet_link = asset.magnet_link
+                rescue.status = asset.status.value.lower()
+
+            try:
+                db.commit()
+            except Exception as e:
+                print(e)
+                print(f"Rescue with rescuer_id='{rescuer_id}' and asset_id='{asset_id}' has not been committed to DB.")
+                not_committed_rescues.append(
+                    {
+                        "asset_id": asset_id,
+                        "rescuer_id": rescuer_id,
+                        "magnet_link": asset.magnet_link,
+                        "status": asset.status.value.lower(),
+                    }
+                )
+            else:
+                committed_rescue = {
+                    "asset_id": asset_id,
+                    "rescuer_id": rescuer_id,
+                    "magnet_link": asset.magnet_link,
+                    "status": asset.status.value.lower(),
+                }
+                if is_insertion_operation:
+                    inserted_rescues.append(committed_rescue)
+                else:
+                    updated_rescues.append(committed_rescue)
+
+        return {
+            "updated_rescues": updated_rescues,
+            "inserted_rescues": inserted_rescues,
+            "not_committed_rescues": not_committed_rescues,
+        }
+
+
+    @staticmethod
+    def _rescuer_exists(rescuer_id: int, db: Session) -> bool:
+        rescuer = db.query(Rescuer).filter(Rescuer.id == rescuer_id).first()
+        return True if rescuer else False
+
+
+    @staticmethod
+    def _are_assets_data_consistent(assets: List[AssetModel], db: Session) -> bool:
+        missing_assets_count = 0
+        asset_inconsistencies_count = 0
+
+        for asset in assets:
+            db_asset = db.query(Asset).filter(Asset.id == int(asset.asset_id)).first()
+            if not db_asset:
+                missing_assets_count += 1
+            elif db_asset.url != asset.url:
+                asset_inconsistencies_count += 1
+
+        if missing_assets_count > 0 or asset_inconsistencies_count > 0:
+            return False
+
+        return True
+
+
+    def upsert_rescues_to_json(self, rescuer_id: int, assets: List[AssetModel]) -> Dict:
         rescues_to_upsert = self._prepare_rescues_to_upsert(rescuer_id=rescuer_id, assets=assets)
         rescues_from_db = self._load_json(self.rescues_file)
 

dispatcher/api/models/payload.py

@@ -48,7 +48,23 @@ class DispatchResponse(BaseModel):
     asset: List[AssetModel] = Field(..., description="Dataset list")
 
 
-class Rescues(BaseModel):
+class RescuesRequest(BaseModel):
     rescuer_id: int = Field(..., description="Rescuer id")
     message: str = Field(..., description="Message given by the application")
     assets: List[AssetModel] = Field(..., description="List of assets for which the rescue is over")
+
+
+class RescuesResponse(DispatchResponse):
+    asset: Optional[List[AssetModel]] = Field(default=None, description="Dataset list")
+    updated_rescues: Optional[List[Dict[str, Any]]] = Field(
+        default=None,
+        description="List of rescued assets that were updated in the database.",
+    )
+    inserted_rescues: Optional[List[Dict[str, Any]]] = Field(
+        default=None,
+        description="List of rescued assets that were inserted in the database.",
+    )
+    not_committed_rescues: Optional[List[Dict[str, Any]]] = Field(
+        default=None,
+        description="List of rescued assets for which an error happened when we tried to commit it to the database.",
+    )

dispatcher/api/routers/dispatch.py

@@ -1,10 +1,14 @@
-from models.state import app_state
-from fastapi import HTTPException, APIRouter
-from fastapi.responses import JSONResponse
-from models.payload import DispatchRequest, DispatchResponse, Rescues
 import json
 from os.path import join, dirname
 
+from models.payload import DispatchRequest, DispatchResponse, RescuesRequest, RescuesResponse
+from models.state import app_state
+from fastapi import HTTPException, APIRouter, Depends
+from fastapi.responses import JSONResponse
+from sqlalchemy.orm import Session
+
+from rescue_api.database import get_db
+
 router = APIRouter()
 
 @router.get('/')
@@ -55,29 +59,51 @@ async def dispatch(request: DispatchRequest):
 
     return dispatch_response
 
-@router.put('/rescues', response_model=DispatchResponse)
-async def update_rescues(rescues: Rescues):
-    app_state._logger.info("________Update rescues after asset downloads")
+@router.post('/assets-downloaded', response_model=RescuesResponse)
+async def upsert_rescues(request: RescuesRequest, db: Session = Depends(get_db)):
+    app_state._logger.info("________Upsert rescues after downloading assets")
+
+    if not request.assets:
+        raise HTTPException(
+            status_code=422,
+            detail="The rescues don't contain any asset, "
+                   "make sure they have at least one asset to be able to upsert them.",
+        )
 
     # Call dispatcher logic
-    result = app_state._dispatcher.upsert_rescues(
-        rescuer_id=rescues.rescuer_id,
-        assets=rescues.assets,
+    result = app_state._dispatcher.upsert_rescues_to_db(
+        rescuer_id=request.rescuer_id,
+        assets=request.assets,
+        db=db,
     )
     app_state._logger.info(result)
 
-    if result["action_status"] == "FAIL":
-        return HTTPException(
+    if not result:
+        raise HTTPException(
+            status_code=422,
+            detail="The error can be the following: the rescuer doesn't exist in the database; or at least one asset "
+                   "doesn't exist in the database; or the asset data you provided don't match the ones in the "
+                   "database. Make sure to provide an existing rescuer, existing assets and correct asset data.",
+        )
+    elif not result["updated_rescues"] and not result["inserted_rescues"]:
+        raise HTTPException(
             status_code=500,
-            detail="An error happened when saving the data, the rescues couldn't be updated, please retry later.",
+            detail={
+                "message": "Something went wrong when saving the data, the rescues below couldn't be upserted, "
+                           "please retry later.",
+                "not_committed_rescues": result["not_committed_rescues"],
+            },
         )
 
-    dispatch_response = DispatchResponse(
+    response = RescuesResponse(
         status="success",
         message="Request received and processed",
-        received_data=rescues.dict(),
-        asset=rescues.assets,
+        received_data=request.dict(),
+        asset=request.assets,
+        updated_rescues=result["updated_rescues"],
+        inserted_rescues=result["inserted_rescues"],
+        not_committed_rescues=result["not_committed_rescues"],
     )
-    print(f"Dispatch response: {dispatch_response}")
+    print(f"Dispatch response: {response}")
 
-    return dispatch_response
+    return response