migrations/versions/initial_migration.py

"""Initial migration

Revision ID: 001
Revises: 
Create Date: 2024-03-19 10:00:00.000000

"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
import enum

# revision identifiers, used by Alembic.
revision = '001'
down_revision = None
branch_labels = None
depends_on = None

class AudiobookStatus(enum.Enum):
    """Status of an audiobook."""
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"

def upgrade() -> None:
    # Create enum type for PostgreSQL
    if op.get_bind().dialect.name == 'postgresql':
        op.execute('CREATE TYPE audiobookstatus AS ENUM (\'pending\', \'processing\', \'completed\', \'failed\')')
        status_type = postgresql.ENUM('pending', 'processing', 'completed', 'failed', name='audiobookstatus')
    else:
        status_type = sa.String(20)

    # Create audiobooks table
    op.create_table(
        'audiobooks',
        sa.Column('id', sa.String(36), primary_key=True),
        sa.Column('title', sa.String(255), nullable=False),
        sa.Column('author', sa.String(255), nullable=False),
        sa.Column('voice_id', sa.Integer, nullable=False),
        sa.Column('status', status_type, nullable=False),
        sa.Column('created_at', sa.DateTime, nullable=False),
        sa.Column('updated_at', sa.DateTime, nullable=False),
        sa.Column('text_content', sa.Text, nullable=True),
        sa.Column('text_file_path', sa.String(255), nullable=True),
        sa.Column('audio_file_path', sa.String(255), nullable=True),
        sa.Column('error_message', sa.Text, nullable=True),
    )

    # Create audiobook_chunks table
    op.create_table(
        'audiobook_chunks',
        sa.Column('id', sa.Integer, primary_key=True),
        sa.Column('audiobook_id', sa.String(36), sa.ForeignKey('audiobooks.id'), nullable=False),
        sa.Column('chunk_number', sa.Integer, nullable=False),
        sa.Column('text_content', sa.Text, nullable=False),
        sa.Column('audio_file_path', sa.String(255), nullable=True),
        sa.Column('status', status_type, nullable=False),
        sa.Column('created_at', sa.DateTime, nullable=False),
        sa.Column('updated_at', sa.DateTime, nullable=False),
    )

    # Create indexes
    op.create_index('ix_audiobooks_created_at', 'audiobooks', ['created_at'])
    op.create_index('ix_audiobooks_status', 'audiobooks', ['status'])
    op.create_index('ix_audiobook_chunks_audiobook_id', 'audiobook_chunks', ['audiobook_id'])
    op.create_index('ix_audiobook_chunks_chunk_number', 'audiobook_chunks', ['chunk_number'])

def downgrade() -> None:
    # Drop tables
    op.drop_table('audiobook_chunks')
    op.drop_table('audiobooks')

    # Drop enum type if using PostgreSQL
    if op.get_bind().dialect.name == 'postgresql':
        op.execute('DROP TYPE audiobookstatus')