128

I'm reading about sqlalchemy and I saw following code:

employees_table = Table('employees', metadata,
    Column('employee_id', Integer, primary_key=True),
    Column('name', String(50)),
    Column('manager_data', String(50)),
    Column('engineer_info', String(50)),
    Column('type', String(20), nullable=False)
)

employee_mapper = mapper(Employee, employees_table, \
    polymorphic_on=employees_table.c.type, polymorphic_identity='employee')
manager_mapper = mapper(Manager, inherits=employee_mapper, polymorphic_identity='manager')
engineer_mapper = mapper(Engineer, inherits=employee_mapper, polymorphic_identity='engineer')

Should I make 'type' an int, with constants in a library? Or should I make just make type an enum?

Phonolog
  • 6,321
  • 3
  • 36
  • 64
Timmy
  • 12,468
  • 20
  • 77
  • 107

8 Answers8

211

Python's enumerated types are directly acceptable by the SQLAlchemy Enum type as of SQLAlchemy 1.1:

import enum
from sqlalchemy import Integer, Enum

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

class MyClass(Base):
    __tablename__ = 'some_table'
    id = Column(Integer, primary_key=True)
    value = Column(Enum(MyEnum))

Note that above, the string values "one", "two", "three" are persisted, not the integer values.

For older versions of SQLAlchemy, I wrote a post which creates its own Enumerated type (http://techspot.zzzeek.org/2011/01/14/the-enum-recipe/)

from sqlalchemy.types import SchemaType, TypeDecorator, Enum
from sqlalchemy import __version__
import re

if __version__ < '0.6.5':
    raise NotImplementedError("Version 0.6.5 or higher of SQLAlchemy is required.")

class EnumSymbol(object):
    """Define a fixed symbol tied to a parent class."""

    def __init__(self, cls_, name, value, description):
        self.cls_ = cls_
        self.name = name
        self.value = value
        self.description = description

    def __reduce__(self):
        """Allow unpickling to return the symbol 
        linked to the DeclEnum class."""
        return getattr, (self.cls_, self.name)

    def __iter__(self):
        return iter([self.value, self.description])

    def __repr__(self):
        return "<%s>" % self.name

class EnumMeta(type):
    """Generate new DeclEnum classes."""

    def __init__(cls, classname, bases, dict_):
        cls._reg = reg = cls._reg.copy()
        for k, v in dict_.items():
            if isinstance(v, tuple):
                sym = reg[v[0]] = EnumSymbol(cls, k, *v)
                setattr(cls, k, sym)
        return type.__init__(cls, classname, bases, dict_)

    def __iter__(cls):
        return iter(cls._reg.values())

class DeclEnum(object):
    """Declarative enumeration."""

    __metaclass__ = EnumMeta
    _reg = {}

    @classmethod
    def from_string(cls, value):
        try:
            return cls._reg[value]
        except KeyError:
            raise ValueError(
                    "Invalid value for %r: %r" % 
                    (cls.__name__, value)
                )

    @classmethod
    def values(cls):
        return cls._reg.keys()

    @classmethod
    def db_type(cls):
        return DeclEnumType(cls)

class DeclEnumType(SchemaType, TypeDecorator):
    def __init__(self, enum):
        self.enum = enum
        self.impl = Enum(
                        *enum.values(), 
                        name="ck%s" % re.sub(
                                    '([A-Z])', 
                                    lambda m:"_" + m.group(1).lower(), 
                                    enum.__name__)
                    )

    def _set_table(self, table, column):
        self.impl._set_table(table, column)

    def copy(self):
        return DeclEnumType(self.enum)

    def process_bind_param(self, value, dialect):
        if value is None:
            return None
        return value.value

    def process_result_value(self, value, dialect):
        if value is None:
            return None
        return self.enum.from_string(value.strip())

if __name__ == '__main__':
    from sqlalchemy.ext.declarative import declarative_base
    from sqlalchemy import Column, Integer, String, create_engine
    from sqlalchemy.orm import Session

    Base = declarative_base()

    class EmployeeType(DeclEnum):
        part_time = "P", "Part Time"
        full_time = "F", "Full Time"
        contractor = "C", "Contractor"

    class Employee(Base):
        __tablename__ = 'employee'

        id = Column(Integer, primary_key=True)
        name = Column(String(60), nullable=False)
        type = Column(EmployeeType.db_type())

        def __repr__(self):
             return "Employee(%r, %r)" % (self.name, self.type)

    e = create_engine('sqlite://', echo=True)
    Base.metadata.create_all(e)

    sess = Session(e)

    sess.add_all([
        Employee(name='e1', type=EmployeeType.full_time),
        Employee(name='e2', type=EmployeeType.full_time),
        Employee(name='e3', type=EmployeeType.part_time),
        Employee(name='e4', type=EmployeeType.contractor),
        Employee(name='e5', type=EmployeeType.contractor),
    ])
    sess.commit()

    print sess.query(Employee).filter_by(type=EmployeeType.contractor).all()
Mark Amery
  • 143,130
  • 81
  • 406
  • 459
zzzeek
  • 72,307
  • 23
  • 193
  • 185
  • Hm, link is still alive, but the runnable code posted there doesn't seem to work with recent versions of SQLA (1.0.9). Is this still your preferred method of handling enums, @zzzeek? – achiang Dec 08 '15 at 16:40
  • I just ran the script at http://techspot.zzzeek.org/files/2011/decl_enum.py exactly as is against master and it runs perfectly. – zzzeek Dec 08 '15 at 21:22
  • Hi @zzzeek, this is what I experienced: https://gist.github.com/achiang/8d4a6e3f495084d6761c – achiang Dec 14 '15 at 04:12
  • the example uses a python2-style metaclass so would need to be ported to Python 3 for its incompatible metaclass syntax. – zzzeek Dec 14 '15 at 14:22
  • @zzzeek would you mind pointing out the metaclass? I'm happy to port this to Python 3 – Anconia Nov 16 '16 at 14:02
  • Does this approach work at the DBMS level? I mean does, for example, MySQL respect this setting for the column, or is it only relevant when interacting with the database through SQLAlchemy? Thanks – pmsoltani Oct 29 '19 at 07:41
  • I have a Pyramid/Python 2 system that implements the `EnumSymbol` and the other classes. I'm currently porting the system to Python 3, and I get error `TypeError: unhashable type: 'EnumSymbol'` when querying some model that has a field type `EnumSymbol` that is a `primary_key`. For instance: `session.query(Model).all()`. This error doesn't occur in Python 2. Would have some complement to your code to handle this case? – Caco Jan 08 '21 at 19:39
  • In this example the response of `MyClass.query.filter_by(value=1).first()` would return a `value: MyEnum.one` for example. How do you marshal that to just be `value: 1`? I've tried setting my default to `MyEnum.one.value`, but it still returns `MyEnum.one` in the Sqlalchemy object – Guus May 25 '21 at 11:20
60

SQLAlchemy has an Enum type since 0.6: http://docs.sqlalchemy.org/en/latest/core/type_basics.html?highlight=enum#sqlalchemy.types.Enum

Although I would only recommend its usage if your database has a native enum type. Otherwise I would personally just use an int.

Steve Piercy
  • 13,693
  • 1
  • 44
  • 57
Wolph
  • 78,177
  • 11
  • 137
  • 148
  • second link is not opening anymore – MajesticRa Oct 01 '12 at 16:33
  • @MajesticRa: thanks for the tip, I've removed the link. It's not relevant anymore anyway. Everyone should have upgraded to 0.6 years ago ;) – Wolph Oct 01 '12 at 17:24
  • How do you give a default type? I'm trying to do `trading_mode = Column(Enum(TradingMode), default=TradingMode.SPOT)` but am getting `E TypeError: Object of type TradingMode is not JSON serializable ` – Sam Aug 03 '21 at 20:01
  • It depends on what kind of object `TradingMode.SPOT` is, but I would assume that a simple string-cast should do the trick: `str(TradingMode.SPOT)` – Wolph Aug 04 '21 at 13:51
  • You said you would use an int. You mean, you would store an int in database and replace it with python enum value dynamically? That gives more opportunities in refactoring it and adding values, doesn't it? Did I understand your thoughts well? – Anton Makarov Aug 30 '23 at 18:40
19

I'm not really knowledgeable in SQLAlchemy but this approach by Paulo seemed much simpler to me.
I didn't need user-friendly descriptions, so I went with it.

Quoting Paulo (I hope he doesn't mind my reposting it here):

Python’s namedtuple collection to the rescue. As the name implies, a namedtuple is a tuple with each item having a name. Like an ordinary tuple, the items are immutable. Unlike an ordinary tuple, an item’s value can be accessed through its name using the dot notation.

Here is a utility function for creating a namedtuple:

from collections import namedtuple

def create_named_tuple(*values):
     return namedtuple('NamedTuple', values)(*values)

The * before the values variable is for “unpacking” the items of the list so that each item is passed as an individual argument to the function.

To create a namedtuple, just invoke the above function with the needed values:

>>> project_version = create_named_tuple('alpha', 'beta', 'prod')
NamedTuple(alpha='alpha', beta='beta', prod='prod')

We can now use the project_version namedtuple to specify the values of the version field.

class Project(Base):
     ...
     version = Column(Enum(*project_version._asdict().values(), name='projects_version'))
     ...

This works great for me and is so much simpler than the other solutions that I previously found.

Dan Abramov
  • 264,556
  • 84
  • 409
  • 511
15

Note: the following is outdated. You should use sqlalchemy.types.Enum now, as recommended by Wolph. It's particularly nice as it complies with PEP-435 since SQLAlchemy 1.1.


I like zzzeek's recipe at http://techspot.zzzeek.org/2011/01/14/the-enum-recipe/, but I changed two things:

  • I'm using the Python name of the EnumSymbol also as the name in the database, instead of using its value. I think that's less confusing. Having a separate value is still useful, e.g. for creating popup menus in the UI. The description can be considered a longer version of the value that can be used e.g. for tooltips.
  • In the original recipe, the order of the EnumSymbols is arbitrary, both when you iterate over them in Python and also when you do an "order by" on the database. But often I want to have a determinate order. So I changed the order to be alphabetic if you set the attributes as strings or tuples, or the order in which the values are declared if you explicitly set the attributes as EnumSymbols - this is using the same trick as SQLAlchemy does when it orders the Columns in DeclarativeBase classes.

Examples:

class EmployeeType(DeclEnum):
    # order will be alphabetic: contractor, part_time, full_time
    full_time = "Full Time"
    part_time = "Part Time"
    contractor = "Contractor"

class EmployeeType(DeclEnum):
    # order will be as stated: full_time, part_time, contractor
    full_time = EnumSymbol("Full Time")
    part_time = EnumSymbol("Part Time")
    contractor = EnumSymbol("Contractor")

Here is the modified recipe; it uses the OrderedDict class available in Python 2.7:

import re

from sqlalchemy.types import SchemaType, TypeDecorator, Enum
from sqlalchemy.util import set_creation_order, OrderedDict


class EnumSymbol(object):
    """Define a fixed symbol tied to a parent class."""

    def __init__(self, value, description=None):
        self.value = value
        self.description = description
        set_creation_order(self)

    def bind(self, cls, name):
        """Bind symbol to a parent class."""
        self.cls = cls
        self.name = name
        setattr(cls, name, self)

    def __reduce__(self):
        """Allow unpickling to return the symbol linked to the DeclEnum class."""
        return getattr, (self.cls, self.name)

    def __iter__(self):
        return iter([self.value, self.description])

    def __repr__(self):
        return "<%s>" % self.name


class DeclEnumMeta(type):
    """Generate new DeclEnum classes."""

    def __init__(cls, classname, bases, dict_):
        reg = cls._reg = cls._reg.copy()
        for k in sorted(dict_):
            if k.startswith('__'):
                continue
            v = dict_[k]
            if isinstance(v, basestring):
                v = EnumSymbol(v)
            elif isinstance(v, tuple) and len(v) == 2:
                v = EnumSymbol(*v)
            if isinstance(v, EnumSymbol):
                v.bind(cls, k)
                reg[k] = v
        reg.sort(key=lambda k: reg[k]._creation_order)
        return type.__init__(cls, classname, bases, dict_)

    def __iter__(cls):
        return iter(cls._reg.values())


class DeclEnum(object):
    """Declarative enumeration.

    Attributes can be strings (used as values),
    or tuples (used as value, description) or EnumSymbols.
    If strings or tuples are used, order will be alphabetic,
    otherwise order will be as in the declaration.

    """

    __metaclass__ = DeclEnumMeta
    _reg = OrderedDict()

    @classmethod
    def names(cls):
        return cls._reg.keys()

    @classmethod
    def db_type(cls):
        return DeclEnumType(cls)


class DeclEnumType(SchemaType, TypeDecorator):
    """DeclEnum augmented so that it can persist to the database."""

    def __init__(self, enum):
        self.enum = enum
        self.impl = Enum(*enum.names(), name="ck%s" % re.sub(
            '([A-Z])', lambda m: '_' + m.group(1).lower(), enum.__name__))

    def _set_table(self, table, column):
        self.impl._set_table(table, column)

    def copy(self):
        return DeclEnumType(self.enum)

    def process_bind_param(self, value, dialect):
        if isinstance(value, EnumSymbol):
            value = value.name
        return value

    def process_result_value(self, value, dialect):
        if value is not None:
            return getattr(self.enum, value.strip())
Cito
  • 5,365
  • 28
  • 30
4

This and related StackOverflow thread answers resort to PostgreSQL or other dialect-specific typing. However, generic support may be easily achieved in SQLAlchemy that is also compatible with Alembic migrations.

If the backend doesn't support Enum, SQLAlchemy and alembic can facilitate enforcing constraints on varchar and similar types to mimic enumerated column types.

First, import the Python enum, the SQLAlchemy Enum, and your SQLAlchemy declarative base wherever you're going to declare your custom SQLAlchemy Enum column type.

import enum
from sqlalchemy import Enum
from sqlalchemy.ext.declarative import declarative_base
Base = declarative_base()

Let's take OP's original Python enumerated class:

class PostStatus(enum.Enum):
    DRAFT='draft'
    APPROVE='approve'
    PUBLISHED='published'

Now we create a SQLAlchemy Enum instantiation:

PostStatusType: Enum = Enum(
    PostStatus,
    name="post_status_type",
    create_constraint=True,
    metadata=Base.metadata,
    validate_strings=True,
)

When you run your Alembic alembic revision --autogenerate -m "Revision Notes" and try to apply the revision with alembic upgrade head, you'll likely get an error about the type not existing. For example:

...
sqlalchemy.exc.ProgrammingError: (psycopg2.errors.UndefinedObject) type "post_status_type" does not exist
LINE 10:  post_status post_status_type NOT NULL,
...

To fix this, import your SQLAlchemy Enum class and add the following to your upgrade() and downgrade() functions in the Alembic autogenerated revision script.

from myproject.database import PostStatusType
...
def upgrade() -> None:
    PostStatusType.create(op.get_bind(), checkfirst=True)
    ... the remainder of the autogen code...
def downgrade() -> None:
    ...the autogen code...
    PostStatusType.drop(op.get_bind(), checkfirst=True)

Finally, be sure to update the auto-generated sa.Column() declaration in the table(s) using the enumerated type to simply reference the SQLAlchemy Enum type instead of using Alembic's attempt to re-declare it. For example in def upgrade() -> None:

op.create_table(
    "my_table",
    sa.Column(
        "post_status",
        PostStatusType,
        nullable=False,
    ),
)
Brent
  • 1,195
  • 10
  • 9
  • This is the definitively the best answer as it properly shows how to initialize the enum in the database using Alembic – amos Mar 01 '23 at 15:57
3

For mysql i use its dialect

from sqlalchemy.dialects.mysql import ENUM

... 

class Videos(Base):
    ...
    video_type  = Column(ENUM('youtube', 'vimeo'))
    ...
Ricky Levi
  • 7,298
  • 1
  • 57
  • 65
2

I wanted something where in my business logic I was using only the Python enum.Enum and this would be automatically marshalled back and forth to/from the DB (PostgreSQL in my case). I did not want to have to worry about manually converting the Enum to some str representation in my business logic.

Here's an example:

from enum import Enum

from sqlalchemy import Column
from sqlalchemy.dialects.postgresql import ENUM, INTEGER
from sqlalchemy import types
from sqlalchemy.orm import declarative_base

Base = declarative_base()


class MarshalledENUM(types.TypeDecorator):
    # postgres dialect enum
    impl = ENUM

    def __init__(self, enum_class, *args, **kwargs):
        """
        Facilitates using Python Enum classes in the business logic while using
        the corresponding Postgres ENUM on the backend.

        NOTE: this assumes that the names and values are identical, e.g.:

            class SomeEnum(Enum):
                INITIALIZED = "INITIALIZED"

        NOTE: you may pass the `name` of the new PostgreSQL ENUM type via the kwargs when defining the `Column`.

        Args:
            enum_class: use this to pass in the Python Enum class you wish to marshal
        """

        self.enum_class = enum_class
        # this will construct the postgres dialect ENUM passing the raw string values enumerated
        super().__init__(*[e.value for e in enum_class], *args, **kwargs)

    def process_bind_param(self, value: Enum, dialect):
        if value is not None:
            # store the string value of the Enum instance in the DB
            return value.value

    def process_result_value(self, value, dialect):
        if value is not None:
            # get the Enum instance from the string value in the DB
            return getattr(self.enum_class, value)


class States(Enum):
    INITIALIZED = "INITIALIZED"
    STARTED = "STARTED"


class StateMachine(Base):
    __tablename__ = "state_machines"

    id = Column(INTEGER, primary_key=True)
    # relevant to the finite state machine
    state = Column(MarshalledENUM(States, name="states_enum"))

which you can use like so:

state_machine = StateMachine()
state_machine.state = States.INITIALIZED

To editorialize a bit, I have found a lot of very helpful use cases for using SQLAlchemy's TypeDecorator so I think it is reasonable be familiar with them and to lean on them a bit.

jsnow
  • 1,399
  • 1
  • 9
  • 7
-1

This is a method I use - using IntEnum

from enum import IntEnum
class GenderType(IntEnum):
    FEMALE:         int = 1
    MALE:           int = 2
    TRANSGENDER:    int = 3

class Citizen(Base):
    __tablename__ = 'citizen'

    user_uuid:      int = Column(UUID(as_uuid=True), primary_key=True)
    gender_type:    int = Column(Integer, nullable=False, default=GenderType.MALE)
    full_name:      str = Column(String(64))
    address:        str = Column(String(128))
Thinkal VB
  • 189
  • 3
  • 12