Handle Errors

Author: UltimaHoarder

helpers/db_helper.py

@@ -29,24 +29,36 @@ def create_database_session(
 
 
 def run_revisions(alembic_directory: str, database_path: str = ""):
-    ini_path = os.path.join(alembic_directory, "alembic.ini")
-    script_location = os.path.join(alembic_directory, "alembic")
-    full_database_path = f"sqlite:///{database_path}"
-    alembic_cfg = Config(ini_path)
-    alembic_cfg.set_main_option("script_location", script_location)
-    alembic_cfg.set_main_option("sqlalchemy.url", full_database_path)
-    command.upgrade(alembic_cfg, "head")
-    command.revision(alembic_cfg, autogenerate=True, message="content")
+    while True:
+        try:
+            ini_path = os.path.join(alembic_directory, "alembic.ini")
+            script_location = os.path.join(alembic_directory, "alembic")
+            full_database_path = f"sqlite:///{database_path}"
+            alembic_cfg = Config(ini_path)
+            alembic_cfg.set_main_option("script_location", script_location)
+            alembic_cfg.set_main_option("sqlalchemy.url", full_database_path)
+            command.upgrade(alembic_cfg, "head")
+            command.revision(alembic_cfg, autogenerate=True, message="content")
+            break
+        except Exception as e:
+            print(e)
+            print
 
 
 def run_migrations(alembic_directory: str, database_path: str) -> None:
-    ini_path = os.path.join(alembic_directory, "alembic.ini")
-    script_location = os.path.join(alembic_directory, "alembic")
-    full_database_path = f"sqlite:///{database_path}"
-    alembic_cfg = Config(ini_path)
-    alembic_cfg.set_main_option("script_location", script_location)
-    alembic_cfg.set_main_option("sqlalchemy.url", full_database_path)
-    command.upgrade(alembic_cfg, "head")
+    while True:
+        try:
+            ini_path = os.path.join(alembic_directory, "alembic.ini")
+            script_location = os.path.join(alembic_directory, "alembic")
+            full_database_path = f"sqlite:///{database_path}"
+            alembic_cfg = Config(ini_path)
+            alembic_cfg.set_main_option("script_location", script_location)
+            alembic_cfg.set_main_option("sqlalchemy.url", full_database_path)
+            command.upgrade(alembic_cfg, "head")
+            break
+        except Exception as e:
+            print(e)
+            print
 
 
 class database_collection(object):
@@ -63,14 +75,13 @@ def database_picker(self, database_name):
         return database
 
 
-
 def create_auth_array(item):
     auth_array = item.__dict__
     auth_array["support_2fa"] = False
     return auth_array
 
 
-def get_or_create(session: Session, model, defaults=None, fbkwargs:dict={}):
+def get_or_create(session: Session, model, defaults=None, fbkwargs: dict = {}):
     fbkwargs2 = fbkwargs.copy()
     instance = session.query(model).filter_by(**fbkwargs2).one_or_none()
     if instance:

helpers/main_helper.py

@@ -136,7 +136,9 @@ async def format_image(filepath: str, timestamp: float):
             break
 
 
-async def async_downloads(download_list: list[template_media_table], subscription: create_user):
+async def async_downloads(
+    download_list: list[template_media_table], subscription: create_user
+):
     async def run(download_list: list[template_media_table]):
         session_m = subscription.session_manager
         proxies = session_m.proxies
@@ -166,7 +168,9 @@ async def run(download_list: list[template_media_table]):
             responses = await asyncio.gather(*tasks)
             tasks.clear()
 
-            async def check(download_item: template_media_table, response: ClientResponse):
+            async def check(
+                download_item: template_media_table, response: ClientResponse
+            ):
                 filepath = os.path.join(download_item.directory, download_item.filename)
                 response_status = False
                 if response.status == 200:
@@ -487,16 +491,21 @@ def export_sqlite2(archive_path, datas, parent_type, legacy_fixer=False):
 
 
 def legacy_sqlite_updater(
-    legacy_metadata_path:str, api_type:str, subscription:create_user, delete_metadatas:list
+    legacy_metadata_path: str,
+    api_type: str,
+    subscription: create_user,
+    delete_metadatas: list,
 ):
     final_result = []
     if os.path.exists(legacy_metadata_path):
         cwd = os.getcwd()
-        alembic_location = os.path.join(cwd, "database", "archived_databases", api_type.lower())
+        alembic_location = os.path.join(
+            cwd, "database", "archived_databases", api_type.lower()
+        )
         db_helper.run_migrations(alembic_location, legacy_metadata_path)
         database_name = "user_data"
         session, engine = db_helper.create_database_session(legacy_metadata_path)
-        database_session:Session = session()
+        database_session: Session = session()
         db_collection = db_helper.database_collection()
         database = db_collection.database_picker(database_name)
         if database:
@@ -526,10 +535,10 @@ def legacy_sqlite_updater(
     return final_result, delete_metadatas
 
 
-def export_sqlite(database_path:str, api_type, datas):
+def export_sqlite(database_path: str, api_type, datas):
     metadata_directory = os.path.dirname(database_path)
     os.makedirs(metadata_directory, exist_ok=True)
-    database_name = os.path.basename(database_path).replace(".db","")
+    database_name = os.path.basename(database_path).replace(".db", "")
     cwd = os.getcwd()
     alembic_location = os.path.join(cwd, "database", "databases", database_name.lower())
     db_helper.run_migrations(alembic_location, database_path)
@@ -1171,3 +1180,21 @@ def module_chooser(domain, json_sites):
         string = f"{domain} not supported"
         site_names = []
     return string, site_names
+
+
+async def move_to_old(
+    folder_directory: str,
+    base_download_directories: list,
+    first_letter: str,
+    model_username: str,
+    source: str,
+):
+    # MOVE TO OLD
+    local_destinations = [
+        os.path.join(x, folder_directory) for x in base_download_directories
+    ]
+    local_destination = check_space(local_destinations, min_size=100)
+    local_destination = os.path.join(local_destination, first_letter, model_username)
+    print(f"Moving {source} -> {local_destination}")
+    shutil.copytree(source, local_destination, dirs_exist_ok=True)
+    shutil.rmtree(source)

modules/onlyfans.py

@@ -10,6 +10,8 @@
 from typing import Any, Optional, Union
 from urllib.parse import urlparse
 
+from sqlalchemy.exc import OperationalError
+
 import extras.OFLogin.start_ofl as oflogin
 import extras.OFRenamer.start_ofr as ofrenamer
 import helpers.db_helper as db_helper
@@ -1244,7 +1246,12 @@ async def prepare_downloads(subscription: create_user):
             if media_set_count:
                 print(string)
                 await main_helper.async_downloads(download_list, subscription)
-            database_session.commit()
+            while True:
+                try:
+                    database_session.commit()
+                    break
+                except OperationalError:
+                    database_session.rollback()
             database_session.close()
         print
     print