ruz-bot/helpers/mongo.py
2022-10-21 18:17:41 +03:00

67 lines
1.8 KiB
Python

import datetime
import zoneinfo
from functools import cached_property
import pymongo
import settings
from helpers import now
from helpers.models import UserSchema, User
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.lessons_collection.create_index([
("discipline", 1),
("auditorium", 1),
("begin", 1),
("hse_user_id", 1)
])
self.lessons_collection.create_index([
("hse_user_id", 1),
("begin", 1)
])
self.lessons_collection.create_index([
("hse_user_id", 1),
("begin", 1),
("notified", 1)
])
def get_user(self, user_id: int) -> User:
return UserSchema().loads(self.users_collection.find_one({"id": user_id}))
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_today_lessons(self, user: User):
today = now()
tomorrow = today + datetime.timedelta(days=1)
tomorrow = datetime.datetime(year=tomorrow.year, month=tomorrow.month, day=tomorrow.day)
lessons = []
for lesson in self.lessons_collection.find({
"hse_user_id": user.hse_id,
"begin": {"$gte": today, "$lte": tomorrow}}
):
lessons.append(lesson)
return lessons
mongo = Mongo()