73 lines
2.1 KiB
Python
73 lines
2.1 KiB
Python
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", 1),
|
|
('finished_at', 1)
|
|
])
|
|
self.dialogs_collection.create_index([
|
|
("chat_id_2", 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()
|