dashboard-nanobot/backend/core/database.py

98 lines
2.9 KiB
Python

from sqlalchemy import inspect, text
from sqlmodel import Session, create_engine
from core.settings import (
DATABASE_ECHO,
DATABASE_MAX_OVERFLOW,
DATABASE_POOL_RECYCLE,
DATABASE_POOL_SIZE,
DATABASE_POOL_TIMEOUT,
DATABASE_URL,
)
_engine_kwargs = {
"echo": DATABASE_ECHO,
"pool_pre_ping": True,
"pool_size": DATABASE_POOL_SIZE,
"max_overflow": DATABASE_MAX_OVERFLOW,
"pool_timeout": DATABASE_POOL_TIMEOUT,
"pool_recycle": DATABASE_POOL_RECYCLE,
}
engine = create_engine(DATABASE_URL, **_engine_kwargs)
BOT_INSTANCE_TABLE = "bot_instance"
BOT_MESSAGE_TABLE = "bot_message"
BOT_IMAGE_TABLE = "bot_image"
BOT_REQUEST_USAGE_TABLE = "bot_request_usage"
BOT_ACTIVITY_EVENT_TABLE = "bot_activity_event"
SYS_LOGIN_LOG_TABLE = "sys_login_log"
SYS_SETTING_TABLE = "sys_setting"
REQUIRED_TABLES = (
BOT_INSTANCE_TABLE,
BOT_MESSAGE_TABLE,
BOT_IMAGE_TABLE,
BOT_REQUEST_USAGE_TABLE,
BOT_ACTIVITY_EVENT_TABLE,
SYS_LOGIN_LOG_TABLE,
SYS_SETTING_TABLE,
"skill_market_item",
"bot_skill_install",
"topic_topic",
"topic_item",
)
REQUIRED_SYS_SETTING_KEYS = (
"page_size",
"chat_pull_page_size",
"command_auto_unlock_seconds",
"auth_token_ttl_hours",
"auth_token_max_active",
"upload_max_mb",
"allowed_attachment_extensions",
"workspace_download_extensions",
"speech_enabled",
"activity_event_retention_days",
)
def _validate_required_tables() -> None:
inspector = inspect(engine)
missing = [table_name for table_name in REQUIRED_TABLES if not inspector.has_table(table_name)]
if missing:
raise RuntimeError(
"Database schema is not initialized. "
f"Missing tables: {', '.join(missing)}. "
"Run scripts/init-full-db.sh or apply scripts/sql/create-tables.sql before starting the backend."
)
def _validate_required_sys_settings() -> None:
placeholders = ", ".join(f":k{i}" for i, _ in enumerate(REQUIRED_SYS_SETTING_KEYS))
params = {f"k{i}": key for i, key in enumerate(REQUIRED_SYS_SETTING_KEYS)}
with engine.connect() as conn:
rows = conn.execute(
text(f'SELECT key FROM "{SYS_SETTING_TABLE}" WHERE key IN ({placeholders})'),
params,
).scalars().all()
present = {str(row or "").strip() for row in rows if str(row or "").strip()}
missing = [key for key in REQUIRED_SYS_SETTING_KEYS if key not in present]
if missing:
raise RuntimeError(
"Database seed data is not initialized. "
f"Missing sys_setting keys: {', '.join(missing)}. "
"Run scripts/init-full-db.sh or apply scripts/sql/init-data.sql before starting the backend."
)
def init_database() -> None:
with engine.connect() as conn:
conn.execute(text("SELECT 1"))
_validate_required_tables()
_validate_required_sys_settings()
def get_session():
with Session(engine) as session:
yield session