grpc

utils/mongo.py

@@ -0,0 +1,46 @@
+import pymongo
+import os
+
+MONGO_USER = os.getenv("MONGO_USER", "mongo")
+MONGO_PASSWORD = os.getenv("MONGO_PASSWORD", "password")
+MONGO_HOST = os.getenv("MONGO_HOST", "localhost")
+
+
+class Mongo:
+    def __init__(self):
+        url = f"mongodb://{MONGO_USER}:{MONGO_PASSWORD}@{MONGO_HOST}:27017/"
+        self.client = pymongo.MongoClient(url)
+        self.database = self.client.get_database("pizda-bot")
+        self.chats_collection.create_index([
+            ("chat_id", 1)
+        ])
+        self.counter_collection.create_index([
+            ("chat_id", 1),
+            ("username", 1),
+        ])
+
+    def __getitem__(self, item):
+        return self.database.get_collection(item)
+
+    @property
+    def chats_collection(self):
+        return self["chats"]
+
+    @property
+    def counter_collection(self):
+        return self["counter"]
+
+    def inc(self, username, chat_id):
+        if self.counter_collection.find_one({"chat_id": chat_id, "username": username}):
+            self.counter_collection.update_one({"chat_id": chat_id, "username": username}, {"$inc": {"count": 1}})
+        else:
+            self.counter_collection.insert_one({"chat_id": chat_id, "username": username, "count": 1})
+
+    def inc_points(self, username, chat_id):
+        if self.counter_collection.find_one({"chat_id": chat_id, "username": username}):
+            self.counter_collection.update_one({"chat_id": chat_id, "username": username}, {"$inc": {"points": 1}})
+        else:
+            self.counter_collection.insert_one({"chat_id": chat_id, "username": username, "points": 1})
+
+
+mongo = Mongo()