tessl/pypi-alembic

A database migration tool for SQLAlchemy.

—

Pending

Quality

Pending

Does it follow best practices?

Impact

Pending

No eval scenarios have been run

Overview

Eval results

Files

Migration Context

Name: tessl/pypi-alembic
Author: tessl

Environment context functions for accessing database connections, configuration, and migration state within migration scripts through the alembic.context module. These functions are primarily used in the env.py file and migration scripts.

Core Imports

from alembic import context

Capabilities

Context Configuration

Configure the migration environment with database connections and settings.

def configure(connection=None, url=None, dialect_name=None, dialect_opts=None, starting_rev=None, tag=None, template_args=None, render_as_batch=None, target_metadata=None, include_name=None, include_object=None, include_schemas=None, process_revision_directives=None, compare_type=None, compare_server_default=None, render_item=None, literal_binds=None, upgrade_token=None, downgrade_token=None, alembic_module_prefix=None, sqlalchemy_module_prefix=None, user_module_prefix=None, **kw):
    """
    Configure the migration context.
    
    Args:
        connection (Connection): Database connection to use
        url (str): Database URL if connection not provided
        dialect_name (str): Name of dialect to use
        dialect_opts (dict): Dialect-specific options
        starting_rev (str): Starting revision for migration
        tag (str): Tag to apply to migration
        template_args (dict): Template variables
        render_as_batch (bool): Render operations as batch
        target_metadata (MetaData): Target metadata for comparison
        include_name (callable): Filter for object names
        include_object (callable): Filter for objects
        include_schemas (bool): Include schema information
        process_revision_directives (callable): Custom revision processing
        compare_type (bool): Enable type comparison
        compare_server_default (bool): Enable server default comparison
        render_item (callable): Custom rendering function
        literal_binds (bool): Use literal parameter binds
        upgrade_token (str): Token for upgrade operations
        downgrade_token (str): Token for downgrade operations
        alembic_module_prefix (str): Prefix for alembic imports
        sqlalchemy_module_prefix (str): Prefix for SQLAlchemy imports
        user_module_prefix (str): Prefix for user imports
        **kw: Additional configuration options
    """

Usage Example:

# In env.py
from alembic import context
from sqlalchemy import engine_from_config, pool

config = context.config
target_metadata = myapp.db.metadata

def run_migrations_online():
    connectable = engine_from_config(
        config.get_section(config.config_ini_section),
        prefix='sqlalchemy.',
        poolclass=pool.NullPool,
    )

    with connectable.connect() as connection:
        context.configure(
            connection=connection,
            target_metadata=target_metadata,
            compare_type=True,
            compare_server_default=True
        )

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

Database Connection Access

Access the current database connection and related objects.

def get_bind():
    """
    Get the current database connection.
    
    Returns:
        Connection: SQLAlchemy connection object bound to current context
    """

def get_context():
    """
    Get the current migration context.
    
    Returns:
        MigrationContext: Current MigrationContext instance
    """

Usage Examples:

# Get connection for custom operations
conn = context.get_bind()
result = conn.execute("SELECT COUNT(*) FROM users")

# Get migration context for advanced operations
ctx = context.get_context()
current_rev = ctx.get_current_revision()

Transaction Management

Manage transactions during migration execution.

def begin_transaction():
    """
    Begin a transaction context for migrations.
    
    Returns:
        Context manager for transaction scope
    """

Usage Example:

# Explicit transaction management
with context.begin_transaction():
    context.run_migrations()

Migration Execution

Execute migration operations and manage the migration process.

def run_migrations(**kw):
    """
    Run the migration operations.
    
    Args:
        **kw: Additional keyword arguments passed to migration functions
    """

def execute(sql, execution_options=None):
    """
    Execute SQL within the migration context.
    
    Args:
        sql (str|ClauseElement): SQL statement to execute
        execution_options (dict): Options for execution
    
    Returns:
        Result of SQL execution
    """

Usage Examples:

# Run migrations with custom arguments
context.run_migrations(custom_arg="value")

# Execute custom SQL
context.execute("CREATE EXTENSION IF NOT EXISTS uuid-ossp")

Revision Information

Access information about revisions and migration state.

def get_revision_argument():
    """
    Get the target revision argument.
    
    Returns:
        str: Target revision identifier
    """

def get_starting_revision_argument():
    """
    Get the starting revision argument.
    
    Returns:
        Union[str, Tuple[str, ...], None]: Starting revision identifier(s)
    """

def get_head_revision():
    """
    Get the head revision from the script directory.
    
    Returns:
        Union[str, Tuple[str, ...], None]: Head revision identifier(s)
    """

def get_head_revisions():
    """
    Get all head revisions from the script directory.
    
    Returns:
        Union[str, Tuple[str, ...], None]: All head revision identifiers
    """

def get_tag_argument():
    """
    Get the --tag argument value.
    
    Returns:
        Optional[str]: Tag argument value
    """

def get_x_argument(as_dictionary=False):
    """
    Get the -x argument values.
    
    Args:
        as_dictionary (bool): Return as dictionary instead of list
        
    Returns:
        Union[List[str], Dict[str, str]]: -x argument values
    """
    """
    Get the starting revision argument.
    
    Returns:
        str: Starting revision identifier
    """

def get_head_revision():
    """
    Get the current head revision.
    
    Returns:
        str: Head revision identifier
    """

def get_head_revisions():
    """
    Get all current head revisions.
    
    Returns:
        tuple: All head revision identifiers
    """

def get_tag_argument():
    """
    Get the tag argument for the migration.
    
    Returns:
        str|None: Tag value if specified
    """

Usage Examples:

# Check target revision
target = context.get_revision_argument()
if target == 'head':
    print("Migrating to latest revision")

# Get current heads
heads = context.get_head_revisions()
print(f"Current heads: {heads}")

Command Line Arguments

Access command-line arguments and custom parameters.

def get_x_argument(**kw):
    """
    Get -x arguments from command line.
    
    Args:
        **kw: Default values for arguments
    
    Returns:
        dict: Dictionary of -x arguments
    """

Usage Example:

# Access custom command line arguments
x_args = context.get_x_argument(dbname="default")
database_name = x_args.get('dbname')

# Usage: alembic -x dbname=testdb upgrade head

Mode Detection

Determine the current migration execution mode.

def is_offline_mode():
    """
    Check if running in offline mode.
    
    Returns:
        bool: True if in offline/SQL generation mode
    """

def is_transactional_ddl():
    """
    Check if DDL operations are transactional.
    
    Returns:
        bool: True if DDL operations support transactions
    """

Usage Examples:

# Different behavior for offline mode
if context.is_offline_mode():
    context.static_output("-- Offline mode migration")
else:
    context.execute("SELECT 1")

# Check DDL transaction support
if context.is_transactional_ddl():
    # Safe to use transactions around DDL
    pass

Output Management

Manage output for offline/SQL generation mode.

def static_output(text):
    """
    Emit text output in offline mode.
    
    Args:
        text (str): Text to output
    """

Usage Example:

# Output comments in offline mode
if context.is_offline_mode():
    context.static_output("-- Custom migration step")
    context.static_output("CREATE INDEX CONCURRENTLY idx_user_email ON users(email);")

Common Patterns

Standard env.py Structure

from alembic import context
from sqlalchemy import engine_from_config, pool
from logging.config import fileConfig

# Alembic Config object
config = context.config

# Setup logging
if config.config_file_name is not None:
    fileConfig(config.config_file_name)

# Target metadata from your application
from myapp import db
target_metadata = db.metadata

def run_migrations_offline():
    """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():
    """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:
        context.configure(
            connection=connection, 
            target_metadata=target_metadata
        )

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

if context.is_offline_mode():
    run_migrations_offline()
else:
    run_migrations_online()

Custom Migration Logic

def run_migrations_online():
    connectable = engine_from_config(
        config.get_section(config.config_ini_section),
        prefix="sqlalchemy.",
        poolclass=pool.NullPool,
    )

    with connectable.connect() as connection:
        # Custom configuration
        context.configure(
            connection=connection,
            target_metadata=target_metadata,
            compare_type=True,
            compare_server_default=True,
            include_object=include_object,
            include_schemas=True,
            process_revision_directives=process_revision_directives
        )

        with context.begin_transaction():
            # Custom pre-migration logic
            if context.get_revision_argument() == 'head':
                context.execute("SET lock_timeout = '2s'")
            
            context.run_migrations()
            
            # Custom post-migration logic  
            context.execute("ANALYZE")

def include_object(object, name, type_, reflected, compare_to):
    """Filter objects during autogenerate."""
    if type_ == "table" and name.startswith("temp_"):
        return False
    return True

def process_revision_directives(context, revision, directives):
    """Custom revision processing."""
    if config.cmd_opts and config.cmd_opts.autogenerate:
        script = directives[0]
        if script.upgrade_ops.is_empty():
            directives[:] = []

Multi-Database Support

def run_migrations_online():
    engines = {
        'primary': engine_from_config(
            config.get_section('primary'),
            prefix='sqlalchemy.'
        ),
        'secondary': engine_from_config(
            config.get_section('secondary'), 
            prefix='sqlalchemy.'
        )
    }
    
    for name, engine in engines.items():
        with engine.connect() as connection:
            context.configure(
                connection=connection,
                target_metadata=get_metadata_for_db(name),
                upgrade_token=f"{name}_upgrade",
                downgrade_token=f"{name}_downgrade"
            )
            
            with context.begin_transaction():
                context.run_migrations()

Error Handling

Context functions may raise:

EnvironmentError: Configuration or setup errors
CommandError: Migration command errors
Database-specific exceptions during execution

Types

# Context objects
class MigrationContext:
    bind: Connection
    dialect: Dialect
    script: ScriptDirectory
    opts: Dict[str, Any]
    
    def get_current_revision(self): ...
    def stamp(self, script_directory, revision): ...

class EnvironmentContext:
    config: Config
    script: ScriptDirectory
    
    def configure(self, **kw): ...
    def run_migrations(self, **kw): ...

Install with Tessl CLI