From ed2d6871356ea94a5d95f9fcbeefc9c2ecd67348 Mon Sep 17 00:00:00 2001 From: Andrej Milicevic Date: Tue, 11 Nov 2025 13:52:34 +0100 Subject: [PATCH] fix: changes based on PR comments --- cognee/context_global_variables.py | 8 ++--- .../databases/utils/constants.py | 4 --- .../utils/get_or_create_dataset_database.py | 36 +++++-------------- 3 files changed, 12 insertions(+), 36 deletions(-) delete mode 100644 cognee/infrastructure/databases/utils/constants.py diff --git a/cognee/context_global_variables.py b/cognee/context_global_variables.py index 6ec467ed9..2d711a8b2 100644 --- a/cognee/context_global_variables.py +++ b/cognee/context_global_variables.py @@ -16,8 +16,8 @@ vector_db_config = ContextVar("vector_db_config", default=None) graph_db_config = ContextVar("graph_db_config", default=None) session_user = ContextVar("session_user", default=None) -vector_dbs_with_multi_user_support = ["lancedb"] -graph_dbs_with_multi_user_support = ["kuzu"] +VECTOR_DBS_WITH_MULTI_USER_SUPPORT = ["lancedb"] +GRAPH_DBS_WITH_MULTI_USER_SUPPORT = ["kuzu"] async def set_session_user_context_variable(user): @@ -28,8 +28,8 @@ def multi_user_support_possible(): graph_db_config = get_graph_context_config() vector_db_config = get_vectordb_context_config() return ( - graph_db_config["graph_database_provider"] in graph_dbs_with_multi_user_support - and vector_db_config["vector_db_provider"] in vector_dbs_with_multi_user_support + graph_db_config["graph_database_provider"] in GRAPH_DBS_WITH_MULTI_USER_SUPPORT + and vector_db_config["vector_db_provider"] in VECTOR_DBS_WITH_MULTI_USER_SUPPORT ) diff --git a/cognee/infrastructure/databases/utils/constants.py b/cognee/infrastructure/databases/utils/constants.py deleted file mode 100644 index fe6390a07..000000000 --- a/cognee/infrastructure/databases/utils/constants.py +++ /dev/null @@ -1,4 +0,0 @@ -VECTOR_DBS_WITH_MULTI_USER_SUPPORT = ["lancedb", "falkor"] -GRAPH_DBS_WITH_MULTI_USER_SUPPORT = ["kuzu", "falkor"] - -HYBRID_DBS = ["falkor"] diff --git a/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py b/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py index a4e50f665..61d7840c0 100644 --- a/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py +++ b/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py @@ -11,12 +11,6 @@ from cognee.infrastructure.databases.graph.config import get_graph_config from cognee.modules.data.methods import get_unique_dataset_id from cognee.modules.users.models import DatasetDatabase from cognee.modules.users.models import User -from .constants import ( - GRAPH_DBS_WITH_MULTI_USER_SUPPORT, - VECTOR_DBS_WITH_MULTI_USER_SUPPORT, - HYBRID_DBS, -) - async def get_or_create_dataset_database( dataset: Union[str, UUID], @@ -42,15 +36,15 @@ async def get_or_create_dataset_database( vector_config = get_vectordb_config() graph_config = get_graph_config() - graph_db_name = f"{dataset_id}.pkl" - - if graph_config.graph_database_provider in HYBRID_DBS: - vector_db_name = graph_db_name + if graph_config.graph_database_provider == "kuzu": + graph_db_name = f"{dataset_id}.pkl" else: - if vector_config.vector_db_provider == "lancedb": - vector_db_name = f"{dataset_id}.lance.db" - else: - vector_db_name = f"{dataset_id}.db" + graph_db_name = dataset_id + + if vector_config.vector_db_provider == "lancedb": + vector_db_name = f"{dataset_id}.lance.db" + else: + vector_db_name = dataset_id async with db_engine.get_async_session() as session: # Create dataset if it doesn't exist @@ -66,20 +60,6 @@ async def get_or_create_dataset_database( if existing: return existing - # Check if we support multi-user for this provider. If not, use default - if graph_config.graph_database_provider not in GRAPH_DBS_WITH_MULTI_USER_SUPPORT: - raise EnvironmentError( - f"Multi-user is currently not supported for the graph database provider: {graph_config.graph_database_provider}. " - f"Supported providers are: {', '.join(GRAPH_DBS_WITH_MULTI_USER_SUPPORT)}. Either use one of these" - f"providers, or disable BACKEND_ACCESS_CONTROL" - ) - if vector_config.vector_db_provider not in VECTOR_DBS_WITH_MULTI_USER_SUPPORT: - raise EnvironmentError( - f"Multi-user is currently not supported for the vector database provider: {vector_config.vector_db_provider}. " - f"Supported providers are: {', '.join(VECTOR_DBS_WITH_MULTI_USER_SUPPORT)}. Either use one of these" - f"providers, or disable BACKEND_ACCESS_CONTROL" - ) - # If there are no existing rows build a new row record = DatasetDatabase( owner_id=user.id,