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:
- Install the
alembic-postgresql-enum library:
pip install alembic-postgresql-enum
or if you’re using Poetry:
poetry add alembic-postgresql-enum
- Add the library import to your Alembic
env.py file:
import alembic_postgresql_enum
- 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.
