[Fix] Close OpenSearch client gracefully (#126)

* Close OpenSearch client gracefully

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

.gitignore

@@ -8,6 +8,9 @@
 .vscode/
 pyrightconfig.json
 
+# Kiro IDE
+.kiro/
+
 # Coverage output
 .coverage
 .coverage.*

CHANGELOG.md

@@ -10,6 +10,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ### Fixed
 - Fix Concurrency: Use Async OpenSearch client to improve concurrency ([#125](https://github.com/opensearch-project/opensearch-mcp-server-py/pull/125))
+- [Fix] Close OpenSearch client gracefully ([#126](https://github.com/opensearch-project/opensearch-mcp-server-py/pull/126))
 
 ### Removed
 

src/mcp_server_opensearch/clusters_information.py

@@ -116,17 +116,9 @@ async def load_clusters_from_yaml(file_path: str) -> None:
                     ssl_verify=cluster_config.get('ssl_verify', None),
                     opensearch_header_auth=cluster_config.get('opensearch_header_auth', None),
                 )
-                # Check if possible to connect to the cluster
-                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}"
-                    )
-                    continue
-                else:
-                    # Add cluster to registry
-                    add_cluster(name=cluster_name, cluster_info=cluster_info)
 
+                # Add cluster to registry without checking connection
+                add_cluster(name=cluster_name, cluster_info=cluster_info)
                 result['loaded_clusters'].append(cluster_name)
 
             except Exception as e:
@@ -141,23 +133,3 @@ async def load_clusters_from_yaml(file_path: str) -> None:
 
     except yaml.YAMLError as e:
         raise yaml.YAMLError(f'Invalid YAML format in {file_path}: {str(e)}')
-
-
-async def check_cluster_connection(cluster_info: ClusterInfo) -> tuple[bool, str]:
-    """Check if the cluster is reachable by attempting to connect.
-
-    Args:
-        cluster_info: ClusterInfo object containing cluster configuration
-
-    Returns:
-        tuple[bool, str]: (True, "") if connection successful, (False, error_message) otherwise
-    """
-    try:
-        # Lazy import to avoid circular dependency
-        from opensearch.client import _initialize_client_multi_mode
-
-        client = _initialize_client_multi_mode(cluster_info)
-        await client.ping()
-        return True, ''
-    except Exception as e:
-        return False, str(e)

src/mcp_server_opensearch/streaming_server.py

@@ -166,6 +166,7 @@ async def serve(
         app=app,
         host=host,
         port=port,
+        timeout_graceful_shutdown=10,
     )
     server = uvicorn.Server(config)
     await server.serve()

src/opensearch/client.py

@@ -11,7 +11,8 @@
 import boto3
 import logging
 import os
-from typing import Any, Dict, Optional
+from contextlib import asynccontextmanager
+from typing import Any, AsyncIterator, Dict, Optional
 from urllib.parse import urlparse
 
 from mcp.server.lowlevel.server import request_ctx
@@ -99,6 +100,43 @@ def initialize_client(args: baseToolArgs) -> AsyncOpenSearch:
         raise ConfigurationError(f'Failed to initialize OpenSearch client: {e}')
 
 
+@asynccontextmanager
+async def get_opensearch_client(args: baseToolArgs) -> AsyncIterator[AsyncOpenSearch]:
+    """Async context manager for OpenSearch client lifecycle management.
+
+    This context manager ensures that OpenSearch clients are properly closed after use,
+    preventing connection leaks and enabling graceful server shutdown.
+
+    Usage:
+        async with get_opensearch_client(args) as client:
+            # Use client for operations
+            result = await client.info()
+
+    Args:
+        args (baseToolArgs): Arguments containing optional opensearch_cluster_name
+
+    Yields:
+        AsyncOpenSearch: An initialized OpenSearch client instance
+
+    Raises:
+        ConfigurationError: If in multi mode but no cluster name provided or invalid mode
+        AuthenticationError: If authentication fails
+    """
+    client = None
+    try:
+        logger.debug('Creating OpenSearch client')
+        client = initialize_client(args)
+        yield client
+    finally:
+        if client is not None:
+            try:
+                logger.debug('Closing OpenSearch client')
+                await client.close()
+            except Exception as e:
+                # Log but don't propagate cleanup errors to avoid masking original errors
+                logger.warning(f'Error closing OpenSearch client: {e}')
+
+
 # Private Implementation Functions
 def _initialize_client_single_mode() -> AsyncOpenSearch:
     """Initialize OpenSearch client for single mode using environment variables.