From 72c291b2fe14be76ac27bb73b0e1c81efcd0fa60 Mon Sep 17 00:00:00 2001
From: Dave Bauman <dbauman@expediagroup.com>
Date: Tue, 9 Jul 2024 12:00:14 -0400
Subject: [PATCH] fix: Metastore sync tables with a single transaction

---
 .../server/lib/metastore/base_metastore_loader.py    | 12 +++++++++---
 querybook/server/logic/metastore.py                  |  4 +++-
 querybook/server/logic/tag.py                        |  6 ++++--
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/querybook/server/lib/metastore/base_metastore_loader.py b/querybook/server/lib/metastore/base_metastore_loader.py
index aec6ee0be..74c548315 100644
--- 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
index ac253de17..81fbbb6ef 100644
--- 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
index a9d96fc7f..02816e3e4 100644
--- 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,
         )