feat: Add database-backed user and subscription management, update project architecture documentation, and outline technical debt.

2026-03-06 19:10:05 +08:00
parent ce2a553e08
commit 14838fde4d
7 changed files with 696 additions and 495 deletions
@@ -0,0 +1,109 @@
+import sqlite3
+import os
+from datetime import datetime, timedelta
+from typing import Optional, Dict, Any
+from loguru import logger
+
+class DBManager:
+    def __init__(self, db_path: str = "polyweather.db"):
+        self.db_path = db_path
+        self._init_db()
+
+    def _get_connection(self):
+        return sqlite3.connect(self.db_path)
+
+    def _init_db(self):
+        """Create tables if they don't exist."""
+        # Ensure directory exists
+        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
+        with self._get_connection() as conn:
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS users (
+                    telegram_id INTEGER PRIMARY KEY,
+                    username TEXT,
+                    is_web_premium BOOLEAN DEFAULT 0,
+                    web_expiry TIMESTAMP,
+                    is_group_premium BOOLEAN DEFAULT 0,
+                    group_expiry TIMESTAMP,
+                    points INTEGER DEFAULT 0,
+                    message_count INTEGER DEFAULT 0,
+                    last_message_at TIMESTAMP,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            """)
+            conn.commit()
+            logger.info("Database initialized successfully.")
+
+    def get_user(self, telegram_id: int) -> Optional[Dict[str, Any]]:
+        with self._get_connection() as conn:
+            conn.row_factory = sqlite3.Row
+            cursor = conn.execute("SELECT * FROM users WHERE telegram_id = ?", (telegram_id,))
+            row = cursor.fetchone()
+            if row:
+                user = dict(row)
+                now = datetime.now()
+                if user['web_expiry']:
+                    expiry = datetime.fromisoformat(user['web_expiry'])
+                    if expiry < now:
+                        user['is_web_premium'] = False
+                if user['group_expiry']:
+                    expiry = datetime.fromisoformat(user['group_expiry'])
+                    if expiry < now:
+                        user['is_group_premium'] = False
+                return user
+        return None
+
+    def upsert_user(self, telegram_id: int, username: str):
+        with self._get_connection() as conn:
+            conn.execute("""
+                INSERT INTO users (telegram_id, username)
+                VALUES (?, ?)
+                ON CONFLICT(telegram_id) DO UPDATE SET
+                username = excluded.username
+            """, (telegram_id, username))
+            conn.commit()
+
+    def add_message_activity(self, telegram_id: int, points_to_add: int = 1):
+        """Update message count and add points with simple anti-spam logic."""
+        now = datetime.now()
+        with self._get_connection() as conn:
+            cursor = conn.execute("SELECT last_message_at FROM users WHERE telegram_id = ?", (telegram_id,))
+            row = cursor.fetchone()
+            if row and row[0]:
+                last_at = datetime.fromisoformat(row[0])
+                if (now - last_at).total_seconds() < 5:  # 5 second cooldown
+                    return False
+            
+            conn.execute("""
+                UPDATE users 
+                SET message_count = message_count + 1,
+                    points = points + ?,
+                    last_message_at = ?
+                WHERE telegram_id = ?
+            """, (points_to_add, now.isoformat(), telegram_id))
+            conn.commit()
+            return True
+
+    def set_premium(self, telegram_id: int, plan: str, months: int = 1):
+        expiry = datetime.now() + timedelta(days=30 * months)
+        col_is = f"is_{plan}_premium"
+        col_expiry = f"{plan}_expiry"
+        with self._get_connection() as conn:
+            conn.execute(f"""
+                UPDATE users 
+                SET {col_is} = 1, {col_expiry} = ?
+                WHERE telegram_id = ?
+            """, (expiry.isoformat(), telegram_id))
+            conn.commit()
+            logger.info(f"User {telegram_id} upgraded to {plan} premium until {expiry}")
+
+    def get_leaderboard(self, limit: int = 10):
+        with self._get_connection() as conn:
+            conn.row_factory = sqlite3.Row
+            cursor = conn.execute("""
+                SELECT username, points, message_count 
+                FROM users 
+                ORDER BY points DESC 
+                LIMIT ?
+            """, (limit,))
+            return [dict(row) for row in cursor.fetchall()]