1
This commit is contained in:
+922
@@ -0,0 +1,922 @@
|
||||
"""unify generation owners and soft delete engines
|
||||
|
||||
Revision ID: e6eac828ff61
|
||||
Revises: 3f47680a71d0
|
||||
Create Date: 2026-07-20 09:19:07.410183
|
||||
|
||||
This revision intentionally contains only the generation-pipeline and engine
|
||||
soft-delete changes. Autogenerate noise from unrelated modules is excluded.
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "e6eac828ff61"
|
||||
down_revision: Union[str, None] = "3f47680a71d0"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
FK_EVENT_GENERATION_RECORD = "fk_chat_generation_task_events_generation_record_id"
|
||||
FK_CALL_LOG_GENERATION_RECORD = "fk_chat_provider_call_logs_generation_record_id"
|
||||
CK_EVENT_OWNER = "ck_chat_generation_task_events_owner"
|
||||
CK_CALL_LOG_OWNER = "ck_chat_provider_call_logs_owner"
|
||||
|
||||
|
||||
def _add_engine_soft_delete_columns() -> None:
|
||||
for table_name in ("image_engines", "model_configs", "video_engines"):
|
||||
op.add_column(
|
||||
table_name,
|
||||
sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True),
|
||||
)
|
||||
op.create_index(
|
||||
op.f(f"ix_{table_name}_deleted_at"),
|
||||
table_name,
|
||||
["deleted_at"],
|
||||
unique=False,
|
||||
)
|
||||
|
||||
|
||||
def _add_shared_log_owner_columns() -> None:
|
||||
op.add_column(
|
||||
"chat_generation_task_events",
|
||||
sa.Column(
|
||||
"owner_type",
|
||||
sa.String(length=32),
|
||||
server_default="chat_generation_task",
|
||||
nullable=False,
|
||||
),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_generation_task_events",
|
||||
sa.Column("generation_record_id", sa.String(length=32), nullable=True),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_generation_task_events",
|
||||
sa.Column(
|
||||
"generation_attempt_no",
|
||||
sa.Integer(),
|
||||
server_default="1",
|
||||
nullable=False,
|
||||
),
|
||||
)
|
||||
op.alter_column(
|
||||
"chat_generation_task_events",
|
||||
"task_id",
|
||||
existing_type=sa.VARCHAR(length=32),
|
||||
nullable=True,
|
||||
)
|
||||
op.alter_column(
|
||||
"chat_generation_task_events",
|
||||
"from_stage",
|
||||
existing_type=sa.VARCHAR(length=32),
|
||||
type_=sa.String(length=48),
|
||||
existing_nullable=True,
|
||||
)
|
||||
op.alter_column(
|
||||
"chat_generation_task_events",
|
||||
"to_stage",
|
||||
existing_type=sa.VARCHAR(length=32),
|
||||
type_=sa.String(length=48),
|
||||
existing_nullable=True,
|
||||
)
|
||||
op.alter_column(
|
||||
"chat_generation_task_events",
|
||||
"message",
|
||||
existing_type=sa.VARCHAR(length=512),
|
||||
type_=sa.Text(),
|
||||
existing_nullable=True,
|
||||
)
|
||||
op.create_index(
|
||||
"idx_chat_generation_task_events_attempt_created",
|
||||
"chat_generation_task_events",
|
||||
["owner_type", "generation_attempt_no", "created_at"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
"idx_chat_generation_task_events_record_created",
|
||||
"chat_generation_task_events",
|
||||
["owner_type", "generation_record_id", "created_at"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
"idx_chat_generation_task_events_task_created",
|
||||
"chat_generation_task_events",
|
||||
["owner_type", "task_id", "created_at"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
op.f("ix_chat_generation_task_events_generation_attempt_no"),
|
||||
"chat_generation_task_events",
|
||||
["generation_attempt_no"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
op.f("ix_chat_generation_task_events_generation_record_id"),
|
||||
"chat_generation_task_events",
|
||||
["generation_record_id"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
op.f("ix_chat_generation_task_events_owner_type"),
|
||||
"chat_generation_task_events",
|
||||
["owner_type"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_foreign_key(
|
||||
FK_EVENT_GENERATION_RECORD,
|
||||
"chat_generation_task_events",
|
||||
"generation_records",
|
||||
["generation_record_id"],
|
||||
["id"],
|
||||
ondelete="CASCADE",
|
||||
)
|
||||
op.create_check_constraint(
|
||||
CK_EVENT_OWNER,
|
||||
"chat_generation_task_events",
|
||||
"(owner_type = 'chat_generation_task' AND task_id IS NOT NULL AND generation_record_id IS NULL) "
|
||||
"OR (owner_type = 'generation_record' AND task_id IS NULL AND generation_record_id IS NOT NULL)",
|
||||
)
|
||||
|
||||
op.add_column(
|
||||
"chat_provider_call_logs",
|
||||
sa.Column(
|
||||
"owner_type",
|
||||
sa.String(length=32),
|
||||
server_default="chat_generation_task",
|
||||
nullable=False,
|
||||
),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_provider_call_logs",
|
||||
sa.Column("generation_record_id", sa.String(length=32), nullable=True),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_provider_call_logs",
|
||||
sa.Column(
|
||||
"generation_attempt_no",
|
||||
sa.Integer(),
|
||||
server_default="1",
|
||||
nullable=False,
|
||||
),
|
||||
)
|
||||
op.alter_column(
|
||||
"chat_provider_call_logs",
|
||||
"task_id",
|
||||
existing_type=sa.VARCHAR(length=32),
|
||||
nullable=True,
|
||||
)
|
||||
op.create_index(
|
||||
"idx_chat_provider_call_logs_attempt_created",
|
||||
"chat_provider_call_logs",
|
||||
["owner_type", "generation_attempt_no", "created_at"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
"idx_chat_provider_call_logs_record_created",
|
||||
"chat_provider_call_logs",
|
||||
["owner_type", "generation_record_id", "created_at"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
"idx_chat_provider_call_logs_task_created",
|
||||
"chat_provider_call_logs",
|
||||
["owner_type", "task_id", "created_at"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
op.f("ix_chat_provider_call_logs_generation_attempt_no"),
|
||||
"chat_provider_call_logs",
|
||||
["generation_attempt_no"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
op.f("ix_chat_provider_call_logs_generation_record_id"),
|
||||
"chat_provider_call_logs",
|
||||
["generation_record_id"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
op.f("ix_chat_provider_call_logs_owner_type"),
|
||||
"chat_provider_call_logs",
|
||||
["owner_type"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_foreign_key(
|
||||
FK_CALL_LOG_GENERATION_RECORD,
|
||||
"chat_provider_call_logs",
|
||||
"generation_records",
|
||||
["generation_record_id"],
|
||||
["id"],
|
||||
ondelete="CASCADE",
|
||||
)
|
||||
op.create_check_constraint(
|
||||
CK_CALL_LOG_OWNER,
|
||||
"chat_provider_call_logs",
|
||||
"(owner_type = 'chat_generation_task' AND task_id IS NOT NULL AND generation_record_id IS NULL) "
|
||||
"OR (owner_type = 'generation_record' AND task_id IS NULL AND generation_record_id IS NOT NULL)",
|
||||
)
|
||||
|
||||
|
||||
def _add_chat_generation_task_columns() -> None:
|
||||
op.add_column(
|
||||
"chat_generation_tasks",
|
||||
sa.Column(
|
||||
"generation_attempt_no",
|
||||
sa.Integer(),
|
||||
server_default="1",
|
||||
nullable=False,
|
||||
),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_generation_tasks",
|
||||
sa.Column(
|
||||
"resource_generation_started_at",
|
||||
sa.DateTime(timezone=True),
|
||||
nullable=True,
|
||||
),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_generation_tasks",
|
||||
sa.Column(
|
||||
"manual_retry_count",
|
||||
sa.Integer(),
|
||||
server_default="0",
|
||||
nullable=False,
|
||||
),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_generation_tasks",
|
||||
sa.Column(
|
||||
"poll_error_count",
|
||||
sa.Integer(),
|
||||
server_default="0",
|
||||
nullable=False,
|
||||
),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_generation_tasks",
|
||||
sa.Column("poll_claim_token", sa.String(length=64), nullable=True),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_generation_tasks",
|
||||
sa.Column("poll_lease_until", sa.DateTime(timezone=True), nullable=True),
|
||||
)
|
||||
op.add_column(
|
||||
"chat_generation_tasks",
|
||||
sa.Column("download_claim_token", sa.String(length=64), nullable=True),
|
||||
)
|
||||
|
||||
op.create_index(
|
||||
op.f("ix_chat_generation_tasks_download_claim_token"),
|
||||
"chat_generation_tasks",
|
||||
["download_claim_token"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
op.f("ix_chat_generation_tasks_poll_claim_token"),
|
||||
"chat_generation_tasks",
|
||||
["poll_claim_token"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
op.f("ix_chat_generation_tasks_poll_lease_until"),
|
||||
"chat_generation_tasks",
|
||||
["poll_lease_until"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
op.f("ix_chat_generation_tasks_resource_generation_started_at"),
|
||||
"chat_generation_tasks",
|
||||
["resource_generation_started_at"],
|
||||
unique=False,
|
||||
)
|
||||
|
||||
|
||||
def _add_generation_record_columns() -> None:
|
||||
columns = [
|
||||
sa.Column("generation_attempt_no", sa.Integer(), server_default="1", nullable=False),
|
||||
sa.Column("resource_generation_started_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("deadline_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("engine_id", sa.String(length=32), nullable=True),
|
||||
sa.Column("engine_snapshot_json", sa.Text(), nullable=True),
|
||||
sa.Column("provider_response_json", sa.Text(), nullable=True),
|
||||
sa.Column("remote_result_url", sa.Text(), nullable=True),
|
||||
sa.Column("provider_create_claim_token", sa.String(length=64), nullable=True),
|
||||
sa.Column("provider_create_lease_until", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("provider_create_started_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("retry_count", sa.Integer(), server_default="0", nullable=False),
|
||||
sa.Column("manual_retry_count", sa.Integer(), server_default="0", nullable=False),
|
||||
sa.Column("poll_error_count", sa.Integer(), server_default="0", nullable=False),
|
||||
sa.Column("poll_count", sa.Integer(), server_default="0", nullable=False),
|
||||
sa.Column("last_poll_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("poll_started_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("next_poll_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("poll_interval_seconds", sa.Integer(), server_default="0", nullable=False),
|
||||
sa.Column("poll_claim_token", sa.String(length=64), nullable=True),
|
||||
sa.Column("poll_lease_until", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("download_celery_task_id", sa.String(length=160), nullable=True),
|
||||
sa.Column("download_enqueued_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("download_started_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("download_claim_token", sa.String(length=64), nullable=True),
|
||||
sa.Column("download_lease_until", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("download_next_retry_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("download_attempt_count", sa.Integer(), server_default="0", nullable=False),
|
||||
sa.Column("download_last_error", sa.Text(), nullable=True),
|
||||
sa.Column("download_storage_date_dir", sa.String(length=16), nullable=True),
|
||||
]
|
||||
for column in columns:
|
||||
op.add_column("generation_records", column)
|
||||
|
||||
op.create_index(
|
||||
"idx_genrec_next_poll_at",
|
||||
"generation_records",
|
||||
["next_poll_at"],
|
||||
unique=False,
|
||||
postgresql_where=sa.text(
|
||||
"deleted_at IS NULL AND status = 'generating' "
|
||||
"AND gen_type = 'video' AND next_poll_at IS NOT NULL"
|
||||
),
|
||||
)
|
||||
for column_name in (
|
||||
"deadline_at",
|
||||
"download_celery_task_id",
|
||||
"download_claim_token",
|
||||
"download_lease_until",
|
||||
"download_next_retry_at",
|
||||
"engine_id",
|
||||
"poll_claim_token",
|
||||
"poll_lease_until",
|
||||
"provider_create_claim_token",
|
||||
"provider_create_lease_until",
|
||||
"resource_generation_started_at",
|
||||
):
|
||||
op.create_index(
|
||||
op.f(f"ix_generation_records_{column_name}"),
|
||||
"generation_records",
|
||||
[column_name],
|
||||
unique=False,
|
||||
)
|
||||
|
||||
|
||||
def _backfill_generation_attempts() -> None:
|
||||
# Existing ChatGenerationTask rows started resource generation when the row
|
||||
# was created. Do not derive this timestamp from poll/download timestamps.
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE chat_generation_tasks
|
||||
SET resource_generation_started_at = created_at
|
||||
WHERE resource_generation_started_at IS NULL
|
||||
"""
|
||||
)
|
||||
|
||||
# Prefer structured owner/attempt fields. The biz_key parser is retained
|
||||
# for older credit rows that were written before those columns were filled.
|
||||
op.execute(
|
||||
"""
|
||||
WITH credit_attempts AS (
|
||||
SELECT
|
||||
COALESCE(
|
||||
NULLIF(owner_id, ''),
|
||||
NULLIF(related_id, ''),
|
||||
substring(biz_key FROM '^chat_generation_task:([^:]+):attempt:')
|
||||
) AS resolved_owner_id,
|
||||
MAX(
|
||||
COALESCE(
|
||||
attempt_no,
|
||||
NULLIF(substring(biz_key FROM ':attempt:([0-9]+):media:charge$'), '')::integer
|
||||
)
|
||||
) AS max_attempt
|
||||
FROM credit_records
|
||||
WHERE type = 'consume'
|
||||
AND (
|
||||
(owner_type = 'chat_generation_task' AND charge_kind = 'media')
|
||||
OR biz_key ~ '^chat_generation_task:[^:]+:attempt:[0-9]+:media:charge$'
|
||||
)
|
||||
GROUP BY 1
|
||||
)
|
||||
UPDATE chat_generation_tasks AS task
|
||||
SET generation_attempt_no = GREATEST(1, attempts.max_attempt)
|
||||
FROM credit_attempts AS attempts
|
||||
WHERE attempts.resolved_owner_id = task.id
|
||||
AND attempts.max_attempt IS NOT NULL
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE chat_generation_tasks
|
||||
SET manual_retry_count = GREATEST(generation_attempt_no - 1, 0),
|
||||
retry_count = GREATEST(generation_attempt_no - 1, 0),
|
||||
poll_error_count = 0
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute(
|
||||
"""
|
||||
WITH generation_charges AS (
|
||||
SELECT
|
||||
resolved_owner_id,
|
||||
resolved_attempt,
|
||||
charge_created_at,
|
||||
engine_id
|
||||
FROM (
|
||||
SELECT
|
||||
COALESCE(
|
||||
NULLIF(owner_id, ''),
|
||||
NULLIF(related_id, ''),
|
||||
substring(biz_key FROM '^generation_record:([^:]+):attempt:')
|
||||
) AS resolved_owner_id,
|
||||
COALESCE(
|
||||
attempt_no,
|
||||
NULLIF(substring(biz_key FROM ':attempt:([0-9]+):media:charge$'), '')::integer,
|
||||
1
|
||||
) AS resolved_attempt,
|
||||
created_at AS charge_created_at,
|
||||
engine_id,
|
||||
ROW_NUMBER() OVER (
|
||||
PARTITION BY COALESCE(
|
||||
NULLIF(owner_id, ''),
|
||||
NULLIF(related_id, ''),
|
||||
substring(biz_key FROM '^generation_record:([^:]+):attempt:')
|
||||
)
|
||||
ORDER BY
|
||||
COALESCE(
|
||||
attempt_no,
|
||||
NULLIF(substring(biz_key FROM ':attempt:([0-9]+):media:charge$'), '')::integer,
|
||||
1
|
||||
) DESC,
|
||||
created_at DESC,
|
||||
id DESC
|
||||
) AS row_no
|
||||
FROM credit_records
|
||||
WHERE type = 'consume'
|
||||
AND (
|
||||
(owner_type = 'generation_record' AND charge_kind = 'media')
|
||||
OR biz_key ~ '^generation_record:[^:]+:attempt:[0-9]+:media:charge$'
|
||||
)
|
||||
) AS ranked
|
||||
WHERE row_no = 1
|
||||
AND resolved_owner_id IS NOT NULL
|
||||
)
|
||||
UPDATE generation_records AS record
|
||||
SET generation_attempt_no = GREATEST(1, charge.resolved_attempt),
|
||||
resource_generation_started_at = COALESCE(
|
||||
record.resource_generation_started_at,
|
||||
charge.charge_created_at
|
||||
),
|
||||
engine_id = COALESCE(record.engine_id, charge.engine_id)
|
||||
FROM generation_charges AS charge
|
||||
WHERE charge.resolved_owner_id = record.id
|
||||
"""
|
||||
)
|
||||
|
||||
# Old GenerationRecord rows do not have a dedicated resource-start field.
|
||||
# Only rows with clear resource-generation evidence are backfilled; prompt-
|
||||
# optimized-only records intentionally remain NULL.
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE generation_records
|
||||
SET resource_generation_started_at = COALESCE(updated_at, created_at)
|
||||
WHERE resource_generation_started_at IS NULL
|
||||
AND (
|
||||
status IN ('generating', 'completed')
|
||||
OR seedance_task_id IS NOT NULL
|
||||
OR image_url IS NOT NULL
|
||||
OR video_url IS NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE generation_records
|
||||
SET manual_retry_count = GREATEST(generation_attempt_no - 1, 0),
|
||||
retry_count = GREATEST(generation_attempt_no - 1, 0),
|
||||
poll_error_count = 0
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _backfill_generation_record_engine() -> None:
|
||||
# For recoverable active rows without a historical billing engine snapshot,
|
||||
# fall back to the current highest-priority active engine of the same type.
|
||||
# Completed/failed history is not assigned a guessed engine.
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE generation_records AS record
|
||||
SET engine_id = engine.id
|
||||
FROM (
|
||||
SELECT id
|
||||
FROM image_engines
|
||||
WHERE is_active IS TRUE AND deleted_at IS NULL
|
||||
ORDER BY priority DESC, id ASC
|
||||
LIMIT 1
|
||||
) AS engine
|
||||
WHERE record.engine_id IS NULL
|
||||
AND record.status = 'generating'
|
||||
AND record.gen_type = 'image'
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE generation_records AS record
|
||||
SET engine_id = engine.id
|
||||
FROM (
|
||||
SELECT id
|
||||
FROM video_engines
|
||||
WHERE is_active IS TRUE AND deleted_at IS NULL
|
||||
ORDER BY priority DESC, id ASC
|
||||
LIMIT 1
|
||||
) AS engine
|
||||
WHERE record.engine_id IS NULL
|
||||
AND record.status = 'generating'
|
||||
AND record.gen_type = 'video'
|
||||
"""
|
||||
)
|
||||
|
||||
# Store a key-free execution snapshot. Runtime code still reads api_key
|
||||
# from the engine row by engine_id, including soft-deleted historical rows.
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE generation_records AS record
|
||||
SET engine_snapshot_json = jsonb_build_object(
|
||||
'engine_type', 'image',
|
||||
'id', engine.id,
|
||||
'name', engine.name,
|
||||
'provider', engine.provider,
|
||||
'api_base', engine.api_base,
|
||||
'api_key_masked', CASE WHEN COALESCE(engine.api_key, '') <> '' THEN '****' ELSE '' END,
|
||||
'model_name', engine.model_name,
|
||||
'generate_url', engine.generate_url,
|
||||
'default_size', engine.default_size,
|
||||
'multi_generation_enabled', engine.multi_generation_enabled,
|
||||
'max_generation_count', engine.max_generation_count,
|
||||
'multi_image_max_images', engine.multi_image_max_images,
|
||||
'max_reference_image_count', engine.max_reference_image_count,
|
||||
'output_format', engine.output_format,
|
||||
'selected_size', COALESCE(record.image_size, engine.default_size, '2K'),
|
||||
'selected_proportion', COALESCE(record.image_proportion, '1:1'),
|
||||
'selected_px', COALESCE(record.image_px, '2048x2048')
|
||||
)::text
|
||||
FROM image_engines AS engine
|
||||
WHERE record.gen_type = 'image'
|
||||
AND record.engine_id = engine.id
|
||||
AND record.engine_snapshot_json IS NULL
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE generation_records AS record
|
||||
SET engine_snapshot_json = jsonb_build_object(
|
||||
'engine_type', 'video',
|
||||
'id', engine.id,
|
||||
'name', engine.name,
|
||||
'provider', engine.provider,
|
||||
'api_base', engine.api_base,
|
||||
'api_key_masked', CASE WHEN COALESCE(engine.api_key, '') <> '' THEN '****' ELSE '' END,
|
||||
'model_name', engine.model_name,
|
||||
'generate_url', engine.generate_url,
|
||||
'query_url', engine.query_url,
|
||||
'max_duration', engine.max_duration,
|
||||
'max_audio_count', engine.max_audio_count,
|
||||
'multi_generation_enabled', engine.multi_generation_enabled,
|
||||
'max_generation_count', engine.max_generation_count,
|
||||
'selected_ratio', COALESCE(record.aspect_ratio, '16:9'),
|
||||
'selected_resolution', COALESCE(record.resolution, '480p'),
|
||||
'selected_duration', COALESCE(record.duration, 4)
|
||||
)::text
|
||||
FROM video_engines AS engine
|
||||
WHERE record.gen_type = 'video'
|
||||
AND record.engine_id = engine.id
|
||||
AND record.engine_snapshot_json IS NULL
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _backfill_generation_pipeline_state() -> None:
|
||||
# Preserve absolute historical provider URLs when they are available.
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE generation_records
|
||||
SET remote_result_url = CASE
|
||||
WHEN gen_type = 'image' AND image_url ~ '^https?://' THEN image_url
|
||||
WHEN gen_type = 'video' AND video_url ~ '^https?://' THEN video_url
|
||||
ELSE remote_result_url
|
||||
END
|
||||
WHERE remote_result_url IS NULL
|
||||
AND (
|
||||
(gen_type = 'image' AND image_url ~ '^https?://')
|
||||
OR (gen_type = 'video' AND video_url ~ '^https?://')
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE generation_records
|
||||
SET pipeline_stage = CASE
|
||||
WHEN remote_result_url IS NOT NULL THEN 'result_ready'
|
||||
WHEN seedance_task_id IS NOT NULL THEN 'waiting_remote'
|
||||
ELSE 'queued'
|
||||
END
|
||||
WHERE status = 'generating'
|
||||
AND pipeline_stage IS NULL
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE generation_records
|
||||
SET pipeline_stage = 'done'
|
||||
WHERE status = 'completed'
|
||||
AND pipeline_stage IS NULL
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE generation_records
|
||||
SET pipeline_stage = 'failed'
|
||||
WHERE status = 'failed'
|
||||
AND pipeline_stage IS NULL
|
||||
"""
|
||||
)
|
||||
|
||||
# Image deadline is now 30 minutes; video remains 24 hours. Only active
|
||||
# tasks are rewritten so terminal historical audit values are preserved.
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE chat_generation_tasks
|
||||
SET deadline_at = resource_generation_started_at + INTERVAL '30 minutes'
|
||||
WHERE status = 'generating'
|
||||
AND gen_type = 'image'
|
||||
AND resource_generation_started_at IS NOT NULL
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE chat_generation_tasks
|
||||
SET deadline_at = resource_generation_started_at + INTERVAL '24 hours'
|
||||
WHERE status = 'generating'
|
||||
AND gen_type = 'video'
|
||||
AND resource_generation_started_at IS NOT NULL
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE generation_records
|
||||
SET deadline_at = resource_generation_started_at
|
||||
+ CASE WHEN gen_type = 'image' THEN INTERVAL '30 minutes' ELSE INTERVAL '24 hours' END
|
||||
WHERE status = 'generating'
|
||||
AND resource_generation_started_at IS NOT NULL
|
||||
"""
|
||||
)
|
||||
|
||||
# Active provider tasks are checked immediately after deployment. The
|
||||
# recovery code still honours Redis execution locks and poll leases.
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE generation_records
|
||||
SET poll_started_at = COALESCE(poll_started_at, resource_generation_started_at),
|
||||
next_poll_at = CURRENT_TIMESTAMP
|
||||
WHERE status = 'generating'
|
||||
AND gen_type = 'video'
|
||||
AND seedance_task_id IS NOT NULL
|
||||
AND remote_result_url IS NULL
|
||||
AND pipeline_stage IN ('waiting_remote', 'polling')
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _create_generated_resource_idempotency_index() -> None:
|
||||
# Do not silently soft-delete duplicates here: doing so without rebuilding
|
||||
# user_resource_*_stats would corrupt capacity totals. Abort with a clear
|
||||
# message so dirty data can be repaired and stats rebuilt deliberately.
|
||||
op.execute(
|
||||
"""
|
||||
DO $$
|
||||
BEGIN
|
||||
IF EXISTS (
|
||||
SELECT 1
|
||||
FROM generated_resources
|
||||
WHERE deleted_at IS NULL
|
||||
GROUP BY source_model, source_id, resource_type
|
||||
HAVING COUNT(*) > 1
|
||||
) THEN
|
||||
RAISE EXCEPTION
|
||||
'Cannot create uq_generated_resources_active_source_type: duplicate active generated_resources exist'
|
||||
USING HINT = 'Query duplicates by source_model/source_id/resource_type, keep one active row, soft-delete the others, then rebuild user resource stats before rerunning this revision.';
|
||||
END IF;
|
||||
END
|
||||
$$;
|
||||
"""
|
||||
)
|
||||
op.create_index(
|
||||
"uq_generated_resources_active_source_type",
|
||||
"generated_resources",
|
||||
["source_model", "source_id", "resource_type"],
|
||||
unique=True,
|
||||
postgresql_where=sa.text("deleted_at IS NULL"),
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
_add_engine_soft_delete_columns()
|
||||
_add_shared_log_owner_columns()
|
||||
_add_chat_generation_task_columns()
|
||||
_add_generation_record_columns()
|
||||
|
||||
_backfill_generation_attempts()
|
||||
_backfill_generation_record_engine()
|
||||
_backfill_generation_pipeline_state()
|
||||
_create_generated_resource_idempotency_index()
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index(
|
||||
"uq_generated_resources_active_source_type",
|
||||
table_name="generated_resources",
|
||||
postgresql_where=sa.text("deleted_at IS NULL"),
|
||||
)
|
||||
|
||||
for column_name in (
|
||||
"resource_generation_started_at",
|
||||
"provider_create_lease_until",
|
||||
"provider_create_claim_token",
|
||||
"poll_lease_until",
|
||||
"poll_claim_token",
|
||||
"engine_id",
|
||||
"download_next_retry_at",
|
||||
"download_lease_until",
|
||||
"download_claim_token",
|
||||
"download_celery_task_id",
|
||||
"deadline_at",
|
||||
):
|
||||
op.drop_index(
|
||||
op.f(f"ix_generation_records_{column_name}"),
|
||||
table_name="generation_records",
|
||||
)
|
||||
op.drop_index("idx_genrec_next_poll_at", table_name="generation_records")
|
||||
|
||||
for column_name in (
|
||||
"download_storage_date_dir",
|
||||
"download_last_error",
|
||||
"download_attempt_count",
|
||||
"download_next_retry_at",
|
||||
"download_lease_until",
|
||||
"download_claim_token",
|
||||
"download_started_at",
|
||||
"download_enqueued_at",
|
||||
"download_celery_task_id",
|
||||
"poll_lease_until",
|
||||
"poll_claim_token",
|
||||
"poll_interval_seconds",
|
||||
"next_poll_at",
|
||||
"poll_started_at",
|
||||
"last_poll_at",
|
||||
"poll_count",
|
||||
"poll_error_count",
|
||||
"manual_retry_count",
|
||||
"retry_count",
|
||||
"provider_create_started_at",
|
||||
"provider_create_lease_until",
|
||||
"provider_create_claim_token",
|
||||
"remote_result_url",
|
||||
"provider_response_json",
|
||||
"engine_snapshot_json",
|
||||
"engine_id",
|
||||
"deadline_at",
|
||||
"resource_generation_started_at",
|
||||
"generation_attempt_no",
|
||||
):
|
||||
op.drop_column("generation_records", column_name)
|
||||
|
||||
# GenerationRecord log rows cannot be represented by the old task_id-only
|
||||
# schema, so they are removed during downgrade before task_id becomes NOT NULL.
|
||||
op.execute(
|
||||
"DELETE FROM chat_provider_call_logs "
|
||||
"WHERE owner_type = 'generation_record' OR generation_record_id IS NOT NULL"
|
||||
)
|
||||
op.drop_constraint(CK_CALL_LOG_OWNER, "chat_provider_call_logs", type_="check")
|
||||
op.drop_constraint(
|
||||
FK_CALL_LOG_GENERATION_RECORD,
|
||||
"chat_provider_call_logs",
|
||||
type_="foreignkey",
|
||||
)
|
||||
op.drop_index(
|
||||
op.f("ix_chat_provider_call_logs_owner_type"),
|
||||
table_name="chat_provider_call_logs",
|
||||
)
|
||||
op.drop_index(
|
||||
op.f("ix_chat_provider_call_logs_generation_record_id"),
|
||||
table_name="chat_provider_call_logs",
|
||||
)
|
||||
op.drop_index(
|
||||
op.f("ix_chat_provider_call_logs_generation_attempt_no"),
|
||||
table_name="chat_provider_call_logs",
|
||||
)
|
||||
op.drop_index("idx_chat_provider_call_logs_task_created", table_name="chat_provider_call_logs")
|
||||
op.drop_index("idx_chat_provider_call_logs_record_created", table_name="chat_provider_call_logs")
|
||||
op.drop_index("idx_chat_provider_call_logs_attempt_created", table_name="chat_provider_call_logs")
|
||||
op.alter_column(
|
||||
"chat_provider_call_logs",
|
||||
"task_id",
|
||||
existing_type=sa.VARCHAR(length=32),
|
||||
nullable=False,
|
||||
)
|
||||
op.drop_column("chat_provider_call_logs", "generation_attempt_no")
|
||||
op.drop_column("chat_provider_call_logs", "generation_record_id")
|
||||
op.drop_column("chat_provider_call_logs", "owner_type")
|
||||
|
||||
op.execute(
|
||||
"DELETE FROM chat_generation_task_events "
|
||||
"WHERE owner_type = 'generation_record' OR generation_record_id IS NOT NULL"
|
||||
)
|
||||
op.drop_constraint(CK_EVENT_OWNER, "chat_generation_task_events", type_="check")
|
||||
op.drop_constraint(
|
||||
FK_EVENT_GENERATION_RECORD,
|
||||
"chat_generation_task_events",
|
||||
type_="foreignkey",
|
||||
)
|
||||
op.drop_index(
|
||||
op.f("ix_chat_generation_task_events_owner_type"),
|
||||
table_name="chat_generation_task_events",
|
||||
)
|
||||
op.drop_index(
|
||||
op.f("ix_chat_generation_task_events_generation_record_id"),
|
||||
table_name="chat_generation_task_events",
|
||||
)
|
||||
op.drop_index(
|
||||
op.f("ix_chat_generation_task_events_generation_attempt_no"),
|
||||
table_name="chat_generation_task_events",
|
||||
)
|
||||
op.drop_index("idx_chat_generation_task_events_task_created", table_name="chat_generation_task_events")
|
||||
op.drop_index("idx_chat_generation_task_events_record_created", table_name="chat_generation_task_events")
|
||||
op.drop_index("idx_chat_generation_task_events_attempt_created", table_name="chat_generation_task_events")
|
||||
op.execute(
|
||||
"UPDATE chat_generation_task_events "
|
||||
"SET message = LEFT(message, 512), "
|
||||
"from_stage = LEFT(from_stage, 32), "
|
||||
"to_stage = LEFT(to_stage, 32)"
|
||||
)
|
||||
op.alter_column(
|
||||
"chat_generation_task_events",
|
||||
"message",
|
||||
existing_type=sa.Text(),
|
||||
type_=sa.VARCHAR(length=512),
|
||||
existing_nullable=True,
|
||||
)
|
||||
op.alter_column(
|
||||
"chat_generation_task_events",
|
||||
"to_stage",
|
||||
existing_type=sa.String(length=48),
|
||||
type_=sa.VARCHAR(length=32),
|
||||
existing_nullable=True,
|
||||
)
|
||||
op.alter_column(
|
||||
"chat_generation_task_events",
|
||||
"from_stage",
|
||||
existing_type=sa.String(length=48),
|
||||
type_=sa.VARCHAR(length=32),
|
||||
existing_nullable=True,
|
||||
)
|
||||
op.alter_column(
|
||||
"chat_generation_task_events",
|
||||
"task_id",
|
||||
existing_type=sa.VARCHAR(length=32),
|
||||
nullable=False,
|
||||
)
|
||||
op.drop_column("chat_generation_task_events", "generation_attempt_no")
|
||||
op.drop_column("chat_generation_task_events", "generation_record_id")
|
||||
op.drop_column("chat_generation_task_events", "owner_type")
|
||||
|
||||
op.drop_index(
|
||||
op.f("ix_chat_generation_tasks_resource_generation_started_at"),
|
||||
table_name="chat_generation_tasks",
|
||||
)
|
||||
op.drop_index(
|
||||
op.f("ix_chat_generation_tasks_poll_lease_until"),
|
||||
table_name="chat_generation_tasks",
|
||||
)
|
||||
op.drop_index(
|
||||
op.f("ix_chat_generation_tasks_poll_claim_token"),
|
||||
table_name="chat_generation_tasks",
|
||||
)
|
||||
op.drop_index(
|
||||
op.f("ix_chat_generation_tasks_download_claim_token"),
|
||||
table_name="chat_generation_tasks",
|
||||
)
|
||||
for column_name in (
|
||||
"download_claim_token",
|
||||
"poll_lease_until",
|
||||
"poll_claim_token",
|
||||
"poll_error_count",
|
||||
"manual_retry_count",
|
||||
"resource_generation_started_at",
|
||||
"generation_attempt_no",
|
||||
):
|
||||
op.drop_column("chat_generation_tasks", column_name)
|
||||
|
||||
for table_name in ("video_engines", "model_configs", "image_engines"):
|
||||
op.drop_index(op.f(f"ix_{table_name}_deleted_at"), table_name=table_name)
|
||||
op.drop_column(table_name, "deleted_at")
|
||||
Reference in New Issue
Block a user