Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
1 task done
evan0greenup opened this issue Jul 11, 2024 · 1 comment
Open
1 task done

Comments

@evan0greenup
Copy link

Privileged issue

  • I'm @tiangolo or he asked me directly to create an issue here.

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
Copy link

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)))

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

2 participants