From 2fff637cda0cb4228a6c855321ece75b13d18128 Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Thu, 24 Jul 2025 12:38:31 +0200
Subject: [PATCH] feat: Add better info when migrating kuzu database

---
 .../databases/graph/kuzu/adapter.py           |  2 +-
 .../databases/graph/kuzu/kuzu_migrate.py      | 21 +++++++++----------
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/cognee/infrastructure/databases/graph/kuzu/adapter.py b/cognee/infrastructure/databases/graph/kuzu/adapter.py
index f277e3cba..386a06dc7 100644
--- a/cognee/infrastructure/databases/graph/kuzu/adapter.py
+++ b/cognee/infrastructure/databases/graph/kuzu/adapter.py
@@ -91,9 +91,9 @@ class KuzuAdapter(GraphDBInterface):
                         kuzu_migration(
                             new_db=self.db_path + "new",
                             old_db=self.db_path,
-                            overwrite=True,
                             new_version=kuzu.__version__,
                             old_version=kuzu_db_version,
+                            overwrite=True,
                         )
                         self.db = Database(
                             self.db_path,
diff --git a/cognee/infrastructure/databases/graph/kuzu/kuzu_migrate.py b/cognee/infrastructure/databases/graph/kuzu/kuzu_migrate.py
index 62ae6f332..53eeb70d0 100644
--- a/cognee/infrastructure/databases/graph/kuzu/kuzu_migrate.py
+++ b/cognee/infrastructure/databases/graph/kuzu/kuzu_migrate.py
@@ -118,10 +118,14 @@ conn.execute(r\"\"\"{cypher}\"\"\")
         sys.exit(proc.returncode)
 
 
-def kuzu_migration(new_db, old_db, overwrite, delete_old, new_version, old_version=None):
+def kuzu_migration(new_db, old_db, new_version, old_version=None, overwrite=None, delete_old=None):
     """
     Main migration function that handles the complete migration process.
     """
+    print(f"🔄 Migrating Kuzu database from {old_version} to {new_version}", file=sys.stderr)
+    print(f"📂 Source: {old_db}", file=sys.stderr)
+    print("", file=sys.stderr)
+
     # If version of old kuzu db is not provided try to determine it based on file info
     if not old_version:
         old_version = read_kuzu_storage_version(old_db)
@@ -167,7 +171,7 @@ def kuzu_migration(new_db, old_db, overwrite, delete_old, new_version, old_versi
     if overwrite or delete_old:
         rename_databases(old_db, old_version, new_db, delete_old)
 
-    print("✅ Migration finished successfully!")
+    print("✅ Kuzu graph database migration finished successfully!")
 
 
 def rename_databases(old_db: str, old_version: str, new_db: str, delete_old: bool):
@@ -243,12 +247,14 @@ to isolate different Kuzu versions.
     )
     p.add_argument(
         "--overwrite",
+        required=False,
         action="store_true",
         default=False,
         help="Rename new-db to the old-db name and location, keeps old-db as backup if delete-old is not True",
     )
     p.add_argument(
         "--delete-old",
+        required=False,
         action="store_true",
         default=False,
         help="When overwrite and delete-old is True old-db will not be stored as backup",
@@ -256,20 +262,13 @@ to isolate different Kuzu versions.
 
     args = p.parse_args()
 
-    print(
-        f"🔄 Migrating Kuzu database from {args.old_version} to {args.new_version}", file=sys.stderr
-    )
-    print(f"📂 Source: {args.old_db}", file=sys.stderr)
-    print(f"📂 Target: {args.new_db}", file=sys.stderr)
-    print("", file=sys.stderr)
-
     kuzu_migration(
         new_db=args.new_db,
         old_db=args.old_db,
-        overwrite=args.overwrite,
-        delete_old=args.delete_old,
         new_version=args.new_version,
         old_version=args.old_version,
+        overwrite=args.overwrite,
+        delete_old=args.delete_old,
     )