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))