mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-06-01 18:49:27 +02:00
66 lines
1.8 KiB
Python
66 lines
1.8 KiB
Python
"""add cloud embedding model and update embedding_model
|
|
|
|
Revision ID: 44f856ae2a4a
|
|
Revises: d716b0791ddd
|
|
Create Date: 2024-06-28 20:01:05.927647
|
|
|
|
"""
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision = "44f856ae2a4a"
|
|
down_revision = "d716b0791ddd"
|
|
branch_labels: None = None
|
|
depends_on: None = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
# Create embedding_provider table
|
|
op.create_table(
|
|
"embedding_provider",
|
|
sa.Column("id", sa.Integer(), nullable=False),
|
|
sa.Column("name", sa.String(), nullable=False),
|
|
sa.Column("api_key", sa.LargeBinary(), nullable=True),
|
|
sa.Column("default_model_id", sa.Integer(), nullable=True),
|
|
sa.PrimaryKeyConstraint("id"),
|
|
sa.UniqueConstraint("name"),
|
|
)
|
|
|
|
# Add cloud_provider_id to embedding_model table
|
|
op.add_column(
|
|
"embedding_model", sa.Column("cloud_provider_id", sa.Integer(), nullable=True)
|
|
)
|
|
|
|
# Add foreign key constraints
|
|
op.create_foreign_key(
|
|
"fk_embedding_model_cloud_provider",
|
|
"embedding_model",
|
|
"embedding_provider",
|
|
["cloud_provider_id"],
|
|
["id"],
|
|
)
|
|
op.create_foreign_key(
|
|
"fk_embedding_provider_default_model",
|
|
"embedding_provider",
|
|
"embedding_model",
|
|
["default_model_id"],
|
|
["id"],
|
|
)
|
|
|
|
|
|
def downgrade() -> None:
|
|
# Remove foreign key constraints
|
|
op.drop_constraint(
|
|
"fk_embedding_model_cloud_provider", "embedding_model", type_="foreignkey"
|
|
)
|
|
op.drop_constraint(
|
|
"fk_embedding_provider_default_model", "embedding_provider", type_="foreignkey"
|
|
)
|
|
|
|
# Remove cloud_provider_id column
|
|
op.drop_column("embedding_model", "cloud_provider_id")
|
|
|
|
# Drop embedding_provider table
|
|
op.drop_table("embedding_provider")
|