diff --git a/cognee/infrastructure/databases/utils/__init__.py b/cognee/infrastructure/databases/utils/__init__.py index f31d1e0dc..3907b4325 100644 --- a/cognee/infrastructure/databases/utils/__init__.py +++ b/cognee/infrastructure/databases/utils/__init__.py @@ -1,2 +1,4 @@ from .get_or_create_dataset_database import get_or_create_dataset_database from .resolve_dataset_database_connection_info import resolve_dataset_database_connection_info +from .get_graph_dataset_database_handler import get_graph_dataset_database_handler +from .get_vector_dataset_database_handler import get_vector_dataset_database_handler diff --git a/cognee/infrastructure/databases/utils/get_graph_dataset_database_handler.py b/cognee/infrastructure/databases/utils/get_graph_dataset_database_handler.py new file mode 100644 index 000000000..d88685b48 --- /dev/null +++ b/cognee/infrastructure/databases/utils/get_graph_dataset_database_handler.py @@ -0,0 +1,10 @@ +from cognee.modules.users.models.DatasetDatabase import DatasetDatabase + + +def get_graph_dataset_database_handler(dataset_database: DatasetDatabase) -> dict: + from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import ( + supported_dataset_database_handlers, + ) + + handler = supported_dataset_database_handlers[dataset_database.graph_dataset_database_handler] + return handler diff --git a/cognee/infrastructure/databases/utils/get_vector_dataset_database_handler.py b/cognee/infrastructure/databases/utils/get_vector_dataset_database_handler.py new file mode 100644 index 000000000..5d1152c04 --- /dev/null +++ b/cognee/infrastructure/databases/utils/get_vector_dataset_database_handler.py @@ -0,0 +1,10 @@ +from cognee.modules.users.models.DatasetDatabase import DatasetDatabase + + +def get_vector_dataset_database_handler(dataset_database: DatasetDatabase) -> dict: + from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import ( + supported_dataset_database_handlers, + ) + + handler = supported_dataset_database_handlers[dataset_database.vector_dataset_database_handler] + return handler diff --git a/cognee/infrastructure/databases/utils/resolve_dataset_database_connection_info.py b/cognee/infrastructure/databases/utils/resolve_dataset_database_connection_info.py index d33169642..561268eaf 100644 --- a/cognee/infrastructure/databases/utils/resolve_dataset_database_connection_info.py +++ b/cognee/infrastructure/databases/utils/resolve_dataset_database_connection_info.py @@ -1,24 +1,12 @@ +from cognee.infrastructure.databases.utils.get_graph_dataset_database_handler import ( + get_graph_dataset_database_handler, +) +from cognee.infrastructure.databases.utils.get_vector_dataset_database_handler import ( + get_vector_dataset_database_handler, +) from cognee.modules.users.models.DatasetDatabase import DatasetDatabase -async def _get_vector_db_connection_info(dataset_database: DatasetDatabase) -> DatasetDatabase: - from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import ( - supported_dataset_database_handlers, - ) - - handler = supported_dataset_database_handlers[dataset_database.vector_dataset_database_handler] - return await handler["handler_instance"].resolve_dataset_connection_info(dataset_database) - - -async def _get_graph_db_connection_info(dataset_database: DatasetDatabase) -> DatasetDatabase: - from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import ( - supported_dataset_database_handlers, - ) - - handler = supported_dataset_database_handlers[dataset_database.graph_dataset_database_handler] - return await handler["handler_instance"].resolve_dataset_connection_info(dataset_database) - - async def resolve_dataset_database_connection_info( dataset_database: DatasetDatabase, ) -> DatasetDatabase: @@ -31,6 +19,12 @@ async def resolve_dataset_database_connection_info( Returns: DatasetDatabase instance with resolved connection info """ - dataset_database = await _get_vector_db_connection_info(dataset_database) - dataset_database = await _get_graph_db_connection_info(dataset_database) + vector_dataset_database_handler = get_vector_dataset_database_handler(dataset_database) + graph_dataset_database_handler = get_graph_dataset_database_handler(dataset_database) + dataset_database = await vector_dataset_database_handler[ + "handler_instance" + ].resolve_dataset_connection_info(dataset_database) + dataset_database = await graph_dataset_database_handler[ + "handler_instance" + ].resolve_dataset_connection_info(dataset_database) return dataset_database diff --git a/cognee/modules/data/deletion/prune_system.py b/cognee/modules/data/deletion/prune_system.py index 645e1a223..22a0fde5f 100644 --- a/cognee/modules/data/deletion/prune_system.py +++ b/cognee/modules/data/deletion/prune_system.py @@ -5,6 +5,10 @@ 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 from cognee.infrastructure.databases.relational import get_relational_engine +from cognee.infrastructure.databases.utils import ( + get_graph_dataset_database_handler, + get_vector_dataset_database_handler, +) from cognee.shared.cache import delete_cache from cognee.modules.users.models import DatasetDatabase from cognee.shared.logging_utils import get_logger @@ -13,22 +17,13 @@ logger = get_logger() async def prune_graph_databases(): - async def _prune_graph_db(dataset_database: DatasetDatabase) -> dict: - from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import ( - supported_dataset_database_handlers, - ) - - handler = supported_dataset_database_handlers[ - dataset_database.graph_dataset_database_handler - ] - return await handler["handler_instance"].delete_dataset(dataset_database) - db_engine = get_relational_engine() try: - data = await db_engine.get_all_data_from_table("dataset_database") + dataset_databases = 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) + for dataset_database in dataset_databases: + handler = get_graph_dataset_database_handler(dataset_database) + await handler["handler_instance"].delete_dataset(dataset_database) except (OperationalError, EntityNotFoundError) as e: logger.debug( "Skipping pruning of graph DB. Error when accessing dataset_database table: %s", @@ -38,22 +33,13 @@ async def prune_graph_databases(): async def prune_vector_databases(): - async def _prune_vector_db(dataset_database: DatasetDatabase) -> dict: - from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import ( - supported_dataset_database_handlers, - ) - - handler = supported_dataset_database_handlers[ - dataset_database.vector_dataset_database_handler - ] - return await handler["handler_instance"].delete_dataset(dataset_database) - db_engine = get_relational_engine() try: - data = await db_engine.get_all_data_from_table("dataset_database") + dataset_databases = 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) + for dataset_database in dataset_databases: + handler = get_vector_dataset_database_handler(dataset_database) + await handler["handler_instance"].delete_dataset(dataset_database) except (OperationalError, EntityNotFoundError) as e: logger.debug( "Skipping pruning of vector DB. Error when accessing dataset_database table: %s",