FEATURE 9 and 12: cache and refactoring

- из админ_командс убрана работа с БД и шелве
- добавлен кеш, он загружается при старте сервиса и обновляется при изменениях, все запросы (антиспам, цитаты) идут уже через него, в БД лишний раз не ходим. В дальнейшем работаем только с ним везде через функции-хелперы
- созданы хелперы для создания/удаления цитат и чатов
- добавлено выражение в спам-лист (вчера пропустило)
- добавлены тесты на спам(выше) и админ_коммандс

src/admin_commands.py

@@ -1,15 +1,14 @@
 import logging
-import shelve
 from datetime import datetime as dt
 
 from telebot import types, logger
 
-from .constants import ADMIN_ID, LOGGING_LEVEL_DEBUG, LOGGING_LEVEL_INFO, DATA
+from .constants import ADMIN_ID, LOGGING_LEVEL_DEBUG, LOGGING_LEVEL_INFO
 from .filters import is_white_id
 from .helpers import me, my_ip, is_admin
-from .models import Chat, Quote, session
+from .models import is_chat_exists, add_chat, delete_chat, update_chat, chat_by_id, all_chat_quotes
 from .reminder import remind, print_get_jobs, create_report_text
-from .report import reset_report_stats
+from .report import reset_report_stats, current_data
 from .robot import bot
 
 
@@ -34,7 +33,7 @@ async def remind_manually(message):
         except ValueError as ve:
             await bot.send_message(message.chat.id, f"Не удалось разобрать дату!\n{ve}")
         else:
-            await remind(message.chat.id, today)
+            remind(message.chat.id, today)
     else:
         await bot.send_message(
             message.chat.id,
@@ -54,15 +53,14 @@ async def list_jobs(_message):
 @bot.message_handler(func=is_admin, commands=["stats"])
 async def send_stats(message: types.Message):
     """Show group statistics."""
-    chat_id = message.chat.id if len(message.text.split()) == 1 else message.text.split()[-1]
+    chat_id = message.chat.id if len(splitted := message.text.split()) == 1 else splitted[-1]
     report_text = create_report_text(chat_id)
     if report_text:
         await bot.send_message(message.chat.id, report_text)
     else:
-        with shelve.open(f"{DATA}{chat_id}") as shelve_db:
-            shelve_messages = {v['User'].first_name: v['Count'] for k, v in shelve_db['Messages'].items()}
-            logging.info(f"[STATS {chat_id}] Невозможно отправить отчёт, недостаточно данных. Счётчик сообщений: {shelve_messages}.\n"
-                         f"Баны: {shelve_db['Banned']=}")
+        shelve_messages, banned = current_data(chat_id)
+        logging.info(f"[STATS {chat_id}] Невозможно отправить отчёт, недостаточно данных. "
+                     f"Счётчик сообщений: {shelve_messages}.\nБаны: {banned}")
 
 
 @bot.message_handler(func=is_admin, commands=["reset_stats"], chat_types=["supergroup", "group"])
@@ -76,12 +74,11 @@ async def send_stats(message: types.Message):
 async def enable_stan(message: types.Message):
     """Add group to database."""
     logging.info(f"[{message.chat.title}] [{message.from_user.id}] {message.from_user.username}: {message.text}")
-    if not session.query(Chat.id).filter(Chat.chat_id == message.chat.id).first():
-        session.add(Chat(chat_id=message.chat.id, title=message.chat.title, antispam=1, report=0, reminder=1))
-        session.commit()
+    if not is_chat_exists(message.chat.id):
+        add_chat(message.chat.id, message.chat.title)
         await bot.send_message(
             message.chat.id,
-            f"""Группа "{message.chat.title} добавлена в БД.
+            f"""Группа "{message.chat.title}" добавлена в БД.
 /get_group_info - узнать текущие настройки""",
         )
     else:
@@ -91,11 +88,9 @@ async def enable_stan(message: types.Message):
 @bot.message_handler(func=is_admin, commands=["disable_stan"], chat_types=["supergroup", "group"])
 async def disable_stan(message: types.Message):
     logging.info(f"[{message.chat.title}] [{message.from_user.id}] {message.from_user.username}: {message.text}")
-    if session.query(Chat.id).filter(Chat.chat_id == message.chat.id).first():
-        chat = session.query(Chat).filter_by(chat_id=message.chat.id).first()
-        session.delete(chat)
-        session.commit()
-        await bot.send_message(message.chat.id, f"Группа {message.chat.title} и все связанные с ней цитаты удалёны!")
+    if is_chat_exists(message.chat.id):
+        title = delete_chat(message.chat.id)
+        await bot.send_message(message.chat.id, f"Группа {title} и все связанные с ней цитаты удалены!")
     else:
         await bot.send_message(message.chat.id, "Отказ. Этой группы нет в БД.")
 
@@ -115,36 +110,32 @@ async def set_antispam_report_reminder(message: types.Message):
         except ValueError:
             logging.info(f"[ERROR] Неверные аргументы {message.text}")
         else:
-            session.query(Chat).filter_by(chat_id=message.chat.id).update(
-                {"antispam": antispam, "report": rep, "reminder": rem}
-            )
-            session.commit()
-            await bot.send_message(message.chat.id, f"""Настройки обновлены. Проверить: /get_group_info""")
+            update_chat(message.chat.id, antispam, rep, rem)
+            await bot.send_message(message.chat.id, "Настройки обновлены. Проверить: /get_group_info")
 
 
 @bot.message_handler(func=is_white_id, commands=["get_quotes"], chat_types=["supergroup", "group"])
 async def get_quotes(message: types.Message):
     logging.info(f"[{message.chat.id}] {message.from_user.first_name} {message.text}.")
-    quotes = session.query(Quote.text).filter_by(chat_id=message.chat.id).all()
+    quotes = all_chat_quotes(message.chat.id)
     if quotes:
-        length = len(session.query(Quote).filter(Quote.chat_id == message.chat.id).all())
-        text = f"Количество цитат: {length}\n\Последние добавленные\n\n"
-        text += "\n".join(f"· {quote[0]}" for quote in quotes[-10:])
+        text = f"Количество цитат: {len(quotes)}\n\Последние добавленные\n\n"
+        text += "\n".join(f"· {quote}" for quote in quotes[-10:])
         await bot.send_message(message.chat.id, f"{text}")
     else:
-        await bot.send_message(message.chat.id, f"Цитаты отсутствуют. Подробнее: /get_group_info")
+        await bot.send_message(message.chat.id, "Цитаты отсутствуют. Подробнее: /get_group_info")
 
 
 @bot.message_handler(func=is_white_id, commands=["get_group_info"], chat_types=["supergroup", "group"])
 async def get_group_info(message: types.Message):
-    group = session.query(Chat).filter_by(chat_id=message.chat.id).first()
+    group = chat_by_id(message.chat.id)
     if group:
         await bot.send_message(
             message.chat.id,
             f"""Группа: {group.title}
 ID группы: {group.chat_id} 
 
-Количество цитат:  {len(session.query(Quote).filter(Quote.chat_id == message.chat.id).all())}
+Количество цитат:  {len(group.quotes)}
 Последние добавленные: /get_quotes
 
 Текущие настройки:
@@ -153,7 +144,7 @@ ID группы: {group.chat_id}
   Праздники: {group.reminder}""",
         )
     else:
-        await bot.send_message(message.chat.id, f"Группа не включена. Включить: /enable_stan")
+        await bot.send_message(message.chat.id, "Группа не включена. Включить: /enable_stan")
 
 
 @bot.message_handler(func=is_white_id, commands=["set_logging_level"], chat_types=["supergroup", "group"])

src/constants.py

@@ -59,7 +59,7 @@ SPAM = {"me.sv/", "tg.sv/", "goo.by/", "go.sv/", "intim.video/", "uclck.ru/", '
         'P2P', 'Идет набор', 'в лс', 'нaбop', 'комaндy', 'kpuпт', 'бupж', '💰', 'пишитe', '💸',
         '$ в неделю', 'обучаю с нуля', 'доходы для всех', 'по трейдингу', 'фальшивые рубли', '‼️',
         'зароботку', 'Набuраю', 'цuфр', 'Обучу вас', 'прuб', 'остались курсы', 'Crypt', 'kр', 'рu',
-        'в лc'}
+        'в лc', "в личные сообщения"}
 
 NON_GRATA = {"дудар", "хауди", "dudar"}
 BAN_WORDS = {"GREEN"}

src/models.py

@@ -1,4 +1,7 @@
+import logging
+import os
 from pathlib import Path
+from typing import Type
 
 from sqlalchemy import Column, String, Integer, Boolean, ForeignKey
 from sqlalchemy import create_engine
@@ -51,27 +54,78 @@ class Chat(Base):
         return f'Chat({self.chat_id=}, {self.title=})'
 
 
+CACHE: dict[int, Type[Chat]] = {}
+
+
+def _fill_cache():
+    if os.environ.get("testing", "FALSE") == "FALSE":
+        for chat in session.query(Chat).all():
+            CACHE[chat.chat_id] = chat
+
+
 def is_antispam_enabled(chat_id: int) -> bool:
-    try:
-        return session.query(Chat.antispam).filter_by(chat_id=chat_id).first()[0]
-    except TypeError:
+    if chat_id not in CACHE:
         return False
+    return CACHE[chat_id].antispam
 
 
 def is_quote_in_chat(text: str, chat_id: int) -> bool:
-    result = session.query(Quote.text).filter_by(text=text, chat_id=chat_id).first()
-    return bool(result)
+    return text in all_chat_quotes(chat_id)
 
 
 def delete_quote_in_chat(text: str, chat_id: int):
     session.execute(delete(Quote).filter_by(text=text, chat_id=chat_id))
     session.commit()
+    _reload_chat(chat_id)
 
 
 def all_chat_quotes(chat_id: int) -> list:
-    return session.query(Quote.text).filter(Quote.chat_id == chat_id).all()
+    if chat_id not in CACHE:
+        return []
+    return [quote.text for quote in CACHE[chat_id].quotes]
+
+
+def is_chat_exists(chat_id: int) -> bool:
+    return chat_id in CACHE
 
 
 def add_quote(chat_id: int, text: str):
     session.add(Quote(chat_id=chat_id, text=text))
     session.commit()
+    _reload_chat(chat_id)
+
+
+def chat_by_id(chat_id: int) -> None | Chat:
+    return CACHE.get(chat_id)
+
+
+def add_chat(chat_id: int, chat_title: str):
+    session.add(Chat(chat_id=chat_id, title=chat_title, antispam=1, report=0, reminder=1))
+    session.commit()
+    _reload_chat(chat_id)
+
+
+def update_chat(chat_id: int, antispam: int = 1, rep: int = 0, rem: int = 1):
+    session.query(Chat).filter_by(chat_id=chat_id).update({"antispam": antispam, "report": rep, "reminder": rem})
+    session.commit()
+    _reload_chat(chat_id)
+
+
+def delete_chat(chat_id: int) -> str:
+    """
+    Удаляем чат из БД вместе с его цитатами, возвращаем его название для логирования
+    """
+    chat = session.query(Chat).filter_by(chat_id=chat_id).first()
+    session.delete(chat)
+    session.commit()
+    del CACHE[chat_id]
+    return chat.title
+
+
+def _reload_chat(chat_id: int):
+    chat = session.query(Chat).filter(Chat.chat_id == chat_id).first()
+    CACHE[chat_id] = chat
+
+
+_fill_cache()
+logging.warning("[START] Load data from db to cache")

src/report.py

@@ -18,6 +18,13 @@ def increment(chat_id: int, banned: bool = True):
         shelve_db[category] += 1
 
 
+def current_data(chat_id: int) -> tuple:
+    with shelve.open(f"{DATA}{chat_id}") as shelve_db:
+        shelve_messages = {v['User'].first_name: v['Count'] for k, v in shelve_db['Messages'].items()}
+        banned = shelve_db['Banned']
+    return shelve_messages, banned
+
+
 def create_report_text(chat_id):
     report = ""
     top_users = ""

src/stan.py

@@ -17,7 +17,7 @@ TYPING_TIMEOUT = 0.13 / 4  # Reading time is quarter of the same text writing ti
 def speak(chance_of: int, chat_id: int) -> None | str:
     number = random.randint(0, chance_of)
     if number == 0:
-        return random.choice([i[0] for i in all_chat_quotes(chat_id)])
+        return random.choice(all_chat_quotes(chat_id))
     return None
 
 
@@ -40,7 +40,7 @@ async def act(message: types.Message):
 async def add_stan_quote(message: types.Message):
     if message.reply_to_message and message.reply_to_message.text:
         quote = message.reply_to_message.text
-        if quote not in {i[0] for i in all_chat_quotes(message.chat.id)}:
+        if not is_quote_in_chat(quote, message.chat.id):
             add_quote(message.chat.id, quote.replace("\n", " "))
             await bot.send_message(message.chat.id, "➕\n  └ " + quote.replace("\n", " "), parse_mode='Markdown')
             await bot.delete_message(message.chat.id, message.id)