✨ Allow to use Enums in custom DB schema

[KrilleGH]

I have a legacy Postgres DB which is receiving an API including new tables in a separate DB schema. The new code must not touch the existing objects in the public schema.

So all models derive from

class SchemaSQLModel(SQLModel):
    __table_args__ = {
        'schema': 'api',
    }

I have some fields which use Enums:

class HeroType(enum.StrEnum):
    normal = enum.auto()
    super = enum.auto()

class Hero(SchemaSQLModel):
    type: HeroType

The DB is PostgreSQL and I'm using the native Enum type (CREATE TYPE). Also these types reside in the custom schema.

When creating a new Hero the generated SQL has explicit casts to the types of the columns. However, for the type column only the plain name of the type is generated. Of course that doesn't work, as the Enum type isn't in the public schema.

The solution is to pass inherit_schema=True when instantiating SQLAlchemy's Enum class. I achived this by making a subclass and using that in get_sqlalchemy_type() instead of the plain one.

I'm aware this PR may not be in it's final form. Please advise what needs to be done to get it merged.

[KrilleGH]

Hey @svlandeg, I've rebased the PR and all checks are green now. I'd be ready for review.

[Seluj78]

LGTM but what would the advantage of this be instead of using typing.Literal ?

[KrilleGH]

@Seluj78, I'm not sure, where would you use typing.Literal?

My database is Postgres (I edited the description to mention this clearly) and it's possible to have different schemas there. This kind of schema is not to be confused with the structure of the database tables, also sometimes called schema. This schema is a way to group tables and other database objects. Without specification everything goes into the default public schema.

I have a lot of stuff in the public schema already which is used by a legacy application. This application is receiving a new API which is also storing some of its own data in the same database. In order to know what belongs to whom I put all the objects for the new part into a separate schema.

While SQLModel normally creates everything in the public schema, this PR allows to specify that objects should be created in a different schema.

[Seluj78]

Oh my, I'm really sorry I completely misread and misunderstood your PR.

After re-reading it, it totally makes sense. Enum fields should be correctly inherited into the correct schema !

Sorry again!

[YuriiMotov]

@KrilleGH, thanks for your interest and efforts!

As I understand it, In SQLAlchemy we need to explicitly specify Enum(MyEnum, inherit_schema=True) to get the same behavior.

My main question is: will it be correct that in SQLModel this will be by-default?

Can it be breaking change if somebody has multiple schemas and wants to have global enum type?

For now you can go the same way as with SQLAlchemy and specify inherit_schema=True explicitly:

class Hero(SchemaSQLModel):
    type: HeroType = Field(sa_type=Enum(HeroType, inherit_schema=True))

Results in:

 CREATE TYPE api.herotype AS ENUM ('normal', 'super')

Runnable code example in the details:

import enum

from sqlalchemy.schema import CreateSchema
from sqlmodel import Enum, Field, SQLModel, create_engine


class SchemaSQLModel(SQLModel):
    __table_args__ = {
        'schema': 'api',
    }


class HeroType(enum.StrEnum):
    normal = enum.auto()
    super = enum.auto()

class Hero(SchemaSQLModel, table=True):
    id: int | None = Field(default=None, primary_key=True)
    type: HeroType = Field(sa_type=Enum(HeroType, inherit_schema=True))


engine = create_engine("postgresql+psycopg2://user:mysecretpassword@localhost:5432/some_db", echo=True)
SQLModel.metadata.drop_all(engine)

with engine.connect() as connection:
    connection.execute(CreateSchema("api", if_not_exists=True))
    connection.commit()

SQLModel.metadata.create_all(engine)

For now I would vote for not changing this and advise to explicitly specify inherit_schema=True.

[KrilleGH]

@YuriiMotov, thank you for the suggestion. That looks like a good alternative I didn't find. I'll try it.

I don't have enough information to comment on your question regarding the difference in behavoir of SQLAlchemy and SQLModel this would create.

I currently have the patch from this PR automatically installed into my Docker container. I think I'd have to put an explicit Fields() declaration into every Enum field definition which isn't there currently, but it allows to remove the patch. It has always been a work-in-progress solution.