fastapi / sqlmodel

SQL databases in Python, designed for simplicity, compatibility, and robustness.
https://sqlmodel.tiangolo.com/
MIT License
14.34k stars 635 forks source link

`enum.IntFlag` field error when the flag is zero or many field enabled. #1010

Open evan0greenup opened 3 months ago

evan0greenup commented 3 months ago

Privileged issue

Issue Content

For enum.Enum it works fine, but when the field is Flag or IntFlag, if the value is not single-only bit enabled (zero or many bits). Then it will raise error.

It just treat Flag as same as Enum. However, it should make all subset of bitwise-or be accepted.

luoshuijs commented 1 month ago

By default, SQLAlchemy, which SQLModel relies on, cannot handle combination values of IntFlag.

Below is an example code:

from enum import IntFlag
from sqlmodel import Field, SQLModel, create_engine, Session, select
from typing import Optional

# 1. Define the IntFlag enum
class Permission(IntFlag):
    READ = 1
    WRITE = 2
    EXECUTE = 4

# 3. Define the SQLModel model
class User(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    permissions: Permission

# 4. Create an in-memory database connection
engine = create_engine("sqlite:///:memory:", echo=True)

# 5. Create the table
SQLModel.metadata.create_all(engine)

# 6. Insert data
def create_user(name: str, permissions: Permission):
    with Session(engine) as session:
        user = User(name=name, permissions=permissions)
        session.add(user)
        session.commit()

# 7. Query data
def get_users_with_permission(permission: Permission):
    with Session(engine) as session:
        statement = select(User).where(User.permissions & permission == permission)
        return session.exec(statement).all()

# Test code

# Insert user data
create_user("Alice", Permission.READ | Permission.WRITE)
create_user("Bob", Permission.READ)
create_user("Charlie", Permission.EXECUTE)

# Query users with READ permission
users_with_read_permission = get_users_with_permission(Permission.READ)
for user in users_with_read_permission:
    print(f"User: {user.name}, Permissions: {user.permissions}")

# Query users with WRITE permission
users_with_write_permission = get_users_with_permission(Permission.WRITE)
for user in users_with_write_permission:
    print(f"User: {user.name}, Permissions: {user.permissions}")

When running the code, the following error occurs:

LookupError: '3' is not among the defined enum values. Enum name: permission. Possible values: READ, WRITE, EXECUTE

Upon investigating how SQLAlchemy handles enum types, particularly in the sqlalchemy/sql/sqltypes.py file, we find that the Enum type has a native_enum parameter, which specifies:

:param native_enum: Use the database's native ENUM type when available. Defaults to True. When False, uses VARCHAR + check constraint for all backends. When False, the VARCHAR length can be controlled with :paramref:.Enum.length; currently "length" is ignored if native_enum=True.

This means that by default, enums are mapped to strings (VARCHAR) in the database, and only the explicitly defined enum values are accepted (i.e., READ, WRITE, EXECUTE in the example above).

When trying to insert combined permissions (e.g., Permission.READ | Permission.WRITE, which equals 3), SQLAlchemy cannot find the corresponding value in the enum, leading to a KeyError, followed by a LookupError.

To resolve this issue, we need to use a TypeDecorator to customize the field type so that it can correctly handle combination values of IntFlag.

from sqlalchemy.types import TypeDecorator, Integer

class IntFlagType(TypeDecorator):
    impl = Integer

    def __init__(self, enum_class, *args, **kwargs):
        self.enum_class = enum_class
        super(IntFlagType, self).__init__(*args, **kwargs)

    def process_bind_param(self, value, dialect):
        if value is not None:
            return int(value)  # Convert Permission to an integer for storage
        else:
            return None

    def process_result_value(self, value, dialect):
        if value is not None:
            return self.enum_class(value)  # Convert the integer back to Permission
        else:
            return None

In the User model, we use the custom IntFlagType to define the permissions field:

from sqlalchemy import Column

class User(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    permissions: Permission = Field(sa_column=Column(IntFlagType(Permission)))