Add progress bar for long running migrations

core/migrations/arangodb.py

@@ -1,5 +1,8 @@
+import logging
 import time
 
+import tqdm
+
 from core.database_arango import ASYNC_JOB_WAIT_TIME, ArangoDatabase
 from core.migrations import migration
 
@@ -51,8 +54,12 @@ def migration_0():
 def migration_1():
     from core.schemas import observable
 
-    for obs in observable.Observable.list():
-        obs.save()
+    total_observables = observable.Observable.count()
+    logging.info(f"Migrating {total_observables} observables. This may take a while...")
+    with tqdm.tqdm(total=total_observables, desc="Migrating observables") as pbar:
+        for obs in observable.Observable.list():
+            obs.save()
+            pbar.update(1)
 
 
 ArangoMigrationManager.register_migration(migration_0)

pyproject.toml

@@ -28,6 +28,7 @@ yara-python = "^4.5.0"
 idstools = "^0.6.5"
 aenum = "^3.1.15"
 boto3 = { version = "^1.35.22", optional = true }
+tqdm = "^4.67.1"
 
 [tool.poetry.group.dev.dependencies]
 pylint = "^2.16.1"