2569718930@qq.com
224 lines
8.7 KiB
Python
224 lines
8.7 KiB
Python
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: Optional[str] = None):
|
|
self.db_path = self._resolve_db_path(db_path)
|
|
self._init_db()
|
|
|
|
def _resolve_db_path(self, db_path: Optional[str]) -> str:
|
|
raw = (db_path or os.getenv("POLYWEATHER_DB_PATH") or "").strip()
|
|
if not raw:
|
|
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
return os.path.join(project_root, "data", "polyweather.db")
|
|
return raw
|
|
|
|
def _get_connection(self):
|
|
return sqlite3.connect(self.db_path)
|
|
|
|
def _init_db(self):
|
|
"""Create tables if they don't exist."""
|
|
# Ensure directory exists
|
|
db_dir = os.path.dirname(self.db_path)
|
|
if db_dir:
|
|
os.makedirs(db_dir, 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,
|
|
daily_points INTEGER DEFAULT 0,
|
|
daily_points_date TEXT,
|
|
message_count INTEGER DEFAULT 0,
|
|
last_message_at TIMESTAMP,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
self._ensure_column(conn, "users", "daily_points", "INTEGER DEFAULT 0")
|
|
self._ensure_column(conn, "users", "daily_points_date", "TEXT")
|
|
self._ensure_column(conn, "users", "supabase_user_id", "TEXT")
|
|
self._ensure_column(conn, "users", "supabase_email", "TEXT")
|
|
conn.commit()
|
|
logger.info(f"Database initialized successfully path={self.db_path}")
|
|
|
|
def _ensure_column(self, conn: sqlite3.Connection, table: str, column: str, ddl: str) -> None:
|
|
existing = {
|
|
row[1]
|
|
for row in conn.execute(f"PRAGMA table_info({table})").fetchall()
|
|
}
|
|
if column not in existing:
|
|
conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {ddl}")
|
|
|
|
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 bind_supabase_identity(
|
|
self,
|
|
telegram_id: int,
|
|
supabase_user_id: str,
|
|
supabase_email: str = "",
|
|
) -> None:
|
|
with self._get_connection() as conn:
|
|
conn.execute(
|
|
"""
|
|
UPDATE users
|
|
SET supabase_user_id = ?, supabase_email = ?
|
|
WHERE telegram_id = ?
|
|
""",
|
|
(supabase_user_id.strip(), supabase_email.strip(), telegram_id),
|
|
)
|
|
conn.commit()
|
|
|
|
def add_message_activity(
|
|
self,
|
|
telegram_id: int,
|
|
text: str,
|
|
points_to_add: int = 1,
|
|
cooldown_sec: int = 30,
|
|
daily_cap: int = 20,
|
|
min_text_length: int = 4,
|
|
) -> Dict[str, Any]:
|
|
"""Award points for valid group activity with cooldown and daily cap."""
|
|
now = datetime.now()
|
|
normalized = "".join((text or "").split())
|
|
if len(normalized) < min_text_length:
|
|
return {"awarded": False, "reason": "too_short"}
|
|
|
|
today_str = now.strftime("%Y-%m-%d")
|
|
with self._get_connection() as conn:
|
|
conn.row_factory = sqlite3.Row
|
|
cursor = conn.execute(
|
|
"""
|
|
SELECT points, daily_points, daily_points_date, last_message_at
|
|
FROM users WHERE telegram_id = ?
|
|
""",
|
|
(telegram_id,),
|
|
)
|
|
row = cursor.fetchone()
|
|
if not row:
|
|
return {"awarded": False, "reason": "user_missing"}
|
|
|
|
last_message_at = row["last_message_at"]
|
|
if last_message_at:
|
|
last_at = datetime.fromisoformat(last_message_at)
|
|
if (now - last_at).total_seconds() < cooldown_sec:
|
|
return {"awarded": False, "reason": "cooldown"}
|
|
|
|
daily_points = int(row["daily_points"] or 0)
|
|
daily_points_date = row["daily_points_date"] or ""
|
|
if daily_points_date != today_str:
|
|
daily_points = 0
|
|
|
|
if daily_points >= daily_cap:
|
|
conn.execute(
|
|
"""
|
|
UPDATE users
|
|
SET last_message_at = ?, daily_points = ?, daily_points_date = ?
|
|
WHERE telegram_id = ?
|
|
""",
|
|
(now.isoformat(), daily_points, today_str, telegram_id),
|
|
)
|
|
conn.commit()
|
|
return {"awarded": False, "reason": "daily_cap", "daily_points": daily_points}
|
|
|
|
conn.execute("""
|
|
UPDATE users
|
|
SET message_count = message_count + 1,
|
|
points = points + ?,
|
|
daily_points = ?,
|
|
daily_points_date = ?,
|
|
last_message_at = ?
|
|
WHERE telegram_id = ?
|
|
""", (points_to_add, daily_points + points_to_add, today_str, now.isoformat(), telegram_id))
|
|
conn.commit()
|
|
return {
|
|
"awarded": True,
|
|
"reason": "ok",
|
|
"daily_points": daily_points + points_to_add,
|
|
}
|
|
|
|
def spend_points(self, telegram_id: int, amount: int) -> Dict[str, Any]:
|
|
if amount <= 0:
|
|
user = self.get_user(telegram_id)
|
|
return {"ok": True, "balance": int((user or {}).get("points") or 0)}
|
|
|
|
with self._get_connection() as conn:
|
|
conn.row_factory = sqlite3.Row
|
|
row = conn.execute(
|
|
"SELECT points FROM users WHERE telegram_id = ?",
|
|
(telegram_id,),
|
|
).fetchone()
|
|
if not row:
|
|
return {"ok": False, "reason": "user_missing", "balance": 0, "required": amount}
|
|
|
|
balance = int(row["points"] or 0)
|
|
if balance < amount:
|
|
return {"ok": False, "reason": "insufficient_points", "balance": balance, "required": amount}
|
|
|
|
new_balance = balance - amount
|
|
conn.execute(
|
|
"UPDATE users SET points = ? WHERE telegram_id = ?",
|
|
(new_balance, telegram_id),
|
|
)
|
|
conn.commit()
|
|
return {"ok": True, "balance": new_balance, "spent": amount}
|
|
|
|
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()]
|