New Database
This commit is contained in:
喆OvO
@@ -0,0 +1,263 @@
|
||||
"""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）")
|
||||
Reference in New Issue
Block a user