migration: preserve startup accounts and system slice

2026-07-31 00:42:06 +08:00
parent 4083dceba3
commit 4002f096f4
37 changed files with 6821 additions and 6327 deletions
@@ -0,0 +1,236 @@
+from __future__ import annotations
+
+import sqlite3
+from datetime import datetime, timezone
+from typing import Any
+
+
+class AccountRepositoryMixin:
+    """Original SQLite account persistence methods, moved without query changes."""
+
+    def count_users(self) -> int:
+        with self.connect() as connection:
+            row = connection.execute("SELECT COUNT(*) AS total FROM users").fetchone()
+        return int(row["total"] if row else 0)
+
+    def first_user_id(self) -> int:
+        with self.connect() as connection:
+            row = connection.execute("SELECT MIN(id) AS id FROM users").fetchone()
+        return int(row["id"] or 0) if row else 0
+
+    def create_user(
+        self,
+        username: str,
+        password_salt: str,
+        password_hash: str,
+    ) -> dict[str, Any]:
+        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
+        try:
+            with self.connect() as connection:
+                role = "admin" if int(connection.execute("SELECT COUNT(*) FROM users").fetchone()[0]) == 0 else "user"
+                cursor = connection.execute(
+                    """
+                    INSERT INTO users
+                        (username, password_salt, password_hash, role, created_at, updated_at)
+                    VALUES (?, ?, ?, ?, ?, ?)
+                    """,
+                    (username, password_salt, password_hash, role, now, now),
+                )
+                user_id = int(cursor.lastrowid)
+        except sqlite3.IntegrityError as exc:
+            raise ValueError("该账号名已被使用。") from exc
+        return {"id": user_id, "username": username, "role": role, "created_at": now}
+
+    def user_by_username(self, username: str) -> dict[str, Any] | None:
+        with self.connect() as connection:
+            row = connection.execute(
+                """
+                SELECT id, username, password_salt, password_hash, role, llm_mode,
+                       membership_status, membership_plan, membership_starts_at,
+                       membership_expires_at, created_at
+                FROM users WHERE username = ? COLLATE NOCASE
+                """,
+                (username,),
+            ).fetchone()
+        return dict(row) if row else None
+
+    def user_password(self, user_id: int) -> dict[str, str] | None:
+        with self.connect() as connection:
+            row = connection.execute(
+                "SELECT password_salt, password_hash FROM users WHERE id = ?",
+                (user_id,),
+            ).fetchone()
+        return dict(row) if row else None
+
+    def update_user_password(self, user_id: int, password_salt: str, password_hash: str) -> bool:
+        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
+        with self.connect() as connection:
+            cursor = connection.execute(
+                "UPDATE users SET password_salt = ?, password_hash = ?, updated_at = ? WHERE id = ?",
+                (password_salt, password_hash, now, user_id),
+            )
+        return cursor.rowcount > 0
+
+    def delete_user(self, user_id: int) -> bool:
+        with self.connect() as connection:
+            cursor = connection.execute("DELETE FROM users WHERE id = ?", (user_id,))
+            return cursor.rowcount > 0
+
+    def create_session(
+        self,
+        session_hash: str,
+        user_id: int,
+        csrf_token: str,
+        expires_at: str,
+    ) -> None:
+        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
+        with self.connect() as connection:
+            connection.execute("DELETE FROM user_sessions WHERE expires_at <= ?", (now,))
+            connection.execute(
+                """
+                INSERT INTO user_sessions
+                    (token_hash, user_id, csrf_token, expires_at, created_at, last_seen_at)
+                VALUES (?, ?, ?, ?, ?, ?)
+                """,
+                (session_hash, user_id, csrf_token, expires_at, now, now),
+            )
+
+    def session_user(self, session_hash: str) -> dict[str, Any] | None:
+        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
+        with self.connect() as connection:
+            row = connection.execute(
+                """
+                SELECT u.id, u.username, u.role, u.llm_mode, u.membership_status,
+                       u.membership_plan, u.membership_starts_at, u.membership_expires_at,
+                       u.created_at, s.csrf_token, s.expires_at
+                FROM user_sessions AS s
+                JOIN users AS u ON u.id = s.user_id
+                WHERE s.token_hash = ? AND s.expires_at > ?
+                """,
+                (session_hash, now),
+            ).fetchone()
+            if row:
+                connection.execute(
+                    "UPDATE user_sessions SET last_seen_at = ? WHERE token_hash = ?",
+                    (now, session_hash),
+                )
+        return dict(row) if row else None
+
+    def delete_session(self, session_hash: str) -> bool:
+        with self.connect() as connection:
+            cursor = connection.execute(
+                "DELETE FROM user_sessions WHERE token_hash = ?",
+                (session_hash,),
+            )
+            return cursor.rowcount > 0
+
+    def get_user_credentials(self, user_id: int) -> str:
+        with self.connect() as connection:
+            row = connection.execute(
+                "SELECT encrypted_payload FROM user_credentials WHERE user_id = ?",
+                (user_id,),
+            ).fetchone()
+        return str(row["encrypted_payload"]) if row else ""
+
+    def save_user_credentials(self, user_id: int, encrypted_payload: str) -> None:
+        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
+        with self.connect() as connection:
+            connection.execute(
+                """
+                INSERT INTO user_credentials (user_id, encrypted_payload, updated_at)
+                VALUES (?, ?, ?)
+                ON CONFLICT(user_id) DO UPDATE SET
+                    encrypted_payload = excluded.encrypted_payload,
+                    updated_at = excluded.updated_at
+                """,
+                (user_id, encrypted_payload, now),
+            )
+
+    def list_user_credentials(self) -> list[dict[str, Any]]:
+        with self.connect() as connection:
+            rows = connection.execute(
+                "SELECT user_id, encrypted_payload FROM user_credentials ORDER BY user_id"
+            ).fetchall()
+        return [dict(row) for row in rows]
+
+    def user_access(self, user_id: int) -> dict[str, Any] | None:
+        with self.connect() as connection:
+            row = connection.execute(
+                """
+                SELECT id, username, role, llm_mode, membership_status, membership_plan,
+                       membership_starts_at, membership_expires_at, created_at
+                FROM users WHERE id = ?
+                """,
+                (user_id,),
+            ).fetchone()
+        return dict(row) if row else None
+
+    def list_users(self) -> list[dict[str, Any]]:
+        with self.connect() as connection:
+            rows = connection.execute(
+                """
+                SELECT id, username, role, llm_mode, membership_status, membership_plan,
+                       membership_starts_at, membership_expires_at, created_at
+                FROM users ORDER BY id
+                """
+            ).fetchall()
+        return [dict(row) for row in rows]
+
+    def update_user_llm_mode(self, user_id: int, mode: str) -> None:
+        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
+        with self.connect() as connection:
+            connection.execute(
+                "UPDATE users SET llm_mode = ?, updated_at = ? WHERE id = ?",
+                (mode, now, user_id),
+            )
+
+    def update_membership(
+        self,
+        user_id: int,
+        status: str,
+        plan: str,
+        starts_at: str | None,
+        expires_at: str | None,
+    ) -> bool:
+        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
+        with self.connect() as connection:
+            cursor = connection.execute(
+                """
+                UPDATE users
+                SET membership_status = ?, membership_plan = ?,
+                    membership_starts_at = ?, membership_expires_at = ?, updated_at = ?
+                WHERE id = ?
+                """,
+                (status, plan, starts_at, expires_at, now, user_id),
+            )
+        return cursor.rowcount > 0
+
+    def get_user_birth_profile(self, user_id: int) -> str:
+        with self.connect() as connection:
+            row = connection.execute(
+                "SELECT encrypted_payload FROM user_birth_profiles WHERE user_id = ?",
+                (user_id,),
+            ).fetchone()
+        return str(row["encrypted_payload"]) if row else ""
+
+    def save_user_birth_profile(self, user_id: int, encrypted_payload: str) -> None:
+        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
+        with self.connect() as connection:
+            connection.execute(
+                """
+                INSERT INTO user_birth_profiles (user_id, encrypted_payload, updated_at)
+                VALUES (?, ?, ?)
+                ON CONFLICT(user_id) DO UPDATE SET
+                    encrypted_payload = excluded.encrypted_payload,
+                    updated_at = excluded.updated_at
+                """,
+                (user_id, encrypted_payload, now),
+            )
+
+    def delete_user_birth_profile(self, user_id: int) -> bool:
+        with self.connect() as connection:
+            cursor = connection.execute(
+                "DELETE FROM user_birth_profiles WHERE user_id = ?",
+                (user_id,),
+            )
+            return cursor.rowcount > 0