From 05e4ef349e82fc385ac53fab663e72a697e29cf9 Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Mon, 21 Oct 2024 14:28:57 +0200
Subject: [PATCH] fix: Fix pruning of postgres database

Fixed issue with newly added tables not being pruned from postgres database

Fix #COG-170
---
 cognee/api/v1/config/config.py                | 35 +++++++++++++++++++
 .../sqlalchemy/SqlAlchemyAdapter.py           |  2 ++
 .../vector/pgvector/PGVectorAdapter.py        |  3 +-
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/cognee/api/v1/config/config.py b/cognee/api/v1/config/config.py
index 7d2c7c6f4..d5acfd08e 100644
--- a/cognee/api/v1/config/config.py
+++ b/cognee/api/v1/config/config.py
@@ -39,6 +39,41 @@ class config():
         cognify_config = get_cognify_config()
         cognify_config.classification_model =  classification_model
 
+    @staticmethod
+    def set_db_name(db_name: str):
+        cognify_config = get_relational_config()
+        cognify_config.db_name = db_name
+
+    @staticmethod
+    def set_db_path(db_path: str):
+        cognify_config = get_relational_config()
+        cognify_config.db_path = db_path
+
+    @staticmethod
+    def set_db_host(db_host: str):
+        cognify_config = get_relational_config()
+        cognify_config.db_host = db_host
+
+    @staticmethod
+    def set_db_port(db_port: str):
+        cognify_config = get_relational_config()
+        cognify_config.db_port = db_port
+
+    @staticmethod
+    def set_db_username(db_username: str):
+        cognify_config = get_relational_config()
+        cognify_config.db_username = db_username
+
+    @staticmethod
+    def set_db_password(db_password: str):
+        cognify_config = get_relational_config()
+        cognify_config.db_password = db_password
+
+    @staticmethod
+    def set_db_provider(db_provider: str):
+        cognify_config = get_relational_config()
+        cognify_config.db_provider = db_provider
+
     @staticmethod
     def set_summarization_model(summarization_model: object):
         cognify_config = get_cognify_config()
diff --git a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
index 36302bce8..81a828bd8 100644
--- a/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
+++ b/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
@@ -119,6 +119,8 @@ class SQLAlchemyAdapter():
                 self.db_path = None
             else:
                 async with self.engine.begin() as connection:
+                    # Load the schema information into the MetaData object
+                    await connection.run_sync(Base.metadata.reflect)
                     for table in Base.metadata.sorted_tables:
                         drop_table_query = text(f"DROP TABLE IF EXISTS {table.name} CASCADE")
                         await connection.execute(drop_table_query)
diff --git a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py
index b9a716663..8b873b2fd 100644
--- a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py
+++ b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py
@@ -135,7 +135,8 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
         with an async engine.
         """
         async with self.engine.begin() as connection:
-            await connection.run_sync(Base.metadata.reflect)  # Reflect the metadata
+            # Load the schema information into the MetaData object
+            await connection.run_sync(Base.metadata.reflect)
             if collection_name in Base.metadata.tables:
                 return Base.metadata.tables[collection_name]
             else: