117 lines
4.3 KiB
Python
117 lines
4.3 KiB
Python
import sqlite3
|
|
from typing import List
|
|
|
|
|
|
class Database:
|
|
def __init__(self):
|
|
self.conn = sqlite3.connect('bot.db')
|
|
self.conn.row_factory = sqlite3.Row
|
|
self.cursor = self.conn.cursor()
|
|
|
|
self.cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS chats (
|
|
"id" INTEGER,
|
|
"active" INTEGER NOT NULL DEFAULT 0,
|
|
PRIMARY KEY("id"))
|
|
""")
|
|
|
|
self.cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
"chat_id" INTEGER,
|
|
"user_id" INTEGER,
|
|
"last_message" INTEGER NOT NULL DEFAULT 0,
|
|
"messages_today" INTEGER NOT NULL DEFAULT 0,
|
|
"messages_month" INTEGER NOT NULL DEFAULT 0,
|
|
"warnings" INTEGER NOT NULL DEFAULT 0,
|
|
PRIMARY KEY("chat_id","user_id"))
|
|
""")
|
|
|
|
self.conn.commit()
|
|
|
|
def get_chat(self, chat_id: int):
|
|
self.cursor.execute("SELECT * FROM chats WHERE id = ?", (chat_id,))
|
|
return self.cursor.fetchone()
|
|
|
|
def add_chat(self, chat_id: int):
|
|
self.cursor.execute("INSERT INTO chats (id) VALUES (?)", (chat_id,))
|
|
self.conn.commit()
|
|
|
|
def chat_update(self, chat_id: int, **kwargs):
|
|
self.cursor.execute("UPDATE chats SET " + ", ".join(f + " = :" + f for f in kwargs) +
|
|
f" WHERE id = {chat_id}", kwargs)
|
|
self.conn.commit()
|
|
|
|
def get_user(self, chat_id: int, user_id: int):
|
|
self.cursor.execute("SELECT * FROM users WHERE chat_id = ? AND user_id = ?", (chat_id, user_id))
|
|
return self.cursor.fetchone()
|
|
|
|
def get_users(self, chat_id: int):
|
|
self.cursor.execute("SELECT * FROM users WHERE chat_id = ?", (chat_id,))
|
|
return self.cursor.fetchall()
|
|
|
|
def add_user(self, chat_id: int, user_id: int):
|
|
self.cursor.execute("INSERT INTO users (chat_id, user_id) VALUES (?, ?)", (chat_id, user_id))
|
|
self.conn.commit()
|
|
|
|
def user_set_last_message(self, chat_id: int, user_id: int, last_message: int):
|
|
self.user_update(chat_id, user_id, last_message=last_message)
|
|
|
|
def user_increment_messages(self, chat_id: int, user_id: int):
|
|
self.user_increment(chat_id, user_id, ['messages_today', 'messages_month'])
|
|
self.conn.commit()
|
|
|
|
def user_increment_warnings(self, chat_id: int, user_id: int):
|
|
self.user_increment(chat_id, user_id, ['warnings'])
|
|
|
|
def user_increment(self, chat_id: int, user_id: int, fields: List[str]):
|
|
self.cursor.execute("UPDATE users SET " + ", ".join(f + " = " + f + " + 1" for f in fields) +
|
|
" WHERE chat_id = ? AND user_id = ?", (chat_id, user_id))
|
|
self.conn.commit()
|
|
|
|
def user_update(self, chat_id: int, user_id: int, **kwargs):
|
|
self.cursor.execute("UPDATE users SET " + ", ".join(f + " = :" + f for f in kwargs) +
|
|
f" WHERE chat_id = {chat_id} AND user_id = {user_id}", kwargs)
|
|
self.conn.commit()
|
|
|
|
def get_top_messages_today(self, chat_id: int):
|
|
self.cursor.execute("""
|
|
SELECT user_id, messages_today FROM users
|
|
WHERE chat_id = ? AND messages_today > 0
|
|
ORDER BY messages_today DESC
|
|
""", (chat_id,))
|
|
return self.cursor.fetchall()
|
|
|
|
def get_top_messages_month(self, chat_id: int):
|
|
self.cursor.execute("""
|
|
SELECT user_id, messages_month FROM users
|
|
WHERE chat_id = ? AND messages_month > 0
|
|
ORDER BY messages_month DESC
|
|
""", (chat_id,))
|
|
return self.cursor.fetchall()
|
|
|
|
def get_top_silent(self, chat_id: int, threshold_time: int):
|
|
self.cursor.execute("""
|
|
SELECT user_id, last_message FROM users
|
|
WHERE chat_id = ? AND last_message < ?
|
|
ORDER BY last_message ASC
|
|
""", (chat_id, threshold_time))
|
|
return self.cursor.fetchall()
|
|
|
|
def get_top_warnings(self, chat_id: int):
|
|
self.cursor.execute("""
|
|
SELECT user_id, warnings FROM users
|
|
WHERE chat_id = ? AND warnings > 0
|
|
ORDER BY warnings DESC
|
|
""", (chat_id,))
|
|
return self.cursor.fetchall()
|
|
|
|
def reset_messages_today(self):
|
|
self.cursor.execute("UPDATE users SET messages_today = 0")
|
|
self.conn.commit()
|
|
|
|
def reset_messages_month(self):
|
|
self.cursor.execute("UPDATE users SET messages_month = 0")
|
|
self.conn.commit()
|
|
|
|
|
|
DB = Database()
|