feat(api): Add `EnumText` column type for SQLAlchemy

Introduce `EnumText`, a custom column type for SQLAlchemy designed
to work seamlessly with enumeration classes based on `StrEnum`.
This type stores enum members as `VARCHAR` in the database and
automatically handles conversion between the enumeration type and
its string representation during reads and writes.

Additionally, it validates that the stored and retrieved values
are valid members of the associated enumeration class.
pull/19737/head
QuantumGhost 1 year ago
parent bbebf9ad3e
commit 601fcbc21a

@ -1,4 +1,7 @@
from sqlalchemy import CHAR, TypeDecorator
import enum
from typing import Generic, TypeVar
from sqlalchemy import CHAR, VARCHAR, TypeDecorator
from sqlalchemy.dialects.postgresql import UUID
@ -24,3 +27,51 @@ class StringUUID(TypeDecorator):
if value is None:
return value
return str(value)
_E = TypeVar("_E", bound=enum.StrEnum)
class EnumText(TypeDecorator, Generic[_E]):
impl = VARCHAR
cache_ok = True
_length: int
_enum_class: type[_E]
def __init__(self, enum_class: type[_E], length: int | None = None):
self._enum_class = enum_class
max_enum_value_len = max(len(e.value) for e in enum_class)
if length is not None:
if length < max_enum_value_len:
raise ValueError("length should be greater than enum value length.")
self._length = length
else:
# leave some rooms for future longer enum values.
self._length = max(max_enum_value_len, 20)
def process_bind_param(self, value: _E | str, dialect):
if value is None:
return value
if isinstance(value, self._enum_class):
return value.value
elif isinstance(value, str):
self._enum_class(value)
return value
else:
raise TypeError(f"expected str or {self._enum_class}, got {type(value)}")
def load_dialect_impl(self, dialect):
return dialect.type_descriptor(VARCHAR(self._length))
def process_result_value(self, value, dialect) -> _E | None:
if value is None:
return value
if not isinstance(value, str):
raise TypeError(f"expected str, got {type(value)}")
return self._enum_class(value)
def compare_values(self, x, y):
if x is None or y is None:
return x is y
return x == y

Loading…
Cancel
Save