From 7dc6df018209965c0e3745951d00496a77eae9ae Mon Sep 17 00:00:00 2001
From: emmatveev <emmtvv@gmail.com>
Date: Mon, 29 Apr 2024 15:05:05 +0300
Subject: [PATCH] redis

---
 .deploy/deploy-dev.yaml  |  2 ++
 .deploy/deploy-prod.yaml |  2 ++
 bot.py                   | 31 +++++++++++++++++--------------
 settings.py              |  3 +++
 tools/redis.py           | 35 +++++++++++++++++++++++++++++++++++
 5 files changed, 59 insertions(+), 14 deletions(-)
 create mode 100644 tools/redis.py

diff --git a/.deploy/deploy-dev.yaml b/.deploy/deploy-dev.yaml
index e1cd7fd..4899bb3 100644
--- a/.deploy/deploy-dev.yaml
+++ b/.deploy/deploy-dev.yaml
@@ -14,6 +14,8 @@ services:
       MINIO_SECRET_KEY: $MINIO_SECRET_KEY_DEV
       PLATFORM_SECURITY_TOKEN: $PLATFORM_SECURITY_TOKEN
       STAGE: "development"
+      REDIS_HOST: "redis.develop.sprinthub.ru"
+      REDIS_PASSWORD: $REDIS_PASSWORD_DEV
     networks:
       - net
     deploy:
diff --git a/.deploy/deploy-prod.yaml b/.deploy/deploy-prod.yaml
index a5d2538..e0a64e2 100644
--- a/.deploy/deploy-prod.yaml
+++ b/.deploy/deploy-prod.yaml
@@ -17,6 +17,8 @@ services:
       MINIO_SECRET_KEY: $MINIO_SECRET_KEY_PROD
       PLATFORM_SECURITY_TOKEN: $PLATFORM_SECURITY_TOKEN
       STAGE: "production"
+      REDIS_HOST: "redis.sprinthub.ru"
+      REDIS_PASSWORD: $REDIS_PASSWORD_PROD
     deploy:
       mode: replicated
       restart_policy:
diff --git a/bot.py b/bot.py
index f9a20d7..3b81e2c 100644
--- a/bot.py
+++ b/bot.py
@@ -8,6 +8,7 @@ from telebot.types import Message, ReplyKeyboardRemove
 from tools.minio import minio_client as minio
 from tools.mongo import mongo
 from tools.sprint_platform import platform
+from tools.redis import redis_client as redis
 
 bot = telebot.TeleBot(os.getenv("TELEGRAM_TOKEN"))
 
@@ -41,13 +42,14 @@ class Core:
             self.set_state('pause')
             if self.state == 'dialog':
                 current_dialog = mongo.get_current_dialog(self.chat_id)
-                mongo.finish_dialog(current_dialog['_id'])
-                if self.chat_id == current_dialog['chat_id_1']:
-                    another_chat_id = current_dialog['chat_id_2']
-                else:
-                    another_chat_id = current_dialog['chat_id_1']
-                self.send_message('🤖 Диалог окончен, жду тебя снова!')
-                self.start_new_dialog([another_chat_id])
+                with redis.lock('search'):
+                    mongo.finish_dialog(current_dialog['_id'])
+                    if self.chat_id == current_dialog['chat_id_1']:
+                        another_chat_id = current_dialog['chat_id_2']
+                    else:
+                        another_chat_id = current_dialog['chat_id_1']
+                    self.send_message('🤖 Диалог окончен, жду тебя снова!')
+                    self.start_new_dialog([another_chat_id])
                 return
             if self.state == 'pause':
                 self.send_message('🤖 Сейчас твой аккаунт не активен. Активируй его с помощью команды /start')
@@ -56,13 +58,14 @@ class Core:
                 self.send_message('🤖 Поиск собеседника окончен, жду тебя снова!')
                 return
         if self.message_text == '/next' or self.message_text == '/start':
-            if self.state == 'dialog':
-                dialog = mongo.get_current_dialog(self.chat_id)
-                self.start_new_dialog([dialog['chat_id_1'], dialog['chat_id_2']])
-                return
-            else:
-                self.start_new_dialog([self.chat_id])
-                return
+            with redis.lock('search'):
+                if self.state == 'dialog':
+                    dialog = mongo.get_current_dialog(self.chat_id)
+                    self.start_new_dialog([dialog['chat_id_1'], dialog['chat_id_2']])
+                    return
+                else:
+                    self.start_new_dialog([self.chat_id])
+                    return
 
     def handle_state_search(self):
         self.send_message('🤖 Поиски собеседника продолжаются')
diff --git a/settings.py b/settings.py
index 3ecd81b..31be338 100644
--- a/settings.py
+++ b/settings.py
@@ -8,3 +8,6 @@ MINIO_HOST = os.getenv("MINIO_HOST", "localhost") + ":9000"
 MINIO_ACCESS_KEY = os.getenv("MINIO_ACCESS_KEY", "serviceminioadmin")
 MINIO_SECRET_KEY = os.getenv("MINIO_SECRET_KEY", "minioadmin")
 MINIO_BUCKET_NAME = 'ruletka'
+
+REDIS_HOST = os.getenv("REDIS_HOST", "localhost")
+REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", None)
diff --git a/tools/redis.py b/tools/redis.py
new file mode 100644
index 0000000..6f3e6d1
--- /dev/null
+++ b/tools/redis.py
@@ -0,0 +1,35 @@
+import contextlib
+
+import redis
+
+import settings
+
+
+class RedisClient:
+
+    def __init__(self, host, password=None):
+        kwargs = {
+            "host": host,
+        }
+        if password:
+            kwargs['password'] = password
+        self.cli = redis.Redis(**kwargs)
+
+    def get(self, key):
+        with self.cli as cli:
+            return cli.get(f"ruletka_{key}")
+
+    def set(self, key, value):
+        with self.cli as cli:
+            cli.set(f"ruletka_{key}", value)
+
+    @contextlib.contextmanager
+    def lock(self, key):
+        with self.cli.lock(f"ruletka_{key}"):
+            yield
+
+
+redis_client = RedisClient(
+    settings.REDIS_HOST,
+    settings.REDIS_PASSWORD
+)