diff --git a/cognee/modules/data/deletion/prune_system.py b/cognee/modules/data/deletion/prune_system.py index 66f93ec48..b5d944dbb 100644 --- a/cognee/modules/data/deletion/prune_system.py +++ b/cognee/modules/data/deletion/prune_system.py @@ -1,3 +1,5 @@ +from sqlalchemy.exc import OperationalError + from cognee.context_global_variables import backend_access_control_enabled from cognee.infrastructure.databases.vector import get_vector_engine from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_engine @@ -6,6 +8,9 @@ from cognee.infrastructure.databases.vector.config import get_vectordb_config from cognee.infrastructure.databases.graph.config import get_graph_config from cognee.shared.cache import delete_cache from cognee.modules.users.models import DatasetDatabase +from cognee.shared.logging_utils import get_logger + +logger = get_logger() async def prune_graph_databases(): @@ -19,11 +24,18 @@ async def prune_graph_databases(): return await handler["handler_instance"].delete_dataset(dataset_database) db_engine = get_relational_engine() - if "dataset_database" in await db_engine.get_table_names(): - data = await db_engine.get_all_data_from_table("dataset_database") - # Go through each dataset database and delete the graph database - for data_item in data: - await _prune_graph_db(data_item) + try: + if "dataset_database" in await db_engine.get_table_names(): + data = await db_engine.get_all_data_from_table("dataset_database") + # Go through each dataset database and delete the graph database + for data_item in data: + await _prune_graph_db(data_item) + except OperationalError as e: + logger.debug( + "Skipping pruning of graph DB. OperationalError when accessing dataset_database table: %s", + e, + ) + return async def prune_vector_databases(): @@ -38,11 +50,18 @@ async def prune_vector_databases(): return await handler["handler_instance"].delete_dataset(dataset_database) db_engine = get_relational_engine() - if "dataset_database" in await db_engine.get_table_names(): - data = await db_engine.get_all_data_from_table("dataset_database") - # Go through each dataset database and delete the vector database - for data_item in data: - await _prune_vector_db(data_item) + try: + if "dataset_database" in await db_engine.get_table_names(): + data = await db_engine.get_all_data_from_table("dataset_database") + # Go through each dataset database and delete the vector database + for data_item in data: + await _prune_vector_db(data_item) + except OperationalError as e: + logger.debug( + "Skipping pruning of vector DB. OperationalError when accessing dataset_database table: %s", + e, + ) + return async def prune_system(graph=True, vector=True, metadata=True, cache=True):