Add progress bar for long running migrations (#1202)

yeti-platform · Dec 4, 2024 · 4b2ec71 · 4b2ec71
1 parent d071847
commit 4b2ec71
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 3 deletions.
diff --git a/core/migrations/arangodb.py b/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)

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/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"