1
This commit is contained in:
@@ -0,0 +1,230 @@
|
||||
"""add client-selectable multi generation and image batch claim
|
||||
|
||||
Revision ID: abae3e1c70f7
|
||||
Revises: 2026070902
|
||||
Create Date: 2026-07-15 10:49:31.803342
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "abae3e1c70f7"
|
||||
down_revision: Union[str, None] = "2026070902"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
FK_CHAT_TASK_PARENT = "fk_chat_generation_tasks_parent_task_id"
|
||||
CK_CHAT_TASK_GENERATION_COUNT = "ck_chat_generation_tasks_generation_count"
|
||||
CK_CHAT_TASK_GENERATION_INDEX = "ck_chat_generation_tasks_generation_index"
|
||||
CK_IMAGE_ENGINE_MAX_GENERATION_COUNT = "ck_image_engines_max_generation_count"
|
||||
CK_IMAGE_ENGINE_MULTI_IMAGE_MAX = "ck_image_engines_multi_image_max_images"
|
||||
CK_IMAGE_ENGINE_MAX_REFERENCE = "ck_image_engines_max_reference_image_count"
|
||||
CK_VIDEO_ENGINE_MAX_GENERATION_COUNT = "ck_video_engines_max_generation_count"
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# ChatGenerationTask：任务级实际生成数量、主子关联和图片批次执行租约。
|
||||
op.add_column(
|
||||
"chat_generation_tasks",
|
||||
sa.Column("parent_task_id", sa.String(length=32), nullable=True),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_generation_tasks",
|
||||
sa.Column("generation_count", sa.Integer(), server_default=sa.text("1"), nullable=False),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_generation_tasks",
|
||||
sa.Column("generation_index", sa.Integer(), nullable=True),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_generation_tasks",
|
||||
sa.Column("provider_create_claim_token", sa.String(length=64), nullable=True),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_generation_tasks",
|
||||
sa.Column("provider_create_lease_until", sa.DateTime(timezone=True), nullable=True),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_generation_tasks",
|
||||
sa.Column("provider_create_started_at", sa.DateTime(timezone=True), nullable=True),
|
||||
)
|
||||
|
||||
op.create_check_constraint(
|
||||
CK_CHAT_TASK_GENERATION_COUNT,
|
||||
"chat_generation_tasks",
|
||||
"generation_count BETWEEN 1 AND 5",
|
||||
)
|
||||
op.create_check_constraint(
|
||||
CK_CHAT_TASK_GENERATION_INDEX,
|
||||
"chat_generation_tasks",
|
||||
"generation_index IS NULL OR generation_index > 0",
|
||||
)
|
||||
op.create_foreign_key(
|
||||
FK_CHAT_TASK_PARENT,
|
||||
"chat_generation_tasks",
|
||||
"chat_generation_tasks",
|
||||
["parent_task_id"],
|
||||
["id"],
|
||||
ondelete="RESTRICT",
|
||||
)
|
||||
op.create_index(
|
||||
"idx_chat_generation_tasks_parent",
|
||||
"chat_generation_tasks",
|
||||
["parent_task_id"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
"idx_chat_generation_tasks_user_mode_created",
|
||||
"chat_generation_tasks",
|
||||
["user_id", "generation_mode", "created_at"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
"ix_chat_generation_tasks_provider_create_claim_token",
|
||||
"chat_generation_tasks",
|
||||
["provider_create_claim_token"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
"ix_chat_generation_tasks_provider_create_lease_until",
|
||||
"chat_generation_tasks",
|
||||
["provider_create_lease_until"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
"uq_chat_generation_tasks_parent_index",
|
||||
"chat_generation_tasks",
|
||||
["parent_task_id", "generation_index"],
|
||||
unique=True,
|
||||
postgresql_where=sa.text(
|
||||
"parent_task_id IS NOT NULL AND generation_index IS NOT NULL"
|
||||
),
|
||||
)
|
||||
op.create_index(
|
||||
"uq_chat_generation_tasks_user_chat_idempotency",
|
||||
"chat_generation_tasks",
|
||||
["user_id", "idempotency_key"],
|
||||
unique=True,
|
||||
postgresql_where=sa.text(
|
||||
"deleted_at IS NULL "
|
||||
"AND idempotency_key IS NOT NULL "
|
||||
"AND generation_mode IN ('chatapi_async', 'chatapi_main')"
|
||||
),
|
||||
)
|
||||
|
||||
# ImageEngine：管理后台只配置是否允许客户端多份生成和数量上限。
|
||||
op.add_column(
|
||||
"image_engines",
|
||||
sa.Column(
|
||||
"multi_generation_enabled",
|
||||
sa.Boolean(),
|
||||
server_default=sa.text("false"),
|
||||
nullable=False,
|
||||
),
|
||||
)
|
||||
op.add_column(
|
||||
"image_engines",
|
||||
sa.Column("max_generation_count", sa.Integer(), server_default=sa.text("1"), nullable=False),
|
||||
)
|
||||
op.add_column(
|
||||
"image_engines",
|
||||
sa.Column("multi_image_max_images", sa.Integer(), server_default=sa.text("15"), nullable=False),
|
||||
)
|
||||
op.add_column(
|
||||
"image_engines",
|
||||
sa.Column("max_reference_image_count", sa.Integer(), server_default=sa.text("14"), nullable=False),
|
||||
)
|
||||
op.add_column(
|
||||
"image_engines",
|
||||
sa.Column("output_format", sa.String(length=16), server_default=sa.text("''"), nullable=False),
|
||||
)
|
||||
op.create_check_constraint(
|
||||
CK_IMAGE_ENGINE_MAX_GENERATION_COUNT,
|
||||
"image_engines",
|
||||
"max_generation_count BETWEEN 1 AND 5",
|
||||
)
|
||||
op.create_check_constraint(
|
||||
CK_IMAGE_ENGINE_MULTI_IMAGE_MAX,
|
||||
"image_engines",
|
||||
"multi_image_max_images BETWEEN 1 AND 15",
|
||||
)
|
||||
op.create_check_constraint(
|
||||
CK_IMAGE_ENGINE_MAX_REFERENCE,
|
||||
"image_engines",
|
||||
"max_reference_image_count BETWEEN 0 AND 14",
|
||||
)
|
||||
|
||||
# VideoEngine：管理后台只配置是否允许客户端多份生成和数量上限。
|
||||
op.add_column(
|
||||
"video_engines",
|
||||
sa.Column(
|
||||
"multi_generation_enabled",
|
||||
sa.Boolean(),
|
||||
server_default=sa.text("false"),
|
||||
nullable=False,
|
||||
),
|
||||
)
|
||||
op.add_column(
|
||||
"video_engines",
|
||||
sa.Column("max_generation_count", sa.Integer(), server_default=sa.text("1"), nullable=False),
|
||||
)
|
||||
op.create_check_constraint(
|
||||
CK_VIDEO_ENGINE_MAX_GENERATION_COUNT,
|
||||
"video_engines",
|
||||
"max_generation_count BETWEEN 1 AND 5",
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_constraint(CK_VIDEO_ENGINE_MAX_GENERATION_COUNT, "video_engines", type_="check")
|
||||
op.drop_column("video_engines", "max_generation_count")
|
||||
op.drop_column("video_engines", "multi_generation_enabled")
|
||||
|
||||
op.drop_constraint(CK_IMAGE_ENGINE_MAX_REFERENCE, "image_engines", type_="check")
|
||||
op.drop_constraint(CK_IMAGE_ENGINE_MULTI_IMAGE_MAX, "image_engines", type_="check")
|
||||
op.drop_constraint(CK_IMAGE_ENGINE_MAX_GENERATION_COUNT, "image_engines", type_="check")
|
||||
op.drop_column("image_engines", "output_format")
|
||||
op.drop_column("image_engines", "max_reference_image_count")
|
||||
op.drop_column("image_engines", "multi_image_max_images")
|
||||
op.drop_column("image_engines", "max_generation_count")
|
||||
op.drop_column("image_engines", "multi_generation_enabled")
|
||||
|
||||
op.drop_index(
|
||||
"uq_chat_generation_tasks_user_chat_idempotency",
|
||||
table_name="chat_generation_tasks",
|
||||
postgresql_where=sa.text(
|
||||
"deleted_at IS NULL "
|
||||
"AND idempotency_key IS NOT NULL "
|
||||
"AND generation_mode IN ('chatapi_async', 'chatapi_main')"
|
||||
),
|
||||
)
|
||||
op.drop_index(
|
||||
"uq_chat_generation_tasks_parent_index",
|
||||
table_name="chat_generation_tasks",
|
||||
postgresql_where=sa.text(
|
||||
"parent_task_id IS NOT NULL AND generation_index IS NOT NULL"
|
||||
),
|
||||
)
|
||||
op.drop_index(
|
||||
"ix_chat_generation_tasks_provider_create_lease_until",
|
||||
table_name="chat_generation_tasks",
|
||||
)
|
||||
op.drop_index(
|
||||
"ix_chat_generation_tasks_provider_create_claim_token",
|
||||
table_name="chat_generation_tasks",
|
||||
)
|
||||
op.drop_index("idx_chat_generation_tasks_user_mode_created", table_name="chat_generation_tasks")
|
||||
op.drop_index("idx_chat_generation_tasks_parent", table_name="chat_generation_tasks")
|
||||
op.drop_constraint(FK_CHAT_TASK_PARENT, "chat_generation_tasks", type_="foreignkey")
|
||||
op.drop_constraint(CK_CHAT_TASK_GENERATION_INDEX, "chat_generation_tasks", type_="check")
|
||||
op.drop_constraint(CK_CHAT_TASK_GENERATION_COUNT, "chat_generation_tasks", type_="check")
|
||||
op.drop_column("chat_generation_tasks", "provider_create_started_at")
|
||||
op.drop_column("chat_generation_tasks", "provider_create_lease_until")
|
||||
op.drop_column("chat_generation_tasks", "provider_create_claim_token")
|
||||
op.drop_column("chat_generation_tasks", "generation_index")
|
||||
op.drop_column("chat_generation_tasks", "generation_count")
|
||||
op.drop_column("chat_generation_tasks", "parent_task_id")
|
||||
Reference in New Issue
Block a user