From b60f2603f4632c1d678d5677a07ba8c8eab4ddc7 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Wed, 20 Nov 2024 17:11:23 +0100 Subject: [PATCH 1/8] test: Add test for pgvector to confirm database deletion is working Added assert to verify all tables in database have been cleared. Added method to SqlAlchemyAdapter to get all table names in database. Test COG-488 --- .../sqlalchemy/SqlAlchemyAdapter.py | 23 +++++++++++++++++++ cognee/tests/test_pgvector.py | 6 ++++- 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py index febfe1931..aa2a022d3 100644 --- a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py +++ b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py @@ -130,6 +130,29 @@ class SQLAlchemyAdapter(): return metadata.tables[full_table_name] raise ValueError(f"Table '{full_table_name}' not found.") + async def get_table_names(self) -> List[str]: + """ + Return a list of all tables names in database + """ + table_names = [] + async with self.engine.begin() as connection: + if self.engine.dialect.name == "sqlite": + await connection.run_sync(Base.metadata.reflect) + for table in Base.metadata.tables: + table_names.append(str(table)) + else: + schema_list = await self.get_schema_list() + # Create a MetaData instance to load table information + metadata = MetaData() + # Drop all tables from all schemas + for schema_name in schema_list: + # Load the schema information into the MetaData object + await connection.run_sync(metadata.reflect, schema=schema_name) + for table in metadata.sorted_tables: + table_names.append(str(table)) + metadata.clear() + return table_names + async def get_data(self, table_name: str, filters: dict = None): async with self.engine.begin() as connection: diff --git a/cognee/tests/test_pgvector.py b/cognee/tests/test_pgvector.py index 1466e195f..18d9e2c91 100644 --- a/cognee/tests/test_pgvector.py +++ b/cognee/tests/test_pgvector.py @@ -87,9 +87,13 @@ async def main(): print(f"{result}\n") history = await cognee.get_search_history() - assert len(history) == 6, "Search history is not correct." + await cognee.prune.prune_data() + await cognee.prune.prune_system(metadata=True) + tables_in_database = await vector_engine.get_table_names() + assert len(tables_in_database) == 0, "The database is not empty" + if __name__ == "__main__": import asyncio From 796bbadb43a977b97afdc2514ccdd9185962033a Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Wed, 20 Nov 2024 17:27:12 +0100 Subject: [PATCH 2/8] test: Add verification of deletion of local files Verify local files used in end to end test are deleted upon completion Test COG-488 --- cognee/tests/test_pgvector.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cognee/tests/test_pgvector.py b/cognee/tests/test_pgvector.py index 18d9e2c91..7e52f7740 100644 --- a/cognee/tests/test_pgvector.py +++ b/cognee/tests/test_pgvector.py @@ -90,6 +90,8 @@ async def main(): assert len(history) == 6, "Search history is not correct." await cognee.prune.prune_data() + assert not os.path.isdir(data_directory_path), "Local files are not deleted" + await cognee.prune.prune_system(metadata=True) tables_in_database = await vector_engine.get_table_names() assert len(tables_in_database) == 0, "The database is not empty" From 27aade2a4394ee4f7fcbfaf8629b3faf9e99637d Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Thu, 21 Nov 2024 10:25:26 +0100 Subject: [PATCH 3/8] refactor: Expand on assert in test Expand on assert fail message in pgvector test Refactor COG-488 --- cognee/tests/test_pgvector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cognee/tests/test_pgvector.py b/cognee/tests/test_pgvector.py index 7e52f7740..c9c841d31 100644 --- a/cognee/tests/test_pgvector.py +++ b/cognee/tests/test_pgvector.py @@ -90,7 +90,7 @@ async def main(): assert len(history) == 6, "Search history is not correct." await cognee.prune.prune_data() - assert not os.path.isdir(data_directory_path), "Local files are not deleted" + assert not os.path.isdir(data_directory_path), "Local data files are not deleted" await cognee.prune.prune_system(metadata=True) tables_in_database = await vector_engine.get_table_names() From b815aeaa2afa5ccd13ff4d06194a154a0c787317 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Thu, 21 Nov 2024 13:44:09 +0100 Subject: [PATCH 4/8] test: Add test for database deletion for LanceDB and SQLite Added database deletion test for LanceDB and SQLite Test COG-488 --- cognee/tests/test_library.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/cognee/tests/test_library.py b/cognee/tests/test_library.py index 66d218c3b..9a49206a1 100755 --- a/cognee/tests/test_library.py +++ b/cognee/tests/test_library.py @@ -57,6 +57,16 @@ async def main(): assert len(history) == 6, "Search history is not correct." + await cognee.prune.prune_data() + assert not os.path.isdir(data_directory_path), "Local data files are not deleted" + + await cognee.prune.prune_system(metadata=True) + connection = await vector_engine.get_connection() + collection_names = await connection.table_names() + assert len(collection_names) == 0, "The vector database is not empty" + from cognee.infrastructure.databases.relational import get_relational_engine + assert not os.path.exists(get_relational_engine().db_path), "The relational database is not empty" + if __name__ == "__main__": import asyncio asyncio.run(main(), debug=True) From ac3f9882787367aeb78306182e24fd0c0f0607f0 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Thu, 21 Nov 2024 16:09:50 +0100 Subject: [PATCH 5/8] test: Add checking of Networkx database deletion Add test to check if networkx database has been cleaned properly Test COG-488 --- cognee/tests/test_library.py | 12 ++++++++++-- cognee/tests/test_pgvector.py | 2 +- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/cognee/tests/test_library.py b/cognee/tests/test_library.py index 9a49206a1..6c9d41800 100755 --- a/cognee/tests/test_library.py +++ b/cognee/tests/test_library.py @@ -57,15 +57,23 @@ async def main(): assert len(history) == 6, "Search history is not correct." + # Assert local data files are cleaned properly await cognee.prune.prune_data() assert not os.path.isdir(data_directory_path), "Local data files are not deleted" + # Assert relational, vector and graph databases have been cleaned properly await cognee.prune.prune_system(metadata=True) + connection = await vector_engine.get_connection() collection_names = await connection.table_names() - assert len(collection_names) == 0, "The vector database is not empty" + assert len(collection_names) == 0, "LanceDB vector database is not empty" + from cognee.infrastructure.databases.relational import get_relational_engine - assert not os.path.exists(get_relational_engine().db_path), "The relational database is not empty" + assert not os.path.exists(get_relational_engine().db_path), "SQLite relational database is not empty" + + from cognee.infrastructure.databases.graph import get_graph_config + graph_config = get_graph_config() + assert not os.path.exists(graph_config.graph_file_path), "Networkx graph database is not empty" if __name__ == "__main__": import asyncio diff --git a/cognee/tests/test_pgvector.py b/cognee/tests/test_pgvector.py index c9c841d31..bd6584cbc 100644 --- a/cognee/tests/test_pgvector.py +++ b/cognee/tests/test_pgvector.py @@ -94,7 +94,7 @@ async def main(): await cognee.prune.prune_system(metadata=True) tables_in_database = await vector_engine.get_table_names() - assert len(tables_in_database) == 0, "The database is not empty" + assert len(tables_in_database) == 0, "PostgreSQL database is not empty" if __name__ == "__main__": import asyncio From 54daa6986cc3d1a884696958ac2fc83a1cbd0c24 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Thu, 21 Nov 2024 16:39:38 +0100 Subject: [PATCH 6/8] test: Add test for deletion of neo4j graph database Added test to verify all nodes and edges from neo4j database have been cleaned Test COG-488 --- cognee/tests/test_neo4j.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/cognee/tests/test_neo4j.py b/cognee/tests/test_neo4j.py index 756b29cc4..02f3eaccd 100644 --- a/cognee/tests/test_neo4j.py +++ b/cognee/tests/test_neo4j.py @@ -61,6 +61,15 @@ async def main(): assert len(history) == 6, "Search history is not correct." + await cognee.prune.prune_data() + assert not os.path.isdir(data_directory_path), "Local data files are not deleted" + + await cognee.prune.prune_system(metadata=True) + from cognee.infrastructure.databases.graph import get_graph_engine + graph_engine = await get_graph_engine() + nodes, edges = await graph_engine.get_graph_data() + assert len(nodes) == 0 and len(edges) == 0, "Neo4j graph database is not empty" + if __name__ == "__main__": import asyncio asyncio.run(main()) From 0f8baeeb0d25a7015a2d313d683863e3c69645e6 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Thu, 21 Nov 2024 18:19:23 +0100 Subject: [PATCH 7/8] test: Add database deletion test for qdrant Added testing of database deletion for qdrant Test COG-488 --- cognee/tests/test_qdrant.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/cognee/tests/test_qdrant.py b/cognee/tests/test_qdrant.py index 680399e60..4c2462c3b 100644 --- a/cognee/tests/test_qdrant.py +++ b/cognee/tests/test_qdrant.py @@ -59,9 +59,16 @@ async def main(): print(f"{result}\n") history = await cognee.get_search_history() - assert len(history) == 6, "Search history is not correct." + await cognee.prune.prune_data() + assert not os.path.isdir(data_directory_path), "Local data files are not deleted" + + await cognee.prune.prune_system(metadata=True) + qdrant_client = get_vector_engine().get_qdrant_client() + collections_response = await qdrant_client.get_collections() + assert len(collections_response.collections) == 0, "QDrant vector database is not empty" + if __name__ == "__main__": import asyncio asyncio.run(main()) From 443133ffbbfa97c2c58f252f41aab8d936af62ce Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Thu, 21 Nov 2024 18:37:09 +0100 Subject: [PATCH 8/8] test: Add database deletion test for Weaviate vector database Added database deletion test for Weaviate vector database Test COG-488 --- cognee/tests/test_weaviate.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/cognee/tests/test_weaviate.py b/cognee/tests/test_weaviate.py index c93dc036a..c352df13e 100644 --- a/cognee/tests/test_weaviate.py +++ b/cognee/tests/test_weaviate.py @@ -57,9 +57,15 @@ async def main(): print(f"{result}\n") history = await cognee.get_search_history() - assert len(history) == 6, "Search history is not correct." + await cognee.prune.prune_data() + assert not os.path.isdir(data_directory_path), "Local data files are not deleted" + + await cognee.prune.prune_system(metadata=True) + collections = get_vector_engine().client.collections.list_all() + assert len(collections) == 0, "Weaviate vector database is not empty" + if __name__ == "__main__": import asyncio asyncio.run(main())