From 9e39613b21edcd3beb893a3144511dffc8bc5185 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juan=20Carlos=20Jos=C3=A9=20Camacho?= Date: Tue, 26 Sep 2023 14:36:12 -0600 Subject: [PATCH] DH-4735 Fix table-descriptions rows query (#186) --- dataherald/db_scanner/repository/base.py | 4 ++-- dataherald/sql_generator/dataherald_sqlagent.py | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/dataherald/db_scanner/repository/base.py b/dataherald/db_scanner/repository/base.py index 00e0b6eb..f69a5972 100644 --- a/dataherald/db_scanner/repository/base.py +++ b/dataherald/db_scanner/repository/base.py @@ -32,8 +32,8 @@ def get_table_info( return TableSchemaDetail(**row) return None - def get_all_tables_by_db(self, db_connection_id: str) -> List[TableSchemaDetail]: - rows = self.storage.find(DB_COLLECTION, {"db_connection_id": db_connection_id}) + def get_all_tables_by_db(self, query: dict) -> List[TableSchemaDetail]: + rows = self.storage.find(DB_COLLECTION, query) tables = [] for row in rows: row["id"] = row["_id"] diff --git a/dataherald/sql_generator/dataherald_sqlagent.py b/dataherald/sql_generator/dataherald_sqlagent.py index f4a3355c..cbec9ca5 100644 --- a/dataherald/sql_generator/dataherald_sqlagent.py +++ b/dataherald/sql_generator/dataherald_sqlagent.py @@ -581,8 +581,10 @@ def generate_response( ) repository = DBScannerRepository(storage) db_scan = repository.get_all_tables_by_db( - db_connection_id=database_connection.id, - status=TableDescriptionStatus.SYNCHRONIZED.value, + { + "db_connection_id": str(database_connection.id), + "status": TableDescriptionStatus.SYNCHRONIZED.value, + } ) if not db_scan: raise ValueError("No scanned tables found for database")