feat(cache): Added LRU (Least Recently Used) cache for prepared statements

Author: Luis Garza

redshift_connector/core.py

@@ -2,7 +2,7 @@
 import os
 import socket
 import typing
-from collections import deque
+from collections import deque, OrderedDict
 from copy import deepcopy
 from datetime import datetime as Datetime
 from datetime import timedelta as Timedelta
@@ -1826,11 +1826,14 @@ def execute(self: "Connection", cursor: Cursor, operation: str, vals) -> None:
  # transforms user provided bind parameters to server friendly bind parameters
  params: typing.Tuple[typing.Optional[typing.Tuple[int, int, typing.Callable]], ...] = ()
  has_bind_parameters: bool = False if vals is None else True
- # multi dimensional dictionary to store the data
- # cache = self._caches[cursor.paramstyle][pid]
- # cache = {'statement': {}, 'ps': {}}
- # statement stores the data of the statement, ps store the data of the prepared statement
- # statement = {operation(query): tuple from 'convert_paramstyle'(statement, make_args)}
+ statements_to_close = []
+ # Cache structure for prepared statements:
+ # self._caches[paramstyle][pid] contains:
+ # - 'statement': stores SQL statements and their parameter processors
+ # - 'ps': stores prepared statement metadata
+ # - 'statement_dict': OrderedDict tracking most recently used prepared statements
+ # (used when max_prepared_statements is set)
+ # Each statement entry contains (processed_statement, parameter_binding_function)
  try:
  cache = self._caches[cursor.paramstyle][pid]
  except KeyError:
@@ -1842,7 +1845,11 @@ def execute(self: "Connection", cursor: Cursor, operation: str, vals) -> None:
  try:
  cache = param_cache[pid]
  except KeyError:
- cache = param_cache[pid] = {"statement": {}, "ps": {}}
+ cache = param_cache[pid] = {
+ "statement": {},
+ "ps": {},
+ "statement_dict": OrderedDict() if self.max_prepared_statements > 0 else None
+ }
 
  try:
  statement, make_args = cache["statement"][operation]
@@ -1863,6 +1870,9 @@ def execute(self: "Connection", cursor: Cursor, operation: str, vals) -> None:
 
  try:
  ps = cache["ps"][key]
+ # If statement exists, move it to end of ordered dict (most recently used)
+ if self.max_prepared_statements > 0 and 'statement_dict' in cache and key in cache["statement_dict"]:
+ cache["statement_dict"].move_to_end(key)
  _logger.debug("Using cached prepared statement")
  cursor.ps = ps
  except KeyError:
@@ -1978,12 +1988,21 @@ def execute(self: "Connection", cursor: Cursor, operation: str, vals) -> None:
 
  ps["bind_2"] = h_pack(len(output_fc)) + pack("!" + "h" * len(output_fc), *output_fc)
 
- if len(cache["ps"]) >= self.max_prepared_statements:
- for p in cache["ps"].values():
- self.close_prepared_statement(p["statement_name_bin"])
- cache["ps"].clear()
  if self.max_prepared_statements > 0:
+ # Ensure consistency between ps and statement_dict
+ if len(cache["ps"]) != len(cache["statement_dict"]):
+ for existing_key in cache["ps"]:
+ cache["statement_dict"][existing_key] = None
+
+ # If cache is full, remove oldest statement
+ if len(cache["ps"]) >= self.max_prepared_statements:
+ oldest_key, _ = cache["statement_dict"].popitem(last=False)
+ statements_to_close.append(cache["ps"][oldest_key]["statement_name_bin"])
+ del cache["ps"][oldest_key]
+
+ # Add new statement to cache and queue
  cache["ps"][key] = ps
+ cache["statement_dict"][key] = None
 
  cursor._cached_rows.clear()
  cursor._row_count = -1
@@ -2031,6 +2050,10 @@ def execute(self: "Connection", cursor: Cursor, operation: str, vals) -> None:
  else:
  self.handle_messages(cursor)
 
+ # Clean up prepared statements after query execution and results are returned
+ for stmt in statements_to_close:
+ self.close_prepared_statement(stmt)
+
  def _send_message(self: "Connection", code: bytes, data: bytes) -> None:
  _logger.debug("Sending message with code %s to BE", code)
  try:

test/integration/test_prepared_statements_cache.py

@@ -0,0 +1,70 @@
+import pytest
+import redshift_connector
+
+def test_lru_prepared_statements_cache(db_kwargs):
+ """
+ Tests LRU (Least Recently Used) behavior of prepared statements cache:
+ 1. Verifies cache size limits
+ 2. Verifies statement ordering
+ 3. Verifies LRU eviction policy
+ 4. Verifies statement reuse behavior
+ """
+ db_kwargs['max_prepared_statements'] = 5
+ conn = redshift_connector.connect(**db_kwargs)
+ cursor = conn.cursor()
+
+ try:
+ # Track statement execution order
+ executed_statements = []
+
+ # Execute 7 unique statements (exceeds cache size of 5)
+ for i in range(7):
+ query = f"SELECT %s::int as col1, %s::int as col2, {i} as unique_id"
+ cursor.execute(query, (i, i + 1))
+ executed_statements.append(query)
+
+ # Get cache and statement queue
+ cache = conn._caches[cursor.paramstyle][cursor.ps['pid']]
+ statement_dict = cache['statement_dict']
+
+ # Basic cache size verification
+ assert len(cache['ps']) == 5, f"Cache size should be 5, but was {len(cache['ps'])}"
+ assert len(statement_dict) == 5, f"Statement dict size should be 5, but was {len(statement_dict)}"
+
+ # Verify the most recent statements are in the queue
+ cached_statements = [key[0] for key in statement_dict.keys()]
+ last_five_statements = executed_statements[-5:]
+ assert all(stmt in cached_statements for stmt in last_five_statements), \
+ "Last 5 statements should be in cache"
+
+ # Verify the first two statements were evicted
+ first_two_statements = executed_statements[:2]
+ for stmt in first_two_statements:
+ assert stmt not in cached_statements, f"Statement should have been evicted: {stmt}"
+
+ # Test statement reuse
+ reuse_stmt = executed_statements[-3]
+ cursor.execute(reuse_stmt, (100, 101))
+
+ # Verify the reused statement is now at the end of the queue
+ assert list(statement_dict.keys())[-1][0] == reuse_stmt, "Reused statement should be most recent"
+
+ # Add new statement and verify LRU behavior
+ new_stmt = "SELECT %s::int as col1, %s::int as col2, 999 as unique_id"
+ # Track which statement should be evicted (least recently used)
+ statements_before_new = [key[0] for key in statement_dict]
+ cursor.execute(new_stmt, (999, 1000))
+
+ # Verify cache size and new statement presence
+ assert len(cache['ps']) == 5, f"Cache size should still be 5, but was {len(cache['ps'])}"
+ assert list(statement_dict.keys())[-1][0] == new_stmt, "New statement should be most recent"
+
+ # Verify LRU eviction - the least recently used statement should be gone
+ current_statements = [key[0] for key in statement_dict]
+ lru_statement = statements_before_new[0] # First statement in queue before adding new one
+ assert lru_statement not in current_statements, \
+ f"Least recently used statement should have been evicted: {lru_statement}"
+
+ finally:
+ cursor.close()
+ conn.close()

test/integration/test_query.py

@@ -474,8 +474,8 @@ def test_handle_COMMAND_COMPLETE_closed_ps(con, mocker) -> None:
  "name": "max_prepared_statements_limit_2",
  "max_prepared_statements": 2,
  "queries": ["SELECT 1", "SELECT 2"],
- "expected_close_calls": 2,
- "expected_cache_size": 1
+ "expected_close_calls": 1,
+ "expected_cache_size": 2
  }
 ])
 def test_max_prepared_statement(con, mocker, test_case) -> None:

test/unit/test_core.py

@@ -178,16 +178,16 @@ def test_handle_command_complete_no_cache_cleanup(command_status, connection):
  "name": "max_prepared_statements_limit_1",
  "max_prepared_statements": 2,
  "queries": ["SELECT 1", "SELECT 2", "SELECT 3"],
- "expected_close_calls": 2
+ "expected_close_calls": 1
  },
-{
+ {
  "name": "max_prepared_statements_limit_2",
  "max_prepared_statements": 2,
  "queries": ["SELECT 1", "SELECT 2"],
  "expected_close_calls": 0
  }
 ])
-def test_max_prepared_statement_zero(mocker, test_case):
+def test_prepared_statement_cache_behavior(mocker, test_case):
  """
  Test prepared statement cache management in execute() with different configurations.
  :type mocker: object
@@ -217,6 +217,53 @@ def test_max_prepared_statement_zero(mocker, test_case):
  assert mock_connection.close_prepared_statement.call_count == test_case["expected_close_calls"]
 
 
+@pytest.mark.parametrize("test_case", [
+ {
+ "name": "statement_reuse",
+ "queries": ["SELECT 1", "SELECT 2", "SELECT 1", "SELECT 3"],
+ "expected_in_cache": ["SELECT 1", "SELECT 3"],
+ "expected_not_in_cache": ["SELECT 2"],
+ "expected_close_calls": 1
+ },
+ {
+ "name": "lru_order",
+ "queries": ["SELECT 1", "SELECT 2", "SELECT 1", "SELECT 2", "SELECT 3"],
+ "expected_in_cache": ["SELECT 2", "SELECT 3"],
+ "expected_not_in_cache": ["SELECT 1"],
+ "expected_close_calls": 1
+ }
+])
+def test_prepared_statement_lru_behavior(mocker, test_case):
+ """Test LRU behavior of prepared statement cache."""
+ from os import getpid
+
+ mock_connection = Connection.__new__(Connection)
+ mock_connection.max_prepared_statements = 2 # Set to 2 for LRU testing
+ mock_connection.merge_socket_read = True
+ mock_connection._caches = {}
+ mock_connection._send_message = mocker.Mock()
+ mock_connection._write = mocker.Mock()
+ mock_connection._flush = mocker.Mock()
+ mock_connection.handle_messages = mocker.Mock()
+ mock_connection.handle_messages_merge_socket_read = mocker.Mock()
+ mock_connection.close_prepared_statement = mocker.Mock()
+
+ mock_cursor = mocker.Mock()
+ mock_cursor.paramstyle = "named"
+
+ for query in test_case["queries"]:
+ mock_connection.execute(mock_cursor, query, None)
+
+ assert mock_connection.close_prepared_statement.call_count == test_case["expected_close_calls"]
+
+ # Verify cache contents
+ pid = getpid()
+ cache = mock_connection._caches["named"][pid]
+ for query in test_case["expected_in_cache"]:
+ assert any(key[0] == query for key in cache["statement_dict"])
+ for query in test_case["expected_not_in_cache"]:
+ assert not any(key[0] == query for key in cache["statement_dict"])
+
 @pytest.mark.parametrize("test_case", [
  {
  "max_prepared_statements" : 0,