fix: Metastore sync tables with a single transaction

pinterest · Aug 2, 2024 · 72c291b · 72c291b
1 parent dd50b5d
commit 72c291b
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 6 deletions.
diff --git a/querybook/server/lib/metastore/base_metastore_loader.py b/querybook/server/lib/metastore/base_metastore_loader.py
@@ -400,7 +400,10 @@ def _create_table_table(
                 )
 
             delete_column_not_in_metastore(
-                table_id, set(map(lambda c: c.name, columns)), session=session
+                table_id,
+                set(map(lambda c: c.name, columns)),
+                commit=False,
+                session=session,
             )
 
             for column in columns:
@@ -630,11 +633,14 @@ def delete_table_not_in_metastore(schema_id, table_names, session=None):
 
 
 @with_session
-def delete_column_not_in_metastore(table_id, column_names, session=None):
+def delete_column_not_in_metastore(table_id, column_names, commit=True, session=None):
     db_columns = get_column_by_table_id(table_id, session=session)
 
     for column in db_columns:
         if column.name not in column_names:
             delete_column(id=column.id, commit=False, session=session)
             LOG.info("deleted column %d" % column.id)
-    session.commit()
+    if commit:
+        session.commit()
+    else:
+        session.flush()
diff --git a/querybook/server/logic/metastore.py b/querybook/server/logic/metastore.py
@@ -351,7 +351,9 @@ def create_table_warnings(
                 "message": message,
                 "severity": severity,
                 "table_id": table_id,
-            }
+            },
+            commit=False,
+            session=session,
         )
     if commit:
         session.commit()

diff --git a/querybook/server/logic/tag.py b/querybook/server/logic/tag.py
@@ -133,12 +133,13 @@ def create_table_tags(
 
         # update or create a new tag if not exist
         create_or_update_tag(
-            tag_name=tag.name, meta=meta, commit=commit, session=session
+            tag_name=tag.name, meta=meta, commit=False, session=session
         )
 
         # add a new tag_item to associate with the table
         TagItem.create(
             {"tag_name": tag.name, "table_id": table_id, "uid": None},
+            commit=False,
             session=session,
         )
 
@@ -176,12 +177,13 @@ def create_column_tags(
 
         # update or create a new tag if not exist
         create_or_update_tag(
-            tag_name=tag.name, meta=meta, commit=commit, session=session
+            tag_name=tag.name, meta=meta, commit=False, session=session
         )
 
         # add a new tag_item to associate with the table
         TagItem.create(
             {"tag_name": tag.name, "column_id": column_id, "uid": None},
+            commit=False,
             session=session,
         )