Create base model for Mongo schemas

database/crud/birthdays.py

@@ -7,7 +7,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import selectinload
 
 from database.crud import users
-from database.models import Birthday, User
+from database.schemas.relational import Birthday, User
 
 __all__ = ["add_birthday", "get_birthday_for_user", "get_birthdays_on_day"]
 

database/crud/currency.py

@@ -5,7 +5,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 
 from database.crud import users
 from database.exceptions import currency as exceptions
-from database.models import Bank, NightlyData
+from database.schemas.relational import Bank, NightlyData
 from database.utils.math.currency import (
     capacity_upgrade_price,
     interest_upgrade_price,

database/crud/custom_commands.py

@@ -5,7 +5,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 
 from database.exceptions.constraints import DuplicateInsertException
 from database.exceptions.not_found import NoResultFoundException
-from database.models import CustomCommand, CustomCommandAlias
+from database.schemas.relational import CustomCommand, CustomCommandAlias
 
 __all__ = [
     "clean_name",

database/crud/dad_jokes.py

@@ -2,7 +2,7 @@ from sqlalchemy import func, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from database.exceptions.not_found import NoResultFoundException
-from database.models import DadJoke
+from database.schemas.relational import DadJoke
 
 __all__ = ["add_dad_joke", "get_random_dad_joke"]
 

database/crud/tasks.py

@@ -5,7 +5,7 @@ from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from database.enums import TaskType
-from database.models import Task
+from database.schemas.relational import Task
 from database.utils.datetime import LOCAL_TIMEZONE
 
 __all__ = ["get_task_by_enum", "set_last_task_execution_time"]

database/crud/ufora_announcements.py

@@ -3,7 +3,7 @@ import datetime
 from sqlalchemy import delete, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from database.models import UforaAnnouncement, UforaCourse
+from database.schemas.relational import UforaAnnouncement, UforaCourse
 
 __all__ = ["create_new_announcement", "get_courses_with_announcements", "remove_old_announcements"]
 

database/crud/ufora_courses.py

@@ -3,7 +3,7 @@ from typing import Optional
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from database.models import UforaCourse, UforaCourseAlias
+from database.schemas.relational import UforaCourse, UforaCourseAlias
 
 __all__ = ["get_all_courses", "get_course_by_name"]
 

database/crud/users.py

@@ -3,7 +3,7 @@ from typing import Optional
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from database.models import Bank, NightlyData, User
+from database.schemas.relational import Bank, NightlyData, User
 
 __all__ = [
     "get_or_add",

database/schemas/mongo.py

@@ -0,0 +1,38 @@
+from bson import ObjectId
+from pydantic import BaseModel, Field
+
+__all__ = []
+
+
+class PyObjectId(str):
+    """Custom type for bson ObjectIds"""
+
+    @classmethod
+    def __get_validators__(cls):
+        yield cls.validate
+
+    @classmethod
+    def validate(cls, value: str):
+        """Check that a string is a valid bson ObjectId"""
+        if not ObjectId.is_valid(value):
+            raise ValueError(f"Invalid ObjectId: '{value}'")
+
+        return ObjectId(value)
+
+    @classmethod
+    def __modify_schema__(cls, field_schema: dict):
+        field_schema.update(type="string")
+
+
+class MongoBase(BaseModel):
+    """Base model that properly sets the _id field, and adds one by default"""
+
+    id: PyObjectId = Field(default_factory=PyObjectId, alias="_id")
+
+    class Config:
+        """Configuration for encoding and construction"""
+
+        allow_population_by_field_name = True
+        arbitrary_types_allowed = True
+        json_encoders = {ObjectId: str, PyObjectId: str}
+        use_enum_values = True