Creating dynamic classes in SQLAlchemy

Question

We have 1 table with a large amount of data and DBA's partitioned it based on a particular parameter. This means I ended up with Employee_TX, Employee_NY kind of table names. Earlier the models.py was simple as in --

class Employee(Base):
    __tablename__ = 'Employee'
    name = Column...
    state = Column...

Now, I don't want to create 50 new classes for the newly partitioned tables as anyways my columns are the same.

Is there a pattern where I can create a single class and then use it in query dynamically? session.query(<Tablename>).filter().all()

Maybe some kind of Factory pattern or something is what I'm looking for.

So far I've tried by running a loop as

for state in ['CA', 'TX', 'NY']:
    class Employee(Base):
        __qualname__ = __tablename__ = 'Employee_{}'.format(state)
        name = Column...
        state = Column...

but this doesn't work and I get a warning as - SAWarning: This declarative base already contains a class with the same class name and module name as app_models.employee, and will be replaced in the string-lookup table.

Also it can't find the generated class when I do from app_models import Employee_TX

This is a flask app with PostgreSQL as a backend and sqlalchemy is used as an ORM

Answer 1

Got it by creating a custom function like -

def get_model(state):
    DynamicBase = declarative_base(class_registry=dict())

    class MyModel(DynamicBase):
        __tablename__ = 'Employee_{}'.format(state)

        name = Column...         
        state = Column...

    return MyModel

And then from my services.py, I just call with get_model(TX)

Answer 2

Whenever you think of dynamically constructing classes think of type() with 3 arguments (see this answer for a demonstration, and the docs more generally).

In your case, it's just a matter of constructing the classes and keeping a reference to them so you can access them again later.

Here's an example:

from sqlalchemy import Column, Integer, String
from sqlalchemy.engine import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker


Base = declarative_base()


# this produces the set of common attributes that each class should have
def attribute_factory():
    return dict(
        id=Column(Integer, primary_key=True),
        name=Column(String, nullable=False),
        state=Column(String, nullable=False),
        CLASS_VAR=12345678,
    )


states = ["CA", "TX", "NY"]


# here we map the state abbreviation to the generated model, notice the templated
# class and table names
model_map = {
    state: type(
        f"Employee_{state}",
        (Base,),
        dict(**attribute_factory(), __tablename__=f"Employee_{state}"),
    )
    for state in states
}


engine = create_engine("sqlite:///", echo=True)

Session = sessionmaker(bind=engine)

Base.metadata.create_all(engine)


if __name__ == "__main__":
    # inserts work
    s = Session()
    for state, model in model_map.items():
        s.add(model(name="something", state=state))
    s.commit()
    s.close()

    # queries work
    s = Session()
    for state, model in model_map.items():
        inst = s.query(model).first()
        print(inst.state, inst.CLASS_VAR)