How to Use ENUM in SQLAlchemy

Introduction

Enumeration, or ENUM, is a data type that consists of a static set of predefined values. In SQLAlchemy, an ORM tool for SQL databases in Python, ENUMs add readability and enforce validity of column values, simplifying database integrity management. This tutorial covers using ENUMs in SQLAlchemy with clear, progressing code examples.

Before getting started, ensure you have installed Python and SQLAlchemy, and have a basic understanding of creating tables and models with SQLAlchemy.

Defining ENUMs using SQLAlchemy

To define an ENUM in SQLAlchemy, you use the Enum type. This can be imported from the sqlalchemy module. You can create an ENUM from a Python enumeration or by directly passing the possible string values it can hold.

First, we’ll define a simple table with an ENUM column using a Python enumeration.

from sqlalchemy import create_engine, Column, Integer, String, Enum
from sqlalchemy.ext.declarative import declarative_base
from enum import Enum as PyEnum

engine = create_engine('sqlite:///example.db')
Base = declarative_base()

class Mood(PyEnum):
    SAD = 'sad'
    HAPPY = 'happy'
    ANGRY = 'angry'

class Person(Base):
    __tablename__ = 'people'
    id = Column(Integer, primary_key=True)
    name = Column(String)
    mood = Column(Enum(Mood))

Base.metadata.create_all(engine)

In the code above, we defined a Person table with a mood column that can only be one of the values defined in the Mood enumeration.

Inserting and Retrieving ENUM Values

Inserting values into the Person table is straightforward. Here’s how to add entries.

from sqlalchemy.orm import sessionmaker

Session = sessionmaker(bind=engine)
session = Session()

new_person = Person(name='Alice', mood=Mood.HAPPY)
session.add(new_person)
session.commit()

# To retrieve the entry
person = session.query(Person).filter_by(name='Alice').first()
print(person.mood, person.mood.name)

When retrieved, SQLAlchemy converts ENUM columns back to their Python representation, in this case, an instance of the Mood enumeration. This allows you to use enum operations and comparisons natively in your code.

Using String-Based ENUMs

If you do not wish to use Python’s Enum class, you can define ENUMs directly as strings.

class Person(Base):
    __tablename__ = 'people'
    id = Column(Integer, primary_key=True)
    name = Column(String)
    mood = Column(Enum('sad', 'happy', 'angry'))

Base.metadata.create_all(engine)

This approach defines the mood column to accept string values ‘sad’, ‘happy’, or ‘angry’. It’s less robust because it lacks the full feature set of the Enum type (e.g., name and value attributes), but it’s acceptable for simpler cases.

ENUMs with Schema Migrations

Managing ENUM types with migrations requires specific attention. If you’re using libraries like Alembic to manage schema changes, changing the elements of an ENUM often requires an operation to alter the type explicitly.

For example, to add a ‘content’ mood to the ENUM:

from alembic import op
import sqlalchemy as sa

# Assuming the revision identifiers are already generated

def upgrade():
    op.alter_column('people', 'mood',
               existing_type=sa.Enum('sad', 'happy', 'angry', name='mood_enum'),
               type_=sa.Enum('sad', 'happy', 'angry', 'content', name='mood_enum'))

def downgrade():
    op.alter_column('people', 'mood',
               existing_type=sa.Enum('sad', 'happy', 'angry', 'content', name='mood_enum'),
               type_=sa.Enum('sad', 'happy', 'angry', name='mood_enum'))

When you add or remove values from an ENUM, you must define both the existing type and the new type in the migration scripts to change the ENUM definition properly.

Advanced SQLAlchemy ENUM practices

As your application grows, you may encounter more complex scenarios, like sharing ENUM definitions across multiple tables or handling NULL values.

Sharing ENUMs

By defining an ENUM at the module level, you can easily share it across multiple table definitions.

from sqlalchemy import MetaData

metadata = MetaData()
mood_enum = Enum('sad', 'happy', 'angry', metadata=metadata)

# Use the 'mood_enum' for multiple table definitions

...

Handling NULL Values

In some cases, you may want to allow an ENUM column to hold a NULL value. Ensure the nullable argument is set to True for the ENUM column definition.

mood = Column(Enum(Mood), nullable=True)

Note that this affects the integrity constraints of your database and should be used thoughtfully to maintain data consistency.

Conclusion

In conclusion, using ENUMs in SQLAlchemy can not only enhance the readability of your code but also bolster the integrity of your database. Move from defining to querying, updating, and handling schema migrations with ENUMs in your SQLAlchemy models to bring robustness to your applications database layer. Happy coding!