Merge branch 'main' of https://git.frik.su/n0one/CatBot

src/Backend/ISwork.py

@@ -13,24 +13,25 @@ def _setClient():
     )
     return minio_client
 
-def getImageName(currentDay):
-    maxFiles = 2
+def getNumberofObjects(client, currentDay, bucket_name):
+    objects = client.list_objects(bucket_name, prefix=currentDay+'/')
+    return sum(1 for _ in objects)
+
+def getImageName(currentDay, client):
+    maxFiles = getNumberofObjects(client, currentDay, config.bucket_name)
     fileNumber = randint(1, maxFiles)
     desiredFile = currentDay + '/' + str(fileNumber) + '.jpeg'
     return desiredFile
 
 def downloadImage(currentDay, username):
-    bucket_name = "cat-images"
     client = _setClient()
-    client.fget_object(bucket_name, getImageName(currentDay), username + '.jpeg')
+    client.fget_object(config.bucket_name, getImageName(currentDay, client), username + '.jpeg')
 
 def downloadForAll(currentDay):
     cur, conn = DBwork.set_connection()
-    counter = 1
-    user = DBwork.get_user(counter, cur)
-    while(user != 'Error'):
+    max_id = DBwork.get_last_id(cur)
+    for id in range(1, max_id + 1):
+        user = DBwork.get_user(id, cur)
         downloadImage(currentDay, user)
-        counter += 1
-        user = DBwork.get_user(counter, cur)
     DBwork.close_connection(conn, cur)
 

src/Backend/dbwork.py

@@ -2,7 +2,7 @@ import psycopg2
 from src import config
 
 
-def _get_last_id(cursor):
+def get_last_id(cursor):
     cursor.execute("SELECT MAX(id) FROM usernames")
     id = cursor.fetchall()[0][0]
     return id
@@ -24,25 +24,18 @@ def close_connection(connection, cursor):
 
 #Functions don't close connection automatically, it has to be closed manually
 def add_user(username, connection, cursor):
-    cursor = connection.cursor()
-    cursor.execute("INSERT INTO usernames VALUES (%s, %s);", (_get_last_id(cursor) + 1, username))
+    cursor.execute("INSERT INTO usernames VALUES (%s, %s);", (get_last_id(cursor) + 1, username))
     connection.commit()
     
 def delete_user(username, connection, cursor):
-    cursor = connection.cursor()
     cursor.execute("DELETE FROM usernames WHERE username = %s;", (username,))
     connection.commit()
 
 def change_name(old_username, new_username, connection, cursor):
-    cursor = connection.cursor()
     cursor.execute("UPDATE usernames SET username = %s WHERE username = %s;", (new_username, old_username))
     connection.commit()
 
 def get_user(id, cursor):
-    max_id = _get_last_id(cursor)
-    if id <= max_id:
-        cursor.execute("SELECT username FROM usernames WHERE id = %s", (id,))
-        username = cursor.fetchall()[0][0]
-    else:
-        username = 'Error'
-    return username, cursor
+    cursor.execute("SELECT username FROM usernames WHERE id = %s", (id,))
+    username = cursor.fetchall()[0][0]
+    return username

src/config.py

@@ -9,6 +9,7 @@ postgres_user = config('postgres_user')
 postgres_password = config('postgres_password')
 host_name = config('host_name')
 port = config('port')
+bucket_name = 'cat-images'
 
 
 TG_TOKEN = config('TG_TOKEN')