diff --git a/migrations/versions/c2d3e4f5a6b7_upstream_provider_base_url_api_key_unique.py b/migrations/versions/c2d3e4f5a6b7_upstream_provider_base_url_api_key_unique.py index 83035eb0..bce53175 100644 --- a/migrations/versions/c2d3e4f5a6b7_upstream_provider_base_url_api_key_unique.py +++ b/migrations/versions/c2d3e4f5a6b7_upstream_provider_base_url_api_key_unique.py @@ -16,35 +16,102 @@ branch_labels = None depends_on = None -def upgrade() -> None: - with op.batch_alter_table("upstream_providers", recreate="always") as batch_op: - batch_op.drop_index("ix_upstream_providers_base_url") - batch_op.alter_column( +def _recreate_table_sqlite(add_base_url_unique: bool) -> None: + conn = op.get_bind() + existing_tables = { + row[0] + for row in conn.exec_driver_sql( + "SELECT name FROM sqlite_master WHERE type='table'" + ).fetchall() + } + if "upstream_providers_old" in existing_tables: + if "upstream_providers" in existing_tables: + op.drop_table("upstream_providers_old") + else: + op.execute( + "ALTER TABLE upstream_providers_old RENAME TO upstream_providers" + ) + existing_tables.add("upstream_providers") + if "upstream_providers" not in existing_tables: + return + + constraints = [ + sa.UniqueConstraint( "base_url", - existing_type=sa.String(), - nullable=False, - unique=False, + "api_key", + name="uq_upstream_providers_base_url_api_key", ) - batch_op.create_unique_constraint( - "uq_upstream_providers_base_url_api_key", - ["base_url", "api_key"], + ] + if add_base_url_unique: + constraints.append( + sa.UniqueConstraint("base_url", name="uq_upstream_providers_base_url") ) + op.execute("ALTER TABLE upstream_providers RENAME TO upstream_providers_old") + op.create_table( + "upstream_providers", + sa.Column( + "id", sa.Integer(), primary_key=True, nullable=False, autoincrement=True + ), + sa.Column("provider_type", sa.String(), nullable=False), + sa.Column("base_url", sa.String(), nullable=False), + sa.Column("api_key", sa.String(), nullable=False), + sa.Column("api_version", sa.String(), nullable=True), + sa.Column("enabled", sa.Boolean(), nullable=False), + sa.Column("provider_fee", sa.Float(), nullable=False, server_default="1.01"), + *constraints, + ) + op.execute( + "INSERT INTO upstream_providers (id, provider_type, base_url, api_key, api_version, enabled, provider_fee) " + "SELECT id, provider_type, base_url, api_key, api_version, enabled, provider_fee " + "FROM upstream_providers_old" + ) + op.drop_table("upstream_providers_old") + + +def upgrade() -> None: + conn = op.get_bind() + if conn.dialect.name == "sqlite": + _recreate_table_sqlite(add_base_url_unique=False) + return + + inspector = sa.inspect(conn) + for constraint in inspector.get_unique_constraints("upstream_providers"): + if constraint.get("column_names") == ["base_url"] and constraint.get("name"): + op.drop_constraint( + constraint["name"], + "upstream_providers", + type_="unique", + ) + index_names = {idx["name"] for idx in inspector.get_indexes("upstream_providers")} + if "ix_upstream_providers_base_url" in index_names: + op.drop_index("ix_upstream_providers_base_url", table_name="upstream_providers") + op.create_unique_constraint( + "uq_upstream_providers_base_url_api_key", + "upstream_providers", + ["base_url", "api_key"], + ) + def downgrade() -> None: - with op.batch_alter_table("upstream_providers", recreate="always") as batch_op: - batch_op.drop_constraint( - "uq_upstream_providers_base_url_api_key", - type_="unique", - ) - batch_op.alter_column( - "base_url", - existing_type=sa.String(), - nullable=False, - unique=True, - ) - batch_op.create_index( - "ix_upstream_providers_base_url", - ["base_url"], - unique=True, - ) + conn = op.get_bind() + if conn.dialect.name == "sqlite": + _recreate_table_sqlite(add_base_url_unique=True) + return + + op.drop_constraint( + "uq_upstream_providers_base_url_api_key", + "upstream_providers", + type_="unique", + ) + op.create_unique_constraint( + "uq_upstream_providers_base_url", + "upstream_providers", + ["base_url"], + ) + op.create_index( + "ix_upstream_providers_base_url", + "upstream_providers", + ["base_url"], + unique=True, + ) diff --git a/migrations/versions/d4e5f6a7b8c9_remove_base_url_unique_constraint.py b/migrations/versions/d4e5f6a7b8c9_remove_base_url_unique_constraint.py deleted file mode 100644 index 31546192..00000000 --- a/migrations/versions/d4e5f6a7b8c9_remove_base_url_unique_constraint.py +++ /dev/null @@ -1,82 +0,0 @@ -"""remove upstream_providers base_url unique constraint - -Revision ID: d4e5f6a7b8c9 -Revises: c2d3e4f5a6b7 -Create Date: 2026-01-25 00:00:00.000000 -""" - -from __future__ import annotations - -import sqlalchemy as sa -from alembic import op - -revision = "d4e5f6a7b8c9" -down_revision = "c2d3e4f5a6b7" -branch_labels = None -depends_on = None - - -def _recreate_table_sqlite(add_base_url_unique: bool) -> None: - constraints = [ - sa.UniqueConstraint( - "base_url", - "api_key", - name="uq_upstream_providers_base_url_api_key", - ) - ] - if add_base_url_unique: - constraints.append( - sa.UniqueConstraint("base_url", name="uq_upstream_providers_base_url") - ) - op.execute("ALTER TABLE upstream_providers RENAME TO upstream_providers_old") - op.create_table( - "upstream_providers", - sa.Column( - "id", sa.Integer(), primary_key=True, nullable=False, autoincrement=True - ), - sa.Column("provider_type", sa.String(), nullable=False), - sa.Column("base_url", sa.String(), nullable=False), - sa.Column("api_key", sa.String(), nullable=False), - sa.Column("api_version", sa.String(), nullable=True), - sa.Column("enabled", sa.Boolean(), nullable=False), - sa.Column("provider_fee", sa.Float(), nullable=False, server_default="1.01"), - *constraints, - ) - op.execute( - "INSERT INTO upstream_providers (id, provider_type, base_url, api_key, api_version, enabled, provider_fee) " - "SELECT id, provider_type, base_url, api_key, api_version, enabled, provider_fee " - "FROM upstream_providers_old" - ) - op.drop_table("upstream_providers_old") - - -def upgrade() -> None: - conn = op.get_bind() - if conn.dialect.name == "sqlite": - _recreate_table_sqlite(add_base_url_unique=False) - return - - inspector = sa.inspect(conn) - for constraint in inspector.get_unique_constraints("upstream_providers"): - if constraint.get("column_names") == ["base_url"] and constraint.get("name"): - op.drop_constraint( - constraint["name"], - "upstream_providers", - type_="unique", - ) - index_names = {idx["name"] for idx in inspector.get_indexes("upstream_providers")} - if "ix_upstream_providers_base_url" in index_names: - op.drop_index("ix_upstream_providers_base_url", table_name="upstream_providers") - - -def downgrade() -> None: - conn = op.get_bind() - if conn.dialect.name == "sqlite": - _recreate_table_sqlite(add_base_url_unique=True) - return - - op.create_unique_constraint( - "uq_upstream_providers_base_url", - "upstream_providers", - ["base_url"], - )