fix(refactor): Alembic migrations

surfsense_backend/alembic/versions/36_remove_fk_constraints_for_global_llm_configs.py

@@ -9,6 +9,7 @@ Create Date: 2025-11-13 23:20:12.912741
 from collections.abc import Sequence
 
 from alembic import op
+from sqlalchemy import text
 
 # revision identifiers, used by Alembic.
 revision: str = "36"
@@ -17,6 +18,20 @@ branch_labels: str | Sequence[str] | None = None
 depends_on: str | Sequence[str] | None = None
 
 
+def constraint_exists(connection, table_name: str, constraint_name: str) -> bool:
+    """Check if a constraint exists on the given table."""
+    result = connection.execute(
+        text(
+            """
+            SELECT 1 FROM information_schema.table_constraints
+            WHERE table_name = :table_name AND constraint_name = :constraint_name
+            """
+        ),
+        {"table_name": table_name, "constraint_name": constraint_name},
+    )
+    return result.fetchone() is not None
+
+
 def upgrade() -> None:
     """
     Remove foreign key constraints on LLM preference columns to allow global configs (negative IDs).
@@ -24,50 +39,48 @@ def upgrade() -> None:
     Global LLM configs use negative IDs and don't exist in the llm_configs table,
     so we need to remove the foreign key constraints that were preventing their use.
     """
-    # Drop the foreign key constraints
+    connection = op.get_bind()
-    op.drop_constraint(
+
+    # Drop the foreign key constraints if they exist
+    constraints_to_drop = [
         "user_search_space_preferences_long_context_llm_id_fkey",
-        "user_search_space_preferences",
-        type_="foreignkey",
-    )
-    op.drop_constraint(
         "user_search_space_preferences_fast_llm_id_fkey",
-        "user_search_space_preferences",
-        type_="foreignkey",
-    )
-    op.drop_constraint(
         "user_search_space_preferences_strategic_llm_id_fkey",
-        "user_search_space_preferences",
+    ]
-        type_="foreignkey",
+
-    )
+    for constraint_name in constraints_to_drop:
+        if constraint_exists(connection, "user_search_space_preferences", constraint_name):
+            op.drop_constraint(
+                constraint_name,
+                "user_search_space_preferences",
+                type_="foreignkey",
+            )
+        else:
+            print(f"Constraint '{constraint_name}' does not exist. Skipping.")
 
 
 def downgrade() -> None:
     """
     Re-add foreign key constraints (will fail if any negative IDs exist in the table).
     """
-    # Re-add the foreign key constraints
+    connection = op.get_bind()
-    op.create_foreign_key(
+
-        "user_search_space_preferences_long_context_llm_id_fkey",
+    # Re-add the foreign key constraints if they don't exist
-        "user_search_space_preferences",
+    constraints_to_create = [
-        "llm_configs",
+        ("user_search_space_preferences_long_context_llm_id_fkey", "long_context_llm_id"),
-        ["long_context_llm_id"],
+        ("user_search_space_preferences_fast_llm_id_fkey", "fast_llm_id"),
-        ["id"],
+        ("user_search_space_preferences_strategic_llm_id_fkey", "strategic_llm_id"),
-        ondelete="SET NULL",
+    ]
-    )
+
-    op.create_foreign_key(
+    for constraint_name, column_name in constraints_to_create:
-        "user_search_space_preferences_fast_llm_id_fkey",
+        if not constraint_exists(connection, "user_search_space_preferences", constraint_name):
-        "user_search_space_preferences",
+            op.create_foreign_key(
-        "llm_configs",
+                constraint_name,
-        ["fast_llm_id"],
+                "user_search_space_preferences",
-        ["id"],
+                "llm_configs",
-        ondelete="SET NULL",
+                [column_name],
-    )
+                ["id"],
-    op.create_foreign_key(
+                ondelete="SET NULL",
-        "user_search_space_preferences_strategic_llm_id_fkey",
+            )
-        "user_search_space_preferences",
+        else:
-        "llm_configs",
+            print(f"Constraint '{constraint_name}' already exists. Skipping.")
-        ["strategic_llm_id"],
-        ["id"],
-        ondelete="SET NULL",
-    )

surfsense_backend/alembic/versions/37_add_system_prompts_to_searchspaces.py

@@ -9,6 +9,7 @@ Create Date: 2025-11-19 00:00:00.000000
 from collections.abc import Sequence
 
 import sqlalchemy as sa
+from sqlalchemy import text
 
 from alembic import op
 
@@ -19,24 +20,55 @@ branch_labels: str | Sequence[str] | None = None
 depends_on: str | Sequence[str] | None = None
 
 
+def column_exists(connection, table_name: str, column_name: str) -> bool:
+    """Check if a column exists on the given table."""
+    result = connection.execute(
+        text(
+            """
+            SELECT 1 FROM information_schema.columns
+            WHERE table_name = :table_name AND column_name = :column_name
+            """
+        ),
+        {"table_name": table_name, "column_name": column_name},
+    )
+    return result.fetchone() is not None
+
+
 def upgrade() -> None:
     """Add QnA configuration columns to searchspaces table."""
+    connection = op.get_bind()
+
     # Add citations_enabled boolean (default True)
-    op.add_column(
+    if not column_exists(connection, "searchspaces", "citations_enabled"):
-        "searchspaces",
+        op.add_column(
-        sa.Column(
+            "searchspaces",
-            "citations_enabled", sa.Boolean(), nullable=False, server_default="true"
+            sa.Column(
-        ),
+                "citations_enabled", sa.Boolean(), nullable=False, server_default="true"
-    )
+            ),
+        )
+    else:
+        print("Column 'citations_enabled' already exists. Skipping.")
 
     # Add custom instructions text field (nullable, defaults to empty)
-    op.add_column(
+    if not column_exists(connection, "searchspaces", "qna_custom_instructions"):
-        "searchspaces",
+        op.add_column(
-        sa.Column("qna_custom_instructions", sa.Text(), nullable=True),
+            "searchspaces",
-    )
+            sa.Column("qna_custom_instructions", sa.Text(), nullable=True),
+        )
+    else:
+        print("Column 'qna_custom_instructions' already exists. Skipping.")
 
 
 def downgrade() -> None:
     """Remove QnA configuration columns from searchspaces table."""
-    op.drop_column("searchspaces", "qna_custom_instructions")
+    connection = op.get_bind()
-    op.drop_column("searchspaces", "citations_enabled")
+
+    if column_exists(connection, "searchspaces", "qna_custom_instructions"):
+        op.drop_column("searchspaces", "qna_custom_instructions")
+    else:
+        print("Column 'qna_custom_instructions' does not exist. Skipping.")
+
+    if column_exists(connection, "searchspaces", "citations_enabled"):
+        op.drop_column("searchspaces", "citations_enabled")
+    else:
+        print("Column 'citations_enabled' does not exist. Skipping.")