From c0a58fdac392d4e6d38aa3d3717555b73c8cd3b1 Mon Sep 17 00:00:00 2001 From: GinHa <15201596918@163.com> Date: Mon, 20 Jul 2026 14:11:45 +0800 Subject: [PATCH] =?UTF-8?q?=E9=A1=B9=E7=9B=AE/AI=E7=94=9F=E6=88=90?= =?UTF-8?q?=E9=93=BE=E8=B7=AF=E5=90=88=E5=B9=B6=20=E7=89=88=E6=9C=AC?= =?UTF-8?q?=E8=BF=81=E7=A7=BB=E4=BF=AE=E5=A4=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...nify_generation_owners_and_soft_delete_.py | 604 ++++++++++++++---- 1 file changed, 479 insertions(+), 125 deletions(-) diff --git a/video-gen-api/alembic/versions/e6eac828ff61_unify_generation_owners_and_soft_delete_.py b/video-gen-api/alembic/versions/e6eac828ff61_unify_generation_owners_and_soft_delete_.py index aa5f7eae..0164e1fa 100644 --- a/video-gen-api/alembic/versions/e6eac828ff61_unify_generation_owners_and_soft_delete_.py +++ b/video-gen-api/alembic/versions/e6eac828ff61_unify_generation_owners_and_soft_delete_.py @@ -7,10 +7,11 @@ 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 typing import Any, Iterable, Mapping, Sequence, Union from alembic import op import sqlalchemy as sa +from sqlalchemy.engine.reflection import Inspector # revision identifiers, used by Alembic. @@ -26,13 +27,255 @@ CK_EVENT_OWNER = "ck_chat_generation_task_events_owner" CK_CALL_LOG_OWNER = "ck_chat_provider_call_logs_owner" +# This revision may be deployed to databases where a previous manual patch or +# interrupted release already created part of the target schema. PostgreSQL +# normally runs Alembic DDL transactionally, but these guards deliberately make +# every schema operation tolerant of existing/missing objects. +def _inspect() -> Inspector: + return sa.inspect(op.get_bind()) + + +def _table_exists(table_name: str) -> bool: + return bool(_inspect().has_table(table_name)) + + +def _column_info(table_name: str, column_name: str) -> dict[str, Any] | None: + if not _table_exists(table_name): + return None + for column in _inspect().get_columns(table_name): + if column.get("name") == column_name: + return column + return None + + +def _column_exists(table_name: str, column_name: str) -> bool: + return _column_info(table_name, column_name) is not None + + +def _columns_exist(table_name: str, column_names: Iterable[str]) -> bool: + if not _table_exists(table_name): + return False + existing = {column.get("name") for column in _inspect().get_columns(table_name)} + return all(column_name in existing for column_name in column_names) + + +def _index_exists(table_name: str, index_name: str) -> bool: + if not _table_exists(table_name): + return False + expected = str(index_name) + return any(index.get("name") == expected for index in _inspect().get_indexes(table_name)) + + +def _foreign_key_exists( + table_name: str, + constraint_name: str, + local_columns: Sequence[str] | None = None, + referent_table: str | None = None, + remote_columns: Sequence[str] | None = None, +) -> bool: + if not _table_exists(table_name): + return False + expected_name = str(constraint_name) + expected_local = list(local_columns or ()) + expected_remote = list(remote_columns or ()) + for foreign_key in _inspect().get_foreign_keys(table_name): + if foreign_key.get("name") == expected_name: + return True + if expected_local and list(foreign_key.get("constrained_columns") or ()) != expected_local: + continue + if referent_table and foreign_key.get("referred_table") != referent_table: + continue + if expected_remote and list(foreign_key.get("referred_columns") or ()) != expected_remote: + continue + if expected_local or referent_table or expected_remote: + return True + return False + + +def _check_constraint_exists(table_name: str, constraint_name: str) -> bool: + if not _table_exists(table_name): + return False + expected = str(constraint_name) + return any( + constraint.get("name") == expected + for constraint in _inspect().get_check_constraints(table_name) + ) + + +def _constraint_exists(table_name: str, constraint_name: str, type_: str | None) -> bool: + if not _table_exists(table_name): + return False + expected = str(constraint_name) + normalized = (type_ or "").lower() + if normalized == "foreignkey": + return any(item.get("name") == expected for item in _inspect().get_foreign_keys(table_name)) + if normalized == "check": + return any(item.get("name") == expected for item in _inspect().get_check_constraints(table_name)) + if normalized == "unique": + return any(item.get("name") == expected for item in _inspect().get_unique_constraints(table_name)) + if normalized == "primary": + primary_key = _inspect().get_pk_constraint(table_name) + return primary_key.get("name") == expected + + return ( + any(item.get("name") == expected for item in _inspect().get_foreign_keys(table_name)) + or any(item.get("name") == expected for item in _inspect().get_check_constraints(table_name)) + or any(item.get("name") == expected for item in _inspect().get_unique_constraints(table_name)) + or _inspect().get_pk_constraint(table_name).get("name") == expected + ) + + +def _normalize_sql_type(type_: sa.types.TypeEngine[Any]) -> str: + compiled = type_.compile(dialect=op.get_bind().dialect) + return " ".join(str(compiled).lower().replace("character varying", "varchar").split()) + + +def _add_column_if_missing( + table_name: str, + column: sa.Column[Any], + *, + schema: str | None = None, + **kwargs: Any, +) -> None: + if _column_exists(table_name, str(column.name)): + return + op.add_column(table_name, column, schema=schema, **kwargs) + + +def _alter_column_if_exists( + table_name: str, + column_name: str, + *, + schema: str | None = None, + **kwargs: Any, +) -> None: + current = _column_info(table_name, column_name) + if current is None: + return + + needs_alter = False + if "nullable" in kwargs and bool(current.get("nullable")) != bool(kwargs["nullable"]): + needs_alter = True + if kwargs.get("type_") is not None: + current_type = current.get("type") + target_type = kwargs["type_"] + if current_type is None or _normalize_sql_type(current_type) != _normalize_sql_type(target_type): + needs_alter = True + if "server_default" in kwargs: + current_default = current.get("default") + target_default = kwargs["server_default"] + if str(current_default) != str(target_default): + needs_alter = True + if "comment" in kwargs and current.get("comment") != kwargs["comment"]: + needs_alter = True + + if not needs_alter: + return + op.alter_column(table_name, column_name, schema=schema, **kwargs) + + +def _create_index_if_missing( + index_name: str, + table_name: str, + columns: Sequence[str | sa.sql.elements.TextClause], + **kwargs: Any, +) -> None: + if not _table_exists(table_name) or _index_exists(table_name, str(index_name)): + return + op.create_index(index_name, table_name, columns, **kwargs) + + +def _drop_index_if_exists( + index_name: str, + *, + table_name: str | None = None, + **kwargs: Any, +) -> None: + if not table_name or not _index_exists(table_name, str(index_name)): + return + op.drop_index(index_name, table_name=table_name, **kwargs) + + +def _create_foreign_key_if_missing( + constraint_name: str, + source_table: str, + referent_table: str, + local_cols: Sequence[str], + remote_cols: Sequence[str], + **kwargs: Any, +) -> None: + if not _table_exists(source_table) or not _table_exists(referent_table): + return + if not _columns_exist(source_table, local_cols) or not _columns_exist(referent_table, remote_cols): + return + if _foreign_key_exists( + source_table, + constraint_name, + local_columns=local_cols, + referent_table=referent_table, + remote_columns=remote_cols, + ): + return + op.create_foreign_key( + constraint_name, + source_table, + referent_table, + local_cols, + remote_cols, + **kwargs, + ) + + +def _create_check_constraint_if_missing( + constraint_name: str, + table_name: str, + condition: str | sa.sql.elements.TextClause, + **kwargs: Any, +) -> None: + if not _table_exists(table_name) or _check_constraint_exists(table_name, constraint_name): + return + op.create_check_constraint(constraint_name, table_name, condition, **kwargs) + + +def _drop_constraint_if_exists( + constraint_name: str, + table_name: str, + *, + type_: str | None = None, + **kwargs: Any, +) -> None: + if not _constraint_exists(table_name, constraint_name, type_): + return + op.drop_constraint(constraint_name, table_name, type_=type_, **kwargs) + + +def _drop_column_if_exists( + table_name: str, + column_name: str, + *, + schema: str | None = None, + **kwargs: Any, +) -> None: + if not _column_exists(table_name, column_name): + return + op.drop_column(table_name, column_name, schema=schema, **kwargs) + + +def _execute_if_columns_exist( + statement: str, + requirements: Mapping[str, Sequence[str]], +) -> None: + if all(_columns_exist(table_name, columns) for table_name, columns in requirements.items()): + op.execute(sa.text(statement)) + + def _add_engine_soft_delete_columns() -> None: for table_name in ("image_engines", "model_configs", "video_engines"): - op.add_column( + _add_column_if_missing( table_name, sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True), ) - op.create_index( + _create_index_if_missing( op.f(f"ix_{table_name}_deleted_at"), table_name, ["deleted_at"], @@ -41,7 +284,7 @@ def _add_engine_soft_delete_columns() -> None: def _add_shared_log_owner_columns() -> None: - op.add_column( + _add_column_if_missing( "chat_generation_task_events", sa.Column( "owner_type", @@ -50,11 +293,11 @@ def _add_shared_log_owner_columns() -> None: nullable=False, ), ) - op.add_column( + _add_column_if_missing( "chat_generation_task_events", sa.Column("generation_record_id", sa.String(length=32), nullable=True), ) - op.add_column( + _add_column_if_missing( "chat_generation_task_events", sa.Column( "generation_attempt_no", @@ -63,70 +306,70 @@ def _add_shared_log_owner_columns() -> None: nullable=False, ), ) - op.alter_column( + _alter_column_if_exists( "chat_generation_task_events", "task_id", existing_type=sa.VARCHAR(length=32), nullable=True, ) - op.alter_column( + _alter_column_if_exists( "chat_generation_task_events", "from_stage", existing_type=sa.VARCHAR(length=32), type_=sa.String(length=48), existing_nullable=True, ) - op.alter_column( + _alter_column_if_exists( "chat_generation_task_events", "to_stage", existing_type=sa.VARCHAR(length=32), type_=sa.String(length=48), existing_nullable=True, ) - op.alter_column( + _alter_column_if_exists( "chat_generation_task_events", "message", existing_type=sa.VARCHAR(length=512), type_=sa.Text(), existing_nullable=True, ) - op.create_index( + _create_index_if_missing( "idx_chat_generation_task_events_attempt_created", "chat_generation_task_events", ["owner_type", "generation_attempt_no", "created_at"], unique=False, ) - op.create_index( + _create_index_if_missing( "idx_chat_generation_task_events_record_created", "chat_generation_task_events", ["owner_type", "generation_record_id", "created_at"], unique=False, ) - op.create_index( + _create_index_if_missing( "idx_chat_generation_task_events_task_created", "chat_generation_task_events", ["owner_type", "task_id", "created_at"], unique=False, ) - op.create_index( + _create_index_if_missing( op.f("ix_chat_generation_task_events_generation_attempt_no"), "chat_generation_task_events", ["generation_attempt_no"], unique=False, ) - op.create_index( + _create_index_if_missing( op.f("ix_chat_generation_task_events_generation_record_id"), "chat_generation_task_events", ["generation_record_id"], unique=False, ) - op.create_index( + _create_index_if_missing( op.f("ix_chat_generation_task_events_owner_type"), "chat_generation_task_events", ["owner_type"], unique=False, ) - op.create_foreign_key( + _create_foreign_key_if_missing( FK_EVENT_GENERATION_RECORD, "chat_generation_task_events", "generation_records", @@ -134,14 +377,14 @@ def _add_shared_log_owner_columns() -> None: ["id"], ondelete="CASCADE", ) - op.create_check_constraint( + _create_check_constraint_if_missing( 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( + _add_column_if_missing( "chat_provider_call_logs", sa.Column( "owner_type", @@ -150,11 +393,11 @@ def _add_shared_log_owner_columns() -> None: nullable=False, ), ) - op.add_column( + _add_column_if_missing( "chat_provider_call_logs", sa.Column("generation_record_id", sa.String(length=32), nullable=True), ) - op.add_column( + _add_column_if_missing( "chat_provider_call_logs", sa.Column( "generation_attempt_no", @@ -163,49 +406,49 @@ def _add_shared_log_owner_columns() -> None: nullable=False, ), ) - op.alter_column( + _alter_column_if_exists( "chat_provider_call_logs", "task_id", existing_type=sa.VARCHAR(length=32), nullable=True, ) - op.create_index( + _create_index_if_missing( "idx_chat_provider_call_logs_attempt_created", "chat_provider_call_logs", ["owner_type", "generation_attempt_no", "created_at"], unique=False, ) - op.create_index( + _create_index_if_missing( "idx_chat_provider_call_logs_record_created", "chat_provider_call_logs", ["owner_type", "generation_record_id", "created_at"], unique=False, ) - op.create_index( + _create_index_if_missing( "idx_chat_provider_call_logs_task_created", "chat_provider_call_logs", ["owner_type", "task_id", "created_at"], unique=False, ) - op.create_index( + _create_index_if_missing( op.f("ix_chat_provider_call_logs_generation_attempt_no"), "chat_provider_call_logs", ["generation_attempt_no"], unique=False, ) - op.create_index( + _create_index_if_missing( op.f("ix_chat_provider_call_logs_generation_record_id"), "chat_provider_call_logs", ["generation_record_id"], unique=False, ) - op.create_index( + _create_index_if_missing( op.f("ix_chat_provider_call_logs_owner_type"), "chat_provider_call_logs", ["owner_type"], unique=False, ) - op.create_foreign_key( + _create_foreign_key_if_missing( FK_CALL_LOG_GENERATION_RECORD, "chat_provider_call_logs", "generation_records", @@ -213,7 +456,7 @@ def _add_shared_log_owner_columns() -> None: ["id"], ondelete="CASCADE", ) - op.create_check_constraint( + _create_check_constraint_if_missing( 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) " @@ -222,7 +465,7 @@ def _add_shared_log_owner_columns() -> None: def _add_chat_generation_task_columns() -> None: - op.add_column( + _add_column_if_missing( "chat_generation_tasks", sa.Column( "generation_attempt_no", @@ -231,7 +474,7 @@ def _add_chat_generation_task_columns() -> None: nullable=False, ), ) - op.add_column( + _add_column_if_missing( "chat_generation_tasks", sa.Column( "resource_generation_started_at", @@ -239,7 +482,7 @@ def _add_chat_generation_task_columns() -> None: nullable=True, ), ) - op.add_column( + _add_column_if_missing( "chat_generation_tasks", sa.Column( "manual_retry_count", @@ -248,7 +491,7 @@ def _add_chat_generation_task_columns() -> None: nullable=False, ), ) - op.add_column( + _add_column_if_missing( "chat_generation_tasks", sa.Column( "poll_error_count", @@ -257,38 +500,38 @@ def _add_chat_generation_task_columns() -> None: nullable=False, ), ) - op.add_column( + _add_column_if_missing( "chat_generation_tasks", sa.Column("poll_claim_token", sa.String(length=64), nullable=True), ) - op.add_column( + _add_column_if_missing( "chat_generation_tasks", sa.Column("poll_lease_until", sa.DateTime(timezone=True), nullable=True), ) - op.add_column( + _add_column_if_missing( "chat_generation_tasks", sa.Column("download_claim_token", sa.String(length=64), nullable=True), ) - op.create_index( + _create_index_if_missing( op.f("ix_chat_generation_tasks_download_claim_token"), "chat_generation_tasks", ["download_claim_token"], unique=False, ) - op.create_index( + _create_index_if_missing( op.f("ix_chat_generation_tasks_poll_claim_token"), "chat_generation_tasks", ["poll_claim_token"], unique=False, ) - op.create_index( + _create_index_if_missing( op.f("ix_chat_generation_tasks_poll_lease_until"), "chat_generation_tasks", ["poll_lease_until"], unique=False, ) - op.create_index( + _create_index_if_missing( op.f("ix_chat_generation_tasks_resource_generation_started_at"), "chat_generation_tasks", ["resource_generation_started_at"], @@ -329,9 +572,9 @@ def _add_generation_record_columns() -> None: sa.Column("download_storage_date_dir", sa.String(length=16), nullable=True), ] for column in columns: - op.add_column("generation_records", column) + _add_column_if_missing("generation_records", column) - op.create_index( + _create_index_if_missing( "idx_genrec_next_poll_at", "generation_records", ["next_poll_at"], @@ -354,7 +597,7 @@ def _add_generation_record_columns() -> None: "provider_create_lease_until", "resource_generation_started_at", ): - op.create_index( + _create_index_if_missing( op.f(f"ix_generation_records_{column_name}"), "generation_records", [column_name], @@ -365,17 +608,18 @@ def _add_generation_record_columns() -> None: 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( + _execute_if_columns_exist( """ UPDATE chat_generation_tasks SET resource_generation_started_at = created_at WHERE resource_generation_started_at IS NULL - """ + """, + {"chat_generation_tasks": ("resource_generation_started_at", "created_at")}, ) # 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( + _execute_if_columns_exist( """ WITH credit_attempts AS ( SELECT @@ -403,18 +647,29 @@ def _backfill_generation_attempts() -> None: FROM credit_attempts AS attempts WHERE attempts.resolved_owner_id = task.id AND attempts.max_attempt IS NOT NULL - """ + """, + { + "credit_records": ( + "owner_id", "related_id", "biz_key", "attempt_no", "type", "owner_type", "charge_kind" + ), + "chat_generation_tasks": ("id", "generation_attempt_no"), + }, ) - op.execute( + _execute_if_columns_exist( """ 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 - """ + """, + { + "chat_generation_tasks": ( + "generation_attempt_no", "manual_retry_count", "retry_count", "poll_error_count" + ) + }, ) - op.execute( + _execute_if_columns_exist( """ WITH generation_charges AS ( SELECT @@ -470,13 +725,22 @@ def _backfill_generation_attempts() -> None: engine_id = COALESCE(record.engine_id, charge.engine_id) FROM generation_charges AS charge WHERE charge.resolved_owner_id = record.id - """ + """, + { + "credit_records": ( + "owner_id", "related_id", "biz_key", "attempt_no", "created_at", "engine_id", + "id", "type", "owner_type", "charge_kind" + ), + "generation_records": ( + "id", "generation_attempt_no", "resource_generation_started_at", "engine_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( + _execute_if_columns_exist( """ UPDATE generation_records SET resource_generation_started_at = COALESCE(updated_at, created_at) @@ -487,15 +751,26 @@ def _backfill_generation_attempts() -> None: OR image_url IS NOT NULL OR video_url IS NOT NULL ) - """ + """, + { + "generation_records": ( + "resource_generation_started_at", "updated_at", "created_at", "status", + "seedance_task_id", "image_url", "video_url" + ) + }, ) - op.execute( + _execute_if_columns_exist( """ 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 - """ + """, + { + "generation_records": ( + "generation_attempt_no", "manual_retry_count", "retry_count", "poll_error_count" + ) + }, ) @@ -503,7 +778,7 @@ 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( + _execute_if_columns_exist( """ UPDATE generation_records AS record SET engine_id = engine.id @@ -517,9 +792,13 @@ def _backfill_generation_record_engine() -> None: WHERE record.engine_id IS NULL AND record.status = 'generating' AND record.gen_type = 'image' - """ + """, + { + "generation_records": ("engine_id", "status", "gen_type"), + "image_engines": ("id", "is_active", "deleted_at", "priority"), + }, ) - op.execute( + _execute_if_columns_exist( """ UPDATE generation_records AS record SET engine_id = engine.id @@ -533,12 +812,16 @@ def _backfill_generation_record_engine() -> None: WHERE record.engine_id IS NULL AND record.status = 'generating' AND record.gen_type = 'video' - """ + """, + { + "generation_records": ("engine_id", "status", "gen_type"), + "video_engines": ("id", "is_active", "deleted_at", "priority"), + }, ) # 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( + _execute_if_columns_exist( """ UPDATE generation_records AS record SET engine_snapshot_json = jsonb_build_object( @@ -564,9 +847,21 @@ def _backfill_generation_record_engine() -> None: WHERE record.gen_type = 'image' AND record.engine_id = engine.id AND record.engine_snapshot_json IS NULL - """ + """, + { + "generation_records": ( + "gen_type", "engine_id", "engine_snapshot_json", "image_size", + "image_proportion", "image_px" + ), + "image_engines": ( + "id", "name", "provider", "api_base", "api_key", "model_name", + "generate_url", "default_size", "multi_generation_enabled", + "max_generation_count", "multi_image_max_images", + "max_reference_image_count", "output_format" + ), + }, ) - op.execute( + _execute_if_columns_exist( """ UPDATE generation_records AS record SET engine_snapshot_json = jsonb_build_object( @@ -591,13 +886,24 @@ def _backfill_generation_record_engine() -> None: WHERE record.gen_type = 'video' AND record.engine_id = engine.id AND record.engine_snapshot_json IS NULL - """ + """, + { + "generation_records": ( + "gen_type", "engine_id", "engine_snapshot_json", "aspect_ratio", + "resolution", "duration" + ), + "video_engines": ( + "id", "name", "provider", "api_base", "api_key", "model_name", + "generate_url", "query_url", "max_duration", "max_audio_count", + "multi_generation_enabled", "max_generation_count" + ), + }, ) def _backfill_generation_pipeline_state() -> None: # Preserve absolute historical provider URLs when they are available. - op.execute( + _execute_if_columns_exist( """ UPDATE generation_records SET remote_result_url = CASE @@ -610,10 +916,11 @@ def _backfill_generation_pipeline_state() -> None: (gen_type = 'image' AND image_url ~ '^https?://') OR (gen_type = 'video' AND video_url ~ '^https?://') ) - """ + """, + {"generation_records": ("remote_result_url", "gen_type", "image_url", "video_url")}, ) - op.execute( + _execute_if_columns_exist( """ UPDATE generation_records SET pipeline_stage = CASE @@ -623,58 +930,80 @@ def _backfill_generation_pipeline_state() -> None: END WHERE status = 'generating' AND pipeline_stage IS NULL - """ + """, + { + "generation_records": ( + "pipeline_stage", "remote_result_url", "seedance_task_id", "status" + ) + }, ) - op.execute( + _execute_if_columns_exist( """ UPDATE generation_records SET pipeline_stage = 'done' WHERE status = 'completed' AND pipeline_stage IS NULL - """ + """, + {"generation_records": ("pipeline_stage", "status")}, ) - op.execute( + _execute_if_columns_exist( """ UPDATE generation_records SET pipeline_stage = 'failed' WHERE status = 'failed' AND pipeline_stage IS NULL - """ + """, + {"generation_records": ("pipeline_stage", "status")}, ) # Image deadline is now 30 minutes; video remains 24 hours. Only active # tasks are rewritten so terminal historical audit values are preserved. - op.execute( + _execute_if_columns_exist( """ 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 - """ + """, + { + "chat_generation_tasks": ( + "deadline_at", "resource_generation_started_at", "status", "gen_type" + ) + }, ) - op.execute( + _execute_if_columns_exist( """ 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 - """ + """, + { + "chat_generation_tasks": ( + "deadline_at", "resource_generation_started_at", "status", "gen_type" + ) + }, ) - op.execute( + _execute_if_columns_exist( """ 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 - """ + """, + { + "generation_records": ( + "deadline_at", "resource_generation_started_at", "gen_type", "status" + ) + }, ) # Active provider tasks are checked immediately after deployment. The # recovery code still honours Redis execution locks and poll leases. - op.execute( + _execute_if_columns_exist( """ UPDATE generation_records SET poll_started_at = COALESCE(poll_started_at, resource_generation_started_at), @@ -684,11 +1013,25 @@ def _backfill_generation_pipeline_state() -> None: AND seedance_task_id IS NOT NULL AND remote_result_url IS NULL AND pipeline_stage IN ('waiting_remote', 'polling') - """ + """, + { + "generation_records": ( + "poll_started_at", "resource_generation_started_at", "next_poll_at", "status", + "gen_type", "seedance_task_id", "remote_result_url", "pipeline_stage" + ) + }, ) def _create_generated_resource_idempotency_index() -> None: + if _index_exists("generated_resources", "uq_generated_resources_active_source_type"): + return + if not _columns_exist( + "generated_resources", + ("source_model", "source_id", "resource_type", "deleted_at"), + ): + return + # 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. @@ -711,7 +1054,7 @@ def _create_generated_resource_idempotency_index() -> None: $$; """ ) - op.create_index( + _create_index_if_missing( "uq_generated_resources_active_source_type", "generated_resources", ["source_model", "source_id", "resource_type"], @@ -733,7 +1076,7 @@ def upgrade() -> None: def downgrade() -> None: - op.drop_index( + _drop_index_if_exists( "uq_generated_resources_active_source_type", table_name="generated_resources", postgresql_where=sa.text("deleted_at IS NULL"), @@ -752,11 +1095,11 @@ def downgrade() -> None: "download_celery_task_id", "deadline_at", ): - op.drop_index( + _drop_index_if_exists( op.f(f"ix_generation_records_{column_name}"), table_name="generation_records", ) - op.drop_index("idx_genrec_next_poll_at", table_name="generation_records") + _drop_index_if_exists("idx_genrec_next_poll_at", table_name="generation_records") for column_name in ( "download_storage_date_dir", @@ -789,120 +1132,131 @@ def downgrade() -> None: "resource_generation_started_at", "generation_attempt_no", ): - op.drop_column("generation_records", column_name) + _drop_column_if_exists("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( + _execute_if_columns_exist( "DELETE FROM chat_provider_call_logs " - "WHERE owner_type = 'generation_record' OR generation_record_id IS NOT NULL" + "WHERE owner_type = 'generation_record' OR generation_record_id IS NOT NULL", + {"chat_provider_call_logs": ("owner_type", "generation_record_id")}, ) - op.drop_constraint(CK_CALL_LOG_OWNER, "chat_provider_call_logs", type_="check") - op.drop_constraint( + _execute_if_columns_exist( + "DELETE FROM chat_provider_call_logs WHERE task_id IS NULL", + {"chat_provider_call_logs": ("task_id",)}, + ) + _drop_constraint_if_exists(CK_CALL_LOG_OWNER, "chat_provider_call_logs", type_="check") + _drop_constraint_if_exists( FK_CALL_LOG_GENERATION_RECORD, "chat_provider_call_logs", type_="foreignkey", ) - op.drop_index( + _drop_index_if_exists( op.f("ix_chat_provider_call_logs_owner_type"), table_name="chat_provider_call_logs", ) - op.drop_index( + _drop_index_if_exists( op.f("ix_chat_provider_call_logs_generation_record_id"), table_name="chat_provider_call_logs", ) - op.drop_index( + _drop_index_if_exists( 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( + _drop_index_if_exists("idx_chat_provider_call_logs_task_created", table_name="chat_provider_call_logs") + _drop_index_if_exists("idx_chat_provider_call_logs_record_created", table_name="chat_provider_call_logs") + _drop_index_if_exists("idx_chat_provider_call_logs_attempt_created", table_name="chat_provider_call_logs") + _alter_column_if_exists( "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") + _drop_column_if_exists("chat_provider_call_logs", "generation_attempt_no") + _drop_column_if_exists("chat_provider_call_logs", "generation_record_id") + _drop_column_if_exists("chat_provider_call_logs", "owner_type") - op.execute( + _execute_if_columns_exist( "DELETE FROM chat_generation_task_events " - "WHERE owner_type = 'generation_record' OR generation_record_id IS NOT NULL" + "WHERE owner_type = 'generation_record' OR generation_record_id IS NOT NULL", + {"chat_generation_task_events": ("owner_type", "generation_record_id")}, ) - op.drop_constraint(CK_EVENT_OWNER, "chat_generation_task_events", type_="check") - op.drop_constraint( + _execute_if_columns_exist( + "DELETE FROM chat_generation_task_events WHERE task_id IS NULL", + {"chat_generation_task_events": ("task_id",)}, + ) + _drop_constraint_if_exists(CK_EVENT_OWNER, "chat_generation_task_events", type_="check") + _drop_constraint_if_exists( FK_EVENT_GENERATION_RECORD, "chat_generation_task_events", type_="foreignkey", ) - op.drop_index( + _drop_index_if_exists( op.f("ix_chat_generation_task_events_owner_type"), table_name="chat_generation_task_events", ) - op.drop_index( + _drop_index_if_exists( op.f("ix_chat_generation_task_events_generation_record_id"), table_name="chat_generation_task_events", ) - op.drop_index( + _drop_index_if_exists( 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( + _drop_index_if_exists("idx_chat_generation_task_events_task_created", table_name="chat_generation_task_events") + _drop_index_if_exists("idx_chat_generation_task_events_record_created", table_name="chat_generation_task_events") + _drop_index_if_exists("idx_chat_generation_task_events_attempt_created", table_name="chat_generation_task_events") + _execute_if_columns_exist( "UPDATE chat_generation_task_events " "SET message = LEFT(message, 512), " "from_stage = LEFT(from_stage, 32), " - "to_stage = LEFT(to_stage, 32)" + "to_stage = LEFT(to_stage, 32)", + {"chat_generation_task_events": ("message", "from_stage", "to_stage")}, ) - op.alter_column( + _alter_column_if_exists( "chat_generation_task_events", "message", existing_type=sa.Text(), type_=sa.VARCHAR(length=512), existing_nullable=True, ) - op.alter_column( + _alter_column_if_exists( "chat_generation_task_events", "to_stage", existing_type=sa.String(length=48), type_=sa.VARCHAR(length=32), existing_nullable=True, ) - op.alter_column( + _alter_column_if_exists( "chat_generation_task_events", "from_stage", existing_type=sa.String(length=48), type_=sa.VARCHAR(length=32), existing_nullable=True, ) - op.alter_column( + _alter_column_if_exists( "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") + _drop_column_if_exists("chat_generation_task_events", "generation_attempt_no") + _drop_column_if_exists("chat_generation_task_events", "generation_record_id") + _drop_column_if_exists("chat_generation_task_events", "owner_type") - op.drop_index( + _drop_index_if_exists( op.f("ix_chat_generation_tasks_resource_generation_started_at"), table_name="chat_generation_tasks", ) - op.drop_index( + _drop_index_if_exists( op.f("ix_chat_generation_tasks_poll_lease_until"), table_name="chat_generation_tasks", ) - op.drop_index( + _drop_index_if_exists( op.f("ix_chat_generation_tasks_poll_claim_token"), table_name="chat_generation_tasks", ) - op.drop_index( + _drop_index_if_exists( op.f("ix_chat_generation_tasks_download_claim_token"), table_name="chat_generation_tasks", ) @@ -915,8 +1269,8 @@ def downgrade() -> None: "resource_generation_started_at", "generation_attempt_no", ): - op.drop_column("chat_generation_tasks", column_name) + _drop_column_if_exists("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") + _drop_index_if_exists(op.f(f"ix_{table_name}_deleted_at"), table_name=table_name) + _drop_column_if_exists(table_name, "deleted_at")