gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: Fix pruning of postgres database

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

Fix #COG-170
This commit is contained in:
Igor Ilic 2024-10-21 14:28:57 +02:00
parent 240c660eac
commit 05e4ef349e
3 changed files with 39 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
cognee/api/v1/config/config.py
View file

@ -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()

2
cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py
View file

@ -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)

3
cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py
View file

@ -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: