n0one/CatBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Bot and DB tweaks

Browse Source 
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
This commit is contained in:
n0body
2025-08-28 19:43:37 +03:00
parent f7bcb89afa
commit e6537b2817
4 changed files with 164 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

BIN
requirements.txt
View File

Binary file not shown.

78
src/Backend/DBwork.py
View File

@ -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)

36
src/Backend/ISwork.py
View File

@ -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,

93
src/Frontend/createbot.py
View File

@ -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 <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.')
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)