Multi-level Polymorphism in SQL Alchemy

Hello 👋, If you are here to visit, this might not be a nice read.

But if you are here searching for a solution for this title, I’ve got you covered.

Nowadays, there are LLMs like ChatGPT and Claude from where you can try to pull out the solution, but I faced this problem pre-LLM era, and I had to dig deep to find/create some solution for “Multi-level polymorphism” in SQLAlchemy.

Let’s get started. I will first explain the use case and then explain the code line-by-line. This can be a quick tutorial for you.

This all started when I was building a Budget module during my work. Picture this: say you are creating a Budget class. Budgets can be created using templates, so to reduce the code repetition, you make a BaseClass and create two polymorphisms: Budget and Budget template identified by type. i.e., type=BUDGET and type=TEMPLATE

class BaseBudget(DBModel): 
  id = ... 
  name = ... 
  start_date = ... 
  end_date = ... 
  type = ... 
  __mapper_args__ = { 
        "polymorphic_on": type 
  } 
 
class Budget(BaseBudget): 
  type = BUDGET 
 
  __mapper_args__ = { 
        "polymorphic_identity": type, 
    } 
 
class BudgetTemplate(BaseBudget): 
  type = TEMPLATE 
 
  __mapper_args__ = { 
          "polymorphic_identity": type, 
      }

All looks good until now.

Now, there are new requirements. The budget will have multiple types: personal, professional, some other, x, y, z, etc.

You cannot add another polymorphic identifier in mapper_args. I searched for many other solutions but found none that fit my needs. So, I came up with one myself.

I hacked the polymorphic identity a bit to send multiple values in there.

First, I added budget_type a field in BaseBudget and modified polymorphic_on to take two values.

class BaseBudget(DBModel): 
  id = ... 
  name = ... 
  start_date = ... 
  end_date = ... 
  type = ... 
  budget_type = fk_to_budget_type 
 
  __mapper_args__ = { 
        "polymorphic_on": cast(type, String) 
        + "_" 
        + cast(func.coalesce(budget_type, ""), String), 
    }

This way, BaseBudget’s polymorphic_on will be Budget_PERSONAL , Budget_X , Budget_Y , Budget_Z , Template_PERSONAL, Template_X ,And so on.

Budget Types are dynamic, so I need a BudgetType class to store them.

class BudgetType(DBModel): 
  name = ...

Now, it’s time to create that hack. I created this compose_identity class to create that value for our BaseBudget’s polymorphic_on mapper arg.

class compose_identity(str): 
 
    """Composed identity. 
 
    Using a string for easy conversion to a string SQL composition. 
    """ 
 
    _tup: Tuple[int, Optional[int]] 
 
    def __new__(cls, d1, d2=None): 
        if callable(d2): 
            d2 = d2() 
        self = super().__new__(cls, f"{d1}_{d2 or ''}") 
        self._tup = d1, d2 
        return self 
 
    def _as_tuple(self): 
        return self._tup

With this, when I create any new type of budget, let’s say PersonalBudget:

class PersonalBudget(Budget): 
    BUDGETTYPE = "Personal" 
 
    __mapper_args__ = { 
        "polymorphic_identity": compose_identity(BUDGET, BUDGETTYPE), 
    }

I can use compose_identity, which will create Budget_Personal The key for us.

The second part of my hack is adding an event listener on our BaseBudget class initialization. Here, I will hack into the argument to consider both type and budget_type while creating polymorphism.

@event.listens_for(BaseBudget, "init", propagate=True) 
def setup_poly(target, args, kw): 
    """receive new BaseBudget objects when they are constructed and 
    set polymorphic identity""" 
 
    # this is the compose_identity() object 
    composed_ident = inspect(target).mapper.polymorphic_identity 
 
    type, budget_type = composed_ident._as_tuple() 
    kw["type"] = type 
    if budget_type: 
        kw["budget_type"] = budget_type

This way, no matter what type and budget_type we create, we can reuse and extend the same BaseBudget properties. 🥳

Here is a full example:

from typing import Optional, Tuple 
 
from sqlalchemy import ( 
    Column, 
    ForeignKey, 
    Integer, 
    String, 
    case, 
    cast, 
    create_engine, 
    event, 
    func, 
    inspect, 
) 
from sqlalchemy.orm import Session, declarative_base, object_session, polymorphic_union 
 
Base = declarative_base() 
 
 
class compose_identity(str): 
 
    """describe a composed identity. 
 
    Using a string for easy conversion to a string SQL composition. 
    """ 
 
    _tup: Tuple[int, Optional[int]] 
 
    def __new__(cls, d1, d2=None): 
        if callable(d2): 
            d2 = d2() 
        self = super().__new__(cls, f"{d1}, {d2 or ''}") 
        self._tup = d1, d2 
        return self 
 
    def _as_tuple(self): 
        return self._tup 
 
 
class BaseBudget(Base): 
    __tablename__ = "budget" 
    id = Column(Integer, primary_key=True) 
 
    name = Column(String(50)) 
    budget_type = Column(String(50)) 
    type = Column(String(50)) 
 
    __mapper_args__ = { 
        "polymorphic_on": cast(type, String) 
        + ", " 
        + cast(func.coalesce(budget_type, ""), String), 
    } 
 
 
class BudgetType(Base): 
    __tablename__ = "budget_type" 
    id = Column(Integer, primary_key=True) 
    name = Column(String(50)) 
 
    def __repr__(self): 
        return f"{self.__class__.__name__}(id={self.id}, name={self.name!r})" 
 
 
class Budget(BaseBudget): 
    ... 
 
    budget_data = Column(String(50)) 
 
    __mapper_args__ = { 
        "polymorphic_identity": compose_identity("budget"), 
    } 
 
 
class XBudget(Budget): 
    BUDGETTYPE = "X" 
 
    __mapper_args__ = { 
        "polymorphic_identity": compose_identity("budget", BUDGETTYPE), 
    } 
 
 
class YBudget(Budget): 
    BUDGETTYPE = "Y" 
 
    __mapper_args__ = { 
        "polymorphic_identity": compose_identity("budget", BUDGETTYPE), 
    } 
 
 
class BudgetTemplate(BaseBudget): 
 
    template_data = Column(String(50)) 
 
    __mapper_args__ = { 
        "polymorphic_identity": compose_identity("template"), 
    } 
 
 
@event.listens_for(BaseBudget, "init", propagate=True) 
def _setup_poly(target, args, kw): 
    """receive new BaseBudget objects when they are constructed and 
    set polymorphic identity""" 
 
    # this is the compose_identity() object 
    composed_identity= inspect(target).mapper.polymorphic_identity 
 
    type, budget_type = composed_identity._as_tuple() 
    kw["type"] = type 
    if budget_type: 
        kw["budget_type"] = budget_type 
 
 
e = create_engine("sqlite://", echo=True) 
Base.metadata.create_all(e) 
s = Session(e) 
 
bt_x = BudgetType(name="X") 
bt_y = BudgetType(name="Y") 
 
x = XBudget(name="foo", budget_data="bar", budget_type=bt_x) 
y = YBudget(name="foo", budget_data="bar", budget_type=bt_y) 
 
x_another = XBudget(name="foo", budget_data="bar", budget_type=bt_x) 
 
t = BudgetTemplate(name="foo", template_data="bar") 
 
s.add(bt_x) 
s.add(bt_y) 
 
s.add(x) 
s.add(x_another) 
s.add(y) 
 
s.add(t) 
 
s.commit() 
s.close() 
 
assert len(s.query(XBudget).all()) == 2 
assert len(s.query(YBudget).all()) == 1 
 
assert len(s.query(BudgetTemplate).all()) == 1

I hope this helps. Happy coding! 🙋‍♂️