awslabs.postgres-mcp-server 1.0.8__py3-none-any.whl → 1.0.11__py3-none-any.whl

awslabs/postgres_mcp_server/server.py

@@ -16,25 +16,45 @@
 
 import argparse
 import asyncio
+import json
 import sys
-from awslabs.postgres_mcp_server.connection import DBConnectionSingleton
+import threading
+import traceback
+from awslabs.postgres_mcp_server.connection.abstract_db_connection import AbstractDBConnection
+from awslabs.postgres_mcp_server.connection.cp_api_connection import (
+    internal_create_serverless_cluster,
+    internal_get_cluster_properties,
+    internal_get_instance_properties,
+    setup_aurora_iam_policy_for_current_user,
+)
+from awslabs.postgres_mcp_server.connection.db_connection_map import (
+    ConnectionMethod,
+    DatabaseType,
+    DBConnectionMap,
+)
 from awslabs.postgres_mcp_server.connection.psycopg_pool_connection import PsycopgPoolConnection
+from awslabs.postgres_mcp_server.connection.rds_api_connection import RDSDataAPIConnection
 from awslabs.postgres_mcp_server.mutable_sql_detector import (
     check_sql_injection_risk,
     detect_mutating_keywords,
 )
-from botocore.exceptions import BotoCoreError, ClientError
+from botocore.exceptions import ClientError
+from datetime import datetime
 from loguru import logger
 from mcp.server.fastmcp import Context, FastMCP
 from pydantic import Field
-from typing import Annotated, Any, Dict, List, Optional
+from typing import Annotated, Any, Dict, List, Optional, Tuple
 
 
+db_connection_map = DBConnectionMap()
+async_job_status: Dict[str, dict] = {}
+async_job_status_lock = threading.Lock()
 client_error_code_key = 'run_query ClientError code'
 unexpected_error_key = 'run_query unexpected error'
 write_query_prohibited_key = 'Your MCP tool only allows readonly query. If you want to write, change the MCP configuration per README.md'
 query_comment_prohibited_key = 'The comment in query is prohibited because of injection risk'
 query_injection_risk_key = 'Your query contains risky injection patterns'
+readonly_query = True
 
 
 class DummyCtx:
@@ -91,7 +111,10 @@ mcp = FastMCP(
 async def run_query(
     sql: Annotated[str, Field(description='The SQL query to run')],
     ctx: Context,
-    db_connection=None,
+    connection_method: Annotated[ConnectionMethod, Field(description='connection method')],
+    cluster_identifier: Annotated[str, Field(description='Cluster identifier')],
+    db_endpoint: Annotated[str, Field(description='database endpoint')],
+    database: Annotated[str, Field(description='database name')],
     query_parameters: Annotated[
         Optional[List[Dict[str, Any]]], Field(description='Parameters for the SQL query')
     ] = None,
@@ -101,7 +124,10 @@ async def run_query(
     Args:
         sql: The sql statement to run
         ctx: MCP context for logging and state management
-        db_connection: DB connection object passed by unit test. It should be None if called by MCP server.
+        connection_method: connection method
+        cluster_identifier: Cluster identifier
+        db_endpoint: database endpoint
+        database: database name
         query_parameters: Parameters for the SQL query
 
     Returns:
@@ -110,25 +136,38 @@ async def run_query(
     global client_error_code_key
     global unexpected_error_key
     global write_query_prohibited_key
+    global db_connection_map
 
-    if db_connection is None:
-        try:
-            # Try to get the connection from the singleton
-            db_connection = DBConnectionSingleton.get().db_connection
-        except RuntimeError:
-            # If the singleton is not initialized, this might be a direct connection
-            logger.error('No database connection available')
-            await ctx.error('No database connection available')
-            return [{'error': 'No database connection available'}]
+    logger.info(
+        f'Entered run_query with '
+        f'method:{connection_method}, cluster_identifier:{cluster_identifier}, '
+        f'db_endpoint:{db_endpoint}, database:{database}, '
+        f'sql:{sql}'
+    )
 
-    if db_connection is None:
-        raise AssertionError('db_connection should never be None')
+    db_connection = db_connection_map.get(
+        method=connection_method,
+        cluster_identifier=cluster_identifier,
+        db_endpoint=db_endpoint,
+        database=database,
+    )
+    if not db_connection:
+        err = (
+            f'No database connection available for method:{connection_method}, '
+            f'cluster_identifier:{cluster_identifier}, db_endpoint:{db_endpoint}, database:{database}'
+        )
+        logger.error(err)
+        await ctx.error(err)
+        return [{'error': err}]
 
     if db_connection.readonly_query:
         matches = detect_mutating_keywords(sql)
         if (bool)(matches):
             logger.info(
-                f'query is rejected because current setting only allows readonly query. detected keywords: {matches}, SQL query: {sql}'
+                (
+                    f'query is rejected because current setting only allows readonly query.'
+                    f'detected keywords: {matches}, SQL query: {sql}'
+                )
             )
             await ctx.error(write_query_prohibited_key)
             return [{'error': write_query_prohibited_key}]
@@ -144,9 +183,15 @@ async def run_query(
         return [{'error': query_injection_risk_key}]
 
     try:
-        logger.info(f'run_query: readonly:{db_connection.readonly_query}, SQL:{sql}')
+        logger.info(
+            (
+                f'run_query: sql:{sql} method:{connection_method}, '
+                f'cluster_identifier:{cluster_identifier} database:{database} '
+                f'db_endpoint:{db_endpoint} '
+                f'readonly:{db_connection.readonly_query} query_parameters:{query_parameters}'
+            )
+        )
 
-        # Execute the query using the abstract connection interface
         response = await db_connection.execute_query(sql, query_parameters)
 
         logger.success(f'run_query successfully executed query:{sql}')
@@ -164,23 +209,34 @@ async def run_query(
         return [{'error': unexpected_error_key}]
 
 
-@mcp.tool(
-    name='get_table_schema',
-    description='Fetch table columns and comments from Postgres',
-)
+@mcp.tool(name='get_table_schema', description='Fetch table columns and comments from Postgres')
 async def get_table_schema(
-    table_name: Annotated[str, Field(description='name of the table')], ctx: Context
+    connection_method: Annotated[ConnectionMethod, Field(description='connection method')],
+    cluster_identifier: Annotated[str, Field(description='Cluster identifier')],
+    db_endpoint: Annotated[str, Field(description='database endpoint')],
+    database: Annotated[str, Field(description='database name')],
+    table_name: Annotated[str, Field(description='name of the table')],
+    ctx: Context,
 ) -> list[dict]:
     """Get a table's schema information given the table name.
 
     Args:
+        connection_method: connection method
+        cluster_identifier: Cluster identifier
+        db_endpoint: database endpoint
+        database: database name
         table_name: name of the table
         ctx: MCP context for logging and state management
 
     Returns:
         List of dictionary that contains query response rows
     """
-    logger.info(f'get_table_schema: {table_name}')
+    logger.info(
+        (
+            f'Entered get_table_schema: table_name:{table_name} connection_method:{connection_method}, '
+            f'cluster_identifier:{cluster_identifier}, db_endpoint:{db_endpoint}, database:{database}'
+        )
+    )
 
     sql = """
         SELECT
@@ -190,7 +246,7 @@ async def get_table_schema(
         FROM
             pg_attribute a
         WHERE
-            a.attrelid = to_regclass(:table_name)
+            a.attrelid = to_regclass(%(table_name)s)
             AND a.attnum > 0
             AND NOT a.attisdropped
         ORDER BY a.attnum
@@ -198,125 +254,511 @@ async def get_table_schema(
 
     params = [{'name': 'table_name', 'value': {'stringValue': table_name}}]
 
-    return await run_query(sql=sql, ctx=ctx, query_parameters=params)
+    return await run_query(
+        sql=sql,
+        ctx=ctx,
+        connection_method=connection_method,
+        cluster_identifier=cluster_identifier,
+        db_endpoint=db_endpoint,
+        database=database,
+        query_parameters=params,
+    )
+
 
+@mcp.tool(
+    name='connect_to_database',
+    description='Connect to a specific database and save the connection internally',
+)
+def connect_to_database(
+    region: Annotated[str, Field(description='region')],
+    database_type: Annotated[DatabaseType, Field(description='database type')],
+    connection_method: Annotated[ConnectionMethod, Field(description='connection method')],
+    cluster_identifier: Annotated[str, Field(description='cluster identifier')],
+    db_endpoint: Annotated[str, Field(description='database endpoint')],
+    port: Annotated[int, Field(description='Postgres port')],
+    database: Annotated[str, Field(description='database name')],
+) -> str:
+    """Connect to a specific database save the connection internally.
 
-def main():
-    """Main entry point for the MCP server application."""
-    global client_error_code_key
+    Args:
+        region: region of the database. Required parametere.
+        database_type: Either APG for Aurora Postgres or RPG for RDS Postgres cluster. Required parameter
+        connection_method: Either RDS_API, PG_WIRE_PROTOCOL, or PG_WIRE_IAM_PROTOCOL. Required parameter
+        cluster_identifier: Either Aurora Postgres cluster identifier or RDS Postgres cluster identifier
+        db_endpoint: database endpoint
+        port: database port
+        database: database name. Required parameter
+
+        Supported scenario:
+        1. Aurora Postgres database with RDS_API + Credential Manager:
+            cluster_identifier must be set
+            db_endpoint and port will be ignored
+        2. Aurora Postgres database with direct connection + IAM:
+            cluster_identifier must be set
+            db_endpoint must be set
+        3. Aurora Postgres database with direct connection + PG_AUTH (Credential Manager):
+            cluster_identifier must be set
+            db_endpoint must be set
+        4. RDS Postgres database with direct connection + PG_AUTH (Credential Manager):
+            credential manager setting is either on instance or cluster
+            db_endpoint must be set
+    """
+    try:
+        db_connection, llm_response = internal_connect_to_database(
+            region=region,
+            database_type=database_type,
+            connection_method=connection_method,
+            cluster_identifier=cluster_identifier,
+            db_endpoint=db_endpoint,
+            port=port,
+            database=database,
+        )
 
-    """Run the MCP server with CLI argument support."""
-    parser = argparse.ArgumentParser(
-        description='An AWS Labs Model Context Protocol (MCP) server for postgres'
+        return str(llm_response)
+
+    except Exception as e:
+        logger.error(f'connect_to_database failed with error: {str(e)}')
+        trace_msg = traceback.format_exc()
+        logger.error(f'Trace:{trace_msg}')
+        llm_response = {'status': 'Failed', 'error': str(e)}
+        return json.dumps(llm_response, indent=2)
+
+
+@mcp.tool(name='is_database_connected', description='Check if a connection has been established')
+def is_database_connected(
+    cluster_identifier: Annotated[str, Field(description='cluster identifier')],
+    db_endpoint: Annotated[str, Field(description='database endpoint')] = '',
+    database: Annotated[str, Field(description='database name')] = 'postgres',
+) -> bool:
+    """Check if a connection has been established.
+
+    Args:
+        cluster_identifier: cluster identifier
+        db_endpoint: database endpoint
+        database: database name
+
+    Returns:
+        result in boolean
+    """
+    global db_connection_map
+    if db_connection_map.get(ConnectionMethod.RDS_API, cluster_identifier, db_endpoint, database):
+        return True
+
+    if db_connection_map.get(
+        ConnectionMethod.PG_WIRE_PROTOCOL, cluster_identifier, db_endpoint, database
+    ):
+        return True
+
+    if db_connection_map.get(
+        ConnectionMethod.PG_WIRE_IAM_PROTOCOL, cluster_identifier, db_endpoint, database
+    ):
+        return True
+
+    return False
+
+
+@mcp.tool(
+    name='get_database_connection_info',
+    description='Get all cached database connection information',
+)
+def get_database_connection_info() -> str:
+    """Get all cached database connection information.
+
+    Return:
+        A list of cached connection information.
+    """
+    global db_connection_map
+    return db_connection_map.get_keys_json()
+
+
+@mcp.tool(name='create_cluster', description='Create an Aurora Postgres cluster')
+def create_cluster(
+    region: Annotated[str, Field(description='region')],
+    cluster_identifier: Annotated[str, Field(description='cluster identifier')],
+    database: Annotated[str, Field(description='default database name')] = 'postgres',
+    engine_version: Annotated[str, Field(description='engine version')] = '17.5',
+) -> str:
+    """Create an RDS/Aurora cluster.
+
+    Args:
+        region: region
+        cluster_identifier: cluster identifier
+        database: database name
+        engine_version: engine version
+
+    Returns:
+        result
+    """
+    logger.info(
+        f'Entered create_cluster with region:{region}, '
+        f'cluster_identifier:{cluster_identifier} '
+        f'database:{database} '
+        f'engine_version:{engine_version}'
     )
 
-    # Connection method 1: RDS Data API
-    parser.add_argument('--resource_arn', help='ARN of the RDS cluster (for RDS Data API)')
+    database_type = DatabaseType.APG
+    connection_method = ConnectionMethod.RDS_API
 
-    # Connection method 2: Psycopg Direct Connection
-    parser.add_argument('--hostname', help='Database hostname (for direct PostgreSQL connection)')
-    parser.add_argument('--port', type=int, default=5432, help='Database port (default: 5432)')
+    job_id = (
+        f'create-cluster-{cluster_identifier}-{datetime.now().isoformat(timespec="milliseconds")}'
+    )
 
-    # Common parameters
-    parser.add_argument(
-        '--secret_arn',
-        required=True,
-        help='ARN of the Secrets Manager secret for database credentials',
+    try:
+        async_job_status_lock.acquire()
+        async_job_status[job_id] = {'state': 'pending', 'result': None}
+    finally:
+        async_job_status_lock.release()
+
+    t = threading.Thread(
+        target=create_cluster_worker,
+        args=(
+            job_id,
+            region,
+            database_type,
+            connection_method,
+            cluster_identifier,
+            engine_version,
+            database,
+        ),
+        daemon=False,
     )
-    parser.add_argument('--database', required=True, help='Database name')
-    parser.add_argument('--region', required=True, help='AWS region')
-    parser.add_argument('--readonly', required=True, help='Enforce readonly SQL statements')
+    t.start()
+
+    logger.info(
+        f'start_create_cluster_job return with job_id:{job_id}'
+        f'region:{region} cluster_identifier:{cluster_identifier} database:{database} '
+        f'engine_version:{engine_version}'
+    )
+
+    result = {
+        'status': 'Pending',
+        'message': 'cluster creation started',
+        'job_id': job_id,
+        'cluster_identifier': cluster_identifier,
+        'check_status_tool': 'get_job_status',
+        'next_action': f"Use get_job_status(job_id='{job_id}') to get results",
+    }
+
+    return json.dumps(result, indent=2)
 
-    args = parser.parse_args()
 
-    # Validate connection parameters
-    if not args.resource_arn and not args.hostname:
-        parser.error(
-            'Either --resource_arn (for RDS Data API) or '
-            '--hostname (for direct PostgreSQL) must be provided'
+@mcp.tool(name='get_job_status', description='get background job status')
+def get_job_status(job_id: str) -> dict:
+    """Get background job status.
+
+    Args:
+        job_id: job id
+    Returns:
+        job status
+    """
+    global async_job_status
+    global async_job_status_lock
+
+    try:
+        async_job_status_lock.acquire()
+        return async_job_status.get(job_id, {'state': 'not_found'})
+    finally:
+        async_job_status_lock.release()
+
+
+def create_cluster_worker(
+    job_id: str,
+    region: str,
+    database_type: DatabaseType,
+    connection_method: ConnectionMethod,
+    cluster_identifier: str,
+    engine_version: str,
+    database: str,
+):
+    """Background worker to create a cluster asynchronously."""
+    global db_connection_map
+    global async_job_status
+    global async_job_status_lock
+    global readonly_query
+
+    try:
+        cluster_result = internal_create_serverless_cluster(
+            region=region,
+            cluster_identifier=cluster_identifier,
+            engine_version=engine_version,
+            database_name=database,
         )
 
-    if args.resource_arn and args.hostname:
-        parser.error(
-            'Cannot specify both --resource_arn and --hostname. Choose one connection method.'
+        setup_aurora_iam_policy_for_current_user(
+            db_user=cluster_result['MasterUsername'],
+            cluster_resource_id=cluster_result['DbClusterResourceId'],
+            cluster_region=region,
         )
 
-    # Convert args to dict for easier handling
-    connection_params = vars(args)
+        internal_connect_to_database(
+            region=region,
+            database_type=database_type,
+            connection_method=connection_method,
+            cluster_identifier=cluster_identifier,
+            db_endpoint=cluster_result['Endpoint'],
+            port=5432,
+            database=database,
+        )
 
-    # Convert readonly string to boolean
-    connection_params['readonly'] = args.readonly.lower() == 'true'
+        try:
+            async_job_status_lock.acquire()
+            async_job_status[job_id]['state'] = 'succeeded'
+        finally:
+            async_job_status_lock.release()
+    except Exception as e:
+        logger.error(f'create_cluster_worker failed with {e}')
+        try:
+            async_job_status_lock.acquire()
+            async_job_status[job_id]['state'] = 'failed'
+            async_job_status[job_id]['result'] = str(e)
+        finally:
+            async_job_status_lock.release()
+
+
+def internal_connect_to_database(
+    region: Annotated[str, Field(description='region')],
+    database_type: Annotated[DatabaseType, Field(description='database type')],
+    connection_method: Annotated[ConnectionMethod, Field(description='connection method')],
+    cluster_identifier: Annotated[str, Field(description='cluster identifier')],
+    db_endpoint: Annotated[str, Field(description='database endpoint')],
+    port: Annotated[int, Field(description='Postgres port')],
+    database: Annotated[str, Field(description='database name')] = 'postgres',
+) -> Tuple:
+    """Connect to a specific database save the connection internally.
 
-    # Log connection information
-    connection_target = args.resource_arn if args.resource_arn else f'{args.hostname}:{args.port}'
+    Args:
+        region: region
+        database_type: database type (APG or RPG)
+        connection_method: connection method (RDS_API, PG_WIRE_PROTOCOL, or PG_WIRE_IAM_PROTOCOL)
+        cluster_identifier: cluster identifier
+        db_endpoint: database endpoint
+        port: database port
+        database: database name
+    """
+    global db_connection_map
+    global readonly_query
+
+    logger.info(
+        f'Enter internal_connect_to_database\n'
+        f'region:{region}\n'
+        f'database_type:{database_type}\n'
+        f'connection_method:{connection_method}\n'
+        f'cluster_identifier:{cluster_identifier}\n'
+        f'db_endpoint:{db_endpoint}\n'
+        f'database:{database}\n'
+        f'readonly_query:{readonly_query}'
+    )
 
-    if args.resource_arn:
-        logger.info(
-            f'Postgres MCP init with RDS Data API: CONNECTION_TARGET:{connection_target}, SECRET_ARN:{args.secret_arn}, REGION:{args.region}, DATABASE:{args.database}, READONLY:{args.readonly}'
+    if not region:
+        raise ValueError("region can't be none or empty")
+
+    if not connection_method:
+        raise ValueError("connection_method can't be none or empty")
+
+    if not database_type:
+        raise ValueError("database_type can't be none or empty")
+
+    if database_type == DatabaseType.APG and not cluster_identifier:
+        raise ValueError("cluster_identifier can't be none or empty for Aurora Postgres Database")
+
+    existing_conn = db_connection_map.get(
+        connection_method, cluster_identifier, db_endpoint, database, port
+    )
+    if existing_conn:
+        llm_response = json.dumps(
+            {
+                'connection_method': connection_method,
+                'cluster_identifier': cluster_identifier,
+                'db_endpoint': db_endpoint,
+                'database': database,
+                'port': port,
+            },
+            indent=2,
+            default=str,
         )
-    else:
-        logger.info(
-            f'Postgres MCP init with psycopg: CONNECTION_TARGET:{connection_target}, PORT:{args.port}, DATABASE:{args.database}, READONLY:{args.readonly}'
+        return (existing_conn, llm_response)
+
+    enable_data_api: bool = False
+    masteruser: str = ''
+    cluster_arn: str = ''
+    secret_arn: str = ''
+
+    if cluster_identifier:
+        # Can be either APG (APG always requires cluster) or RPG multi-AZ cluster deployment case
+        cluster_properties = internal_get_cluster_properties(
+            cluster_identifier=cluster_identifier, region=region
         )
 
-    # Create the appropriate database connection based on the provided parameters
+        enable_data_api = cluster_properties.get('HttpEndpointEnabled', False)
+        masteruser = cluster_properties.get('MasterUsername', '')
+        cluster_arn = cluster_properties.get('DBClusterArn', '')
+        secret_arn = cluster_properties.get('MasterUserSecret', {}).get('SecretArn')
+
+        if not db_endpoint:
+            # if db_endpoint not set, we will use cluster's endpoint
+            db_endpoint = cluster_properties.get('Endpoint', '')
+            port = int(cluster_properties.get('Port', ''))
+    else:
+        # Must be RPG instance only deployment case (i.e. without cluster)
+        instance_properties = internal_get_instance_properties(db_endpoint, region)
+        masteruser = instance_properties.get('MasterUsername', '')
+        secret_arn = instance_properties.get('MasterUserSecret', {}).get('SecretArn')
+        port = int(instance_properties.get('Endpoint', {}).get('Port'))
+
+    logger.info(
+        f'About to create internal DB connections with:'
+        f'enable_data_api:{enable_data_api}\n'
+        f'masteruser:{masteruser}\n'
+        f'cluster_arn:{cluster_arn}\n'
+        f'secret_arn:{secret_arn}\n'
+        f'db_endpoint:{db_endpoint}\n'
+        f'port:{port}\n'
+        f'region:{region}\n'
+        f'readonly:{readonly_query}'
+    )
+
     db_connection = None
+    if connection_method == ConnectionMethod.PG_WIRE_IAM_PROTOCOL:
+        db_connection = PsycopgPoolConnection(
+            host=db_endpoint,
+            port=port,
+            database=database,
+            readonly=readonly_query,
+            secret_arn='',
+            db_user=masteruser,
+            region=region,
+            is_iam_auth=True,
+        )
 
-    try:
-        if args.resource_arn:
-            # Use RDS Data API with singleton pattern
-            try:
-                # Initialize the RDS Data API connection singleton
-                DBConnectionSingleton.initialize(
-                    resource_arn=args.resource_arn,
-                    secret_arn=args.secret_arn,
-                    database=args.database,
-                    region=args.region,
-                    readonly=connection_params['readonly'],
-                )
+    elif connection_method == ConnectionMethod.RDS_API:
+        db_connection = RDSDataAPIConnection(
+            cluster_arn=cluster_arn,
+            secret_arn=str(secret_arn),
+            database=database,
+            region=region,
+            readonly=readonly_query,
+        )
+    else:
+        # must be connection_method == ConnectionMethod.PG_WIRE_PROTOCOL
+        db_connection = PsycopgPoolConnection(
+            host=db_endpoint,
+            port=port,
+            database=database,
+            readonly=readonly_query,
+            secret_arn=secret_arn,
+            db_user='',
+            region=region,
+            is_iam_auth=False,
+        )
 
-                # Get the connection from the singleton
-                db_connection = DBConnectionSingleton.get().db_connection
-            except Exception as e:
-                logger.exception(f'Failed to create RDS Data API connection: {str(e)}')
-                sys.exit(1)
-
-        else:
-            # Use Direct PostgreSQL connection using psycopg connection pool
-            try:
-                # Create a direct PostgreSQL connection pool
-                db_connection = PsycopgPoolConnection(
-                    host=args.hostname,
-                    port=args.port,
-                    database=args.database,
-                    readonly=connection_params['readonly'],
-                    secret_arn=args.secret_arn,
-                    region=args.region,
-                )
-            except Exception as e:
-                logger.exception(f'Failed to create PostgreSQL connection: {str(e)}')
-                sys.exit(1)
-
-    except BotoCoreError as e:
-        logger.exception(f'Failed to create database connection: {str(e)}')
-        sys.exit(1)
-
-    # Test database connection
-    ctx = DummyCtx()
-    response = asyncio.run(run_query('SELECT 1', ctx, db_connection))
-    if (
-        isinstance(response, list)
-        and len(response) == 1
-        and isinstance(response[0], dict)
-        and 'error' in response[0]
-    ):
-        logger.error('Failed to validate database connection to Postgres. Exit the MCP server')
-        sys.exit(1)
+    if db_connection:
+        db_connection_map.set(
+            connection_method, cluster_identifier, db_endpoint, database, db_connection
+        )
+        llm_response = json.dumps(
+            {
+                'connection_method': connection_method,
+                'cluster_identifier': cluster_identifier,
+                'db_endpoint': db_endpoint,
+                'database': database,
+                'port': port,
+            },
+            indent=2,
+            default=str,
+        )
+        return (db_connection, llm_response)
 
-    logger.success('Successfully validated database connection to Postgres')
+    raise ValueError("Can't create connection because invalid input parameter combination")
 
-    logger.info('Starting Postgres MCP server')
-    mcp.run()
+
+def main():
+    """Main entry point for the MCP server application.
+
+    Runs the MCP server with CLI argument support for PostgreSQL connections.
+    """
+    global db_connection_map
+    global readonly_query
+
+    parser = argparse.ArgumentParser(
+        description='An AWS Labs Model Context Protocol (MCP) server for postgres'
+    )
+
+    parser.add_argument(
+        '--connection_method',
+        help='Connection method to the database. It can be RDS_API, PG_WIRE_PROTOCOL OR PG_WIRE_IAM_PROTOCOL)',
+    )
+    parser.add_argument('--db_cluster_arn', help='ARN of the RDS or Aurora Postgres cluster')
+    parser.add_argument('--db_type', help='APG for Aurora Postgres or RPG for RDS Postgres')
+    parser.add_argument('--db_endpoint', help='Instance endpoint address')
+    parser.add_argument('--region', help='AWS region')
+    parser.add_argument(
+        '--allow_write_query', action='store_true', help='Enforce readonly SQL statements'
+    )
+    parser.add_argument('--database', help='Database name')
+    parser.add_argument('--port', type=int, default=5432, help='Database port (default: 5432)')
+    args = parser.parse_args()
+
+    logger.info(
+        f'MCP configuration:\n'
+        f'db_type:{args.db_type}\n'
+        f'db_cluster_arn:{args.db_cluster_arn}\n'
+        f'connection_method:{args.connection_method}\n'
+        f'db_endpoint:{args.db_endpoint}\n'
+        f'region:{args.region}\n'
+        f'allow_write_query:{args.allow_write_query}\n'
+        f'database:{args.database}\n'
+        f'port:{args.port}\n'
+    )
+
+    readonly_query = not args.allow_write_query
+
+    try:
+        if args.db_type:
+            # Create the appropriate database connection based on the provided parameters
+            db_connection: Optional[AbstractDBConnection] = None
+
+            cluster_identifier = args.db_cluster_arn.split(':')[-1]
+            db_connection, llm_response = internal_connect_to_database(
+                region=args.region,
+                database_type=DatabaseType[args.db_type],
+                connection_method=ConnectionMethod[args.connection_method],
+                cluster_identifier=cluster_identifier,
+                db_endpoint=args.hostname,
+                port=args.port,
+                database=args.database,
+            )
+
+            # Test database connection
+            if db_connection:
+                ctx = DummyCtx()
+                response = asyncio.run(
+                    run_query(
+                        'SELECT 1',
+                        ctx,
+                        ConnectionMethod[args.connection_method],
+                        cluster_identifier,
+                        args.db_endpoint,
+                        args.database,
+                    )
+                )
+                if (
+                    isinstance(response, list)
+                    and len(response) == 1
+                    and isinstance(response[0], dict)
+                    and 'error' in response[0]
+                ):
+                    logger.error(
+                        'Failed to validate database connection to Postgres. Exit the MCP server'
+                    )
+                    sys.exit(1)
+                else:
+                    logger.success('Successfully validated database connection to Postgres')
+
+        logger.info('Postgres MCP server started')
+        mcp.run()
+        logger.info('Postgres MCP server stopped')
+    finally:
+        db_connection_map.close_all()
 
 
 if __name__ == '__main__':