Change sqlmodel to sqlalchemy

Author: ChristopherSpelt

amt/api/routes/project.py

@@ -21,6 +21,7 @@ async def get_root(
     projects_service: Annotated[ProjectsService, Depends(ProjectsService)],
     tasks_service: Annotated[TasksService, Depends(TasksService)],
 ) -> HTMLResponse:
+    logger.info(f"getting project with id {project_id}")
     project = projects_service.get(project_id)
 
     context = {

amt/core/db.py

@@ -1,9 +1,11 @@
 import logging
 
+from sqlalchemy import create_engine, select
 from sqlalchemy.engine import Engine
-from sqlmodel import Session, SQLModel, create_engine, select
+from sqlalchemy.orm import Session
 
 from amt.core.config import get_settings
+from amt.models.base import Base
 
 logger = logging.getLogger(__name__)
 
@@ -22,7 +24,7 @@ def get_engine() -> Engine:
 def check_db() -> None:
     logger.info("Checking database connection")
     with Session(get_engine()) as session:
-        session.exec(select(1))
+        session.execute(select(1))
 
     logger.info("Finish Checking database connection")
 
@@ -32,6 +34,6 @@ def init_db() -> None:
 
     if get_settings().AUTO_CREATE_SCHEMA:  # pragma: no cover
         logger.info("Creating database schema")
-        SQLModel.metadata.create_all(get_engine())
+        Base.metadata.create_all(get_engine())
 
     logger.info("Finished initializing database")

amt/migrations/env.py

@@ -4,14 +4,14 @@
 from alembic import context
 from amt.models import *  # noqa
 from sqlalchemy import engine_from_config, pool
-from sqlmodel import SQLModel
+from sqlalchemy.schema import MetaData
 
 config = context.config
 
 if config.config_file_name is not None:
     fileConfig(config.config_file_name)
 
-target_metadata = SQLModel.metadata
+target_metadata = MetaData()
 
 
 def get_url() -> str:

amt/migrations/versions/9ce2341f2922_remove_the_status_table.py

@@ -9,7 +9,6 @@
 from collections.abc import Sequence
 
 import sqlalchemy as sa
-import sqlmodel.sql.sqltypes
 from alembic import op
 
 # revision identifiers, used by Alembic.
@@ -45,7 +44,7 @@ def downgrade() -> None:
     op.create_table(
         "status",
         sa.Column("id", sa.Integer(), nullable=False),
-        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("name", sa.String(255), nullable=False),
         sa.Column("sort_order", sa.Float(), nullable=False),
         sa.PrimaryKeyConstraint("id"),
     )

amt/migrations/versions/b62dbd9468e4_create_status_user_and_task_table.py

@@ -9,7 +9,6 @@
 from collections.abc import Sequence
 
 import sqlalchemy as sa
-import sqlmodel.sql.sqltypes
 from alembic import op
 
 # revision identifiers, used by Alembic.
@@ -24,23 +23,23 @@ def upgrade() -> None:
     status = op.create_table(
         "status",
         sa.Column("id", sa.Integer(), nullable=False),
-        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("name", sa.String(255), nullable=False),
         sa.Column("sort_order", sa.Float(), nullable=False),
         sa.PrimaryKeyConstraint("id"),
     )
 
     op.create_table(
         "user",
         sa.Column("id", sa.Integer(), nullable=False),
-        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
-        sa.Column("avatar", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+        sa.Column("name", sa.String(255), nullable=False),
+        sa.Column("avatar", sa.String(255), nullable=True),
         sa.PrimaryKeyConstraint("id"),
     )
     op.create_table(
         "task",
         sa.Column("id", sa.Integer(), nullable=False),
-        sa.Column("title", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
-        sa.Column("description", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("title", sa.String(255), nullable=False),
+        sa.Column("description", sa.String(255), nullable=False),
         sa.Column("sort_order", sa.Float(), nullable=False),
         sa.Column("status_id", sa.Integer(), nullable=True),
         sa.Column("user_id", sa.Integer(), nullable=True),

amt/migrations/versions/c5254dc6083f_add_project.py

@@ -9,7 +9,6 @@
 from collections.abc import Sequence
 
 import sqlalchemy as sa
-import sqlmodel.sql.sqltypes
 from alembic import op
 
 # revision identifiers, used by Alembic.
@@ -24,8 +23,8 @@ def upgrade() -> None:
     op.create_table(
         "project",
         sa.Column("id", sa.Integer(), nullable=False),
-        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(length=255), nullable=False),
-        sa.Column("model_card", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("name",sa.String(255), nullable=False),
+        sa.Column("model_card", sa.String(255), nullable=False),
         sa.PrimaryKeyConstraint("id"),
     )
     # ### end Alembic commands ###

amt/models/base.py

@@ -0,0 +1,5 @@
+from sqlalchemy.orm import DeclarativeBase
+
+
+class Base(DeclarativeBase):
+    pass

amt/models/project.py

@@ -1,21 +1,16 @@
-from pathlib import Path
 from typing import TypeVar
 
-from pydantic import field_validator
-from sqlmodel import Field, SQLModel  # pyright: ignore [reportUnknownVariableType]
+from sqlalchemy import String
+from sqlalchemy.orm import Mapped, mapped_column
 
-T = TypeVar("T", bound="Project")
+from amt.models.base import Base
 
+T = TypeVar("T", bound="Project")
 
-class Project(SQLModel, table=True):
-    id: int | None = Field(default=None, primary_key=True)
-    name: str = Field(max_length=255, min_items=3)
-    model_card: str | None = Field(description="Model card storage location", default=None)
 
-    @field_validator("model_card")
-    @classmethod
-    def validate_model_card(cls: type[T], model_card: str) -> str:
-        if not Path(model_card).is_file():
-            raise ValueError("Model card must be a file")
+class Project(Base):
+    __tablename__ = "project"
 
-        return model_card
+    id: Mapped[int] = mapped_column(primary_key=True)
+    name: Mapped[str] = mapped_column(String(255))  # TODO: (Christopher) how to set min_length?
+    model_card: Mapped[str | None] = mapped_column(default=None)

amt/models/task.py

@@ -1,16 +1,20 @@
-from sqlmodel import Field as SQLField  # pyright: ignore [reportUnknownVariableType]
-from sqlmodel import SQLModel
+from sqlalchemy import ForeignKey
+from sqlalchemy.orm import Mapped, mapped_column
 
+from amt.models.base import Base
 
-class Task(SQLModel, table=True):
-    id: int | None = SQLField(default=None, primary_key=True)
-    title: str
-    description: str
-    sort_order: float
-    status_id: int | None = SQLField(default=None)
-    user_id: int | None = SQLField(default=None, foreign_key="user.id")
+
+class Task(Base):
+    __tablename__ = "task"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    title: Mapped[str]
+    description: Mapped[str]
+    sort_order: Mapped[float]
+    status_id: Mapped[int | None] = mapped_column(default=None)
+    user_id: Mapped[int | None] = mapped_column(ForeignKey("user.id"))
     # TODO: (Christopher) SQLModel does not allow to give the below restraint an name
     #       which is needed for alembic. This results in changing the migration file
     #       manually to give the restrain a name.
-    project_id: int | None = SQLField(default=None, foreign_key="project.id")
+    project_id: Mapped[int | None] = mapped_column(ForeignKey("project.id"))
     # todo(robbert) Tasks probably are grouped (and sub-grouped), so we probably need a reference to a group_id

amt/models/user.py

@@ -1,7 +1,11 @@
-from sqlmodel import Field, SQLModel  # pyright: ignore [reportUnknownVariableType]
+from sqlalchemy.orm import Mapped, mapped_column
 
+from amt.models.base import Base
 
-class User(SQLModel, table=True):
-    id: int = Field(default=None, primary_key=True)
-    name: str
-    avatar: str | None
+
+class User(Base):
+    __tablename__ = "user"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    name: Mapped[str]
+    avatar: Mapped[str | None] = mapped_column(default=None)