vk_chat_bot/database.py

130 lines
5 KiB
Python

import sqlite3
from typing import List
class BasicDatabase:
def __init__(self, filename: str):
self.conn = sqlite3.connect(filename)
self.conn.row_factory = sqlite3.Row
self.cursor = self.conn.cursor()
def get_chats(self):
self.cursor.execute("SELECT * FROM chats")
return self.cursor.fetchall()
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 chat_delete(self, chat_id: int):
self.cursor.execute("DELETE FROM chats WHERE id = ?", (chat_id,))
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 delete_user(self, chat_id: int, user_id: int):
self.cursor.execute("DELETE FROM users WHERE chat_id = ? AND user_id = ?", (chat_id, user_id))
self.conn.commit()
def get_top_messages_today(self, chat_id: int):
self.cursor.execute("""
SELECT user_id, messages_today AS value 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 AS value 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_days: int):
# noinspection SpellCheckingInspection
self.cursor.execute("""
SELECT user_id,
CASE
WHEN last_message = 0 THEN 'никогда'
ELSE (unixepoch() - last_message) / 86400
END AS value
FROM users
WHERE chat_id = ? AND value >= ?
ORDER BY value DESC
""", (chat_id, threshold_days))
return self.cursor.fetchall()
def get_top_warnings(self, chat_id: int):
self.cursor.execute("""
SELECT user_id, warnings AS value FROM users
WHERE chat_id = ? AND warnings > 0
ORDER BY warnings DESC
""", (chat_id,))
return self.cursor.fetchall()
def delete_users(self, chat_id: int):
self.cursor.execute("DELETE FROM users WHERE chat_id = ?", (chat_id,))
self.conn.commit()
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()
def create_chat_if_not_exists(self, chat_id: int):
chat = self.get_chat(chat_id)
if chat is None:
self.add_chat(chat_id)
chat = self.get_chat(chat_id)
return chat
def create_user_if_not_exists(self, chat_id: int, user_id: int):
user = self.get_user(chat_id, user_id)
if user is None:
self.add_user(chat_id, user_id)
user = self.get_user(chat_id, user_id)
return user