fix: check for column before adding in migrations (#31185)

superset/migrations/shared/utils.py

@@ -21,6 +21,7 @@ from collections.abc import Iterator
 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")

superset/migrations/versions/2024-04-01_22-44_c22cb5c2e546_user_attr_avatar_url.py

@@ -24,14 +24,17 @@ Create Date: 2024-04-01 22:44:40.386543
 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),
     )
 
 

superset/migrations/versions/2024-04-11_15-41_5f57af97bc3f_add_catalog_column.py

@@ -25,6 +25,8 @@ Create Date: 2024-04-11 15:41:34.663989
 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 @@ tables = ["tables", "query", "saved_query", "tab_state", "table_schema"]
 
 def upgrade():
     for table in tables:
-        op.add_column(
+        add_column_if_not_exists(
             table,
             sa.Column("catalog", sa.String(length=256), nullable=True),
         )

superset/migrations/versions/2024-05-01_10-52_58d051681a3b_add_catalog_perm_to_tables.py

@@ -29,6 +29,7 @@ from superset.migrations.shared.catalogs import (
     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 @@ down_revision = "4a33124c18ad"
 
 
 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),
     )