Bot and DB tweaks

Changes:
	General:
		- Implemented logging(loguru) throughout back and front
		  ends(primarily error handlers)
	Frontend:
		- Implemented setting commands for telegram bot command
		  menu
	Backend:
		- A schema and a table for bot are created in PGSQL DB
		  if there are no any

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(
+    try:
-    dbname = config.db_name,
+        connection = psycopg2.connect(
-    user = config.postgres_user,
+        dbname = config.db_name,
-    password = config.postgres_password,
+        user = config.postgres_user,
-    host = config.host_name,
+        password = config.postgres_password,
-    port = config.port
+        host = config.host_name,
-    )
+        port = config.port
-    cursor = connection.cursor()
+        )
-    return cursor, connection
+        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()
+    try:
-    connection.close()
+        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)

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(
+    try:
-        config.IS_address,
+        minio_client = Minio(
-        access_key = config.acc_key,
+            config.IS_address,
-        secret_key = config.sec_key,
+            access_key = config.acc_key,
-        secure = False
+            secret_key = config.sec_key,
-    )
+            secure = False
-    return minio_client
+        )
+        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,

src/Frontend/createbot.py

@@ -1,31 +1,37 @@
+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
-from aiogram import Bot, Dispatcher, Router
+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')
 
-
+logger.add(
-logging.basicConfig(level = logging.INFO, format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+    "sys.stdout",
-logger = logging.getLogger(__name__)
+    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):
@@ -37,17 +43,24 @@ List of available commands:
 /subscribe - subscribe to daily cat images sent at 12:00 UTC+3
 /subscription_modify <number> - 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.')
+        if psycopg2.errors.UniqueViolation:
-        print(e.pgerror)
+            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:
+    except psycopg2.Error as e:
-        await message.answer('You are not yet subscribed.')
+        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)
-