Agno 模板预配置了 SqlAlchemyalembic 来管理数据库。添加表的通用工作流程是:
  1. 将表定义添加到 db/tables 目录。
  2. db/tables/__init__.py 文件中导入表类。
  3. 创建数据库迁移。
  4. 运行数据库迁移。

表定义

我们来创建一个 UsersTable,将以下代码复制到 db/tables/user.py
db/tables/user.py
from datetime import datetime
from typing import Optional

from sqlalchemy.orm import Mapped, mapped_column
from sqlalchemy.sql.expression import text
from sqlalchemy.types import BigInteger, DateTime, String

from db.tables.base import Base


class UsersTable(Base):
    """存储用户数据的表。"""

    __tablename__ = "dim_users"

    id_user: Mapped[int] = mapped_column(
        BigInteger, primary_key=True, autoincrement=True, nullable=False, index=True
    )
    email: Mapped[str] = mapped_column(String)
    is_active: Mapped[bool] = mapped_column(default=True)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=text("now()")
    )
    updated_at: Mapped[Optional[datetime]] = mapped_column(
        DateTime(timezone=True), onupdate=text("now()")
    )
更新 db/tables/__init__.py 文件:
db/tables/__init__.py
from db.tables.base import Base
from db.tables.user import UsersTable

创建数据库修订

在开发容器中运行 alembic 命令来创建数据库迁移:
docker exec -it ai-api alembic -c db/alembic.ini revision --autogenerate -m "Initialize DB"

迁移开发数据库

运行 alembic 命令来迁移开发数据库:
docker exec -it ai-api alembic -c db/alembic.ini upgrade head

可选:添加测试用户

现在我们来添加一个测试用户。将以下代码复制到 db/tables/test_add_user.py
db/tables/test_add_user.py
from typing import Optional
from sqlalchemy.orm import Session

from db.session import SessionLocal
from db.tables.user import UsersTable
from utils.log import logger


def create_user(db_session: Session, email: str) -> UsersTable:
    """创建新用户。"""
    new_user = UsersTable(email=email)
    db_session.add(new_user)
    return new_user


def get_user(db_session: Session, email: str) -> Optional[UsersTable]:
    """通过电子邮件获取用户。"""
    return db_session.query(UsersTable).filter(UsersTable.email == email).first()


if __name__ == "__main__":
    test_user_email = "test@test.com"
    with SessionLocal() as sess, sess.begin():
        logger.info(f"正在创建用户: {test_user_email}")
        create_user(db_session=sess, email=test_user_email)
        logger.info(f"正在获取用户: {test_user_email}")
        user = get_user(db_session=sess, email=test_user_email)
        if user:
            logger.info(f"用户已创建: {user.id_user}")
        else:
            logger.info(f"用户未找到: {test_user_email}")

运行脚本以添加测试用户:
docker exec -it ai-api python db/tables/test_add_user.py

迁移生产数据库

我们建议通过设置环境变量 MIGRATE_DB = True 并重新启动生产服务来迁移生产数据库。这将从入口点脚本在容器启动时运行 alembic -c db/alembic.ini upgrade head

更新 workspace/prd_resources.py 文件

workspace/prd_resources.py
...
# -*- 构建容器环境
container_env = {
    ...
    # 启动时使用 alembic 迁移数据库
    "MIGRATE_DB": ws_settings.prd_db_enabled,
}
...

更新 ECS 任务定义

由于我们更新了环境变量,我们需要更新任务定义:
ag ws patch --env prd --infra aws --name td

更新 ECS 服务

更新任务定义后,重新部署生产应用程序:
ag ws patch --env prd --infra aws --name service

手动迁移生产数据库

另一种方法是 SSH 进入生产容器以手动运行迁移。您的 ECS 任务已启用 SSH 访问。运行 alembic 命令来迁移生产数据库:
ECS_CLUSTER=ai-app-prd-cluster
TASK_ARN=$(aws ecs list-tasks --cluster ai-app-prd-cluster --query "taskArns[0]" --output text)
CONTAINER_NAME=ai-api-prd

aws ecs execute-command --cluster $ECS_CLUSTER \
    --task $TASK_ARN \
    --container $CONTAINER_NAME \
    --interactive \
    --command "alembic -c db/alembic.ini upgrade head"

migrations 目录是如何创建的

这些命令已经运行过,并在此描述以供完整性参考
migrations 目录是使用以下命令创建的:
docker exec -it ai-api cd db && alembic init migrations
  • 运行上述命令后,应会创建 db/migrations 目录。
  • 更新 alembic.ini
    • script_location = db/migrations 设置为
    • 取消注释 [post_write_hooks] 中的 black hook
  • 按照 此链接 更新 db/migrations/env.py 文件
  • 将以下函数添加到 configure 中,以仅包含 target_metadata 中的表
db/migrations/env.py
# -*- 只包含 target_metadata 中的表
def include_name(name, type_, parent_names):
    if type_ == "table":
        return name in target_metadata.tables
    else:
        return True
...