gateway/scripts/script_db_init_chatbotv2.py

102 lines
3.1 KiB
Python

#!/usr/bin/env python3
"""
Initialize poweron_chatbotv2 database for the Chatbot V2 feature.
Creates the poweron_chatbotv2 database if it does not exist.
Uses DB_CHATBOTV2_* config (falls back to DB_*).
Tables (ChatbotV2Conversation, ChatbotV2ContextFile, ChatbotV2ExtractedContext,
ChatbotV2Message, ChatbotV2Document, ChatbotV2Log) are auto-created by the
connector on first use.
Usage:
python script_db_init_chatbotv2.py [--dry-run]
"""
import os
import sys
import argparse
import logging
from pathlib import Path
scriptPath = Path(__file__).resolve()
gatewayPath = scriptPath.parent.parent
sys.path.insert(0, str(gatewayPath))
os.chdir(str(gatewayPath))
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)
import psycopg2
from modules.shared.configuration import APP_CONFIG
DB_NAME = "poweron_chatbotv2"
CONFIG_PREFIX = "DB_CHATBOTV2"
def _get_config():
"""Get DB config: DB_CHATBOTV2_* with fallback to DB_*."""
host = APP_CONFIG.get(f"{CONFIG_PREFIX}_HOST") or APP_CONFIG.get("DB_HOST", "localhost")
port = int(APP_CONFIG.get(f"{CONFIG_PREFIX}_PORT") or APP_CONFIG.get("DB_PORT", "5432"))
user = APP_CONFIG.get(f"{CONFIG_PREFIX}_USER") or APP_CONFIG.get("DB_USER")
password = (
APP_CONFIG.get(f"{CONFIG_PREFIX}_PASSWORD_SECRET")
or APP_CONFIG.get(f"{CONFIG_PREFIX}_PASSWORD")
or APP_CONFIG.get("DB_PASSWORD_SECRET")
or APP_CONFIG.get("DB_PASSWORD")
)
return {"host": host, "port": port, "user": user, "password": password}
def init_chatbotv2_db(dry_run: bool = False) -> bool:
"""Create poweron_chatbotv2 database if it does not exist."""
config = _get_config()
if not config["user"] or not config["password"]:
logger.error("DB_USER and DB_PASSWORD (or DB_CHATBOTV2_*) required")
return False
try:
conn = psycopg2.connect(
host=config["host"],
port=config["port"],
database="postgres",
user=config["user"],
password=config["password"],
)
conn.autocommit = True
with conn.cursor() as cur:
cur.execute(
"SELECT 1 FROM pg_database WHERE datname = %s",
(DB_NAME,),
)
exists = cur.fetchone() is not None
if exists:
logger.info(f"Database {DB_NAME} already exists")
else:
if dry_run:
logger.info(f"[DRY-RUN] Would create database {DB_NAME}")
else:
cur.execute(f'CREATE DATABASE "{DB_NAME}"')
logger.info(f"Created database {DB_NAME}")
conn.close()
return True
except Exception as e:
logger.error(f"Failed to init {DB_NAME}: {e}")
return False
def main():
parser = argparse.ArgumentParser(description="Initialize poweron_chatbotv2 database")
parser.add_argument("--dry-run", action="store_true", help="Do not create, only report")
args = parser.parse_args()
ok = init_chatbotv2_db(dry_run=args.dry_run)
sys.exit(0 if ok else 1)
if __name__ == "__main__":
main()