diff --git a/.env.example b/.env.example
index 00adfda..75680fc 100644
--- a/.env.example
+++ b/.env.example
@@ -1,4 +1,6 @@
TG_BOT_TOKEN=""
MONGODB_USERNAME=""
MONGODB_PASSWORD=""
-MONGODB_PORT=27017
\ No newline at end of file
+MONGODB_PORT=27017
+
+MONGODB_HOST="mongodb://login:password@ip:port" # If you don't use docker compose
\ No newline at end of file
diff --git a/main.py b/main.py
index f0920d9..5194639 100644
--- a/main.py
+++ b/main.py
@@ -1,10 +1,6 @@
import asyncio
-from src.bot import bot, dp
-
-
-async def main():
- await dp.start_polling(bot)
+from src.bot import main
if __name__ == "__main__":
asyncio.run(main())
diff --git a/random_prizes.py b/random_prizes.py
index e10f0b6..0013c99 100644
--- a/random_prizes.py
+++ b/random_prizes.py
@@ -1,24 +1,19 @@
-from dotenv import load_dotenv
-import os
import argparse
-import pymongo
+import os
import random
+import pymongo
+from dotenv import load_dotenv
+
load_dotenv()
-MONGODB_USERNAME = os.getenv("MONGODB_USERNAME")
-MONGODB_PASSWORD = os.getenv("MONGODB_PASSWORD")
-MONGODB_PORT = os.getenv("MONGODB_PORT")
-MONGODB_IP = os.getenv("MONGODB_IP")
+MONGODB_HOST = os.getenv("MONGODB_HOST")
parser = argparse.ArgumentParser()
parser.add_argument("count", type=int)
args = parser.parse_args()
-MONGO_HOST = (
- f"mongodb://{MONGODB_USERNAME}:{MONGODB_PASSWORD}@{MONGODB_IP}:{MONGODB_PORT}"
-)
-client = pymongo.MongoClient(MONGO_HOST)
+client = pymongo.MongoClient(MONGODB_HOST)
database = client.get_database("cu_graph_bot")
collection = database.get_collection("users")
users = ["@" + i["username"] for i in collection.find({"_links.4": {"$exists": True}})]
diff --git a/src/bot.py b/src/bot.py
index 03d5c9e..621d890 100644
--- a/src/bot.py
+++ b/src/bot.py
@@ -1,25 +1,26 @@
import os
-from typing import List
-from aiogram import Bot, Dispatcher, F, html, types
+from aiogram import Bot, Dispatcher, F, types
+from aiogram.client.default import DefaultBotProperties
from aiogram.enums import ParseMode
-from aiogram.filters import Command, CommandStart, callback_data
+from aiogram.filters import CommandStart
from aiogram.fsm.context import FSMContext
from aiogram.fsm.state import State, StatesGroup
+from aiogram.fsm.storage.pymongo import PyMongoStorage
from aiogram.types import (
- BotCommand,
CallbackQuery,
InlineKeyboardButton,
InlineKeyboardMarkup,
KeyboardButton,
ReplyKeyboardMarkup,
- user,
)
-from aiogram.utils.formatting import Bold, CustomEmoji, Text
+from aiogram.utils.formatting import as_list
from dotenv import load_dotenv
+from pymongo import AsyncMongoClient
+from . import callbacks, templates
from .models import Link, User, check_tg_username
-from .userdb import userdb
+from .userdb import UserDB
load_dotenv()
@@ -27,29 +28,17 @@
if TOKEN is None:
raise Exception("Couldn't find TG_BOT_TOKEN")
-bot = Bot(token=TOKEN)
-dp = Dispatcher()
+MONGODB_HOST = os.getenv("MONGODB_HOST")
+client = AsyncMongoClient(MONGODB_HOST)
+bot = Bot(token=TOKEN, default=DefaultBotProperties(parse_mode=ParseMode.HTML))
+dp = Dispatcher(storage=PyMongoStorage(client, db_name="cu_graph_bot"))
-class StartingCallback(callback_data.CallbackData, prefix="start"):
- pass
-
-class LinkCallback(callback_data.CallbackData, prefix="link"):
- username_to: str
- rating: int
-
-
-class SexCallback(callback_data.CallbackData, prefix="sex"):
- sex: str
-
-
-class CourseCallback(callback_data.CallbackData, prefix="course"):
- course: int
-
-
-class LivingCallback(callback_data.CallbackData, prefix="living"):
- living: str
+async def main():
+ global userdb, bot, dp
+ userdb = UserDB(client)
+ await dp.start_polling(bot)
rkb = ReplyKeyboardMarkup(
@@ -60,50 +49,6 @@ class LivingCallback(callback_data.CallbackData, prefix="living"):
]
)
-starting_message = """🧬 ДОБРО ПОЖАЛОВАТЬ В CAMPUS DNA!
-
-Мы создаём первую карту социальных связей нашего университета.
-Это исследование научной студии, и каждый участник получит:
-• Личный анализ социального типа
-• Место на интерактивной карте универа
-• Шанс выиграть КРУТЫЕ ПРИЗЫ 🎁
-
-🏆 УСЛОВИЯ УЧАСТИЯ В РОЗЫГРЫШЕ:
-
-1. ✅ Подписаться на канал @campusdna
-2. ✅ Отметь 5+ друзей и оцени вашу близость
-
-🎁 ПРИЗОВОЙ ФОНД:
-• 20 ПИЦЦ (1 пицца = 1 победитель)
-• ИГРУШКИ-МИНЬОНЫ
-• МЕРЧ ОТ ЦУ И ПАРТНЁРОВ
-
-📢 Следи за розыгрышами в канале: @campusdna
-
-🧭 ЧТО ДЕЛАТЬ ДАЛЬШЕ:
-
-Сначала тебе нужно ввести базовые сведения: пол, курс, общежитие.
-Затем я попрошу тебя ввести юзернеймы твоих друзей в Telegram
-и оценить вашу близость по шкале от 1 до 3.
-
-Чем больше друзей ты отметишь — тем точнее будет твой
-социальный портрет и тем ценнее твой вклад в исследование!
-
-Готов начать и узнать, кто ты в социальной сети университета?"""
-
-explaining_links_message = """⁉️ НА КАКИЕ ГРУППЫ МЫ ДЕЛИМ СВЯЗИ?
-🔴 1 — Друзья
-«С ними я провожу больше всего времени»
-Постоянное общение в универе и в мессенджерах. Видимся почти каждый день. Делимся личными новостями, поддерживаем друг друга.
-
-🟡 2 — Приятели
-«Всегда подойду спросить: "Как дела? Как жизнь?"»
-Видимся несколько раз в неделю. Общаемся и про учебу, и про жизнь, иногда затрагиваем личное (но не глубокое). Можем вместе пообедать или поиграть в пин-понг.
-
-🔵 3 — Знакомые
-«Мы здороваемся в коридоре»
-Видимся изредка, общение короткое и ситуативное. В основном на учебные/повседневные темы."""
-
class AddingUser(StatesGroup):
sex = State()
@@ -114,26 +59,25 @@ class AddingUser(StatesGroup):
@dp.message(CommandStart())
async def start_handler(message: types.Message, state: FSMContext):
await message.answer(
- starting_message,
+ templates.starting_message,
reply_markup=InlineKeyboardMarkup(
inline_keyboard=[
[
InlineKeyboardButton(
- text="Далее", callback_data=StartingCallback().pack()
+ text="Далее", callback_data=callbacks.StartingCallback().pack()
)
]
]
),
- parse_mode=ParseMode.HTML,
)
-@dp.callback_query(StartingCallback.filter())
+@dp.callback_query(callbacks.StartingCallback.filter())
async def next_handler(
- query: CallbackQuery, callback_data: StartingCallback, state: FSMContext
+ query: CallbackQuery, callback_data: callbacks.StartingCallback, state: FSMContext
):
if await userdb.get_user(query.from_user.username) is not None:
- await start_survey(query.message)
+ await explaining_links(query.message)
return
await state.set_state(AddingUser.sex)
await question_sex(query.message, state)
@@ -148,10 +92,12 @@ async def question_sex(message: types.Message, state: FSMContext):
inline_keyboard=[
[
InlineKeyboardButton(
- text="Мужской", callback_data=SexCallback(sex="male").pack()
+ text="Мужской",
+ callback_data=callbacks.SexCallback(sex="male").pack(),
),
InlineKeyboardButton(
- text="Женский", callback_data=SexCallback(sex="female").pack()
+ text="Женский",
+ callback_data=callbacks.SexCallback(sex="female").pack(),
),
]
]
@@ -159,9 +105,9 @@ async def question_sex(message: types.Message, state: FSMContext):
)
-@dp.callback_query(SexCallback.filter())
+@dp.callback_query(callbacks.SexCallback.filter())
async def process_sex(
- query: CallbackQuery, callback_data: SexCallback, state: FSMContext
+ query: CallbackQuery, callback_data: callbacks.SexCallback, state: FSMContext
):
await state.update_data(sex=callback_data.sex)
await question_course(query.message, state)
@@ -176,10 +122,12 @@ async def question_course(message: types.Message, state: FSMContext):
inline_keyboard=[
[
InlineKeyboardButton(
- text="1", callback_data=CourseCallback(course=1).pack()
+ text="1",
+ callback_data=callbacks.CourseCallback(course=1).pack(),
),
InlineKeyboardButton(
- text="2", callback_data=CourseCallback(course=2).pack()
+ text="2",
+ callback_data=callbacks.CourseCallback(course=2).pack(),
),
]
]
@@ -187,9 +135,9 @@ async def question_course(message: types.Message, state: FSMContext):
)
-@dp.callback_query(CourseCallback.filter())
+@dp.callback_query(callbacks.CourseCallback.filter())
async def process_course(
- query: CallbackQuery, callback_data: CourseCallback, state: FSMContext
+ query: CallbackQuery, callback_data: callbacks.CourseCallback, state: FSMContext
):
await state.update_data(course=callback_data.course)
await question_living(query.message, state)
@@ -205,25 +153,27 @@ async def question_living(message: types.Message, state: FSMContext):
[
InlineKeyboardButton(
text="В Облаке",
- callback_data=LivingCallback(living="Cloud").pack(),
+ callback_data=callbacks.LivingCallback(living="Cloud").pack(),
),
],
[
InlineKeyboardButton(
text="В Космосе",
- callback_data=LivingCallback(living="Cosmos").pack(),
+ callback_data=callbacks.LivingCallback(living="Cosmos").pack(),
),
],
[
InlineKeyboardButton(
text="В Байкале",
- callback_data=LivingCallback(living="Baikal").pack(),
+ callback_data=callbacks.LivingCallback(living="Baikal").pack(),
),
],
[
InlineKeyboardButton(
text="Не в общаге",
- callback_data=LivingCallback(living="Homeless").pack(),
+ callback_data=callbacks.LivingCallback(
+ living="Homeless"
+ ).pack(),
),
],
]
@@ -231,9 +181,9 @@ async def question_living(message: types.Message, state: FSMContext):
)
-@dp.callback_query(LivingCallback.filter())
+@dp.callback_query(callbacks.LivingCallback.filter())
async def process_living(
- query: CallbackQuery, callback_data: LivingCallback, state: FSMContext
+ query: CallbackQuery, callback_data: callbacks.LivingCallback, state: FSMContext
):
state_data = await state.get_data()
await userdb.add_user(
@@ -244,12 +194,28 @@ async def process_living(
living=callback_data.living,
)
)
- await start_survey(query.message)
+ await state.clear()
+ await explaining_links(query.message)
-async def start_survey(message: types.Message):
- await message.answer(explaining_links_message, parse_mode=ParseMode.HTML)
+async def explaining_links(message: types.Message):
await message.answer(
+ templates.explaining_links_message,
+ reply_markup=InlineKeyboardMarkup(
+ inline_keyboard=[
+ [
+ InlineKeyboardButton(
+ text="Далее", callback_data=callbacks.TypeInfoCallback().pack()
+ )
+ ]
+ ]
+ ),
+ )
+
+
+@dp.callback_query(callbacks.TypeInfoCallback.filter())
+async def start_survey(query: CallbackQuery, callback_data: callbacks.TypeInfoCallback):
+ await query.message.answer(
"Напиши юзернейм (@username) и я предложу тебе выбрать его категорию",
reply_markup=rkb,
)
@@ -267,6 +233,9 @@ def rating_to_text(rating: int) -> str:
@dp.message(F.text == "Мои контакты")
async def get_usS(message: types.Message):
+ await userdb.add_ids_to_user(
+ message.from_user.username, message.from_user.id, message.chat.id
+ )
links = await userdb.get_links(message.from_user.username)
if len(links) == 0:
await message.answer("Ты ещё не добавил связи!\nВведи юзернейм (@username)")
@@ -277,22 +246,11 @@ async def get_usS(message: types.Message):
await message.answer(all_users_and_rating, reply_markup=rkb)
-def make_type_str(type: str, profile: str, strong_sides: List[str], recomendation: str):
- return f"""🎯ТИП: «{type}»
-
-📊 Ваш профиль:
-{profile}
-
-💪 Ваши сильные стороны:
-• {"\n• ".join(strong_sides)}
-
-🌟 Рекомендация:
-{recomendation}
-"""
-
-
@dp.message(F.text == "Узнать тип личности")
async def get_summary(message: types.Message):
+ await userdb.add_ids_to_user(
+ message.from_user.username, message.from_user.id, message.chat.id
+ )
links = await userdb.get_links(message.from_user.username)
ratings = [i.rating for i in links]
p1 = ratings.count(1) / len(ratings)
@@ -304,7 +262,7 @@ async def get_summary(message: types.Message):
)
elif p3 >= 0.5:
await message.answer(
- make_type_str(
+ templates.make_type_str(
"Сердце компании",
"Вы создаете глубокие, осознанные отношения. Для вас важно не количество контактов, а их качество и надежность",
[
@@ -315,11 +273,10 @@ async def get_summary(message: types.Message):
],
'Попробуйте иногда быть "социальным мостом" — знакомить своих друзей из разных кругов. Ваша глубина общения может стать основой для новых интересных компаний',
),
- parse_mode=ParseMode.HTML,
)
elif p2 >= 0.4:
await message.answer(
- make_type_str(
+ templates.make_type_str(
"Социальный организатор",
"Вы — мастер поддерживать ровные, комфортные отношения. С вами легко и приятно общаться на повседневные темы",
[
@@ -330,11 +287,10 @@ async def get_summary(message: types.Message):
],
"Попробуйте выбрать 1-2 самых интересных вам приятеля и предложить им более тесное общение — совместный проект или регулярные встречи. Ваши легкие связи могут перерасти в нечто большее",
),
- parse_mode=ParseMode.HTML,
)
elif p3 >= 0.25 and p2 >= 0.25 and p1 >= 0.25:
await message.answer(
- make_type_str(
+ templates.make_type_str(
"Универсальный коннектор",
"Вы легко перемещаетесь между разными социальными слоями. От тактических знакомств до близкой дружбы — вы чувствуете себя комфортно на любом уровне",
[
@@ -345,11 +301,10 @@ async def get_summary(message: types.Message):
],
"Используйте свой дар соединять людей! Организуйте мини-встречи людей из разных ваших кругов — возможно, вы создадите новые интересные коллаборации",
),
- parse_mode=ParseMode.HTML,
)
elif p3 >= 0.35 and p1 >= 0.25:
await message.answer(
- make_type_str(
+ templates.make_type_str(
"Стратегический коммуникатор",
"Вы сочетаете глубокую эмоциональную привязанность с широким кругом полезных контактов. Это редкий и ценный навык!",
[
@@ -360,11 +315,10 @@ async def get_summary(message: types.Message):
],
'Подумайте, как ваши "знакомые" могут помочь вашим "друзьям" (и наоборот). Вы идеально positioned для создания синергии между разными частями вашей сети',
),
- parse_mode=ParseMode.HTML,
)
elif abs(p3 - p2) <= 0.3 and abs(p2 - p1) / len(ratings) <= 0.3:
await message.answer(
- make_type_str(
+ templates.make_type_str(
"Стабильный якорь",
"Вы выстраиваете гармоничную социальную экосистему, где каждому типу отношений находится свое место",
[
@@ -375,7 +329,6 @@ async def get_summary(message: types.Message):
],
'Ваша сила — в стабильности. Подумайте, не хотите ли вы немного "сдвинуть баланс" в какую-то сторону: углубить несколько связей или, наоборот, расширить круг тактических контактов',
),
- parse_mode=ParseMode.HTML,
)
else:
await message.answer(
@@ -385,6 +338,9 @@ async def get_summary(message: types.Message):
@dp.message(F.text == "Кол-во пользователей")
async def get_count(message: types.Message):
+ await userdb.add_ids_to_user(
+ message.from_user.username, message.from_user.id, message.chat.id
+ )
count = await userdb.count_users()
await message.answer(
f"Ботом уже воспользовались {count} человек{'а' if count % 10 >= 2 and count % 10 < 5 else ''}!\nНапоминаю, что для участия в розыгрыше нужно подписаться на @campusdna"
@@ -393,6 +349,9 @@ async def get_count(message: types.Message):
@dp.message(F.text[0] == "@")
async def user_name_checker(message: types.Message):
+ userdb.add_ids_to_user(
+ message.from_user.username, message.from_user.id, message.chat.id
+ )
msg = (message.text).strip()
try:
username_to = check_tg_username(msg)
@@ -405,7 +364,7 @@ async def user_name_checker(message: types.Message):
[
InlineKeyboardButton(
text="Близкий друг",
- callback_data=LinkCallback(
+ callback_data=callbacks.LinkCallback(
username_to=username_to, rating=3
).pack(),
)
@@ -413,7 +372,7 @@ async def user_name_checker(message: types.Message):
[
InlineKeyboardButton(
text="Приятель",
- callback_data=LinkCallback(
+ callback_data=callbacks.LinkCallback(
username_to=username_to, rating=2
).pack(),
),
@@ -421,7 +380,7 @@ async def user_name_checker(message: types.Message):
[
InlineKeyboardButton(
text="Знакомый",
- callback_data=LinkCallback(
+ callback_data=callbacks.LinkCallback(
username_to=username_to, rating=1
).pack(),
),
@@ -432,23 +391,21 @@ async def user_name_checker(message: types.Message):
await message.answer("Кто он для тебя?", reply_markup=kb)
-@dp.callback_query(LinkCallback.filter())
-async def process_data(query: CallbackQuery, callback_data: LinkCallback):
+@dp.callback_query(callbacks.LinkCallback.filter())
+async def process_data(query: CallbackQuery, callback_data: callbacks.LinkCallback):
from_username = query.from_user.username
- if from_username is None:
- raise Exception("WTF?")
if await userdb.get_user(from_username) is None:
- await query.answer("Похоже тебе нужно перезапустить бота: /start")
+ await query.message.answer("Похоже тебе нужно перезапустить бота: /start")
return
await userdb.add_link(
from_username,
Link(username_to=callback_data.username_to, rating=callback_data.rating),
)
await query.message.edit_text(
- **Text(
+ as_list(
f"✅ @{callback_data.username_to} добавлен как {rating_to_text(callback_data.rating).lower()}",
"\n📝 Чтобы добавить ещё друга — просто введи следующий юзернейм.",
"\n🔁 Чем больше друзей ты добавишь — тем точнее будет твой социальный портрет!",
- ).as_kwargs(),
+ ),
reply_markup=None,
)
diff --git a/src/callbacks.py b/src/callbacks.py
new file mode 100644
index 0000000..b978398
--- /dev/null
+++ b/src/callbacks.py
@@ -0,0 +1,26 @@
+from aiogram.filters import callback_data
+
+
+class StartingCallback(callback_data.CallbackData, prefix="start"):
+ pass
+
+
+class LinkCallback(callback_data.CallbackData, prefix="link"):
+ username_to: str
+ rating: int
+
+
+class SexCallback(callback_data.CallbackData, prefix="sex"):
+ sex: str
+
+
+class CourseCallback(callback_data.CallbackData, prefix="course"):
+ course: int
+
+
+class LivingCallback(callback_data.CallbackData, prefix="living"):
+ living: str
+
+
+class TypeInfoCallback(callback_data.CallbackData, prefix="typeinfo"):
+ pass
diff --git a/src/models.py b/src/models.py
index a6fcfe5..8dda885 100644
--- a/src/models.py
+++ b/src/models.py
@@ -5,26 +5,33 @@
def check_tg_username(username: str) -> str:
- username = username.strip().strip('@')
+ username = username.strip().strip("@")
pattern = r"^[A-z0-9_]+$"
if re.match(pattern, username):
return username
- raise ValueError(f'{username} is not valid username')
+ raise ValueError(f"{username} is not valid username")
+
Username = Annotated[str, AfterValidator(check_tg_username)]
Sex = Literal["male", "female"]
Living = Literal["Cloud", "Cosmos", "Baikal", "Homeless"]
+
class Link(BaseModel):
username_to: Username
- rating: int = Field(ge=1,le=3)
+ rating: int = Field(ge=1, le=3)
+
+
+class User(BaseModel):
+ userid: int
+ chatid: int
-class User(BaseModel) :
username: Username
- # TODO: Add sex, course, etc.
+
sex: Sex
- course: int = Field(ge=1,le=2)
+ course: int = Field(ge=1, le=2)
living: Living
+
_links: list[Link] = []
def set_link(self, link: Link):
@@ -32,4 +39,4 @@ def set_link(self, link: Link):
if i.username_to == link.username_to:
i = link
return
- self._links.append(link)
\ No newline at end of file
+ self._links.append(link)
diff --git a/src/templates.py b/src/templates.py
new file mode 100644
index 0000000..a6ac6f9
--- /dev/null
+++ b/src/templates.py
@@ -0,0 +1,59 @@
+from typing import List
+
+starting_message = """🧬 ДОБРО ПОЖАЛОВАТЬ В CAMPUS DNA!
+
+Мы создаём первую карту социальных связей нашего университета.
+Это исследование научной студии, и каждый участник получит:
+• Личный анализ социального типа
+• Место на интерактивной карте универа
+• Шанс выиграть КРУТЫЕ ПРИЗЫ 🎁
+
+🏆 УСЛОВИЯ УЧАСТИЯ В РОЗЫГРЫШЕ:
+
+1. ✅ Подписаться на канал @campusdna
+2. ✅ Отметь 5+ друзей и оцени вашу близость
+
+🎁 ПРИЗОВОЙ ФОНД:
+• 20 ПИЦЦ (1 пицца = 1 победитель)
+• ИГРУШКИ-МИНЬОНЫ
+• МЕРЧ ОТ ЦУ И ПАРТНЁРОВ
+
+📢 Следи за розыгрышами в канале: @campusdna
+
+🧭 ЧТО ДЕЛАТЬ ДАЛЬШЕ:
+
+Сначала тебе нужно ввести базовые сведения: пол, курс, общежитие.
+Затем я попрошу тебя ввести юзернеймы твоих друзей в Telegram
+и оценить вашу близость по шкале от 1 до 3.
+
+Чем больше друзей ты отметишь — тем точнее будет твой
+социальный портрет и тем ценнее твой вклад в исследование!
+
+Готов начать и узнать, кто ты в социальной сети университета?"""
+
+explaining_links_message = """⁉️ НА КАКИЕ ГРУППЫ МЫ ДЕЛИМ СВЯЗИ?
+🔴 1 — Друзья
+«С ними я провожу больше всего времени»
+Постоянное общение в универе и в мессенджерах. Видимся почти каждый день. Делимся личными новостями, поддерживаем друг друга.
+
+🟡 2 — Приятели
+«Всегда подойду спросить: "Как дела? Как жизнь?"»
+Видимся несколько раз в неделю. Общаемся и про учебу, и про жизнь, иногда затрагиваем личное (но не глубокое). Можем вместе пообедать или поиграть в пин-понг.
+
+🔵 3 — Знакомые
+«Мы здороваемся в коридоре»
+Видимся изредка, общение короткое и ситуативное. В основном на учебные/повседневные темы."""
+
+
+def make_type_str(type: str, profile: str, strong_sides: List[str], recomendation: str):
+ return f"""🎯ТИП: «{type}»
+
+📊 Ваш профиль:
+{profile}
+
+💪 Ваши сильные стороны:
+• {"\n• ".join(strong_sides)}
+
+🌟 Рекомендация:
+{recomendation}
+"""
diff --git a/src/userdb.py b/src/userdb.py
index 95e1d9b..7fa169c 100644
--- a/src/userdb.py
+++ b/src/userdb.py
@@ -1,68 +1,31 @@
import asyncio
import os
-from abc import ABC, abstractmethod
from typing import List, Optional
import pymongo.asynchronous.collection as pymongo_collection
import pymongo.asynchronous.database as pymongo_database
from dotenv import load_dotenv
from pymongo import AsyncMongoClient
+from pymongo.asynchronous.cursor import AsyncCursor
from pymongo.errors import ServerSelectionTimeoutError
from .models import Link, User, Username
-load_dotenv()
-
-MONGODB_HOST = os.getenv("MONGODB_HOST")
-
-
-class UserDB(ABC):
- @abstractmethod
- def add_user(self, user: User) -> None:
- pass
-
- @abstractmethod
- def get_user(self, username: Username) -> Optional[User]:
- pass
-
- @abstractmethod
- def get_links(self, username: Username) -> Optional[list[User]]:
- pass
-
- @abstractmethod
- def add_link(self, username: Username, link: Link) -> None:
- pass
-
class UserNotExist(Exception):
pass
-class MongoUserDB(UserDB):
- client: Optional[AsyncMongoClient] = None
- db: Optional[pymongo_database.AsyncDatabase] = None
- collection: Optional[pymongo_collection.AsyncCollection] = None
+class UserDB:
+ db: pymongo_database.AsyncDatabase
+ collection: pymongo_collection.AsyncCollection = None
- def __init__(self):
- # Не создаем клиент в __init__, а делаем это лениво
- pass
-
- async def _ensure_connection(self):
- """Ленивая инициализация подключения к MongoDB"""
- if self.client is None:
- try:
- self.client = AsyncMongoClient(
- MONGODB_HOST, serverSelectionTimeoutMS=5000
- )
- self.db = self.client.get_database("cu_graph_bot")
- self.collection = self.db.get_collection("users")
- await self.collection.create_index("username", unique=True)
- except ServerSelectionTimeoutError as e:
- print("Ошибка подключения к MongoDB:", e)
- raise e
+ def __init__(self, client):
+ self.db = client.get_database("cu_graph_bot")
+ self.collection = self.db.get_collection("users")
+ asyncio.create_task(self.collection.create_index("username", unique=True))
async def add_user(self, user: User) -> None:
- await self._ensure_connection()
current_user_to_add = await self.collection.find_one(
{"username": user.username}
)
@@ -72,7 +35,6 @@ async def add_user(self, user: User) -> None:
pass
async def get_links(self, username: Username) -> List[Link]:
- await self._ensure_connection()
all_friends = await self.collection.find_one({"username": username})
if all_friends is None:
raise UserNotExist
@@ -81,12 +43,16 @@ async def get_links(self, username: Username) -> List[Link]:
return [Link.model_validate(i) for i in all_friends["_links"]]
async def get_user(self, username: Username) -> Optional[User]:
- await self._ensure_connection()
all_user_data = await self.collection.find_one({"username": username})
return all_user_data
+ async def get_users(self, links_less_than: int) -> AsyncCursor:
+ all_user_data = self.collection.find(
+ {f"_links.{links_less_than}": {"$exists": False}}
+ )
+ return all_user_data
+
async def add_link(self, username: Username, link: Link) -> None:
- await self._ensure_connection()
current_user = await self.collection.find_one({"username": username})
if current_user is None:
raise UserNotExist
@@ -102,42 +68,22 @@ async def add_link(self, username: Username, link: Link) -> None:
{"username": username}, {"$push": {"_links": link.model_dump()}}
)
- async def count_users(self) -> int:
- await self._ensure_connection()
+ async def count_users(self, links_more_than: int = 4) -> int:
count = len(
- [i async for i in self.collection.find({"_links.4": {"$exists": True}})]
+ [
+ i
+ async for i in self.collection.find(
+ {f"_links.{links_more_than}": {"$exists": True}}
+ )
+ ]
)
return count
-
-class ListUserDB(UserDB):
- _users: list[User]
-
- def __init__(self) -> None:
- self._users = []
-
- def add_user(self, user: User) -> None:
- if user in self._users:
- raise ValueError
- self._users.append(user)
-
- def get_user(self, username: Username) -> Optional[User]:
- for i in self._users:
- if i.username == username:
- return i
- return None
-
- def get_links(self, username: Username) -> Optional[list[Link]]:
- for i in self._users:
- if i.username == username:
- return i._links
- return None
-
- def add_link(self, username: Username, link: Link) -> None:
- for i in self._users:
- if i.username == username:
- i.set_link(link)
- break
+ # NOTE: This is temporary function for fixing current database
+ # And should be deleted someday because of obvious performance loss
+ async def add_ids_to_user(self, username: str, userid: int, chatid: int) -> None:
+ await self.collection.update_one(
+ {"username": username}, {"$set": {"userid": userid, "chatid": chatid}}
+ )
-userdb = MongoUserDB()