From 5aca3f091e3f740ece70ef08ec153339d96efce3 Mon Sep 17 00:00:00 2001 From: Vasilije <8619304+Vasilije1990@users.noreply.github.com> Date: Sat, 26 Apr 2025 00:00:45 +0200 Subject: [PATCH] fix: Doesn't drop entire PG database, just cleans public schema - Cog 1947 (#760) ## Description ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin. --------- Co-authored-by: Boris Co-authored-by: Igor Ilic <30923996+dexters1@users.noreply.github.com> Co-authored-by: Igor Ilic --- .../databases/relational/sqlalchemy/SqlAlchemyAdapter.py | 6 +++--- .../databases/vector/exceptions/exceptions.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py index f20904ba7..ade2a0821 100644 --- a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py +++ b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py @@ -69,7 +69,7 @@ class SQLAlchemyAdapter: async def delete_table(self, table_name: str, schema_name: Optional[str] = "public"): async with self.engine.begin() as connection: if self.engine.dialect.name == "sqlite": - # SQLite doesn’t support schema namespaces and the CASCADE keyword. + # SQLite doesn't support schema namespaces and the CASCADE keyword. # However, foreign key constraint can be defined with ON DELETE CASCADE during table creation. await connection.execute(text(f'DROP TABLE IF EXISTS "{table_name}";')) else: @@ -327,10 +327,10 @@ class SQLAlchemyAdapter: file.write("") else: async with self.engine.begin() as connection: - schema_list = await self.get_schema_list() # Create a MetaData instance to load table information metadata = MetaData() - # Drop all tables from all schemas + # Drop all tables from the public schema + schema_list = ["public", "public_staging"] for schema_name in schema_list: # Load the schema information into the MetaData object await connection.run_sync(metadata.reflect, schema=schema_name) diff --git a/cognee/infrastructure/databases/vector/exceptions/exceptions.py b/cognee/infrastructure/databases/vector/exceptions/exceptions.py index 8f7ff17cc..4ad14f716 100644 --- a/cognee/infrastructure/databases/vector/exceptions/exceptions.py +++ b/cognee/infrastructure/databases/vector/exceptions/exceptions.py @@ -6,7 +6,7 @@ class CollectionNotFoundError(CriticalError): def __init__( self, message, - name: str = "DatabaseNotCreatedError", + name: str = "CollectionNotFoundError", status_code: int = status.HTTP_422_UNPROCESSABLE_ENTITY, ): super().__init__(message, name, status_code)