RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: Ensure alembic migrations use a session pooler (#247) 

* fix: Ensure alembic migrations use a session pooler

* fix: bump batch size; fix document delete in 08894082221a (#249)

* chore: add alembic logging

* fix (alembic): remove order by in batches

* fix: fkey -> fk

* chore: code rabbit

---------

Co-authored-by: Rajat Ahuja <rahuja445@gmail.com>
This commit is contained in:
Vineeth Voruganti 2025-10-28 17:41:49 -04:00 committed by GitHub
parent 8a37a9570b
commit 6df41265ed
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 171 additions and 67 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
migrations/env.py
View File

@ -2,9 +2,10 @@ import logging
import sys
from logging.config import fileConfig
from pathlib import Path
from urllib.parse import urlparse, urlunparse
from alembic import context
from sqlalchemy import engine_from_config, pool, text
from sqlalchemy import engine_from_config, text
from src.config import settings
@ -59,7 +60,7 @@ def run_migrations_offline() -> None:
script output.
"""
url = get_url()
url = ensure_session_pooler(get_url())
context.configure(
url=url,
@ -74,6 +75,51 @@ def run_migrations_offline() -> None:
context.run_migrations()
def ensure_session_pooler(connection_uri: str) -> str:
"""
Ensure a PostgreSQL connection URI uses the session pooler port (5432).
Converts transaction pooler port (6543) to session pooler port (5432).
Leaves other ports unchanged.
Args:
connection_uri: PostgreSQL connection URI
Returns:
Connection URI with session pooler port (5432)
Examples:
>>> ensure_session_pooler("postgresql://user:pass@host:6543/db")
'postgresql://user:pass@host:5432/db'
>>> ensure_session_pooler("postgresql://user:pass@host:5432/db")
'postgresql://user:pass@host:5432/db'
>>> ensure_session_pooler("postgresql+psycopg://user:pass@host.supabase.co:6543/postgres")
'postgresql+psycopg://user:pass@host.supabase.co:5432/postgres'
"""
parsed = urlparse(connection_uri)
# Get current port, default to 5432 if not specified
current_port = parsed.port or 5432
# If using transaction pooler port (6543), switch to session pooler (5432)
if current_port == 6543:
# Replace the port in the netloc
if parsed.port:
# If port is explicitly in the URL, replace it
new_netloc = parsed.netloc.replace(f":{current_port}", ":5432")
else:
# If port not in URL but somehow detected, add it
new_netloc = f"{parsed.netloc}:5432"
# Reconstruct the URL with new port
new_parsed = parsed._replace(netloc=new_netloc)
return urlunparse(new_parsed)
return connection_uri
def run_migrations_online() -> None:
"""Run migrations in 'online' mode.
@ -87,13 +133,13 @@ def run_migrations_online() -> None:
configuration = {}
url = get_url()
configuration["sqlalchemy.url"] = url
validated_url = ensure_session_pooler(url)
configuration["sqlalchemy.url"] = validated_url
connectable = engine_from_config(
configuration,
prefix="sqlalchemy.",
echo=False,
poolclass=pool.NullPool,
echo=True,
connect_args={
"prepare_threshold": None,
"options": "-c statement_timeout=300000", # 5 minutes in milliseconds

160
migrations/versions/08894082221a_replace_collection_name_with_observer_.py
View File

@ -22,6 +22,9 @@ down_revision: str | None = "564ba40505c5"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
# Batch size for bulk operations
BATCH_SIZE = 10000
def upgrade() -> None:
"""Replace collections.name and documents.collection_name with observer and observed fields."""
@ -56,7 +59,6 @@ def upgrade() -> None:
{"session_id": session_id, "workspace_name": workspace_name},
)
# Update all documents with NULL session_name in batches
batch_size = 5000
while True:
result = connection.execute(
text(
@ -65,17 +67,15 @@ def upgrade() -> None:
SELECT id
FROM {schema}.documents
WHERE session_name IS NULL
ORDER BY id
LIMIT :batch_size
)
UPDATE {schema}.documents d
SET session_name = '__global_observations__'
FROM batch
WHERE d.id = batch.id
AND d.session_name IS NULL
"""
),
{"batch_size": batch_size},
{"batch_size": BATCH_SIZE},
)
if result.rowcount == 0:
break
@ -117,31 +117,37 @@ def upgrade() -> None:
# Step 2b: Delete documents that reference collections marked for deletion (in batches)
if collections_to_delete:
# Delete documents in batches
batch_size = 5000
for i in range(0, len(collections_to_delete), batch_size):
batch = collections_to_delete[i : i + batch_size]
collection_ids = [row.id for row in batch]
collection_ids = [row.id for row in collections_to_delete]
connection.execute(
# Delete in smaller chunks of actual documents to avoid locking issues
while True:
result = connection.execute(
text(
f"""
DELETE FROM {schema}.documents d
USING {schema}.collections c
WHERE d.collection_name = c.name
AND d.peer_name = c.peer_name
AND d.workspace_name = c.workspace_name
AND c.id = ANY(:collection_ids)
"""
WITH to_delete AS (
SELECT d.id
FROM {schema}.documents d
JOIN {schema}.collections c
ON d.collection_name = c.name
AND d.peer_name = c.peer_name
AND d.workspace_name = c.workspace_name
WHERE c.id = ANY(:collection_ids)
LIMIT :batch_size
)
DELETE FROM {schema}.documents
WHERE id IN (SELECT id FROM to_delete)
"""
),
{"collection_ids": collection_ids},
{"collection_ids": collection_ids, "batch_size": BATCH_SIZE},
)
if result.rowcount == 0:
break # No more rows to delete
# Step 2c: Delete the collections identified in step 2a (in batches)
if collections_to_delete:
batch_size = 5000
for i in range(0, len(collections_to_delete), batch_size):
batch = collections_to_delete[i : i + batch_size]
for i in range(0, len(collections_to_delete), BATCH_SIZE):
batch = collections_to_delete[i : i + BATCH_SIZE]
collection_ids = [row.id for row in batch]
connection.execute(
@ -161,7 +167,17 @@ def upgrade() -> None:
# - If name starts with peer_name + "_", extract the observed part (pattern: observer_observed)
# - If name ends with "_" + peer_name, extract the first part (pattern: observed_observer)
# - Any legacy edge cases will have been deleted in step 2a.
batch_size = 5000
# Create temporary index to speed up batching
if not index_exists("collections", "idx_temp_collections_null_observer"):
op.create_index(
"idx_temp_collections_null_observer",
"collections",
["id"],
postgresql_where=text("observer IS NULL OR observed IS NULL"),
schema=schema,
)
while True:
result = connection.execute(
text(
@ -170,7 +186,6 @@ def upgrade() -> None:
SELECT id
FROM {schema}.collections
WHERE observer IS NULL OR observed IS NULL
ORDER BY id
LIMIT :batch_size
)
UPDATE {schema}.collections c
@ -186,11 +201,17 @@ def upgrade() -> None:
WHERE c.id = batch.id
"""
),
{"batch_size": batch_size},
{"batch_size": BATCH_SIZE},
)
if result.rowcount == 0:
break
# Drop temporary index after batching is complete
if index_exists("collections", "idx_temp_collections_null_observer", inspector):
op.drop_index(
"idx_temp_collections_null_observer", "collections", schema=schema
)
# Step 3: Make collections observer and observed NOT NULL
op.alter_column("collections", "observer", nullable=False, schema=schema)
op.alter_column("collections", "observed", nullable=False, schema=schema)
@ -213,14 +234,24 @@ def upgrade() -> None:
# Step 5: Populate documents observer and observed from collections table
# Join to the already-populated collections table to get authoritative values
# Process in batches of 1000 to reduce query size
batch_size = 1000
# Process in batches to reduce query size
# Create temporary index to speed up batching
if not index_exists("documents", "idx_temp_docs_null_observer"):
op.create_index(
"idx_temp_docs_null_observer",
"documents",
["id"],
postgresql_where=text("observer IS NULL OR observed IS NULL"),
schema=schema,
)
while True:
result = connection.execute(
text(
f"""
WITH batch AS (
SELECT d.ctid
SELECT d.id
FROM {schema}.documents d
WHERE d.observer IS NULL OR d.observed IS NULL
LIMIT :batch_size
@ -230,18 +261,21 @@ def upgrade() -> None:
observer = c.observer,
observed = c.observed
FROM {schema}.collections c, batch
WHERE d.ctid = batch.ctid
WHERE d.id = batch.id
AND d.collection_name = c.name
AND d.peer_name = c.peer_name
AND d.workspace_name = c.workspace_name
AND (d.observer IS NULL OR d.observed IS NULL)
"""
),
{"batch_size": batch_size},
{"batch_size": BATCH_SIZE},
)
if result.rowcount == 0:
break
# Drop temporary index after batching is complete
if index_exists("documents", "idx_temp_docs_null_observer", inspector):
op.drop_index("idx_temp_docs_null_observer", "documents", schema=schema)
# Step 6: Make documents observer and observed NOT NULL
op.alter_column("documents", "observer", nullable=False, schema=schema)
op.alter_column("documents", "observed", nullable=False, schema=schema)
@ -308,10 +342,10 @@ def upgrade() -> None:
# Step 10: Add composite foreign key constraint for observer peer on collections
if not fk_exists(
"collections", "collections_observer_workspace_name_fkey", inspector
"collections", "fk_collections_observer_workspace_name_peers", inspector
):
op.create_foreign_key(
"collections_observer_workspace_name_fkey",
"fk_collections_observer_workspace_name_peers",
"collections",
"peers",
["observer", "workspace_name"],
@ -322,10 +356,10 @@ def upgrade() -> None:
# Step 11: Add composite foreign key constraint for observed peer on collections
if not fk_exists(
"collections", "collections_observed_workspace_name_fkey", inspector
"collections", "fk_collections_observed_workspace_name_peers", inspector
):
op.create_foreign_key(
"collections_observed_workspace_name_fkey",
"fk_collections_observed_workspace_name_peers",
"collections",
"peers",
["observed", "workspace_name"],
@ -336,10 +370,12 @@ def upgrade() -> None:
# Step 12: Add composite foreign key constraint from documents to collections using observer/observed
if not fk_exists(
"documents", "documents_observer_observed_workspace_name_fkey", inspector
"documents",
"fk_documents_observer_observed_workspace_name_collections",
inspector,
):
op.create_foreign_key(
"documents_observer_observed_workspace_name_fkey",
"fk_documents_observer_observed_workspace_name_collections",
"documents",
"collections",
["observer", "observed", "workspace_name"],
@ -349,9 +385,11 @@ def upgrade() -> None:
)
# Step 13: Add composite foreign key constraint for observer peer on documents
if not fk_exists("documents", "documents_observer_workspace_name_fkey", inspector):
if not fk_exists(
"documents", "fk_documents_observer_workspace_name_peers", inspector
):
op.create_foreign_key(
"documents_observer_workspace_name_fkey",
"fk_documents_observer_workspace_name_peers",
"documents",
"peers",
["observer", "workspace_name"],
@ -361,9 +399,11 @@ def upgrade() -> None:
)
# Step 14: Add composite foreign key constraint for observed peer on documents
if not fk_exists("documents", "documents_observed_workspace_name_fkey", inspector):
if not fk_exists(
"documents", "fk_documents_observed_workspace_name_peers", inspector
):
op.create_foreign_key(
"documents_observed_workspace_name_fkey",
"fk_documents_observed_workspace_name_peers",
"documents",
"peers",
["observed", "workspace_name"],
@ -498,7 +538,6 @@ def downgrade() -> None:
)
# Step 5: Populate documents collection_name from observer and observed in batches
batch_size = 5000
while True:
result = connection.execute(
text(
@ -507,7 +546,6 @@ def downgrade() -> None:
SELECT id
FROM {schema}.documents
WHERE collection_name IS NULL
ORDER BY id
LIMIT :batch_size
)
UPDATE {schema}.documents d
@ -520,7 +558,7 @@ def downgrade() -> None:
AND d.collection_name IS NULL
"""
),
{"batch_size": batch_size},
{"batch_size": BATCH_SIZE},
)
if result.rowcount == 0:
break
@ -537,7 +575,6 @@ def downgrade() -> None:
)
# Populate peer_name with observed value in batches
batch_size = 5000
while True:
result = connection.execute(
text(
@ -546,7 +583,6 @@ def downgrade() -> None:
SELECT id
FROM {schema}.documents
WHERE peer_name IS NULL
ORDER BY id
LIMIT :batch_size
)
UPDATE {schema}.documents d
@ -556,7 +592,7 @@ def downgrade() -> None:
AND d.peer_name IS NULL
"""
),
{"batch_size": batch_size},
{"batch_size": BATCH_SIZE},
)
if result.rowcount == 0:
break
@ -565,9 +601,11 @@ def downgrade() -> None:
op.alter_column("documents", "peer_name", nullable=False, schema=schema)
# Recreate the foreign key constraint for peer_name on documents
if not fk_exists("documents", "documents_peer_name_workspace_name_fkey", inspector):
if not fk_exists(
"documents", "fk_documents_peer_name_workspace_name_peers", inspector
):
op.create_foreign_key(
"documents_peer_name_workspace_name_fkey",
"fk_documents_peer_name_workspace_name_peers",
"documents",
"peers",
["peer_name", "workspace_name"],
@ -588,26 +626,28 @@ def downgrade() -> None:
# CONSTRAINTS AND INDEXES
# Step 8: Drop new foreign key constraints from documents
if fk_exists(
"documents", "documents_observer_observed_workspace_name_fkey", inspector
"documents",
"fk_documents_observer_observed_workspace_name_collections",
inspector,
):
op.drop_constraint(
"documents_observer_observed_workspace_name_fkey",
"fk_documents_observer_observed_workspace_name_collections",
"documents",
type_="foreignkey",
schema=schema,
)
if fk_exists("documents", "documents_observer_workspace_name_fkey", inspector):
if fk_exists("documents", "fk_documents_observer_workspace_name_peers", inspector):
op.drop_constraint(
"documents_observer_workspace_name_fkey",
"fk_documents_observer_workspace_name_peers",
"documents",
type_="foreignkey",
schema=schema,
)
if fk_exists("documents", "documents_observed_workspace_name_fkey", inspector):
if fk_exists("documents", "fk_documents_observed_workspace_name_peers", inspector):
op.drop_constraint(
"documents_observed_workspace_name_fkey",
"fk_documents_observed_workspace_name_peers",
"documents",
type_="foreignkey",
schema=schema,
@ -652,17 +692,21 @@ def downgrade() -> None:
)
# Step 12: Drop foreign key constraints from collections
if fk_exists("collections", "collections_observer_workspace_name_fkey", inspector):
if fk_exists(
"collections", "fk_collections_observer_workspace_name_peers", inspector
):
op.drop_constraint(
"collections_observer_workspace_name_fkey",
"fk_collections_observer_workspace_name_peers",
"collections",
type_="foreignkey",
schema=schema,
)
if fk_exists("collections", "collections_observed_workspace_name_fkey", inspector):
if fk_exists(
"collections", "fk_collections_observed_workspace_name_peers", inspector
):
op.drop_constraint(
"collections_observed_workspace_name_fkey",
"fk_collections_observed_workspace_name_peers",
"collections",
type_="foreignkey",
schema=schema,

18
migrations/versions/564ba40505c5_add_session_name_column_to_documents.py
View File

@ -38,7 +38,17 @@ def upgrade() -> None:
# Process in batches to avoid timeout with large datasets
# Only migrate documents that have 'session_name' key with a non-null, non-empty value
bind = op.get_bind()
batch_size = 5000
batch_size = 10000
# Create temporary index to speed up batching
if not index_exists("documents", "idx_temp_docs_null_session", inspector):
op.create_index(
"idx_temp_docs_null_session",
"documents",
["id"],
postgresql_where=sa.text("session_name IS NULL"),
schema=schema,
)
while True:
result = bind.execute(
@ -51,14 +61,12 @@ def upgrade() -> None:
AND internal_metadata ? 'session_name'
AND internal_metadata->>'session_name' IS NOT NULL
AND internal_metadata->>'session_name' != ''
ORDER BY id
LIMIT :batch_size
)
UPDATE {schema}.documents d
SET session_name = d.internal_metadata->>'session_name'
FROM batch b
WHERE d.id = b.id
AND d.session_name IS NULL
"""
),
{"batch_size": batch_size},
@ -67,6 +75,10 @@ def upgrade() -> None:
if result.rowcount == 0:
break
# Drop temporary index after batching is complete
if index_exists("documents", "idx_temp_docs_null_session"):
op.drop_index("idx_temp_docs_null_session", "documents", schema=schema)
# Step 3: Create index on session_name for efficient querying
if not index_exists("documents", "idx_documents_session_name", inspector):
op.create_index(

4
tests/alembic/revisions/test_08894082221a_replace_collection_name_with_observer_.py
View File

@ -262,7 +262,9 @@ def verify_observer_observed_migration(verifier: MigrationVerifier) -> None:
"collections", "unique_observer_observed_collection", "unique"
)
verifier.assert_constraint_exists(
"documents", "documents_observer_observed_workspace_name_fkey", "foreign_key"
"documents",
"fk_documents_observer_observed_workspace_name_collections",
"foreign_key",
)
collection = verifier.conn.execute(