Fix schema_name filter in table-description repository

Dataherald · Apr 24, 2024 · 4eb7a3e · 4eb7a3e
1 parent ef9e4e6
commit 4eb7a3e
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/dataherald/db_scanner/repository/base.py b/dataherald/db_scanner/repository/base.py
@@ -53,14 +53,18 @@ def save_table_info(self, table_info: TableDescription) -> TableDescription:
         table_info_dict = {
             k: v for k, v in table_info_dict.items() if v is not None and v != []
         }
+
+        query = {
+            "db_connection_id": table_info_dict["db_connection_id"],
+            "table_name": table_info_dict["table_name"],
+        }
+        if "schema_name" in table_info_dict:
+            query["schema_name"] = table_info_dict["schema_name"]
+
         table_info.id = str(
             self.storage.update_or_create(
                 DB_COLLECTION,
-                {
-                    "db_connection_id": table_info_dict["db_connection_id"],
-                    "table_name": table_info_dict["table_name"],
-                    "schema_name": table_info_dict["schema_name"],
-                },
+                query,
                 table_info_dict,
             )
         )