| import json
|
| import secrets
|
| from typing import List, Optional, Dict, Any
|
| import numpy as np
|
| from sentence_transformers import SentenceTransformer
|
| import psycopg2
|
| import psycopg2.pool
|
|
|
| from .config import config
|
|
|
|
|
| class DatabaseManager:
|
| def __init__(self):
|
| self._pool = psycopg2.pool.ThreadedConnectionPool(
|
| minconn=1,
|
| maxconn=10,
|
| dsn=config.DATABASE_URL,
|
| )
|
| self._init_schema()
|
| self.embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
|
|
|
|
|
|
|
|
|
| class _Conn:
|
| def __init__(self, pool):
|
| self._pool = pool
|
| self.conn = None
|
|
|
| def __enter__(self):
|
| self.conn = self._pool.getconn()
|
| return self.conn
|
|
|
| def __exit__(self, exc_type, *_):
|
| if self.conn:
|
| if exc_type:
|
| self.conn.rollback()
|
| self._pool.putconn(self.conn)
|
|
|
| def _conn(self):
|
| return self._Conn(self._pool)
|
|
|
|
|
|
|
|
|
| def _init_schema(self):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute("""
|
| CREATE TABLE IF NOT EXISTS users (
|
| id SERIAL PRIMARY KEY,
|
| username TEXT UNIQUE NOT NULL,
|
| password_hash TEXT NOT NULL,
|
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
| )
|
| """)
|
| cur.execute("""
|
| CREATE TABLE IF NOT EXISTS conversations (
|
| id SERIAL PRIMARY KEY,
|
| user_id INTEGER NOT NULL REFERENCES users(id),
|
| title TEXT,
|
| started_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
| )
|
| """)
|
| cur.execute("""
|
| CREATE TABLE IF NOT EXISTS messages (
|
| id SERIAL PRIMARY KEY,
|
| conversation_id INTEGER NOT NULL REFERENCES conversations(id),
|
| role TEXT NOT NULL CHECK(role IN ('user','assistant')),
|
| content TEXT NOT NULL,
|
| timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
| )
|
| """)
|
| cur.execute("""
|
| CREATE TABLE IF NOT EXISTS user_profiles (
|
| user_id INTEGER PRIMARY KEY REFERENCES users(id),
|
| summary_text TEXT DEFAULT '',
|
| last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
| )
|
| """)
|
| cur.execute("""
|
| CREATE TABLE IF NOT EXISTS memories (
|
| id SERIAL PRIMARY KEY,
|
| user_id INTEGER NOT NULL REFERENCES users(id),
|
| content TEXT NOT NULL,
|
| embedding TEXT NOT NULL,
|
| metadata TEXT DEFAULT '{}',
|
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
| )
|
| """)
|
| cur.execute(
|
| "ALTER TABLE users ADD COLUMN IF NOT EXISTS is_guest BOOLEAN DEFAULT FALSE"
|
| )
|
| cur.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS guest_ip TEXT")
|
| cur.execute("""
|
| CREATE TABLE IF NOT EXISTS guest_ip_log (
|
| id SERIAL PRIMARY KEY,
|
| ip TEXT NOT NULL,
|
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
| )
|
| """)
|
| cur.execute(
|
| "CREATE INDEX IF NOT EXISTS idx_guest_ip_log_ip ON guest_ip_log(ip)"
|
| )
|
| cur.execute("""
|
| CREATE TABLE IF NOT EXISTS shared_conversations (
|
| id SERIAL PRIMARY KEY,
|
| token TEXT UNIQUE NOT NULL,
|
| user_id INTEGER NOT NULL REFERENCES users(id),
|
| conversation_id INTEGER NOT NULL REFERENCES conversations(id),
|
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
| )
|
| """)
|
| cur.execute(
|
| "CREATE INDEX IF NOT EXISTS idx_shared_token ON shared_conversations(token)"
|
| )
|
| cur.execute(
|
| "ALTER TABLE users ADD COLUMN IF NOT EXISTS is_admin BOOLEAN DEFAULT FALSE"
|
| )
|
| cur.execute(
|
| "ALTER TABLE users ADD COLUMN IF NOT EXISTS stripe_customer_id TEXT"
|
| )
|
| cur.execute(
|
| "ALTER TABLE users ADD COLUMN IF NOT EXISTS subscription_status TEXT"
|
| )
|
| cur.execute(
|
| "ALTER TABLE users ADD COLUMN IF NOT EXISTS subscription_plan TEXT"
|
| )
|
| cur.execute(
|
| "ALTER TABLE users ADD COLUMN IF NOT EXISTS subscription_id TEXT"
|
| )
|
| cur.execute(
|
| "ALTER TABLE users ADD COLUMN IF NOT EXISTS current_period_end TIMESTAMP"
|
| )
|
| cur.execute(
|
| "ALTER TABLE users ADD COLUMN IF NOT EXISTS free_messages_used INTEGER DEFAULT 0"
|
| )
|
| cur.execute(
|
| "ALTER TABLE users ADD COLUMN IF NOT EXISTS sub_messages_used INTEGER DEFAULT 0"
|
| )
|
| cur.execute(
|
| "ALTER TABLE users ADD COLUMN IF NOT EXISTS sub_period_start TIMESTAMP"
|
| )
|
| cur.execute("""
|
| CREATE TABLE IF NOT EXISTS payment_history (
|
| id SERIAL PRIMARY KEY,
|
| user_id INTEGER NOT NULL REFERENCES users(id),
|
| stripe_payment_id TEXT,
|
| amount_cents INTEGER,
|
| currency TEXT DEFAULT 'usd',
|
| status TEXT,
|
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
| )
|
| """)
|
| cur.execute(
|
| """CREATE UNIQUE INDEX IF NOT EXISTS idx_payment_history_stripe_payment_id
|
| ON payment_history(stripe_payment_id)
|
| WHERE stripe_payment_id IS NOT NULL"""
|
| )
|
|
|
| cur.execute("""
|
| CREATE TABLE IF NOT EXISTS admin_audit_log (
|
| id SERIAL PRIMARY KEY,
|
| admin_id INTEGER REFERENCES users(id),
|
| action TEXT NOT NULL,
|
| target_user_id INTEGER,
|
| ip_address TEXT,
|
| details TEXT,
|
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
| )
|
| """)
|
| cur.execute(
|
| "CREATE INDEX IF NOT EXISTS idx_audit_log_admin ON admin_audit_log(admin_id)"
|
| )
|
| cur.execute(
|
| "CREATE INDEX IF NOT EXISTS idx_audit_log_created ON admin_audit_log(created_at)"
|
| )
|
|
|
| cur.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS totp_secret TEXT")
|
| cur.execute(
|
| "ALTER TABLE users ADD COLUMN IF NOT EXISTS totp_pending_secret TEXT"
|
| )
|
| cur.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS email TEXT")
|
|
|
| cur.execute("""
|
| CREATE TABLE IF NOT EXISTS oauth_accounts (
|
| id SERIAL PRIMARY KEY,
|
| user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
| provider TEXT NOT NULL,
|
| provider_user_id TEXT NOT NULL,
|
| email TEXT,
|
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
| UNIQUE(provider, provider_user_id)
|
| )
|
| """)
|
| cur.execute(
|
| "CREATE INDEX IF NOT EXISTS idx_oauth_user_id ON oauth_accounts(user_id)"
|
| )
|
| cur.execute(
|
| "CREATE INDEX IF NOT EXISTS idx_users_email ON users(email) WHERE email IS NOT NULL"
|
| )
|
|
|
|
|
| cur.execute("""
|
| CREATE TABLE IF NOT EXISTS conversation_summaries (
|
| id SERIAL PRIMARY KEY,
|
| conversation_id INTEGER NOT NULL REFERENCES conversations(id),
|
| user_id INTEGER NOT NULL REFERENCES users(id),
|
| summary_text TEXT NOT NULL,
|
| message_count INTEGER DEFAULT 0,
|
| updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
| )
|
| """)
|
| cur.execute(
|
| "CREATE INDEX IF NOT EXISTS idx_conv_summaries_conv ON conversation_summaries(conversation_id)"
|
| )
|
| cur.execute(
|
| "CREATE UNIQUE INDEX IF NOT EXISTS idx_conv_summaries_unique_conv ON conversation_summaries(conversation_id)"
|
| )
|
| cur.execute(
|
| "CREATE INDEX IF NOT EXISTS idx_conv_summaries_user ON conversation_summaries(user_id)"
|
| )
|
|
|
| cur.execute("""
|
| CREATE TABLE IF NOT EXISTS user_facts (
|
| id SERIAL PRIMARY KEY,
|
| user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
| category TEXT NOT NULL DEFAULT 'general',
|
| content TEXT NOT NULL,
|
| importance FLOAT DEFAULT 0.5,
|
| confidence FLOAT DEFAULT 0.8,
|
| source_conv_id INTEGER REFERENCES conversations(id),
|
| created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
| last_reinforced TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
| is_pinned BOOLEAN DEFAULT FALSE
|
| )
|
| """)
|
| cur.execute(
|
| "CREATE INDEX IF NOT EXISTS idx_user_facts_user ON user_facts(user_id)"
|
| )
|
| cur.execute(
|
| "CREATE INDEX IF NOT EXISTS idx_user_facts_importance ON user_facts(user_id, importance DESC)"
|
| )
|
| cur.execute(
|
| "CREATE UNIQUE INDEX IF NOT EXISTS idx_user_facts_unique_content ON user_facts(user_id, content)"
|
| )
|
|
|
| cur.execute(
|
| "ALTER TABLE conversations ADD COLUMN IF NOT EXISTS openwebui_chat_id TEXT"
|
| )
|
| cur.execute(
|
| "CREATE UNIQUE INDEX IF NOT EXISTS idx_conv_openwebui_chat ON conversations(openwebui_chat_id) WHERE openwebui_chat_id IS NOT NULL"
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
|
|
|
|
|
|
| def _embed(self, text: str) -> List[float]:
|
| return self.embedding_model.encode([text])[0].tolist()
|
|
|
| @staticmethod
|
| def _cosine(a: List[float], b: List[float]) -> float:
|
| va, vb = np.array(a), np.array(b)
|
| na, nb = np.linalg.norm(va), np.linalg.norm(vb)
|
| return float(np.dot(va, vb) / (na * nb)) if na and nb else 0.0
|
|
|
|
|
|
|
|
|
| def get_user_by_username(self, username: str) -> Optional[Dict]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT id, username, password_hash FROM users WHERE username = %s",
|
| (username,),
|
| )
|
| row = cur.fetchone()
|
| cur.close()
|
| return (
|
| {"id": row[0], "username": row[1], "password_hash": row[2]} if row else None
|
| )
|
|
|
| def create_user(self, username: str, password_hash: str) -> int:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "INSERT INTO users (username, password_hash) VALUES (%s, %s) RETURNING id",
|
| (username, password_hash),
|
| )
|
| row = cur.fetchone()
|
| user_id: int = row[0] if row else 0
|
| cur.execute(
|
| "INSERT INTO user_profiles (user_id, summary_text) VALUES (%s, '')",
|
| (user_id,),
|
| )
|
| conn.commit()
|
| cur.close()
|
| return user_id
|
|
|
| def verify_password(self, username: str, password: str) -> Optional[int]:
|
| from .auth import verify_password as _verify
|
|
|
| user = self.get_user_by_username(username)
|
| if user and _verify(password, user["password_hash"]):
|
| return user["id"]
|
| return None
|
|
|
|
|
|
|
|
|
| def create_conversation(self, user_id: int) -> int:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "INSERT INTO conversations (user_id) VALUES (%s) RETURNING id",
|
| (user_id,),
|
| )
|
| row = cur.fetchone()
|
| conv_id: int = row[0] if row else 0
|
| conn.commit()
|
| cur.close()
|
| return conv_id
|
|
|
|
|
|
|
|
|
| def get_or_create_openwebui_user(self, email: str) -> int:
|
| """Look up or auto-provision a Montaigne user for an OpenWebUI user."""
|
| existing = self.get_user_by_email(email)
|
| if existing and not existing.get("is_guest"):
|
| return existing["id"]
|
|
|
| username = email.split("@")[0]
|
| if self.get_user_by_username(username):
|
| username = f"{username}_{secrets.token_hex(4)}"
|
| return self.create_oauth_user(username, email)
|
|
|
| def get_or_create_conversation_by_openwebui_chat(
|
| self, user_id: int, openwebui_chat_id: str
|
| ) -> int:
|
| """Map an OpenWebUI chat ID to a Montaigne conversation (create if new)."""
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT id FROM conversations WHERE openwebui_chat_id = %s",
|
| (openwebui_chat_id,),
|
| )
|
| row = cur.fetchone()
|
| if row:
|
| conv_id: int = row[0]
|
| else:
|
| cur.execute(
|
| "INSERT INTO conversations (user_id, openwebui_chat_id) VALUES (%s, %s) RETURNING id",
|
| (user_id, openwebui_chat_id),
|
| )
|
| row = cur.fetchone()
|
| conv_id = row[0] if row else 0
|
|
|
| cur.execute(
|
| "INSERT INTO user_profiles (user_id, summary_text) VALUES (%s, '') ON CONFLICT (user_id) DO NOTHING",
|
| (user_id,),
|
| )
|
| conn.commit()
|
| cur.close()
|
| return conv_id
|
|
|
| def get_user_conversations(self, user_id: int, limit: int = 20) -> List[Dict]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT id, title, started_at FROM conversations WHERE user_id = %s ORDER BY started_at DESC LIMIT %s",
|
| (user_id, limit),
|
| )
|
| rows = cur.fetchall()
|
| cur.close()
|
| return [{"id": r[0], "title": r[1], "started_at": str(r[2])} for r in rows]
|
|
|
| def update_conversation_title(self, conversation_id: int, title: str):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "UPDATE conversations SET title = %s WHERE id = %s",
|
| (title, conversation_id),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
|
|
|
|
|
|
| def add_message(self, conversation_id: int, role: str, content: str) -> int:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "INSERT INTO messages (conversation_id, role, content) VALUES (%s, %s, %s) RETURNING id",
|
| (conversation_id, role, content),
|
| )
|
| row = cur.fetchone()
|
| msg_id: int = row[0] if row else 0
|
| conn.commit()
|
| cur.close()
|
| return msg_id
|
|
|
| def get_conversation_messages(
|
| self, conversation_id: int, limit: int = 50
|
| ) -> List[Dict]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT role, content, timestamp FROM messages WHERE conversation_id = %s ORDER BY timestamp ASC LIMIT %s",
|
| (conversation_id, limit),
|
| )
|
| rows = cur.fetchall()
|
| cur.close()
|
| return [{"role": r[0], "content": r[1], "timestamp": str(r[2])} for r in rows]
|
|
|
| def get_recent_messages(self, user_id: int, limit: int = 10) -> List[Dict]:
|
| """Last N messages across ALL of a user's conversations, in chronological order."""
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """SELECT m.role, m.content, m.timestamp
|
| FROM messages m
|
| JOIN conversations c ON m.conversation_id = c.id
|
| WHERE c.user_id = %s
|
| ORDER BY m.timestamp DESC
|
| LIMIT %s""",
|
| (user_id, limit),
|
| )
|
| rows = cur.fetchall()
|
| cur.close()
|
| msgs = [{"role": r[0], "content": r[1], "timestamp": str(r[2])} for r in rows]
|
| return list(reversed(msgs))
|
|
|
| def count_conversation_messages(self, conversation_id: int) -> int:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT COUNT(*) FROM messages WHERE conversation_id = %s",
|
| (conversation_id,),
|
| )
|
| count = cur.fetchone()[0]
|
| cur.close()
|
| return count
|
|
|
| def get_conversation_turns(
|
| self, conversation_id: int, limit: int = 20, offset_from_end: int = 0
|
| ) -> List[Dict]:
|
| """Last N turns within a single conversation, in chronological order."""
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """SELECT role, content, timestamp, id
|
| FROM messages
|
| WHERE conversation_id = %s
|
| ORDER BY id DESC
|
| LIMIT %s OFFSET %s""",
|
| (conversation_id, limit, offset_from_end),
|
| )
|
| rows = cur.fetchall()
|
| cur.close()
|
| msgs = [
|
| {"role": r[0], "content": r[1], "timestamp": str(r[2]), "id": r[3]}
|
| for r in reversed(rows)
|
| ]
|
| return msgs
|
|
|
|
|
|
|
|
|
|
|
| def get_conversation_summary(self, conversation_id: int) -> Optional[str]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """SELECT summary_text FROM conversation_summaries
|
| WHERE conversation_id = %s ORDER BY updated_at DESC LIMIT 1""",
|
| (conversation_id,),
|
| )
|
| row = cur.fetchone()
|
| cur.close()
|
| return row[0] if row else None
|
|
|
| def save_conversation_summary(
|
| self,
|
| conversation_id: int,
|
| user_id: int,
|
| summary_text: str,
|
| message_count: int = 0,
|
| ):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """INSERT INTO conversation_summaries
|
| (conversation_id, user_id, summary_text, message_count)
|
| VALUES (%s, %s, %s, %s)
|
| ON CONFLICT (conversation_id)
|
| DO UPDATE SET
|
| summary_text = EXCLUDED.summary_text,
|
| message_count = EXCLUDED.message_count,
|
| updated_at = CURRENT_TIMESTAMP""",
|
| (conversation_id, user_id, summary_text, message_count),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
|
|
|
|
|
|
|
|
| def add_user_fact(
|
| self,
|
| user_id: int,
|
| content: str,
|
| category: str = "general",
|
| importance: float = 0.5,
|
| source_conv_id: Optional[int] = None,
|
| ):
|
| """Insert a fact if it doesn't already exist (same user + same content)."""
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """INSERT INTO user_facts
|
| (user_id, category, content, importance, source_conv_id)
|
| VALUES (%s, %s, %s, %s, %s)
|
| ON CONFLICT (user_id, content)
|
| DO UPDATE SET
|
| category = EXCLUDED.category,
|
| importance = GREATEST(user_facts.importance, EXCLUDED.importance),
|
| source_conv_id = COALESCE(EXCLUDED.source_conv_id, user_facts.source_conv_id),
|
| last_reinforced = CURRENT_TIMESTAMP""",
|
| (user_id, category, content, importance, source_conv_id),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
| def get_user_facts(
|
| self, user_id: int, limit: int = 8, category: Optional[str] = None
|
| ) -> List[Dict]:
|
| """Top facts for a user, ordered by importance then recency."""
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| if category:
|
| cur.execute(
|
| """SELECT category, content, importance, is_pinned
|
| FROM user_facts WHERE user_id = %s AND category = %s
|
| ORDER BY is_pinned DESC, importance DESC, last_reinforced DESC
|
| LIMIT %s""",
|
| (user_id, category, limit),
|
| )
|
| else:
|
| cur.execute(
|
| """SELECT category, content, importance, is_pinned
|
| FROM user_facts WHERE user_id = %s
|
| ORDER BY is_pinned DESC, importance DESC, last_reinforced DESC
|
| LIMIT %s""",
|
| (user_id, limit),
|
| )
|
| rows = cur.fetchall()
|
| cur.close()
|
| return [
|
| {"category": r[0], "content": r[1], "importance": r[2], "is_pinned": r[3]}
|
| for r in rows
|
| ]
|
|
|
| def delete_user_fact(self, user_id: int, content: str) -> bool:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "DELETE FROM user_facts WHERE user_id = %s AND content = %s",
|
| (user_id, content),
|
| )
|
| deleted = cur.rowcount > 0
|
| conn.commit()
|
| cur.close()
|
| return deleted
|
|
|
| def get_conversation_title(self, conversation_id: int) -> Optional[str]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT title FROM conversations WHERE id = %s",
|
| (conversation_id,),
|
| )
|
| row = cur.fetchone()
|
| cur.close()
|
| return row[0] if row else None
|
|
|
| def get_title_sample_messages(self, conversation_id: int) -> List[Dict]:
|
| """First message + last 4 messages for title generation sampling."""
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """SELECT role, content, timestamp FROM messages
|
| WHERE conversation_id = %s
|
| ORDER BY timestamp ASC""",
|
| (conversation_id,),
|
| )
|
| rows = cur.fetchall()
|
| cur.close()
|
| if not rows:
|
| return []
|
| msgs = [{"role": r[0], "content": r[1], "timestamp": str(r[2])} for r in rows]
|
| if len(msgs) <= 5:
|
| return msgs
|
| return [msgs[0]] + msgs[-4:]
|
|
|
|
|
|
|
|
|
| def get_user_profile(self, user_id: int) -> str:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT summary_text FROM user_profiles WHERE user_id = %s",
|
| (user_id,),
|
| )
|
| row = cur.fetchone()
|
| cur.close()
|
| return row[0] if row else ""
|
|
|
| def update_user_profile(self, user_id: int, summary_text: str):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "UPDATE user_profiles SET summary_text = %s, last_updated = CURRENT_TIMESTAMP WHERE user_id = %s",
|
| (summary_text, user_id),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
| def update_username(self, user_id: int, username: str) -> bool:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| try:
|
| cur.execute(
|
| "UPDATE users SET username = %s WHERE id = %s",
|
| (username, user_id),
|
| )
|
| conn.commit()
|
| return True
|
| except Exception:
|
| conn.rollback()
|
| return False
|
| finally:
|
| cur.close()
|
|
|
| def update_password(self, user_id: int, password_hash: str) -> bool:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "UPDATE users SET password_hash = %s WHERE id = %s",
|
| (password_hash, user_id),
|
| )
|
| success = cur.rowcount > 0
|
| conn.commit()
|
| cur.close()
|
| return success
|
|
|
| def update_email(self, user_id: int, email: Optional[str]) -> bool:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "UPDATE users SET email = %s WHERE id = %s",
|
| (email, user_id),
|
| )
|
| success = cur.rowcount > 0
|
| conn.commit()
|
| cur.close()
|
| return success
|
|
|
|
|
|
|
|
|
| def add_memory(
|
| self, user_id: int, text: str, metadata: Optional[Dict[str, Any]] = None
|
| ):
|
| embedding_json = json.dumps(self._embed(text))
|
| metadata_json = json.dumps(metadata or {})
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "INSERT INTO memories (user_id, content, embedding, metadata) VALUES (%s, %s, %s, %s)",
|
| (user_id, text, embedding_json, metadata_json),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
| def search_memories(
|
| self, user_id: int, query: str, n_results: int = 5
|
| ) -> List[Dict]:
|
| query_vec = self._embed(query)
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT content, embedding, metadata FROM memories WHERE user_id = %s",
|
| (user_id,),
|
| )
|
| rows = cur.fetchall()
|
| cur.close()
|
|
|
| scored = []
|
| for content, emb_json, meta_json in rows:
|
| sim = self._cosine(query_vec, json.loads(emb_json))
|
| scored.append(
|
| {
|
| "content": content,
|
| "metadata": json.loads(meta_json) if meta_json else {},
|
| "similarity": sim,
|
| }
|
| )
|
|
|
| scored.sort(key=lambda x: x["similarity"], reverse=True)
|
| return [m for m in scored[:n_results] if m["similarity"] > 0.3]
|
|
|
| def create_guest_user(self, ip: str) -> int:
|
| import uuid
|
|
|
| username = f"guest_{uuid.uuid4().hex[:12]}"
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "INSERT INTO users (username, password_hash, is_guest, guest_ip) VALUES (%s, '', TRUE, %s) RETURNING id",
|
| (username, ip),
|
| )
|
| row = cur.fetchone()
|
| user_id: int = row[0] if row else 0
|
| cur.execute(
|
| "INSERT INTO user_profiles (user_id, summary_text) VALUES (%s, '')",
|
| (user_id,),
|
| )
|
| conn.commit()
|
| cur.close()
|
| return user_id
|
|
|
| def get_or_create_guest_user(self, ip: str) -> int:
|
| """Look up or create a guest user for a given IP address.
|
|
|
| Reuses the most recent guest user for this IP so that the per-user
|
| message limit is enforced across requests from the same visitor.
|
| """
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT id FROM users WHERE is_guest = TRUE AND guest_ip = %s ORDER BY created_at DESC LIMIT 1",
|
| (ip,),
|
| )
|
| row = cur.fetchone()
|
| if row:
|
| user_id: int = row[0]
|
| cur.close()
|
| return user_id
|
| cur.close()
|
| return self.create_guest_user(ip)
|
|
|
| def convert_guest_to_user(
|
| self, user_id: int, username: str, password_hash: str
|
| ) -> bool:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "UPDATE users SET username = %s, password_hash = %s, is_guest = FALSE, guest_ip = NULL WHERE id = %s AND is_guest = TRUE",
|
| (username, password_hash, user_id),
|
| )
|
| success = cur.rowcount > 0
|
| conn.commit()
|
| cur.close()
|
| return success
|
|
|
| def is_guest_user(self, user_id: int) -> bool:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute("SELECT is_guest FROM users WHERE id = %s", (user_id,))
|
| row = cur.fetchone()
|
| cur.close()
|
| return row[0] if row else False
|
|
|
| def count_ip_guest_sessions(self, ip: str, hours: int = 24) -> int:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """SELECT COUNT(*) FROM users
|
| WHERE guest_ip = %s AND is_guest = TRUE
|
| AND created_at > NOW() - INTERVAL '%s hours'""",
|
| (ip, hours),
|
| )
|
| count = cur.fetchone()[0]
|
| cur.close()
|
| return count
|
|
|
| def log_guest_ip(self, ip: str):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute("INSERT INTO guest_ip_log (ip) VALUES (%s)", (ip,))
|
| conn.commit()
|
| cur.close()
|
|
|
| def count_user_messages(self, user_id: int) -> int:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """SELECT COUNT(*) FROM messages m
|
| JOIN conversations c ON m.conversation_id = c.id
|
| WHERE c.user_id = %s AND m.role = 'user'""",
|
| (user_id,),
|
| )
|
| count = cur.fetchone()[0]
|
| cur.close()
|
| return count
|
|
|
| def cleanup_expired_guests(self, hours: int = 48):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """DELETE FROM users
|
| WHERE is_guest = TRUE AND created_at < NOW() - INTERVAL '%s hours'""",
|
| (hours,),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
| def create_share_token(self, user_id: int, conversation_id: int) -> str:
|
| import uuid
|
|
|
| token = uuid.uuid4().hex
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "INSERT INTO shared_conversations (token, user_id, conversation_id) VALUES (%s, %s, %s)",
|
| (token, user_id, conversation_id),
|
| )
|
| conn.commit()
|
| cur.close()
|
| return token
|
|
|
| def get_shared_conversation(self, token: str) -> Optional[Dict]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """SELECT c.title, c.id FROM shared_conversations sc
|
| JOIN conversations c ON sc.conversation_id = c.id
|
| WHERE sc.token = %s""",
|
| (token,),
|
| )
|
| row = cur.fetchone()
|
| cur.close()
|
| if not row:
|
| return None
|
| conv_id = row[1]
|
| title = row[0]
|
| messages = self.get_conversation_messages(conv_id)
|
| return {"title": title, "messages": messages}
|
|
|
| def get_user_by_id(self, user_id: int) -> Optional[Dict]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """SELECT id, username, is_guest, is_admin, stripe_customer_id,
|
| subscription_status, subscription_plan, subscription_id,
|
| current_period_end, free_messages_used, created_at, email
|
| FROM users WHERE id = %s""",
|
| (user_id,),
|
| )
|
| row = cur.fetchone()
|
| cur.close()
|
| if not row:
|
| return None
|
| return {
|
| "id": row[0],
|
| "username": row[1],
|
| "is_guest": row[2],
|
| "is_admin": row[3],
|
| "stripe_customer_id": row[4],
|
| "subscription_status": row[5],
|
| "subscription_plan": row[6],
|
| "subscription_id": row[7],
|
| "current_period_end": str(row[8]) if row[8] else None,
|
| "free_messages_used": row[9] or 0,
|
| "created_at": str(row[10]) if row[10] else None,
|
| "email": row[11],
|
| }
|
|
|
| def get_free_messages_used(self, user_id: int) -> int:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT free_messages_used FROM users WHERE id = %s", (user_id,)
|
| )
|
| row = cur.fetchone()
|
| cur.close()
|
| return row[0] if row and row[0] is not None else 0
|
|
|
| def increment_free_messages_used(self, user_id: int) -> int:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """UPDATE users SET free_messages_used = COALESCE(free_messages_used, 0) + 1
|
| WHERE id = %s RETURNING free_messages_used""",
|
| (user_id,),
|
| )
|
| row = cur.fetchone()
|
| conn.commit()
|
| cur.close()
|
| return row[0] if row else 0
|
|
|
| def try_consume_free_message(self, user_id: int) -> bool:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """UPDATE users
|
| SET free_messages_used = COALESCE(free_messages_used, 0) + 1
|
| WHERE id = %s AND COALESCE(free_messages_used, 0) < %s
|
| RETURNING free_messages_used""",
|
| (user_id, config.FREE_MESSAGE_LIMIT),
|
| )
|
| success = cur.fetchone() is not None
|
| conn.commit()
|
| cur.close()
|
| return success
|
|
|
| def refund_free_message(self, user_id: int):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """UPDATE users
|
| SET free_messages_used = GREATEST(COALESCE(free_messages_used, 0) - 1, 0)
|
| WHERE id = %s""",
|
| (user_id,),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
| def set_stripe_customer_id(self, user_id: int, customer_id: str):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "UPDATE users SET stripe_customer_id = %s WHERE id = %s",
|
| (customer_id, user_id),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
| def try_consume_subscriber_message(self, user_id: int) -> bool:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """UPDATE users
|
| SET sub_messages_used = CASE
|
| WHEN sub_period_start IS NULL
|
| OR sub_period_start < DATE_TRUNC('month', CURRENT_DATE)
|
| THEN 1
|
| ELSE COALESCE(sub_messages_used, 0) + 1
|
| END,
|
| sub_period_start = CASE
|
| WHEN sub_period_start IS NULL
|
| OR sub_period_start < DATE_TRUNC('month', CURRENT_DATE)
|
| THEN DATE_TRUNC('month', CURRENT_DATE)
|
| ELSE sub_period_start
|
| END
|
| WHERE id = %s
|
| RETURNING sub_messages_used""",
|
| (user_id,),
|
| )
|
| row = cur.fetchone()
|
| used = row[0] if row else 0
|
| success = used <= config.SUBSCRIBER_MONTHLY_LIMIT
|
| if success:
|
| conn.commit()
|
| else:
|
| conn.rollback()
|
| cur.close()
|
| return success
|
|
|
| def refund_subscriber_message(self, user_id: int):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """UPDATE users
|
| SET sub_messages_used = GREATEST(COALESCE(sub_messages_used, 0) - 1, 0)
|
| WHERE id = %s""",
|
| (user_id,),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
| def get_subscriber_messages_used(self, user_id: int) -> int:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """SELECT sub_messages_used FROM users
|
| WHERE id = %s
|
| AND (sub_period_start IS NULL
|
| OR sub_period_start >= DATE_TRUNC('month', CURRENT_DATE))""",
|
| (user_id,),
|
| )
|
| row = cur.fetchone()
|
| cur.close()
|
| return row[0] if row and row[0] is not None else 0
|
|
|
| def update_subscription(
|
| self,
|
| user_id: int,
|
| status: Optional[str] = None,
|
| plan: Optional[str] = None,
|
| subscription_id: Optional[str] = None,
|
| current_period_end: Optional[str] = None,
|
| ):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| updates = []
|
| params = []
|
| if status is not None:
|
| updates.append("subscription_status = %s")
|
| params.append(status)
|
| if plan is not None:
|
| updates.append("subscription_plan = %s")
|
| params.append(plan)
|
| if subscription_id is not None:
|
| updates.append("subscription_id = %s")
|
| params.append(subscription_id)
|
| if current_period_end is not None:
|
| updates.append("current_period_end = %s")
|
| params.append(current_period_end)
|
| if not updates:
|
| return
|
| params.append(user_id)
|
| cur.execute(
|
| f"UPDATE users SET {', '.join(updates)} WHERE id = %s",
|
| params,
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
| def clear_subscription(self, user_id: int):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """UPDATE users SET subscription_status = NULL, subscription_plan = NULL,
|
| subscription_id = NULL, current_period_end = NULL WHERE id = %s""",
|
| (user_id,),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
| def get_user_by_stripe_customer_id(self, customer_id: str) -> Optional[int]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT id FROM users WHERE stripe_customer_id = %s",
|
| (customer_id,),
|
| )
|
| row = cur.fetchone()
|
| cur.close()
|
| return row[0] if row else None
|
|
|
| def add_payment(
|
| self,
|
| user_id: int,
|
| stripe_payment_id: str,
|
| amount_cents: int,
|
| status: str,
|
| currency: str = "usd",
|
| ):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """INSERT INTO payment_history (user_id, stripe_payment_id, amount_cents, currency, status)
|
| VALUES (%s, %s, %s, %s, %s)
|
| ON CONFLICT (stripe_payment_id) DO NOTHING""",
|
| (user_id, stripe_payment_id, amount_cents, currency, status),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
| def get_admin_users(
|
| self, page: int = 1, per_page: int = 20, search: str = ""
|
| ) -> Dict:
|
| offset = (page - 1) * per_page
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| search_clause = ""
|
| params = []
|
| if search:
|
| search_clause = "WHERE username ILIKE %s"
|
| params.append(f"%{search}%")
|
| params.extend([per_page, offset])
|
| cur.execute(
|
| f"""SELECT id, username, is_guest, is_admin, subscription_status, subscription_plan,
|
| free_messages_used, created_at, current_period_end
|
| FROM users {search_clause}
|
| ORDER BY created_at DESC
|
| LIMIT %s OFFSET %s""",
|
| params,
|
| )
|
| rows = cur.fetchall()
|
| cur.execute(
|
| f"SELECT COUNT(*) FROM users {search_clause}",
|
| params[:1] if search else [],
|
| )
|
| total = cur.fetchone()[0]
|
| cur.close()
|
| users = [
|
| {
|
| "id": r[0],
|
| "username": r[1],
|
| "is_guest": r[2],
|
| "is_admin": r[3],
|
| "subscription_status": r[4],
|
| "subscription_plan": r[5],
|
| "free_messages_used": r[6] or 0,
|
| "created_at": str(r[7]) if r[7] else None,
|
| "current_period_end": str(r[8]) if r[8] else None,
|
| }
|
| for r in rows
|
| ]
|
| return {"users": users, "total": total, "page": page, "per_page": per_page}
|
|
|
| def get_admin_stats(self) -> Dict:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute("SELECT COUNT(*) FROM users WHERE is_guest = FALSE")
|
| total_users = cur.fetchone()[0]
|
| cur.execute("SELECT COUNT(*) FROM users WHERE is_guest = TRUE")
|
| total_guests = cur.fetchone()[0]
|
| cur.execute(
|
| "SELECT COUNT(*) FROM users WHERE subscription_status = 'active'"
|
| )
|
| active_subscribers = cur.fetchone()[0]
|
| cur.execute(
|
| "SELECT COUNT(*) FROM users WHERE subscription_plan = 'monthly'"
|
| )
|
| monthly_subscribers = cur.fetchone()[0]
|
| cur.execute("SELECT COUNT(*) FROM users WHERE subscription_plan = 'annual'")
|
| annual_subscribers = cur.fetchone()[0]
|
| cur.execute("SELECT COUNT(*) FROM messages WHERE role = 'user'")
|
| total_messages = cur.fetchone()[0]
|
| cur.execute("SELECT COUNT(*) FROM conversations")
|
| total_conversations = cur.fetchone()[0]
|
| cur.execute(
|
| """SELECT COALESCE(SUM(amount_cents), 0) FROM payment_history
|
| WHERE status = 'succeeded'"""
|
| )
|
| total_revenue_cents = cur.fetchone()[0]
|
| cur.close()
|
| return {
|
| "total_users": total_users,
|
| "total_guests": total_guests,
|
| "active_subscribers": active_subscribers,
|
| "monthly_subscribers": monthly_subscribers,
|
| "annual_subscribers": annual_subscribers,
|
| "total_messages": total_messages,
|
| "total_conversations": total_conversations,
|
| "total_revenue_cents": total_revenue_cents,
|
| }
|
|
|
| def set_user_admin(self, user_id: int, is_admin: bool):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "UPDATE users SET is_admin = %s WHERE id = %s",
|
| (is_admin, user_id),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
| def delete_user(self, user_id: int) -> bool:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute("DELETE FROM user_profiles WHERE user_id = %s", (user_id,))
|
| cur.execute("DELETE FROM memories WHERE user_id = %s", (user_id,))
|
| cur.execute(
|
| "DELETE FROM shared_conversations WHERE user_id = %s", (user_id,)
|
| )
|
| cur.execute(
|
| """DELETE FROM messages WHERE conversation_id IN
|
| (SELECT id FROM conversations WHERE user_id = %s)""",
|
| (user_id,),
|
| )
|
| cur.execute("DELETE FROM conversations WHERE user_id = %s", (user_id,))
|
| cur.execute("DELETE FROM payment_history WHERE user_id = %s", (user_id,))
|
| cur.execute("DELETE FROM users WHERE id = %s", (user_id,))
|
| success = cur.rowcount > 0
|
| conn.commit()
|
| cur.close()
|
| return success
|
|
|
| def get_user_detail(self, user_id: int) -> Optional[Dict]:
|
| user = self.get_user_by_id(user_id)
|
| if not user:
|
| return None
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """SELECT id, title, started_at FROM conversations
|
| WHERE user_id = %s ORDER BY started_at DESC LIMIT 10""",
|
| (user_id,),
|
| )
|
| conv_rows = cur.fetchall()
|
| cur.execute(
|
| "SELECT summary_text FROM user_profiles WHERE user_id = %s", (user_id,)
|
| )
|
| profile_row = cur.fetchone()
|
| cur.execute(
|
| """SELECT COUNT(*) FROM messages m
|
| JOIN conversations c ON m.conversation_id = c.id
|
| WHERE c.user_id = %s AND m.role = 'user'""",
|
| (user_id,),
|
| )
|
| total_messages = cur.fetchone()[0]
|
| cur.close()
|
| conversations = [
|
| {"id": r[0], "title": r[1], "started_at": str(r[2])} for r in conv_rows
|
| ]
|
| return {
|
| **user,
|
| "conversations": conversations,
|
| "profile_summary": profile_row[0] if profile_row else "",
|
| "total_messages": total_messages,
|
| }
|
|
|
|
|
|
|
|
|
| def count_admins(self) -> int:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute("SELECT COUNT(*) FROM users WHERE is_admin = TRUE")
|
| count = cur.fetchone()[0]
|
| cur.close()
|
| return count
|
|
|
| def get_admin_ids(self, exclude_user_id: Optional[int] = None) -> List[int]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| if exclude_user_id:
|
| cur.execute(
|
| "SELECT id FROM users WHERE is_admin = TRUE AND id != %s",
|
| (exclude_user_id,),
|
| )
|
| else:
|
| cur.execute("SELECT id FROM users WHERE is_admin = TRUE")
|
| ids = [r[0] for r in cur.fetchall()]
|
| cur.close()
|
| return ids
|
|
|
| def is_last_admin(self, user_id: int) -> bool:
|
| """Return True if user_id is the sole remaining admin."""
|
| user = self.get_user_by_id(user_id)
|
| if not user or not user.get("is_admin"):
|
| return False
|
| return self.count_admins() <= 1
|
|
|
|
|
|
|
|
|
| def log_admin_action(
|
| self,
|
| admin_id: Optional[int],
|
| action: str,
|
| target_user_id: Optional[int] = None,
|
| ip_address: Optional[str] = None,
|
| details: Optional[str] = None,
|
| ):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """INSERT INTO admin_audit_log
|
| (admin_id, action, target_user_id, ip_address, details)
|
| VALUES (%s, %s, %s, %s, %s)""",
|
| (admin_id, action, target_user_id, ip_address, details),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
| def get_audit_logs(
|
| self,
|
| admin_id: Optional[int] = None,
|
| action: Optional[str] = None,
|
| limit: int = 50,
|
| ) -> List[Dict]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| conditions = []
|
| params: list = []
|
| if admin_id:
|
| conditions.append("al.admin_id = %s")
|
| params.append(admin_id)
|
| if action:
|
| conditions.append("al.action = %s")
|
| params.append(action)
|
| where = ("WHERE " + " AND ".join(conditions)) if conditions else ""
|
| params.append(limit)
|
| cur.execute(
|
| f"""SELECT al.id, al.admin_id, u.username, al.action,
|
| al.target_user_id, al.ip_address, al.details, al.created_at
|
| FROM admin_audit_log al
|
| LEFT JOIN users u ON al.admin_id = u.id
|
| {where}
|
| ORDER BY al.created_at DESC
|
| LIMIT %s""",
|
| params,
|
| )
|
| rows = cur.fetchall()
|
| cur.close()
|
| return [
|
| {
|
| "id": r[0],
|
| "admin_id": r[1],
|
| "admin_username": r[2],
|
| "action": r[3],
|
| "target_user_id": r[4],
|
| "ip_address": r[5],
|
| "details": r[6],
|
| "created_at": str(r[7]) if r[7] else None,
|
| }
|
| for r in rows
|
| ]
|
|
|
|
|
|
|
|
|
| def set_pending_totp_secret(self, user_id: int, secret: str):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "UPDATE users SET totp_pending_secret = %s WHERE id = %s",
|
| (secret, user_id),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
| def get_totp_secret(self, user_id: int) -> Optional[str]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute("SELECT totp_secret FROM users WHERE id = %s", (user_id,))
|
| row = cur.fetchone()
|
| cur.close()
|
| return row[0] if row else None
|
|
|
| def get_pending_totp_secret(self, user_id: int) -> Optional[str]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT totp_pending_secret FROM users WHERE id = %s", (user_id,)
|
| )
|
| row = cur.fetchone()
|
| cur.close()
|
| return row[0] if row else None
|
|
|
| def enable_totp_from_pending(self, user_id: int) -> bool:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| """UPDATE users
|
| SET totp_secret = totp_pending_secret,
|
| totp_pending_secret = NULL
|
| WHERE id = %s AND totp_pending_secret IS NOT NULL""",
|
| (user_id,),
|
| )
|
| success = cur.rowcount > 0
|
| conn.commit()
|
| cur.close()
|
| return success
|
|
|
| def clear_totp_secret(self, user_id: int):
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "UPDATE users SET totp_secret = NULL, totp_pending_secret = NULL WHERE id = %s",
|
| (user_id,),
|
| )
|
| conn.commit()
|
| cur.close()
|
|
|
|
|
|
|
|
|
| def get_oauth_account(self, provider: str, provider_user_id: str) -> Optional[Dict]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT id, user_id, provider, provider_user_id, email FROM oauth_accounts WHERE provider = %s AND provider_user_id = %s",
|
| (provider, provider_user_id),
|
| )
|
| row = cur.fetchone()
|
| cur.close()
|
| if not row:
|
| return None
|
| return {
|
| "id": row[0],
|
| "user_id": row[1],
|
| "provider": row[2],
|
| "provider_user_id": row[3],
|
| "email": row[4],
|
| }
|
|
|
| def get_user_by_email(self, email: str) -> Optional[Dict]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT id, username, password_hash, is_guest, email FROM users WHERE email = %s",
|
| (email,),
|
| )
|
| row = cur.fetchone()
|
| cur.close()
|
| if not row:
|
| return None
|
| return {
|
| "id": row[0],
|
| "username": row[1],
|
| "password_hash": row[2],
|
| "is_guest": row[3],
|
| "email": row[4],
|
| }
|
|
|
| def create_oauth_account(
|
| self,
|
| user_id: int,
|
| provider: str,
|
| provider_user_id: str,
|
| email: Optional[str] = None,
|
| ) -> int:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "INSERT INTO oauth_accounts (user_id, provider, provider_user_id, email) VALUES (%s, %s, %s, %s) RETURNING id",
|
| (user_id, provider, provider_user_id, email),
|
| )
|
| row = cur.fetchone()
|
| oauth_id: int = row[0] if row else 0
|
| conn.commit()
|
| cur.close()
|
| return oauth_id
|
|
|
| def get_oauth_accounts_for_user(self, user_id: int) -> List[Dict]:
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT id, provider, provider_user_id, email, created_at FROM oauth_accounts WHERE user_id = %s",
|
| (user_id,),
|
| )
|
| rows = cur.fetchall()
|
| cur.close()
|
| return [
|
| {
|
| "id": r[0],
|
| "provider": r[1],
|
| "provider_user_id": r[2],
|
| "email": r[3],
|
| "created_at": str(r[4]) if r[4] else None,
|
| }
|
| for r in rows
|
| ]
|
|
|
| def create_oauth_user(self, username: str, email: Optional[str] = None) -> int:
|
| """Create a new user for OAuth sign-in (no password)."""
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "INSERT INTO users (username, password_hash, email) VALUES (%s, '', %s) RETURNING id",
|
| (username, email),
|
| )
|
| row = cur.fetchone()
|
| user_id: int = row[0] if row else 0
|
| cur.execute(
|
| "INSERT INTO user_profiles (user_id, summary_text) VALUES (%s, '')",
|
| (user_id,),
|
| )
|
| conn.commit()
|
| cur.close()
|
| return user_id
|
|
|
| def convert_guest_to_oauth(
|
| self, user_id: int, username: str, email: Optional[str] = None
|
| ) -> bool:
|
| """Convert a guest user to a full user via OAuth (no password)."""
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "UPDATE users SET username = %s, password_hash = '', email = COALESCE(%s, email), is_guest = FALSE, guest_ip = NULL WHERE id = %s AND is_guest = TRUE",
|
| (username, email, user_id),
|
| )
|
| success = cur.rowcount > 0
|
| conn.commit()
|
| cur.close()
|
| return success
|
|
|
| def migrate_guest_to_user(self, guest_user_id: int, target_user_id: int) -> bool:
|
| if guest_user_id == target_user_id:
|
| return False
|
|
|
| with self._conn() as conn:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT is_guest FROM users WHERE id = %s",
|
| (guest_user_id,),
|
| )
|
| row = cur.fetchone()
|
| if not row or not row[0]:
|
| cur.close()
|
| return False
|
|
|
| cur.execute(
|
| "UPDATE conversations SET user_id = %s WHERE user_id = %s",
|
| (target_user_id, guest_user_id),
|
| )
|
| cur.execute(
|
| "UPDATE memories SET user_id = %s WHERE user_id = %s",
|
| (target_user_id, guest_user_id),
|
| )
|
| cur.execute(
|
| "UPDATE shared_conversations SET user_id = %s WHERE user_id = %s",
|
| (target_user_id, guest_user_id),
|
| )
|
|
|
| cur.execute(
|
| "SELECT summary_text FROM user_profiles WHERE user_id = %s",
|
| (guest_user_id,),
|
| )
|
| guest_profile_row = cur.fetchone()
|
| guest_summary = (
|
| guest_profile_row[0].strip()
|
| if guest_profile_row and guest_profile_row[0]
|
| else ""
|
| )
|
|
|
| cur.execute(
|
| "SELECT summary_text FROM user_profiles WHERE user_id = %s",
|
| (target_user_id,),
|
| )
|
| target_profile_row = cur.fetchone()
|
| target_summary = (
|
| target_profile_row[0].strip()
|
| if target_profile_row and target_profile_row[0]
|
| else ""
|
| )
|
|
|
| if guest_summary:
|
| merged_summary = guest_summary
|
| if target_summary and target_summary != guest_summary:
|
| merged_summary = f"{target_summary}\n\n{guest_summary}"
|
| cur.execute(
|
| "UPDATE user_profiles SET summary_text = %s, last_updated = CURRENT_TIMESTAMP WHERE user_id = %s",
|
| (merged_summary, target_user_id),
|
| )
|
|
|
| cur.execute(
|
| "DELETE FROM user_profiles WHERE user_id = %s", (guest_user_id,)
|
| )
|
| cur.execute(
|
| "DELETE FROM users WHERE id = %s AND is_guest = TRUE", (guest_user_id,)
|
| )
|
| conn.commit()
|
| cur.close()
|
| return True
|
|
|
| def close(self):
|
| self._pool.closeall()
|
|
|
