diff --git a/requirements.txt b/requirements.txt index cfaeed3..6be3002 100644 Binary files a/requirements.txt and b/requirements.txt differ diff --git a/src/Backend/DBwork.py b/src/Backend/DBwork.py index 1f5edc2..2d5d535 100644 --- a/src/Backend/DBwork.py +++ b/src/Backend/DBwork.py @@ -1,5 +1,16 @@ +import logging + import psycopg2 from src import config +from loguru import logger + + +logger.add( + "sys.stdout", + format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {level} | {file}:{line} - {message}", + colorize=True, + level="INFO" +) def get_last_id(cursor): @@ -11,20 +22,27 @@ def get_last_id(cursor): def set_connection(): - connection = psycopg2.connect( - dbname = config.db_name, - user = config.postgres_user, - password = config.postgres_password, - host = config.host_name, - port = config.port - ) - cursor = connection.cursor() - return cursor, connection + try: + connection = psycopg2.connect( + dbname = config.db_name, + user = config.postgres_user, + password = config.postgres_password, + host = config.host_name, + port = config.port + ) + cursor = connection.cursor() + logger.info('Successfully set connection to the PostgreSQL DB') + return cursor, connection + except psycopg2.Error as e: + logger.error(f'Failed to set connection to the PostgreSQL DB: {e.pgerror}') def close_connection(connection, cursor): - cursor.close() - connection.close() + try: + cursor.close() + connection.close() + except psycopg2.Error as e: + logger.error(f'Failed to close PostgreSQL connection: {e.pgerror}') #Functions don't close connection automatically, it has to be closed manually @@ -53,3 +71,41 @@ def get_chat_id(id, cursor): cursor.execute("SELECT chat_id FROM Users WHERE id = %s", (id,)) chat_id = cursor.fetchall()[0][0] return chat_id + + +def schema_creator(schema_name): + cur, conn = set_connection() + try: + cur.execute(f'CREATE SCHEMA IF NOT EXISTS {schema_name};') + conn.commit() + logger.info(f'Successfully created schema {schema_name} if it didn\'t exist yet') + except psycopg2.Error as e: + logger.error(f'Error during schema creation: {e}') + finally: + close_connection(conn, cur) + + +def table_creator(schema_name, table_name): + cur, conn = set_connection() + try: + cur.execute(f''' +CREATE TABLE IF NOT EXISTS {schema_name}.{table_name} +( + id integer NOT NULL DEFAULT nextval('users_id_seq'::regclass), + chat_id bigint NOT NULL, + images_amount bigint, + CONSTRAINT users_pkey PRIMARY KEY (id), + CONSTRAINT chat_id_unique UNIQUE (chat_id) +) + +TABLESPACE pg_default; + +ALTER TABLE IF EXISTS {schema_name}.users + OWNER to {config.postgres_user}; +''') + conn.commit() + logger.info(f'Successfully created table {table_name} in schema {schema_name} if it didn\'t exist yet') + except psycopg2.Error as e: + logging.error(f'Error during table creation: {e}') + finally: + close_connection(conn, cur) \ No newline at end of file diff --git a/src/Backend/ISwork.py b/src/Backend/ISwork.py index bfa2c14..e960ec1 100644 --- a/src/Backend/ISwork.py +++ b/src/Backend/ISwork.py @@ -1,24 +1,35 @@ -from minio import Minio +from minio import Minio, S3Error from random import randint from datetime import timedelta from src import config +from loguru import logger + + +logger.add( + "sys.stdout", + format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {level} | {file}:{line} - {message}", + colorize=True, + level="INFO" +) def _setClient(): - minio_client = Minio( - config.IS_address, - access_key = config.acc_key, - secret_key = config.sec_key, - secure = False - ) - return minio_client + try: + minio_client = Minio( + config.IS_address, + access_key = config.acc_key, + secret_key = config.sec_key, + secure = False + ) + logger.info('Successfully set connection to the MinIO bucket') + return minio_client + except S3Error as e: + logger.error(f'S3 error during connection to bucket. Code: {e.code}, Message: {e.message}') def getNumberofObjects(client, currentDay): objects = client.list_objects(config.bucket_name, prefix=str(currentDay) + '/') numberOfObjects = sum(1 for _ in objects) - if numberOfObjects == 0: - return 'No objects in the folder' return numberOfObjects @@ -35,6 +46,8 @@ def getObjectExtension(client, currentDay, fileNumber): def getImageName(currentDay, client): maxFiles = getNumberofObjects(client, currentDay) + if maxFiles == 0: + return None fileNumber = randint(1, maxFiles) fileExtension = '.' + getObjectExtension(client, currentDay, fileNumber) desiredFile = str(currentDay) + '/' + str(fileNumber) + fileExtension @@ -44,6 +57,9 @@ def getImageName(currentDay, client): def getDownloadURL(currentDay): client = _setClient() object_name = getImageName(currentDay, client) + if object_name is None: + logger.error(f"Can't generate a URL: no files in current MinIO directory({currentDay})") + return None url = client.presigned_get_object( config.bucket_name, object_name, diff --git a/src/Frontend/createbot.py b/src/Frontend/createbot.py index 40d5a00..f00224d 100644 --- a/src/Frontend/createbot.py +++ b/src/Frontend/createbot.py @@ -1,53 +1,66 @@ +from loguru import logger import psycopg2 import Backend.ISwork, Backend.DBwork from src import config from src.Backend import DBwork from src.Backend import ISwork -import logging -from aiogram import Bot, Dispatcher, Router +from aiogram import Bot, Dispatcher +from aiogram.fsm.storage.memory import MemoryStorage from aiogram.filters import Command, CommandObject -from aiogram.types import Message, URLInputFile +from aiogram.types import Message, URLInputFile, BotCommand, BotCommandScopeDefault, BotCommandScopeAllPrivateChats, BotCommandScopeAllGroupChats from aiogram.client.default import DefaultBotProperties from aiogram.enums import ParseMode -from aiogram.fsm.storage.memory import MemoryStorage +from aiogram.methods import DeleteMyCommands from apscheduler.schedulers.asyncio import AsyncIOScheduler from datetime import datetime current_day = datetime.now().weekday() -start_router = Router() scheduler = AsyncIOScheduler(timezone = 'Europe/Moscow') - -logging.basicConfig(level = logging.INFO, format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s') -logger = logging.getLogger(__name__) +logger.add( + "sys.stdout", + format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {level} | {file}:{line} - {message}", + colorize=True, + level="INFO" +) bot = Bot(token = config.TG_token , default = DefaultBotProperties(parse_mode = ParseMode.HTML)) dp = Dispatcher(storage = MemoryStorage()) +schema_name = 'catbot' +table_name = 'Users' + @dp.message(Command('start')) async def cmd_start(message: Message): await message.answer(''' This is a bot that sends images of cats. - + List of available commands: /cat - request an image of a cat from today's pool of images /subscribe - subscribe to daily cat images sent at 12:00 UTC+3 /subscription_modify - change the amount of images sent daily /unsubscribe - cancel your subscription(but why would you want to? :3) -''') +''', parse_mode=None) + logger.info(f'Command /start executed successfully. ChatID: {message.chat.id}') @dp.message(Command('cat')) async def cmd_cat(message: Message): + chat_id = message.chat.id image_link = URLInputFile(ISwork.getDownloadURL(current_day), filename=datetime.now().strftime('%Y_%m_%d_%H_%M_%S')) - await message.answer_photo(image_link, caption='Look, a cat :3') + if image_link is None: + await bot.send_message(chat_id=chat_id, text="We are sorry, but there seems to be a problem with finding images for today.") + else: + await message.answer_photo(image_link, caption='Look, a cat :3') + logger.info(f'Command /cat executed successfully. ChatID: {chat_id}') @dp.message(Command('subscription_modify')) async def subscription_modify(message: Message, command: CommandObject): + chat_id = message.chat.id if command.args is None or command.args.isdigit() == False: await message.answer('Please write the number of images you would like\n' 'to receive in the same message as a command.\n' @@ -56,45 +69,60 @@ async def subscription_modify(message: Message, command: CommandObject): return try: cursor, connection = DBwork.set_connection() - chat_id = message.chat.id amount = command.args DBwork.change_images_amount(chat_id, amount, connection, cursor) DBwork.close_connection(connection, cursor) except psycopg2.Error: - await message.answer('You are not yet subscribed.') + if psycopg2.errors.IntegrityError: + await message.answer('You are not yet subscribed.') + logger.warning(f'A non-subscribed user in chat {chat_id} tried to modify subscription') + else: + await message.answer('There seems to be a problem on our side.') return await message.answer('Amount of daily images was changed successfully!') + logger.info(f'Command /subscription_modify executed successfully. ChatID: {chat_id}') @dp.message(Command('subscribe')) async def cmd_subscribe(message: Message): + chat_id = message.chat.id try: cursor, connection = DBwork.set_connection() - chat_id = message.chat.id DBwork.add_user(chat_id, connection, cursor) DBwork.close_connection(connection, cursor) except psycopg2.Error as e: - await message.answer('You are already subscribed.') - print(e.pgerror) + if psycopg2.errors.UniqueViolation: + await message.answer('You are already subscribed.') + logger.warning(f'An already subscribed user in chat {chat_id} tried to subscribe again') + else: + await message.answer('There seems to be a problem on our side.') + logger.error(f'PostgreSQL error occurred. ChatID: {chat_id}. Error: {str(e.pgerror)}') return await message.answer(''' You have successfully subscribed to daily cat photos! You will get 1 photo a day by default, use /subscription_modify to change that amount. ''') + logger.info(f'Command /subscribe executed successfully. ChatID: {chat_id}') @dp.message(Command('unsubscribe')) async def cmd_unsubscribe(message: Message): + chat_id = message.chat.id try: cursor, connection = DBwork.set_connection() - chat_id = message.chat.id DBwork.delete_user(chat_id, connection, cursor) DBwork.close_connection(connection, cursor) - except psycopg2.Error: - await message.answer('You are not yet subscribed.') + except psycopg2.Error as e: + if psycopg2.errors.NoData: + await message.answer('You are not yet subscribed.') + logger.warning(f'A non-subscribed user in chat {chat_id} tried to unsubscribe') + else: + await message.answer('There seems to be a problem on our side.') + logger.error(f'PostgreSQL error occurred. ChatID: {chat_id}. Error: {str(e.pgerror)}') return await message.answer('You have successfully unsubscribed.') + logger.info(f'Command /unsubscribe executed successfully. ChatID: {chat_id}') async def send_daily_images(): @@ -105,14 +133,35 @@ async def send_daily_images(): images_amount = DBwork.get_images_amount(chat_id, cursor) for _ in range(images_amount): image_link = URLInputFile(ISwork.getDownloadURL(current_day), filename=datetime.now().strftime('%Y_%m_%d_%H_%M_%S')) - await bot.send_photo(chat_id = chat_id, photo = image_link) + if image_link is None: + await bot.send_message(chat_id=chat_id, text="We are sorry, but there seems to be a problem with finding images for today.") + else: + await bot.send_photo(chat_id = chat_id, photo = image_link) DBwork.close_connection(connection, cursor) + logger.info('Daily mass sending to subscribers has finished') + + +async def set_commands_for_menu(): + # await bot(DeleteMyCommands(scope=BotCommandScopeDefault())) + logger.info('Bot command list cleared') + commands = [ + BotCommand(command='start', description='Get info about the bot and its commands'), + BotCommand(command='cat', description='Request an image of a cat from today\'s pool of images'), + BotCommand(command='subscribe', description='Subscribe to daily cat images sent at 12:00 UTC+3'), + BotCommand(command='subscription_modify', description='Change the amount of images sent daily'), + BotCommand(command='unsubscribe', description='Cancel your subscription') + ] + await bot.set_my_commands(commands=commands, scope=BotCommandScopeAllPrivateChats()) + await bot.set_my_commands(commands=commands, scope=BotCommandScopeAllGroupChats()) + logger.info('Command menu for bot has been set') scheduler.add_job(send_daily_images, 'cron', hour = 12, minute = 0) async def main(): + DBwork.schema_creator(schema_name) + DBwork.table_creator(schema_name, table_name) + await set_commands_for_menu() scheduler.start() - await dp.start_polling(bot) - + await dp.start_polling(bot) \ No newline at end of file