SQLModel cheat sheet

sqlmodel.py

"""
MIT License
Copyright (c) 202X Valentyn Solonechnyi
===========================================
"""

"""
Created/updated timestamp fields similar to 
Django's db models.DateTimeField(auto_now_add=True) 
and models.DateTimeField(auto_now=True)
"""
from sqlmodel import TIMESTAMP, Column, Field
from sqlalchemy.orm import declarative_mixin
from datetime import datetime, timezone
from pydantic import BaseModel

@declarative_mixin
class IDMixin(BaseModel):
    id: Optional[int] = Field(default=None, primary_key=True, nullable=False)


@declarative_mixin
class CreatedAtFieldsMixin(BaseModel):
    created_at: datetime = Field(default_factory=datetime.utcnow, nullable=False)


@declarative_mixin
class UpdatedAtFieldsMixin(BaseModel):
    updated_at: datetime | None = Field(
        sa_column=Column(
            TIMESTAMP,
            server_default=func.now(),
            onupdate=datetime.utcnow,
        ),
        nullable=False,
    )
      
class ExampleSQLTableModel(SQLModel, IDMixin, CreatedAtFieldsMixin, UpdatedAtFieldsMixin, table=True):
   __tablename__ = "demotable"
   ref_id: int
  
# To bo extended


"""
SQLModel ORM wrapper doesn't offer a simple solution for a field with predefined choices.
This short example demonstrates the integration of SQLAlchemy-Utils Choice Field with SQLModel Pydentic models
"""

from sqlmodel import Field, SQLModel
from pydantic import BaseModel, validator
from sqlalchemy_utils.types.choice import ChoiceType
from sqlalchemy_utils import Choice

# Define choices in tuples (Django way)

USER_ROLES = [
    (1, 'owner'),
    (2, 'administrator'),
    (3, 'commentator'),
    (4, 'employee'),
]

def transform_choice(obj: Choice) -> str:
    if isinstance(obj, Choice):
        return obj.value
    return obj
  
class UserBase(SQLModel):
  # store choice code as int in the db
  role: int = Field(sa_column=Column(ChoiceType(USER_ROLES, impl=Integer())))
    
    
class UserRead(SQLModel):
  # return choice value as a str
  role: str
    
  _extract_value = validator(
      'role', 
      # other fields with ChoiceType
      pre=True, 
      allow_reuse=True
  )(transform_choice)

  class Config:
      arbitrary_types_allowed = True