Issue/9/insert cmodels in database

.gitignore

@@ -151,3 +151,6 @@ cython_debug/
 
 # Experiments
 experiments/
+
+# MacOS
+.DS_Store

main.py

@@ -1,11 +1,14 @@
 import uvicorn
+
 from src.config import settings
+from src.use_cases.cmodels import CmodelUseCases
 
-if __name__ == '__main__':
 
+if __name__ == '__main__':
+    CmodelUseCases.update_cmodels_collection()
     uvicorn.run(
         "src.api:app",
-        host=settings['HOST'],
+        host=settings.get('HOST'),
         port=int(settings['PORT']),
         reload=True,
         debug=True,

src/api.py

@@ -5,6 +5,7 @@
 from src.config import settings
 from src.routers.main import main_router
 
+
 app = FastAPI()
 
 

src/db/mongo.py

@@ -1,6 +1,8 @@
 from typing import Generator, Tuple
+
 from pymongo import MongoClient
 from pymongo.database import Database
+from pymongo.collection import Collection
 
 from src.config import db_config
 
@@ -45,3 +47,17 @@ def get_db(
     db: Database = db_connection[db_name]
 
     return db_connection, db
+
+
+def get_collection(
+    coll_name: str = db_config['CMODELS_COLL']
+) -> Tuple[MongoClient, Collection]:
+    """
+    Returns
+    -------
+    db_connection: pymongo.MongoClient
+    coll: pymongo.collection.Collection
+    """
+    db_connection, db = get_db()
+    coll: Collection = db[coll_name]
+    return db_connection, coll

src/interfaces/cmodels.py

@@ -0,0 +1,59 @@
+from datetime import datetime
+
+from pymongo.collection import Collection
+from pymongo.mongo_client import MongoClient
+
+from src.models.db.cmodels import (
+    CompartmentalModel,
+    CompartmentalModelEnum
+)
+from src.interfaces.crud import MongoCRUD
+
+
+class CModelsInterface:
+
+    def __init__(
+        self,
+        db_connection: MongoClient,
+        collection: Collection
+    ) -> None:
+        self.crud = MongoCRUD(db_connection, collection)
+
+    def insert_one_cmodel_document(self, model: CompartmentalModel):
+
+        cmodel_document = model.dict(by_alias=True)
+
+        existent_model = self.crud.read(model.id)
+
+        if existent_model:
+            pruned_existent_model = CModelsInterface._prune_db_document(
+                existent_model
+            )
+            pruned_current_model = CModelsInterface._prune_db_document(
+                cmodel_document
+            )
+            if pruned_existent_model == pruned_current_model:
+                # TODO: log cmodel exists f'Cmodel exists: {model.name}'
+                return existent_model
+            else:
+                model.updated_at = datetime.utcnow()
+                updated_model = self.crud.update(
+                    model.id,
+                    model.dict(by_alias=True)
+                )
+                # TODO log updated cmodel f'Updated cmodel: {model.name}'
+                return updated_model
+        else:
+            model_inserted = self.crud.insert(cmodel_document)
+            # TODO: log created cmodel f'Created cmodel: {model.name}'
+            return model_inserted
+
+    def insert_all_cmodel_documents(self):
+        return [self.insert_one_cmodel_document(model) for model in CompartmentalModelEnum.values()]
+
+    @staticmethod
+    def _prune_db_document(model_in_db: dict) -> dict:
+        if model_in_db:
+            model_in_db.pop('inserted_at')
+            model_in_db.pop('updated_at')
+            return model_in_db

src/interfaces/crud.py

@@ -0,0 +1,102 @@
+from typing import Any, Union
+from bson.objectid import ObjectId
+
+from pymongo.collection import Collection
+from pymongo.mongo_client import MongoClient
+
+
+class MongoCRUD():
+    def __init__(
+        self,
+        db_connection: MongoClient,
+        collection: Collection
+    ) -> None:
+        self.db_connection = db_connection
+        self.collection = collection
+
+    def insert(self, document: dict):
+        """
+        Parameters
+        ----------
+        model
+            The compartmental models that will be save in the database
+
+        Return
+        ----------
+        model: pymongo object
+        """
+        with self.db_connection:
+            try:
+                document['_id']
+            except KeyError:
+                raise ValueError('self.insert(): document must have key "_id"')
+
+            existent_document = self.collection.find_one(
+                self._id_to_dict(document['_id'])
+            )
+
+            if existent_document:
+                raise ValueError(
+                    f'Document with _id={document["_id"]} already exists in'
+                    'collection {self.collection}. if you want to change the'
+                    'fields\' values please use self.update'
+                )
+                # TODO: log
+
+            return self.collection.insert_one(document)
+
+    def read(self, _id: ObjectId) -> Union[Any, None]:
+        """Search for a specific model in ``self.collection``.
+
+        Parameters
+        ----------
+        query
+            Document's id. ``dict`` schema: ``{'_id': bson.ObjectID}``
+
+        Return
+        ----------
+        model: pymongo object
+            Object containing the results of the search
+        """
+        with self.db_connection:
+            return self.collection.find_one(self._id_to_dict(_id))
+
+    def update(self, _id: ObjectId, new_data: dict) -> bool:
+        """Update document in ``self.collection``.
+
+        Parameters
+        ----------
+        query
+            Document's id. ``dict`` schema: ``{'_id': bson.ObjectID}``
+        data
+            Updated document fields
+
+        Return
+        ----------
+        False:
+            * If query has no information
+            * If is not possible to update the document's status
+            * If the query id doesn't match the one associated to ``new_data``
+        True:
+            If the model has valid data and its status can be updated
+        """
+        _id_dict = self._id_to_dict(_id)
+        with self.db_connection:
+            if not _id_dict:
+                return False
+            cmodel = self.collection.find_one(_id_dict)
+            if cmodel:
+                update_model = self.collection.update_one(
+                    _id_dict,
+                    {"$set": new_data},
+                )
+                if update_model:
+                    return True
+            return False
+
+    def delete(self, _id: ObjectId):
+        with self.db_connection:
+            return self.collection.delete_one(self._id_to_dict(_id))
+
+    def _id_to_dict(self, _id: ObjectId):
+        return {'_id': _id} if _id else None

src/models/db/__init__.py

@@ -0,0 +1,13 @@
+from .cmodels import (
+    CompartmentalModelBase,
+    CompartmentalModel,
+    CModel,
+    AllCModels
+)
+
+__all__ = [
+    'CompartmentalModelBase',
+    'CompartmentalModel',
+    'CModel',
+    'AllCModels'
+]

src/models/db/base_model.py

@@ -0,0 +1,26 @@
+from datetime import datetime
+
+from pydantic import BaseModel, Field
+from bson.objectid import ObjectId as BsonObjectId
+
+
+class PydanticObjectId(BsonObjectId):
+    @classmethod
+    def __get_validators__(cls):
+        yield cls.validate
+
+    @classmethod
+    def validate(cls, v):
+        if not isinstance(v, BsonObjectId):
+            raise TypeError('ObjectId required')
+        return v
+
+
+class MetadataBaseDoc(BaseModel):
+    id: PydanticObjectId = Field(..., alias='_id')
+    inserted_at: datetime = datetime.utcnow()
+    updated_at: datetime = datetime.utcnow()
+
+    class Config:
+        allow_population_by_field_name = True
+        extra = 'forbid'

src/models/db/cmodels.py

@@ -0,0 +1,101 @@
+from enum import Enum
+from typing import Dict, List
+from bson.objectid import ObjectId
+
+from pydantic import BaseModel
+
+from .base_model import MetadataBaseDoc
+
+class CompartmentalModelBase(BaseModel):
+    """Base Model for Compartmental Models
+    """
+    name: str
+    """Name of the compartmental model"""
+    state_variables: List[str]
+    """Name of state variables of corresponding model"""
+    state_variables_units: Dict[str, str]
+    """Units of each state variable. The keys are the 'state_variables' array
+    elements
+    """
+    parameters: List[str]
+    """Parameters of the corresponding model"""
+    parameters_units: Dict[str, str]
+    """Units of each parameter. The keys are the 'parameters' array elements"""
+
+
+class CompartmentalModel(MetadataBaseDoc, CompartmentalModelBase):
+    pass
+
+
+class CompartmentalModelEnum(Enum):
+    """Compartmental Models' Data.
+
+    Each element of this ``Enum`` class contains all the essential information
+    on the corresponding Compartmental Model. Each element is a
+    :class:``CompartmentalModelBase`` object
+    """
+    sir: CompartmentalModelBase = CompartmentalModel(
+        id=ObjectId('6083175ea91f5aacea234423'),
+        name='SIR',
+        state_variables=['S', 'I', 'R'],
+        state_variables_units={
+            'S': 'persons',
+            'I': 'persons',
+            'R': 'persons',
+        },
+        parameters=['a', 'b'],
+        parameters_units={
+            'a': 'units of a',
+            'b': 'units of b',
+        },
+    )
+
+    seir: CompartmentalModelBase = CompartmentalModel(
+        id=ObjectId('6083176ca91f5aacea234424'),
+        name='SEIR',
+        state_variables=['S', 'E', 'I', 'R'],
+        state_variables_units={
+            'S': 'persons',
+            'E': 'persons',
+            'I': 'persons',
+            'R': 'persons',
+        },
+        parameters=['a', 'b'],
+        parameters_units={
+            'a': 'units of a',
+            'b': 'units of b',
+        },
+    )
+
+    seirv: CompartmentalModelBase = CompartmentalModel(
+        id=ObjectId('608317d0a91f5aacea234426'),
+        name='SEIRV',
+        state_variables=['S', 'E', 'I', 'R', 'V'],
+        state_variables_units={
+            'S': 'persons',
+            'E': 'persons',
+            'I': 'persons',
+            'R': 'persons',
+            'V': 'persons',
+        },
+        parameters=['a', 'b'],
+        parameters_units={
+            'a': 'units of a',
+            'b': 'units of b',
+            'c': 'units of c'
+        },
+    )
+
+    @classmethod
+    def values(cls) -> List[CompartmentalModel]:
+        return [m.value for m in cls]
+
+
+class CModel(BaseModel):
+    model: CompartmentalModelEnum
+
+
+class AllCModels(BaseModel):
+    models: List[CompartmentalModelEnum] = [
+        model.value for model in CompartmentalModelEnum
+    ]