DH-4735 Fix table-descriptions rows query

Dataherald · Sep 26, 2023 · b4e2039 · b4e2039
1 parent 9e2d119
commit b4e2039
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git 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
@@ -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")