CheckConstraint on a table that takes a combination of Fields: is there a way to enforce one of two Fields not being null?

[charlie-corus]

First Check

• I added a very descriptive title here.
• I used the GitHub search to find a similar question and didn't find it.
• I searched the SQLModel documentation, with the integrated search.
• I already searched in Google "How to X in SQLModel" and didn't find any information.
• I already read and followed all the tutorial in the docs and didn't find an answer.
• I already checked if it is not related to SQLModel but to Pydantic.
• I already checked if it is not related to SQLModel but to SQLAlchemy.

Commit to Help

• I commit to help with one of those options 👆

Example Code

"""
This is my best attempt at a working example in SQLModel
"""

from typing import Optional

from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, CheckConstraint, column


class Hero(SQLModel, table=True):
    model_config = {"arbitrary_types_allowed": True}

    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    secret_name: str

    base_id: Optional[int] = Field(default=None, foreign_key="base.id")
    base: Optional["Base"] = Relationship()  #back_populates="hero")

    team_id: Optional[int] = Field(default=None, foreign_key="team.id")
    team: Optional["Team"] = Relationship() #back_populates="hero")

    """
    This fails with:
        CheckConstraint(team_id.isnot(None) | base_id.isnot(None),
                        ^^^^^^^^^^^^^
    AttributeError: 'FieldInfo' object has no attribute 'isnot'. Did you mean: 'init'?
    """
    __table_args__ = (
        CheckConstraint(team_id.isnot(None) | base_id.isnot(None),
                        name="A Hero needs a place to hang their cape")
    )


class Base(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    

class Team(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    
    base_id: int = Field(foreign_key="base.id")
    base: Base = Relationship()


solo_base = Base(name="The Lillypad")
team_base = Base(name="The Duck Pond")

pond_team = Team(name="Pond Crew", base=team_base)

solo_hero = Hero(name="Deadpond", 
              secret_name="Dive Wilson", 
              base=solo_base)

team_hero = Hero(name="Frogling",
              secret_name="Salvador von Schmuckle", 
              team=pond_team)

not_a_hero = Hero(name="The Vagrant",
                  secret_name="Carmichael Exogenes")

heroes = solo_hero, team_hero, not_a_hero

engine = create_engine("sqlite:///xor_sqlmodel_example.db")

SQLModel.metadata.clear()
SQLModel.metadata.create_all(engine)

with Session(engine) as session:
    for hero in heroes:
        session.add(hero)
    session.commit()
    
    for hero in heroes:
        session.refresh(hero)
        print(hero)

Description

I'd like to know if it's possible to set a CheckConstraint to a table that can apply some sort of combinatorial logic to multiple fields. Specifically, I'd like to enforce that exactly one of two fields is not null (or None).

Let's say I want to make sure that if a hero isn't in a team then they have a base of their own to call home. Equally, if they are in a team, they shouldn't be wasting valuable hero resources on a separate base.

In SQL this would look something like

CREATE TABLE Hero (
  PRIMARY KEY id NOT NULL,
  FOREIGN KEY base_id REFERENCES Base.id,
  FOREIGN KEY team_id REFERENCES Team.id,
  CHECK (base_id IS NOT NULL XOR team_id IS NOT NULL) -- This is where the magic happens
)

In SQLAlchemy it's possible to add a check constraint to the table. Here's a working example (with "base" renamed to "homebase" to avoid clashing with SQLAlchemy convention):

from sqlalchemy import Column, Integer, String, ForeignKey, CheckConstraint, create_engine
from sqlalchemy.orm import relationship, sessionmaker, declarative_base

Base = declarative_base()


class Homebase(Base):
    __tablename__ = 'homebase'
    id = Column(Integer, primary_key=True)
    name = Column(String)

class Team(Base):
    __tablename__ = 'team'
    id = Column(Integer, primary_key=True)
    name = Column(String)
    
    homebase_id = Column(Integer, ForeignKey('homebase.id'))
    homebase = relationship("Homebase")

class Hero(Base):
    __tablename__ = 'hero'
    id = Column(Integer, primary_key=True)
    name = Column(String)
    secret_name = Column(String)

    homebase_id = Column(Integer, ForeignKey('homebase.id'))
    homebase = relationship("Homebase")

    team_id = Column(Integer, ForeignKey('team.id'))
    team = relationship("Team")

    def __repr__(self):
        return f"Hero(name='{self.name}', secret_name='{self.secret_name}', homebase='{self.homebase.name if self.homebase else None}', team='{self.team.name if self.team else None}')"

    # This does the job!
    __table_args__ = (
        CheckConstraint(team_id.isnot(None) | homebase_id.isnot(None),
                        name="A Hero needs a place to hang their cape!"),
    )

engine = create_engine("sqlite:///xor_sqlalchemy.db")
Base.metadata.create_all(engine)


solo_base = Homebase(name="The Lillypad")
team_base = Homebase(name="The Duck Pond")

pond_team = Team(name="Pond Crew", homebase=team_base)

solo_hero = Hero(name="Deadpond", 
              secret_name="Dive Wilson", 
              homebase=solo_base)

team_hero = Hero(name="Frogling",
              secret_name="Salvador von Schmuckle", 
              team=pond_team)

# This hero fails the CheckConstraint
not_a_hero = Hero(name="The Vagrant",
                  secret_name="Carmichael Exogenes")

heroes = solo_hero, team_hero, not_a_hero

Session = sessionmaker()
Session.configure(bind=engine)

with Session() as session:
    for hero in heroes:
        session.add(hero)
    session.commit()
    
    for hero in heroes:
        session.refresh(hero)
        print(hero)

As expected, this fails with the helpful message sqlalchemy.exc.IntegrityError: (sqlite3.IntegrityError) CHECK constraint failed: A Hero needs a place to hang their cape!. Note that the failure is triggered from the backend, at commit time.

I see that SQLAlchemy.CheckConstraint is exposed by the SQLModel API, and a CheckConstraint can be passed straight to a Field constructor as of PR #436. This doesn't work for me - I need the CheckConstraint to take two Fields, so it needs to operate on the table. The approach that works in the SQLAlchemy ORM - creating a CheckConstraint in a field called __table_args__ - passes type and syntax checks in my editor, but fails at runtime because the Field objects don't seem to have access to the same logical operators that exist for a SQLAlchemy column (like isnot(); see Example Code above).

I can get SQLModel to create the constraint on the database by passing a string construct as follows:

    __table_args__ = (
        CheckConstraint(
            "base_id IS NOT NULL <> team_id IS NOT NULL",
            name='A Hero needs a place to hang their cape'
        ),
    )

This is less than ideal because the constraint is only applied in the backend when data is committed to the DB, so I wouldn't get to use it for validating inputs. I can't see a way to apply this sort of logic in a Pydantic model, but I'm very new to this whole ecosystem so apologies if I've overlooked something there.

Does SQLModel support, or plan to support, this sort of business? Do I have to restructure my database to fit into the ORM?

Operating System

macOS

Operating System Details

No response

SQLModel Version

0.0.14

Python Version

Python 3.12.1

Additional Context

No response

[charlie-corus]

One solution is to use three different models: a BasicHero, then SoloHero and TeamHero which both inherit from it:

class BasicHero(SQLModel):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    secret_name: str

class SoloHero(BasicHero, table=True):
    homebase_id = Column(Integer, ForeignKey('homebase.id'))
    homebase = relationship("Homebase")

class TeamHero(BasicHero, table=True):
    team_id = Column(Integer, ForeignKey('team.id'))
    team = relationship("Team")

The constraints can then be enforced separately. The problem with this is that I have to restructure my database. I guess there are areas where the ORM and relational way of thinking just differ, and we have to adapt our design patterns if we want to take advantage of the ORM niceness.

[charlie-corus]

OK I think I've found my solution, implementing custom validation logic:

from pydantic import ValidationError

class Hero(SQLModel, table=True):
    model_config = {"arbitrary_types_allowed": True}

    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    secret_name: str

    base_id: Optional[int] = Field(default=None, foreign_key="base.id")
    base: Optional["Base"] = Relationship()  #back_populates="hero")

    team_id: Optional[int] = Field(default=None, foreign_key="team.id")
    team: Optional["Team"] = Relationship() #back_populates="hero")

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.has_base()

    def update(self, **kwargs):
        super().update(**kwargs)
        self.has_base()

    def has_base(self):
        if not (self.base_id is not None) ^ (self.team_id is not None):
            raise ValidationError("A hero needs a place to hang their cape!")