fix: Resolve DB caching issues when deleting databases

This commit is contained in:
Igor Ilic 2025-12-03 18:05:47 +01:00
parent fd84edeb74
commit f1c5b9a55f
3 changed files with 11 additions and 11 deletions

View file

@ -121,6 +121,8 @@ async def set_database_global_context_variables(dataset: Union[str, UUID], user_
) )
# Set vector and graph database configuration based on dataset database information # Set vector and graph database configuration based on dataset database information
# TODO: Add better handling of vector and graph config accross Cognee.
# LRU_CACHE takes into account order of inputs, if order of inputs is changed it will be registered as a new DB adapter
vector_config = { vector_config = {
"vector_db_provider": dataset_database.vector_database_provider, "vector_db_provider": dataset_database.vector_database_provider,
"vector_db_url": dataset_database.vector_database_url, "vector_db_url": dataset_database.vector_database_url,
@ -142,6 +144,8 @@ async def set_database_global_context_variables(dataset: Union[str, UUID], user_
"graph_database_password": dataset_database.graph_database_connection_info.get( "graph_database_password": dataset_database.graph_database_connection_info.get(
"graph_database_password", "" "graph_database_password", ""
), ),
"graph_dataset_database_handler": "",
"graph_database_port": "",
} }
storage_config = { storage_config = {

View file

@ -6,7 +6,6 @@ from cognee.infrastructure.databases.graph.get_graph_engine import create_graph_
from cognee.base_config import get_base_config from cognee.base_config import get_base_config
from cognee.modules.users.models import User from cognee.modules.users.models import User
from cognee.modules.users.models import DatasetDatabase from cognee.modules.users.models import DatasetDatabase
from cognee.infrastructure.databases.graph.config import get_graph_config
from cognee.infrastructure.databases.dataset_database_handler import DatasetDatabaseHandlerInterface from cognee.infrastructure.databases.dataset_database_handler import DatasetDatabaseHandlerInterface
@ -56,7 +55,6 @@ class KuzuDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
@classmethod @classmethod
async def delete_dataset(cls, dataset_database: DatasetDatabase): async def delete_dataset(cls, dataset_database: DatasetDatabase):
graph_config = get_graph_config()
base_config = get_base_config() base_config = get_base_config()
databases_directory_path = os.path.join( databases_directory_path = os.path.join(
base_config.system_root_directory, "databases", str(dataset_database.owner_id) base_config.system_root_directory, "databases", str(dataset_database.owner_id)
@ -66,16 +64,17 @@ class KuzuDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
) )
graph_engine = create_graph_engine( graph_engine = create_graph_engine(
graph_database_provider=dataset_database.graph_database_provider, graph_database_provider=dataset_database.graph_database_provider,
graph_file_path=graph_file_path,
graph_database_url=dataset_database.graph_database_url, graph_database_url=dataset_database.graph_database_url,
graph_database_name=dataset_database.graph_database_name, graph_database_name=dataset_database.graph_database_name,
graph_database_key=dataset_database.graph_database_key,
graph_file_path=graph_file_path,
graph_database_username=dataset_database.graph_database_connection_info.get( graph_database_username=dataset_database.graph_database_connection_info.get(
"graph_database_username", "" "graph_database_username", ""
), ),
graph_database_password=dataset_database.graph_database_connection_info.get( graph_database_password=dataset_database.graph_database_connection_info.get(
"graph_database_password", "" "graph_database_password", ""
), ),
graph_database_port=graph_config.graph_database_port, graph_dataset_database_handler="",
graph_database_key=dataset_database.graph_database_key, graph_database_port="",
) )
await graph_engine.delete_graph() await graph_engine.delete_graph()

View file

@ -32,21 +32,18 @@ class LanceDBDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
vector_db_name = f"{dataset_id}.lance.db" vector_db_name = f"{dataset_id}.lance.db"
return { return {
"vector_database_name": vector_db_name,
"vector_database_url": os.path.join(databases_directory_path, vector_db_name),
"vector_database_provider": vector_config.vector_db_provider, "vector_database_provider": vector_config.vector_db_provider,
"vector_database_url": os.path.join(databases_directory_path, vector_db_name),
"vector_database_key": vector_config.vector_db_key, "vector_database_key": vector_config.vector_db_key,
"vector_database_name": vector_db_name,
} }
@classmethod @classmethod
async def delete_dataset(cls, dataset_database: DatasetDatabase): async def delete_dataset(cls, dataset_database: DatasetDatabase):
vector_config = get_vectordb_config()
vector_engine = create_vector_engine( vector_engine = create_vector_engine(
vector_db_provider=dataset_database.vector_database_provider, vector_db_provider=dataset_database.vector_database_provider,
vector_db_url=dataset_database.vector_database_url, vector_db_url=dataset_database.vector_database_url,
vector_db_name=dataset_database.vector_database_name,
vector_db_port=vector_config.vector_db_port,
vector_db_key=dataset_database.vector_database_key, vector_db_key=dataset_database.vector_database_key,
vector_dataset_database_handler=vector_config.vector_dataset_database_handler, vector_db_name=dataset_database.vector_database_name,
) )
await vector_engine.prune() await vector_engine.prune()