fix migration version 40 | upgrade

This commit is contained in:
CREDO23 2025-12-13 21:47:20 +02:00
parent c0b414fefa
commit 85f4068735

View file

@ -12,6 +12,7 @@ of a search space, rather than being per-user.
""" """
import sqlalchemy as sa import sqlalchemy as sa
from sqlalchemy import inspect
from alembic import op from alembic import op
@ -23,25 +24,29 @@ depends_on = None
def upgrade(): def upgrade():
# Add LLM preference columns to searchspaces table connection = op.get_bind()
op.add_column( inspector = inspect(connection)
"searchspaces", columns = [col["name"] for col in inspector.get_columns("searchspaces")]
sa.Column("long_context_llm_id", sa.Integer(), nullable=True),
) # Add LLM preference columns to searchspaces table if they don't exist
op.add_column( if "long_context_llm_id" not in columns:
"searchspaces", op.add_column(
sa.Column("fast_llm_id", sa.Integer(), nullable=True), "searchspaces",
) sa.Column("long_context_llm_id", sa.Integer(), nullable=True),
op.add_column( )
"searchspaces", if "fast_llm_id" not in columns:
sa.Column("strategic_llm_id", sa.Integer(), nullable=True), 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 # Migrate existing preferences from user_search_space_preferences to searchspaces
# We take the owner's preferences (the user who created the search space) # 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
connection.execute( connection.execute(
sa.text(""" sa.text("""
UPDATE searchspaces ss UPDATE searchspaces ss
@ -57,7 +62,14 @@ def upgrade():
def downgrade(): def downgrade():
# Remove LLM preference columns from searchspaces table connection = op.get_bind()
op.drop_column("searchspaces", "strategic_llm_id") inspector = inspect(connection)
op.drop_column("searchspaces", "fast_llm_id") columns = [col["name"] for col in inspector.get_columns("searchspaces")]
op.drop_column("searchspaces", "long_context_llm_id")
# 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")