feat(pool): add connection pooling support

- Introduce ConnectionPool class for managing SQLite connections.
- Update SQLiteVecClient to support optional connection pooling.
- Add example demonstrating connection pooling usage.
- Implement comprehensive tests for connection pooling functionality.

Author: atasoglu

CHANGELOG.md

@@ -5,6 +5,20 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [2.1.0] - 2025-01-31
+
+### Added
+- Connection pooling support via `ConnectionPool` class
+- Thread-safe connection reuse for concurrent access scenarios
+- Optional `pool` parameter in `SQLiteVecClient` constructor
+- `examples/connection_pool_example.py` demonstrating pooled connections
+- Comprehensive tests for connection pooling in `tests/test_pool.py`
+
+### Improved
+- Better resource management for multi-threaded applications
+- Connection reuse reduces overhead in high-concurrency scenarios
+- Backward compatible - pooling is optional
+
 ## [2.0.0] - 2025-01-30
 
 ### Changed

TODO

@@ -62,7 +62,7 @@
 ## 🔵 Low Priority (New Features)
 
 ### Performance
-- [ ] Connection pooling support
+- [x] Connection pooling support
 - [x] Batch update operation
 - [x] Lazy loading option
 - [ ] Index strategy documentation

examples/connection_pool_example.py

@@ -0,0 +1,78 @@
+"""Example demonstrating connection pooling for concurrent access."""
+
+import concurrent.futures
+import time
+
+from sqlite_vec_client import ConnectionPool, SQLiteVecClient
+
+
+def worker_task(worker_id: int, pool: ConnectionPool) -> tuple[int, float]:
+    """Simulate a worker performing database operations.
+
+    Args:
+        worker_id: Worker identifier
+        pool: Connection pool to use
+
+    Returns:
+        Tuple of (worker_id, execution_time)
+    """
+    start = time.time()
+
+    # Create client with pooled connection (no db_path needed)
+    client = SQLiteVecClient(table=f"docs_{worker_id}", pool=pool)
+
+    # Create table if needed
+    client.create_table(dim=384)
+
+    # Add some data
+    texts = [f"Document {i} from worker {worker_id}" for i in range(10)]
+    embeddings = [[0.1 * i] * 384 for i in range(10)]
+    client.add(texts, embeddings)
+
+    # Perform similarity search
+    query = [0.5] * 384
+    _ = client.similarity_search(query, top_k=5)
+
+    # Close (returns connection to pool)
+    client.close()
+
+    elapsed = time.time() - start
+    return worker_id, elapsed
+
+
+def main() -> None:
+    """Demonstrate connection pooling with concurrent workers."""
+    print("Connection Pooling Example")
+    print("=" * 50)
+
+    # Create connection pool
+    pool = ConnectionPool(
+        connection_factory=lambda: SQLiteVecClient.create_connection("./pooled.db"),
+        pool_size=5,
+    )
+
+    print("Created connection pool with size=5\n")
+
+    # Run concurrent workers
+    num_workers = 10
+    print(f"Running {num_workers} concurrent workers...")
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
+        futures = [executor.submit(worker_task, i, pool) for i in range(num_workers)]
+        results = [f.result() for f in concurrent.futures.as_completed(futures)]
+
+    # Print results
+    print("\nWorker execution times:")
+    for worker_id, elapsed in sorted(results):
+        print(f"  Worker {worker_id}: {elapsed:.3f}s")
+
+    avg_time = sum(t for _, t in results) / len(results)
+    print(f"\nAverage execution time: {avg_time:.3f}s")
+
+    # Cleanup
+    pool.close_all()
+    print("\nPool closed successfully")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "sqlite-vec-client"
-version = "2.0.0"
+version = "2.1.0"
 description = "A tiny Python client around sqlite-vec for CRUD and similarity search."
 readme = "README.md"
 requires-python = ">=3.9"

sqlite_vec_client/__init__.py

@@ -13,9 +13,11 @@
     VecClientError,
 )
 from .logger import get_logger
+from .pool import ConnectionPool
 
 __all__ = [
     "SQLiteVecClient",
+    "ConnectionPool",
     "VecClientError",
     "ValidationError",
     "TableNameError",

sqlite_vec_client/base.py

@@ -12,7 +12,7 @@
 from collections.abc import Generator
 from contextlib import contextmanager
 from types import TracebackType
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 import sqlite_vec
 
@@ -29,6 +29,9 @@
     validate_top_k,
 )
 
+if TYPE_CHECKING:
+    from .pool import ConnectionPool
+
 logger = get_logger()
 
 
@@ -91,22 +94,34 @@ def rows_to_results(rows: list[sqlite3.Row]) -> list[Result]:
             for row in rows
         ]
 
-    def __init__(self, table: str, db_path: str) -> None:
+    def __init__(
+        self, table: str, db_path: str | None = None, pool: ConnectionPool | None = None
+    ) -> None:
         """Initialize the client for a given base table and database file.
 
         Args:
             table: Name of the base table
-            db_path: Path to SQLite database file
+            db_path: Path to SQLite database file (required if pool is None)
+            pool: Optional connection pool for connection reuse
 
         Raises:
             TableNameError: If table name is invalid
             VecConnectionError: If connection fails
+            ValueError: If both db_path and pool are None
         """
         validate_table_name(table)
         self.table = table
         self._in_transaction = False
+        self._pool = pool
         logger.debug(f"Initializing SQLiteVecClient for table: {table}")
-        self.connection = self.create_connection(db_path)
+
+        if pool:
+            self.connection = pool.get_connection()
+            logger.debug("Using connection from pool")
+        elif db_path:
+            self.connection = self.create_connection(db_path)
+        else:
+            raise ValueError("Either db_path or pool must be provided")
 
     def __enter__(self) -> SQLiteVecClient:
         """Support context manager protocol and return `self`."""
@@ -518,10 +533,14 @@ def transaction(self) -> Generator[None, None, None]:
             self._in_transaction = False
 
     def close(self) -> None:
-        """Close the underlying SQLite connection, suppressing close errors."""
+        """Close or return the connection to pool, suppressing close errors."""
         try:
             logger.debug(f"Closing connection for table '{self.table}'")
-            self.connection.close()
-            logger.info(f"Connection closed for table '{self.table}'")
+            if self._pool:
+                self._pool.return_connection(self.connection)
+                logger.info(f"Connection returned to pool for table '{self.table}'")
+            else:
+                self.connection.close()
+                logger.info(f"Connection closed for table '{self.table}'")
         except Exception as e:
             logger.warning(f"Error closing connection: {e}")

sqlite_vec_client/pool.py

@@ -0,0 +1,84 @@
+"""Connection pooling for SQLite connections."""
+
+from __future__ import annotations
+
+import sqlite3
+import threading
+from typing import Callable
+
+from .logger import get_logger
+
+logger = get_logger()
+
+
+class ConnectionPool:
+    """Thread-local connection pool for SQLite databases.
+
+    Note: SQLite connections are not thread-safe. This pool creates
+    one connection per thread and reuses it within that thread.
+    """
+
+    def __init__(
+        self,
+        connection_factory: Callable[[], sqlite3.Connection],
+        pool_size: int = 5,
+    ) -> None:
+        """Initialize connection pool.
+
+        Args:
+            connection_factory: Function that creates new connections
+            pool_size: Maximum number of connections (one per thread)
+        """
+        if pool_size < 1:
+            raise ValueError("pool_size must be at least 1")
+
+        self._factory = connection_factory
+        self._pool_size = pool_size
+        self._local = threading.local()
+        self._lock = threading.Lock()
+        self._created = 0
+        self._all_connections: list[sqlite3.Connection] = []
+        logger.debug(f"Initialized connection pool with size={pool_size}")
+
+    def get_connection(self) -> sqlite3.Connection:
+        """Get a connection for the current thread.
+
+        Returns:
+            SQLite connection for this thread
+        """
+        if hasattr(self._local, "connection"):
+            conn: sqlite3.Connection = self._local.connection
+            logger.debug("Reused thread-local connection")
+            return conn
+
+        with self._lock:
+            if self._created >= self._pool_size:
+                raise RuntimeError(f"Connection pool exhausted (max={self._pool_size})")
+            conn = self._factory()
+            self._created += 1
+            self._all_connections.append(conn)
+            self._local.connection = conn
+            logger.debug(f"Created new connection ({self._created}/{self._pool_size})")
+            return conn
+
+    def return_connection(self, conn: sqlite3.Connection) -> None:
+        """Return a connection (no-op for thread-local pool).
+
+        Args:
+            conn: Connection to return (kept in thread-local storage)
+        """
+        logger.debug("Connection kept in thread-local storage")
+
+    def close_all(self) -> None:
+        """Close all connections in the pool."""
+        logger.debug("Closing all connections in pool")
+        with self._lock:
+            for conn in self._all_connections:
+                try:
+                    conn.close()
+                except Exception as e:
+                    logger.warning(f"Error closing connection: {e}")
+            closed = len(self._all_connections)
+            self._all_connections.clear()
+            self._created = 0
+        logger.info(f"Closed {closed} connections from pool")