fix(cockroachdb): Handle 30s timeout and ensure vector index usage

Two critical fixes for multi-node CockroachDB clusters:

1. Connection Timeout Handling:
   On multi-node v25.4 clusters, CREATE VECTOR INDEX from subprocess
   contexts experiences a 30-second connection timeout. The index
   creation continues successfully in the background. This fix detects
   the timeout and polls for completion (up to 5 minutes).

2. Vector Index Usage:
   Fixed vector_search_beam_size not being set on pooled connections,
   causing queries to use full table scan instead of the vector index.
   Now configures every connection from the pool with proper beam size.

Testing:
- Single-node: Works without timeout (178s index creation)
- Multi-node: Successfully handles timeout and completes (131s total)
- Vector index: Now properly used for all searches (verified with EXPLAIN)
- Both achieve ~83% recall with good QPS

Fixes issues where:
- Benchmarks would fail despite successful index creation
- Searches were slow due to full table scans instead of index usage

Author: viragtripathi

vectordb_bench/backend/clients/cockroachdb/cockroachdb.py

@@ -35,7 +35,7 @@ class CockroachDB(VectorDB):
         FilterOp.StrEqual,
     ]
 
-    def __init__(
+    def __init__(  # noqa: PLR0915
         self,
         dim: int,
         db_config: dict,
@@ -155,7 +155,7 @@ def _create_connection_pool(self) -> ConnectionPool:
         conninfo += " options='-c statement_timeout=600s'"
 
         # Configure each connection with vector support and search parameters
-        def configure_connection(conn):
+        def configure_connection(conn: Connection) -> None:
             register_vector(conn)
             # Set vector_search_beam_size on every connection for index usage
             if self.case_config is not None:
@@ -315,6 +315,56 @@ def _create_index(self):
         finally:
             conn.close()
 
+    def _wait_for_index_creation(self, start_time: float) -> None:
+        """Wait for background index creation to complete after connection timeout."""
+        import time
+
+        max_wait = 300  # 5 minutes max
+        poll_interval = 5
+        waited = 0
+
+        while waited < max_wait:
+            time.sleep(poll_interval)
+            waited += poll_interval
+
+            # Create fresh connection to check status
+            try:
+                check_conn = psycopg.connect(**self.connect_config)
+                check_cursor = check_conn.cursor()
+                try:
+                    # Check if index exists
+                    check_cursor.execute(
+                        "SELECT 1 FROM pg_indexes WHERE tablename = %s AND indexname = %s",
+                        (self.table_name, self._index_name),
+                    )
+                    if check_cursor.fetchone():
+                        # Index exists! Verify it's usable by doing a quick test query
+                        try:
+                            from psycopg import sql
+
+                            check_cursor.execute(
+                                sql.SQL("SELECT 1 FROM {} LIMIT 1").format(sql.Identifier(self.table_name))
+                            )
+                            check_cursor.fetchone()
+                            total_time = time.time() - start_time
+                            log.info(f"Index {self._index_name} created successfully (total time: {total_time:.1f}s)")
+                        except Exception as query_error:
+                            # Index not yet usable
+                            log.info(f"Index exists but not yet usable... ({waited}s elapsed, error: {query_error})")
+                        else:
+                            return
+                finally:
+                    check_cursor.close()
+                    check_conn.close()
+            except Exception as check_error:
+                log.warning(f"Error checking index status: {check_error}")
+                # Continue waiting
+
+        # Timeout waiting for index
+        msg = f"Timeout waiting for index {self._index_name} after {waited}s"
+        log.error(msg)
+        raise RuntimeError(msg)
+
     def optimize(self, data_size: int | None = None):
         """Post-insert optimization: create index if needed.
 
@@ -365,58 +415,15 @@ def optimize(self, data_size: int | None = None):
                 # Check if this is the expected 30s timeout on multi-node clusters
                 if "server closed the connection" in str(e) or "connection" in str(e).lower():
                     log.warning(f"Connection closed after {elapsed:.1f}s during index creation: {e}")
-                    log.info(f"This is expected on multi-node clusters - checking if index was created...")
+                    log.info("This is expected on multi-node clusters - checking if index was created...")
                     connection_closed = True
                 else:
                     # Unexpected error, re-raise
                     raise
 
             # Connection closed - wait for background index creation to complete
             if connection_closed:
-                max_wait = 300  # 5 minutes max
-                poll_interval = 5
-                waited = 0
-
-                while waited < max_wait:
-                    time.sleep(poll_interval)
-                    waited += poll_interval
-
-                    # Create fresh connection to check status
-                    try:
-                        check_conn = psycopg.connect(**self.connect_config)
-                        check_cursor = check_conn.cursor()
-                        try:
-                            # Check if index exists
-                            check_cursor.execute(
-                                "SELECT 1 FROM pg_indexes WHERE tablename = %s AND indexname = %s",
-                                (self.table_name, self._index_name),
-                            )
-                            if check_cursor.fetchone():
-                                # Index exists! Verify it's usable by doing a quick test query
-                                try:
-                                    check_cursor.execute(f"SELECT 1 FROM {self.table_name} LIMIT 1")
-                                    check_cursor.fetchone()
-                                    total_time = time.time() - start_time
-                                    log.info(
-                                        f"✅ Index {self._index_name} created successfully (total time: {total_time:.1f}s)"
-                                    )
-                                    return
-                                except Exception as query_error:
-                                    # Index not yet usable
-                                    log.info(
-                                        f"Index exists but not yet usable... ({waited}s elapsed, error: {query_error})"
-                                    )
-                        finally:
-                            check_cursor.close()
-                            check_conn.close()
-                    except Exception as check_error:
-                        log.warning(f"Error checking index status: {check_error}")
-                        # Continue waiting
-
-                # Timeout waiting for index
-                msg = f"Timeout waiting for index {self._index_name} after {waited}s"
-                log.error(msg)
-                raise RuntimeError(msg)
+                self._wait_for_index_creation(start_time)
 
     @db_retry(max_attempts=3, initial_delay=0.5, backoff_factor=2.0)
     def insert_embeddings(