461 lines
17 KiB
Python
461 lines
17 KiB
Python
import asyncpg
|
||
import json
|
||
import uuid
|
||
import os
|
||
from dotenv import load_dotenv # 🌟 1. 引入 load_dotenv
|
||
from typing import Dict, List, Optional, Any
|
||
from logger import get_logger
|
||
|
||
# 🌟 2. 明確指示 Python 讀取同目錄下的 .env 檔案
|
||
load_dotenv()
|
||
|
||
# ==========================================
|
||
# 💡 PostgreSQL 連線與操作 (高可用性版)
|
||
# ==========================================
|
||
|
||
logger = get_logger("app.database")
|
||
|
||
# 🌟 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")
|
||
}
|
||
|
||
_pool: Optional[asyncpg.Pool] = None
|
||
|
||
async def init_db_pool():
|
||
"""初始化非同步連線池"""
|
||
global _pool
|
||
try:
|
||
if _pool is None:
|
||
_pool = await asyncpg.create_pool(
|
||
database=DB_CONFIG["database"],
|
||
user=DB_CONFIG["user"],
|
||
password=DB_CONFIG["password"],
|
||
host=DB_CONFIG["host"],
|
||
port=int(DB_CONFIG["port"]),
|
||
min_size=1,
|
||
max_size=10
|
||
)
|
||
logger.info("✅ PostgreSQL Connection Pool Initialized.")
|
||
except Exception as e:
|
||
logger.error(f"❌ Failed to initialize DB Pool: {e}")
|
||
|
||
async def close_db_pool():
|
||
"""關閉非同步連線池"""
|
||
global _pool
|
||
if _pool:
|
||
await _pool.close()
|
||
_pool = None
|
||
logger.info("🔌 PostgreSQL Connection Pool Closed.")
|
||
|
||
async def get_pool() -> Optional[asyncpg.Pool]:
|
||
if _pool is None:
|
||
await init_db_pool()
|
||
return _pool
|
||
|
||
# ------------------------------------------
|
||
# CRUD Functions for cmts_options
|
||
# ------------------------------------------
|
||
|
||
async def upsert_leaf_option(host: str, path: str, data: dict) -> bool:
|
||
pool = await get_pool()
|
||
if not pool: return False
|
||
query = """
|
||
INSERT INTO cmts_options (host, path, data)
|
||
VALUES ($1, $2, $3)
|
||
ON CONFLICT (host, path)
|
||
DO UPDATE SET data = EXCLUDED.data, updated_at = CURRENT_TIMESTAMP;
|
||
"""
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
await conn.execute(query, host, path, json.dumps(data))
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"❌ DB Error (upsert_leaf_option): {e}")
|
||
return False
|
||
|
||
async def get_all_leaf_options(host: str) -> Optional[Dict[str, Any]]:
|
||
pool = await get_pool()
|
||
if not pool: return None
|
||
query = "SELECT path, data FROM cmts_options WHERE host = $1;"
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
records = await conn.fetch(query, host)
|
||
result = {}
|
||
for record in records:
|
||
data_val = record['data']
|
||
result[record['path']] = json.loads(data_val) if isinstance(data_val, str) else data_val
|
||
return result
|
||
except Exception as e:
|
||
logger.error(f"❌ DB Error (get_all_leaf_options): {e}")
|
||
return None
|
||
|
||
async def delete_leaf_options(host: str, paths: List[str]) -> int:
|
||
pool = await get_pool()
|
||
if not pool or not paths: return -1
|
||
query = "DELETE FROM cmts_options WHERE host = $1 AND path = ANY($2);"
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
status = await conn.execute(query, host, paths)
|
||
return int(status.split()[-1])
|
||
except Exception as e:
|
||
logger.error(f"❌ DB Error (delete_leaf_options): {e}")
|
||
return -1
|
||
|
||
# ------------------------------------------
|
||
# CRUD Functions for device_status
|
||
# ------------------------------------------
|
||
|
||
async def upsert_device_status(host: str, metadata: dict) -> bool:
|
||
pool = await get_pool()
|
||
if not pool: return False
|
||
cmts_version = metadata.get("cmts_version")
|
||
last_scanned = metadata.get("last_scanned", None)
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
if cmts_version and cmts_version != "unknown":
|
||
query = """
|
||
INSERT INTO device_status (host, cmts_version, last_scanned)
|
||
VALUES ($1, $2, $3)
|
||
ON CONFLICT (host)
|
||
DO UPDATE SET cmts_version = EXCLUDED.cmts_version, last_scanned = EXCLUDED.last_scanned, updated_at = CURRENT_TIMESTAMP;
|
||
"""
|
||
await conn.execute(query, host, cmts_version, last_scanned)
|
||
else:
|
||
query = """
|
||
INSERT INTO device_status (host, last_scanned)
|
||
VALUES ($1, $2)
|
||
ON CONFLICT (host)
|
||
DO UPDATE SET last_scanned = EXCLUDED.last_scanned, updated_at = CURRENT_TIMESTAMP;
|
||
"""
|
||
await conn.execute(query, host, last_scanned)
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"❌ DB Error (upsert_device_status): {e}")
|
||
return False
|
||
|
||
async def get_device_status(host: str) -> Optional[Dict[str, Any]]:
|
||
pool = await get_pool()
|
||
if not pool: return None
|
||
query = "SELECT cmts_version, last_scanned FROM device_status WHERE host = $1;"
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
record = await conn.fetchrow(query, host)
|
||
if record:
|
||
return {"cmts_version": record["cmts_version"], "last_scanned": record["last_scanned"]}
|
||
return None
|
||
except Exception as e:
|
||
logger.error(f"❌ DB Error (get_device_status): {e}")
|
||
return None
|
||
|
||
# ------------------------------------------
|
||
# CRUD Functions for system_filters (Tree Filters)
|
||
# ------------------------------------------
|
||
async def upsert_tree_filters(config_type: str, hidden_keys: List[str]) -> bool:
|
||
"""寫入全域的樹狀圖隱藏節點名單"""
|
||
pool = await get_pool()
|
||
if not pool:
|
||
return False
|
||
|
||
# Use 'global' as a dummy host to keep schema simple if needed,
|
||
# but a dedicated system_filters table is better.
|
||
query = """
|
||
INSERT INTO system_filters (config_type, hidden_keys)
|
||
VALUES ($1, $2)
|
||
ON CONFLICT (config_type)
|
||
DO UPDATE SET hidden_keys = EXCLUDED.hidden_keys, updated_at = CURRENT_TIMESTAMP;
|
||
"""
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
await conn.execute(query, config_type, hidden_keys)
|
||
return True
|
||
except asyncpg.PostgresError as e:
|
||
logger.error(f"❌ DB Error (upsert_tree_filters): {e}")
|
||
return False
|
||
except Exception as e:
|
||
logger.error(f"❌ Unknown Error (upsert_tree_filters): {e}")
|
||
return False
|
||
|
||
async def get_tree_filters(config_type: str) -> Optional[List[str]]:
|
||
"""讀取全域的樹狀圖隱藏節點名單"""
|
||
pool = await get_pool()
|
||
if not pool:
|
||
return None
|
||
|
||
query = """
|
||
SELECT hidden_keys FROM system_filters
|
||
WHERE config_type = $1;
|
||
"""
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
record = await conn.fetchrow(query, config_type)
|
||
if record:
|
||
return record["hidden_keys"]
|
||
return []
|
||
except asyncpg.PostgresError as e:
|
||
logger.error(f"❌ DB Error (get_tree_filters): {e}")
|
||
return None
|
||
except Exception as e:
|
||
logger.error(f"❌ Unknown Error (get_tree_filters): {e}")
|
||
return None
|
||
|
||
# ==========================================
|
||
# CRUD Functions for config_backups (Phase 1 & 2)
|
||
# ==========================================
|
||
|
||
async def insert_config_backup(
|
||
host: str,
|
||
config_type: str,
|
||
raw_cli: str,
|
||
parsed_tree: dict,
|
||
snapshot_name: Optional[str] = None,
|
||
description: str = "", # 🟢 新增描述參數 (預設為空字串)
|
||
is_auto: bool = False
|
||
) -> Optional[str]:
|
||
"""新增一筆設備配置備份，回傳產生的 Backup ID"""
|
||
pool = await get_pool()
|
||
if not pool:
|
||
return None
|
||
|
||
backup_id = str(uuid.uuid4())
|
||
|
||
# 🟢 SQL 語句加入 description 與對應的 $5
|
||
query = """
|
||
INSERT INTO config_backups
|
||
(id, host, config_type, snapshot_name, description, is_auto, raw_cli, parsed_tree)
|
||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8::jsonb)
|
||
"""
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
await conn.execute(
|
||
query,
|
||
backup_id,
|
||
host,
|
||
config_type,
|
||
snapshot_name,
|
||
description, # 🟢 傳入描述參數
|
||
is_auto,
|
||
raw_cli,
|
||
json.dumps(parsed_tree)
|
||
)
|
||
return backup_id
|
||
except asyncpg.PostgresError as e:
|
||
logger.error(f"❌ DB Error (insert_config_backup): {e}")
|
||
return None
|
||
except Exception as e:
|
||
logger.error(f"❌ Unknown Error (insert_config_backup): {e}")
|
||
return None
|
||
|
||
async def get_config_backup_list(host: str, config_type: str) -> Optional[List[Dict[str, Any]]]:
|
||
"""取得歷史快照列表 (支援 config_type='all' 撈取全部，包含 is_pinned 狀態)"""
|
||
pool = await get_pool()
|
||
if not pool:
|
||
return None
|
||
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
if config_type == "all":
|
||
# 🌟 SQL 查詢補上 is_pinned
|
||
query = """
|
||
SELECT id, host, config_type, timestamp, snapshot_name, description, is_auto, is_pinned
|
||
FROM config_backups
|
||
WHERE host = $1
|
||
ORDER BY timestamp DESC;
|
||
"""
|
||
records = await conn.fetch(query, host)
|
||
else:
|
||
# 🌟 SQL 查詢補上 is_pinned
|
||
query = """
|
||
SELECT id, host, config_type, timestamp, snapshot_name, description, is_auto, is_pinned
|
||
FROM config_backups
|
||
WHERE host = $1 AND config_type = $2
|
||
ORDER BY timestamp DESC;
|
||
"""
|
||
records = await conn.fetch(query, host, config_type)
|
||
|
||
return [
|
||
{
|
||
"id": str(r["id"]),
|
||
"host": r["host"],
|
||
"config_type": r["config_type"],
|
||
"timestamp": r["timestamp"].isoformat(),
|
||
"snapshot_name": r["snapshot_name"],
|
||
"description": r["description"],
|
||
"is_auto": r["is_auto"],
|
||
"is_pinned": r["is_pinned"] # 🌟 放入回傳字典
|
||
}
|
||
for r in records
|
||
]
|
||
except asyncpg.PostgresError as e:
|
||
logger.error(f"❌ DB Error (get_config_backup_list): {e}")
|
||
return None
|
||
except Exception as e:
|
||
logger.error(f"❌ Unknown Error (get_config_backup_list): {e}")
|
||
return None
|
||
|
||
async def get_config_backup_detail(backup_id: str) -> Optional[Dict[str, Any]]:
|
||
"""取得特定快照的完整內容 (包含 parsed_tree，供 Phase 3 還原使用)"""
|
||
pool = await get_pool()
|
||
if not pool:
|
||
return None
|
||
|
||
# 🟢 SELECT 加入 description
|
||
query = """
|
||
SELECT id, host, timestamp, snapshot_name, description, is_auto, parsed_tree
|
||
FROM config_backups
|
||
WHERE id = $1;
|
||
"""
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
record = await conn.fetchrow(query, backup_id)
|
||
if record:
|
||
data_val = record['parsed_tree']
|
||
parsed_tree = json.loads(data_val) if isinstance(data_val, str) else data_val
|
||
|
||
return {
|
||
"id": str(record["id"]),
|
||
"host": record["host"],
|
||
"timestamp": record["timestamp"].isoformat(),
|
||
"snapshot_name": record["snapshot_name"],
|
||
"description": record["description"], # 🟢 將資料庫的描述放入回傳字典
|
||
"is_auto": record["is_auto"],
|
||
"parsed_tree": parsed_tree
|
||
}
|
||
return None
|
||
except asyncpg.PostgresError as e:
|
||
logger.error(f"❌ DB Error (get_config_backup_detail): {e}")
|
||
return None
|
||
except Exception as e:
|
||
logger.error(f"❌ Unknown Error (get_config_backup_detail): {e}")
|
||
return None
|
||
|
||
async def delete_config_backup(backup_id: str) -> bool:
|
||
"""刪除指定的快照"""
|
||
pool = await get_pool()
|
||
if not pool:
|
||
return False
|
||
|
||
query = "DELETE FROM config_backups WHERE id = $1;"
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
status = await conn.execute(query, backup_id)
|
||
# status 會是 'DELETE 1' 或 'DELETE 0'
|
||
return int(status.split()[-1]) > 0
|
||
except asyncpg.PostgresError as e:
|
||
logger.error(f"❌ DB Error (delete_config_backup): {e}")
|
||
return False
|
||
except Exception as e:
|
||
logger.error(f"❌ Unknown Error (delete_config_backup): {e}")
|
||
return False
|
||
|
||
# ============================================================================
|
||
# 🧹 備份滾動淘汰機制 (Retention Policy)
|
||
# ============================================================================
|
||
|
||
async def cleanup_old_backups(pool, host: str) -> None:
|
||
"""
|
||
執行手動備份的滾動淘汰機制 (Retention Policy)
|
||
- 每個設備 (host) 保留最新 20 筆未釘選 (is_pinned=FALSE) 的手動備份
|
||
- 釘選的備份 (is_pinned=True) 永遠不刪除
|
||
"""
|
||
if not pool:
|
||
logger.warning("⚠️ [Retention] 無法執行備份清理：Database Pool 未初始化。")
|
||
return
|
||
|
||
# 使用 PostgreSQL CTE (Common Table Expression) 語法
|
||
# 先選出該設備最新 20 筆需要保留的 ID，再將其餘未釘選的舊備份一次性刪除
|
||
query = """
|
||
WITH kept_backups AS (
|
||
SELECT id FROM config_backups
|
||
WHERE host = $1 AND is_pinned = FALSE
|
||
ORDER BY timestamp DESC
|
||
LIMIT 20
|
||
)
|
||
DELETE FROM config_backups
|
||
WHERE host = $1 AND is_pinned = FALSE
|
||
AND id NOT IN (SELECT id FROM kept_backups)
|
||
RETURNING id;
|
||
"""
|
||
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
deleted_records = await conn.fetch(query, host)
|
||
total_deleted = len(deleted_records)
|
||
|
||
if total_deleted > 0:
|
||
logger.info(f"🧹 [Retention Policy] 設備 {host} 清理完畢：已自動淘汰 {total_deleted} 筆過期手動備份。")
|
||
else:
|
||
logger.debug(f"ℹ️ [Retention Policy] 設備 {host} 備份數量未達 20 筆上限，無需清理。")
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ [Retention Policy] 清理設備 {host} 過期備份時發生錯誤: {e}", exc_info=True)
|
||
|
||
# ============================================================================
|
||
# 📌 釘選防護與容量指標計算 (Pinning & Metrics)
|
||
# ============================================================================
|
||
|
||
async def toggle_config_backup_pin(backup_id: str) -> Optional[bool]:
|
||
"""
|
||
切換特定備份的釘選狀態 (True -> False, False -> True)
|
||
回傳更新後的 is_pinned 狀態，若失敗則回傳 None
|
||
"""
|
||
pool = await get_pool()
|
||
if not pool:
|
||
return None
|
||
|
||
# 使用 PostgreSQL 的 UPDATE ... RETURNING 語法，一步完成切換與讀取，保證原子性
|
||
query = """
|
||
UPDATE config_backups
|
||
SET is_pinned = NOT is_pinned
|
||
WHERE id = $1
|
||
RETURNING is_pinned;
|
||
"""
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
new_state = await conn.fetchval(query, backup_id)
|
||
return new_state
|
||
except Exception as e:
|
||
logger.error(f"❌ DB Error (toggle_config_backup_pin): {e}")
|
||
return None
|
||
|
||
|
||
async def get_backup_metrics(pool, host: str) -> dict:
|
||
"""
|
||
計算特定設備的備份容量指標
|
||
- total: 總備份數 (含釘選與未釘選)
|
||
- pinned: 已釘選保護的數量
|
||
- unpinned: 未釘選的數量 (上限為 20 筆)
|
||
- remaining: 剩餘可用手動備份額度 (20 - unpinned)
|
||
"""
|
||
if not pool:
|
||
return {"total": 0, "pinned": 0, "unpinned": 0, "remaining": 20}
|
||
|
||
query = """
|
||
SELECT
|
||
COUNT(*) as total,
|
||
COUNT(*) FILTER (WHERE is_pinned = TRUE) as pinned,
|
||
COUNT(*) FILTER (WHERE is_pinned = FALSE) as unpinned
|
||
FROM config_backups
|
||
WHERE host = $1;
|
||
"""
|
||
try:
|
||
async with pool.acquire() as conn:
|
||
r = await conn.fetchrow(query, host)
|
||
total = r["total"] or 0
|
||
pinned = r["pinned"] or 0
|
||
unpinned = r["unpinned"] or 0
|
||
remaining = max(0, 20 - unpinned)
|
||
|
||
return {
|
||
"total": total,
|
||
"pinned": pinned,
|
||
"unpinned": unpinned,
|
||
"remaining": remaining
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"❌ DB Error (get_backup_metrics): {e}")
|
||
return {"total": 0, "pinned": 0, "unpinned": 0, "remaining": 20} |