44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
"""数据库引擎与初始化。
|
||
|
||
init_db():
|
||
- 若目标库不存在则自动创建（连到默认 postgres 库执行 CREATE DATABASE）
|
||
- 建表
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import psycopg2
|
||
from sqlalchemy import create_engine
|
||
from sqlalchemy.orm import sessionmaker
|
||
|
||
import config
|
||
from models import Base
|
||
|
||
engine = create_engine(config.DB_URL, pool_pre_ping=True, future=True)
|
||
SessionLocal = sessionmaker(bind=engine, autoflush=False, expire_on_commit=False, future=True)
|
||
|
||
|
||
def _ensure_database():
|
||
try:
|
||
conn = psycopg2.connect(
|
||
host=config.PG_HOST, port=config.PG_PORT, user=config.PG_USER,
|
||
password=config.PG_PASSWORD, dbname="postgres", connect_timeout=5)
|
||
conn.autocommit = True
|
||
cur = conn.cursor()
|
||
cur.execute("SELECT 1 FROM pg_database WHERE datname = %s", (config.PG_DB,))
|
||
if not cur.fetchone():
|
||
cur.execute(f'CREATE DATABASE "{config.PG_DB}"')
|
||
cur.close()
|
||
conn.close()
|
||
except Exception as e: # 数据库已存在或无权限时不阻断
|
||
print("[db] ensure_database:", repr(e)[:120])
|
||
|
||
|
||
def init_db():
|
||
_ensure_database()
|
||
Base.metadata.create_all(engine)
|
||
print("[db] tables ready")
|
||
|
||
|
||
def get_session():
|
||
return SessionLocal()
|