✨ Adopt SQLModel, create models, start using it (fastapi#559)

* 🔥 Remove old SQLAlchemy models

* ✨ Add new SQLModel models

* 🔧 Update Alembic configs to work with SQLModel

* ✨ Re-generate initial Alembic migration

* 🔧 Update PostgreSQL driver connection string URL

* ✨ Create new SQLModel engine

* 🔥 Remove old unneeded SQLAlchemy-specific files

* ♻️ Update init_db

* ♻️ Use new SQLModel session

* ♻️ Update conftest with new DB Session

* ♻️ Update pre-start scripts to use SQLModel session

* ♻️ Import new SQLModel models

* ✨ Create new simplified create_user crud util

* ♻️ Update import in CRUDBase class (soon to be removed)

* 🙈 Update .gitignore with Python files

src/backend/.gitignore

@@ -1,2 +1,3 @@
 __pycache__
 app.egg-info
+*.pyc

src/backend/app/alembic/env.py

@@ -20,9 +20,9 @@
 # target_metadata = mymodel.Base.metadata
 # target_metadata = None
 
-from app.db.base import Base  # noqa
+from app.models import SQLModel  # noqa
 
-target_metadata = Base.metadata
+target_metadata = SQLModel.metadata
 
 # other values from the config, defined by the needs of env.py,
 # can be acquired:
@@ -35,7 +35,7 @@ def get_url():
     password = os.getenv("POSTGRES_PASSWORD", "")
     server = os.getenv("POSTGRES_SERVER", "db")
     db = os.getenv("POSTGRES_DB", "app")
-    return f"postgresql://{user}:{password}@{server}/{db}"
+    return f"postgresql+psycopg://{user}:{password}@{server}/{db}"
 
 
 def run_migrations_offline():

src/backend/app/alembic/script.py.mako

@@ -7,6 +7,7 @@ Create Date: ${create_date}
 """
 from alembic import op
 import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
 ${imports if imports else ""}
 
 # revision identifiers, used by Alembic.

src/backend/app/alembic/versions/e2412789c190_initialize_models.py

@@ -0,0 +1,48 @@
+"""Initialize models
+
+Revision ID: e2412789c190
+Revises: 
+Create Date: 2023-11-24 22:55:43.195942
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+
+
+# revision identifiers, used by Alembic.
+revision = 'e2412789c190'
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('user',
+    sa.Column('email', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+    sa.Column('is_active', sa.Boolean(), nullable=False),
+    sa.Column('is_superuser', sa.Boolean(), nullable=False),
+    sa.Column('full_name', sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('hashed_password', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+    sa.PrimaryKeyConstraint('id')
+    )
+    op.create_index(op.f('ix_user_email'), 'user', ['email'], unique=True)
+    op.create_table('item',
+    sa.Column('description', sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('title', sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+    sa.Column('owner_id', sa.Integer(), nullable=False),
+    sa.ForeignKeyConstraint(['owner_id'], ['user.id'], ),
+    sa.PrimaryKeyConstraint('id')
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table('item')
+    op.drop_index(op.f('ix_user_email'), table_name='user')
+    op.drop_table('user')
+    # ### end Alembic commands ###

src/backend/app/app/api/deps.py

@@ -9,19 +9,16 @@
 from app import crud, models, schemas
 from app.core import security
 from app.core.config import settings
-from app.db.session import SessionLocal
+from app.db.engine import engine
 
 reusable_oauth2 = OAuth2PasswordBearer(
     tokenUrl=f"{settings.API_V1_STR}/login/access-token"
 )
 
 
 def get_db() -> Generator:
-    try:
-        db = SessionLocal()
-        yield db
-    finally:
-        db.close()
+    with Session(engine) as session:
+        yield session
 
 
 def get_current_user(

src/backend/app/app/backend_pre_start.py

@@ -1,8 +1,9 @@
 import logging
 
+from sqlmodel import Session, select
 from tenacity import after_log, before_log, retry, stop_after_attempt, wait_fixed
 
-from app.db.session import SessionLocal
+from app.db.engine import engine
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -19,9 +20,9 @@
 )
 def init() -> None:
     try:
-        db = SessionLocal()
-        # Try to create session to check if DB is awake
-        db.execute("SELECT 1")
+        with Session(engine) as session:
+            # Try to create session to check if DB is awake
+            session.exec(select(1))
     except Exception as e:
         logger.error(e)
         raise e

src/backend/app/app/celeryworker_pre_start.py

@@ -1,8 +1,9 @@
 import logging
 
+from sqlmodel import Session, select
 from tenacity import after_log, before_log, retry, stop_after_attempt, wait_fixed
 
-from app.db.session import SessionLocal
+from app.db.engine import engine
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -20,8 +21,8 @@
 def init() -> None:
     try:
         # Try to create session to check if DB is awake
-        db = SessionLocal()
-        db.execute("SELECT 1")
+        with Session(engine) as session:
+            session.exec(select(1))
     except Exception as e:
         logger.error(e)
         raise e

src/backend/app/app/core/config.py

@@ -44,7 +44,7 @@ def assemble_db_connection(cls, v: Optional[str], values: Dict[str, Any]) -> Any
         if isinstance(v, str):
             return v
         return PostgresDsn.build(
-            scheme="postgresql",
+            scheme="postgresql+psycopg",
             user=values.get("POSTGRES_USER"),
             password=values.get("POSTGRES_PASSWORD"),
             host=values.get("POSTGRES_SERVER"),

src/backend/app/app/crud/__init__.py

@@ -8,3 +8,16 @@
 # from app.schemas.item import ItemCreate, ItemUpdate
 
 # item = CRUDBase[Item, ItemCreate, ItemUpdate](Item)
+from sqlmodel import Session
+from app.core.security import get_password_hash
+from app.models import UserCreate, User
+
+
+def create_user(session: Session, *, user_create: UserCreate) -> User:
+    db_obj = User.from_orm(
+        user_create, update={"hashed_password": get_password_hash(user_create.password)}
+    )
+    session.add(db_obj)
+    session.commit()
+    session.refresh(db_obj)
+    return db_obj

src/backend/app/app/crud/base.py

@@ -4,9 +4,7 @@
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
 
-from app.db.base_class import Base
-
-ModelType = TypeVar("ModelType", bound=Base)
+ModelType = TypeVar("ModelType", bound=Any)
 CreateSchemaType = TypeVar("CreateSchemaType", bound=BaseModel)
 UpdateSchemaType = TypeVar("UpdateSchemaType", bound=BaseModel)
 

src/backend/app/app/crud/crud_item.py

@@ -4,7 +4,7 @@
 from sqlalchemy.orm import Session
 
 from app.crud.base import CRUDBase
-from app.models.item import Item
+from app.models import Item
 from app.schemas.item import ItemCreate, ItemUpdate
 
 

src/backend/app/app/crud/crud_user.py

@@ -4,7 +4,7 @@
 
 from app.core.security import get_password_hash, verify_password
 from app.crud.base import CRUDBase
-from app.models.user import User
+from app.models import User
 from app.schemas.user import UserCreate, UserUpdate
 
 

src/backend/app/app/db/engine.py

@@ -0,0 +1,5 @@
+from sqlmodel import create_engine
+
+from app.core.config import settings
+
+engine = create_engine(settings.SQLALCHEMY_DATABASE_URI)

src/backend/app/app/db/init_db.py

@@ -1,25 +1,26 @@
-from sqlalchemy.orm import Session
+from sqlmodel import Session, select
 
-from app import crud, schemas
+from app import crud
 from app.core.config import settings
-from app.db import base  # noqa: F401
+from app.models import User, UserCreate  # noqa: F401
 
-# make sure all SQL Alchemy models are imported (app.db.base) before initializing DB
-# otherwise, SQL Alchemy might fail to initialize relationships properly
+# make sure all SQLModel models are imported (app.models) before initializing DB
+# otherwise, SQLModel might fail to initialize relationships properly
 # for more details: https://github.com/tiangolo/full-stack-fastapi-postgresql/issues/28
 
 
-def init_db(db: Session) -> None:
+def init_db(session: Session) -> None:
     # Tables should be created with Alembic migrations
     # But if you don't want to use migrations, create
     # the tables un-commenting the next line
     # Base.metadata.create_all(bind=engine)
-
-    user = crud.user.get_by_email(db, email=settings.FIRST_SUPERUSER)
+    user = session.exec(
+        select(User).where(User.email == settings.FIRST_SUPERUSER)
+    ).first()
     if not user:
-        user_in = schemas.UserCreate(
+        user_in = UserCreate(
             email=settings.FIRST_SUPERUSER,
             password=settings.FIRST_SUPERUSER_PASSWORD,
             is_superuser=True,
         )
-        user = crud.user.create(db, obj_in=user_in)  # noqa: F841
+        user = crud.create_user(session, user_create=user_in)

src/backend/app/app/initial_data.py

@@ -1,15 +1,17 @@
 import logging
 
+from sqlmodel import Session
+
+from app.db.engine import engine
 from app.db.init_db import init_db
-from app.db.session import SessionLocal
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
 def init() -> None:
-    db = SessionLocal()
-    init_db(db)
+    with Session(engine) as session:
+        init_db(session)
 
 
 def main() -> None: