Using Enums in SQLModel with Postgres

enums_in_sqlmodel_with_postgres.md

Working with Enums in SQLModel and PostgreSQL: A Tasty Example with Ice Cream Flavors

Introduction:
Enums are a great way to represent a fixed set of values in your database. In this example, we'll model an IceCreamShop that serves various flavors of ice cream using SQLModel and PostgreSQL. We'll demonstrate how to store these flavors using enums and how to query for specific flavors using SQLAlchemy's powerful querying capabilities.

import os
from enum import Enum
from typing import List, Optional

from sqlmodel import Session, SQLModel, create_engine, select, Field, Enum as SQLModelEnum, ARRAY
from sqlalchemy import func

# Set up the database connection
PG_CONNECTION_STRING = os.environ["PG_CONNECTION_STRING"]
engine = create_engine(PG_CONNECTION_STRING)

# Define an enum for ice cream flavors
class FlavorEnum(str, Enum):
    VANILLA = "vanilla"
    CHOCOLATE = "chocolate"
    STRAWBERRY = "strawberry"
    ALL = "all"

# Define the IceCreamShop model using SQLModel
class IceCreamShop(SQLModel, table=True):
    id: int = Field(primary_key=True)
    flavors_served: Optional[List[FlavorEnum]] = Field(
        sa_type=ARRAY(SQLModelEnum(FlavorEnum, name="flavorenum")),
        description="The flavors that this shop serves. Your ONLY options are: 'vanilla', 'chocolate', 'strawberry', 'all'. Choose 'all' if all flavors are served.",
    )

# Function to query for a specific flavor using unnest
def query_for_flavor(session: Session):
    flavor_value = FlavorEnum.CHOCOLATE.name  # Get the name of the enum for comparison
    unnest_stmt = (
        select(func.unnest(IceCreamShop.flavors_served).label("flavor"))
        .where(IceCreamShop.id == 393)
        .subquery()
    )
    stmt = select(unnest_stmt.c.flavor).where(unnest_stmt.c.flavor == flavor_value)
    results = session.exec(stmt)
    return results.all()

# Example usage:
with Session(engine) as session:
    results = query_for_flavor(session)
    for result in results:
        print(result)

Learnings:

1. Enum Storage in PostgreSQL: When using SQLModel backed by PostgreSQL, enums are stored by their name (e.g., CHOCOLATE, VANILLA). This means you should use the name attribute of the enum when querying.
2. Querying Enums with SQLAlchemy: To filter based on an enum, use SQLAlchemy's powerful querying functions like func.unnest to expand arrays and filter by specific enum values.
3. Working with Arrays: PostgreSQL's ARRAY and unnest functions make it easy to store and query multiple enum values, such as the flavors an ice cream shop serves.

Conclusion:

Enums are a powerful tool for working with fixed sets of values in your database. By using SQLModel with PostgreSQL, you can easily define, store, and query enums, making your code more robust and easier to maintain. This example with IceCreamShop and FlavorEnum is just one of the many ways you can leverage enums in your database models.