From 8020729cedd17fde21ace71335113d55bbac4f53 Mon Sep 17 00:00:00 2001
From: Joe Li <joe@preset.io>
Date: Mon, 2 Dec 2024 13:47:55 -0800
Subject: [PATCH] fix: check for column before adding in migrations (#31185)

---
 superset/migrations/shared/utils.py               | 15 +++++++++++++++
 ...-01_22-44_c22cb5c2e546_user_attr_avatar_url.py |  7 +++++--
 ...04-11_15-41_5f57af97bc3f_add_catalog_column.py |  4 +++-
 ...-52_58d051681a3b_add_catalog_perm_to_tables.py |  5 +++--
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/superset/migrations/shared/utils.py b/superset/migrations/shared/utils.py
index dba9ceef6393c..f62f8a858a2b4 100644
--- a/superset/migrations/shared/utils.py
+++ b/superset/migrations/shared/utils.py
@@ -21,6 +21,7 @@
 from typing import Any, Callable, Optional, Union
 from uuid import uuid4
 
+import sqlalchemy as sa
 from alembic import op
 from sqlalchemy import inspect
 from sqlalchemy.dialects.mysql.base import MySQLDialect
@@ -182,3 +183,17 @@ def has_table(table_name: str) -> bool:
     table_exists = insp.has_table(table_name)
 
     return table_exists
+
+
+def add_column_if_not_exists(table_name: str, column: sa.Column) -> None:
+    """
+    Adds a column to a table if it does not already exist.
+
+    :param table_name: Name of the table.
+    :param column: SQLAlchemy Column object.
+    """
+    if not table_has_column(table_name, column.name):
+        print(f"Adding column '{column.name}' to table '{table_name}'.\n")
+        op.add_column(table_name, column)
+    else:
+        print(f"Column '{column.name}' already exists in table '{table_name}'.\n")
diff --git a/superset/migrations/versions/2024-04-01_22-44_c22cb5c2e546_user_attr_avatar_url.py b/superset/migrations/versions/2024-04-01_22-44_c22cb5c2e546_user_attr_avatar_url.py
index 8e07f1b6c38b1..126280f338227 100644
--- a/superset/migrations/versions/2024-04-01_22-44_c22cb5c2e546_user_attr_avatar_url.py
+++ b/superset/migrations/versions/2024-04-01_22-44_c22cb5c2e546_user_attr_avatar_url.py
@@ -24,14 +24,17 @@
 import sqlalchemy as sa
 from alembic import op
 
+from superset.migrations.shared.utils import add_column_if_not_exists
+
 # revision identifiers, used by Alembic.
 revision = "c22cb5c2e546"
 down_revision = "678eefb4ab44"
 
 
 def upgrade():
-    op.add_column(
-        "user_attribute", sa.Column("avatar_url", sa.String(length=100), nullable=True)
+    add_column_if_not_exists(
+        "user_attribute",
+        sa.Column("avatar_url", sa.String(length=100), nullable=True),
     )
 
 
diff --git a/superset/migrations/versions/2024-04-11_15-41_5f57af97bc3f_add_catalog_column.py b/superset/migrations/versions/2024-04-11_15-41_5f57af97bc3f_add_catalog_column.py
index 7ed95324ddac2..b535867d64a4b 100644
--- a/superset/migrations/versions/2024-04-11_15-41_5f57af97bc3f_add_catalog_column.py
+++ b/superset/migrations/versions/2024-04-11_15-41_5f57af97bc3f_add_catalog_column.py
@@ -25,6 +25,8 @@
 import sqlalchemy as sa
 from alembic import op
 
+from superset.migrations.shared.utils import add_column_if_not_exists
+
 # revision identifiers, used by Alembic.
 revision = "5f57af97bc3f"
 down_revision = "d60591c5515f"
@@ -34,7 +36,7 @@
 
 def upgrade():
     for table in tables:
-        op.add_column(
+        add_column_if_not_exists(
             table,
             sa.Column("catalog", sa.String(length=256), nullable=True),
         )
diff --git a/superset/migrations/versions/2024-05-01_10-52_58d051681a3b_add_catalog_perm_to_tables.py b/superset/migrations/versions/2024-05-01_10-52_58d051681a3b_add_catalog_perm_to_tables.py
index 88b3d63527dc4..6dfc2845bcb7d 100644
--- a/superset/migrations/versions/2024-05-01_10-52_58d051681a3b_add_catalog_perm_to_tables.py
+++ b/superset/migrations/versions/2024-05-01_10-52_58d051681a3b_add_catalog_perm_to_tables.py
@@ -29,6 +29,7 @@
     downgrade_catalog_perms,
     upgrade_catalog_perms,
 )
+from superset.migrations.shared.utils import add_column_if_not_exists
 
 # revision identifiers, used by Alembic.
 revision = "58d051681a3b"
@@ -36,11 +37,11 @@
 
 
 def upgrade():
-    op.add_column(
+    add_column_if_not_exists(
         "tables",
         sa.Column("catalog_perm", sa.String(length=1000), nullable=True),
     )
-    op.add_column(
+    add_column_if_not_exists(
         "slices",
         sa.Column("catalog_perm", sa.String(length=1000), nullable=True),
     )