Skip to content

Commit

Permalink
feat: database migrations, so long db resets (#858)
Browse files Browse the repository at this point in the history
* feat: database migrations, so long db resets

* chore: remove README, styling, reset_pre_migration cli and env option

---------

Co-authored-by: Gaisberg <None>
  • Loading branch information
Gaisberg authored Nov 5, 2024
1 parent 27c8534 commit 14e818f
Show file tree
Hide file tree
Showing 8 changed files with 489 additions and 86 deletions.
12 changes: 12 additions & 0 deletions src/alembic.ini
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
# A generic, single database configuration.

[alembic]

script_location = %(here)s/alembic
file_template = %%(year)d%%(month).2d%%(day).2d_%%(hour).2d%%(minute).2d_%%(rev)s_%%(slug)s
prepend_sys_path = .
truncate_slug_length = 40
version_locations = %(here)s/alembic/versions
version_path_separator = os
output_encoding = utf-8
sqlalchemy.url = driver://user:pass@localhost/dbname
124 changes: 124 additions & 0 deletions src/alembic/env.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,124 @@
import logging

from loguru import logger
from sqlalchemy import engine_from_config, pool, text
from sqlalchemy.exc import OperationalError, ProgrammingError

from alembic import context
from program.db.db import db
from program.settings.manager import settings_manager


# Loguru handler for alembic logs
class LoguruHandler(logging.Handler):
def emit(self, record):
logger.opt(depth=1, exception=record.exc_info).log("DATABASE", record.getMessage())

if settings_manager.settings.debug_database:
# Configure only alembic and SQLAlchemy loggers
logging.getLogger("alembic").handlers = [LoguruHandler()]
logging.getLogger("alembic").propagate = False
logging.getLogger("sqlalchemy").handlers = [LoguruHandler()]
logging.getLogger("sqlalchemy").propagate = False

# Set log levels
logging.getLogger("alembic").setLevel(logging.DEBUG if settings_manager.settings.debug else logging.FATAL)
logging.getLogger("sqlalchemy").setLevel(logging.DEBUG if settings_manager.settings.debug else logging.FATAL)

# Alembic configuration
config = context.config
config.set_main_option("sqlalchemy.url", settings_manager.settings.database.host)

# Set MetaData object for autogenerate support
target_metadata = db.Model.metadata

def reset_database(connection) -> bool:
"""Reset database if needed"""
try:
# Drop and recreate schema
if db.engine.name == "postgresql":
connection.execute(text("""
SELECT pg_terminate_backend(pid)
FROM pg_stat_activity
WHERE datname = current_database()
AND pid <> pg_backend_pid()
"""))
connection.execute(text("DROP SCHEMA public CASCADE"))
connection.execute(text("CREATE SCHEMA public"))
connection.execute(text("GRANT ALL ON SCHEMA public TO public"))

logger.debug("DATABASE", "Database reset complete")
return True
except Exception as e:
logger.error(f"Database reset failed: {e}")
return False

def run_migrations_offline() -> None:
"""Run migrations in 'offline' mode."""
url = config.get_main_option("sqlalchemy.url")
context.configure(
url=url,
target_metadata=target_metadata,
literal_binds=True,
dialect_opts={"paramstyle": "named"},
)

with context.begin_transaction():
context.run_migrations()

def run_migrations_online() -> None:
"""Run migrations in 'online' mode."""
connectable = engine_from_config(
config.get_section(config.config_ini_section),
prefix="sqlalchemy.",
poolclass=pool.NullPool,
)

with connectable.connect() as connection:
connection = connection.execution_options(isolation_level="AUTOCOMMIT")
try:
context.configure(
connection=connection,
target_metadata=target_metadata,
compare_type=True, # Compare column types
compare_server_default=True, # Compare default values
include_schemas=True, # Include schema in migrations
render_as_batch=True, # Enable batch operations
)

with context.begin_transaction():
logger.debug("Starting migrations...")
context.run_migrations()
logger.debug("Migrations completed successfully")

except (OperationalError, ProgrammingError) as e:
logger.error(f"Database error during migration: {e}")
logger.warning("Attempting database reset...")

if reset_database(connection):
# Configure alembic again after reset
context.configure(
connection=connection,
target_metadata=target_metadata,
compare_type=True,
compare_server_default=True,
include_schemas=True,
render_as_batch=True,
)

# Try migrations again
with context.begin_transaction():
logger.debug("Rerunning migrations after reset...")
context.run_migrations()
logger.debug("Migrations completed successfully")
else:
raise Exception("Migration recovery failed")

except Exception as e:
logger.error(f"Unexpected error during migration: {e}")
raise

if context.is_offline_mode():
run_migrations_offline()
else:
run_migrations_online()
26 changes: 26 additions & 0 deletions src/alembic/script.py.mako
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
"""${message}

Revision ID: ${up_revision}
Revises: ${down_revision | comma,n}
Create Date: ${create_date}

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
${imports if imports else ""}

# revision identifiers, used by Alembic.
revision: str = ${repr(up_revision)}
down_revision: Union[str, None] = ${repr(down_revision)}
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}


def upgrade() -> None:
${upgrades if upgrades else "pass"}


def downgrade() -> None:
${downgrades if downgrades else "pass"}
179 changes: 179 additions & 0 deletions src/alembic/versions/20241105_1300_c99709e3648f_baseline_schema.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,179 @@
"""baseline_schema
Revision ID: c99709e3648f
Revises:
Create Date: 2024-11-05 13:00:06.356164
"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'c99709e3648f'
down_revision: Union[str, None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.create_table('MediaItem',
sa.Column('id', sa.String(), nullable=False),
sa.Column('imdb_id', sa.String(), nullable=True),
sa.Column('tvdb_id', sa.String(), nullable=True),
sa.Column('tmdb_id', sa.String(), nullable=True),
sa.Column('number', sa.Integer(), nullable=True),
sa.Column('type', sa.String(), nullable=False),
sa.Column('requested_at', sa.DateTime(), nullable=True),
sa.Column('requested_by', sa.String(), nullable=True),
sa.Column('requested_id', sa.Integer(), nullable=True),
sa.Column('indexed_at', sa.DateTime(), nullable=True),
sa.Column('scraped_at', sa.DateTime(), nullable=True),
sa.Column('scraped_times', sa.Integer(), nullable=True),
sa.Column('active_stream', sa.JSON(), nullable=True),
sa.Column('symlinked', sa.Boolean(), nullable=True),
sa.Column('symlinked_at', sa.DateTime(), nullable=True),
sa.Column('symlinked_times', sa.Integer(), nullable=True),
sa.Column('symlink_path', sa.String(), nullable=True),
sa.Column('file', sa.String(), nullable=True),
sa.Column('folder', sa.String(), nullable=True),
sa.Column('alternative_folder', sa.String(), nullable=True),
sa.Column('aliases', sa.JSON(), nullable=True),
sa.Column('is_anime', sa.Boolean(), nullable=True),
sa.Column('title', sa.String(), nullable=True),
sa.Column('network', sa.String(), nullable=True),
sa.Column('country', sa.String(), nullable=True),
sa.Column('language', sa.String(), nullable=True),
sa.Column('aired_at', sa.DateTime(), nullable=True),
sa.Column('year', sa.Integer(), nullable=True),
sa.Column('genres', sa.JSON(), nullable=True),
sa.Column('key', sa.String(), nullable=True),
sa.Column('guid', sa.String(), nullable=True),
sa.Column('update_folder', sa.String(), nullable=True),
sa.Column('overseerr_id', sa.Integer(), nullable=True),
sa.Column('last_state', sa.Enum('Unknown', 'Unreleased', 'Ongoing', 'Requested', 'Indexed', 'Scraped', 'Downloaded', 'Symlinked', 'Completed', 'PartiallyCompleted', 'Failed', name='states'), nullable=True),
sa.PrimaryKeyConstraint('id')
)
op.create_index('ix_mediaitem_aired_at', 'MediaItem', ['aired_at'], unique=False)
op.create_index('ix_mediaitem_country', 'MediaItem', ['country'], unique=False)
op.create_index('ix_mediaitem_imdb_id', 'MediaItem', ['imdb_id'], unique=False)
op.create_index('ix_mediaitem_language', 'MediaItem', ['language'], unique=False)
op.create_index('ix_mediaitem_network', 'MediaItem', ['network'], unique=False)
op.create_index('ix_mediaitem_overseerr_id', 'MediaItem', ['overseerr_id'], unique=False)
op.create_index('ix_mediaitem_requested_by', 'MediaItem', ['requested_by'], unique=False)
op.create_index('ix_mediaitem_title', 'MediaItem', ['title'], unique=False)
op.create_index('ix_mediaitem_tmdb_id', 'MediaItem', ['tmdb_id'], unique=False)
op.create_index('ix_mediaitem_tvdb_id', 'MediaItem', ['tvdb_id'], unique=False)
op.create_index('ix_mediaitem_type', 'MediaItem', ['type'], unique=False)
op.create_index('ix_mediaitem_type_aired_at', 'MediaItem', ['type', 'aired_at'], unique=False)
op.create_index('ix_mediaitem_year', 'MediaItem', ['year'], unique=False)
op.create_table('Stream',
sa.Column('id', sa.Integer(), nullable=False),
sa.Column('infohash', sa.String(), nullable=False),
sa.Column('raw_title', sa.String(), nullable=False),
sa.Column('parsed_title', sa.String(), nullable=False),
sa.Column('rank', sa.Integer(), nullable=False),
sa.Column('lev_ratio', sa.Float(), nullable=False),
sa.PrimaryKeyConstraint('id')
)
op.create_index('ix_stream_infohash', 'Stream', ['infohash'], unique=False)
op.create_index('ix_stream_parsed_title', 'Stream', ['parsed_title'], unique=False)
op.create_index('ix_stream_rank', 'Stream', ['rank'], unique=False)
op.create_index('ix_stream_raw_title', 'Stream', ['raw_title'], unique=False)
op.create_table('Movie',
sa.Column('id', sa.String(), nullable=False),
sa.ForeignKeyConstraint(['id'], ['MediaItem.id'], ),
sa.PrimaryKeyConstraint('id')
)
op.create_table('Show',
sa.Column('id', sa.String(), nullable=False),
sa.ForeignKeyConstraint(['id'], ['MediaItem.id'], ),
sa.PrimaryKeyConstraint('id')
)
op.create_table('StreamBlacklistRelation',
sa.Column('id', sa.Integer(), nullable=False),
sa.Column('media_item_id', sa.String(), nullable=False),
sa.Column('stream_id', sa.Integer(), nullable=False),
sa.ForeignKeyConstraint(['media_item_id'], ['MediaItem.id'], ondelete='CASCADE'),
sa.ForeignKeyConstraint(['stream_id'], ['Stream.id'], ondelete='CASCADE'),
sa.PrimaryKeyConstraint('id')
)
op.create_index('ix_streamblacklistrelation_media_item_id', 'StreamBlacklistRelation', ['media_item_id'], unique=False)
op.create_index('ix_streamblacklistrelation_stream_id', 'StreamBlacklistRelation', ['stream_id'], unique=False)
op.create_table('StreamRelation',
sa.Column('id', sa.Integer(), nullable=False),
sa.Column('parent_id', sa.String(), nullable=False),
sa.Column('child_id', sa.Integer(), nullable=False),
sa.ForeignKeyConstraint(['child_id'], ['Stream.id'], ondelete='CASCADE'),
sa.ForeignKeyConstraint(['parent_id'], ['MediaItem.id'], ondelete='CASCADE'),
sa.PrimaryKeyConstraint('id')
)
op.create_index('ix_streamrelation_child_id', 'StreamRelation', ['child_id'], unique=False)
op.create_index('ix_streamrelation_parent_id', 'StreamRelation', ['parent_id'], unique=False)
op.create_table('Subtitle',
sa.Column('id', sa.Integer(), nullable=False),
sa.Column('language', sa.String(), nullable=False),
sa.Column('file', sa.String(), nullable=True),
sa.Column('parent_id', sa.String(), nullable=False),
sa.ForeignKeyConstraint(['parent_id'], ['MediaItem.id'], ondelete='CASCADE'),
sa.PrimaryKeyConstraint('id')
)
op.create_index('ix_subtitle_file', 'Subtitle', ['file'], unique=False)
op.create_index('ix_subtitle_language', 'Subtitle', ['language'], unique=False)
op.create_index('ix_subtitle_parent_id', 'Subtitle', ['parent_id'], unique=False)
op.create_table('Season',
sa.Column('id', sa.String(), nullable=False),
sa.Column('parent_id', sa.String(), nullable=False),
sa.ForeignKeyConstraint(['id'], ['MediaItem.id'], ),
sa.ForeignKeyConstraint(['parent_id'], ['Show.id'], ),
sa.PrimaryKeyConstraint('id')
)
op.create_table('Episode',
sa.Column('id', sa.String(), nullable=False),
sa.Column('parent_id', sa.String(), nullable=False),
sa.ForeignKeyConstraint(['id'], ['MediaItem.id'], ),
sa.ForeignKeyConstraint(['parent_id'], ['Season.id'], ),
sa.PrimaryKeyConstraint('id')
)
# ### end Alembic commands ###


def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.drop_table('Episode')
op.drop_table('Season')
op.drop_index('ix_subtitle_parent_id', table_name='Subtitle')
op.drop_index('ix_subtitle_language', table_name='Subtitle')
op.drop_index('ix_subtitle_file', table_name='Subtitle')
op.drop_table('Subtitle')
op.drop_index('ix_streamrelation_parent_id', table_name='StreamRelation')
op.drop_index('ix_streamrelation_child_id', table_name='StreamRelation')
op.drop_table('StreamRelation')
op.drop_index('ix_streamblacklistrelation_stream_id', table_name='StreamBlacklistRelation')
op.drop_index('ix_streamblacklistrelation_media_item_id', table_name='StreamBlacklistRelation')
op.drop_table('StreamBlacklistRelation')
op.drop_table('Show')
op.drop_table('Movie')
op.drop_index('ix_stream_raw_title', table_name='Stream')
op.drop_index('ix_stream_rank', table_name='Stream')
op.drop_index('ix_stream_parsed_title', table_name='Stream')
op.drop_index('ix_stream_infohash', table_name='Stream')
op.drop_table('Stream')
op.drop_index('ix_mediaitem_year', table_name='MediaItem')
op.drop_index('ix_mediaitem_type_aired_at', table_name='MediaItem')
op.drop_index('ix_mediaitem_type', table_name='MediaItem')
op.drop_index('ix_mediaitem_tvdb_id', table_name='MediaItem')
op.drop_index('ix_mediaitem_tmdb_id', table_name='MediaItem')
op.drop_index('ix_mediaitem_title', table_name='MediaItem')
op.drop_index('ix_mediaitem_requested_by', table_name='MediaItem')
op.drop_index('ix_mediaitem_overseerr_id', table_name='MediaItem')
op.drop_index('ix_mediaitem_network', table_name='MediaItem')
op.drop_index('ix_mediaitem_language', table_name='MediaItem')
op.drop_index('ix_mediaitem_imdb_id', table_name='MediaItem')
op.drop_index('ix_mediaitem_country', table_name='MediaItem')
op.drop_index('ix_mediaitem_aired_at', table_name='MediaItem')
op.drop_table('MediaItem')
# ### end Alembic commands ###
Loading

0 comments on commit 14e818f

Please sign in to comment.