Pine
ee1dfdfed9
- infrastructure/db.py：MySQL 方言编译钩子（String 无长度→VARCHAR(255)，Text→MEDIUMTEXT）； 新增 sync_database_url/make_sync_engine（sqlite+aiosqlite→sqlite，mysql+asyncmy→mysql+pymysql） - models.py：sessions.token 显式 String(2048)、system_configs.value 改 Text（存量数据超 255，MySQL 需兼容） - seed.py：_add_agent 幂等插入按方言分支（SQLite=INSERT OR IGNORE / MySQL=INSERT IGNORE） - park/tenants.py：同步只读引擎改走 make_sync_engine（原 +aiosqlite replace 对 MySQL 无效） - 依赖补 pymysql；.env.example 补 PINEAGENTS_DEMO_DATABASE_URL 说明
106 lines
3.4 KiB
Python
106 lines
3.4 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""基础设施层 · 异步 SQLAlchemy 引擎 / 会话 / 声明基类。
|
||
|
||
四层架构：基础设施层（含数据模型层）。全异步：
|
||
- 开发/默认：SQLite + aiosqlite
|
||
- 生产：MySQL + asyncmy（经 config.DATABASE_URL 切换）
|
||
每请求一个 async session（``get_session`` 依赖），事件循环内无阻塞。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from collections.abc import AsyncGenerator
|
||
|
||
from sqlalchemy import Text, create_engine, types
|
||
from sqlalchemy.ext.asyncio import (
|
||
AsyncSession,
|
||
async_sessionmaker,
|
||
create_async_engine,
|
||
)
|
||
from sqlalchemy.ext.compiler import compiles
|
||
from sqlalchemy.orm import DeclarativeBase
|
||
|
||
from .. import config
|
||
|
||
|
||
class Base(DeclarativeBase):
|
||
"""SQLAlchemy 声明式基类（数据模型层继承）。"""
|
||
|
||
|
||
@compiles(types.String, "mysql")
|
||
def _mysql_string_ddl(type_: types.String, compiler, **kw) -> str:
|
||
"""MySQL 方言适配：String() 未显式给长度时按 VARCHAR(255) 建表（MySQL 要求
|
||
VARCHAR 必须带长度；SQLite 忽略长度不受影响），Text 渲染 MEDIUMTEXT
|
||
（富文本/JSON 大字段，避免 MySQL TEXT 64KB 上限）。"""
|
||
if isinstance(type_, Text):
|
||
return "MEDIUMTEXT"
|
||
if not type_.length:
|
||
return "VARCHAR(255)"
|
||
return f"VARCHAR({int(type_.length)})"
|
||
|
||
|
||
def sync_database_url(url: str | None = None) -> str:
|
||
"""DATABASE_URL → 同步驱动 URL（培训同步旁路 / 迁移脚本用）。
|
||
|
||
sqlite+aiosqlite → sqlite（pysqlite）；mysql+asyncmy → mysql+pymysql。
|
||
"""
|
||
url = url or config.DATABASE_URL
|
||
if url.startswith("sqlite+aiosqlite"):
|
||
return url.replace("sqlite+aiosqlite", "sqlite", 1)
|
||
if url.startswith("mysql+asyncmy"):
|
||
return url.replace("mysql+asyncmy", "mysql+pymysql", 1)
|
||
return url
|
||
|
||
|
||
def make_sync_engine(url: str | None = None):
|
||
"""按 DATABASE_URL 创建同步引擎（asyncio.to_thread / 脚本场景）。"""
|
||
u = sync_database_url(url)
|
||
kwargs: dict = {"pool_pre_ping": True, "future": True}
|
||
if u.startswith("mysql"):
|
||
kwargs["pool_recycle"] = 3600
|
||
if "charset=" not in u:
|
||
kwargs["connect_args"] = {"charset": "utf8mb4"}
|
||
elif u.startswith("sqlite"):
|
||
kwargs["connect_args"] = {"check_same_thread": False}
|
||
return create_engine(u, **kwargs)
|
||
|
||
|
||
engine = create_async_engine(
|
||
config.DATABASE_URL,
|
||
pool_pre_ping=True,
|
||
)
|
||
|
||
AsyncSessionLocal = async_sessionmaker(
|
||
bind=engine,
|
||
class_=AsyncSession,
|
||
autoflush=False,
|
||
expire_on_commit=False,
|
||
)
|
||
|
||
|
||
async def init_db() -> None:
|
||
"""幂等地创建全部表（create_all 对已存在表无操作）。"""
|
||
from . import models # noqa: F401
|
||
|
||
async with engine.begin() as conn:
|
||
await conn.run_sync(Base.metadata.create_all)
|
||
|
||
|
||
async def get_session() -> AsyncGenerator[AsyncSession, None]:
|
||
"""FastAPI 依赖：请求级异步 SQLAlchemy 会话，请求结束自动关闭。"""
|
||
async with AsyncSessionLocal() as session:
|
||
yield session
|
||
|
||
|
||
def make_async_engine(db_url: str):
|
||
"""按给定 URL 创建独立异步引擎（测试用 tmp 数据库）。"""
|
||
return create_async_engine(db_url, pool_pre_ping=True)
|
||
|
||
|
||
def make_async_session_factory(db_url: str) -> async_sessionmaker:
|
||
return async_sessionmaker(
|
||
bind=make_async_engine(db_url),
|
||
class_=AsyncSession,
|
||
autoflush=False,
|
||
expire_on_commit=False,
|
||
)
|