diff --git a/alembic/versions/b9274c27a25a_kuzu_11_migration.py b/alembic/versions/b9274c27a25a_kuzu_11_migration.py index fa1c0a8d5..7602cb4b5 100644 --- a/alembic/versions/b9274c27a25a_kuzu_11_migration.py +++ b/alembic/versions/b9274c27a25a_kuzu_11_migration.py @@ -9,6 +9,12 @@ Create Date: 2025-07-24 17:11:52.174737 import os from typing import Sequence, Union +from cognee.infrastructure.databases.graph.kuzu.kuzu_migrate import ( + kuzu_migration, + read_kuzu_storage_version, +) +import kuzu + # revision identifiers, used by Alembic. revision: str = "b9274c27a25a" down_revision: Union[str, None] = "e4ebee1091e7" @@ -18,34 +24,43 @@ depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: # This migration is only for multi-user Cognee mode - if not os.getenv("ENABLE_BACKEND_ACCESS_CONTROL", "false").lower() == "true": - return + if os.getenv("ENABLE_BACKEND_ACCESS_CONTROL", "false").lower() == "true": + from cognee.base_config import get_base_config - from cognee.base_config import get_base_config + base_config = get_base_config() - base_config = get_base_config() + databases_root = os.path.join(base_config.system_root_directory, "databases") + if not os.path.isdir(databases_root): + raise FileNotFoundError(f"Directory not found: {databases_root}") - databases_root = os.path.join(base_config.system_root_directory, "databases") - if not os.path.isdir(databases_root): - raise FileNotFoundError(f"Directory not found: {databases_root}") + for current_path, dirnames, _ in os.walk(databases_root): + # If file is kuzu graph database + if ".pkl" in current_path[-4:]: + kuzu_db_version = read_kuzu_storage_version(current_path) + if ( + kuzu_db_version == "0.9.0" or kuzu_db_version == "0.8.2" + ) and kuzu_db_version != kuzu.__version__: + # Try to migrate kuzu database to latest version + kuzu_migration( + new_db=current_path + "_new", + old_db=current_path, + new_version=kuzu.__version__, + old_version=kuzu_db_version, + overwrite=True, + ) + else: + from cognee.infrastructure.databases.graph import get_graph_config - for current_path, dirnames, _ in os.walk(databases_root): - # If file is kuzu graph database - if ".pkl" in current_path[-4:]: - from cognee.infrastructure.databases.graph.kuzu.kuzu_migrate import ( - kuzu_migration, - read_kuzu_storage_version, - ) - import kuzu - - kuzu_db_version = read_kuzu_storage_version(current_path) + graph_config = get_graph_config() + if graph_config.graph_database_provider.lower() == "kuzu": + kuzu_db_version = read_kuzu_storage_version(graph_config.graph_file_path) if ( kuzu_db_version == "0.9.0" or kuzu_db_version == "0.8.2" ) and kuzu_db_version != kuzu.__version__: # Try to migrate kuzu database to latest version kuzu_migration( - new_db=current_path + "new", - old_db=current_path, + new_db=graph_config.graph_file_path + "_new", + old_db=graph_config.graph_file_path, new_version=kuzu.__version__, old_version=kuzu_db_version, overwrite=True, diff --git a/cognee/infrastructure/databases/graph/kuzu/adapter.py b/cognee/infrastructure/databases/graph/kuzu/adapter.py index 5bc03ed36..4775e4b48 100644 --- a/cognee/infrastructure/databases/graph/kuzu/adapter.py +++ b/cognee/infrastructure/databases/graph/kuzu/adapter.py @@ -86,12 +86,11 @@ class KuzuAdapter(GraphDBInterface): if ( kuzu_db_version == "0.9.0" or kuzu_db_version == "0.8.2" ) and kuzu_db_version != kuzu.__version__: - # TODO: Write migration script that will handle all user graph databases in multi-user mode # Try to migrate kuzu database to latest version from .kuzu_migrate import kuzu_migration kuzu_migration( - new_db=self.db_path + "new", + new_db=self.db_path + "_new", old_db=self.db_path, new_version=kuzu.__version__, old_version=kuzu_db_version, @@ -1464,11 +1463,8 @@ class KuzuAdapter(GraphDBInterface): It raises exceptions for failures occurring during deletion processes. """ try: - # Use DETACH DELETE to remove both nodes and their relationships in one operation - await self.query("MATCH (n:Node) DETACH DELETE n") - logger.info("Cleared all data from graph while preserving structure") - if self.connection: + self.connection.close() self.connection = None if self.db: self.db.close() diff --git a/cognee/infrastructure/databases/graph/kuzu/kuzu_migrate.py b/cognee/infrastructure/databases/graph/kuzu/kuzu_migrate.py index e958374ad..3f930d91f 100644 --- a/cognee/infrastructure/databases/graph/kuzu/kuzu_migrate.py +++ b/cognee/infrastructure/databases/graph/kuzu/kuzu_migrate.py @@ -94,6 +94,7 @@ def ensure_env(version: str, export_dir) -> str: print(f"→ Setting up venv for Kùzu {version}...", file=sys.stderr) # Create venv + # NOTE: Running python in debug mode can cause issues with creating a virtual environment from that python instance 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)