king
114 lines
3.2 KiB
Python
114 lines
3.2 KiB
Python
import logging
|
|
import os
|
|
from urllib.parse import quote_plus
|
|
|
|
from sqlalchemy import create_engine, text
|
|
from sqlalchemy.orm import declarative_base, sessionmaker
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def load_env_file(path: str = ".env"):
|
|
if not os.path.exists(path):
|
|
return
|
|
|
|
with open(path, "r", encoding="utf-8") as env_file:
|
|
for raw_line in env_file:
|
|
line = raw_line.strip()
|
|
if not line or line.startswith("#") or "=" not in line:
|
|
continue
|
|
|
|
key, value = line.split("=", 1)
|
|
key = key.strip()
|
|
value = value.strip().strip('"').strip("'")
|
|
if key and key not in os.environ:
|
|
os.environ[key] = value
|
|
|
|
|
|
load_env_file()
|
|
|
|
|
|
def _build_pg_url(db_name: str) -> str:
|
|
user = os.getenv("DB_USER", "")
|
|
password = os.getenv("DB_PASSWORD", "")
|
|
host = os.getenv("DB_HOST", "localhost")
|
|
port = os.getenv("DB_PORT", "5432")
|
|
return (
|
|
f"postgresql+psycopg2://{quote_plus(user)}:{quote_plus(password)}"
|
|
f"@{host}:{port}/{db_name}"
|
|
)
|
|
|
|
|
|
def _resolve_database_url() -> str:
|
|
explicit_url = os.getenv("DATABASE_URL")
|
|
if explicit_url:
|
|
return explicit_url
|
|
|
|
order_db = os.getenv("ORDER_DB_NAME")
|
|
if order_db and os.getenv("DB_HOST"):
|
|
return _build_pg_url(order_db)
|
|
|
|
return "sqlite:///./app.db"
|
|
|
|
|
|
DATABASE_URL = _resolve_database_url()
|
|
|
|
connect_args = {}
|
|
engine_kwargs = {}
|
|
if DATABASE_URL.startswith("sqlite"):
|
|
connect_args = {"check_same_thread": False, "timeout": 30}
|
|
else:
|
|
engine_kwargs.update(pool_pre_ping=True, pool_recycle=1800)
|
|
|
|
engine = create_engine(DATABASE_URL, connect_args=connect_args, **engine_kwargs)
|
|
|
|
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
|
|
|
Base = declarative_base()
|
|
|
|
|
|
# --- Secondary engines for the other PostgreSQL databases ----------------
|
|
# These are exposed for code that needs to query item / return / checklist
|
|
# databases. They are created lazily and only when PostgreSQL is configured.
|
|
_secondary_engines: dict[str, object] = {}
|
|
|
|
|
|
def get_secondary_engine(env_var: str):
|
|
"""Get (or create) an engine for ITEM_DB_NAME / RETURN_DB_NAME / CHECKLIST_DB_NAME."""
|
|
db_name = os.getenv(env_var)
|
|
if not db_name or not os.getenv("DB_HOST"):
|
|
return None
|
|
if env_var in _secondary_engines:
|
|
return _secondary_engines[env_var]
|
|
eng = create_engine(_build_pg_url(db_name), pool_pre_ping=True, pool_recycle=1800)
|
|
_secondary_engines[env_var] = eng
|
|
return eng
|
|
|
|
|
|
# Dependency for FastAPI to get DB session
|
|
def get_db():
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def check_db_connection() -> tuple[bool, str]:
|
|
"""Ping the primary database. Returns (ok, message)."""
|
|
try:
|
|
with engine.connect() as conn:
|
|
conn.execute(text("SELECT 1"))
|
|
return True, "ok"
|
|
except Exception as exc: # noqa: BLE001
|
|
return False, str(exc)
|
|
|
|
|
|
def log_db_startup_status() -> None:
|
|
safe_url = engine.url.render_as_string(hide_password=True)
|
|
ok, msg = check_db_connection()
|
|
if ok:
|
|
logger.warning("[DB] Connected: %s", safe_url)
|
|
else:
|
|
logger.error("[DB] Connection FAILED: %s -> %s", safe_url, msg)
|