Sync-schemas endpoint let adding ids from different db connection

Dataherald · Apr 23, 2024 · a6eea8a · a6eea8a
1 parent 5e76e53
commit a6eea8a
Showing 1 changed file with 21 additions and 13 deletions.
diff --git a/dataherald/api/fastapi.py b/dataherald/api/fastapi.py
@@ -138,9 +138,18 @@ def scan_db(
             table_description = scanner_repository.find_by_id(id)
             if not table_description:
                 raise Exception("Table description not found")
-            if table_description.schema_name not in data.keys():
-                data[table_description.schema_name] = []
-            data[table_description.schema_name].append(table_description)
+            if table_description.db_connection_id not in data.keys():
+                data[table_description.db_connection_id] = {}
+            if (
+                table_description.schema_name
+                not in data[table_description.db_connection_id].keys()
+            ):
+                data[table_description.db_connection_id][
+                    table_description.schema_name
+                ] = []
+            data[table_description.db_connection_id][
+                table_description.schema_name
+            ].append(table_description)
 
         db_connection_repository = DatabaseConnectionRepository(self.storage)
         scanner = self.system.instance(Scanner)
@@ -149,17 +158,16 @@ def scan_db(
             TableDescriptionRepository(self.storage),
         )
         database_connection_service = DatabaseConnectionService(scanner, self.storage)
-        for schema, table_descriptions in data.items():
-            db_connection = db_connection_repository.find_by_id(
-                table_descriptions[0].db_connection_id
-            )
-            database = database_connection_service.get_sql_database(
-                db_connection, schema
-            )
+        for db_connection_id, schemas_and_table_descriptions in data.items():
+            for schema, table_descriptions in schemas_and_table_descriptions.items():
+                db_connection = db_connection_repository.find_by_id(db_connection_id)
+                database = database_connection_service.get_sql_database(
+                    db_connection, schema
+                )
 
-            background_tasks.add_task(
-                async_scanning, scanner, database, table_descriptions, self.storage
-            )
+                background_tasks.add_task(
+                    async_scanning, scanner, database, table_descriptions, self.storage
+                )
         return [TableDescriptionResponse(**row.dict()) for row in rows]
 
     @override