From 85f4068735de4e78063ad0912b4f064e3243e630 Mon Sep 17 00:00:00 2001
From: CREDO23 <thierrybakera12@gmail.com>
Date: Sat, 13 Dec 2025 21:47:20 +0200
Subject: [PATCH] fix migration version 40 | upgrade

---
 .../40_move_llm_preferences_to_searchspace.py | 54 +++++++++++--------
 1 file changed, 33 insertions(+), 21 deletions(-)

diff --git a/surfsense_backend/alembic/versions/40_move_llm_preferences_to_searchspace.py b/surfsense_backend/alembic/versions/40_move_llm_preferences_to_searchspace.py
index 1067cffcc..d582f9681 100644
--- a/surfsense_backend/alembic/versions/40_move_llm_preferences_to_searchspace.py
+++ b/surfsense_backend/alembic/versions/40_move_llm_preferences_to_searchspace.py
@@ -12,6 +12,7 @@ of a search space, rather than being per-user.
 """
 
 import sqlalchemy as sa
+from sqlalchemy import inspect
 
 from alembic import op
 
@@ -23,25 +24,29 @@ depends_on = None
 
 
 def upgrade():
-    # Add LLM preference columns to searchspaces table
-    op.add_column(
-        "searchspaces",
-        sa.Column("long_context_llm_id", sa.Integer(), nullable=True),
-    )
-    op.add_column(
-        "searchspaces",
-        sa.Column("fast_llm_id", sa.Integer(), nullable=True),
-    )
-    op.add_column(
-        "searchspaces",
-        sa.Column("strategic_llm_id", sa.Integer(), nullable=True),
-    )
+    connection = op.get_bind()
+    inspector = inspect(connection)
+    columns = [col["name"] for col in inspector.get_columns("searchspaces")]
+
+    # Add LLM preference columns to searchspaces table if they don't exist
+    if "long_context_llm_id" not in columns:
+        op.add_column(
+            "searchspaces",
+            sa.Column("long_context_llm_id", sa.Integer(), nullable=True),
+        )
+    if "fast_llm_id" not in columns:
+        op.add_column(
+            "searchspaces",
+            sa.Column("fast_llm_id", sa.Integer(), nullable=True),
+        )
+    if "strategic_llm_id" not in columns:
+        op.add_column(
+            "searchspaces",
+            sa.Column("strategic_llm_id", sa.Integer(), nullable=True),
+        )
 
     # Migrate existing preferences from user_search_space_preferences to searchspaces
-    # We take the owner's preferences (the user who created the search space)
-    connection = op.get_bind()
-
-    # Get all search spaces and their owner's preferences
+    # Take the owner's preferences (the user who created the search space)
     connection.execute(
         sa.text("""
             UPDATE searchspaces ss
@@ -57,7 +62,14 @@ def upgrade():
 
 
 def downgrade():
-    # Remove LLM preference columns from searchspaces table
-    op.drop_column("searchspaces", "strategic_llm_id")
-    op.drop_column("searchspaces", "fast_llm_id")
-    op.drop_column("searchspaces", "long_context_llm_id")
+    connection = op.get_bind()
+    inspector = inspect(connection)
+    columns = [col["name"] for col in inspector.get_columns("searchspaces")]
+
+    # Remove columns only if they exist
+    if "strategic_llm_id" in columns:
+        op.drop_column("searchspaces", "strategic_llm_id")
+    if "fast_llm_id" in columns:
+        op.drop_column("searchspaces", "fast_llm_id")
+    if "long_context_llm_id" in columns:
+        op.drop_column("searchspaces", "long_context_llm_id")