dev_1.0.0_yuyr_2：重新提交 PR，增加 master/agent 以及系统集成测试 (#17)
Reviewed-on: #17 Reviewed-by: sundapeng <sundp@mail.zgclab.edu.cn> Reviewed-by: xuxt <xuxt@zgclab.edu.cn>
This commit was merged in pull request #17.
This commit is contained in:
332
src/master/app/storage.py
Normal file
332
src/master/app/storage.py
Normal file
@@ -0,0 +1,332 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
import threading
|
||||
from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple
|
||||
|
||||
from .models import serialize_node_row, serialize_node_summary
|
||||
from .util import ensure_parent, to_iso, utcnow
|
||||
|
||||
|
||||
class Storage:
|
||||
def __init__(self, db_path: str, node_id_prefix: str) -> None:
|
||||
self._db_path = db_path
|
||||
self._node_id_prefix = node_id_prefix
|
||||
ensure_parent(db_path)
|
||||
self._lock = threading.Lock()
|
||||
self._conn = sqlite3.connect(db_path, detect_types=sqlite3.PARSE_DECLTYPES, check_same_thread=False)
|
||||
self._conn.row_factory = sqlite3.Row
|
||||
with self._lock:
|
||||
self._conn.execute("PRAGMA foreign_keys = ON;")
|
||||
self._ensure_schema()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# schema & helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _ensure_schema(self) -> None:
|
||||
"""初始化表结构，确保服务启动时数据库结构就绪。"""
|
||||
with self._lock:
|
||||
self._conn.executescript(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS nodes (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL UNIQUE,
|
||||
type TEXT NOT NULL,
|
||||
version TEXT,
|
||||
status TEXT NOT NULL,
|
||||
config_json TEXT,
|
||||
labels_json TEXT,
|
||||
meta_json TEXT,
|
||||
health_json TEXT,
|
||||
register_time TEXT,
|
||||
last_report TEXT,
|
||||
agent_last_report TEXT,
|
||||
last_updated TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS kv (
|
||||
key TEXT PRIMARY KEY,
|
||||
value TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_nodes_status ON nodes(status);
|
||||
CREATE INDEX IF NOT EXISTS idx_nodes_name ON nodes(name);
|
||||
"""
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
self._conn.close()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Node ID allocation
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def allocate_node_id(self) -> str:
|
||||
"""在 kv 表里维护自增序列，为新节点生成形如 A1 的 ID。"""
|
||||
with self._lock:
|
||||
cur = self._conn.execute("SELECT value FROM kv WHERE key = ?", ("node_id_seq",))
|
||||
row = cur.fetchone()
|
||||
if row is None:
|
||||
next_id = 1
|
||||
self._conn.execute("INSERT INTO kv(key, value) VALUES(?, ?)", ("node_id_seq", str(next_id)))
|
||||
else:
|
||||
next_id = int(row["value"]) + 1
|
||||
self._conn.execute("UPDATE kv SET value = ? WHERE key = ?", (str(next_id), "node_id_seq"))
|
||||
self._conn.commit()
|
||||
return f"{self._node_id_prefix}{next_id}"
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Query helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def list_nodes(self) -> List[Dict[str, Any]]:
|
||||
with self._lock:
|
||||
cur = self._conn.execute(
|
||||
"SELECT id, name, status, type, version FROM nodes ORDER BY id ASC"
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
return [serialize_node_summary(row) for row in rows]
|
||||
|
||||
def get_node(self, node_id: str) -> Optional[Dict[str, Any]]:
|
||||
with self._lock:
|
||||
cur = self._conn.execute("SELECT * FROM nodes WHERE id = ?", (node_id,))
|
||||
row = cur.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return serialize_node_row(row)
|
||||
|
||||
def get_node_raw(self, node_id: str) -> Optional[sqlite3.Row]:
|
||||
with self._lock:
|
||||
cur = self._conn.execute("SELECT * FROM nodes WHERE id = ?", (node_id,))
|
||||
row = cur.fetchone()
|
||||
return row
|
||||
|
||||
def get_node_by_name(self, name: str) -> Optional[Dict[str, Any]]:
|
||||
with self._lock:
|
||||
cur = self._conn.execute("SELECT * FROM nodes WHERE name = ?", (name,))
|
||||
row = cur.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return serialize_node_row(row)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Mutation helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def create_node(
|
||||
self,
|
||||
node_id: str,
|
||||
name: str,
|
||||
node_type: str,
|
||||
version: str | None,
|
||||
meta_data: Mapping[str, Any],
|
||||
status: str,
|
||||
register_time_iso: str,
|
||||
last_updated_iso: str,
|
||||
) -> Dict[str, Any]:
|
||||
"""插入节点初始记录，默认 config/label/health 为空。"""
|
||||
now_iso = last_updated_iso
|
||||
with self._lock:
|
||||
self._conn.execute(
|
||||
"""
|
||||
INSERT INTO nodes (
|
||||
id, name, type, version, status, config_json, labels_json, meta_json,
|
||||
health_json, register_time, last_report, agent_last_report, last_updated
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
node_id,
|
||||
name,
|
||||
node_type,
|
||||
version,
|
||||
status,
|
||||
json.dumps({}),
|
||||
json.dumps([]),
|
||||
json.dumps(dict(meta_data)),
|
||||
json.dumps({}),
|
||||
register_time_iso,
|
||||
None,
|
||||
None,
|
||||
now_iso,
|
||||
),
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
created = self.get_node(node_id)
|
||||
if created is None:
|
||||
raise RuntimeError("Failed to read back created node")
|
||||
return created
|
||||
|
||||
def update_node_meta(
|
||||
self,
|
||||
node_id: str,
|
||||
*,
|
||||
name: Optional[str] = None,
|
||||
node_type: Optional[str] = None,
|
||||
version: Optional[str | None] = None,
|
||||
meta_data: Optional[Mapping[str, Any]] = None,
|
||||
last_updated_iso: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""重注册时更新节点静态信息，缺省字段保持不变。"""
|
||||
updates: List[str] = []
|
||||
params: List[Any] = []
|
||||
if name is not None:
|
||||
updates.append("name = ?")
|
||||
params.append(name)
|
||||
if node_type is not None:
|
||||
updates.append("type = ?")
|
||||
params.append(node_type)
|
||||
if version is not None:
|
||||
updates.append("version = ?")
|
||||
params.append(version)
|
||||
if meta_data is not None:
|
||||
updates.append("meta_json = ?")
|
||||
params.append(json.dumps(dict(meta_data)))
|
||||
if last_updated_iso is not None:
|
||||
updates.append("last_updated = ?")
|
||||
params.append(last_updated_iso)
|
||||
|
||||
if not updates:
|
||||
result = self.get_node(node_id)
|
||||
if result is None:
|
||||
raise KeyError(node_id)
|
||||
return result
|
||||
|
||||
params.append(node_id)
|
||||
with self._lock:
|
||||
self._conn.execute(
|
||||
f"UPDATE nodes SET {', '.join(updates)} WHERE id = ?",
|
||||
tuple(params),
|
||||
)
|
||||
self._conn.commit()
|
||||
updated = self.get_node(node_id)
|
||||
if updated is None:
|
||||
raise KeyError(node_id)
|
||||
return updated
|
||||
|
||||
def update_config_and_labels(
|
||||
self, node_id: str, *, config: Optional[Mapping[str, Any]] = None, labels: Optional[Iterable[str]] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""部分更新 config/label，并刷新 last_updated 时间戳。"""
|
||||
updates: List[str] = []
|
||||
params: List[Any] = []
|
||||
if config is not None:
|
||||
updates.append("config_json = ?")
|
||||
params.append(json.dumps(dict(config)))
|
||||
if labels is not None:
|
||||
updates.append("labels_json = ?")
|
||||
params.append(json.dumps(list(labels)))
|
||||
updates.append("last_updated = ?")
|
||||
params.append(to_iso(utcnow()))
|
||||
params.append(node_id)
|
||||
with self._lock:
|
||||
self._conn.execute(
|
||||
f"UPDATE nodes SET {', '.join(updates)} WHERE id = ?",
|
||||
tuple(params),
|
||||
)
|
||||
if self._conn.total_changes == 0:
|
||||
self._conn.rollback()
|
||||
raise KeyError(node_id)
|
||||
self._conn.commit()
|
||||
updated = self.get_node(node_id)
|
||||
if updated is None:
|
||||
raise KeyError(node_id)
|
||||
return updated
|
||||
|
||||
def update_last_report(
|
||||
self,
|
||||
node_id: str,
|
||||
*,
|
||||
server_timestamp_iso: str,
|
||||
agent_timestamp_iso: str,
|
||||
health: Mapping[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
"""记录最新上报时间和健康信息，用于后续状态计算。"""
|
||||
with self._lock:
|
||||
self._conn.execute(
|
||||
"""
|
||||
UPDATE nodes
|
||||
SET last_report = ?,
|
||||
agent_last_report = ?,
|
||||
health_json = ?,
|
||||
last_updated = ?
|
||||
WHERE id = ?
|
||||
""",
|
||||
(
|
||||
server_timestamp_iso,
|
||||
agent_timestamp_iso,
|
||||
json.dumps(health),
|
||||
server_timestamp_iso,
|
||||
node_id,
|
||||
),
|
||||
)
|
||||
if self._conn.total_changes == 0:
|
||||
self._conn.rollback()
|
||||
raise KeyError(node_id)
|
||||
self._conn.commit()
|
||||
updated = self.get_node(node_id)
|
||||
if updated is None:
|
||||
raise KeyError(node_id)
|
||||
return updated
|
||||
|
||||
def update_status(self, node_id: str, status: str, *, last_updated_iso: str) -> None:
|
||||
with self._lock:
|
||||
self._conn.execute(
|
||||
"UPDATE nodes SET status = ?, last_updated = ? WHERE id = ?",
|
||||
(status, last_updated_iso, node_id),
|
||||
)
|
||||
self._conn.commit()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Reporting helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""统计节点总数及按状态聚合的数量。"""
|
||||
with self._lock:
|
||||
cur = self._conn.execute("SELECT COUNT(*) AS total FROM nodes")
|
||||
total_row = cur.fetchone()
|
||||
cur = self._conn.execute("SELECT status, COUNT(*) AS count FROM nodes GROUP BY status")
|
||||
status_rows = cur.fetchall()
|
||||
return {
|
||||
"total": total_row["total"] if total_row else 0,
|
||||
"status_statistics": [
|
||||
{"status": row["status"], "count": row["count"]}
|
||||
for row in status_rows
|
||||
],
|
||||
}
|
||||
|
||||
def fetch_nodes_for_scheduler(self) -> List[sqlite3.Row]:
|
||||
with self._lock:
|
||||
cur = self._conn.execute(
|
||||
"SELECT id, last_report, status FROM nodes"
|
||||
)
|
||||
return cur.fetchall()
|
||||
|
||||
def get_online_nodes(self) -> List[Dict[str, Any]]:
|
||||
"""返回在线节点列表，用于生成 nodes.json。"""
|
||||
with self._lock:
|
||||
cur = self._conn.execute(
|
||||
"SELECT id, meta_json, labels_json, name FROM nodes WHERE status = ? ORDER BY id ASC",
|
||||
("online",),
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
|
||||
result: List[Dict[str, Any]] = []
|
||||
for row in rows:
|
||||
meta = json.loads(row["meta_json"]) if row["meta_json"] else {}
|
||||
labels = json.loads(row["labels_json"]) if row["labels_json"] else []
|
||||
result.append(
|
||||
{
|
||||
"node_id": row["id"],
|
||||
"user_id": meta.get("user"),
|
||||
"ip": meta.get("ip"),
|
||||
"hostname": meta.get("hostname", row["name"]),
|
||||
"labels": labels if isinstance(labels, list) else [],
|
||||
}
|
||||
)
|
||||
return result
|
||||
Reference in New Issue
Block a user