Search code examples
pythonsqlalchemyfastapisqlmodel

How to define a CheckConstraint for an integer field in SQLModel


Assume we have such atable. How we can define a CheckConstraint for age?

from sqlmodel import SQLModel, Field, Column, String, Integer
from typing import Optional


class Hero(SQLModel, table=True):
    __tablename__ = "heroes"
    id: Optional[int] = Field(default=None, primary_key=True, nullable=False)
    name: str = Field(sa_column=Column("name", String(32), unique=True), nullable=False)
    age: int = Field(sa_column=Column("age", ???Integer(gt=0)???))

Solution

  • Since you are already using sa_column, you can construct your Column as you like. Meaning you just do Column("age", Integer, CheckConstraint("age>0")).

    Note that using sa_column will mean that the field ignores all other SQLAlchemy-related paramters including nullable, which means in your example the name field will actually end up being nullable because on a Column it is nullable=True by default. You need to specify nullable=False on the Column itself, not on the Field.

    For age you could also just pass the CheckConstraint as a positional argument via sa_column_args instead of using sa_column. The SQL type Integer will be assigned by SQLModel by default because of the int annotation.

    There is also no need for the name to be explicitly set as "name" on that Column. The attribute name is chosen by default.

    Lastly, you don't need nullable=False on a field that has primary_key=True. The latter always implies the former.

    Here is the code with my suggestions:

    from sqlmodel import SQLModel, Field, Column, String, CheckConstraint
    from typing import Optional
    
    
    class Hero(SQLModel, table=True):
        __tablename__ = "heroes"
        id: Optional[int] = Field(default=None, primary_key=True)
        name: str = Field(sa_column=Column(String(32), unique=True, nullable=False))
        age: int = Field(sa_column_args=(CheckConstraint("age>0"),))
    

    Resulting SQL:

    CREATE TABLE heroes (
        name VARCHAR(32) NOT NULL, 
        id INTEGER NOT NULL, 
        age INTEGER NOT NULL CHECK (age>0), 
        PRIMARY KEY (id), 
        UNIQUE (name)
    )