FIX Batch v1.0 conversation migration backfills

pyrit/memory/alembic/versions/b2f4c6a8d1e3_add_conversations_table.py

@@ -21,10 +21,14 @@
 
 import logging
 from collections.abc import Sequence  # noqa: TC003
+from typing import TYPE_CHECKING
 
 import sqlalchemy as sa
 from alembic import op
 
+if TYPE_CHECKING:
+    from sqlalchemy.engine import Connection
+
 # revision identifiers, used by Alembic.
 revision: str = "b2f4c6a8d1e3"
 down_revision: str | None = "f1a2b3c4d5e6"
@@ -34,6 +38,24 @@
 
 logger = logging.getLogger(__name__)
 
+_CONVERSATION_INSERT_BATCH_SIZE = 400
+_CONVERSATION_INSERT_PROGRESS_INTERVAL = 25
+_CONVERSATION_INSERT_PREFIX = 'INSERT INTO "Conversations" (conversation_id, target_identifier, pyrit_version) VALUES '
+
+
+def _report_progress(message: str) -> None:
+    """Write migration progress to Alembic stdout, or the logger outside a migration context."""
+    try:
+        context = op.get_context()
+    except (AttributeError, NameError):
+        logger.info(message)
+        return
+    config = context.config
+    if config is not None:
+        config.print_stdout(message)
+    else:
+        logger.info(message)
+
 
 def upgrade() -> None:
     """Apply this schema upgrade."""
@@ -118,22 +140,39 @@ def _backfill_conversations() -> None:
         if conversation_id is not None and conversation_id not in targets_by_conversation:
             targets_by_conversation[conversation_id] = None
 
-    insert_stmt = sa.text(
-        'INSERT INTO "Conversations" (conversation_id, target_identifier, pyrit_version) '
-        "VALUES (:cid, :target, :version)"
-    )
-
-    inserted = 0
-    for conversation_id, target_identifier in targets_by_conversation.items():
-        if conversation_id in existing_ids:
-            continue
-        bind.execute(
-            insert_stmt,
-            {"cid": conversation_id, "target": target_identifier, "version": None},
-        )
-        inserted += 1
+    rows_to_insert = [
+        (conversation_id, target_identifier)
+        for conversation_id, target_identifier in targets_by_conversation.items()
+        if conversation_id not in existing_ids
+    ]
+    _insert_conversation_rows(bind=bind, rows=rows_to_insert)
 
+    inserted = len(rows_to_insert)
     if inserted or conflict_warnings:
         logger.info(
             f"Conversations backfill: inserted {inserted} row(s); {conflict_warnings} target-conflict warning(s)."
         )
+
+
+def _insert_conversation_rows(*, bind: Connection, rows: Sequence[tuple[str, str | None]]) -> None:
+    """Insert conversation rows in bounded multi-value statements."""
+    batch_count = (len(rows) + _CONVERSATION_INSERT_BATCH_SIZE - 1) // _CONVERSATION_INSERT_BATCH_SIZE
+    if not batch_count:
+        return
+
+    _report_progress(f"Conversations backfill: inserting {len(rows)} row(s) in {batch_count} batch(es).")
+    for batch_number, start in enumerate(range(0, len(rows), _CONVERSATION_INSERT_BATCH_SIZE), start=1):
+        _insert_conversation_batch(bind=bind, rows=rows[start : start + _CONVERSATION_INSERT_BATCH_SIZE])
+        if batch_number % _CONVERSATION_INSERT_PROGRESS_INTERVAL == 0 or batch_number == batch_count:
+            _report_progress(f"Conversations backfill: completed insert batch {batch_number}/{batch_count}.")
+
+
+def _insert_conversation_batch(*, bind: Connection, rows: Sequence[tuple[str, str | None]]) -> None:
+    """Insert one non-empty batch of conversation rows."""
+    value_clauses: list[str] = []
+    parameters: dict[str, str | None] = {}
+    for index, (conversation_id, target_identifier) in enumerate(rows):
+        value_clauses.append(f"(:cid_{index}, :target_{index}, NULL)")
+        parameters[f"cid_{index}"] = conversation_id
+        parameters[f"target_{index}"] = target_identifier
+    bind.execute(sa.text(_CONVERSATION_INSERT_PREFIX + ", ".join(value_clauses)), parameters)