How to model enums backed by integers with sqlachemy?

Question

I am using sqlalchemy as a readable way to model my database, I'm only interested in generating the database definition for several engines from my model.

Some of the columns in my tables have type Enum, which works fine in engines such as MySQL, since it has native enum support. However for SQL Server, it generates the column as VARCHAR and sets a constraint to check the values are within the expected enum values I specify.

I'd like to replace this alternative with a numeric based fallback, so that the column type is actually numeric and the constraint checks the numeric values are within the range of the enum size (assumes sequential values starting with 0).

I have tried creating a TypeDecorator with Enum as impl, but this was not enough or I did not know how to make it work. I also tried to just copy the code for the Boolean type and mix it with the Enum type to create my own type, but it seems that database compiler support is required too.

Is there a way in which I can achieve this without having to patch sqlalchemy itself?

Note that I am not interested in querying the database with python, after it's generated, I'm done, so that might simplify, perhaps.

Answer 1

Here's what you need:

import sqlalchemy as sa

class IntEnum(sa.types.TypeDecorator):
    impl = sa.Integer
    def __init__(self, enumtype, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._enumtype = enumtype

    def process_bind_param(self, value, dialect):
        return value.value

    def process_result_value(self, value, dialect):
        return self._enumtype(value)

And then you use it like this:

from enum import Enum
from sqlalchemy.ext.declarative import declarative_base

class MyEnum(Enum):
    one = 1
    two = 2
    three = 3

engine = sa.create_engine('sqlite:///:memory:')
session = sa.orm.sessionmaker(bind=engine)()
Base = declarative_base()

class Stuff(Base):
    __tablename__ = 'stuff'

    id = sa.Column('id', sa.Integer, primary_key=True)
    thing = sa.Column('num', IntEnum(MyEnum))

Base.metadata.create_all(engine)

session.add(Stuff(thing=MyEnum.one))
session.add(Stuff(thing=MyEnum.two))
session.add(Stuff(thing=MyEnum.three))
session.commit()
engine.execute(sa.text('insert into stuff values(4, 2);'))

for thing in session.query(Stuff):
    print(thing.id, thing.thing)

Really your only problem is that impl needed to be a sa.Integer, as that's what's actually backing the enum, not enum.Enum.