File size: 5,375 Bytes
9d382e5 b5e074c 9d382e5 b5e074c 9d382e5 b5e074c 9d382e5 8fd932e 9d382e5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | import aiosqlite
import os
from config import DATA_DIR
DB_PATH = os.path.join(DATA_DIR, "bot.db")
async def init_db():
os.makedirs(DATA_DIR, exist_ok=True)
async with aiosqlite.connect(DB_PATH) as db:
db.row_factory = aiosqlite.Row
await db.executescript("""
CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY,
bale_id TEXT UNIQUE NOT NULL,
username TEXT,
first_name TEXT,
is_admin INTEGER DEFAULT 0,
chat_mode TEXT DEFAULT 'none',
created_at TEXT DEFAULT (datetime('now'))
);
CREATE TABLE IF NOT EXISTS conversations (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER REFERENCES users(id),
role TEXT NOT NULL,
content TEXT NOT NULL,
created_at TEXT DEFAULT (datetime('now'))
);
CREATE TABLE IF NOT EXISTS stats (
id INTEGER PRIMARY KEY AUTOINCREMENT,
action TEXT NOT NULL,
user_id INTEGER REFERENCES users(id),
details TEXT,
created_at TEXT DEFAULT (datetime('now'))
);
CREATE TABLE IF NOT EXISTS settings (
key TEXT PRIMARY KEY,
value TEXT NOT NULL
);
""")
await db.commit()
async def get_user(bale_id: str):
async with aiosqlite.connect(DB_PATH) as db:
db.row_factory = aiosqlite.Row
cursor = await db.execute("SELECT * FROM users WHERE bale_id = ?", (bale_id,))
return await cursor.fetchone()
async def create_user(bale_id: str, username: str = None, first_name: str = None):
async with aiosqlite.connect(DB_PATH) as db:
await db.execute(
"INSERT OR IGNORE INTO users (bale_id, username, first_name) VALUES (?, ?, ?)",
(bale_id, username, first_name),
)
await db.commit()
async def set_chat_mode(bale_id: str, mode: str):
async with aiosqlite.connect(DB_PATH) as db:
await db.execute("UPDATE users SET chat_mode = ? WHERE bale_id = ?", (mode, bale_id))
await db.commit()
async def get_chat_mode(bale_id: str):
user = await get_user(bale_id)
return user["chat_mode"] if user else "none"
async def add_conversation(bale_id: str, role: str, content: str):
async with aiosqlite.connect(DB_PATH) as db:
user = await get_user(bale_id)
if user:
await db.execute(
"INSERT INTO conversations (user_id, role, content) VALUES (?, ?, ?)",
(user["id"], role, content),
)
await db.commit()
async def get_conversation_history(bale_id: str, limit: int = 20):
async with aiosqlite.connect(DB_PATH) as db:
db.row_factory = aiosqlite.Row
user = await get_user(bale_id)
if not user:
return []
cursor = await db.execute(
"SELECT role, content FROM conversations WHERE user_id = ? ORDER BY created_at DESC LIMIT ?",
(user["id"], limit),
)
rows = await cursor.fetchall()
return list(reversed(rows))
async def clear_conversation(bale_id: str):
async with aiosqlite.connect(DB_PATH) as db:
user = await get_user(bale_id)
if user:
await db.execute("DELETE FROM conversations WHERE user_id = ?", (user["id"],))
await db.commit()
async def log_stat(action: str, bale_id: str, details: str = None):
async with aiosqlite.connect(DB_PATH) as db:
user = await get_user(bale_id)
if user:
await db.execute(
"INSERT INTO stats (action, user_id, details) VALUES (?, ?, ?)",
(action, user["id"], details),
)
await db.commit()
async def get_all_users():
async with aiosqlite.connect(DB_PATH) as db:
db.row_factory = aiosqlite.Row
cursor = await db.execute("SELECT * FROM users ORDER BY created_at DESC")
return await cursor.fetchall()
async def get_stats_summary():
async with aiosqlite.connect(DB_PATH) as db:
db.row_factory = aiosqlite.Row
total = await (await db.execute("SELECT COUNT(*) as c FROM users")).fetchone()
cursor = await db.execute("SELECT action, COUNT(*) as count FROM stats GROUP BY action")
stats_rows = await cursor.fetchall()
return total, stats_rows
async def get_setting(key: str, default: str = None):
async with aiosqlite.connect(DB_PATH) as db:
cursor = await db.execute("SELECT value FROM settings WHERE key = ?", (key,))
row = await cursor.fetchone()
return row[0] if row else default
async def set_setting(key: str, value: str):
async with aiosqlite.connect(DB_PATH) as db:
await db.execute(
"INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)", (key, value)
)
await db.commit()
async def set_admin(bale_id: str, is_admin: bool = True):
async with aiosqlite.connect(DB_PATH) as db:
await db.execute("UPDATE users SET is_admin = ? WHERE bale_id = ?", (int(is_admin), bale_id))
await db.commit()
async def is_admin(bale_id: str):
user = await get_user(bale_id)
return user and user["is_admin"] == 1
|