Create base model for Mongo schemas

alembic/env.py

@@ -5,7 +5,7 @@ from sqlalchemy.ext.asyncio import AsyncEngine
 
 from alembic import context
 from database.engine import postgres_engine
-from database.models import Base
+from database.schemas.relational import Base
 
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.

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

didier/data/embeds/ufora/announcements.py

@@ -13,7 +13,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 
 import settings
 from database.crud import ufora_announcements as crud
-from database.models import UforaCourse
+from database.schemas.relational import UforaCourse
 from didier.data.embeds.base import EmbedBaseModel
 from didier.utils.types.datetime import int_to_weekday
 from didier.utils.types.string import leading

tests/test_database/conftest.py

@@ -3,7 +3,13 @@ import datetime
 import pytest
 
 from database.crud import users
-from database.models import Bank, UforaAnnouncement, UforaCourse, UforaCourseAlias, User
+from database.schemas.relational import (
+    Bank,
+    UforaAnnouncement,
+    UforaCourse,
+    UforaCourseAlias,
+    User,
+)
 
 
 @pytest.fixture(scope="session")

tests/test_database/test_crud/test_birthdays.py

@@ -4,7 +4,7 @@ from freezegun import freeze_time
 
 from database.crud import birthdays as crud
 from database.crud import users
-from database.models import User
+from database.schemas.relational import User
 
 
 async def test_add_birthday_not_present(postgres, user: User):

tests/test_database/test_crud/test_currency.py

@@ -5,7 +5,7 @@ from freezegun import freeze_time
 
 from database.crud import currency as crud
 from database.exceptions import currency as exceptions
-from database.models import Bank
+from database.schemas.relational import Bank
 
 
 async def test_add_dinks(postgres, bank: Bank):

tests/test_database/test_crud/test_custom_commands.py

@@ -4,7 +4,7 @@ from sqlalchemy import select
 from database.crud import custom_commands as crud
 from database.exceptions.constraints import DuplicateInsertException
 from database.exceptions.not_found import NoResultFoundException
-from database.models import CustomCommand
+from database.schemas.relational import CustomCommand
 
 
 async def test_create_command_non_existing(postgres):

tests/test_database/test_crud/test_dad_jokes.py

@@ -1,7 +1,7 @@
 from sqlalchemy import select
 
 from database.crud import dad_jokes as crud
-from database.models import DadJoke
+from database.schemas.relational import DadJoke
 
 
 async def test_add_dad_joke(postgres):

tests/test_database/test_crud/test_tasks.py

@@ -6,7 +6,7 @@ from sqlalchemy import select
 
 from database.crud import tasks as crud
 from database.enums import TaskType
-from database.models import Task
+from database.schemas.relational import Task
 
 
 @pytest.fixture

tests/test_database/test_crud/test_ufora_announcements.py

@@ -1,7 +1,7 @@
 import datetime
 
 from database.crud import ufora_announcements as crud
-from database.models import UforaAnnouncement, UforaCourse
+from database.schemas.relational import UforaAnnouncement, UforaCourse
 
 
 async def test_get_courses_with_announcements_none(postgres):

tests/test_database/test_crud/test_ufora_courses.py

@@ -1,5 +1,5 @@
 from database.crud import ufora_courses as crud
-from database.models import UforaCourse
+from database.schemas.relational import UforaCourse
 
 
 async def test_get_course_by_name_exact(postgres, ufora_course: UforaCourse):

tests/test_database/test_crud/test_users.py

@@ -1,7 +1,7 @@
 from sqlalchemy import select
 
 from database.crud import users as crud
-from database.models import User
+from database.schemas.relational import User
 
 
 async def test_get_or_add_non_existing(postgres):

tests/test_database/test_utils/test_caches.py

@@ -1,4 +1,4 @@
-from database.models import UforaCourse
+from database.schemas.relational import UforaCourse
 from database.utils.caches import UforaCourseCache