From e6a39b98bd1db874d54193ada9b0903ea8d18269 Mon Sep 17 00:00:00 2001 From: csasq Date: Tue, 13 Aug 2024 17:46:44 +0300 Subject: [PATCH] =?UTF-8?q?=D0=94=D0=BE=D0=B1=D0=B0=D0=B2=D0=BB=D0=B5?= =?UTF-8?q?=D0=BD=D1=8B=20=D0=BC=D0=BE=D0=B4=D1=83=D0=BB=D0=B8=20=D0=B1?= =?UTF-8?q?=D0=B0=D0=B7=D1=8B=20=D0=B4=D0=B0=D0=BD=D0=BD=D1=8B=D1=85=20?= =?UTF-8?q?=D0=B8=20=D0=BC=D0=BE=D0=B4=D0=B5=D0=BB=D0=B5=D0=B9=20=D0=B4?= =?UTF-8?q?=D0=B0=D0=BD=D0=BD=D1=8B=D1=85,=20=D0=BE=D1=80=D0=B3=D0=B0?= =?UTF-8?q?=D0=BD=D0=B8=D0=B7=D0=BE=D0=B2=D0=B0=D0=BD=D0=B0=20=D0=B0=D1=80?= =?UTF-8?q?=D1=85=D0=B8=D1=82=D0=B5=D0=BA=D1=82=D1=83=D1=80=D0=B0=20=D1=81?= =?UTF-8?q?=D0=B8=D1=81=D1=82=D0=B5=D0=BC=D1=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bot/main.py | 246 ++++++++++++++++++++++---------------- config/__init__.py | 2 +- config/main.py | 88 +++++++++++--- database.sql | 74 ++++++++++-- database/__init__.py | 1 + database/main.py | 278 +++++++++++++++++++++++++++++++++++++++++++ models/__init__.py | 1 + models/main.py | 46 +++++++ requirements.txt | 1 + 9 files changed, 614 insertions(+), 123 deletions(-) create mode 100644 database/__init__.py create mode 100644 database/main.py create mode 100644 models/__init__.py create mode 100644 models/main.py diff --git a/bot/main.py b/bot/main.py index 107167b..17d62c2 100644 --- a/bot/main.py +++ b/bot/main.py @@ -1,13 +1,15 @@ -import asyncio +import os from aiogram import Bot, Dispatcher from aiogram.client.default import DefaultBotProperties from aiogram.enums import ParseMode from aiogram.filters import CommandStart -from aiogram.types import Message, Poll, PollAnswer, InputPollOption +from aiogram.types import User, Message, PollAnswer, InputPollOption from apscheduler.schedulers.asyncio import AsyncIOScheduler import config +import database +import models dp = Dispatcher() @@ -17,7 +19,7 @@ dp = Dispatcher() async def command_start_handler( message: Message, ): - config.Dynamic.set('chat_id', message.chat.id) + config.Redis.set('chat_id', message.chat.id) await message.answer('Чат успешно зарегистрирован!') @@ -31,124 +33,166 @@ bot = Bot( scheduler = AsyncIOScheduler() -async def send_mood_poll() -> Poll: +async def create_or_update_local_user_data( + telegram_user: User, +) -> models.User: + database_user = await database.Users.insert_or_update_user( + telegram_id=telegram_user.id, + first_name=telegram_user.first_name, + last_name=telegram_user.last_name, + username=telegram_user.username, + ) + user_profile_photos = await telegram_user.get_profile_photos( + limit=1, + ) + try: + photo_size = user_profile_photos.photos[0][0] + file = await bot.get_file(photo_size.file_id) + dir_path = os.path.join( + config.Main.cwd, + 'static', + 'images', + 'users', + ) + file_path = os.path.join( + dir_path, + '%s%s' % ( + database_user.id, + os.path.splitext(file.file_path)[1], + ), + ) + os.makedirs( + name=dir_path, + exist_ok=True, + ) + await bot.download_file( + file_path=file.file_path, + destination=file_path, + ) + except Exception: + pass + return database_user + + +async def send_mood_poll(): + poll_schema = await database.PollSchemas.get_poll_schema_by_name( + name='mood', + ) + poll_options = await database.PollOptions.get_poll_options( + poll_schema=poll_schema, + ) message = await bot.send_poll( - chat_id=config.Dynamic.get('chat_id'), - question='Оцените свое состояние на текущую минуту', + chat_id=config.Redis.get('chat_id'), + question=poll_schema.question, options=[ InputPollOption( - text='😄', - ), - InputPollOption( - text='🤪', - ), - InputPollOption( - text='🫠', - ), - InputPollOption( - text='☠️', - ), - InputPollOption( - text='🤡', - ), - InputPollOption( - text='😟', - ), - InputPollOption( - text='😩', - ), - InputPollOption( - text='😡', - ), + text=poll_option.name, + ) + for poll_option in poll_options ], is_anonymous=False, allows_multiple_answers=False, is_closed=False, disable_notification=True, ) - return message.poll - - -async def send_lunch_poll(): - message = await bot.send_poll( - chat_id=config.Dynamic.get('chat_id'), - question='Какие у вас планы на обед?', - options=[ - InputPollOption( - text='🍽️ Пойду в общепит', - ), - InputPollOption( - text='📦 Хочу заказать в офис', - ), - InputPollOption( - text='🥪 Всё своё ношу с собой', - ), - InputPollOption( - text='😴 Хочу спать', - ), - ], - is_anonymous=False, - allows_multiple_answers=False, - is_closed=False, - disable_notification=True, + await database.Polls.insert_poll( + telegram_message_id=message.message_id, + telegram_poll_id=message.poll.id, + poll_schema=poll_schema, ) - config.Dynamic.set('lunch_poll', message.poll.id) + + +# async def send_lunch_poll(): +# poll_schema = await database.PollSchemas.get_poll_schema_by_name( +# name='lunch', +# ) +# message = await bot.send_poll( +# chat_id=config.Redis.get('chat_id'), +# question=poll_schema.question, +# options=get_poll_options(poll_schema), +# is_anonymous=False, +# allows_multiple_answers=False, +# is_closed=False, +# disable_notification=True, +# ) +# await database.Polls.insert_poll( +# telegram_message_id=message.message_id, +# telegram_poll_id=message.poll.id, +# poll_schema=poll_schema, +# ) + + +# async def send_lunch_delivery_poll() -> Poll: +# message = await bot.send_poll( +# chat_id=config.Redis.get('chat_id'), +# question='Где будем заказывать?', +# options=[ +# InputPollOption( +# text='Dark Side', +# ), +# InputPollOption( +# text='Самокат', +# ), +# InputPollOption( +# text='...', +# ), +# InputPollOption( +# text='...', +# ), +# ], +# is_anonymous=False, +# allows_multiple_answers=False, +# is_closed=False, +# disable_notification=True, +# ) +# return message.poll @dp.poll_answer() -async def get_lunch_poll_result( +async def get_poll_answer( poll_answer: PollAnswer, ): - if poll_answer.poll_id == config.Dynamic.get('lunch_poll'): - await bot.send_message( - chat_id=config.Dynamic.get('chat_id'), - text='%s, %s' % (poll_answer.user.first_name, poll_answer.option_ids), - ) - - -async def send_lunch_delivery_poll() -> Poll: - message = await bot.send_poll( - chat_id=config.Dynamic.get('chat_id'), - question='Где будем заказывать?', - options=[ - InputPollOption( - text='Dark Side', - ), - InputPollOption( - text='Самокат', - ), - InputPollOption( - text='...', - ), - InputPollOption( - text='...', - ), - ], - is_anonymous=False, - allows_multiple_answers=False, - is_closed=False, - disable_notification=True, - ) - return message.poll + user = await create_or_update_local_user_data(poll_answer.user) + for poll in await database.Polls.get_polls(): + if poll.telegram_poll_id == poll_answer.poll_id: + if poll.is_complete: + await bot.delete_message( + chat_id=config.Redis.get('chat_id'), + message_id=poll.telegram_message_id, + ) + break + poll_options = await database.PollOptions.get_poll_options( + poll_schema=poll.poll_schema, + ordinals=poll_answer.option_ids, + ) + await database.PollAnswers.insert_or_update_poll_answer( + poll=poll, + user=user, + poll_options=poll_options, + ) + break async def on_startup( dispatcher: Dispatcher, ): - scheduler.add_job( - func=send_mood_poll, - trigger='cron', - day_of_week='mon-fri', - hour=22, - minute=35, - ) - scheduler.add_job( - func=send_lunch_poll, - trigger='cron', - day_of_week='mon-fri', - hour=23, - minute=55, - ) + # ### TEST + await send_mood_poll() + # TEST ### + # scheduler.add_job( + # func=send_mood_poll, + # trigger='cron', + # day_of_week='mon-fri', + # hour=22, + # minute=35, + # ) + # scheduler.add_job( + # func=send_lunch_poll, + # trigger='cron', + # day_of_week='mon-fri', + # hour=23, + # minute=55, + # ) # scheduler.add_job( # func=get_lunch_poll_result, # trigger='cron', diff --git a/config/__init__.py b/config/__init__.py index 76f014b..f5c2da8 100644 --- a/config/__init__.py +++ b/config/__init__.py @@ -1 +1 @@ -from .main import Telegram, Dynamic +from .main import Main, Postgres, Redis, Telegram diff --git a/config/main.py b/config/main.py index 72abd9f..636b588 100644 --- a/config/main.py +++ b/config/main.py @@ -1,7 +1,7 @@ from configparser import RawConfigParser import os -from redis import Redis +import redis cwd = os.getcwd() @@ -14,23 +14,73 @@ config.read( ) -class Telegram: - token = config.get( - section='Telegram', - option='token', +class Main: + cwd = config.get( + section='Main', + option='cwd', + fallback=os.getcwd(), ) -class Dynamic: +class Postgres: + host = config.get( + section='Postgres', + option='host', + fallback='localhost', + ) + port = config.getint( + section='Postgres', + option='port', + fallback=5432, + ) + user = config.get( + section='Postgres', + option='user', + ) + password = config.get( + section='Postgres', + option='password', + ) + dbname = config.get( + section='Postgres', + option='dbname', + ) + + +class Redis: + host = config.get( + section='Redis', + option='host', + fallback='localhost', + ) + port = config.getint( + section='Redis', + option='port', + fallback=6379, + ) + db = config.get( + section='Redis', + option='db', + ) + password = config.get( + section='Redis', + option='password', + fallback=None, + ) + @classmethod def get( cls, key: str, ): - with Redis( - db=3, - ) as redis: - return redis.get( + with redis.Redis( + host=cls.host, + port=cls.port, + db=cls.db, + password=cls.password, + decode_responses=True, + ) as connection: + return connection.get( name=key, ) @@ -40,11 +90,21 @@ class Dynamic: key: str, value, ): - with Redis( - db=3, + with redis.Redis( + host=cls.host, + port=cls.port, + db=cls.db, + password=cls.password, decode_responses=True, - ) as redis: - redis.set( + ) as connection: + connection.set( name=key, value=value, ) + + +class Telegram: + token = config.get( + section='Telegram', + option='token', + ) diff --git a/database.sql b/database.sql index 967e951..d155b9b 100644 --- a/database.sql +++ b/database.sql @@ -1,9 +1,69 @@ -create table polls ( - id bigint not null, - date date not null, - is_complete boolean default false not null, - primary key (id) +create table users ( + id bigserial not null, + telegram_user_id bigint not null, + first_name character varying (64) not null, + last_name character varying (64), + username character varying (32), + primary key (id), + unique (telegram_user_id) ); -create table mood_polls () inherits (polls); -create table lunch_polls () inherits (polls); +create table poll_schemas ( + id bigserial not null, + name character varying (32) not null, + question character varying (255) not null, + primary key (id), + unique (name) +); + +create table poll_options ( + id bigserial not null, + poll_schema_id bigint not null, + name character varying (100) not null, + ordinal bigint, + primary key (id), + foreign key (poll_schema_id) references poll_schemas on delete cascade on update cascade, + unique (poll_schema_id, name), + unique (poll_schema_id, ordinal) +); + +create table polls ( + id bigserial not null, + telegram_message_id bigint not null, + telegram_poll_id text not null, + poll_schema_id bigint not null, + created_at timestamp default now() not null, + is_complete boolean default false not null, + primary key (id), + foreign key (poll_schema_id) references poll_schemas on delete cascade on update cascade, + unique (telegram_message_id), + unique (telegram_poll_id) +); + +create table poll_answers ( + id bigserial not null, + poll_id bigint not null, + user_id bigint not null, + poll_option_id bigint not null, + primary key (id), + foreign key (poll_id) references polls on delete cascade on update cascade, + foreign key (user_id) references users on delete cascade on update cascade, + foreign key (poll_option_id) references poll_options on delete cascade on update cascade, + unique (poll_id, user_id, poll_option_id) +); + +insert into poll_schemas (name, question) values ('mood', 'Оцените свое состояние на текущую минуту'); +insert into poll_options (poll_schema_id, name, ordinal) values (1, '😄', 0); +insert into poll_options (poll_schema_id, name, ordinal) values (1, '🤪', 1); +insert into poll_options (poll_schema_id, name, ordinal) values (1, '🫠', 2); +insert into poll_options (poll_schema_id, name, ordinal) values (1, '☠️', 3); +insert into poll_options (poll_schema_id, name, ordinal) values (1, '🤡', 4); +insert into poll_options (poll_schema_id, name, ordinal) values (1, '😟', 5); +insert into poll_options (poll_schema_id, name, ordinal) values (1, '😩', 6); +insert into poll_options (poll_schema_id, name, ordinal) values (1, '😡', 7); + +insert into poll_schemas (name, question) values ('lunch', 'Какие у вас планы на обед?'); +insert into poll_options (poll_schema_id, name, ordinal) values (2, '🍽️ Пойду в общепит', 0); +insert into poll_options (poll_schema_id, name, ordinal) values (2, '📦 Хочу заказать в офис', 1); +insert into poll_options (poll_schema_id, name, ordinal) values (2, '🥪 Всё своё ношу с собой', 2); +insert into poll_options (poll_schema_id, name, ordinal) values (2, '😴 Хочу спать', 3); diff --git a/database/__init__.py b/database/__init__.py new file mode 100644 index 0000000..b8ca897 --- /dev/null +++ b/database/__init__.py @@ -0,0 +1 @@ +from .main import Users, PollSchemas, PollOptions, Polls, PollAnswers diff --git a/database/main.py b/database/main.py new file mode 100644 index 0000000..4d167f5 --- /dev/null +++ b/database/main.py @@ -0,0 +1,278 @@ +from psycopg import AsyncConnection + +import config +import models + + +conninfo = 'host=%(host)s port=%(port)s user=%(user)s password=%(password)s dbname=%(dbname)s' % { + 'host': config.Postgres.host, + 'port': config.Postgres.port, + 'user': config.Postgres.user, + 'password': config.Postgres.password, + 'dbname': config.Postgres.dbname, +} + + +class NotFoundError(Exception): + pass + + +class Users: + @staticmethod + async def insert_or_update_user( + telegram_id: int, + first_name: str, + last_name: str = None, + username: str = None, + ) -> models.User: + async with await AsyncConnection.connect(conninfo) as connection: + async with connection.cursor() as cursor: + sql = ''' + insert into users ( + telegram_id, + first_name, + last_name, + username + ) values ( + %(telegram_id)s, + %(first_name)s, + %(last_name)s, + %(username)s + ) on conflict (telegram_id) do update set + first_name = excluded.first_name, + last_name = excluded.last_name, + username = excluded.username + returning + users.id; + ''' + await cursor.execute( + sql, + { + 'telegram_id': telegram_id, + 'first_name': first_name, + 'last_name': last_name, + 'username': username, + }, + ) + user_id, = await cursor.fetchone() + return models.User( + id=user_id, + telegram_id=telegram_id, + first_name=first_name, + last_name=last_name, + username=username, + ) + + +class PollSchemas: + @staticmethod + async def get_poll_schema_by_name( + name: str, + ) -> models.PollSchema: + async with await AsyncConnection.connect(conninfo) as connection: + async with connection.cursor() as cursor: + sql = ''' + select + poll_schemas.id, + poll_schemas.name, + poll_schemas.question + from + poll_schemas + where + poll_schemas.name = %(name)s; + ''' + await cursor.execute( + sql, + { + 'name': name, + }, + ) + try: + poll_schema_id, name, question = await cursor.fetchone() + except TypeError: + raise NotFoundError() + return models.PollSchema( + id=poll_schema_id, + name=name, + question=question, + ) + + +class PollOptions: + @staticmethod + async def get_poll_options( + poll_schema: models.PollSchema, + ordinals: list[int] = None, + ) -> list[models.PollOption]: + async with await AsyncConnection.connect(conninfo) as connection: + async with connection.cursor() as cursor: + sql = ''' + select + poll_options.id, + poll_options.name, + poll_options.ordinal + from + poll_options + where + poll_options.poll_schema_id = %(poll_schema_id)s + and poll_options.ordinal is not null; + ''' if ordinals is not None else ''' + select + poll_options.id, + poll_options.name, + poll_options.ordinal + from + poll_options + where + poll_options.poll_schema_id = %(poll_schema_id)s + and poll_options.ordinal = any(%(ordinals)s); + ''' + await cursor.execute( + sql, + { + 'poll_schema_id': poll_schema.id, + 'ordinals': ordinals, + }, + ) + records = await cursor.fetchall() + return [ + models.PollOption( + id=poll_option_id, + name=name, + ordinal=ordinal, + ) + for poll_option_id, name, ordinal + in records + ] + + +class Polls: + @staticmethod + async def get_polls() -> list[models.Poll]: + async with await AsyncConnection.connect(conninfo) as connection: + async with connection.cursor() as cursor: + sql = ''' + select + polls.id, + polls.telegram_message_id, + polls.telegram_poll_id, + poll_schemas.id, + poll_schemas.name, + poll_schemas.question, + poll_schemas.options, + polls.created_at, + polls.is_active + from + polls + inner join poll_schemas on + polls.schema_id = poll_schemas.id; + ''' + await cursor.execute(sql) + records = await cursor.fetchall() + return [ + models.Poll( + id=poll_id, + telegram_message_id=telegram_message_id, + telegram_poll_id=telegram_poll_id, + poll_schema=models.PollSchema( + id=poll_schema_id, + name=name, + question=question, + options=options, + ), + created_at=created_at, + is_active=is_active, + ) + for poll_id, telegram_message_id, telegram_poll_id, poll_schema_id, name, question, options, created_at, is_active + in records + ] + + @staticmethod + async def insert_poll( + telegram_message_id: int, + telegram_poll_id: str, + poll_schema: models.PollSchema, + ) -> models.Poll: + async with await AsyncConnection.connect(conninfo) as connection: + async with connection.cursor() as cursor: + sql = ''' + insert into polls ( + telegram_message_id, + telegram_poll_id, + poll_schema + ) values ( + %(telegram_message_id)s, + %(telegram_poll_id)s, + %(poll_schema_id)s + ) returning + polls.id, + polls.created_at, + polls.is_active; + ''' + await cursor.execute( + sql, + { + 'telegram_message_id': telegram_message_id, + 'telegram_poll_id': telegram_poll_id, + 'poll_schema': poll_schema.id, + }, + ) + try: + poll_id, created_at, is_active = await cursor.fetchone() + except TypeError: + raise NotFoundError() + return models.Poll( + id=poll_id, + telegram_message_id=telegram_message_id, + telegram_poll_id=telegram_poll_id, + poll_schema=poll_schema, + created_at=created_at, + is_active=is_active, + ) + + +class PollAnswers: + @staticmethod + async def insert_or_update_poll_answer( + poll: models.Poll, + user: models.User, + poll_options: list[models.PollOption], + ) -> list[models.PollAnswer]: + async with await AsyncConnection.connect(conninfo) as connection: + async with connection.cursor() as cursor: + sql = ''' + insert into poll_answers ( + poll_id, + user_id, + poll_option_id + ) values ( + %(poll_id)s, + %(user_id)s, + %(poll_option_id)s + ) on conflict (poll_id, user_id, poll_option_id) do nothing + returning + poll_answers.id; + ''' + await cursor.executemany( + sql, + [ + { + 'poll_id': poll.id, + 'user_id': user.id, + 'poll_option_id': poll_option.id, + } + for poll_option + in poll_options + ], + ) + records = await cursor.fetchall() + return [ + models.PollAnswer( + id=poll_answer_id, + poll=poll, + user=user, + poll_option=poll_options, + ) + for poll_answer_id, + in records + ] diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000..b0caa36 --- /dev/null +++ b/models/__init__.py @@ -0,0 +1 @@ +from .main import User, PollSchema, PollOption, Poll, PollAnswer diff --git a/models/main.py b/models/main.py new file mode 100644 index 0000000..0811bc2 --- /dev/null +++ b/models/main.py @@ -0,0 +1,46 @@ +from datetime import datetime + +from pydantic import BaseModel, Field + + +class User(BaseModel): + id: int + telegram_id: int = Field( + exclude=True, + ) + first_name: str + last_name: str + username: str + + +class PollSchema(BaseModel): + id: int + name: str + question: str + + +class PollOption(BaseModel): + id: int + poll_schema: PollSchema + name: str + ordinal: int + + +class Poll(BaseModel): + id: int + telegram_message_id: int = Field( + exclude=True, + ) + telegram_poll_id: str = Field( + exclude=True, + ) + poll_schema: PollSchema + created_at: datetime + is_complete: bool + + +class PollAnswer(BaseModel): + id: int + poll: Poll + user: User + poll_option: list[PollOption] diff --git a/requirements.txt b/requirements.txt index b2ee50f..85efd53 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ aiogram APScheduler psycopg[binary] +pydantic redis