diff --git a/README.md b/README.md index 65f3859..7baa632 100644 --- a/README.md +++ b/README.md @@ -65,7 +65,11 @@ Database.from_config({"type": "postgresql", "url": "postgresql://user:pwd@127.0. | 方法 | 签名 | 说明 | | ------------- | ------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------- | | `__init__` | `(db_path=None, url=None)` | 二选一必填，否则抛 `ValueError`；`db_path` 自动创建父目录 | -| `initialize` | `(generate_schemas=True, echo=False, **kwargs)` | 创建异步引擎，`SQLModel.metadata.create_all` 自动建表（`generate_schemas=False` 可跳过，交由外部迁移工具），构建 `sessionmaker` | +| `initialize` | `(generate_schemas=True, echo=False, **kwargs)` | 创建异步引擎，`SQLModel.metadata.create_all` 自动建表（`generate_schemas=False` 可跳过，交由外部迁移工具），构建 `sessionmaker` | + +> PostgreSQL：`initialize()` 检测到目标库不存在时（如全新 PG 实例），会以相同凭据连接 +> `postgres` 维护库自动 `CREATE DATABASE`（多进程并发启动安全），因此无需在部署流程中 +> 手工建库；要求连接账号具备 `CREATEDB` 权限（默认超级用户 `postgres` 满足）。 | `get_session` | `() -> AsyncSession` | 返回新的异步会话；未初始化则抛 `RuntimeError` | | `dispose` | `async () -> None` | 关闭异步引擎并清除内部状态 | diff --git a/pyproject.toml b/pyproject.toml index b18c2f3..0f4cd65 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ requires = ["hatchling"] [project] name = "chongming-database" -version = "0.2.0" +version = "0.2.1" description = "异步数据库封装：SQLite（aiosqlite）+ PostgreSQL（asyncpg），基于 SQLModel/SQLAlchemy" readme = "README.md" requires-python = ">=3.12" diff --git a/src/chongming_database/__init__.py b/src/chongming_database/__init__.py index c48a235..7e212a1 100644 --- a/src/chongming_database/__init__.py +++ b/src/chongming_database/__init__.py @@ -1,9 +1,15 @@ +import logging from pathlib import Path + +import asyncpg +from sqlalchemy.engine import make_url from sqlalchemy.exc import OperationalError from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine from sqlalchemy.orm import sessionmaker from sqlmodel import SQLModel +logger = logging.getLogger("chongming_database") + class Database: """异步数据库封装，支持 SQLite 与 PostgreSQL 两种后端。 @@ -66,6 +72,47 @@ class Database: return self._url return f"sqlite+aiosqlite:///{self._db_path}" + async def _ensure_database(self) -> None: + """PostgreSQL 目标库不存在时自动创建（SQLite 文件本就自动生成，无需此步骤）。 + + 全新 PG 实例（如部署环境刚起容器）里只有 postgres 维护库，若不先建库， + initialize() 首次连接即报 ``database "xxx" does not exist``。 + 这里以相同凭据连接 postgres 维护库检查并执行 CREATE DATABASE。 + + 说明： + - CREATE DATABASE 不能运行在事务块内，因此不走 SQLAlchemy 引擎 + （其连接默认开事务），直接用 asyncpg 裸连接（语句自动提交）。 + - 要求连接账号具备 CREATEDB 权限；部署中通常以超级用户（postgres）运行。 + - 并发启动时多个进程可能同时判断"库不存在"，DuplicateDatabaseError 忽略即可。 + """ + url = make_url(self._build_url()) + if url.get_backend_name() != "postgresql": + return + db_name = url.database + if not db_name or db_name in ("postgres", "template0", "template1"): + return + + conn = await asyncpg.connect( + host=url.host, + port=url.port, + user=url.username, + password=url.password, + database="postgres", + ) + try: + exists = await conn.fetchval( + "SELECT 1 FROM pg_database WHERE datname = $1", db_name + ) + if exists is None: + quoted = db_name.replace('"', '""') + await conn.execute(f'CREATE DATABASE "{quoted}"') + logger.info("PostgreSQL 数据库 %s 不存在，已自动创建", db_name) + except asyncpg.DuplicateDatabaseError: + # 多进程并发启动的竞态：另一进程刚创建成功 + pass + finally: + await conn.close() + async def initialize( self, generate_schemas: bool = True, @@ -86,11 +133,17 @@ class Database: PostgreSQL：DDL 是事务性的，用 begin() 包裹 create_all， 同样 checkfirst=True + "already exists" 兜底。 + PostgreSQL 建库说明：目标库不存在时（全新 PG 实例），会先以相同凭据 + 连接 postgres 维护库自动 CREATE DATABASE（见 _ensure_database）， + 因此无需在部署流程中手工建库；要求连接账号具备 CREATEDB 权限。 + Args: generate_schemas: 是否自动建表（False 表示 schema 由外部迁移工具管理） echo: 是否打印 SQL（调试用） engine_kwargs: 透传给 create_async_engine（如 pool_size 等） """ + if self._url is not None: + await self._ensure_database() self._engine = create_async_engine( self._build_url(), echo=echo, diff --git a/tests/test_postgres.py b/tests/test_postgres.py index 381642d..7bd563d 100644 --- a/tests/test_postgres.py +++ b/tests/test_postgres.py @@ -3,6 +3,7 @@ import os import uuid +import asyncpg import pytest from sqlmodel import Field, SQLModel, select @@ -72,3 +73,37 @@ async def test_postgres_generate_schemas_false(): await session.execute(select(1)) finally: await db.dispose() + + +async def test_postgres_auto_create_missing_database(): + """目标库不存在时 initialize() 应自动建库（回归：database "x" does not exist）。""" + from sqlalchemy.engine import make_url + + base = make_url(PG_URL) + db_name = f"chongming_test_autocreate_{uuid.uuid4().hex[:8]}" + url = base.set(database=db_name).render_as_string(hide_password=False) + + db = Database.from_url(url) + try: + await db.initialize() # 此前这里会因库不存在而抛 OperationalError + async with db.get_session() as session: + cur = await session.execute(select(1)) + assert cur.scalar() == 1 + # 幂等：库已存在时重复 initialize 也应成功（先释放旧引擎连接，便于清理） + await db.dispose() + await db.initialize() + finally: + await db.dispose() + # 清理：连 postgres 维护库删除测试库 + maint = base.set(database="postgres") + conn = await asyncpg.connect( + host=maint.host, + port=maint.port, + user=maint.username, + password=maint.password, + database="postgres", + ) + try: + await conn.execute(f'DROP DATABASE IF EXISTS "{db_name}"') + finally: + await conn.close()