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
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")