126 lines
5.0 KiB
Python
126 lines
5.0 KiB
Python
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))
|
||
|