在SQLModel中实现带自定义逻辑的自增字符串主键

要把自定义主键生成逻辑整合到Query类内部，可以借助SQLAlchemy的插入前事件钩子（SQLModel基于SQLAlchemy实现），结合数据库查询获取最新递增序号，具体实现如下：

完整代码实现 #

from typing import Optional
from sqlmodel import Field, SQLModel, Session, create_engine
from sqlalchemy import event, Integer
from sqlalchemy.sql import func

class Query(SQLModel, table=True):
    id: Optional[str] = Field(default=None, primary_key=True)
    text: str = Field(default=None, unique=True)
    task: str 
    use_case: str

# 绑定插入前事件，自动生成自定义主键
@event.listens_for(Query, 'before_insert')
def generate_custom_id(mapper, connection, target):
    # 提取task和use_case的前缀，统一大写避免大小写差异
    task_prefix = target.task[:3].upper()
    use_case_prefix = target.use_case[:3].upper()
    prefix_len = len(f"{task_prefix}-{use_case_prefix}-")
    
    # 查询当前表中对应前缀的最大序号
    max_num_result = connection.execute(
        func.max(
            func.cast(
                func.substr(Query.id, prefix_len + 1),
                Integer
            )
        ).label('max_num')
    ).fetchone()
    
    # 确定最新递增序号，无数据则从1开始
    latest_id = max_num_result.max_num if max_num_result.max_num is not None else 0
    target.id = f'{task_prefix}-{use_case_prefix}-{latest_id + 1}'

# 示例使用
if __name__ == "__main__":
    engine = create_engine("sqlite:///test.db")
    SQLModel.metadata.create_all(engine)
    
    with Session(engine) as session:
        # 创建实例时无需手动传入id
        new_query = Query(
            text="测试查询内容",
            task="文本分类",
            use_case="用户意图识别"
        )
        session.add(new_query)
        session.commit()
        print(f"自动生成的主键：{new_query.id}")  # 输出示例：WBF-YHT-1

关键逻辑说明 #

• 事件钩子触发：通过@event.listens_for注册插入前处理函数，每次插入Query实例时自动执行主键生成逻辑。
• 前缀处理：截取task和use_case的前3个字符并统一大写，避免因大小写导致前缀匹配异常。
• 序号递增：利用SQL函数提取现有主键的数字部分，转换为整数后取最大值，在此基础上加1得到新序号。
• 主键赋值：拼接前缀与新序号，直接赋值给当前实例的id字段。

注意事项 #

1. 主流数据库（SQLite、MySQL、PostgreSQL）均支持代码中用到的substr、cast等SQL函数，无需额外适配。
2. 高并发场景下，建议配合数据库锁或序列机制（如PostgreSQL的SERIAL）避免序号重复。
3. 若task或use_case长度不足3位，代码会自动截取全部字符，无需额外判断。

内容的提问来源于stack exchange，提问作者Madaray