README.md

SQLAlchemy & SQLModel table constraint demonstration

create env

python -m venv venv
source ./venv/bin/activate
pip install sqlmodel psycopg2-binary

start the database

docker-compose up

run the python files

python constraint_sqlalchemy.py

python constraint_sqlmodel.py

stop the database:

docker-compose down

This was tested on 2023/11/01 with python3.11 on MacOS (ARM) and the following package versions:

psycopg2-binary==2.9.9
pydantic==1.10.13
SQLAlchemy==1.4.50
sqlalchemy2-stubs==0.0.2a36
sqlmodel==0.0.11
typing_extensions==4.8.0

constraint_sqlalchemy.py

"""Example of creating multi-column constraints with SQLAlchemy. """
from sqlalchemy import Column, Integer, CheckConstraint
from sqlalchemy.engine import URL
from sqlalchemy.exc import IntegrityError
from sqlalchemy.future import create_engine
from sqlalchemy.orm import declarative_base, Session


Base = declarative_base()


class Record(Base):
    """Example record.

    Holds two numbers, `a` & `b`. Constraints: `a, b > 0`, `b > a`.
    """

    __tablename__ = "sqla_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),
    )


def main():
    """Demonstration with Postgres:
    - connect to DB
    - create table
    - add valid records
    - try adding invalid records: verify that they fail
    """
    engine = create_engine(
        URL.create(
            drivername="postgresql",
            username="hello_user",
            password="hello_pw",
            host="localhost",
            port=5432,
            database="hello",
        ),
        echo=True,
    )

    Base.metadata.drop_all(engine)
    Base.metadata.create_all(engine)

    # these are valid
    with Session(engine) as session:
        record0 = Record(a=1, b=2)
        record1 = Record(a=5, b=11)
        session.add_all([record0, record1])
        session.commit()

    # these are not valid
    engine.echo = False
    with Session(engine) as session:
        for a, b in [
            (None, 10),
            (10, None),
            (0, 1),
            (1, 0),
            (1, 1),
            (2, 2),
        ]:
            record = Record(a=a, b=b)
            try:
                session.add(record)
                session.commit()
            except IntegrityError:
                print(record, "failed - good")
                session.rollback()
            else:
                print(record, "succeeded - boo")


if __name__ == "__main__":
    main()

constraint_sqlmodel.py

"""Example of creating multi-column constraints with SQLModel."""
from sqlalchemy.engine import URL
from sqlalchemy.exc import IntegrityError
from sqlmodel import (
    SQLModel,
    Field,
    create_engine,
    Session,
    CheckConstraint,
    Column,
    Integer,
)


class Record(SQLModel, table=True):
    """Example record.

    Holds two numbers, `a` & `b`. Constraints: `a, b > 0`, `b > a`.
    """

    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),
    )


def main():
    """Demonstration with Postgres:
    - connect to DB
    - create table
    - add valid records
    - try adding invalid records: verify that they fail
    """
    engine = create_engine(
        URL.create(
            drivername="postgresql",
            username="hello_user",
            password="hello_pw",
            host="localhost",
            port=5432,
            database="hello",
        ),
        echo=True,
    )

    SQLModel.metadata.drop_all(engine)
    SQLModel.metadata.create_all(engine)

    # these are valid
    with Session(engine) as session:
        record0 = Record(a=1, b=2)
        record1 = Record(a=5, b=11)
        session.add_all([record0, record1])
        session.commit()
    engine.echo = False
    with Session(engine) as session:
        for a, b in [
            (None, 10),
            (10, None),
            (0, 1),
            (1, 0),
            (1, 1),
            (2, 2),
        ]:
            record = Record(a=a, b=b)
            try:
                session.add(record)
                session.commit()
            except IntegrityError:
                print(record, "failed - good")
                session.rollback()
            else:
                print(record, "succeeded - boo")


if __name__ == "__main__":
    main()

docker-compose.yaml

# Supporting Docker compose file: creates a postgres database on port 5432
version: '3.8'

services:

  postgres:
    image: postgres:15-alpine
    volumes:
      - postgresdata:/var/lib/postgresql/data/
    ports:
      - '5432:5432'
    environment:
      - POSTGRES_USER=hello_user
      - POSTGRES_PASSWORD=hello_pw
      - POSTGRES_DB=hello

volumes:
  postgresdata: