Sqlalchemy database int to python enum

Question

I have lots of integers in my existing database which signify an enum. In my python project I have created the corresponding enums like so:

class Status(Enum):
    OK = 0
    PENDING = 1
    CANCELED = 2
    DUPLICATE = 3
    INCOMING = 4
    BLOCKED = 5

In my SQLalchemy models I'd like to be able to do the following

status = Column(Status, nullable=False)

Which should convert the database integers to python enums and back transparantly

So far I've tried to inherit from the sqlalchemy.types.TypeDecorator and overwrite the process_bind_param and process_result_value to cast the int back and forwards.

This would however mean I would have to repeat this for every Enum in violation of the DRY principle.

I would thus prefer a Mixin class which inherits sqlalchemy.types.TypeDecorator and overwrites the two methods to cast to and from the Enum which inherits this mixin.

Any ideas on how to approach this?

Answer 1

The TypeDecorator can simply take the enum type as an argument.

class EnumAsInteger(sqlalchemy.types.TypeDecorator):
    """Column type for storing Python enums in a database INTEGER column.

    This will behave erratically if a database value does not correspond to
    a known enum value.
    """
    impl = sqlalchemy.types.Integer # underlying database type

    def __init__(self, enum_type):
        super(EnumAsInteger, self).__init__()
        self.enum_type = enum_type

    def process_bind_param(self, value, dialect):
        if isinstance(value, self.enum_type):
            return value.value
        raise ValueError('expected %s value, got %s'
            % (self.enum_type.__name__, value.__class__.__name__))

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

    def copy(self, **kwargs):
        return EnumAsInteger(self.enum_type)


class MyTable(Base):
    status = sqlalchemy.Column(EnumAsInteger(Status), nullable=False)

One thing you may need to consider is how to deal with unknown enum values in the database. The above solution will simply throw ValueError when such values are seen, e.g. during a query.

Answer 2

You can use ChoiceType from SQLAlchemy-Utils.

status = Column(ChoiceType(Status, impl=db.Integer()), nullable=False)

This maps the column to an integer type on the database side.