add new features

tools/minio.py

@@ -0,0 +1,45 @@
+import io
+
+from minio import Minio
+from minio.error import MinioException
+
+import settings
+
+
+class Client:
+
+    def __init__(self, host: str, access_key: str, secret_key: str, bucket_name: str):
+        self.bucket_name = bucket_name
+        self.cli = Minio(
+            host,
+            access_key=access_key,
+            secret_key=secret_key,
+            secure=False
+        )
+        try:
+            self.cli.make_bucket(bucket_name)
+        except MinioException:
+            pass
+
+    def put_object(self, name: str, data: bytes):
+        self.cli.put_object(self.bucket_name, name, io.BytesIO(data), len(data))
+
+    def get_object(self, name: str) -> bytes:
+        try:
+            return self.cli.get_object(self.bucket_name, name).data
+        except MinioException:
+            return b""
+
+    def delete_object(self, name: str):
+        try:
+            self.cli.remove_object(self.bucket_name, name)
+        except MinioException:
+            pass
+
+
+minio_client = Client(
+    settings.MINIO_HOST,
+    settings.MINIO_ACCESS_KEY,
+    settings.MINIO_SECRET_KEY,
+    settings.MINIO_BUCKET_NAME
+)

tools/mongo.py

@@ -0,0 +1,68 @@
+import datetime
+import random
+
+import pymongo
+
+import settings
+
+
+class Mongo:
+    def __init__(self):
+        url = f"mongodb://{settings.MONGO_USER}:{settings.MONGO_PASSWORD}@{settings.MONGO_HOST}:27017/"
+        self.client = pymongo.MongoClient(url)
+        self.database = self.client.get_database("roulette-bot")
+        self.chats_collection.create_index([
+            ("chat_id", 1),
+            ('state', 1)
+        ])
+        self.dialogs_collection.create_index([
+            ("chat_id", 1),
+            ('finished_at', 1)
+        ])
+
+    def __getitem__(self, item):
+        return self.database.get_collection(item)
+
+    @property
+    def chats_collection(self):
+        return self["chats"]
+
+    @property
+    def dialogs_collection(self):
+        return self["dialogs"]
+
+    @property
+    def messages_collection(self):
+        return self["messages"]
+
+    def find_searching(self, except_ids):
+        chats = list(self.chats_collection.find({"state": 'search', 'chat_id': {'$nin': except_ids}}))
+        if not chats:
+            return None
+        return random.choice(chats)
+
+    def get_current_dialog(self, chat_id):
+        return self.dialogs_collection.find_one({'$or': [{'chat_id_1': chat_id}, {'chat_id_2': chat_id}], 'finished_at': None})
+
+    def create_message(self, message_type, text, dialog_id, sender):
+        return self.messages_collection.insert_one({
+            'message_type': message_type,
+            'dialog_id': dialog_id,
+            'text': text,
+            'sender': sender,
+            'sent_at': datetime.datetime.now(),
+        })
+
+    def create_dialog(self, chat_id_1, chat_id_2):
+        self.dialogs_collection.insert_one({
+            'chat_id_1': chat_id_1,
+            'chat_id_2': chat_id_2,
+            'started_at': datetime.datetime.now(),
+            'finished_at': None,
+        })
+
+    def finish_dialog(self, dialog_id):
+        self.dialogs_collection.update_one({'_id': dialog_id}, {'$set': {'finished_at': datetime.datetime.now()}})
+
+
+mongo = Mongo()