alembic upgrade head. NameError: name 'sqlmodel' is not defined

简介

fastapi中，我通过SQLModel定义了一个users model(如下), 然后通过alembic进行表迁移，
执行alembic revision --autogenerate -m "Initial migration"正常，
然后再执行alembic upgrade head，
就报错了NameError: name 'sqlmodel' is not defined.

import datetime

from pydantic import EmailStr
from typing import Optional
from sqlmodel import Field, SQLModel
from sqlalchemy import Column, DateTime, func


class User(SQLModel, table=True):
    __tablename__ = "user"
    __table_args__ = {"comment": "用户表"}

    id: int = Field(primary_key=True, sa_column_kwargs={"autoincrement": True})

    name: str = Field(max_length=100, description="名称")
    password: str = Field(max_length=255, description="密码")  
    ...

问题: NameError: name 'sqlmodel' is not defined

原因

alembic revision --autogenerate -m "Initial migration"会生成 version_initial_migration.py文件，
文件里有字段是通过sqlmodel进行定义，但version_initial_migration.py里却没有导入sqlmodel库，
导致执行下一步迁移时报错 NameError: name 'sqlmodel' is not defined

处理方法

方案1：自己手动在version_initial_migration.py文件头部添加import sqlmodel，但这样后续每生成一个版本文件都要手动添加。

方案2（推荐）：找到/alembic/script.py.mako文件，然后在文件头部中添加上import sqlmodel。
这样每次执行迁移指令生成的每个版本文件都会自动添加import sqlmodel了。

# ./alembic/script.py.mako


from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
import sqlmodel # 新增
${imports if imports else ""}

参考内容

https://github.com/sqlalchemy/alembic/discussions/1404