From c98c3a3d5379890b33c29155de680eccd624b57c Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Wed, 23 Jul 2025 14:43:44 +0200 Subject: [PATCH] feat: Add kuzu migration script --- examples/kuzu_migrate.py | 161 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 161 insertions(+) create mode 100644 examples/kuzu_migrate.py diff --git a/examples/kuzu_migrate.py b/examples/kuzu_migrate.py new file mode 100644 index 000000000..bf76bcf9a --- /dev/null +++ b/examples/kuzu_migrate.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +""" +Kuzu Database Migration Script + +This script migrates Kuzu databases between different versions by: +1. Setting up isolated Python environments for each Kuzu version +2. Exporting data from the source database using the old version +3. Importing data into the target database using the new version +4. Handling edge cases like empty databases gracefully + +The script automatically handles: +- Empty databases (creates new database with standard Cognee schema) +- Environment setup (creates virtual environments as needed) +- Export/import validation +- Error handling and reporting + +Usage Examples: + # Basic migration from 0.9.0 to 0.11.0 + python kuzu_migrate.py --old-version 0.9.0 --new-version 0.11.0 \\ + --old-db /path/to/old/database --new-db /path/to/new/database + + # Migrate Cognee's default Kuzu database + python kuzu_migrate.py --old-version 0.9.0 --new-version 0.11.0 \\ + --old-db ~/.cognee_system/databases/cognee_graph_kuzu \\ + --new-db ~/.cognee_system/databases/cognee_graph + +Requirements: +- Python 3.7+ +- Internet connection (to download Kuzu packages) +- Sufficient disk space for virtual environments and temporary exports + +Author: Cognee Team +""" + +import tempfile +import sys +import subprocess +import argparse +import os + + +def ensure_env(version: str) -> str: + """ + Create (if needed) a venv at .kuzu_envs/{version} and install kuzu=={version}. + Returns the path to the venv's python executable. + """ + base = os.path.join(".kuzu_envs", version) + py_bin = os.path.join(base, "bin", "python") + if not os.path.isfile(py_bin): + print(f"→ Setting up venv for Kùzu {version}...", file=sys.stderr) + # Create venv + subprocess.run([sys.executable, "-m", "venv", base], check=True) + # Install the specific Kùzu version + subprocess.run([py_bin, "-m", "pip", "install", "--upgrade", "pip"], check=True) + subprocess.run([py_bin, "-m", "pip", "install", f"kuzu=={version}"], check=True) + return py_bin + + +def run_migration_step(python_exe: str, db_path: str, cypher: str): + """ + Uses the given python_exe to execute a short snippet that + connects to the Kùzu database and runs a Cypher command. + """ + snippet = f""" +import kuzu +db = kuzu.Database(r"{db_path}") +conn = kuzu.Connection(db) +conn.execute(r\"\"\"{cypher}\"\"\") +""" + proc = subprocess.run([python_exe, "-c", snippet], capture_output=True, text=True) + if proc.returncode != 0: + print(f"[ERROR] {cypher} failed:\n{proc.stderr}", file=sys.stderr) + sys.exit(proc.returncode) + + +def migrate(old_ver, new_ver, old_db, new_db): + """ + Main migration function that handles the complete migration process. + """ + # Check if old database exists + if not os.path.exists(old_db): + print(f"Source database '{old_db}' does not exist.", file=sys.stderr) + sys.exit(1) + + # Prepare target - ensure parent directory exists but remove target if it exists + parent_dir = os.path.dirname(new_db) + if parent_dir: + os.makedirs(parent_dir, exist_ok=True) + + # Remove target database if it exists (Kuzu 0.11.0 requires clean slate) + if os.path.exists(new_db): + if os.path.isdir(new_db): + import shutil + + shutil.rmtree(new_db) + print(f"Removed existing target directory: {new_db}", file=sys.stderr) + else: + os.remove(new_db) + print(f"Removed existing target file: {new_db}", file=sys.stderr) + + # Set up environments + print(f"Setting up Kuzu {old_ver} environment...", file=sys.stderr) + old_py = ensure_env(old_ver) + print(f"Setting up Kuzu {new_ver} environment...", file=sys.stderr) + new_py = ensure_env(new_ver) + + with tempfile.TemporaryDirectory() as export_dir: + export_file = os.path.join(export_dir, "kuzu_export") + print(f"Exporting old DB → {export_dir}", file=sys.stderr) + run_migration_step(old_py, old_db, f"EXPORT DATABASE '{export_file}'") + print("Export complete.", file=sys.stderr) + + # Check if export files were created and have content + schema_file = os.path.join(export_file, "schema.cypher") + if not os.path.exists(schema_file) or os.path.getsize(schema_file) == 0: + raise ValueError(f"Schema file not found: {schema_file}") + + print(f"Importing into new DB at {new_db}", file=sys.stderr) + run_migration_step(new_py, new_db, f"IMPORT DATABASE '{export_file}'") + print("Import complete.", file=sys.stderr) + + print("✅ Migration finished successfully!") + + +def main(): + p = argparse.ArgumentParser( + description="Migrate Kùzu DB via PyPI versions", + epilog=""" +Examples: + %(prog)s --old-version 0.9.0 --new-version 0.11.0 \\ + --old-db /path/to/old/db --new-db /path/to/new/db + + %(prog)s --old-version 0.9.0 --new-version 0.11.0 \\ + --old-db ~/.cognee_system/databases/cognee_graph_kuzu \\ + --new-db ~/.cognee_system/databases/cognee_graph + +Note: This script will create virtual environments in .kuzu_envs/ directory +to isolate different Kuzu versions. + """, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + p.add_argument("--old-version", required=True, help="Source Kuzu version (e.g., 0.9.0)") + p.add_argument("--new-version", required=True, help="Target Kuzu version (e.g., 0.11.0)") + p.add_argument("--old-db", required=True, help="Path to source database directory") + p.add_argument("--new-db", required=True, help="Path to target database directory") + + 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) + + migrate(args.old_version, args.new_version, args.old_db, args.new_db) + # migrate("0.9.0", "0.11.0", "/Users/igorilic/Desktop/cognee/cognee/.cognee_system/databases/cognee_graph_kuzu", "/Users/igorilic/Desktop/cognee/cognee/.cognee_system/databases/cognee_graph") + + +if __name__ == "__main__": + main()