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_library.py b/cognee/tests/test_library.py index 66d218c3b..6c9d41800 100755 --- a/cognee/tests/test_library.py +++ b/cognee/tests/test_library.py @@ -57,6 +57,24 @@ 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, "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), "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 asyncio.run(main(), debug=True) 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()) diff --git a/cognee/tests/test_pgvector.py b/cognee/tests/test_pgvector.py index 1466e195f..bd6584cbc 100644 --- a/cognee/tests/test_pgvector.py +++ b/cognee/tests/test_pgvector.py @@ -87,9 +87,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) + tables_in_database = await vector_engine.get_table_names() + assert len(tables_in_database) == 0, "PostgreSQL database is not empty" + if __name__ == "__main__": import asyncio 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()) 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())