ruz-bot/helpers/mongo.py
2023-08-13 19:48:28 +00:00

87 lines
2.6 KiB
Python

import datetime
from functools import cached_property
import pymongo
import settings
from helpers import now
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("ruz-bot")
self.users_collection.create_index([
("chat_id", 1)
])
self.users_collection.create_index([
("notify_minutes", 1)
])
self.users_collection.create_index([
("next_daily_notify_time", 1)
])
self.lessons_collection.create_index([
("discipline", 1),
("auditorium", 1),
("begin", 1),
("user_email", 1),
("link", 1)
])
self.lessons_collection.create_index([
("user_email", 1),
("begin", 1)
])
self.lessons_collection.create_index([
("user_email", 1),
("begin", 1),
("notified", 1)
])
def __getitem__(self, item):
return self.database.get_collection(item)
@cached_property
def users_collection(self):
return self["users"]
@cached_property
def lessons_collection(self):
return self["lessons"]
def _get_lessons_on_date(self, user: dict, date: datetime.datetime):
date = datetime.datetime(year=date.year, month=date.month, day=date.day)
next_date = date + datetime.timedelta(days=1)
lessons = []
for lesson in self.lessons_collection.find({
"user_email": user['email'],
"begin": {"$gte": date, "$lte": next_date}}
):
lessons.append(lesson)
lessons.sort(key=lambda les: les["begin"])
return lessons
def get_today_lessons(self, user: dict):
return self._get_lessons_on_date(user, now(user))
def get_tomorrow_lessons(self, user: dict):
return self._get_lessons_on_date(user, now(user) + datetime.timedelta(days=1))
def get_week_lessons(self, user: dict):
date = now(user)
date = datetime.datetime(year=date.year, month=date.month, day=date.day)
weekday = date.weekday()
next_date = date + datetime.timedelta(days=(6 - weekday))
lessons = []
for lesson in self.lessons_collection.find({
"user_email": user['email'],
"begin": {"$gte": date, "$lte": next_date}}
):
lessons.append(lesson)
lessons.sort(key=lambda les: les["begin"])
return lessons
mongo = Mongo()