✨ 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

Author: tiangolo

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/d4867f3a4c0a_first_revision.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/alembic/versions/e2412789c190_initialize_models.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/api/deps.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/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__)
@@ -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/celeryworker_pre_start.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/core/config.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