PostgreSQL Enum Types with SQLModel and Alembic


While working on a product that uses FastAPI, SQLModel, Alembic, and PostgreSQL, I encountered a situation where I needed to add an enum column to an existing table. Since it took me some time to figure out the correct approach, I decided to document the process to help others who might face similar challenges.

Let’s start with a basic scenario. Assume you have a data model called Task as shown below:

import uuid
from datetime import datetime
from typing import Optional
from sqlmodel import SQLModel, Field

class Task(SQLModel, table=True):
    __tablename__ = "tasks"
    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    title: str = default=None
    description: str | None = Field(default=None)
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))

Using Alembic, you can generate the initial migration script with these commands:

alembic revision --autogenerate -m "Created task table"
alembic upgrade head

Now, let’s say you want to add a status field that should be an enum with two values – OPEN and CLOSED. First, define the enum class:

import enum
class TaskStatus(str, enum.Enum):
    OPEN = "open"
    CLOSED = "closed"

Then, add the status field to the Task class:

class Task(SQLModel, table=True):
    __tablename__ = "tasks"
    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    title: str = default=None
    description: str | None = Field(default=None)
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    status: Optional[TaskStatus] = None

If you run the Alembic migration commands at this point, it will define the status column as text. However, if you want to create a proper PostgreSQL enum type instead of storing the data as text, you’ll need to follow these additional steps:

  1. Install the alembic-postgresql-enum library:
pip install alembic-postgresql-enum

or if you’re using Poetry:

poetry add alembic-postgresql-enum
  1. Add the library import to your Alembic env.py file:
import alembic_postgresql_enum
  1. Modify the status field declaration in your Task class to explicitly use the enum type:
from sqlmodel import SQLModel, Field, Enum, Column


class Task(SQLModel, table=True):
    __tablename__ = "tasks"
    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    title: str = default=None
    description: str | None = Field(default=None)
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    status: Optional[TaskStatus] = Field(default=None, sa_column=Column(Enum(TaskStatus)))

Now you can run the Alembic commands to create a new PostgreSQL type for TaskStatus and use it for the column type:

alembic revision --autogenerate -m "Added status column in tasks table"
alembic upgrade head

To verify that the enum type was created correctly, connect to your PostgreSQL instance using psql and run the \dT+ command:

taskdb=# \dT+
                                          List of data types
 Schema |    Name     | Internal name | Size | Elements  |  Owner   | Access privileges | Description
--------+-------------+---------------+------+-----------+----------+-------------------+-------------
 public | taskstatus  | taskstatus    | 4    | OPEN +    | postgres |                   |
        |             |               |      | CLOSED   +|          |                   |

This approach ensures that your enum values are properly constrained at the database level, providing better data integrity than using a simple text field.

Discover more from Shekhar Gulati

Subscribe to get the latest posts sent to your email.

Leave a comment