From f3bdcdc79b4c6654043fb2ec72d7a7ba661f485d Mon Sep 17 00:00:00 2001
From: John Bodley <4567245+john-bodley@users.noreply.github.com>
Date: Tue, 14 Feb 2023 06:47:56 +1300
Subject: [PATCH] chore(datasets): Refactor DatasetDAO update to leverage bulk
 logic for create, update, and delete operations (#22957)

---
 superset/datasets/dao.py | 111 +++++++++++++++++++++++----------------
 1 file changed, 66 insertions(+), 45 deletions(-)

diff --git a/superset/datasets/dao.py b/superset/datasets/dao.py
index d260df3610..1c55723b47 100644
--- a/superset/datasets/dao.py
+++ b/superset/datasets/dao.py
@@ -187,39 +187,50 @@ class DatasetDAO(BaseDAO):  # pylint: disable=too-many-public-methods
         then we delete.
         """
 
-        column_by_id = {column.id: column for column in model.columns}
-        seen = set()
-        original_cols = {obj.id for obj in model.columns}
-
         if override_columns:
-            for id_ in original_cols:
-                DatasetDAO.delete_column(column_by_id[id_], commit=False)
+            db.session.query(TableColumn).filter(
+                TableColumn.table_id == model.id
+            ).delete(synchronize_session="fetch")
 
-            db.session.flush()
-
-            for properties in property_columns:
-                DatasetDAO.create_column(
-                    {**properties, "table_id": model.id},
-                    commit=False,
-                )
+            db.session.bulk_insert_mappings(
+                TableColumn,
+                [
+                    {**properties, "table_id": model.id}
+                    for properties in property_columns
+                ],
+            )
         else:
-            for properties in property_columns:
-                if "id" in properties:
-                    seen.add(properties["id"])
+            columns_by_id = {column.id: column for column in model.columns}
 
-                    DatasetDAO.update_column(
-                        column_by_id[properties["id"]],
-                        properties,
-                        commit=False,
-                    )
-                else:
-                    DatasetDAO.create_column(
-                        {**properties, "table_id": model.id},
-                        commit=False,
-                    )
+            property_columns_by_id = {
+                properties["id"]: properties
+                for properties in property_columns
+                if "id" in properties
+            }
 
-            for id_ in {obj.id for obj in model.columns} - seen:
-                DatasetDAO.delete_column(column_by_id[id_], commit=False)
+            db.session.bulk_insert_mappings(
+                TableColumn,
+                [
+                    {**properties, "table_id": model.id}
+                    for properties in property_columns
+                    if not "id" in properties
+                ],
+            )
+
+            db.session.bulk_update_mappings(
+                TableColumn,
+                [
+                    {**columns_by_id[properties["id"]].__dict__, **properties}
+                    for properties in property_columns_by_id.values()
+                ],
+            )
+
+            db.session.query(TableColumn).filter(
+                TableColumn.id.in_(
+                    {column.id for column in model.columns}
+                    - property_columns_by_id.keys()
+                )
+            ).delete(synchronize_session="fetch")
 
         if commit:
             db.session.commit()
@@ -241,26 +252,36 @@ class DatasetDAO(BaseDAO):  # pylint: disable=too-many-public-methods
         then we delete.
         """
 
-        metric_by_id = {metric.id: metric for metric in model.metrics}
-        seen = set()
+        metrics_by_id = {metric.id: metric for metric in model.metrics}
 
-        for properties in property_metrics:
-            if "id" in properties:
-                seen.add(properties["id"])
+        property_metrics_by_id = {
+            properties["id"]: properties
+            for properties in property_metrics
+            if "id" in properties
+        }
 
-                DatasetDAO.update_metric(
-                    metric_by_id[properties["id"]],
-                    properties,
-                    commit=False,
-                )
-            else:
-                DatasetDAO.create_metric(
-                    {**properties, "table_id": model.id},
-                    commit=False,
-                )
+        db.session.bulk_insert_mappings(
+            SqlMetric,
+            [
+                {**properties, "table_id": model.id}
+                for properties in property_metrics
+                if not "id" in properties
+            ],
+        )
 
-        for id_ in {obj.id for obj in model.metrics} - seen:
-            DatasetDAO.delete_column(metric_by_id[id_], commit=False)
+        db.session.bulk_update_mappings(
+            SqlMetric,
+            [
+                {**metrics_by_id[properties["id"]].__dict__, **properties}
+                for properties in property_metrics_by_id.values()
+            ],
+        )
+
+        db.session.query(SqlMetric).filter(
+            SqlMetric.id.in_(
+                {metric.id for metric in model.metrics} - property_metrics_by_id.keys()
+            )
+        ).delete(synchronize_session="fetch")
 
         if commit:
             db.session.commit()