from alembic import op
import sqlalchemy as sa
revision = "20240612_add_status"
down_revision = "20240515_create_orders"
branch_labels = None
depends_on = None
BATCH_SIZE = 5000
def upgrade():
op.add_column(
"orders",
sa.Column(
"status",
sa.String(length=32),
nullable=True,
server_default="pending",
),
)
_backfill_status()
op.alter_column("orders", "status", nullable=False)
op.create_check_constraint(
"ck_orders_status_allowed",
"orders",
"status IN ('pending', 'paid', 'shipped', 'cancelled')",
)
def _backfill_status():
bind = op.get_bind()
stmt = sa.text(
"""
UPDATE orders
SET status = 'pending'
WHERE id IN (
SELECT id FROM orders
WHERE status IS NULL
ORDER BY id
LIMIT :limit
)
RETURNING id
"""
)
while True:
result = bind.execute(stmt, {"limit": BATCH_SIZE})
if result.rowcount == 0:
break
def downgrade():
op.drop_constraint("ck_orders_status_allowed", "orders", type_="check")
op.drop_column("orders", "status")
import enum
from sqlalchemy import (
BigInteger,
CheckConstraint,
Column,
DateTime,
String,
func,
)
from sqlalchemy.orm import declarative_base
Base = declarative_base()
class OrderStatus(str, enum.Enum):
pending = "pending"
paid = "paid"
shipped = "shipped"
cancelled = "cancelled"
class Order(Base):
__tablename__ = "orders"
__table_args__ = (
CheckConstraint(
"status IN ('pending', 'paid', 'shipped', 'cancelled')",
name="ck_orders_status_allowed",
),
)
id = Column(BigInteger, primary_key=True)
status = Column(
String(32),
nullable=False,
server_default="pending",
index=True,
)
created_at = Column(DateTime(timezone=True), server_default=func.now())
def mark(self, new_status):
if not isinstance(new_status, OrderStatus):
new_status = OrderStatus(new_status)
self.status = new_status.value
return self
Adding a NOT NULL column to a large, live table is a classic migration trap: a single statement that sets the constraint while backfilling every existing row takes a long exclusive lock and can stall writes for the duration of the rewrite. The safe pattern is to split the change into phases, and these files show that split across an Alembic migration and the SQLAlchemy model that eventually depends on it.
In versions/20240612_add_status.py, the upgrade function first adds status as a plain nullable column with a server-side default. Adding a nullable column with a constant default is a fast, metadata-only operation on modern Postgres, so it does not rewrite the table or block other writers. The migration then backfills existing rows in bounded chunks inside _backfill_status, committing after each batch of BATCH_SIZE rows. Batching keeps each transaction short, so autovacuum and concurrent queries are never starved by one giant transaction holding locks and bloating the WAL.
Only after every row has a value does the migration issue ALTER COLUMN ... SET NOT NULL. On Postgres 12+ this validation is cheap when a matching CHECK (status IS NOT NULL) constraint is already VALID, but even the plain form here runs quickly because the backfill guarantees no NULLs remain. The downgrade simply drops the column, which is why the schema change is reversible.
The _backfill_status helper uses op.get_bind() to run raw UPDATE statements through the same connection Alembic owns, filtering on status IS NULL and a primary-key window so each pass touches a disjoint slice. The RETURNING id clause lets the loop detect when a batch came up empty and stop.
In models.py, the Order model declares status with nullable=False and a server_default, mirroring the final migration state so the ORM and database agree. The server_default matters: without it, rows inserted by paths that bypass the application default would violate the constraint. The Enum-backed OrderStatus and the __table_args__ check constraint document the allowed values and let the database reject bad data independently of Python. The trade-off is operational complexity — three steps instead of one — but it is the difference between a migration that runs invisibly and one that causes an outage.
Related snips
package dbutil
import (
"context"
"github.com/jackc/pgconn"Retry Postgres serialization failures with bounded attempts
require "csv"
class PeopleCsvStream
include Enumerable
HEADERS = %w[id full_name email signed_up_at plan].freezeResilient CSV Export as a Streamed Response
type User {
id: ID!
name: String!
email: String!
posts: [Post!]!
createdAt: String!GraphQL API with Spring Boot
# Vulnerable: user input is concatenated directly into SQL.
email = params[:email]
password = params[:password]
sql = "SELECT * FROM users WHERE email = '#{email}' AND password_hash = '#{password}'"
user = ActiveRecord::Base.connection.execute(sql).firstSQL injection prevention with unsafe and safe query patterns
class AddSettingsToAccounts < ActiveRecord::Migration[7.1]
disable_ddl_transaction!
def change
add_column :accounts, :settings, :jsonb, null: false, default: {}Postgres JSONB Partial Index for Feature Flags
-- EXPLAIN ANALYZE (actual execution statistics)
EXPLAIN ANALYZE
SELECT u.username, COUNT(o.id) AS order_count
FROM users u
LEFT JOIN orders o ON u.id = o.user_id
WHERE u.created_at >= '2024-01-01'Advanced query optimization techniques
Share this code
Here's the card — post it anywhere.
