scm-harmonic-cmts-admin/init_db.py

import asyncio
import asyncpg
import logging
import sys
import os
from dotenv import load_dotenv  # 🌟 1. 引入 load_dotenv

# 🌟 2. 明確指示 Python 讀取同目錄下的 .env 檔案
load_dotenv()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 🌟 2. 改用 os.getenv 讀取環境變數，並保留原本的設定作為安全預設值
DB_CONFIG = {
    "database": os.getenv("DB_NAME", "cmts_nms"),
    "user": os.getenv("DB_USER", "postgres"),  # 本地開發常用的預設帳號，或留空 ""
    "password": os.getenv("DB_PASS", ""),      # 🌟 絕對機密：預設留空！
    "host": os.getenv("DB_HOST", "127.0.0.1"),
    "port": os.getenv("DB_PORT", "5432")
}

async def init_database(force_reset: bool = False):
    try:
        logger.info("🔄 正在連線到 PostgreSQL (asyncpg)...")
        conn = await asyncpg.connect(
            database=DB_CONFIG["database"],
            user=DB_CONFIG["user"],
            password=DB_CONFIG["password"],
            host=DB_CONFIG["host"],
            port=int(DB_CONFIG["port"])
        )

        # ==========================================
        # 💣 核彈模式：清除舊有資料表
        # ==========================================
        if force_reset:
            logger.warning("⚠️ 警告：啟動強制重建模式，正在刪除現有資料表...")
            await conn.execute("DROP TABLE IF EXISTS cmts_options;")
            await conn.execute("DROP TABLE IF EXISTS device_status;")
            await conn.execute("DROP TABLE IF EXISTS system_filters;")
            # 💡 備份表通常極度重要，即使 reset 也不建議輕易 DROP，除非你確定要連備份一起砍
            # await conn.execute("DROP TABLE IF EXISTS config_backups;") 
            logger.info("🗑️ 舊資料表已清除完畢。")

        # ==========================================
        # 🏗️ 安全建置模式：建立資料表 (IF NOT EXISTS)
        # ==========================================
        # 1. 建立選項快取表 cmts_options (無 config_type)
        logger.info("🛠️ 正在檢查/建立 cmts_options 資料表...")
        await conn.execute("""
            CREATE TABLE IF NOT EXISTS cmts_options (
                host VARCHAR(255) NOT NULL,
                path VARCHAR(500) NOT NULL,
                data JSONB NOT NULL,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (host, path)
            );
        """)

        # 2. 建立設備狀態表 device_status (無 config_type)
        logger.info("🛠️ 正在檢查/建立 device_status 資料表...")
        await conn.execute("""
            CREATE TABLE IF NOT EXISTS device_status (
                host VARCHAR(255) NOT NULL PRIMARY KEY,
                cmts_version VARCHAR(100) DEFAULT 'unknown',
                last_scanned VARCHAR(100),
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );
        """)

        # 3. 建立過濾器設定表 system_filters
        logger.info("🛠️ 正在檢查/建立 system_filters 資料表...")
        await conn.execute("""
            CREATE TABLE IF NOT EXISTS system_filters (
                config_type VARCHAR(50) PRIMARY KEY,
                hidden_keys TEXT[] DEFAULT '{}',
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );
        """)

        # 4. 建立設備配置備份表 config_backups
        logger.info("🛠️ 正在檢查/建立 config_backups 資料表與索引...")
        await conn.execute("""
            CREATE TABLE IF NOT EXISTS config_backups (
                id UUID PRIMARY KEY,
                host VARCHAR(255) NOT NULL,
                config_type VARCHAR(50) NOT NULL DEFAULT 'running',
                timestamp TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
                snapshot_name VARCHAR(255),
                description TEXT DEFAULT '',
                is_auto BOOLEAN NOT NULL DEFAULT FALSE,
                raw_cli TEXT,
                parsed_tree JSONB,
                CONSTRAINT uq_snapshot_name UNIQUE (host, config_type, snapshot_name)
            );
        """)
        
        await conn.execute("""
            CREATE INDEX IF NOT EXISTS idx_config_backups_host_type_ts 
            ON config_backups (host, config_type, timestamp DESC);
        """)

        await conn.close()
        logger.info("✅ 資料庫初始化完成！所有資料表已準備就緒。")

    except Exception as e:
        logger.error(f"❌ 資料庫初始化失敗: {e}")

if __name__ == "__main__":
    # 偵測命令列參數是否包含 --reset
    is_reset = "--reset" in sys.argv
    
    if is_reset:
        print("\n" + "="*50)
        print("🚨 你正在執行資料庫重置 (--reset) 🚨")
        print("這將會清空所有的選項快取與系統設定！")
        print("="*50 + "\n")
        confirm = input("確定要繼續嗎？(輸入 yes 繼續): ")
        if confirm.lower() != "yes":
            print("已取消操作。")
            sys.exit(0)
            
    asyncio.run(init_database(force_reset=is_reset))