diff --git a/surfsense_backend/alembic/versions/36_remove_fk_constraints_for_global_llm_configs.py b/surfsense_backend/alembic/versions/36_remove_fk_constraints_for_global_llm_configs.py index fa4c929ce..a750d6455 100644 --- a/surfsense_backend/alembic/versions/36_remove_fk_constraints_for_global_llm_configs.py +++ b/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 - op.drop_constraint( + connection = op.get_bind() + + # 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 - op.create_foreign_key( - "user_search_space_preferences_long_context_llm_id_fkey", - "user_search_space_preferences", - "llm_configs", - ["long_context_llm_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - "user_search_space_preferences_fast_llm_id_fkey", - "user_search_space_preferences", - "llm_configs", - ["fast_llm_id"], - ["id"], - ondelete="SET NULL", - ) - op.create_foreign_key( - "user_search_space_preferences_strategic_llm_id_fkey", - "user_search_space_preferences", - "llm_configs", - ["strategic_llm_id"], - ["id"], - ondelete="SET NULL", - ) + connection = op.get_bind() + + # Re-add the foreign key constraints if they don't exist + constraints_to_create = [ + ("user_search_space_preferences_long_context_llm_id_fkey", "long_context_llm_id"), + ("user_search_space_preferences_fast_llm_id_fkey", "fast_llm_id"), + ("user_search_space_preferences_strategic_llm_id_fkey", "strategic_llm_id"), + ] + + for constraint_name, column_name in constraints_to_create: + if not constraint_exists(connection, "user_search_space_preferences", constraint_name): + op.create_foreign_key( + constraint_name, + "user_search_space_preferences", + "llm_configs", + [column_name], + ["id"], + ondelete="SET NULL", + ) + else: + print(f"Constraint '{constraint_name}' already exists. Skipping.") diff --git a/surfsense_backend/alembic/versions/37_add_system_prompts_to_searchspaces.py b/surfsense_backend/alembic/versions/37_add_system_prompts_to_searchspaces.py index afdee4942..6c82b4bad 100644 --- a/surfsense_backend/alembic/versions/37_add_system_prompts_to_searchspaces.py +++ b/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( - "searchspaces", - sa.Column( - "citations_enabled", sa.Boolean(), nullable=False, server_default="true" - ), - ) + if not column_exists(connection, "searchspaces", "citations_enabled"): + op.add_column( + "searchspaces", + 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( - "searchspaces", - sa.Column("qna_custom_instructions", sa.Text(), nullable=True), - ) + if not column_exists(connection, "searchspaces", "qna_custom_instructions"): + op.add_column( + "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") - op.drop_column("searchspaces", "citations_enabled") + connection = op.get_bind() + + 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.")