Set session conneciton to lazily instantiate with async function calls

Author: Michael Christensen

src/amazon-keyspaces-mcp-server/awslabs/amazon_keyspaces_mcp_server/client.py

@@ -25,7 +25,6 @@
 
 from cassandra.auth import PlainTextAuthProvider
 from cassandra.cluster import Cluster, Session
-from cassandra.io.asyncioreactor import AsyncioConnection
 
 from .consts import (
     CERT_DIRECTORY,
@@ -70,19 +69,23 @@ def __init__(self, database_config: DatabaseConfig):
         """Initialize the client with the given configuration."""
         self.database_config = database_config
         self.is_keyspaces = database_config.use_keyspaces
-
-        # Initialize session for the configured database type (Keyspaces or Cassandra)
-        try:
-            if self.is_keyspaces:
-                self.session = self._create_keyspaces_session()
-                logger.info('Connected to Amazon Keyspaces')
-            else:
-                self.session = self._create_cassandra_session()
-                logger.info('Connected to Cassandra cluster')
-        except Exception as e:
-            target = 'Amazon Keyspaces' if self.is_keyspaces else 'Cassandra cluster'
-            logger.error('Failed to connect to %s: %s', target, str(e))
-            raise RuntimeError(f'Failed to connect to {target}: {str(e)}') from e
+        self._session: Optional[Session] = None
+
+    async def get_session(self) -> Session:
+        """Get or create database session lazily."""
+        if self._session is None:
+            try:
+                if self.is_keyspaces:
+                    self._session = self._create_keyspaces_session()
+                    logger.info('Connected to Amazon Keyspaces')
+                else:
+                    self._session = self._create_cassandra_session()
+                    logger.info('Connected to Cassandra cluster')
+            except Exception as e:
+                target = 'Amazon Keyspaces' if self.is_keyspaces else 'Cassandra cluster'
+                logger.error('Failed to connect to %s: %s', target, str(e))
+                raise RuntimeError(f'Failed to connect to {target}: {str(e)}') from e
+        return self._session
 
     def _create_cassandra_session(self) -> Session:
         """Create a session for Apache Cassandra."""
@@ -102,8 +105,6 @@ def _create_cassandra_session(self) -> Session:
             connect_timeout=int(CONNECTION_TIMEOUT),
         )
 
-        cluster.connection_class = AsyncioConnection
-
         return cluster.connect()
 
     def _create_keyspaces_session(self) -> Session:
@@ -144,8 +145,6 @@ def _create_keyspaces_session(self) -> Session:
                 connect_timeout=int(CONNECTION_TIMEOUT),
             )
 
-        cluster.connection_class = AsyncioConnection
-
         return cluster.connect()
 
     def _create_ssl_context_for_keyspaces(self) -> ssl.SSLContext:
@@ -171,13 +170,14 @@ def is_using_keyspaces(self) -> bool:
         """Check if the client is using Amazon Keyspaces."""
         return self.is_keyspaces
 
-    def list_keyspaces(self) -> List[KeyspaceInfo]:
+    async def list_keyspaces(self) -> List[KeyspaceInfo]:
         """List all keyspaces in the database."""
         keyspaces = []
 
         try:
             query = 'SELECT keyspace_name, replication FROM system_schema.keyspaces'
-            rows = self.session.execute(query)
+            session = await self.get_session()
+            rows = session.execute(query)
 
             for row in rows:
                 name = row.keyspace_name
@@ -199,13 +199,14 @@ def list_keyspaces(self) -> List[KeyspaceInfo]:
             logger.error('Error listing keyspaces: %s', str(e))
             raise RuntimeError(f'Failed to list keyspaces: {str(e)}') from e
 
-    def list_tables(self, keyspace_name: str) -> List[TableInfo]:
+    async def list_tables(self, keyspace_name: str) -> List[TableInfo]:
         """List all tables in a keyspace."""
         tables = []
 
         try:
             query = 'SELECT table_name FROM system_schema.tables WHERE keyspace_name = %s'
-            rows = self.session.execute(query, [keyspace_name])
+            session = await self.get_session()
+            rows = session.execute(query, [keyspace_name])
 
             for row in rows:
                 name = row.table_name
@@ -218,11 +219,13 @@ def list_tables(self, keyspace_name: str) -> List[TableInfo]:
                 f'Failed to list tables for keyspace {keyspace_name}: {str(e)}'
             ) from e
 
-    def describe_keyspace(self, keyspace_name: str) -> Dict[str, Any]:
+    async def describe_keyspace(self, keyspace_name: str) -> Dict[str, Any]:
         """Get detailed information about a keyspace."""
         try:
             query = 'SELECT * FROM system_schema.keyspaces WHERE keyspace_name = %s'
-            row = self.session.execute(query, [keyspace_name]).one()
+            session = await self.get_session()
+
+            row = session.execute(query, [keyspace_name]).one()
 
             if not row:
                 raise RuntimeError(f'Keyspace not found: {keyspace_name}')
@@ -245,14 +248,16 @@ def describe_keyspace(self, keyspace_name: str) -> Dict[str, Any]:
             logger.error('Error describing keyspace %s: %s', keyspace_name, str(e))
             raise RuntimeError(f'Failed to describe keyspace {keyspace_name}: {str(e)}') from e
 
-    def describe_table(self, keyspace_name: str, table_name: str) -> Dict[str, Any]:
+    async def describe_table(self, keyspace_name: str, table_name: str) -> Dict[str, Any]:
         """Get detailed information about a table."""
         try:
             query = (
                 'SELECT * FROM system_schema.tables WHERE '
                 'keyspace_name = %s AND table_name = %s'
             )
-            table_row = self.session.execute(query, [keyspace_name, table_name]).one()
+            session = await self.get_session()
+
+            table_row = session.execute(query, [keyspace_name, table_name]).one()
 
             if not table_row:
                 raise RuntimeError(f'Table not found: {keyspace_name}.{table_name}')
@@ -266,7 +271,9 @@ def describe_table(self, keyspace_name: str, table_name: str) -> Dict[str, Any]:
             query = (
                 'SELECT * FROM system_schema.columns WHERE keyspace_name = %s AND table_name = %s'
             )
-            column_rows = self.session.execute(query, [keyspace_name, table_name])
+            session = await self.get_session()
+
+            column_rows = session.execute(query, [keyspace_name, table_name])
 
             columns = []
             for column_row in column_rows:
@@ -284,7 +291,9 @@ def describe_table(self, keyspace_name: str, table_name: str) -> Dict[str, Any]:
             query = (
                 'SELECT * FROM system_schema.indexes WHERE keyspace_name = %s AND table_name = %s'
             )
-            index_rows = self.session.execute(query, [keyspace_name, table_name])
+            session = await self.get_session()
+
+            index_rows = session.execute(query, [keyspace_name, table_name])
 
             indexes = []
             for index_row in index_rows:
@@ -307,7 +316,9 @@ def describe_table(self, keyspace_name: str, table_name: str) -> Dict[str, Any]:
                         'SELECT custom_properties FROM system_schema_mcs.tables '
                         'WHERE keyspace_name = %s AND table_name = %s'
                     )
-                    capacity_row = self.session.execute(
+                    session = await self.get_session()
+
+                    capacity_row = session.execute(
                         query, [keyspace_name, table_name]
                     ).one()
 
@@ -339,7 +350,7 @@ def describe_table(self, keyspace_name: str, table_name: str) -> Dict[str, Any]:
                 f'Failed to describe table {keyspace_name}.{table_name}: {str(e)}'
             ) from e
 
-    def execute_read_only_query(
+    async def execute_read_only_query(
         self, query: str, params: Optional[List[Any]] = None
     ) -> Dict[str, Any]:
         """Execute a read-only SELECT query against the database."""
@@ -360,9 +371,13 @@ def execute_read_only_query(
 
             # Execute the query
             if params:
-                rs = self.session.execute(query, params)
+                session = await self.get_session()
+
+                rs = session.execute(query, params)
             else:
-                rs = self.session.execute(query)
+                session = await self.get_session()
+
+                rs = session.execute(query)
 
             # Process the results
             rows = []

src/amazon-keyspaces-mcp-server/awslabs/amazon_keyspaces_mcp_server/server.py

@@ -18,7 +18,6 @@
 
 from fastmcp import Context, FastMCP
 from loguru import logger
-from pydantic import Field
 
 from .client import UnifiedCassandraClient
 from .config import AppConfig
@@ -42,13 +41,15 @@
     build_query_result_context,
     build_table_details_context,
 )
+from .models import KeyspaceInput, QueryInput, TableInput
 from .services import DataService, QueryAnalysisService, SchemaService
 
 
 # Remove all default handlers then add our own
 logger.remove()
 logger.add(sys.stderr, level='INFO')
 
+
 mcp = FastMCP(
     name=SERVER_NAME,
     version=SERVER_VERSION,
@@ -79,10 +80,10 @@
 _PROXY = None
 
 
-def get_proxy():
+async def get_proxy():
     """Returns a singleton instance of the main Keyspaces MCP server implementation.
 
-    The singleton is initialized lazily.
+    The singleton is initialized lazily when first accessed (ensuring event loop is running).
     """
     global _PROXY  # pylint: disable=global-statement
     if _PROXY is None:
@@ -110,72 +111,74 @@ async def list_keyspaces(
     ctx: Optional[Context] = None,
 ) -> str:
     """Lists all keyspaces in the Cassandra/Keyspaces database."""
-    return await get_proxy()._handle_list_keyspaces(ctx)  # pylint: disable=protected-access
+    proxy = await get_proxy()
+    return await proxy._handle_list_keyspaces(ctx)  # pylint: disable=protected-access
 
 
 @mcp.tool(
     name='listTables',
     description='Lists all tables in a specified keyspace - args: keyspace',
 )
 async def list_tables(
-    keyspace: str = Field(..., description='The keyspace to list tables from.'),
+    input: KeyspaceInput,
     ctx: Optional[Context] = None,
 ) -> str:
     """Lists all tables in a specified keyspace."""
-    return await get_proxy()._handle_list_tables(keyspace, ctx)  # pylint: disable=protected-access
+    proxy = await get_proxy()
+    return await proxy._handle_list_tables(input.keyspace, ctx)  # pylint: disable=protected-access
 
 
 @mcp.tool(
     name='describeKeyspace',
     description='Gets detailed information about a keyspace - args: keyspace',
 )
 async def describe_keyspace(
-    keyspace: str = Field(..., description='The keyspace to retrieve metadata for.'),
+    input: KeyspaceInput,
     ctx: Optional[Context] = None,
 ) -> str:
     """Gets detailed information about a keyspace."""
-    return await get_proxy()._handle_describe_keyspace(keyspace, ctx)  # pylint: disable=protected-access
+    proxy = await get_proxy()
+    return await proxy._handle_describe_keyspace(input.keyspace, ctx)  # pylint: disable=protected-access
 
 
 @mcp.tool(
     name='describeTable',
     description='Gets detailed information about a table - args: keyspace, table',
 )
 async def describe_table(
-    keyspace: str = Field(..., description='The keyspace containing the table'),
-    table: str = Field(..., description='The name of the table to describe'),
+    input: TableInput,
     ctx: Optional[Context] = None,
 ) -> str:
     """Gets detailed information about a table."""
-    return await get_proxy()._handle_describe_table(keyspace, table, ctx)  # pylint: disable=protected-access
+    proxy = await get_proxy()
+    return await proxy._handle_describe_table(input.keyspace, input.table, ctx)  # pylint: disable=protected-access
 
 
 @mcp.tool(
     name='executeQuery',
     description='Executes a read-only SELECT query against the database - args: keyspace, query',
 )
 async def execute_query(
-    keyspace: str = Field(..., description='The keyspace to execute the query against'),
-    query: str = Field(..., description='The CQL SELECT query to execute'),
+    input: QueryInput,
     ctx: Optional[Context] = None,
 ) -> str:
     """Executes a read-only (SELECT) query against the database."""
-    return await get_proxy()._handle_execute_query(keyspace, query, ctx)  # pylint: disable=protected-access
+    proxy = await get_proxy()
+    return await proxy._handle_execute_query(input.keyspace, input.query, ctx)  # pylint: disable=protected-access
 
 
 @mcp.tool(
     name='analyzeQueryPerformance',
     description='Analyzes the performance characteristics of a CQL query - args: keyspace, query',
 )
 async def analyze_query_performance(
-    keyspace: str = Field(..., description='The keyspace to analyze the query against'),
-    query: str = Field(..., description='The CQL query to analyze for performance'),
+    input: QueryInput,
     ctx: Optional[Context] = None,
 ) -> str:
     """Analyzes the performance characteristics of a CQL query."""
-    proxy = get_proxy()
+    proxy = await get_proxy()
     return await proxy._handle_analyze_query_performance(  # pylint: disable=protected-access
-        keyspace, query, ctx
+        input.keyspace, input.query, ctx
     )
 
 
@@ -196,7 +199,7 @@ def __init__(
     async def _handle_list_keyspaces(self, ctx: Optional[Any] = None) -> str:
         """Handle the listKeyspaces tool."""
         try:
-            keyspaces = self.schema_service.list_keyspaces()
+            keyspaces = await self.schema_service.list_keyspaces()
 
             # Format keyspace names as a markdown list for better display
             keyspace_names = [k.name for k in keyspaces]
@@ -225,7 +228,7 @@ async def _handle_list_tables(self, keyspace: str, ctx: Optional[Context] = None
             if not keyspace:
                 raise ValidationError('Keyspace name is required')
 
-            tables = self.schema_service.list_tables(keyspace)
+            tables = await self.schema_service.list_tables(keyspace)
 
             # Format table names as a markdown list for better display
             table_names = [t.name for t in tables]
@@ -256,7 +259,7 @@ async def _handle_describe_keyspace(
             if not keyspace:
                 raise ValidationError('Keyspace name is required')
 
-            keyspace_details = self.schema_service.describe_keyspace(keyspace)
+            keyspace_details = await self.schema_service.describe_keyspace(keyspace)
 
             # Format keyspace details as markdown
             formatted_text = f'## Keyspace: `{keyspace}`\n\n'
@@ -303,7 +306,7 @@ async def _handle_describe_table(
             if not table:
                 raise ValidationError('Table name is required')
 
-            table_details = self.schema_service.describe_table(keyspace, table)
+            table_details = await self.schema_service.describe_table(keyspace, table)
 
             # Format table details as markdown
             formatted_text = f'## Table: `{keyspace}.{table}`\n\n'
@@ -382,7 +385,7 @@ async def _handle_execute_query(
                 raise QuerySecurityError('Query contains potentially unsafe operations')
 
             # Execute the query using the DataService
-            query_results = self.data_service.execute_read_only_query(keyspace, query)
+            query_results = await self.data_service.execute_read_only_query(keyspace, query)
 
             # Format the results for display
             formatted_text = '## Query Results\n\n'
@@ -447,7 +450,7 @@ async def _handle_analyze_query_performance(
             if not query:
                 raise ValidationError('Query is required')
 
-            analysis_result = self.query_analysis_service.analyze_query(keyspace, query)
+            analysis_result = await self.query_analysis_service.analyze_query(keyspace, query)
 
             # Build a user-friendly response
             formatted_text = '## Query Analysis Results\n\n'
@@ -478,6 +481,17 @@ async def _handle_analyze_query_performance(
 
 def main():
     """Run the MCP server."""
+    import asyncio
+    
+    # Validate connection before starting server
+    try:
+        proxy = asyncio.run(get_proxy())
+        asyncio.run(proxy.schema_service.cassandra_client.get_session())
+        logger.success('Successfully validated database connection')
+    except Exception as e:
+        logger.error(f'Failed to connect to database: {e}')
+        sys.exit(1)
+    
     mcp.run()