Python 与 AI 开发教程

Postgres 和数据库迁移

SQLite 撑不住多实例和高并发写入时,就该换 Postgres。代码改动很小,要注意的是类型、连接池和上线时的迁移顺序。

  • Python 服务端 · 第 6 章
  • 18 分钟
  • 4 组对比
  • 基于 PostgreSQL 18 · Psycopg 3.3 · SQLAlchemy 2.1 · Alembic 1.20 · Spring Boot 4.1

什么时候换 Postgres

SQLite 适合单机、写入不多的服务。出现下面任何一种情况,就该换成 Postgres:

  • 服务要部署多个实例,它们不能共用一个 SQLite 文件
  • 写入并发高,经常出现 database is locked
  • 需要 JSONB 查询、全文检索或向量检索这类能力

SQLModel 和 SQLAlchemy 屏蔽了大部分差异,换数据库主要是换连接地址。Android 应用不会直接连接服务端数据库,所以这一章的 Kotlin 示例是服务端写法,和 Java 一样基于 Spring。

本地用 Docker 启动一个就够了:

终端
docker run --name pg -e POSTGRES_PASSWORD=secret -p 5432:5432 -d postgres:18

Python 这边要安装驱动。binary 附带了编译好的 C 库,不用在系统里安装 libpq。

终端
pip install "psycopg[binary]"

连接数据库

连接地址的开头写 postgresql+psycopg,表示用 Psycopg 3 驱动。其余代码和 SQLite 相同。

Postgres 的每个连接都对应服务端的一个进程,创建成本比较高,所以要用连接池。SQLAlchemy 默认已经有连接池,下面这几个参数常要调整。

配置连接池

Python
from sqlalchemy import text
from sqlmodel import create_engine

DATABASE_URL = "postgresql+psycopg://postgres:secret@localhost:5432/postgres"

engine = create_engine(
    DATABASE_URL,
    pool_size=5,         # 保持的连接数
    max_overflow=10,     # 高峰时最多多开几个
    pool_pre_ping=True,  # 使用前检查连接是否还活着
)

with engine.connect() as connection:
    print(connection.scalar(text("SELECT 1 + 1")))  # => 2
@Bean
DataSource dataSource() {
    var config = new HikariConfig();
    config.setJdbcUrl("jdbc:postgresql://localhost:5432/postgres");
    config.setUsername("postgres");
    config.setPassword("secret");
    config.setMaximumPoolSize(15);
    return new HikariDataSource(config);
}
对比项PythonJava
连接池SQLAlchemy 自带HikariCP
最大连接数pool_size + max_overflowmaximumPoolSize
检查失效连接pool_pre_ping=True默认检查

多个实例的连接数要加起来算。Postgres 的 max_connections 默认只有 100,四个实例各开 30 个就超了。

每个请求一个 Session、接口写成 def 的做法和 SQLite 一样。check_same_thread 是 SQLite 的参数,换成 Postgres 后要删掉。

和 SQLite 的差异

SQLite 对类型很宽松,往整数列写字符串也能成功。Postgres 会直接报错,这些问题常在换库时才暴露。

开发环境也用 Postgres,不要开发用 SQLite、线上用 Postgres。两边的差异会让问题只在线上出现。

存 JSON

Agent 应用常要保存结构不固定的数据,比如工具调用的参数。Postgres 的 JSONB 把 JSON 解析后存储,可以查询内部字段,也能建索引。

JSONB 字段

Python
from sqlalchemy import Column
from sqlalchemy.dialects.postgresql import JSONB
from sqlmodel import Field, Session, SQLModel, create_engine, select

class ToolCall(SQLModel, table=True):
    id: int | None = Field(default=None, primary_key=True)
    name: str
    arguments: dict = Field(default_factory=dict, sa_column=Column(JSONB))

engine = create_engine(DATABASE_URL)
SQLModel.metadata.create_all(engine)

with Session(engine) as session:
    session.add(ToolCall(name="weather", arguments={"city": "北京"}))
    session.commit()

    query = select(ToolCall).where(ToolCall.arguments["city"].astext == "北京")
    print(session.exec(query).one().arguments)  # => {'city': '北京'}
@Entity
class ToolCall {
    @Id @GeneratedValue Long id;
    String name;

    @JdbcTypeCode(SqlTypes.JSON)
    @Column(columnDefinition = "jsonb")
    Map<String, Object> arguments;
}

interface ToolCallRepository extends JpaRepository<ToolCall, Long> {
    @Query(value = "SELECT * FROM tool_call "
        + "WHERE arguments ->> 'city' = :city", nativeQuery = true)
    List<ToolCall> findByCity(String city);
}
对比项PythonJava
字段声明sa_column=Column(JSONB)@JdbcTypeCode(SqlTypes.JSON)
读出来的类型dictMap
查询内部字段字段["city"].astext原生 SQL 的 ->>

字段是 dict 时,直接改字典内容不会被 Session 发现。要重新赋值,或者用 SQLAlchemy 的 MutableDict。

并发写入

多个请求同时修改同一行时,容易出现「先读后写」的丢失更新。两个请求都读到余额 100,各扣 30,最后余额是 70,而不是 40。

有两种常用的解决办法:

  • 行锁:SELECT ... FOR UPDATE 读取时锁住这一行,其他事务要等当前事务结束
  • 原子更新:一条 SQL 直接完成,比如 UPDATE ... SET n = n - 30

行锁

Python
from sqlmodel import Field, Session, SQLModel, create_engine, select

class Account(SQLModel, table=True):
    id: int | None = Field(default=None, primary_key=True)
    balance: int

engine = create_engine(DATABASE_URL)
SQLModel.metadata.create_all(engine)
with Session(engine) as session:
    session.add(Account(id=1, balance=100))
    session.commit()

def withdraw(session: Session, account_id: int, amount: int) -> int:
    query = select(Account).where(Account.id == account_id).with_for_update()
    account = session.exec(query).one()
    account.balance -= amount
    session.commit()
    return account.balance

with Session(engine) as session:
    print(withdraw(session, 1, 30))  # => 70
interface AccountRepository extends JpaRepository<Account, Long> {
    @Lock(LockModeType.PESSIMISTIC_WRITE)
    @Query("select a from Account a where a.id = :id")
    Account findForUpdate(long id);
}

@Transactional
int withdraw(AccountRepository accounts, long id, int amount) {
    var account = accounts.findForUpdate(id);
    account.balance -= amount;
    return account.balance;
}
对比项PythonJava
加行锁.with_for_update()@Lock(PESSIMISTIC_WRITE)
释放commit 或回滚时事务结束时

事务要尽量短,不要在持有行锁时调用大模型。一次调用几十秒,这一行就被锁住几十秒。

「有就更新,没有就插入」用 ON CONFLICT,一条语句完成,不会有竞争。

Python
from sqlalchemy.dialects.postgresql import insert
from sqlmodel import Field, Session, SQLModel, create_engine, select

class Usage(SQLModel, table=True):
    user_id: str = Field(primary_key=True)
    tokens: int = 0

engine = create_engine(DATABASE_URL)
SQLModel.metadata.create_all(engine)

def add_tokens(session: Session, user_id: str, tokens: int) -> None:
    statement = insert(Usage).values(user_id=user_id, tokens=tokens)
    session.exec(statement.on_conflict_do_update(
        index_elements=["user_id"],
        set_={"tokens": Usage.tokens + tokens},
    ))
    session.commit()

with Session(engine) as session:
    add_tokens(session, "ada", 10)
    add_tokens(session, "ada", 5)
    print(session.exec(select(Usage)).one().tokens)  # => 15

迁移怎么上线

「SQLite 与 SQLModel」一章讲了 Alembic 的基本用法。换成 Postgres 后,render_as_batch 可以关掉,上线时要多注意几件事。

**密码不要写进 alembic.ini。** 在 migrations/env.py 里从环境变量读取:

Python
config.set_main_option("sqlalchemy.url", os.environ["DATABASE_URL"])

先看再执行。 这几个命令适合放进发布流程:

终端
alembic upgrade head --sql   # 只输出 SQL,不执行
alembic current              # 数据库当前的版本
alembic check                # 模型和迁移不一致时报错,适合放进 CI

新旧代码要能同时运行。 滚动发布时,新旧两个版本的服务同时连着数据库,所以破坏性修改不能一步完成。比如把 name 改成 title,要分三次发布:

  1. 新增 title 列,代码同时写两列,再把旧数据复制过来
  2. 代码改成只读写 title
  3. 删除 name 列

注意锁。 ALTER TABLE 大多需要表上最强的锁,等待期间后面的查询全部排队。给迁移设置 lock_timeout,等不到锁就失败,再找时间重试。大表上建索引要用 CONCURRENTLY,它不能放在事务里,所以要包一层 autocommit_block:

Python
def upgrade() -> None:
    op.execute("SET lock_timeout = '5s'")
    with op.get_context().autocommit_block():
        op.create_index(
            "ix_message_user_id",
            "message",
            ["user_id"],
            postgresql_concurrently=True,
        )

迁移工具

Python
# alembic upgrade head
def upgrade() -> None:
    op.add_column("task", sa.Column("priority", sa.Integer(), server_default="0"))

def downgrade() -> None:
    op.drop_column("task", "priority")
public class V3__AddPriority extends BaseJavaMigration {
    @Override
    public void migrate(Context context) throws Exception {
        try (var statement = context.getConnection().createStatement()) {
            statement.execute(
                "ALTER TABLE task ADD COLUMN priority INT DEFAULT 0");
        }
    }
}
对比项PythonJava
工具AlembicFlyway
生成对比模型自动生成,再人工检查自己写
触发单独运行命令默认在应用启动时运行
多实例发布流程里只运行一次用数据库锁保证只有一个实例执行

不要在每个实例启动时运行 alembic upgrade head。多个实例同时迁移会互相冲突,「用 Docker 部署」一章用单独的任务完成这一步。

异步驱动选读

create_async_engine 配合 postgresql+psycopg 或 asyncpg 可以在 async def 接口里访问数据库。代价是所有数据库代码都要写 await,访问关联属性也不能再隐式触发查询。

大多数服务用同步 Session 配 def 接口就够了。只有并发连接很多、数据库操作又很频繁时再考虑换成异步。

服务端数据库写法对照
PythonJava说明
postgresql+psycopg://...jdbc:postgresql://...连接地址
create_engine(pool_size=5)HikariCP都要限制连接数
Column(JSONB)@JdbcTypeCode(SqlTypes.JSON)可查询的 JSON
.with_for_update()@Lock(PESSIMISTIC_WRITE)行锁
on_conflict_do_update原生 SQL一条语句完成 upsert
AlembicFlywayFlyway 默认随应用启动

本章要点

  • 多实例、高并发写入或需要 JSONB 时换 Postgres,开发环境也用它
  • 连接地址换成 postgresql+psycopg,删掉 check_same_thread
  • 配好连接池,所有实例的连接数加起来不能超过 max_connections
  • 先读后写要加行锁或改成原子更新,事务里不要调用大模型
  • 破坏性迁移分多次发布,让新旧代码能同时运行
  • 迁移时设置 lock_timeout,大表索引用 CONCURRENTLY,只由一个任务执行迁移

参考资料:PostgreSQL JSON 类型显式锁Psycopg 3SQLAlchemy 连接池Alembic 自动生成Alembic 操作参考