264 lines
9.5 KiB
Python
264 lines
9.5 KiB
Python
"""Database administration: export / import / format detection.
|
|
|
|
Only Root may call these (enforced at the route layer). All operations are
|
|
atomic: an import runs inside a single transaction — any failing row rolls the
|
|
whole import back, leaving the current database untouched.
|
|
|
|
Formats:
|
|
- ``sqlite`` — a legacy SQLite database file (magic bytes ``SQLite format 3``).
|
|
Tables not present in the current SCHEMA are dropped (e.g. camera_config);
|
|
the legacy single-user `auth_config` row (keyed by id=1) is attributed to
|
|
Root (user 0); missing columns are filled with their SCHEMA DEFAULT.
|
|
- ``json`` — the self-contained backup produced by :func:`export_database`
|
|
(``{"format": "iot-dboard-backup", ...}``).
|
|
"""
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import tempfile
|
|
from datetime import datetime, timezone
|
|
|
|
from backend.models import (
|
|
get_db, SCHEMA, DB_VERSION, _CREATE_TABLE_RE, _table_columns,
|
|
_read_db_version, _ensure_defaults, _fix_identity_sequences,
|
|
)
|
|
|
|
# Tables whose contents are exported/imported (everything except the schema
|
|
# version row and transient session tokens).
|
|
_EXCLUDE_TABLES = {"meta", "auth_tokens"}
|
|
|
|
|
|
def _schema_table_names() -> list:
|
|
return [m.group(1) for m in _CREATE_TABLE_RE.finditer(SCHEMA)]
|
|
|
|
|
|
def detect_format(data: bytes) -> str:
|
|
"""Detect the database file format: 'sqlite', 'json' or 'unknown'."""
|
|
if data[:16] == b"SQLite format 3\x00":
|
|
return "sqlite"
|
|
try:
|
|
obj = json.loads(data.decode("utf-8"))
|
|
if isinstance(obj, dict) and obj.get("format") == "iot-dboard-backup":
|
|
return "json"
|
|
except Exception:
|
|
pass
|
|
return "unknown"
|
|
|
|
|
|
def export_database(db_path: str) -> dict:
|
|
"""Snapshot every data table (except auth_tokens) into a JSON-serializable dict."""
|
|
conn = get_db(db_path)
|
|
try:
|
|
tables = _schema_table_names()
|
|
out = {
|
|
"format": "iot-dboard-backup",
|
|
"backup_version": 1,
|
|
"db_version": _read_db_version(conn) or DB_VERSION,
|
|
"exported_at": datetime.now(timezone.utc).isoformat(),
|
|
"tables": {},
|
|
}
|
|
for t in tables:
|
|
if t in _EXCLUDE_TABLES:
|
|
continue
|
|
cols = _table_columns(conn, t)
|
|
if not cols:
|
|
continue
|
|
rows = conn.execute(
|
|
f"SELECT {', '.join(cols)} FROM {t} ORDER BY 1"
|
|
).fetchall()
|
|
out["tables"][t] = {
|
|
"columns": cols,
|
|
"rows": [[r[c] for c in cols] for r in rows],
|
|
}
|
|
return out
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def export_database_json(db_path: str) -> str:
|
|
"""Export as a compact JSON string (files are small; no pretty printing)."""
|
|
return json.dumps(export_database(db_path), ensure_ascii=False, default=str)
|
|
|
|
|
|
# ─── SQLite import ───────────────────────────
|
|
|
|
def _sqlite_tables(conn_old) -> list:
|
|
return [
|
|
r["name"] for r in conn_old.execute(
|
|
"SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
|
|
).fetchall()
|
|
]
|
|
|
|
|
|
def _sqlite_rows(conn_old, table: str, cols: list) -> list:
|
|
if not cols:
|
|
return []
|
|
col_str = ", ".join(cols)
|
|
return conn_old.execute(f"SELECT {col_str} FROM {table}").fetchall()
|
|
|
|
|
|
def _read_sqlite(data: bytes) -> dict:
|
|
"""Read a legacy SQLite file into {table: {"columns": [...], "rows": [...]}}.
|
|
auth_config special case: legacy single-user row (id=1) belongs to Root."""
|
|
tmp_path = None
|
|
try:
|
|
fd, tmp_path = tempfile.mkstemp(suffix=".db")
|
|
with os.fdopen(fd, "wb") as f:
|
|
f.write(data)
|
|
conn_old = sqlite3.connect(f"file:{tmp_path}?mode=ro", uri=True)
|
|
conn_old.row_factory = sqlite3.Row
|
|
try:
|
|
tables = {}
|
|
for t in _sqlite_tables(conn_old):
|
|
cols = [r["name"] for r in conn_old.execute(f"PRAGMA table_info({t})").fetchall()]
|
|
if not cols:
|
|
continue
|
|
rows = []
|
|
for row in _sqlite_rows(conn_old, t, cols):
|
|
rows.append([row[c] for c in cols])
|
|
tables[t] = {"columns": cols, "rows": rows}
|
|
return tables
|
|
finally:
|
|
conn_old.close()
|
|
finally:
|
|
if tmp_path and os.path.exists(tmp_path):
|
|
os.unlink(tmp_path)
|
|
|
|
|
|
def _import_table(conn, table: str, data: dict, where: str, skipped: list) -> int:
|
|
"""Insert rows of one table into the target PG database (strict mode).
|
|
Returns the number of rows inserted."""
|
|
src_cols = data.get("columns") or []
|
|
rows = data.get("rows") or []
|
|
if not src_cols or not rows:
|
|
return 0
|
|
dst_cols = _table_columns(conn, table)
|
|
if not dst_cols:
|
|
skipped.append(table) # table not in schema — dropped (e.g. camera_config)
|
|
return 0
|
|
|
|
if table == "auth_config":
|
|
return _import_auth_config(conn, data, where)
|
|
|
|
# Only common columns are copied; missing target columns get their DEFAULT.
|
|
common = [c for c in dst_cols if c in src_cols]
|
|
if not common:
|
|
skipped.append(table)
|
|
return 0
|
|
col_str = ", ".join(common)
|
|
placeholders = ", ".join("?" * len(common))
|
|
sql = f"INSERT INTO {table} ({col_str}) VALUES ({placeholders})"
|
|
n = 0
|
|
for i, row in enumerate(rows):
|
|
vals = []
|
|
for c in common:
|
|
vals.append(row[src_cols.index(c)])
|
|
try:
|
|
conn.execute(sql, tuple(vals))
|
|
n += 1
|
|
except Exception as e:
|
|
raise ValueError(
|
|
f"{where}: 表 {table} 第 {i + 1} 行导入失败: {e}"
|
|
) from e
|
|
return n
|
|
|
|
|
|
def _import_auth_config(conn, data: dict, where: str) -> int:
|
|
"""Special case: the legacy auth_config table is keyed by `id` (single user,
|
|
id=1); the new table is keyed by `user_id`. The legacy row belongs to Root."""
|
|
src_cols = data.get("columns") or []
|
|
rows = data.get("rows") or []
|
|
if not src_cols or not rows:
|
|
return 0
|
|
dst_cols = _table_columns(conn, "auth_config")
|
|
n = 0
|
|
for i, row in enumerate(rows):
|
|
d = {c: row[src_cols.index(c)] for c in src_cols}
|
|
if "user_id" in src_cols:
|
|
uid = d.get("user_id") or 0
|
|
else:
|
|
old_id = d.get("id", 1)
|
|
uid = 0 if old_id == 1 else (old_id or 0)
|
|
common = [c for c in dst_cols if c in src_cols and c not in ("id", "user_id")]
|
|
if not common:
|
|
continue
|
|
col_str = ", ".join(common)
|
|
placeholders = ", ".join("?" * len(common))
|
|
sql = f"INSERT OR IGNORE INTO auth_config (user_id, {col_str}) VALUES (?, {placeholders})"
|
|
try:
|
|
conn.execute(sql, (uid, *[d[c] for c in common]))
|
|
n += 1
|
|
except Exception as e:
|
|
raise ValueError(
|
|
f"{where}: auth_config 第 {i + 1} 行导入失败: {e}"
|
|
) from e
|
|
return n
|
|
|
|
|
|
# ─── Restore core (shared by sqlite + json imports) ─────
|
|
|
|
def _restore(db_path: str, tables: dict, where: str) -> dict:
|
|
"""Atomic restore: TRUNCATE every data table, insert the given rows, then
|
|
re-bootstrap defaults and stamp the schema version. Any failure rolls back
|
|
the whole import (the shim rolls back on error, keeping it atomic)."""
|
|
conn = get_db(db_path)
|
|
stats: dict[str, int] = {}
|
|
skipped: list = []
|
|
try:
|
|
# 1. Wipe current data (TRUNCATE is transactional in PostgreSQL).
|
|
all_tables = _schema_table_names()
|
|
truncate_list = [t for t in all_tables if t != "meta"]
|
|
if truncate_list:
|
|
conn.execute("TRUNCATE TABLE " + ", ".join(truncate_list) + " CASCADE")
|
|
|
|
# 2. Insert data (strict: first failing row aborts everything).
|
|
for t in all_tables:
|
|
if t in _EXCLUDE_TABLES:
|
|
continue
|
|
if t in tables:
|
|
stats[t] = _import_table(conn, t, tables[t], where, skipped)
|
|
|
|
# 3. Re-bootstrap defaults (Root user, config rows) + stamp version.
|
|
_ensure_defaults(conn)
|
|
conn.execute("UPDATE meta SET version=? WHERE id=1", (DB_VERSION,))
|
|
conn.execute("INSERT OR IGNORE INTO meta (id, version) VALUES (1, ?)", (DB_VERSION,))
|
|
_fix_identity_sequences(conn)
|
|
conn.commit()
|
|
if skipped:
|
|
stats["_skipped_tables"] = skipped
|
|
return stats
|
|
except Exception:
|
|
try:
|
|
conn.rollback()
|
|
except Exception:
|
|
pass
|
|
raise
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def import_sqlite(db_path: str, data: bytes) -> dict:
|
|
"""Import a legacy SQLite database file (atomic)."""
|
|
tables = _read_sqlite(data)
|
|
return _restore(db_path, tables, "SQLite 导入")
|
|
|
|
|
|
def import_json(db_path: str, data: dict) -> dict:
|
|
"""Import a JSON backup produced by export_database (atomic)."""
|
|
if not isinstance(data, dict) or data.get("format") != "iot-dboard-backup":
|
|
raise ValueError("不是有效的备份文件（缺少 iot-dboard-backup 标记）")
|
|
tables = data.get("tables") or {}
|
|
if not isinstance(tables, dict):
|
|
raise ValueError("备份文件结构无效：tables 缺失")
|
|
return _restore(db_path, tables, "备份导入")
|
|
|
|
|
|
def import_blob(db_path: str, data: bytes) -> dict:
|
|
"""Import a file, auto-detecting sqlite vs json backup."""
|
|
fmt = detect_format(data)
|
|
if fmt == "sqlite":
|
|
return import_sqlite(db_path, data)
|
|
if fmt == "json":
|
|
return import_json(db_path, json.loads(data.decode("utf-8")))
|
|
raise ValueError("无法识别的数据库文件格式（仅支持 SQLite 数据库或本系统导出的备份 JSON）")
|