Fix Concurrency: Use Async OpenSearch client to improve concurrency (#125)

* Fix Concurrency: Use Async OpenSearch client to improve concurrency

Signed-off-by: rithin-pullela-aws <[email protected]>

* Add changelog

Signed-off-by: rithin-pullela-aws <[email protected]>

---------

Signed-off-by: rithin-pullela-aws <[email protected]>

Author: rithin-pullela-aws

CHANGELOG.md

@@ -9,6 +9,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Add header-based authentication + Code Clean up ([#117](https://github.com/opensearch-project/opensearch-mcp-server-py/pull/117))
 
 ### Fixed
+- Fix Concurrency: Use Async OpenSearch client to improve concurrency ([#125](https://github.com/opensearch-project/opensearch-mcp-server-py/pull/125))
 
 ### Removed
 

src/mcp_server_opensearch/clusters_information.py

@@ -52,7 +52,7 @@ def get_cluster(name: str) -> Optional[ClusterInfo]:
     return cluster_registry.get(name)
 
 
-def load_clusters_from_yaml(file_path: str) -> None:
+async def load_clusters_from_yaml(file_path: str) -> None:
     """Load cluster configurations from a YAML file and populate the global registry.
 
     Args:
@@ -117,7 +117,7 @@ def load_clusters_from_yaml(file_path: str) -> None:
                     opensearch_header_auth=cluster_config.get('opensearch_header_auth', None),
                 )
                 # Check if possible to connect to the cluster
-                is_connected, error_message = check_cluster_connection(cluster_info)
+                is_connected, error_message = await check_cluster_connection(cluster_info)
                 if not is_connected:
                     result['errors'].append(
                         f"Error connecting to cluster '{cluster_name}': {error_message}"
@@ -143,7 +143,7 @@ def load_clusters_from_yaml(file_path: str) -> None:
         raise yaml.YAMLError(f'Invalid YAML format in {file_path}: {str(e)}')
 
 
-def check_cluster_connection(cluster_info: ClusterInfo) -> tuple[bool, str]:
+async def check_cluster_connection(cluster_info: ClusterInfo) -> tuple[bool, str]:
     """Check if the cluster is reachable by attempting to connect.
 
     Args:
@@ -157,7 +157,7 @@ def check_cluster_connection(cluster_info: ClusterInfo) -> tuple[bool, str]:
         from opensearch.client import _initialize_client_multi_mode
 
         client = _initialize_client_multi_mode(cluster_info)
-        client.ping()
+        await client.ping()
         return True, ''
     except Exception as e:
         return False, str(e)

src/mcp_server_opensearch/stdio_server.py

@@ -34,7 +34,7 @@ async def serve(
     server = Server('opensearch-mcp-server')
     # Load clusters from YAML file
     if mode == 'multi':
-        load_clusters_from_yaml(config_file_path)
+        await load_clusters_from_yaml(config_file_path)
 
     # Call tool generator
     await generate_tools_from_openapi()
@@ -43,7 +43,9 @@ async def serve(
         TOOL_REGISTRY, config_file_path, cli_tool_overrides or {}
     )
     # Get enabled tools (tool filter)
-    enabled_tools = get_tools(tool_registry=customized_registry, config_file_path=config_file_path)
+    enabled_tools = await get_tools(
+        tool_registry=customized_registry, config_file_path=config_file_path
+    )
     logging.info(f'Enabled tools: {list(enabled_tools.keys())}')
 
     @server.list_tools()

src/mcp_server_opensearch/streaming_server.py

@@ -41,7 +41,7 @@ async def create_mcp_server(
 
     # Load clusters from YAML file
     if mode == 'multi':
-        load_clusters_from_yaml(config_file_path)
+        await load_clusters_from_yaml(config_file_path)
 
     server = Server('opensearch-mcp-server')
     # Call tool generator
@@ -51,7 +51,9 @@ async def create_mcp_server(
         TOOL_REGISTRY, config_file_path, cli_tool_overrides or {}
     )
     # Get enabled tools (tool filter)
-    enabled_tools = get_tools(tool_registry=customized_registry, config_file_path=config_file_path)
+    enabled_tools = await get_tools(
+        tool_registry=customized_registry, config_file_path=config_file_path
+    )
     logging.info(f'Enabled tools: {list(enabled_tools.keys())}')
 
     @server.list_tools()
@@ -141,6 +143,7 @@ def create_app(self) -> Starlette:
                 Route('/health', endpoint=self.handle_health, methods=['GET']),
                 Mount('/messages/', app=self.sse.handle_post_message),
                 Mount('/mcp', app=self.handle_streamable_http),
+                Mount('/mcp/', app=self.handle_streamable_http),
             ],
             lifespan=self.lifespan,
         )

src/opensearch/client.py

@@ -19,9 +19,9 @@
 
 from mcp_server_opensearch.clusters_information import ClusterInfo, get_cluster
 from mcp_server_opensearch.global_state import get_mode, get_profile
-from opensearchpy import OpenSearch, RequestsHttpConnection
-from requests_aws4auth import AWS4Auth
+from opensearchpy import AsyncOpenSearch, AsyncHttpConnection, AWSV4SignerAsyncAuth
 from tools.tool_params import baseToolArgs
+from botocore.credentials import Credentials
 
 # Configure logging
 logger = logging.getLogger(__name__)
@@ -53,7 +53,7 @@ class ConfigurationError(OpenSearchClientError):
 
 
 # Public API Functions
-def initialize_client(args: baseToolArgs) -> OpenSearch:
+def initialize_client(args: baseToolArgs) -> AsyncOpenSearch:
     """Initialize and return an OpenSearch client based on the current mode.
 
     Behavior depends on the global mode:
@@ -100,7 +100,7 @@ def initialize_client(args: baseToolArgs) -> OpenSearch:
 
 
 # Private Implementation Functions
-def _initialize_client_single_mode() -> OpenSearch:
+def _initialize_client_single_mode() -> AsyncOpenSearch:
     """Initialize OpenSearch client for single mode using environment variables.
 
     Single mode uses environment variables for connection, with optional header-based auth:
@@ -200,7 +200,7 @@ def _initialize_client_single_mode() -> OpenSearch:
         raise ConfigurationError(f'Failed to initialize single mode client: {e}')
 
 
-def _initialize_client_multi_mode(cluster_info: ClusterInfo) -> OpenSearch:
+def _initialize_client_multi_mode(cluster_info: ClusterInfo) -> AsyncOpenSearch:
     """Initialize OpenSearch client for multi mode using cluster configuration.
 
     Multi mode uses cluster configuration from the provided ClusterInfo object.
@@ -303,7 +303,7 @@ def _create_opensearch_client(
     aws_access_key_id: Optional[str] = None,
     aws_secret_access_key: Optional[str] = None,
     aws_session_token: Optional[str] = None,
-) -> OpenSearch:
+) -> AsyncOpenSearch:
     """Common function to create OpenSearch client with authentication.
 
     This function handles the common authentication logic used by both
@@ -362,7 +362,7 @@ def _create_opensearch_client(
         'hosts': [opensearch_url],
         'use_ssl': (parsed_url.scheme == 'https'),
         'verify_certs': ssl_verify,
-        'connection_class': RequestsHttpConnection,
+        'connection_class': AsyncHttpConnection,
         'timeout': timeout,
     }
 
@@ -379,7 +379,7 @@ def _create_opensearch_client(
         if opensearch_no_auth:
             logger.info('[NO AUTH] Attempting connection without authentication')
             try:
-                return OpenSearch(**client_kwargs)
+                return AsyncOpenSearch(**client_kwargs)
             except Exception as e:
                 logger.error(f'[NO AUTH] Failed to connect without authentication: {e}')
                 raise AuthenticationError(f'Failed to connect without authentication: {e}')
@@ -392,16 +392,16 @@ def _create_opensearch_client(
                     raise AuthenticationError(
                         'AWS region is required for header-based authentication'
                     )
-
-                aws_auth = AWS4Auth(
-                    aws_access_key_id,
-                    aws_secret_access_key,
-                    aws_region.strip(),
-                    service_name,
-                    session_token=aws_session_token,
+                credentials = Credentials(
+                    access_key=aws_access_key_id,
+                    secret_key=aws_secret_access_key,
+                    token=aws_session_token,
+                )
+                aws_auth = AWSV4SignerAsyncAuth(
+                    credentials=credentials, region=aws_region.strip(), service=service_name
                 )
                 client_kwargs['http_auth'] = aws_auth
-                return OpenSearch(**client_kwargs)
+                return AsyncOpenSearch(**client_kwargs)
             except Exception as e:
                 logger.error(f'[HEADER AUTH] Failed to authenticate with header credentials: {e}')
                 raise AuthenticationError(f'Failed to authenticate with header credentials: {e}')
@@ -419,15 +419,11 @@ def _create_opensearch_client(
                 )
                 credentials = assumed_role['Credentials']
 
-                aws_auth = AWS4Auth(
-                    credentials['AccessKeyId'],
-                    credentials['SecretAccessKey'],
-                    aws_region,
-                    service_name,
-                    session_token=credentials['SessionToken'],
+                aws_auth = AWSV4SignerAsyncAuth(
+                    credentials=credentials, region=aws_region.strip(), service=service_name
                 )
                 client_kwargs['http_auth'] = aws_auth
-                return OpenSearch(**client_kwargs)
+                return AsyncOpenSearch(**client_kwargs)
             except Exception as e:
                 logger.error(f'[IAM AUTH] Failed to assume IAM role {iam_arn}: {e}')
                 raise AuthenticationError(f'Failed to assume IAM role {iam_arn}: {e}')
@@ -437,7 +433,7 @@ def _create_opensearch_client(
             logger.info(f'[BASIC AUTH] Using basic authentication for user: {opensearch_username}')
             try:
                 client_kwargs['http_auth'] = (opensearch_username.strip(), opensearch_password)
-                return OpenSearch(**client_kwargs)
+                return AsyncOpenSearch(**client_kwargs)
             except Exception as e:
                 logger.error(f'[BASIC AUTH] Failed to connect with basic authentication: {e}')
                 raise AuthenticationError(f'Failed to connect with basic authentication: {e}')
@@ -454,13 +450,11 @@ def _create_opensearch_client(
             if not credentials:
                 raise AuthenticationError('No AWS credentials found in session')
 
-            aws_auth = AWS4Auth(
-                refreshable_credentials=credentials,
-                service=service_name,
-                region=aws_region,
+            aws_auth = AWSV4SignerAsyncAuth(
+                credentials=credentials, region=aws_region.strip(), service=service_name
             )
             client_kwargs['http_auth'] = aws_auth
-            return OpenSearch(**client_kwargs)
+            return AsyncOpenSearch(**client_kwargs)
         except Exception as e:
             logger.error(f'[AWS CREDS] Failed to authenticate with AWS credentials: {e}')
             raise AuthenticationError(f'Failed to authenticate with AWS credentials: {e}')