如何使用sqlachemy为整数支持的枚举建模？

Question

我使用sqlalchemy作为模型化数据库的可读方式,我只对从我的模型生成多个引擎的数据库定义感兴趣.

我的表中的一些列具有类型Enum,在MySQL等引擎中工作正常,因为它具有本机枚举支持.但是对于SQL Server,它会生成列,VARCHAR并设置约束以检查值是否在我指定的预期枚举值范围内.

我想用基于数字的回退替换此替代方法,以便列类型实际上是数字,并且约束检查数值是否在枚举大小的范围内(假设以0开头的连续值).

我曾尝试创建一个TypeDecorator与Enum作为impl,但是这是不够的或我不知道如何使它发挥作用.我还尝试只复制布尔类型的代码并将其与Enum类型混合以创建我自己的类型,但似乎也需要数据库编译器支持.

有没有一种方法可以实现这一点,而无需修补sqlalchemy本身？

请注意,我对使用python查询数据库不感兴趣,在生成之后,我已经完成了,所以这可能会简化.

Answer 1

这就是你需要的:

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)

然后你像这样使用它:

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)

真的,你唯一的问题是impl需要成为一个sa.Integer,因为那是实际支持枚举的,而不是enum.Enum.