195 lines
12 KiB
Python
195 lines
12 KiB
Python
"""add api v3 tables (api_keys, api_generation_tasks, api_usage_logs, api_key_upscale_configs, api_upscale_links)
|
|
|
|
Revision ID: a1b2c3d4e5f6
|
|
Revises: 20da1d353914
|
|
Create Date: 2026-07-28 12:00:00.000000
|
|
"""
|
|
from typing import Sequence, Union
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = 'a1b2c3d4e5f6g'
|
|
down_revision: Union[str, None] = '20da1d353914'
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
# === 1. api_keys ===
|
|
op.create_table(
|
|
'api_keys',
|
|
sa.Column('id', sa.String(32), primary_key=True),
|
|
sa.Column('company_name', sa.String(128), nullable=False),
|
|
sa.Column('api_key_hash', sa.String(64), nullable=False, unique=True),
|
|
sa.Column('api_key_prefix', sa.String(16), nullable=False),
|
|
sa.Column('description', sa.Text, nullable=True),
|
|
sa.Column('callable_models', sa.Text, nullable=False, server_default='[]'),
|
|
sa.Column('quota_limit', sa.Float, nullable=True),
|
|
sa.Column('quota_cycle', sa.String(16), nullable=True),
|
|
sa.Column('quota_used', sa.Float, nullable=False, server_default='0.0'),
|
|
sa.Column('valid_from', sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column('valid_until', sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column('max_concurrent_video_tasks', sa.Integer, nullable=True),
|
|
sa.Column('is_active', sa.Boolean, nullable=False, server_default='true'),
|
|
sa.Column('last_used_at', sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
|
|
sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
|
|
sa.Column('deleted_at', sa.DateTime(timezone=True), nullable=True),
|
|
)
|
|
op.create_index('idx_api_key_hash', 'api_keys', ['api_key_hash'], unique=True)
|
|
op.create_index('idx_api_keys_active', 'api_keys', ['is_active'])
|
|
op.create_index('idx_api_keys_company', 'api_keys', ['company_name'])
|
|
|
|
# === 2. api_generation_tasks ===
|
|
op.create_table(
|
|
'api_generation_tasks',
|
|
sa.Column('id', sa.String(32), primary_key=True),
|
|
sa.Column('api_key_id', sa.String(32), sa.ForeignKey('api_keys.id', ondelete='CASCADE'), nullable=False),
|
|
sa.Column('external_idempotency_key', sa.String(64), nullable=True),
|
|
sa.Column('original_prompt', sa.Text, nullable=False),
|
|
sa.Column('optimized_prompt', sa.Text, nullable=True),
|
|
sa.Column('gen_type', sa.String(16), nullable=False, default='video'),
|
|
sa.Column('duration', sa.Integer, nullable=True),
|
|
sa.Column('aspect_ratio', sa.String(8), nullable=True),
|
|
sa.Column('resolution', sa.String(8), nullable=True),
|
|
sa.Column('provider_generation_resolution', sa.String(16), nullable=True),
|
|
sa.Column('image_size', sa.String(16), nullable=True),
|
|
sa.Column('image_proportion', sa.String(8), nullable=True),
|
|
sa.Column('image_px', sa.String(16), nullable=True),
|
|
sa.Column('generation_count', sa.Integer, nullable=False, default=1, server_default='1'),
|
|
sa.Column('engine_id', sa.String(32), nullable=True),
|
|
sa.Column('media_references', sa.Text, nullable=True),
|
|
sa.Column('engine_snapshot_json', sa.Text, nullable=True),
|
|
sa.Column('request_params_json', sa.Text, nullable=True),
|
|
sa.Column('status', sa.String(32), nullable=False, default='pending'),
|
|
sa.Column('pipeline_stage', sa.String(32), nullable=True),
|
|
sa.Column('generation_attempt_no', sa.Integer, nullable=False, default=1, server_default='1'),
|
|
sa.Column('resource_generation_started_at', sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column('deadline_at', sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column('provider_task_id', sa.String(128), nullable=True),
|
|
sa.Column('remote_result_url', sa.Text, nullable=True),
|
|
sa.Column('provider_response_json', sa.Text, nullable=True),
|
|
sa.Column('image_url', sa.String(512), nullable=True),
|
|
sa.Column('video_url', sa.String(512), nullable=True),
|
|
sa.Column('video_cover_url', sa.String(512), nullable=True),
|
|
sa.Column('video_upscale_enabled_snapshot', sa.Boolean, nullable=False, default=False, server_default='false'),
|
|
sa.Column('video_upscale_snapshot_json', sa.Text, nullable=True),
|
|
sa.Column('credits_cost', sa.Float, nullable=False, default=0.0, server_default='0.0'),
|
|
sa.Column('video_tokens_used', sa.Integer, nullable=False, default=0, server_default='0'),
|
|
sa.Column('image_tokens_used', sa.Integer, nullable=False, default=0, server_default='0'),
|
|
sa.Column('error_message', sa.Text, nullable=True),
|
|
sa.Column('generated_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, nullable=False, default=30, server_default='30'),
|
|
sa.Column('poll_count', sa.Integer, nullable=False, default=0, server_default='0'),
|
|
sa.Column('last_poll_at', sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column('provider_create_claim_token', sa.String(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('poll_started_at', sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column('poll_claim_token', sa.String(64), nullable=True),
|
|
sa.Column('poll_lease_until', sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column('poll_error_count', sa.Integer, nullable=False, default=0, server_default='0'),
|
|
sa.Column('download_celery_task_id', sa.String(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(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, nullable=False, default=0, server_default='0'),
|
|
sa.Column('download_last_error', sa.Text, nullable=True),
|
|
sa.Column('download_storage_date_dir', sa.String(16), nullable=True),
|
|
sa.Column('local_path', sa.Text, nullable=True),
|
|
sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
|
|
sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
|
|
sa.Column('deleted_at', sa.DateTime(timezone=True), nullable=True),
|
|
)
|
|
op.create_index('idx_api_generation_tasks_api_key', 'api_generation_tasks', ['api_key_id'])
|
|
op.create_index('idx_api_generation_tasks_status', 'api_generation_tasks', ['status'])
|
|
op.create_index('idx_api_generation_tasks_provider_task_id', 'api_generation_tasks', ['provider_task_id'])
|
|
op.create_index('idx_api_generation_tasks_next_poll_at', 'api_generation_tasks', ['next_poll_at'])
|
|
op.create_index('idx_api_generation_tasks_api_key_created', 'api_generation_tasks', ['api_key_id', 'created_at'])
|
|
op.create_index(
|
|
'uq_api_generation_tasks_key_idempotency',
|
|
'api_generation_tasks',
|
|
['api_key_id', 'external_idempotency_key'],
|
|
unique=True,
|
|
postgresql_where=sa.text("deleted_at IS NULL AND external_idempotency_key IS NOT NULL"),
|
|
)
|
|
|
|
# === 3. api_usage_logs ===
|
|
op.create_table(
|
|
'api_usage_logs',
|
|
sa.Column('id', sa.String(32), primary_key=True),
|
|
sa.Column('api_key_id', sa.String(32), sa.ForeignKey('api_keys.id', ondelete='CASCADE'), nullable=False),
|
|
sa.Column('api_generation_task_id', sa.String(32), sa.ForeignKey('api_generation_tasks.id', ondelete='SET NULL'), nullable=True),
|
|
sa.Column('request_type', sa.String(32), nullable=False),
|
|
sa.Column('model_name', sa.String(128), nullable=False),
|
|
sa.Column('gen_type', sa.String(16), nullable=False),
|
|
sa.Column('credits_cost', sa.Float, nullable=False, default=0.0, server_default='0.0'),
|
|
sa.Column('tokens_used', sa.Integer, nullable=False, default=0, server_default='0'),
|
|
sa.Column('request_duration_ms', sa.Integer, nullable=False, default=0, server_default='0'),
|
|
sa.Column('status', sa.String(32), nullable=False),
|
|
sa.Column('error_message', sa.Text, nullable=True),
|
|
sa.Column('error_code', sa.String(64), nullable=True),
|
|
sa.Column('request_payload_json', sa.Text, nullable=True),
|
|
sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
|
|
sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
|
|
)
|
|
op.create_index('idx_api_usage_logs_api_key', 'api_usage_logs', ['api_key_id'])
|
|
op.create_index('idx_api_usage_logs_task_id', 'api_usage_logs', ['api_generation_task_id'])
|
|
op.create_index('idx_api_usage_logs_api_key_created', 'api_usage_logs', ['api_key_id', 'created_at'])
|
|
|
|
# === 4. api_key_upscale_configs ===
|
|
op.create_table(
|
|
'api_key_upscale_configs',
|
|
sa.Column('id', sa.String(32), primary_key=True),
|
|
sa.Column('api_key_id', sa.String(32), sa.ForeignKey('api_keys.id', ondelete='CASCADE'), nullable=False, unique=True),
|
|
sa.Column('enabled', sa.Boolean, nullable=False, default=False, server_default='false'),
|
|
sa.Column('delete_source_after_success', sa.Boolean, nullable=False, default=True, server_default='true'),
|
|
sa.Column('rules_json', sa.Text, nullable=False, server_default='[]'),
|
|
sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
|
|
sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
|
|
)
|
|
|
|
# === 5. api_upscale_links ===
|
|
op.create_table(
|
|
'api_upscale_links',
|
|
sa.Column('id', sa.String(32), primary_key=True),
|
|
sa.Column('api_generation_task_id', sa.String(32), sa.ForeignKey('api_generation_tasks.id', ondelete='CASCADE'), nullable=False),
|
|
sa.Column('video_upscale_task_id', sa.String(32), sa.ForeignKey('video_upscale_tasks.id', ondelete='CASCADE'), nullable=False),
|
|
sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
|
|
sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
|
|
)
|
|
op.create_index('idx_api_upscale_links_api_task', 'api_upscale_links', ['api_generation_task_id'])
|
|
op.create_index('idx_api_upscale_links_video_task', 'api_upscale_links', ['video_upscale_task_id'])
|
|
|
|
|
|
def downgrade() -> None:
|
|
op.drop_index('idx_api_upscale_links_video_task', table_name='api_upscale_links')
|
|
op.drop_index('idx_api_upscale_links_api_task', table_name='api_upscale_links')
|
|
op.drop_table('api_upscale_links')
|
|
|
|
op.drop_table('api_key_upscale_configs')
|
|
|
|
op.drop_index('idx_api_usage_logs_api_key_created', table_name='api_usage_logs')
|
|
op.drop_index('idx_api_usage_logs_task_id', table_name='api_usage_logs')
|
|
op.drop_index('idx_api_usage_logs_api_key', table_name='api_usage_logs')
|
|
op.drop_table('api_usage_logs')
|
|
|
|
op.drop_index('uq_api_generation_tasks_key_idempotency', table_name='api_generation_tasks')
|
|
op.drop_index('idx_api_generation_tasks_api_key_created', table_name='api_generation_tasks')
|
|
op.drop_index('idx_api_generation_tasks_next_poll_at', table_name='api_generation_tasks')
|
|
op.drop_index('idx_api_generation_tasks_provider_task_id', table_name='api_generation_tasks')
|
|
op.drop_index('idx_api_generation_tasks_status', table_name='api_generation_tasks')
|
|
op.drop_index('idx_api_generation_tasks_api_key', table_name='api_generation_tasks')
|
|
op.drop_table('api_generation_tasks')
|
|
|
|
op.drop_index('idx_api_keys_company', table_name='api_keys')
|
|
op.drop_index('idx_api_keys_active', table_name='api_keys')
|
|
op.drop_index('idx_api_key_hash', table_name='api_keys')
|
|
op.drop_table('api_keys')
|