1193 lines
43 KiB
Python
1193 lines
43 KiB
Python
"""PostgreSQL database models, initialization, and user-scoped helpers.
|
|
|
|
Multi-user design (unchanged from the SQLite era):
|
|
- The `users` table stores every account. Root (ID 0) is the administrator and
|
|
is created automatically on first boot (fresh DB) or when migrating an
|
|
existing single-user database (all pre-existing data becomes Root's data).
|
|
- Every per-user table carries a `user_id INTEGER NOT NULL DEFAULT 0` column,
|
|
so legacy single-user rows are automatically attributed to Root on import
|
|
and any NEW column added in the future gets a DEFAULT (the migration fills
|
|
missing columns with their SCHEMA defaults).
|
|
- `auth_config` is keyed by `user_id` (legacy `id=1` row maps to user 0).
|
|
|
|
The data-access layer talks PostgreSQL through ``backend.pg_compat``, a thin
|
|
sqlite3-compatible shim (``?`` placeholders, ``INSERT OR IGNORE``,
|
|
``datetime('now')``, ``cursor.lastrowid``, ``conn.total_changes`` are all
|
|
translated transparently).
|
|
"""
|
|
import re
|
|
import sys
|
|
import time
|
|
|
|
from backend.pg_compat import connect as _pg_connect
|
|
|
|
# ─── PostgreSQL schema ───────────────────────
|
|
# All new columns must carry a DEFAULT (migration contract: missing columns are
|
|
# filled with their SCHEMA default when the DB is upgraded).
|
|
SCHEMA = """
|
|
-- Global schema version (single row).
|
|
CREATE TABLE IF NOT EXISTS meta (
|
|
id INTEGER PRIMARY KEY DEFAULT 1,
|
|
version INTEGER NOT NULL DEFAULT 1
|
|
);
|
|
|
|
-- Auth tokens (dual-token auth): access (short-lived, per-request) and
|
|
-- refresh (long-lived, stored in an httpOnly cookie). Only hashes are stored.
|
|
CREATE TABLE IF NOT EXISTS auth_tokens (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
token_hash TEXT NOT NULL,
|
|
kind TEXT NOT NULL DEFAULT 'access' CHECK(kind IN ('access','refresh')),
|
|
view TEXT NOT NULL DEFAULT 'display' CHECK(view IN ('control','display')),
|
|
temp INTEGER NOT NULL DEFAULT 0,
|
|
session_version INTEGER NOT NULL DEFAULT 0,
|
|
reauth_at INTEGER NOT NULL DEFAULT 0,
|
|
created_at TEXT DEFAULT (to_char(now(), 'YYYY-MM-DD HH24:MI:SS')),
|
|
expires_at INTEGER NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_auth_tokens_hash ON auth_tokens(token_hash);
|
|
CREATE INDEX IF NOT EXISTS idx_auth_tokens_user ON auth_tokens(user_id);
|
|
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
remark TEXT NOT NULL DEFAULT '',
|
|
role TEXT NOT NULL DEFAULT 'user' CHECK(role IN ('admin','user')),
|
|
is_temp INTEGER NOT NULL DEFAULT 0,
|
|
temp_totp_secret TEXT DEFAULT '',
|
|
status TEXT NOT NULL DEFAULT 'active' CHECK(status IN ('active','pending')),
|
|
session_version INTEGER NOT NULL DEFAULT 0,
|
|
theme TEXT NOT NULL DEFAULT 'dark' CHECK(theme IN ('dark','light','auto')),
|
|
created_at TEXT DEFAULT (to_char(now(), 'YYYY-MM-DD HH24:MI:SS'))
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS auth_config (
|
|
user_id INTEGER PRIMARY KEY DEFAULT 0,
|
|
passkey_credential_id TEXT DEFAULT '',
|
|
passkey_public_key TEXT DEFAULT '',
|
|
totp_secret TEXT DEFAULT '',
|
|
totp_enabled INTEGER DEFAULT 0,
|
|
recovery_codes TEXT DEFAULT '',
|
|
temp_password TEXT DEFAULT '',
|
|
must_configure_mfa INTEGER DEFAULT 1
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS calendar_groups (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
name TEXT NOT NULL,
|
|
color TEXT DEFAULT '#4A90D9',
|
|
enabled INTEGER DEFAULT 1,
|
|
created_at TEXT DEFAULT (to_char(now(), 'YYYY-MM-DD HH24:MI:SS'))
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS calendar_sources (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
group_id INTEGER NOT NULL,
|
|
type TEXT NOT NULL CHECK(type IN ('upload', 'subscription', 'manual')),
|
|
name TEXT NOT NULL,
|
|
url TEXT,
|
|
ics_data TEXT,
|
|
last_sync TEXT,
|
|
FOREIGN KEY (group_id) REFERENCES calendar_groups(id)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS calendar_events (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
source_id INTEGER NOT NULL,
|
|
uid TEXT,
|
|
summary TEXT,
|
|
description TEXT,
|
|
location TEXT,
|
|
dtstart TEXT,
|
|
dtend TEXT,
|
|
is_all_day INTEGER DEFAULT 0,
|
|
FOREIGN KEY (source_id) REFERENCES calendar_sources(id)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS reminders (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
calendar_group_id INTEGER NOT NULL,
|
|
minutes_before INTEGER DEFAULT 15,
|
|
enabled INTEGER DEFAULT 1,
|
|
FOREIGN KEY (calendar_group_id) REFERENCES calendar_groups(id)
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS weather_config (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
api_key TEXT DEFAULT '',
|
|
city TEXT DEFAULT 'Beijing',
|
|
lat REAL,
|
|
lon REAL,
|
|
units TEXT DEFAULT 'metric',
|
|
update_interval INTEGER DEFAULT 1800,
|
|
jwt_private_key TEXT DEFAULT '',
|
|
jwt_public_key TEXT DEFAULT '',
|
|
credential_id TEXT DEFAULT '',
|
|
project_id TEXT DEFAULT '',
|
|
qweather_host TEXT DEFAULT ''
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS weather_cache (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
data TEXT,
|
|
updated_at TEXT DEFAULT (to_char(now(), 'YYYY-MM-DD HH24:MI:SS'))
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS weather_hourly_cache (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
data TEXT,
|
|
fetched_at TEXT DEFAULT (to_char(now(), 'YYYY-MM-DD HH24:MI:SS'))
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS weather_minutely_cache (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
data TEXT,
|
|
fetched_at TEXT
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS weather_warning_cache (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
data TEXT,
|
|
fetched_at TEXT
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS clock_config (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
city TEXT DEFAULT 'Beijing',
|
|
timezone TEXT DEFAULT 'Asia/Shanghai',
|
|
is_primary INTEGER DEFAULT 1
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS alarms (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
label TEXT DEFAULT '',
|
|
time TEXT NOT NULL,
|
|
days TEXT DEFAULT '1,2,3,4,5,6,7',
|
|
enabled INTEGER DEFAULT 1,
|
|
sound TEXT DEFAULT 'default',
|
|
once INTEGER DEFAULT 0
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS countdowns (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
label TEXT DEFAULT '',
|
|
duration_seconds INTEGER DEFAULT 300,
|
|
end_time TEXT DEFAULT '',
|
|
enabled INTEGER DEFAULT 1,
|
|
sound TEXT DEFAULT 'default'
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS display_config (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
mode TEXT NOT NULL DEFAULT 'default' CHECK(mode IN ('default', 'custom')),
|
|
custom_html TEXT DEFAULT '',
|
|
custom_html_1 TEXT DEFAULT '',
|
|
custom_html_2 TEXT DEFAULT '',
|
|
custom_html_3 TEXT DEFAULT ''
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS pomodoro_state (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
user_id INTEGER NOT NULL DEFAULT 0,
|
|
running INTEGER DEFAULT 0,
|
|
phase TEXT DEFAULT 'work' CHECK(phase IN ('work', 'break')),
|
|
work_minutes INTEGER DEFAULT 25,
|
|
break_minutes INTEGER DEFAULT 5,
|
|
long_break_minutes INTEGER DEFAULT 15,
|
|
cycles_before_long_break INTEGER DEFAULT 4,
|
|
end_time TEXT,
|
|
remaining_seconds INTEGER DEFAULT 1500,
|
|
cycles INTEGER DEFAULT 0,
|
|
enabled INTEGER DEFAULT 0,
|
|
display_controls_enabled INTEGER DEFAULT 1,
|
|
cycles_date TEXT DEFAULT '',
|
|
work_label TEXT DEFAULT '工作中',
|
|
break_label TEXT DEFAULT '休息中'
|
|
);
|
|
|
|
CREATE UNIQUE INDEX IF NOT EXISTS idx_pomodoro_state_user ON pomodoro_state(user_id);
|
|
|
|
-- WebAuthn challenge store: one-time, short-TTL challenges shared across
|
|
-- Gunicorn workers via PostgreSQL (in-memory stores were per-process and
|
|
-- caused 'challenge was not expected' failures under multiple workers).
|
|
CREATE TABLE IF NOT EXISTS webauthn_challenges (
|
|
id INTEGER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
|
key TEXT NOT NULL,
|
|
challenge TEXT NOT NULL,
|
|
created_at INTEGER NOT NULL,
|
|
expires_at INTEGER NOT NULL
|
|
);
|
|
CREATE UNIQUE INDEX IF NOT EXISTS idx_webauthn_challenges_key ON webauthn_challenges(key);
|
|
"""
|
|
|
|
|
|
# Ordered list of all user-data tables (children first) for FK-safe deletion.
|
|
_DELETE_ORDER = [
|
|
"calendar_events", "calendar_sources", "reminders", "calendar_groups",
|
|
"weather_cache", "weather_hourly_cache", "weather_minutely_cache",
|
|
"weather_warning_cache", "weather_config", "clock_config",
|
|
"alarms", "countdowns", "display_config", "pomodoro_state", "auth_config",
|
|
"auth_tokens",
|
|
]
|
|
|
|
|
|
def get_db(db_path: str = "data.db"):
|
|
"""Open a PostgreSQL connection (sqlite3-compatible wrapper).
|
|
|
|
`db_path` is a PostgreSQL connection string (kept as the historical
|
|
parameter name to minimise churn across the codebase).
|
|
"""
|
|
return _pg_connect(db_path)
|
|
|
|
|
|
def _db_path(preferred: str = None) -> str:
|
|
"""Get the actual database connection string, preferring current_app config."""
|
|
if preferred:
|
|
return preferred
|
|
try:
|
|
from flask import current_app
|
|
return current_app.config.get("DB_PATH", "data.db")
|
|
except Exception:
|
|
return "data.db"
|
|
|
|
|
|
# ─── Users ───────────────────────────────────
|
|
|
|
def get_users(db_path: str = None) -> list:
|
|
"""Return all users ordered by id."""
|
|
conn = get_db(_db_path(db_path))
|
|
rows = [dict(r) for r in conn.execute("SELECT * FROM users ORDER BY id").fetchall()]
|
|
conn.close()
|
|
return rows
|
|
|
|
|
|
def get_active_user_ids(db_path: str = None) -> list:
|
|
"""Return ids of all active (non-pending) users. Used by the scheduler."""
|
|
conn = get_db(_db_path(db_path))
|
|
rows = conn.execute("SELECT id FROM users WHERE status='active' ORDER BY id").fetchall()
|
|
conn.close()
|
|
return [r["id"] for r in rows]
|
|
|
|
|
|
def get_user(user_id: int, db_path: str = None) -> dict | None:
|
|
conn = get_db(_db_path(db_path))
|
|
row = conn.execute("SELECT * FROM users WHERE id=?", (user_id,)).fetchone()
|
|
conn.close()
|
|
return dict(row) if row else None
|
|
|
|
|
|
def user_exists(user_id: int, db_path: str = None) -> bool:
|
|
conn = get_db(_db_path(db_path))
|
|
row = conn.execute("SELECT 1 FROM users WHERE id=?", (user_id,)).fetchone()
|
|
conn.close()
|
|
return row is not None
|
|
|
|
|
|
def is_admin_user(user_id: int, db_path: str = None) -> bool:
|
|
"""Root is always admin; otherwise check the role column."""
|
|
if user_id == 0:
|
|
return True
|
|
conn = get_db(_db_path(db_path))
|
|
row = conn.execute("SELECT role FROM users WHERE id=?", (user_id,)).fetchone()
|
|
conn.close()
|
|
return bool(row and row["role"] == "admin")
|
|
|
|
|
|
def create_user(remark: str = "", role: str = "user", is_temp: int = 0,
|
|
temp_totp_secret: str = "", status: str = "active",
|
|
db_path: str = None) -> int:
|
|
"""Create a user. Returns the new user id.
|
|
The remark defaults to the user's own id (the only place an id is shown).
|
|
"""
|
|
conn = get_db(_db_path(db_path))
|
|
cur = conn.execute(
|
|
"INSERT INTO users (remark, role, is_temp, temp_totp_secret, status) VALUES (?, ?, ?, ?, ?)",
|
|
(remark, role, 1 if is_temp else 0, temp_totp_secret, status)
|
|
)
|
|
uid = cur.lastrowid
|
|
if not remark:
|
|
conn.execute("UPDATE users SET remark=? WHERE id=?", (str(uid), uid))
|
|
conn.commit()
|
|
conn.close()
|
|
return uid
|
|
|
|
|
|
def update_user(user_id: int, remark: str = None, role: str = None,
|
|
is_temp: int = None, temp_totp_secret: str = None,
|
|
status: str = None, theme: str = None, db_path: str = None):
|
|
conn = get_db(_db_path(db_path))
|
|
fields = []
|
|
vals = []
|
|
for key, val in [("remark", remark), ("role", role), ("is_temp", is_temp),
|
|
("temp_totp_secret", temp_totp_secret), ("status", status),
|
|
("theme", theme)]:
|
|
if val is not None:
|
|
fields.append(f"{key}=?")
|
|
vals.append(val)
|
|
if fields:
|
|
vals.append(user_id)
|
|
conn.execute(f"UPDATE users SET {', '.join(fields)} WHERE id=?", vals)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
|
|
def delete_user(user_id: int, db_path: str = None):
|
|
"""Delete a user and ALL of their data. Root (id 0) cannot be deleted.
|
|
|
|
Deletion order respects foreign keys (children before parents).
|
|
"""
|
|
if user_id == 0:
|
|
return
|
|
conn = get_db(_db_path(db_path))
|
|
try:
|
|
for table in _DELETE_ORDER:
|
|
conn.execute(f"DELETE FROM {table} WHERE user_id=?", (user_id,))
|
|
conn.execute("DELETE FROM users WHERE id=?", (user_id,))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def get_user_session_version(user_id: int, db_path: str = None) -> int:
|
|
conn = get_db(_db_path(db_path))
|
|
row = conn.execute("SELECT session_version FROM users WHERE id=?", (user_id,)).fetchone()
|
|
conn.close()
|
|
return row["session_version"] if row else -1
|
|
|
|
|
|
def bump_user_session_version(user_id: int, db_path: str = None) -> int:
|
|
"""Invalidate all sessions of one user by incrementing their version."""
|
|
conn = get_db(_db_path(db_path))
|
|
conn.execute("UPDATE users SET session_version = session_version + 1 WHERE id=?", (user_id,))
|
|
conn.commit()
|
|
row = conn.execute("SELECT session_version FROM users WHERE id=?", (user_id,)).fetchone()
|
|
conn.close()
|
|
return row["session_version"] if row else -1
|
|
|
|
|
|
def bump_all_session_versions(db_path: str = None) -> int:
|
|
"""Invalidate every session (admin 'kick all' action)."""
|
|
conn = get_db(_db_path(db_path))
|
|
conn.execute("UPDATE users SET session_version = session_version + 1")
|
|
conn.commit()
|
|
conn.close()
|
|
return 0
|
|
|
|
|
|
def is_temp_account(user_id: int, db_path: str = None) -> bool:
|
|
"""True if the account is a temp account (awaiting first login-method setup)."""
|
|
conn = get_db(_db_path(db_path))
|
|
row = conn.execute("SELECT is_temp FROM users WHERE id=?", (user_id,)).fetchone()
|
|
conn.close()
|
|
return bool(row and row["is_temp"])
|
|
|
|
|
|
def get_temp_totp_secret(user_id: int, db_path: str = None) -> str:
|
|
conn = get_db(_db_path(db_path))
|
|
row = conn.execute("SELECT temp_totp_secret FROM users WHERE id=?", (user_id,)).fetchone()
|
|
conn.close()
|
|
return (row["temp_totp_secret"] or "") if row else ""
|
|
|
|
|
|
def finalize_temp_account(user_id: int, db_path: str = None):
|
|
"""Convert a temp account into a formal account:
|
|
invalidates the temp TOTP, clears the temp flag, and revokes every
|
|
temp=1 token (a temp refresh must never mint formal tokens afterwards).
|
|
"""
|
|
conn = get_db(_db_path(db_path))
|
|
conn.execute(
|
|
"UPDATE users SET is_temp=0, temp_totp_secret='' WHERE id=?",
|
|
(user_id,)
|
|
)
|
|
conn.execute(
|
|
"UPDATE auth_config SET must_configure_mfa=0 WHERE user_id=?",
|
|
(user_id,)
|
|
)
|
|
conn.execute(
|
|
"DELETE FROM auth_tokens WHERE user_id=? AND temp=1",
|
|
(user_id,)
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
|
|
# ─── Display config (per user) ──────────────
|
|
|
|
def _ensure_display_row(conn, user_id: int):
|
|
row = conn.execute("SELECT id FROM display_config WHERE user_id=?", (user_id,)).fetchone()
|
|
if row is None:
|
|
conn.execute(
|
|
"INSERT INTO display_config (user_id, mode) VALUES (?, 'default')",
|
|
(user_id,)
|
|
)
|
|
|
|
|
|
def get_display_config(db_path: str = None, user_id: int = 0) -> dict:
|
|
db = get_db(_db_path(db_path))
|
|
row = db.execute(
|
|
"SELECT mode, custom_html, custom_html_1, custom_html_2, custom_html_3 "
|
|
"FROM display_config WHERE user_id=?",
|
|
(user_id,)
|
|
).fetchone()
|
|
if row is None:
|
|
db.execute(
|
|
"INSERT INTO display_config (user_id, mode, custom_html, custom_html_1, custom_html_2, custom_html_3) "
|
|
"VALUES (?, 'default', '', '', '', '')",
|
|
(user_id,)
|
|
)
|
|
db.commit()
|
|
cfg = {"mode": "default", "custom_html": "", "custom_html_1": "", "custom_html_2": "", "custom_html_3": ""}
|
|
else:
|
|
cfg = {
|
|
"mode": row["mode"],
|
|
"custom_html": row["custom_html"] or "",
|
|
"custom_html_1": row["custom_html_1"] or "",
|
|
"custom_html_2": row["custom_html_2"] or "",
|
|
"custom_html_3": row["custom_html_3"] or "",
|
|
}
|
|
db.close()
|
|
return cfg
|
|
|
|
|
|
def set_display_mode(mode: str, db_path: str = None, user_id: int = 0) -> bool:
|
|
if mode not in ("default", "custom"):
|
|
return False
|
|
db = get_db(_db_path(db_path))
|
|
_ensure_display_row(db, user_id)
|
|
db.execute("UPDATE display_config SET mode=? WHERE user_id=?", (mode, user_id))
|
|
db.commit()
|
|
db.close()
|
|
return True
|
|
|
|
|
|
def set_custom_html(html: str, db_path: str = None, user_id: int = 0) -> bool:
|
|
db = get_db(_db_path(db_path))
|
|
_ensure_display_row(db, user_id)
|
|
db.execute("UPDATE display_config SET custom_html=? WHERE user_id=?", (html, user_id))
|
|
db.commit()
|
|
db.close()
|
|
return True
|
|
|
|
|
|
def set_custom_html_template(index: int, html: str, db_path: str = None, user_id: int = 0) -> bool:
|
|
if index not in (1, 2, 3):
|
|
return False
|
|
db = get_db(_db_path(db_path))
|
|
_ensure_display_row(db, user_id)
|
|
db.execute(
|
|
f"UPDATE display_config SET custom_html_{index}=? WHERE user_id=?",
|
|
(html, user_id)
|
|
)
|
|
db.commit()
|
|
db.close()
|
|
return True
|
|
|
|
|
|
def apply_custom_html(index: int, db_path: str = None, user_id: int = 0) -> bool:
|
|
"""Copy template N (0=main custom_html, 1-3=custom_html_1-3) into custom_html and set mode to custom."""
|
|
db = get_db(_db_path(db_path))
|
|
row = db.execute(
|
|
"SELECT custom_html, custom_html_1, custom_html_2, custom_html_3 "
|
|
"FROM display_config WHERE user_id=?",
|
|
(user_id,)
|
|
).fetchone()
|
|
if row is None:
|
|
db.close()
|
|
return False
|
|
if index == 0:
|
|
html = row["custom_html"] or ""
|
|
elif index == 1:
|
|
html = row["custom_html_1"] or ""
|
|
elif index == 2:
|
|
html = row["custom_html_2"] or ""
|
|
elif index == 3:
|
|
html = row["custom_html_3"] or ""
|
|
else:
|
|
db.close()
|
|
return False
|
|
db.execute("UPDATE display_config SET mode='custom', custom_html=? WHERE user_id=?", (html, user_id))
|
|
db.commit()
|
|
db.close()
|
|
return True
|
|
|
|
|
|
# ─── Auth tokens (dual-token auth) ───────────
|
|
# Access tokens are short-lived and validated on every request; refresh tokens
|
|
# are long-lived and stored in an httpOnly cookie. Only SHA-256 hashes are
|
|
# stored in the DB. Kicking a user/device = delete their tokens + bump the
|
|
# user's session_version (which invalidates every outstanding access token).
|
|
|
|
ACCESS_TOKEN_TTL = 15 * 60 # 15 minutes
|
|
REFRESH_TOKEN_TTL = 30 * 24 * 3600 # 30 days (sliding, renewed on each refresh)
|
|
|
|
|
|
def _hash_token(raw: str) -> str:
|
|
import hashlib
|
|
return hashlib.sha256(raw.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def _prune_expired_tokens(conn):
|
|
conn.execute("DELETE FROM auth_tokens WHERE expires_at < ?", (int(time.time()),))
|
|
|
|
|
|
def create_auth_token(user_id: int, kind: str, view: str = "display",
|
|
temp: bool = False, ttl: int = None,
|
|
db_path: str = None) -> str:
|
|
"""Create an auth token row and return the raw (unhashed) token.
|
|
Only the SHA-256 hash is persisted."""
|
|
import secrets as _secrets
|
|
if ttl is None:
|
|
ttl = ACCESS_TOKEN_TTL if kind == "access" else REFRESH_TOKEN_TTL
|
|
raw = _secrets.token_urlsafe(48 if kind == "refresh" else 32)
|
|
conn = get_db(_db_path(db_path))
|
|
try:
|
|
_prune_expired_tokens(conn)
|
|
ver_row = conn.execute(
|
|
"SELECT session_version FROM users WHERE id=?", (user_id,)
|
|
).fetchone()
|
|
ver = ver_row["session_version"] if ver_row else 0
|
|
conn.execute(
|
|
"INSERT INTO auth_tokens (user_id, token_hash, kind, view, temp, "
|
|
"session_version, expires_at) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
|
(user_id, _hash_token(raw), kind, view, 1 if temp else 0, ver,
|
|
int(time.time()) + ttl)
|
|
)
|
|
conn.commit()
|
|
return raw
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def _token_lookup(raw: str, kind: str, db_path: str = None) -> dict | None:
|
|
"""Look up a token by its raw value. For access tokens the row is only
|
|
valid while it matches the user's current session_version (kick support).
|
|
Returns None when invalid/expired/user inactive/version bumped."""
|
|
if not raw:
|
|
return None
|
|
conn = get_db(_db_path(db_path))
|
|
try:
|
|
h = _hash_token(raw)
|
|
row = conn.execute(
|
|
"SELECT t.*, u.status AS user_status FROM auth_tokens t "
|
|
"JOIN users u ON u.id = t.user_id WHERE t.token_hash=? AND t.kind=?",
|
|
(h, kind)
|
|
).fetchone()
|
|
if not row:
|
|
return None
|
|
if row["expires_at"] <= int(time.time()):
|
|
return None
|
|
if row["user_status"] != "active":
|
|
return None
|
|
if kind == "access" and row["session_version"] != _cur_session_version(row["user_id"], conn):
|
|
return None
|
|
return dict(row)
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def _cur_session_version(user_id: int, conn) -> int:
|
|
r = conn.execute("SELECT session_version FROM users WHERE id=?", (user_id,)).fetchone()
|
|
return r["session_version"] if r else -1
|
|
|
|
|
|
def find_access_token(raw: str, db_path: str = None) -> dict | None:
|
|
"""Validate an access token. Returns the row dict or None."""
|
|
return _token_lookup(raw, "access", db_path)
|
|
|
|
|
|
def rotate_refresh_token(raw: str, db_path: str = None) -> dict | None:
|
|
"""Atomically validate and rotate a refresh token (concurrency-safe).
|
|
Returns a fresh {access_token, refresh_token} pair or None. The old refresh
|
|
token is consumed inside a single transaction, so a replayed/concurrent
|
|
duplicate can never mint two valid pairs."""
|
|
if not raw:
|
|
return None
|
|
import secrets as _secrets
|
|
conn = get_db(_db_path(db_path))
|
|
try:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
row = conn.execute(
|
|
"SELECT t.*, u.status AS user_status FROM auth_tokens t "
|
|
"JOIN users u ON u.id = t.user_id WHERE t.token_hash=? AND t.kind='refresh'",
|
|
(_hash_token(raw),)
|
|
).fetchone()
|
|
if (not row or row["expires_at"] <= int(time.time())
|
|
or row["user_status"] != "active"):
|
|
conn.rollback()
|
|
return None
|
|
cur = conn.execute("DELETE FROM auth_tokens WHERE id=?", (row["id"],))
|
|
if cur.rowcount == 0:
|
|
conn.rollback()
|
|
return None
|
|
now = int(time.time())
|
|
acc_raw = _secrets.token_urlsafe(32)
|
|
ref_raw = _secrets.token_urlsafe(48)
|
|
ver = row["session_version"]
|
|
conn.execute(
|
|
"INSERT INTO auth_tokens (user_id, token_hash, kind, view, temp, "
|
|
"session_version, expires_at) VALUES (?, ?, 'access', ?, ?, ?, ?)",
|
|
(row["user_id"], _hash_token(acc_raw), row["view"], row["temp"],
|
|
ver, now + ACCESS_TOKEN_TTL)
|
|
)
|
|
conn.execute(
|
|
"INSERT INTO auth_tokens (user_id, token_hash, kind, view, temp, "
|
|
"session_version, expires_at) VALUES (?, ?, 'refresh', ?, ?, ?, ?)",
|
|
(row["user_id"], _hash_token(ref_raw), row["view"], row["temp"],
|
|
ver, now + REFRESH_TOKEN_TTL)
|
|
)
|
|
conn.commit()
|
|
return {"access_token": acc_raw, "refresh_token": ref_raw}
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def find_refresh_token(raw: str, db_path: str = None) -> dict | None:
|
|
"""Validate a refresh token (status check only; version is not compared
|
|
here because a refresh token belongs to one device and is deleted on
|
|
kick, not merely invalidated)."""
|
|
if not raw:
|
|
return None
|
|
conn = get_db(_db_path(db_path))
|
|
try:
|
|
h = _hash_token(raw)
|
|
row = conn.execute(
|
|
"SELECT t.*, u.status AS user_status FROM auth_tokens t "
|
|
"JOIN users u ON u.id = t.user_id WHERE t.token_hash=? AND t.kind='refresh'",
|
|
(h,)
|
|
).fetchone()
|
|
if not row:
|
|
return None
|
|
if row["expires_at"] <= int(time.time()):
|
|
return None
|
|
if row["user_status"] != "active":
|
|
return None
|
|
return dict(row)
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def mark_token_reauth(token_row_id: int, db_path: str = None):
|
|
"""Stamp the current access token row as re-authenticated (per-device)."""
|
|
if not token_row_id:
|
|
return
|
|
conn = get_db(_db_path(db_path))
|
|
try:
|
|
conn.execute("UPDATE auth_tokens SET reauth_at=? WHERE id=?",
|
|
(int(time.time()), token_row_id))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def delete_auth_token(raw: str, kind: str = None, db_path: str = None):
|
|
"""Delete a single token row by raw value (logout of one device)."""
|
|
if not raw:
|
|
return
|
|
conn = get_db(_db_path(db_path))
|
|
try:
|
|
h = _hash_token(raw)
|
|
if kind:
|
|
conn.execute("DELETE FROM auth_tokens WHERE token_hash=? AND kind=?", (h, kind))
|
|
else:
|
|
conn.execute("DELETE FROM auth_tokens WHERE token_hash=?", (h,))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def revoke_user_tokens(user_id: int, db_path: str = None):
|
|
"""Delete every token of a user (kick all devices of one account).
|
|
Callers should also bump_user_session_version() to invalidate access tokens."""
|
|
conn = get_db(_db_path(db_path))
|
|
try:
|
|
conn.execute("DELETE FROM auth_tokens WHERE user_id=?", (user_id,))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def revoke_all_tokens(db_path: str = None):
|
|
"""Delete every token of every user (Root: kick all users)."""
|
|
conn = get_db(_db_path(db_path))
|
|
try:
|
|
conn.execute("DELETE FROM auth_tokens")
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
# ─── Auth / MFA helpers (per user) ──────────
|
|
|
|
def get_auth_config(db_path: str = None, user_id: int = 0) -> dict:
|
|
"""Return the full auth_config row for a user, with defaults for missing fields."""
|
|
conn = get_db(_db_path(db_path))
|
|
row = conn.execute("SELECT * FROM auth_config WHERE user_id = ?", (user_id,)).fetchone()
|
|
if row is None:
|
|
conn.execute(
|
|
"""INSERT INTO auth_config (user_id, passkey_credential_id, passkey_public_key,
|
|
totp_secret, totp_enabled, recovery_codes, temp_password, must_configure_mfa)
|
|
VALUES (?, '', '', '', 0, '', '', 1)""",
|
|
(user_id,)
|
|
)
|
|
conn.commit()
|
|
cfg = {
|
|
"user_id": user_id,
|
|
"passkey_credential_id": "",
|
|
"passkey_public_key": "",
|
|
"totp_secret": "",
|
|
"totp_enabled": 0,
|
|
"recovery_codes": "",
|
|
"temp_password": "",
|
|
"must_configure_mfa": 1,
|
|
}
|
|
else:
|
|
cfg = {
|
|
"user_id": row["user_id"],
|
|
"passkey_credential_id": row["passkey_credential_id"] or "",
|
|
"passkey_public_key": row["passkey_public_key"] or "",
|
|
"totp_secret": row["totp_secret"] or "",
|
|
"totp_enabled": bool(row["totp_enabled"]),
|
|
"recovery_codes": row["recovery_codes"] or "",
|
|
"temp_password": row["temp_password"] or "",
|
|
"must_configure_mfa": bool(row["must_configure_mfa"]),
|
|
}
|
|
conn.close()
|
|
return cfg
|
|
|
|
|
|
def update_auth_config(user_id: int = 0, **kwargs):
|
|
"""Update auth_config fields. Keys match DB column names."""
|
|
db_path = _db_path(kwargs.pop("_db_path", None))
|
|
conn = get_db(db_path)
|
|
row = conn.execute("SELECT user_id FROM auth_config WHERE user_id = ?", (user_id,)).fetchone()
|
|
fields = []
|
|
vals = []
|
|
for key, val in kwargs.items():
|
|
if key in ("passkey_credential_id", "passkey_public_key", "totp_secret",
|
|
"recovery_codes", "temp_password", "must_configure_mfa", "totp_enabled"):
|
|
fields.append(f"{key}=?")
|
|
vals.append(val)
|
|
if not fields:
|
|
conn.close()
|
|
return
|
|
if row is None:
|
|
cols = ["user_id"] + [f.split("=")[0] for f in fields]
|
|
col_str = ", ".join(cols)
|
|
ph = ", ".join("?" * len(cols))
|
|
conn.execute(f"INSERT INTO auth_config ({col_str}) VALUES ({ph})", (user_id, *vals))
|
|
else:
|
|
vals.append(user_id)
|
|
conn.execute(f"UPDATE auth_config SET {', '.join(fields)} WHERE user_id=?", vals)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
|
|
def get_temp_password(db_path: str = None, user_id: int = 0) -> str:
|
|
cfg = get_auth_config(db_path, user_id)
|
|
return cfg["temp_password"]
|
|
|
|
|
|
def set_temp_password(password: str, db_path: str = None, user_id: int = 0):
|
|
update_auth_config(user_id, temp_password=password, _db_path=db_path)
|
|
|
|
|
|
def clear_temp_password(db_path: str = None, user_id: int = 0):
|
|
update_auth_config(user_id, temp_password="", must_configure_mfa=0, _db_path=db_path)
|
|
|
|
|
|
def get_recovery_codes(db_path: str = None, user_id: int = 0) -> list:
|
|
cfg = get_auth_config(db_path, user_id)
|
|
raw = cfg["recovery_codes"]
|
|
if not raw:
|
|
return []
|
|
return raw.split(",")
|
|
|
|
|
|
def set_recovery_codes(codes: list, db_path: str = None, user_id: int = 0):
|
|
update_auth_config(user_id, recovery_codes=",".join(codes), _db_path=db_path)
|
|
|
|
|
|
def has_mfa_configured(db_path: str = None, user_id: int = 0) -> bool:
|
|
cfg = get_auth_config(db_path, user_id)
|
|
return bool(cfg["passkey_credential_id"]) or bool(cfg["totp_enabled"])
|
|
|
|
|
|
def find_user_by_credential_id(credential_id: str, db_path: str = None) -> int | None:
|
|
"""Find the user who owns a given WebAuthn credential id (for login)."""
|
|
conn = get_db(_db_path(db_path))
|
|
row = conn.execute(
|
|
"SELECT a.user_id FROM auth_config a "
|
|
"JOIN users u ON u.id = a.user_id "
|
|
"WHERE a.passkey_credential_id = ? AND u.status = 'active'",
|
|
(credential_id,)
|
|
).fetchone()
|
|
conn.close()
|
|
return row["user_id"] if row else None
|
|
|
|
|
|
# ─── WebAuthn challenge store ────────────────
|
|
# Challenges are persisted in PostgreSQL so registration/login ceremonies work
|
|
# across Gunicorn workers (in-memory stores were per-process). Consumed
|
|
# atomically (DELETE ... RETURNING) and pruned by TTL.
|
|
|
|
WEBAUTHN_CHALLENGE_TTL = 300 # 5 minutes
|
|
|
|
|
|
def save_webauthn_challenge(key: str, challenge: str,
|
|
ttl: int = WEBAUTHN_CHALLENGE_TTL,
|
|
db_path: str = None):
|
|
"""Store a challenge (base64url string) under `key`, replacing any previous."""
|
|
conn = get_db(_db_path(db_path))
|
|
try:
|
|
conn.execute("DELETE FROM webauthn_challenges WHERE expires_at < ?", (int(time.time()),))
|
|
now = int(time.time())
|
|
conn.execute(
|
|
"INSERT INTO webauthn_challenges (key, challenge, created_at, expires_at) "
|
|
"VALUES (?, ?, ?, ?) "
|
|
"ON CONFLICT (key) DO UPDATE SET challenge=excluded.challenge, "
|
|
"created_at=excluded.created_at, expires_at=excluded.expires_at",
|
|
(key, challenge, now, now + ttl)
|
|
)
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def pop_webauthn_challenge(key: str, db_path: str = None) -> str | None:
|
|
"""Atomically consume and return the challenge for `key` (None when absent/expired)."""
|
|
conn = get_db(_db_path(db_path))
|
|
try:
|
|
conn.execute("DELETE FROM webauthn_challenges WHERE expires_at < ?", (int(time.time()),))
|
|
row = conn.execute(
|
|
"DELETE FROM webauthn_challenges WHERE key=? AND expires_at >= ? RETURNING challenge",
|
|
(key, int(time.time()))
|
|
).fetchone()
|
|
conn.commit()
|
|
return row["challenge"] if row else None
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
# ─── DB init / migration ─────────────────────
|
|
|
|
# Stores passwords that were just generated, for the startup script to print.
|
|
_EMITTED_PASSWORDS: list = []
|
|
|
|
# Current schema version supported by this program. Bump it whenever the
|
|
# SCHEMA changes (new table, column, index, constraint, ...) so existing
|
|
# databases get upgraded on next start. Rule: ANY database change ⇒ bump.
|
|
DB_VERSION = 3
|
|
|
|
|
|
def _ensure_defaults(conn):
|
|
"""Insert default rows for config tables and the Root user when missing.
|
|
Generates a random 6-digit temp password when Root's auth_config is created
|
|
(or when a migrated DB has no MFA configured yet).
|
|
"""
|
|
import secrets
|
|
import string
|
|
|
|
# Schema-version row (single row, id=1).
|
|
conn.execute("INSERT OR IGNORE INTO meta (id, version) VALUES (1, ?)", (DB_VERSION,))
|
|
|
|
# Root user (id 0)
|
|
root = conn.execute("SELECT id FROM users WHERE id=0").fetchone()
|
|
if root is None:
|
|
conn.execute(
|
|
"INSERT OR IGNORE INTO users (id, remark, role, status) VALUES (0, 'Root', 'admin', 'active')"
|
|
)
|
|
# Keep the users identity sequence ahead of explicit id=0 insertion.
|
|
conn.execute(
|
|
"SELECT setval(pg_get_serial_sequence('users', 'id'), "
|
|
"GREATEST((SELECT COALESCE(MAX(id), 1) FROM users), 1))"
|
|
)
|
|
|
|
# weather_config for Root
|
|
count = conn.execute("SELECT COUNT(*) FROM weather_config WHERE user_id=0").fetchone()[0]
|
|
if count == 0:
|
|
conn.execute("INSERT INTO weather_config (user_id, city, lat, lon) "
|
|
"VALUES (0, 'Beijing', 39.9042, 116.4074)")
|
|
|
|
# clock_config for Root
|
|
count = conn.execute("SELECT COUNT(*) FROM clock_config WHERE user_id=0").fetchone()[0]
|
|
if count == 0:
|
|
conn.execute("INSERT INTO clock_config (user_id, city, timezone, is_primary) "
|
|
"VALUES (0, 'Beijing', 'Asia/Shanghai', 1)")
|
|
|
|
# auth_config for Root
|
|
row = conn.execute("SELECT * FROM auth_config WHERE user_id=0").fetchone()
|
|
mfa_configured = bool(row and (row["passkey_credential_id"] or row["totp_enabled"]))
|
|
if row is None:
|
|
temp_pw = "".join(secrets.choice(string.digits) for _ in range(6))
|
|
conn.execute(
|
|
"INSERT INTO auth_config (user_id, passkey_credential_id, passkey_public_key, "
|
|
"totp_secret, totp_enabled, recovery_codes, temp_password, must_configure_mfa) "
|
|
"VALUES (0, '', '', '', 0, '', ?, 1)",
|
|
(temp_pw,)
|
|
)
|
|
_EMITTED_PASSWORDS.append(temp_pw)
|
|
print(f"\n{'='*50}", file=sys.stderr)
|
|
print(f" 🔐 Root 初始临时密码: {temp_pw}", file=sys.stderr)
|
|
print(f" ⚠️ 请使用此密码登录 Root 并配置双因素认证", file=sys.stderr)
|
|
print(f"{'='*50}\n", file=sys.stderr)
|
|
else:
|
|
# Migrated DB. If Root already has MFA, no temp password is needed.
|
|
if mfa_configured:
|
|
if row["temp_password"] or row["must_configure_mfa"]:
|
|
conn.execute(
|
|
"UPDATE auth_config SET temp_password='', must_configure_mfa=0 WHERE user_id=0"
|
|
)
|
|
elif (not row["temp_password"]) and row["must_configure_mfa"]:
|
|
# Migrated DB without MFA — generate a fresh temp password.
|
|
temp_pw = "".join(secrets.choice(string.digits) for _ in range(6))
|
|
conn.execute(
|
|
"UPDATE auth_config SET temp_password=? WHERE user_id=0",
|
|
(temp_pw,)
|
|
)
|
|
_EMITTED_PASSWORDS.append(temp_pw)
|
|
print(f"\n{'='*50}", file=sys.stderr)
|
|
print(f" 🔐 迁移已存在数据库 — Root 临时密码: {temp_pw}", file=sys.stderr)
|
|
print(f" ⚠️ 请使用此密码登录 Root 并配置双因素认证", file=sys.stderr)
|
|
print(f"{'='*50}\n", file=sys.stderr)
|
|
|
|
conn.commit()
|
|
|
|
|
|
def _table_columns(conn, table_name: str) -> list:
|
|
"""Return the column names of a table (information_schema)."""
|
|
rows = conn.execute(
|
|
"SELECT column_name FROM information_schema.columns "
|
|
"WHERE table_schema = current_schema() AND table_name = ? "
|
|
"ORDER BY ordinal_position",
|
|
(table_name,)
|
|
).fetchall()
|
|
return [r["column_name"] for r in rows]
|
|
|
|
|
|
def _read_db_version(conn) -> int | None:
|
|
"""Read the schema version from the `meta` table.
|
|
Returns None when the DB is legacy (no meta table / no version row).
|
|
"""
|
|
try:
|
|
row = conn.execute("SELECT version FROM meta WHERE id=1").fetchone()
|
|
return int(row["version"]) if row else None
|
|
except Exception:
|
|
conn._safe_rollback()
|
|
return None
|
|
|
|
|
|
# ─── SCHEMA parsing for column-level migration ──────
|
|
_CREATE_TABLE_RE = re.compile(
|
|
r"CREATE TABLE IF NOT EXISTS\s+(\w+)\s*\((.*?)\)\s*;",
|
|
re.IGNORECASE | re.DOTALL,
|
|
)
|
|
|
|
|
|
def _split_top_level(text: str) -> list:
|
|
"""Split `text` on top-level commas (ignoring commas inside parentheses
|
|
AND inside single-quoted strings, e.g. DEFAULT '1,2,3,4,5,6,7')."""
|
|
parts = []
|
|
depth = 0
|
|
cur = []
|
|
in_str = False
|
|
i = 0
|
|
while i < len(text):
|
|
ch = text[i]
|
|
if ch == "'":
|
|
# SQL 单引号转义用两个单引号表示
|
|
if in_str and i + 1 < len(text) and text[i + 1] == "'":
|
|
cur.append("''")
|
|
i += 2
|
|
continue
|
|
in_str = not in_str
|
|
cur.append(ch)
|
|
elif not in_str:
|
|
if ch == "(":
|
|
depth += 1
|
|
elif ch == ")":
|
|
depth -= 1
|
|
if ch == "," and depth == 0:
|
|
parts.append("".join(cur))
|
|
cur = []
|
|
i += 1
|
|
continue
|
|
cur.append(ch)
|
|
else:
|
|
cur.append(ch)
|
|
i += 1
|
|
if cur:
|
|
parts.append("".join(cur))
|
|
return parts
|
|
|
|
|
|
def _strip_check(s: str) -> str:
|
|
"""Remove every CHECK(...) clause (parenthesis-balanced)."""
|
|
out = s
|
|
while True:
|
|
idx = out.upper().find("CHECK(")
|
|
if idx == -1:
|
|
break
|
|
depth = 0
|
|
end = -1
|
|
for j in range(idx, len(out)):
|
|
if out[j] == "(":
|
|
depth += 1
|
|
elif out[j] == ")":
|
|
depth -= 1
|
|
if depth == 0:
|
|
end = j
|
|
break
|
|
if end == -1:
|
|
break
|
|
out = out[:idx] + out[end + 1:]
|
|
return out
|
|
|
|
|
|
def _migrate_missing_columns(conn):
|
|
"""Add columns that exist in SCHEMA but are missing from existing tables.
|
|
New columns are created with their SCHEMA DEFAULT (migration contract).
|
|
Best-effort: a failing column does not block startup.
|
|
"""
|
|
for m in _CREATE_TABLE_RE.finditer(SCHEMA):
|
|
table = m.group(1)
|
|
try:
|
|
existing = set(_table_columns(conn, table))
|
|
except Exception:
|
|
conn._safe_rollback()
|
|
continue
|
|
for col_def in _split_top_level(m.group(2)):
|
|
col_def = col_def.strip()
|
|
if not col_def:
|
|
continue
|
|
upper = col_def.upper()
|
|
if upper.startswith(("PRIMARY KEY", "UNIQUE", "CHECK", "FOREIGN KEY", "CONSTRAINT")):
|
|
continue
|
|
name = col_def.split()[0].strip('"')
|
|
if not name or name in existing:
|
|
continue
|
|
# Make the ADD COLUMN safe for rows that already exist:
|
|
# drop NOT NULL / PK / UNIQUE / IDENTITY; keep TYPE + DEFAULT.
|
|
safe = re.sub(r"\s+NOT\s+NULL\b", "", col_def, flags=re.IGNORECASE)
|
|
safe = re.sub(r"\s+PRIMARY\s+KEY\b", "", safe, flags=re.IGNORECASE)
|
|
safe = re.sub(r"\s+UNIQUE\b", "", safe, flags=re.IGNORECASE)
|
|
safe = re.sub(r"\s+GENERATED\s+BY\s+DEFAULT\s+AS\s+IDENTITY\b", "", safe, flags=re.IGNORECASE)
|
|
safe = _strip_check(safe).strip().rstrip(",")
|
|
try:
|
|
conn.execute(f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS {safe}")
|
|
except Exception:
|
|
# rollback done by the shim; the next column starts a fresh tx
|
|
pass
|
|
|
|
|
|
def _fix_identity_sequences(conn):
|
|
"""Sync every IDENTITY sequence to max(id) (safe after explicit id inserts)."""
|
|
try:
|
|
rows = conn.execute(
|
|
"SELECT table_name, column_name FROM information_schema.columns "
|
|
"WHERE table_schema = current_schema() AND is_identity = 'YES'"
|
|
).fetchall()
|
|
except Exception:
|
|
conn._safe_rollback()
|
|
return
|
|
for r in rows:
|
|
t, c = r["table_name"], r["column_name"]
|
|
try:
|
|
conn.execute(
|
|
f"SELECT setval(pg_get_serial_sequence('{t}', '{c}'), "
|
|
f"GREATEST((SELECT COALESCE(MAX({c}), 1) FROM {t}), 1))"
|
|
)
|
|
except Exception:
|
|
conn._safe_rollback()
|
|
|
|
|
|
# Timestamp columns whose DEFAULT was generated by early builds with the wrong
|
|
# to_char pattern ('YYYY-MM-DD HH:MM:SS' — HH is 12-hour, MM is the month).
|
|
_TIMESTAMP_DEFAULT_FIX = [
|
|
("auth_tokens", "created_at"),
|
|
("users", "created_at"),
|
|
("calendar_groups", "created_at"),
|
|
("weather_cache", "updated_at"),
|
|
("weather_hourly_cache", "fetched_at"),
|
|
]
|
|
|
|
|
|
def _fix_timestamp_defaults(conn):
|
|
"""Repair legacy DEFAULT expressions (idempotent, best-effort)."""
|
|
for t, c in _TIMESTAMP_DEFAULT_FIX:
|
|
try:
|
|
conn.execute(
|
|
f"ALTER TABLE {t} ALTER COLUMN {c} "
|
|
"SET DEFAULT to_char(now(), 'YYYY-MM-DD HH24:MI:SS')"
|
|
)
|
|
except Exception:
|
|
conn._safe_rollback()
|
|
|
|
|
|
def init_db(db_path: str = "data.db"):
|
|
"""Initialize the database with schema-version aware migration (PostgreSQL).
|
|
|
|
- Fresh database (no `meta` row) -> create the latest SCHEMA + defaults.
|
|
- Existing database:
|
|
* DB version == DB_VERSION -> use as-is (fast path, defaults ensured)
|
|
* no version / DB version < DB_VERSION
|
|
-> create missing tables/columns; missing
|
|
columns are filled with their SCHEMA
|
|
DEFAULT; stamp the new version
|
|
* DB version > DB_VERSION -> abort with an error (DB is newer than
|
|
this program)
|
|
"""
|
|
conn = get_db(db_path)
|
|
try:
|
|
old_version = _read_db_version(conn)
|
|
|
|
if old_version is not None and old_version > DB_VERSION:
|
|
raise SystemExit(
|
|
f"[FATAL] 数据库 schema 版本 (v{old_version}) 高于当前程序支持的版本 "
|
|
f"(v{DB_VERSION})。请先升级程序，或删除/降级数据库后再启动。"
|
|
)
|
|
|
|
if old_version is not None and old_version == DB_VERSION:
|
|
# Same version — no upgrade needed. Run the idempotent DDL anyway so
|
|
# any newly-added tables (e.g. webauthn_challenges) exist on
|
|
# existing databases without bumping the schema version.
|
|
conn.executescript(SCHEMA)
|
|
conn.commit()
|
|
_ensure_defaults(conn)
|
|
_fix_timestamp_defaults(conn)
|
|
conn.commit()
|
|
return
|
|
|
|
# Create any missing tables (CREATE TABLE IF NOT EXISTS is idempotent),
|
|
# then add any columns missing from older databases. The DDL is
|
|
# committed first so a later best-effort column migration (which may
|
|
# roll back its own transaction) can never undo the schema creation.
|
|
conn.executescript(SCHEMA)
|
|
conn.commit()
|
|
_migrate_missing_columns(conn)
|
|
|
|
# Bootstrap defaults (Root user, meta row, config rows).
|
|
_ensure_defaults(conn)
|
|
|
|
# Repair legacy timestamp DEFAULT expressions (old to_char HH:MM pattern).
|
|
_fix_timestamp_defaults(conn)
|
|
|
|
# Stamp the database with the program's schema version.
|
|
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()
|
|
finally:
|
|
conn.close()
|