sqlmodel icon indicating copy to clipboard operation
sqlmodel copied to clipboard

How to make enum column to work with SQLModel?

Open haja-k opened this issue 3 years ago • 11 comments

First Check

  • [X] I added a very descriptive title to this issue.
  • [X] I used the GitHub search to find a similar issue and didn't find it.
  • [X] I searched the SQLModel documentation, with the integrated search.
  • [X] I already searched in Google "How to X in SQLModel" and didn't find any information.
  • [X] I already read and followed all the tutorial in the docs and didn't find an answer.
  • [X] I already checked if it is not related to SQLModel but to Pydantic.
  • [X] I already checked if it is not related to SQLModel but to SQLAlchemy.

Commit to Help

  • [X] I commit to help with one of those options 👆

Example Code

from sqlmodel import SQLModel, Field, JSON, Enum, Column
from typing import Optional
from pydantic import BaseModel
from datetime import datetime

class TrainingStatus(str, enum.Enum):
    scheduled_for_training = "scheduled_for_training"
    training = "training"
    trained = "trained"

class model_profile(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    training_status: Column(Enum(TrainingStatus))
    model_version: str

Description

I am trying to create an enum column and use the value in the database. But I am getting this error:

RuntimeError: error checking inheritance of Column(None, Enum('scheduled_for_training', 'training', 'trained', name='trainingstatus'), table=None) (type: Column)

Does anyone knows how to help me? I tried

training_status: Column('value', Enum(TrainingStatus))

but it doesn't seem to work as I don't understand where the 'value' should be coming from 😓 I would really appreciate any input

Operating System

macOS

Operating System Details

No response

SQLModel Version

0.0.4

Python Version

3.7.9

Additional Context

No response

haja-k avatar Sep 14 '21 09:09 haja-k

training_status: TrainingStatus = Field(sa_column=Column(Enum(TrainingStatus)))

Just make sure that if you're using Alembic migration autogeneration and you require values to be stored as Enum in the database and not String, you modify the column type in the generated migration script.

yasamoka avatar Sep 16 '21 19:09 yasamoka

training_status: TrainingStatus = Field(sa_column=Column(Enum(TrainingStatus)))

Just make sure that if you're using Alembic migration autogeneration and you require values to be stored as Enum in the database and not String, you modify the column type in the generated migration script.

This worked for me. However, I think it would be a bit nicer if we could just specify

training_status: Enum[TrainingStatus]

or similar as the type.

JLHasson avatar Nov 20 '21 22:11 JLHasson

Thank you!

haja-k avatar Dec 08 '21 02:12 haja-k

training_status: TrainingStatus = Field(sa_column=Column(Enum(TrainingStatus)))

Just make sure that if you're using Alembic migration autogeneration and you require values to be stored as Enum in the database and not String, you modify the column type in the generated migration script.

Thanks @yasamoka This solution worked for me. Just one thing Enum on training_status: TrainingStatus = Field(sa_column=Column(Enum(TrainingStatus))) is imported from sqlmodel and Enum on TrainingStatus is imported from enum.Enum I lost some time until I realized this.

jonra1993 avatar May 17 '22 18:05 jonra1993

training_status: TrainingStatus = Field(sa_column=Column(Enum(TrainingStatus)))

Just make sure that if you're using Alembic migration autogeneration and you require values to be stored as Enum in the database and not String, you modify the column type in the generated migration script.

I do not get how I should "require values to be stored as Enum" when using postgresql engine for example. Is there a flag somewhere or how does that work?

I'm getting an error which looks like sqlmodel is trying to save the enum as a string:

sqlalchemy.exc.DataError: (psycopg2.errors.InvalidTextRepresentation) invalid input value for enum objecttypes: "type1"

alhoo avatar May 31 '22 08:05 alhoo

I do not get how I should "require values to be stored as Enum" when using postgresql engine for example. Is there a flag somewhere or how does that work?

I'm getting an error which looks like sqlmodel is trying to save the enum as a string:

sqlalchemy.exc.DataError: (psycopg2.errors.InvalidTextRepresentation) invalid input value for enum objecttypes: "type1"

Here is a sample migration from a project of mine that shows how this is done. This is a migration script that changes a column from String to Enum.

"""language_fluency

Revision ID: 5342a2459462
Revises: 919490ea9497
Create Date: 2021-09-17 03:49:14.251234

"""
from alembic import op
import sqlalchemy as sa
import sqlmodel
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision = '5342a2459462'
down_revision = '919490ea9497'
branch_labels = None
depends_on = None

fluency = sa.Enum('WEAK', 'MODERATE', 'GOOD', name='fluency')


def upgrade():
    fluency.create(op.get_bind())
    op.alter_column('language', 'fluency', type_=fluency, postgresql_using='fluency::text::fluency')


def downgrade():
    op.alter_column('language', 'fluency', type_=sqlmodel.sql.sqltypes.AutoString())
    fluency.drop(op.get_bind())

yasamoka avatar May 31 '22 12:05 yasamoka

do I really need a working alembic setup if I start from scratch? I can see that the table is created correctly but for some reason sqlmodel / psycopg2==2.9.3 are creating wrong queries. This is what the db looks like:

                    Table "public.items"
 Column |       Type        | Collation | Nullable | Default 
--------+-------------------+-----------+----------+---------
 type   | itemtypes         |           |          | 
 id     | uuid              |           | not null | 
 name   | character varying |           | not null | 

Does alembic change how sqlmodel/sqlalchemy generates statements and passes the parameters to that?

alhoo avatar May 31 '22 15:05 alhoo

What do you mean by starting from scratch? Do you mean that you have an existing database and you have yet to setup Alembic for migrations? If so, then no, you do not need a working alembic setup if you can manually define your column type correctly on the database schema itself directly.

Alembic doesn't change anything related to how SQLModel / SQLAlchemy generates statements. Alembic compares your existing database against the schema it reads from your code and generates schema migrations based on that. Regardless of whether those migrations are applied, it is SQLAlchemy that ultimately generates statements based on your models and tables as defined in your code. If there is still a mismatch between your actual database and the models / tables you have defined, then you get an error.

yasamoka avatar May 31 '22 16:05 yasamoka

From scratch I mean that I create an empty postgresql server and use the SQLModel.metadata.create_all(engine) to create tables. I got this working with mysql so I guess the main issue I'm having is that sqlmodel does not create the postgresql enum correctly. Postgres shows

docker compose exec db psql -U postgres -c "select enum_range(null::itemtypes);"
 enum_range 
------------
 {}
(1 row)

Maybe alembic works by correcting that enum.

alhoo avatar May 31 '22 17:05 alhoo

From scratch I mean that I create an empty postgresql server and use the SQLModel.metadata.create_all(engine) to create tables. I got this working with mysql so I guess the main issue I'm having is that sqlmodel does not create the postgresql enum correctly. Postgres shows

docker compose exec db psql -U postgres -c "select enum_range(null::itemtypes);"
 enum_range 
------------
 {}
(1 row)

Maybe alembic works by correcting that enum.

The Alembic migration I showed you was not autogenerated. It was written manually (alembic revision -m "message").

yasamoka avatar May 31 '22 17:05 yasamoka

Hello, @alhoo I had the same problem and I noticed that the error happens if the table you are adding the new enum column already exists. In my case, I deleted the table and created a new migration and it worked, new types were created. I am not sure why alembic does not create a new enum type automatically on existing tables. This is not the best solution but works. image

jonra1993 avatar Jun 02 '22 01:06 jonra1993

The following solution worked fine for me:

import enum

from sqlalchemy import (Column, Integer)
from sqlalchemy_utils import ChoiceType
from sqlmodel import SQLModel, Field

class LightType(enum.IntEnum):
    off = 0
    on = 1
    flashing = 2

class Events(SQLModel, table=True):
    green: LightType = Field(sa_column=Column(ChoiceType(LightType, impl=Integer()), nullable=False))

malf2 avatar Sep 22 '22 14:09 malf2

Thanks for the help here everyone! 👏 🙇

If that solves the original problem, then you can close this issue @haja-k ✔️

tiangolo avatar Nov 14 '22 09:11 tiangolo

Thanks for the help everyone :)

haja-k avatar Nov 15 '22 00:11 haja-k

@jonra1993 Using this approach how would you set a default value, either/both at the database level and sqlmodel level

EDIT: Got it, you can just set a default inside the Field(), ie, Field(default=SomeEnum.foo)

Pk13055 avatar Feb 23 '23 04:02 Pk13055

The solution provided here doesn't automatically create the enum type in case it is missing with postgresql, at least with my setting:

sqlmodel: 0.0.8
python: 3.10.11
macOS: Ventura 13.3.1 (22E261)

I am calling create manually (Enum(TrainingStatus).create(engine)) to fix this

turangojayev avatar Jul 26 '23 15:07 turangojayev

@turangojayev could you please provide more details for your proposed fix? E.g. when and where exactly do you call create?

christianholland avatar Jul 27 '23 18:07 christianholland

sorry for late response, well, I basically do


sa_TrainingStatus = Enum(TrainingStatus)

def init_db():
    sa_TrainingStatus.create(engine, checkfirst=True)
    SQLModel.metadata.create_all(engine)


@asynccontextmanager
async def lifespan(app: FastAPI):
    create_db_and_tables()
    yield


app = FastAPI(lifespan=lifespan)

turangojayev avatar Aug 17 '23 13:08 turangojayev

If you want enum types to be recognized by alembic, do the following. Installation

  SQLModel: 0.0.11 
  Python: 3.11

pip install alembic-postgresql-enum or poetry add alembic-postgresql-enum

Add import to envy.py import alembic_postgresql_enum

Now, when you update enum types, alembic will detect the changes and migrations will be automatically generated.

def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.sync_enum_values('public', 'channel', ['USSD', 'WHATSAPP'],
                        [('subscription', 'channel')],
                        enum_values_to_rename=[])
    # ### end Alembic commands ###


def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.sync_enum_values('public', 'channel', ['WHATSAPP'],
                    [('subscription', 'channel')],
                    enum_values_to_rename=[])
# ### end Alembic commands ###

tjeaneric avatar Nov 14 '23 10:11 tjeaneric

Can we change the scheme of where the Enum type is stored?

dipsx avatar Apr 04 '24 00:04 dipsx