Creating table constraints with SQLModel

(image from https://github.com/tiangolo/sqlmodel)

In this short article we will go through how to set up table constraints using SQLModel. This is not entirely out of the box right now, but can be achieved in a simple and logical way --- this article tells you how.

First, an example scenario with a minimal business logic example, then shortly explore how one can do this in SQL and SQLAlchemy and the steps for reproducing the same with SQLModel. Finally, a couple remarks and discussion with links for further reading.

If you are here for a TLDR version or have the same problem and just want to see the solution, scroll to the end.

SQLModel: ORM for FastAPI

SQLModel is a package providing ORM features and object serialisation functionality all in one, by the same author as FastAPI. Under the hood the package is using SQLAlchemy and Pydantic for the two respective functionalities, allowing for defining one set of models that translate to both, heavily utilising type annotations to achieve this.

Example scenarios & business logic

Consider any case with CHECK constraints, or UNIQUE constraints including more than one column. More specifically, we will consider a table with two integer columns: a & b, both positive and b greater than a.

Alternatively, you can implement the same for other concrete business logic, e.g.:

• two columns with only one of them allowed to be null
• a "unique together" constraint
• XOR relationship between conditions on two columns

All of the above you can implement with the solutions below. (exercise for the reader)

Photo by brooklyn on Unsplash

Plain SQL solution

This is here as a reference, we will be checking for this (or something isomorphic) being created by the ORM

CREATE TABLE record (
pk SERIAL NOT NULL,
a INTEGER NOT NULL,
b INTEGER NOT NULL,
PRIMARY KEY (pk),
CHECK (a > 0),
CHECK (b > a)
)

n.b. we only need 2 constraints: b > a & a > 0 implies b > 0

Important note: from now on we categorically dismiss the solution of writing pure SQL and passing that to the ORM or any constraint class. If you are using an ORM then that should be (and see that it is!) able to handle this. Not to mention the pitfalls when refactoring your code, e.g. not noticing that the column name is used in the constraints as well.

SQLAlchemy solution

With SQLAlchemy's ORM, you need to define the columns and add a constraint under __table_args__:

from sqlalchemy import Column, Integer, CheckConstraint  
from sqlalchemy.orm import declarative_base  
  
Base = declarative_base()  
  
  
class Record(Base):  
    __tablename__ = "record"  
    pk: int = Column(Integer, primary_key=True)  
  
    a: int = Column(Integer, nullable=False)  
    b: int = Column(Integer, nullable=False)  
  
    __table_args__ = (  
        CheckConstraint(a > 0),  
        CheckConstraint(b > a),  
    )

n.b. The above uses SQLAlchemy 1.4, the previous major version, but SQLModel is not supporting SQLAlchemy 2.0 yet so I have included this example. As of the time of writing this article, see SQLModel's roadmap for more detail.

SQLModel solution

The above table definition can be converted to SQLModel, let's do this step by step.

First, create the columns.

from sqlmodel import SQLModel, Field  
  
class Record(SQLModel, table=True):
    pk: int | None = Field(default=None, primary_key=True)  

    # only columns, unconstrained
    a: int  
    b: int

This satisfies that the two columns need to be filled, and that they are integers. None of the constraints are enforced yet.

Let's add the constraints!

Even though SQLModel's class definitions are very similar to those of SQlAlchemy, if one tries to just go for it then will encounter a couple problems. Let's note a couple gotchas before we look at the working solution:

1. Type annotations in class definitions don't initialise class members. So with the above you cannot use a & b as variables in the class body, unless you assign something to them. (You'd get a NameError if you used them inside the class body but outside of a method.. rookie mistake)
2. While SQLModel's fields look a lot like SQLAlchemy's columns, especially like version 2.0's column definitions, they are not. They are field definitions (FieldInfo objects later), meaning we need to access the associated column for applying constraints on.
3. As such, in the body of the table class the members need to be initialised as Fields (from SQLModel), which gives you a FieldInfo object that stores the associated SQLAlchemy column in it's .sa_column attribute. This is only filled in (at this stage) if you pass a Column to the Field constructor, otherwise it's pydantic's UndefinedType object for the time being. The column definition needs to include the type and other settings of the column here, so you need to define that twice, but on the same line at least.

In a nutshell: we need to initialise the fields involved in the constraints with a Column attached, and use the .sa_column member of theirs for our constraints and voilà.

Let's see the end result:

from sqlmodel import SQLModel, Field, CheckConstraint, Column, Integer
  
class Record(SQLModel, table=True):  
    pk: int | None = Field(default=None, primary_key=True)  
  
    a: int = Field(sa_column=Column(Integer, nullable=False))  
    b: int = Field(sa_column=Column(Integer, nullable=False))  
  
    __table_args__ = (  
        CheckConstraint(a.sa_column > 0),  
        CheckConstraint(b.sa_column > a.sa_column),  
    )

This produces the following SQL when executed (execute it with an engine with echo=True or inspect the database)

CREATE TABLE record (
 a INTEGER NOT NULL, 
 b INTEGER NOT NULL, 
 pk SERIAL NOT NULL, 
 PRIMARY KEY (pk), 
 CHECK (a > 0), 
 CHECK (b > a)
)

We can verify this with a couple test cases:

Record(a=1, b=2)
Record(a=5, b=11)

can be created, the database will be happy. (tested with PostgreSQL)

The following cases will not be accepted by the database, as expected:

Record(a=None, b=10)  # a is None
Record(a=10, b=None)  # b is NOne
Record(a=0, b=1)  # a is 0
Record(a=1, b=0)  # b is 0
Record(a=1, b=1)  # a == b

Recap of the key points

• Initialise the variables used in constraints
• Specify properties of the column associated with the field
• Use the .sa_column attribute of the fields in the constraint definitions

Further remarks

Looking in the source code of SQLModel you will find that it's exposing a large number of classes & functions from SQLAlchemy and adds type annotations or reasonable defaults to them. The constraints are only exposed, they are not overwritten at all to support the internal field definitions. Perhaps something to be implemented later on?

Discussion

SQLModel is very useful in what it's trying to do, and reduces code duplication or need for other tooling for FastAPI applications using a SQL backend. While not entirely without pain, the framework is logical and allows for fulfilling commonly encountered requirements like constraints involving more than one column.

A question for the future is how these frameworks will evolve and expand. SQLModel is not providing the seamless unity of SQLAlchemy and pydantic yet, but is on a good track. As it seems from the documentation (as of 2023 Oct) there are some missing pages about more advanced features, which are yet to be filled out. Some developments and additions may invalidate the current article as well, so if you are reading this further down the line do take a look and check for relevance.

Code availability

See the supporting files and full working examples with SQLAlchemy and SQLModel as well on Github Gist.