From 6df41265ed495e3c2d775589f027b50e3c09271f Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Tue, 28 Oct 2025 17:41:49 -0400 Subject: [PATCH] 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 --- migrations/env.py | 56 +++++- ..._replace_collection_name_with_observer_.py | 160 +++++++++++------- ...c5_add_session_name_column_to_documents.py | 18 +- ..._replace_collection_name_with_observer_.py | 4 +- 4 files changed, 171 insertions(+), 67 deletions(-) diff --git a/migrations/env.py b/migrations/env.py index 6787437e..f0a8197c 100644 --- a/migrations/env.py +++ b/migrations/env.py @@ -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 diff --git a/migrations/versions/08894082221a_replace_collection_name_with_observer_.py b/migrations/versions/08894082221a_replace_collection_name_with_observer_.py index 87067d67..481ae33f 100644 --- a/migrations/versions/08894082221a_replace_collection_name_with_observer_.py +++ b/migrations/versions/08894082221a_replace_collection_name_with_observer_.py @@ -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, diff --git a/migrations/versions/564ba40505c5_add_session_name_column_to_documents.py b/migrations/versions/564ba40505c5_add_session_name_column_to_documents.py index d7238385..36274e12 100644 --- a/migrations/versions/564ba40505c5_add_session_name_column_to_documents.py +++ b/migrations/versions/564ba40505c5_add_session_name_column_to_documents.py @@ -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( diff --git a/tests/alembic/revisions/test_08894082221a_replace_collection_name_with_observer_.py b/tests/alembic/revisions/test_08894082221a_replace_collection_name_with_observer_.py index 30f9aeb6..2cf7951b 100644 --- a/tests/alembic/revisions/test_08894082221a_replace_collection_name_with_observer_.py +++ b/tests/alembic/revisions/test_08894082221a_replace_collection_name_with_observer_.py @@ -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(