Kuzu migration notes (#1149)

<!-- .github/pull_request_template.md -->

## Description
<!-- Provide a clear description of the changes in this PR -->

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin.
This commit is contained in:
Igor Ilic 2025-07-25 14:38:10 +02:00 committed by GitHub
parent 9157d3c2dd
commit 93e9a9aa1c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 37 additions and 25 deletions

View file

@ -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,

View file

@ -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()

View file

@ -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)