如何用SQLAlchemy映射含特殊字符的PostgreSQL枚举列？

问题背景 #

PostgreSQL数据库中定义了如下枚举类型和表：

CREATE TYPE enum_type AS ENUM('Hired', 'Possible Hire', 'Non-hired');

CREATE TABLE Status (
    ID VARCHAR(63) PRIMARY KEY,
    type enum_type NOT NULL
);

最初尝试定义Python枚举类和SQLAlchemy模型：

import enum

class TypeEnum(str, enum.Enum):
    Hired = 'hired'
    PossibleHire = 'possible hire'
    Nonhired = 'non-hired'

from sqlalchemy import Column, Enum, String
from sqlalchemy.ext.declarative import declarative_base

Base = declarative_base()

class StatusModel(Base):
    __tablename__ = "status"

    id = Column(String(63), primary_key=True, index=True)
    type = Column("type", Enum(TypeEnum))

但遇到无法将'Non-hired'作为Python类变量名的问题，改用Enum函数式API后：

TypeEnum = enum.Enum(
    value="TypeEnum",
    names=[
        ('Hired', 'hired'),
        ('Possible hire', 'possible hire'),
        ('Non-hired', 'non-hired'),
    ]
)

SQLAlchemy会直接传递枚举值（如'hired'）到数据库，不符合PostgreSQL枚举要求的'Hired'格式，导致写入失败。

后续尝试通过Pydantic Schema处理：

from enum import Enum
from pydantic import BaseModel

TypeEnumSchema = Enum(
    value="TypeEnum",
    names=[
       ('Hired', 'hired'),
       ('PossibleHire', 'possible hire'),
       ('Non-hired', 'non-hired'),
    ]
)

class StatusSchema(BaseModel):
    id: str
    type: TypeEnumSchema
    class Config:
        orm_mode = True

数据流程为：读取JSON数据→Pydantic校验→转换为SQLAlchemy模型→批量插入，但出现报错：

sqlalchemy.exc.StatementError: (builtins.LookupError) 'TypeEnum.Hired' is not among the defined enum values. Enum name: typeenum. Possible values: Hired, Possible hired, Non-hired

解决方案 #

方案1：自定义SQLAlchemy枚举类型，实现双向映射 #

通过SQLAlchemy的TypeDecorator自定义类型，处理Python枚举与PostgreSQL枚举的格式转换：

1. 定义Python枚举和映射字典：

import enum
from sqlalchemy.types import TypeDecorator, Enum as SAEnum

class TypeEnum(str, enum.Enum):
    HIRED = 'hired'
    POSSIBLE_HIRE = 'possible hire'
    NON_HIRED = 'non-hired'

# 双向映射字典：Python枚举→数据库枚举，数据库枚举→Python枚举
PY_TO_DB = {
    TypeEnum.HIRED: 'Hired',
    TypeEnum.POSSIBLE_HIRE: 'Possible Hire',
    TypeEnum.NON_HIRED: 'Non-hired'
}
DB_TO_PY = {v: k for k, v in PY_TO_DB.items()}

2. 自定义SQLAlchemy类型：

class DBTypeEnum(TypeDecorator):
    impl = SAEnum(TypeEnum)
    cache_ok = True

    def process_bind_param(self, value, dialect):
        # 写入数据库时，将Python枚举转换为数据库需要的格式
        return PY_TO_DB.get(value) if value else None

    def process_result_value(self, value, dialect):
        # 读取数据时，将数据库返回值转换为Python枚举
        return DB_TO_PY.get(value) if value else None

3. 修改SQLAlchemy模型：

class StatusModel(Base):
    __tablename__ = "status"

    id = Column(String(63), primary_key=True, index=True)
    type = Column("type", DBTypeEnum)

方案2：统一枚举值与数据库一致，用别名处理变量名 #

让Python枚举的值与PostgreSQL枚举完全匹配，用合法的变量名作为枚举成员名，同时添加转换方法处理输入的小写字符串：

1. 定义Python枚举：

import enum

class TypeEnum(str, enum.Enum):
    HIRED = 'Hired'
    POSSIBLE_HIRE = 'Possible Hire'
    NON_HIRED = 'Non-hired'

    @classmethod
    def from_lowercase(cls, input_str):
        """从输入的小写字符串转换为枚举"""
        mapping = {
            'hired': cls.HIRED,
            'possible hire': cls.POSSIBLE_HIRE,
            'non-hired': cls.NON_HIRED
        }
        return mapping[input_str.strip().lower()]

2. Pydantic Schema添加验证器：

from pydantic import BaseModel, validator

class StatusSchema(BaseModel):
    id: str
    type: TypeEnum

    @validator('type', pre=True)
    def convert_input(cls, value):
        if isinstance(value, str):
            return TypeEnum.from_lowercase(value)
        return value

    class Config:
        orm_mode = True

3. SQLAlchemy模型直接使用原生Enum：

class StatusModel(Base):
    __tablename__ = "status"

    id = Column(String(63), primary_key=True, index=True)
    type = Column("type", Enum(TypeEnum))

方案3：利用SQLAlchemy的values_callable参数 #

通过values_callable指定SQLAlchemy使用枚举的value属性而非成员名进行映射：

1. 定义与数据库值一致的Python枚举：

import enum

class TypeEnum(str, enum.Enum):
    HIRED = 'Hired'
    POSSIBLE_HIRE = 'Possible Hire'
    NON_HIRED = 'Non-hired'

2. 修改SQLAlchemy模型：

class StatusModel(Base):
    __tablename__ = "status"

    id = Column(String(63), primary_key=True, index=True)
    # 指定values_callable，使用枚举的value属性映射数据库枚举
    type = Column("type", Enum(TypeEnum, values_callable=lambda e: [item.value for item in e]))

3. Pydantic Schema处理输入转换：

from pydantic import BaseModel, validator

class StatusSchema(BaseModel):
    id: str
    type: TypeEnum

    @validator('type', pre=True)
    def normalize_input(cls, value):
        if isinstance(value, str):
            # 将输入的小写字符串转换为数据库枚举格式
            if value == 'hired':
                return TypeEnum.HIRED
            elif value == 'possible hire':
                return TypeEnum.POSSIBLE_HIRE
            elif value == 'non-hired':
                return TypeEnum.NON_HIRED
            raise ValueError(f"Invalid status type: {value}")
        return value

    class Config:
        orm_mode = True

内容的提问来源于stack exchange，提问作者Carlos Figueredo Triana