244 lines
8.2 KiB
Python
244 lines
8.2 KiB
Python
import asyncpg
|
|
import json
|
|
import logging
|
|
from typing import Dict, List, Optional, Any
|
|
|
|
# ==========================================
|
|
# 💡 PostgreSQL 連線與操作 (高可用性版)
|
|
# ==========================================
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# DB 設定 (從您的 init_db.py 中提取)
|
|
DB_CONFIG = {
|
|
"database": "cmts_nms",
|
|
"user": "swpa",
|
|
"password": "swpa4920",
|
|
"host": "127.0.0.1",
|
|
"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, config_type: str, path: str, data: dict) -> bool:
|
|
"""將選項寫入或更新至資料庫"""
|
|
pool = await get_pool()
|
|
if not pool:
|
|
return False
|
|
|
|
query = """
|
|
INSERT INTO cmts_options (host, config_type, path, data)
|
|
VALUES ($1, $2, $3, $4)
|
|
ON CONFLICT (host, config_type, path)
|
|
DO UPDATE SET data = EXCLUDED.data, updated_at = CURRENT_TIMESTAMP;
|
|
"""
|
|
try:
|
|
async with pool.acquire() as conn:
|
|
await conn.execute(query, host, config_type, path, json.dumps(data))
|
|
return True
|
|
except asyncpg.PostgresError as e:
|
|
logger.error(f"❌ DB Error (upsert_leaf_option): {e}")
|
|
return False
|
|
except Exception as e:
|
|
logger.error(f"❌ Unknown Error (upsert_leaf_option): {e}")
|
|
return False
|
|
|
|
async def get_all_leaf_options(host: str, config_type: str) -> Optional[Dict[str, Any]]:
|
|
"""取得特定 host 與 config_type 的所有選項"""
|
|
pool = await get_pool()
|
|
if not pool:
|
|
return None
|
|
|
|
query = """
|
|
SELECT path, data FROM cmts_options
|
|
WHERE host = $1 AND config_type = $2;
|
|
"""
|
|
try:
|
|
async with pool.acquire() as conn:
|
|
records = await conn.fetch(query, host, config_type)
|
|
result = {}
|
|
for record in records:
|
|
# asyncpg returns strings for JSON if not explicitly configured with type mapping,
|
|
# but usually it's fine to just json.loads it.
|
|
data_val = record['data']
|
|
if isinstance(data_val, str):
|
|
result[record['path']] = json.loads(data_val)
|
|
else:
|
|
result[record['path']] = data_val
|
|
return result
|
|
except asyncpg.PostgresError as e:
|
|
logger.error(f"❌ DB Error (get_all_leaf_options): {e}")
|
|
return None
|
|
except Exception as e:
|
|
logger.error(f"❌ Unknown Error (get_all_leaf_options): {e}")
|
|
return None
|
|
|
|
async def delete_leaf_options(host: str, config_type: 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 config_type = $2 AND path = ANY($3);
|
|
"""
|
|
try:
|
|
async with pool.acquire() as conn:
|
|
# Execute returns the command tag, e.g., 'DELETE 5'
|
|
status = await conn.execute(query, host, config_type, paths)
|
|
deleted_count = int(status.split()[-1])
|
|
return deleted_count
|
|
except asyncpg.PostgresError as e:
|
|
logger.error(f"❌ DB Error (delete_leaf_options): {e}")
|
|
return -1
|
|
except Exception as e:
|
|
logger.error(f"❌ Unknown Error (delete_leaf_options): {e}")
|
|
return -1
|
|
|
|
# ------------------------------------------
|
|
# CRUD Functions for device_status
|
|
# ------------------------------------------
|
|
|
|
async def upsert_device_status(host: str, config_type: str, metadata: dict) -> bool:
|
|
"""更新設備的 metadata"""
|
|
pool = await get_pool()
|
|
if not pool:
|
|
return False
|
|
|
|
# Extract known fields
|
|
cmts_version = metadata.get("cmts_version", "unknown")
|
|
last_scanned = metadata.get("last_scanned", None)
|
|
|
|
query = """
|
|
INSERT INTO device_status (host, config_type, cmts_version, last_scanned)
|
|
VALUES ($1, $2, $3, $4)
|
|
ON CONFLICT (host, config_type)
|
|
DO UPDATE SET
|
|
cmts_version = EXCLUDED.cmts_version,
|
|
last_scanned = EXCLUDED.last_scanned,
|
|
updated_at = CURRENT_TIMESTAMP;
|
|
"""
|
|
try:
|
|
async with pool.acquire() as conn:
|
|
await conn.execute(query, host, config_type, cmts_version, last_scanned)
|
|
return True
|
|
except asyncpg.PostgresError as e:
|
|
logger.error(f"❌ DB Error (upsert_device_status): {e}")
|
|
return False
|
|
except Exception as e:
|
|
logger.error(f"❌ Unknown Error (upsert_device_status): {e}")
|
|
return False
|
|
|
|
async def get_device_status(host: str, config_type: str) -> Optional[Dict[str, Any]]:
|
|
"""獲取設備 metadata"""
|
|
pool = await get_pool()
|
|
if not pool:
|
|
return None
|
|
|
|
query = """
|
|
SELECT cmts_version, last_scanned FROM device_status
|
|
WHERE host = $1 AND config_type = $2;
|
|
"""
|
|
try:
|
|
async with pool.acquire() as conn:
|
|
record = await conn.fetchrow(query, host, config_type)
|
|
if record:
|
|
return {
|
|
"cmts_version": record["cmts_version"],
|
|
"last_scanned": record["last_scanned"]
|
|
}
|
|
return None
|
|
except asyncpg.PostgresError as e:
|
|
logger.error(f"❌ DB Error (get_device_status): {e}")
|
|
return None
|
|
except Exception as e:
|
|
logger.error(f"❌ Unknown 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
|