feat(redshift-mcp-server): long poll the Data API for statement status (#4560)

* feat(redshift-mcp-server): long poll the Data API for statement status

Every statement waited for completion by calling DescribeStatement once per
second, and each execute_query runs three statements (BEGIN READ ONLY, user
SQL, ROLLBACK), so a warm read-only query cost at least 8 Data API calls: 3
ExecuteStatement, 4 DescribeStatement and 1 GetStatementResult.

The Data API now supports long polling. WaitTimeSeconds (1-30) holds the
request open until the statement reaches a terminal status or the wait
elapses, and is accepted on both ExecuteStatement and DescribeStatement.

- Send WaitTimeSeconds, from the new QUERY_LONG_POLL constant, on
  ExecuteStatement and on every DescribeStatement poll. Setting it to 0
  restores plain QUERY_POLL_INTERVAL polling.
- Collapse the completion wait into a single loop that evaluates whichever
  response came last, since both operations report Status identically.
- Return the terminal statement response from _execute_statement rather than
  just the id, so _create_session_with_app_name reads SessionId and
  _execute_protected_statement reads HasResultSet off it. Both callers drop
  their own DescribeStatement call.
- Track the timeout with time.monotonic(), since accumulated sleeps stop
  approximating elapsed time once a call blocks server-side.
- Fall back to plain polling on ActiveWaitingRequestsExceededException, which
  the service raises when too many requests wait on one statement.
- Require boto3/botocore >= 1.43.55, the first release carrying
  WaitTimeSeconds in the redshift-data service model.
- Cover long polling, the terminal-status shortcut, the disabled path and the
  waiter-limit fallback in TestExecuteStatement.

A warm read-only query now costs 4 calls, and a statement with no result set
costs 1. Only a failing statement still needs a DescribeStatement, to read the
Error field that ExecuteStatement does not carry.

Verified end-to-end against live provisioned and serverless clusters.

* fix(redshift-mcp-server): run Data API calls off the event loop

boto3 is synchronous, so long polling turned each Data API call into an event
loop stall of up to WaitTimeSeconds. Measured against a live cluster, a 79s
query blocked the loop for 48s in one stretch and let a 50ms heartbeat tick
once rather than the ~1579 times expected, which serialized every other
in-flight tool call and stalled MCP protocol traffic. Before long polling the
same calls blocked only for a short HTTP round trip.

- Run execute_statement and both describe_statement calls in _execute_statement
  through asyncio.to_thread.
- Add a regression test that fails when a blocking Data API call starves a
  concurrent coroutine.

The same query now yields 1552 heartbeat ticks with a 0.05s worst gap.

Reported by Copilot on #4560.

---------

Co-authored-by: Sergey Konoplev <sergkono@amazon.com>
Co-authored-by: Laith Al-Saadoon <9553966+theagenticguy@users.noreply.github.com>

Author: 3 people

src/redshift-mcp-server/awslabs/redshift_mcp_server/consts.py

@@ -22,6 +22,7 @@
 DEFAULT_LOG_LEVEL = 'WARNING'
 QUERY_TIMEOUT = 3600
 QUERY_POLL_INTERVAL = 1
+QUERY_LONG_POLL = 30
 SESSION_KEEPALIVE = 600
 
 # SQL discovery commands. Results are read positionally; {placeholders} are

src/redshift-mcp-server/awslabs/redshift_mcp_server/redshift.py

@@ -26,6 +26,7 @@
     CLIENT_USER_AGENT_NAME,
     COLUMNS_SQL,
     DATABASES_SQL,
+    QUERY_LONG_POLL,
     QUERY_POLL_INTERVAL,
     QUERY_TIMEOUT,
     SCHEMAS_SQL,
@@ -206,18 +207,15 @@ async def _create_session_with_app_name(
         app_name_sql = f"SET application_name TO '{self._app_name}';"
 
         # Execute statement to create session
-        statement_id = await _execute_statement(
+        response = await _execute_statement(
             cluster_info=cluster_info,
             cluster_identifier=cluster_identifier,
             database_name=database_name,
             sql=app_name_sql,
             session_keepalive=self._session_keepalive,
         )
 
-        # Get session ID from the response
-        data_client = client_manager.redshift_data_client()
-        status_response = data_client.describe_statement(Id=statement_id)
-        session_id = status_response['SessionId']
+        session_id = response['SessionId']
 
         logger.debug(f'Created session with application name: {session_id}')
         return session_id
@@ -296,7 +294,7 @@ async def _execute_protected_statement(
         if allow_read_write:
             # Read-write: run the single guarded statement directly (autocommit). No
             # transaction wrapper. Any error propagates.
-            user_query_id = await _execute_statement(
+            user_statement = await _execute_statement(
                 cluster_info=cluster_info,
                 cluster_identifier=cluster_identifier,
                 database_name=database_name,
@@ -316,11 +314,11 @@ async def _execute_protected_statement(
             )
 
             # Execute user SQL with parameters, ensuring the transaction is always closed.
-            user_query_id = None
+            user_statement = None
             user_sql_error: Exception | None = None
 
             try:
-                user_query_id = await _execute_statement(
+                user_statement = await _execute_statement(
                     cluster_info=cluster_info,
                     cluster_identifier=cluster_identifier,
                     database_name=database_name,
@@ -357,14 +355,14 @@ async def _execute_protected_statement(
                 raise user_sql_error
 
     # Get results from user query (shared by both modes); runs outside the lock.
-    # describe_statement / get_statement_result are keyed by query_id, not session-bound,
-    # so the lock is not held during the (potentially unbounded) results wait.
-    data_client = client_manager.redshift_data_client()
-    assert user_query_id is not None, 'user_query_id should not be None at this point'
+    # get_statement_result is keyed by query_id, not session-bound, so the lock is not
+    # held during the (potentially unbounded) results wait.
+    assert user_statement is not None, 'user_statement should not be None at this point'
+    user_query_id = user_statement['Id']
 
     # Only fetch results when the statement produced a result set (e.g. SET does not).
-    describe_response = data_client.describe_statement(Id=user_query_id)
-    if describe_response.get('HasResultSet'):
+    if user_statement.get('HasResultSet'):
+        data_client = client_manager.redshift_data_client()
         results_response = data_client.get_statement_result(Id=user_query_id)
     else:
         results_response = {'Records': [], 'ColumnMetadata': []}
@@ -381,7 +379,8 @@ async def _execute_statement(
     session_keepalive: int | None = None,
     query_poll_interval: float = QUERY_POLL_INTERVAL,
     query_timeout: float = QUERY_TIMEOUT,
-) -> str:
+    query_long_poll: int = QUERY_LONG_POLL,
+) -> dict:
     """Execute a single statement with optional session support and parameters.
 
     Args:
@@ -394,9 +393,13 @@ async def _execute_statement(
         session_keepalive: Optional session keepalive seconds (only used when session_id is None).
         query_poll_interval: Polling interval in seconds for checking query status.
         query_timeout: Maximum time in seconds to wait for query completion.
+        query_long_poll: Data API WaitTimeSeconds, 1-30, or 0 to disable long polling.
 
     Returns:
-        Statement ID from the ExecuteStatement response.
+        The terminal statement response, carrying Id, Status, SessionId and HasResultSet.
+
+    Raises:
+        Exception: If the statement fails, is aborted, or times out.
     """
     data_client = client_manager.redshift_data_client()
 
@@ -423,35 +426,54 @@ async def _execute_statement(
     elif session_keepalive is not None:
         request_params['SessionKeepAliveSeconds'] = session_keepalive
 
-    response = data_client.execute_statement(**request_params)
+    long_poll_params = {'WaitTimeSeconds': query_long_poll} if query_long_poll else {}
+
+    # boto3 is synchronous and a long poll holds the caller for up to query_long_poll
+    # seconds, so every Data API call here runs off the event loop.
+    response = await asyncio.to_thread(
+        data_client.execute_statement, **request_params, **long_poll_params
+    )
     statement_id = response['Id']
 
     logger.debug(
         f'Executed statement: {statement_id}' + (f' in session {session_id}' if session_id else '')
     )
 
-    # Wait for statement completion
-    wait_time = 0
-    while wait_time < query_timeout:
-        status_response = data_client.describe_statement(Id=statement_id)
-        status = status_response['Status']
+    # ExecuteStatement and DescribeStatement report status alike, so one loop settles the
+    # long-polled submit and every later poll. Wall clock, since a long poll blocks server-side.
+    deadline = time.monotonic() + query_timeout
+    while True:
+        status = response.get('Status')
 
         if status == 'FINISHED':
             logger.debug(f'Statement completed: {statement_id}')
-            break
+            return response
         elif status in ['FAILED', 'ABORTED']:
-            error_msg = status_response.get('Error', 'Unknown error')
+            error_msg = response.get('Error')
+            if error_msg is None:
+                # ExecuteStatement carries no Error field, so the reason takes a describe.
+                described = await asyncio.to_thread(
+                    data_client.describe_statement, Id=statement_id
+                )
+                error_msg = described.get('Error', 'Unknown error')
             logger.error(f'Statement failed: {error_msg}')
             raise Exception(f'Statement failed: {error_msg}')
 
-        await asyncio.sleep(query_poll_interval)
-        wait_time += query_poll_interval
+        if time.monotonic() >= deadline:
+            logger.error(f'Statement timed out: {statement_id}')
+            raise Exception(f'Statement timed out after {query_timeout} seconds')
 
-    if wait_time >= query_timeout:
-        logger.error(f'Statement timed out: {statement_id}')
-        raise Exception(f'Statement timed out after {wait_time} seconds')
+        await asyncio.sleep(query_poll_interval)
 
-    return statement_id
+        try:
+            response = await asyncio.to_thread(
+                data_client.describe_statement, Id=statement_id, **long_poll_params
+            )
+        except ClientError as e:
+            if e.response.get('Error', {}).get('Code') != 'ActiveWaitingRequestsExceededException':
+                raise
+            logger.warning(f'Long polling limit reached, polling instead: {statement_id}')
+            long_poll_params = {}
 
 
 async def discover_clusters() -> list[RedshiftCluster]:

src/redshift-mcp-server/pyproject.toml

@@ -8,8 +8,9 @@ description = "An AWS Labs Model Context Protocol (MCP) server for Redshift"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
-    "boto3>=1.38.39",
-    "botocore>=1.38.39",
+    # Data API long polling (WaitTimeSeconds) entered the service model in 1.43.55.
+    "boto3>=1.43.55",
+    "botocore>=1.43.55",
     "loguru>=0.7.0",
     "mcp[cli]>=2.0.0,<3.0.0",
     "pydantic>=2.10.6",