migrate

Migration tool for generating and applying database migrations using LLMs.

apply_migration_file(conn, db, migration_name, filepath)

Apply a single migration file.

Source code in src/embar/tools/migrate.py

def apply_migration_file(conn: psycopg.Connection, db: PgDb, migration_name: str, filepath: str) -> None:
    """Apply a single migration file."""
    # Read the SQL file
    with open(filepath, "r") as f:
        sql_content = f.read()

    # Extract SQL statements (skip comment lines starting with --)
    current_statement: list[str] = []
    for line in sql_content.split("\n"):
        stripped = line.strip()
        # Skip comment lines and empty lines
        if stripped.startswith("--") or not stripped:
            continue
        current_statement.append(line)

    # Join all non-comment lines
    full_sql = "\n".join(current_statement).strip()

    if not full_sql:
        print(f"  {yellow('⊘ No SQL to execute (comments only)')}")
        return

    # Record migration start
    with conn.cursor() as cur:
        cur.execute(
            """
            INSERT INTO _embar_migrations (migration_name, started_at)
            VALUES (%s, NOW())
        """,
            (migration_name,),
        )
    conn.commit()

    try:
        # Execute the SQL
        db.execute(QuerySingle(full_sql))

        # Record migration completion
        with conn.cursor() as cur:
            cur.execute(
                """
                UPDATE _embar_migrations
                SET finished_at = NOW()
                WHERE migration_name = %s
            """,
                (migration_name,),
            )
        conn.commit()

        print(f"  {green('✓ Applied successfully')}")

    except Exception as e:
        print(f"  {red_bold(f'✗ Error: {e}')}")
        conn.rollback()
        raise

check_migration_state(conn)

Check if any migrations are in an invalid state (started but not finished).

Source code in src/embar/tools/migrate.py

def check_migration_state(conn: psycopg.Connection) -> None:
    """Check if any migrations are in an invalid state (started but not finished)."""
    with conn.cursor() as cur:
        cur.execute("""
            SELECT migration_name, started_at
            FROM _embar_migrations
            WHERE started_at IS NOT NULL AND finished_at IS NULL
            ORDER BY started_at
        """)
        incomplete = cur.fetchall()

    if incomplete:
        print(red_bold("Error: Database is in an invalid state!"))
        print("The following migrations were started but not completed:")
        for name, started_at in incomplete:
            print(f"  - {name} (started at {started_at})")
        print("\nPlease resolve this manually before running new migrations.")
        sys.exit(1)

ensure_migrations_table(conn)

Create _embar_migrations table if it doesn't exist.

Source code in src/embar/tools/migrate.py

def ensure_migrations_table(conn: psycopg.Connection) -> None:
    """Create _embar_migrations table if it doesn't exist."""
    with conn.cursor() as cur:
        cur.execute("""
            CREATE TABLE IF NOT EXISTS _embar_migrations (
                migration_name TEXT PRIMARY KEY,
                started_at TIMESTAMP NOT NULL DEFAULT NOW(),
                finished_at TIMESTAMP
            )
        """)
    conn.commit()

execute_migrations(diffs, db)

Execute migrations with user confirmation for each one.

Source code in src/embar/tools/migrate.py

def execute_migrations(diffs: list[MigrationDiff], db: PgDb) -> None:
    """Execute migrations with user confirmation for each one."""
    print(f"\n{yellow('EXECUTE MODE ENABLED')}")
    print("You will be prompted to confirm each migration.\n")

    for i, diff in enumerate(diffs, 1):
        if not _confirm_migration(diff, i, len(diffs)):
            print(red_bold("\n✗ Migration cancelled by user. Exiting."))
            sys.exit(0)

        # Execute the migration
        print("Executing...")
        try:
            db.execute(QuerySingle(diff.sql))
            print(green("✓ Migration executed successfully"))
        except Exception as e:
            print(red_bold(f"✗ Error executing migration: {e}"))
            sys.exit(1)

    print(f"\n{green('✓ All migrations executed successfully!')}")

generate_diffs(config, api_key, llm)

Generate migration diffs from database and schema comparison.

Source code in src/embar/tools/migrate.py

def generate_diffs(config: MigrateConfig, api_key: str, llm: Llm) -> list[MigrationDiff]:
    """Generate migration diffs from database and schema comparison."""
    print(f"Connecting to database: {config.db_url}")
    print(f"Loading schema from: {config.schema_path}")
    print("")

    conn = psycopg.connect(config.db_url)

    try:
        diffs = _create_migrations(config, api_key, conn, llm)
    except Exception as e:
        print(f"Error generating migrations: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)

    return diffs

get_applied_migrations(conn)

Get set of migration names that have been successfully applied.

Source code in src/embar/tools/migrate.py

def get_applied_migrations(conn: psycopg.Connection) -> set[str]:
    """Get set of migration names that have been successfully applied."""
    with conn.cursor() as cur:
        cur.execute("""
            SELECT migration_name
            FROM _embar_migrations
            WHERE finished_at IS NOT NULL
            ORDER BY finished_at
        """)
        return {row[0] for row in cur.fetchall()}

get_migration_files(migrations_dir)

Get list of migration files sorted by timestamp. Returns (filename, filepath) tuples.

Source code in src/embar/tools/migrate.py

def get_migration_files(migrations_dir: str) -> list[tuple[str, str]]:
    """Get list of migration files sorted by timestamp. Returns (filename, filepath) tuples."""
    if not os.path.exists(migrations_dir):
        return []

    files: list[tuple[str, str]] = []
    for filename in sorted(os.listdir(migrations_dir)):
        if filename.endswith(".sql"):
            filepath = os.path.join(migrations_dir, filename)
            # Extract migration name (remove .sql extension)
            migration_name = filename[:-4]
            files.append((migration_name, filepath))

    return files

get_schema_from_db(conn)

Extract current database schema as DDL objects.

Source code in src/embar/tools/migrate.py

def get_schema_from_db(conn: psycopg.Connection) -> list[Ddl]:
    """Extract current database schema as DDL objects."""
    results: list[Ddl] = []

    # Get enums
    enum_query = """
        SELECT
            t.typname as enum_name,
            string_agg(e.enumlabel, ', ' ORDER BY e.enumsortorder) as enum_values
        FROM pg_type t
        JOIN pg_enum e ON t.oid = e.enumtypid
        JOIN pg_namespace n ON t.typnamespace = n.oid
        WHERE n.nspname = 'public'
        GROUP BY t.typname
        ORDER BY t.typname
    """

    with conn.cursor() as cur:
        cur.execute(enum_query)
        for enum_name, enum_values in cur.fetchall():
            ddl = f"CREATE TYPE {enum_name} AS ENUM ({', '.join(f"'{v}'" for v in enum_values.split(', '))});"
            results.append(Ddl(name=enum_name, ddl=ddl))

    # Get tables
    table_query = """
        SELECT tablename
        FROM pg_tables
        WHERE schemaname = 'public'
        ORDER BY tablename
    """

    with conn.cursor() as cur:
        cur.execute(table_query)
        tables = [row[0] for row in cur.fetchall()]

    for table in tables:
        # Get column definitions
        column_query = """
            SELECT
                a.attname,
                pg_catalog.format_type(a.atttypid, a.atttypmod) as data_type,
                a.attnotnull,
                pg_get_expr(d.adbin, d.adrelid) as default_value
            FROM pg_attribute a
            LEFT JOIN pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum
            WHERE a.attrelid = %s::regclass
                AND a.attnum > 0
                AND NOT a.attisdropped
            ORDER BY a.attnum
        """

        with conn.cursor() as cur:
            cur.execute(column_query, (table,))
            columns = cur.fetchall()

        # Build CREATE TABLE statement
        col_defs: list[str] = []
        for col_name, data_type, not_null, default in columns:
            col_def = f"    {col_name} {data_type}"
            if default:
                col_def += f" DEFAULT {default}"
            if not_null:
                col_def += " NOT NULL"
            col_defs.append(col_def)

        ddl = f"CREATE TABLE {table} (\n" + ",\n".join(col_defs) + "\n);"

        # Get constraints
        constraint_query = """
            SELECT conname, pg_get_constraintdef(oid)
            FROM pg_constraint
            WHERE conrelid = %s::regclass
            ORDER BY contype, conname
        """

        constraints: list[str] = []
        with conn.cursor() as cur:
            cur.execute(constraint_query, (table,))
            for con_name, con_def in cur.fetchall():
                constraints.append(f"ALTER TABLE {table} ADD CONSTRAINT {con_name} {con_def};")

        results.append(Ddl(name=table, ddl=ddl, constraints=constraints))

    return results

save_migration_to_file(diffs, migrations_dir, migration_name)

Save migrations to a file in the migrations directory.

Source code in src/embar/tools/migrate.py

def save_migration_to_file(diffs: list[MigrationDiff], migrations_dir: str, migration_name: str) -> str:
    """Save migrations to a file in the migrations directory."""
    # Create migrations directory if it doesn't exist
    os.makedirs(migrations_dir, exist_ok=True)

    # Generate filename with timestamp
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{timestamp}_{migration_name}.sql"
    filepath = os.path.join(migrations_dir, filename)

    # Write migration to file
    content = format_migration_output(diffs)
    with open(filepath, "w") as f:
        f.write(content)

    return filepath