81 lines
2.2 KiB
Python
81 lines
2.2 KiB
Python
"""Postgres 连接 helper。
|
||
|
||
配置全部走环境变量，默认指向 docker-compose 起的本地实例：
|
||
|
||
PMDA_PG_HOST (默认 localhost)
|
||
PMDA_PG_PORT (默认 5432)
|
||
PMDA_PG_DB (默认 pmda)
|
||
PMDA_PG_USER (默认 pmda)
|
||
PMDA_PG_PASSWORD (默认 pmda_local_dev — 仅本地开发，生产由 secret 注入)
|
||
|
||
`connect()` 返回 psycopg3 connection（autocommit=False）。
|
||
长时跑批时使用 `pool()` 取 ConnectionPool。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from contextlib import contextmanager
|
||
from typing import Iterator
|
||
|
||
import psycopg
|
||
from psycopg import Connection
|
||
from psycopg_pool import ConnectionPool
|
||
|
||
PG_HOST = os.environ.get("PMDA_PG_HOST", "localhost")
|
||
PG_PORT = int(os.environ.get("PMDA_PG_PORT", "5432"))
|
||
PG_DB = os.environ.get("PMDA_PG_DB", "pmda")
|
||
PG_USER = os.environ.get("PMDA_PG_USER", "pmda")
|
||
PG_PASSWORD = os.environ.get("PMDA_PG_PASSWORD", "pmda_local_dev")
|
||
|
||
|
||
def conninfo() -> str:
|
||
return (
|
||
f"host={PG_HOST} port={PG_PORT} dbname={PG_DB} "
|
||
f"user={PG_USER} password={PG_PASSWORD}"
|
||
)
|
||
|
||
|
||
def connect(*, autocommit: bool = False) -> Connection:
|
||
"""Open a single connection. Caller is responsible for closing."""
|
||
return psycopg.connect(conninfo(), autocommit=autocommit)
|
||
|
||
|
||
@contextmanager
|
||
def session(*, autocommit: bool = False) -> Iterator[Connection]:
|
||
"""`with session() as conn:` — auto close on exit."""
|
||
conn = connect(autocommit=autocommit)
|
||
try:
|
||
yield conn
|
||
if not autocommit:
|
||
conn.commit()
|
||
except Exception:
|
||
if not autocommit:
|
||
conn.rollback()
|
||
raise
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
_pool: ConnectionPool | None = None
|
||
|
||
|
||
def pool(min_size: int = 1, max_size: int = 8) -> ConnectionPool:
|
||
"""Lazy-init module-level pool. Use for batch / agent-loop hot path."""
|
||
global _pool
|
||
if _pool is None:
|
||
_pool = ConnectionPool(
|
||
conninfo(),
|
||
min_size=min_size,
|
||
max_size=max_size,
|
||
kwargs={"autocommit": False},
|
||
open=True,
|
||
)
|
||
return _pool
|
||
|
||
|
||
def close_pool() -> None:
|
||
global _pool
|
||
if _pool is not None:
|
||
_pool.close()
|
||
_pool = None
|