gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: Handle scenario when there is no relational database on prune time

Browse source
This commit is contained in:
Igor Ilic 2025-12-03 14:27:06 +01:00
parent f4078d1247
commit 1961efcc33
1 changed files with 29 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

39
cognee/modules/data/deletion/prune_system.py
View file

@ -1,3 +1,5 @@
from sqlalchemy.exc import OperationalError
from cognee.context_global_variables import backend_access_control_enabled from cognee.context_global_variables import backend_access_control_enabled
from cognee.infrastructure.databases.vector import get_vector_engine from cognee.infrastructure.databases.vector import get_vector_engine
from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_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.infrastructure.databases.graph.config import get_graph_config
from cognee.shared.cache import delete_cache from cognee.shared.cache import delete_cache
from cognee.modules.users.models import DatasetDatabase from cognee.modules.users.models import DatasetDatabase
from cognee.shared.logging_utils import get_logger
logger = get_logger()
async def prune_graph_databases(): async def prune_graph_databases():
@ -19,11 +24,18 @@ async def prune_graph_databases():
return await handler["handler_instance"].delete_dataset(dataset_database) return await handler["handler_instance"].delete_dataset(dataset_database)
db_engine = get_relational_engine() db_engine = get_relational_engine()
if "dataset_database" in await db_engine.get_table_names(): try:
data = await db_engine.get_all_data_from_table("dataset_database") if "dataset_database" in await db_engine.get_table_names():
# Go through each dataset database and delete the graph database data = await db_engine.get_all_data_from_table("dataset_database")
for data_item in data: # Go through each dataset database and delete the graph database
await _prune_graph_db(data_item) 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(): async def prune_vector_databases():
@ -38,11 +50,18 @@ async def prune_vector_databases():
return await handler["handler_instance"].delete_dataset(dataset_database) return await handler["handler_instance"].delete_dataset(dataset_database)
db_engine = get_relational_engine() db_engine = get_relational_engine()
if "dataset_database" in await db_engine.get_table_names(): try:
data = await db_engine.get_all_data_from_table("dataset_database") if "dataset_database" in await db_engine.get_table_names():
# Go through each dataset database and delete the vector database data = await db_engine.get_all_data_from_table("dataset_database")
for data_item in data: # Go through each dataset database and delete the vector database
await _prune_vector_db(data_item) 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): async def prune_system(graph=True, vector=True, metadata=True, cache=True):