Xiaoxia AI
84 lines
2.4 KiB
Python
84 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
from sqlalchemy import create_engine, text
|
|
from sqlalchemy.engine import URL, make_url
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from packages.adapters.sqlalchemy_impl.models import Base
|
|
|
|
SCHEMA_INIT_LOCK_ID = 2026061501
|
|
SessionLocal = None
|
|
|
|
|
|
def build_engine(
|
|
database_url: str,
|
|
*,
|
|
pool_size: int = 20,
|
|
max_overflow: int = 40,
|
|
pool_timeout: int = 30,
|
|
pool_recycle: int = 3600,
|
|
):
|
|
return create_engine(
|
|
database_url,
|
|
pool_size=pool_size,
|
|
max_overflow=max_overflow,
|
|
pool_timeout=pool_timeout,
|
|
pool_recycle=pool_recycle,
|
|
)
|
|
|
|
|
|
def build_session_factory(
|
|
database_url: str,
|
|
*,
|
|
pool_size: int = 20,
|
|
max_overflow: int = 40,
|
|
pool_timeout: int = 30,
|
|
pool_recycle: int = 3600,
|
|
):
|
|
engine = build_engine(
|
|
database_url,
|
|
pool_size=pool_size,
|
|
max_overflow=max_overflow,
|
|
pool_timeout=pool_timeout,
|
|
pool_recycle=pool_recycle,
|
|
)
|
|
session_factory = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
|
global SessionLocal
|
|
SessionLocal = session_factory
|
|
return engine, session_factory
|
|
|
|
|
|
def _build_admin_url(database_url: str) -> URL:
|
|
url = make_url(database_url)
|
|
return url.set(database="postgres")
|
|
|
|
|
|
def ensure_database_exists(database_url: str) -> None:
|
|
target_url = make_url(database_url)
|
|
admin_engine = create_engine(_build_admin_url(database_url), isolation_level="AUTOCOMMIT")
|
|
try:
|
|
with admin_engine.connect() as connection:
|
|
exists = connection.execute(
|
|
text("SELECT 1 FROM pg_database WHERE datname = :database_name"),
|
|
{"database_name": target_url.database},
|
|
).scalar()
|
|
if exists:
|
|
return
|
|
connection.execute(text(f'CREATE DATABASE "{target_url.database}"'))
|
|
finally:
|
|
admin_engine.dispose()
|
|
|
|
|
|
def initialize_database(engine) -> None:
|
|
with engine.connect() as connection:
|
|
connection.execute(text("SELECT pg_advisory_lock(:lock_id)"), {"lock_id": SCHEMA_INIT_LOCK_ID})
|
|
try:
|
|
Base.metadata.create_all(bind=connection)
|
|
connection.commit()
|
|
finally:
|
|
connection.execute(
|
|
text("SELECT pg_advisory_unlock(:lock_id)"),
|
|
{"lock_id": SCHEMA_INIT_LOCK_ID},
|
|
)
|
|
connection.commit()
|