Open evan0greenup opened 3 months 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)))
Privileged issue
Issue Content
For
enum.Enum
it works fine, but when the field isFlag
orIntFlag
, if the value is not single-only bit enabled (zero or many bits). Then it will raise error.It just treat
Flag
as same asEnum
. However, it should make all subset of bitwise-or be accepted.