awslabs.postgres-mcp-server 1.0.3__py3-none-any.whl → 1.0.4__py3-none-any.whl

awslabs/postgres_mcp_server/__init__.py

@@ -14,4 +14,4 @@
 
 """awslabs.postgres-mcp-server"""
 
-__version__ = '0.0.4'
+__version__ = '1.0.4'

awslabs/postgres_mcp_server/connection/__init__.py

@@ -0,0 +1,18 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""aws.postgres-mcp-server.connection"""
+
+from awslabs.postgres_mcp_server.connection.db_connection_singleton import DBConnectionSingleton
+from awslabs.postgres_mcp_server.connection.abstract_db_connection import AbstractDBConnection

awslabs/postgres_mcp_server/connection/abstract_db_connection.py

@@ -0,0 +1,68 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Abstract database connection interface for postgres MCP Server."""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+
+
+class AbstractDBConnection(ABC):
+    """Abstract base class for database connections."""
+
+    def __init__(self, readonly: bool):
+        """Initialize the database connection.
+
+        Args:
+            readonly: Whether the connection should be read-only
+        """
+        self._readonly = readonly
+
+    @property
+    def readonly_query(self) -> bool:
+        """Get whether this connection is read-only.
+
+        Returns:
+            bool: True if the connection is read-only, False otherwise
+        """
+        return self._readonly
+
+    @abstractmethod
+    async def execute_query(
+        self, sql: str, parameters: Optional[List[Dict[str, Any]]] = None
+    ) -> Dict[str, Any]:
+        """Execute a SQL query.
+
+        Args:
+            sql: The SQL query to execute
+            parameters: Optional parameters for the query
+
+        Returns:
+            Dict containing query results with column metadata and records
+        """
+        pass
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Close the database connection."""
+        pass
+
+    @abstractmethod
+    async def check_connection_health(self) -> bool:
+        """Check if the database connection is healthy.
+
+        Returns:
+            bool: True if the connection is healthy, False otherwise
+        """
+        pass

awslabs/postgres_mcp_server/connection/db_connection_singleton.py

@@ -0,0 +1,117 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Database connection singleton for postgres MCP Server."""
+
+import asyncio
+from awslabs.postgres_mcp_server.connection.rds_api_connection import RDSDataAPIConnection
+from loguru import logger
+
+
+class DBConnectionSingleton:
+    """Manages a single RDS Data API connection instance across the application."""
+
+    _instance = None
+
+    def __init__(
+        self,
+        resource_arn: str,
+        secret_arn: str,
+        database: str,
+        region: str,
+        readonly: bool = True,
+        is_test: bool = False,
+    ):
+        """Initialize a new DB connection singleton for RDS Data API.
+
+        Args:
+            resource_arn: The ARN of the RDS cluster
+            secret_arn: The ARN of the secret containing credentials
+            database: The name of the database to connect to
+            region: The AWS region where the RDS instance is located
+            readonly: Whether the connection should be read-only (default: True)
+            is_test: Whether this is a test connection (default: False)
+        """
+        if not all([resource_arn, secret_arn, database, region]):
+            raise ValueError(
+                'Missing required connection parameters for RDS Data API. '
+                'Please provide resource_arn, secret_arn, database, and region.'
+            )
+
+        self._db_connection = RDSDataAPIConnection(
+            cluster_arn=resource_arn,
+            secret_arn=secret_arn,
+            database=database,
+            region=region,
+            readonly=readonly,
+            is_test=is_test,
+        )
+
+    @classmethod
+    def initialize(
+        cls,
+        resource_arn: str,
+        secret_arn: str,
+        database: str,
+        region: str,
+        readonly: bool = True,
+        is_test: bool = False,
+    ):
+        """Initialize the singleton instance if it doesn't exist.
+
+        Args:
+            resource_arn: The ARN of the RDS cluster
+            secret_arn: The ARN of the secret containing credentials
+            database: The name of the database to connect to
+            region: The AWS region where the RDS instance is located
+            readonly: Whether the connection should be read-only (default: True)
+            is_test: Whether this is a test connection (default: False)
+        """
+        if cls._instance is None:
+            cls._instance = cls(
+                resource_arn=resource_arn,
+                secret_arn=secret_arn,
+                database=database,
+                region=region,
+                readonly=readonly,
+                is_test=is_test,
+            )
+
+    @classmethod
+    def get(cls):
+        """Get the singleton instance."""
+        if cls._instance is None:
+            raise RuntimeError('DBConnectionSingleton is not initialized.')
+        return cls._instance
+
+    @property
+    def db_connection(self):
+        """Get the database connection."""
+        return self._db_connection
+
+    @classmethod
+    def cleanup(cls):
+        """Clean up resources when shutting down."""
+        if cls._instance and cls._instance._db_connection:
+            # Handle calling async close method from sync context
+            try:
+                loop = asyncio.get_event_loop()
+                if loop.is_running():
+                    # If we're in an async context, create a task
+                    asyncio.create_task(cls._instance._db_connection.close())
+                else:
+                    # If we're in a sync context, run the coroutine to completion
+                    loop.run_until_complete(cls._instance._db_connection.close())
+            except Exception as e:
+                logger.error(f'Error during connection cleanup: {str(e)}')

awslabs/postgres_mcp_server/connection/psycopg_pool_connection.py

@@ -0,0 +1,287 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Psycopg connector for postgres MCP Server.
+
+This connector provides direct connection to PostgreSQL databases using psycopg.
+It supports both Aurora PostgreSQL and RDS PostgreSQL instances via direct connection
+parameters (host, port, database, user, password) or via AWS Secrets Manager.
+"""
+
+import boto3
+import json
+from awslabs.postgres_mcp_server.connection.abstract_db_connection import AbstractDBConnection
+from loguru import logger
+from psycopg_pool import AsyncConnectionPool
+from typing import Any, Dict, List, Optional, Tuple
+
+
+class PsycopgPoolConnection(AbstractDBConnection):
+    """Class that wraps DB connection using psycopg connection pool.
+
+    This class can connect directly to any PostgreSQL database, including:
+    - Aurora PostgreSQL (using the cluster endpoint)
+    - RDS PostgreSQL (using the instance endpoint)
+    - Self-hosted PostgreSQL
+
+    It uses AWS Secrets Manager (secret_arn and region) for authentication.
+    """
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        database: str,
+        readonly: bool,
+        secret_arn: str,
+        region: str,
+        min_size: int = 1,
+        max_size: int = 10,
+        is_test: bool = False,
+    ):
+        """Initialize a new DB connection pool.
+
+        Args:
+            host: Database host (Aurora cluster endpoint or RDS instance endpoint)
+            port: Database port
+            database: Database name
+            readonly: Whether connections should be read-only
+            secret_arn: ARN of the secret containing credentials
+            region: AWS region for Secrets Manager
+            min_size: Minimum number of connections in the pool
+            max_size: Maximum number of connections in the pool
+            is_test: Whether this is a test connection
+        """
+        super().__init__(readonly)
+        self.host = host
+        self.port = port
+        self.database = database
+        self.min_size = min_size
+        self.max_size = max_size
+        self.pool: Optional['AsyncConnectionPool[Any]'] = None
+
+        # Get credentials from Secrets Manager
+        logger.info(f'Retrieving credentials from Secrets Manager: {secret_arn}')
+        self.user, self.password = self._get_credentials_from_secret(secret_arn, region, is_test)
+        logger.info(f'Successfully retrieved credentials for user: {self.user}')
+
+        # Store connection info
+        if not is_test:
+            self.conninfo = f'host={host} port={port} dbname={database} user={self.user} password={self.password}'
+            logger.info('Connection parameters stored')
+
+    async def initialize_pool(self):
+        """Initialize the connection pool."""
+        if self.pool is None:
+            logger.info(
+                f'Initializing connection pool with min_size={self.min_size}, max_size={self.max_size}'
+            )
+            self.pool = AsyncConnectionPool(
+                self.conninfo, min_size=self.min_size, max_size=self.max_size, open=True
+            )
+            logger.info('Connection pool initialized successfully')
+
+            # Set read-only mode if needed
+            if self.readonly_query:
+                await self._set_all_connections_readonly()
+
+    async def _get_connection(self):
+        """Get a database connection from the pool."""
+        if self.pool is None:
+            await self.initialize_pool()
+
+        if self.pool is None:
+            raise ValueError('Failed to initialize connection pool')
+
+        return self.pool.connection(timeout=15.0)
+
+    async def _set_all_connections_readonly(self):
+        """Set all connections in the pool to read-only mode."""
+        if self.pool is None:
+            logger.warning('Connection pool is not initialized, cannot set read-only mode')
+            return
+
+        try:
+            async with self.pool.connection(timeout=15.0) as conn:
+                await conn.execute(
+                    'ALTER ROLE CURRENT_USER SET default_transaction_read_only = on'
+                )  # type: ignore
+                logger.info('Successfully set connection to read-only mode')
+        except Exception as e:
+            logger.warning(f'Failed to set connections to read-only mode: {str(e)}')
+            logger.warning('Continuing without setting read-only mode')
+
+    async def execute_query(
+        self, sql: str, parameters: Optional[List[Dict[str, Any]]] = None
+    ) -> Dict[str, Any]:
+        """Execute a SQL query using async connection."""
+        try:
+            async with await self._get_connection() as conn:
+                async with conn.transaction():
+                    if self.readonly_query:
+                        await conn.execute('SET TRANSACTION READ ONLY')  # type: ignore
+
+                    # Create a cursor for better control
+                    async with conn.cursor() as cursor:
+                        # Execute the query
+                        if parameters:
+                            params = self._convert_parameters(parameters)
+                            await cursor.execute(sql, params)
+                        else:
+                            await cursor.execute(sql)
+
+                        # Check if there are results to fetch by examining the cursor's description
+                        if cursor.description:
+                            # Get column names
+                            columns = [desc[0] for desc in cursor.description]
+
+                            # Fetch all rows
+                            rows = await cursor.fetchall()
+
+                            # Structure the response to match the interface contract required by server.py
+                            column_metadata = [{'name': col} for col in columns]
+                            records = []
+
+                            # Convert each row to the expected format
+                            for row in rows:
+                                record = []
+                                for value in row:
+                                    if value is None:
+                                        record.append({'isNull': True})
+                                    elif isinstance(value, str):
+                                        record.append({'stringValue': value})
+                                    elif isinstance(value, int):
+                                        record.append({'longValue': value})
+                                    elif isinstance(value, float):
+                                        record.append({'doubleValue': value})
+                                    elif isinstance(value, bool):
+                                        record.append({'booleanValue': value})
+                                    elif isinstance(value, bytes):
+                                        record.append({'blobValue': value})
+                                    else:
+                                        # Convert other types to string
+                                        record.append({'stringValue': str(value)})
+                                records.append(record)
+
+                            return {'columnMetadata': column_metadata, 'records': records}
+                        else:
+                            # No results (e.g., for INSERT, UPDATE, etc.)
+                            return {'columnMetadata': [], 'records': []}
+
+        except Exception as e:
+            logger.error(f'Database connection error: {str(e)}')
+            raise e
+
+    def _convert_parameters(self, parameters: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """Transform structured parameter format to psycopg's native parameter format."""
+        result = {}
+        for param in parameters:
+            name = param.get('name')
+            value = param.get('value', {})
+
+            # Extract the value based on its type
+            if 'stringValue' in value:
+                result[name] = value['stringValue']
+            elif 'longValue' in value:
+                result[name] = value['longValue']
+            elif 'doubleValue' in value:
+                result[name] = value['doubleValue']
+            elif 'booleanValue' in value:
+                result[name] = value['booleanValue']
+            elif 'blobValue' in value:
+                result[name] = value['blobValue']
+            elif 'isNull' in value and value['isNull']:
+                result[name] = None
+
+        return result
+
+    def _get_credentials_from_secret(
+        self, secret_arn: str, region: str, is_test: bool = False
+    ) -> Tuple[str, str]:
+        """Get database credentials from AWS Secrets Manager."""
+        if is_test:
+            return 'test_user', 'test_password'
+
+        try:
+            # Create a Secrets Manager client
+            logger.info(f'Creating Secrets Manager client in region {region}')
+            session = boto3.Session()
+            client = session.client(service_name='secretsmanager', region_name=region)
+
+            # Get the secret value
+            logger.info(f'Retrieving secret value for {secret_arn}')
+            get_secret_value_response = client.get_secret_value(SecretId=secret_arn)
+            logger.info('Successfully retrieved secret value')
+
+            # Parse the secret string
+            if 'SecretString' in get_secret_value_response:
+                secret = json.loads(get_secret_value_response['SecretString'])
+                logger.info(f'Secret keys: {", ".join(secret.keys())}')
+
+                # Extract username and password
+                username = secret.get('username') or secret.get('user') or secret.get('Username')
+                password = secret.get('password') or secret.get('Password')
+
+                if not username:
+                    logger.error(
+                        f'Username not found in secret. Available keys: {", ".join(secret.keys())}'
+                    )
+                    raise ValueError(
+                        f'Secret does not contain username. Available keys: {", ".join(secret.keys())}'
+                    )
+
+                if not password:
+                    logger.error('Password not found in secret')
+                    raise ValueError(
+                        f'Secret does not contain password. Available keys: {", ".join(secret.keys())}'
+                    )
+
+                logger.info(f'Successfully extracted credentials for user: {username}')
+                return username, password
+            else:
+                logger.error('Secret does not contain a SecretString')
+                raise ValueError('Secret does not contain a SecretString')
+        except Exception as e:
+            logger.error(f'Error retrieving secret: {str(e)}')
+            raise ValueError(f'Failed to retrieve credentials from Secrets Manager: {str(e)}')
+
+    async def close(self) -> None:
+        """Close all connections in the pool."""
+        if self.pool is not None:
+            logger.info('Closing connection pool')
+            await self.pool.close()
+            self.pool = None
+            logger.info('Connection pool closed successfully')
+
+    async def check_connection_health(self) -> bool:
+        """Check if the connection is healthy."""
+        try:
+            result = await self.execute_query('SELECT 1')
+            return len(result.get('records', [])) > 0
+        except Exception as e:
+            logger.error(f'Connection health check failed: {str(e)}')
+            return False
+
+    def get_pool_stats(self) -> Dict[str, int]:
+        """Get current connection pool statistics."""
+        if not hasattr(self, 'pool') or self.pool is None:
+            return {'size': 0, 'min_size': self.min_size, 'max_size': self.max_size, 'idle': 0}
+
+        # Access pool attributes safely
+        size = getattr(self.pool, 'size', 0)
+        min_size = getattr(self.pool, 'min_size', self.min_size)
+        max_size = getattr(self.pool, 'max_size', self.max_size)
+        idle = getattr(self.pool, 'idle', 0)
+
+        return {'size': size, 'min_size': min_size, 'max_size': max_size, 'idle': idle}

awslabs/postgres_mcp_server/connection/rds_api_connection.py

@@ -0,0 +1,157 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""RDS Data API connector for postgres MCP Server."""
+
+import asyncio
+import boto3
+from awslabs.postgres_mcp_server.connection.abstract_db_connection import AbstractDBConnection
+from loguru import logger
+from typing import Any, Dict, List, Optional
+
+
+class RDSDataAPIConnection(AbstractDBConnection):
+    """Class that wraps DB connection client by RDS API."""
+
+    def __init__(
+        self,
+        cluster_arn: str,
+        secret_arn: str,
+        database: str,
+        region: str,
+        readonly: bool,
+        is_test: bool = False,
+    ):
+        """Initialize a new DB connection.
+
+        Args:
+            cluster_arn: The ARN of the RDS cluster
+            secret_arn: The ARN of the secret containing credentials
+            database: The name of the database to connect to
+            region: The AWS region where the RDS instance is located
+            readonly: Whether the connection should be read-only
+            is_test: Whether this is a test connection
+        """
+        super().__init__(readonly)
+        self.cluster_arn = cluster_arn
+        self.secret_arn = secret_arn
+        self.database = database
+        if not is_test:
+            self.data_client = boto3.client('rds-data', region_name=region)
+
+    async def execute_query(
+        self, sql: str, parameters: Optional[List[Dict[str, Any]]] = None
+    ) -> Dict[str, Any]:
+        """Execute a SQL query using RDS Data API.
+
+        Args:
+            sql: The SQL query to execute
+            parameters: Optional parameters for the query
+
+        Returns:
+            Dict containing query results with column metadata and records
+        """
+        if self.readonly_query:
+            return await asyncio.to_thread(self._execute_readonly_query, sql, parameters)
+        else:
+            execute_params = {
+                'resourceArn': self.cluster_arn,
+                'secretArn': self.secret_arn,
+                'database': self.database,
+                'sql': sql,
+                'includeResultMetadata': True,
+            }
+
+            if parameters:
+                execute_params['parameters'] = parameters
+
+            return await asyncio.to_thread(self.data_client.execute_statement, **execute_params)
+
+    def _execute_readonly_query(
+        self, query: str, parameters: Optional[List[Dict[str, Any]]] = None
+    ) -> Dict[str, Any]:
+        """Execute a query under readonly transaction.
+
+        Args:
+            query: query to run
+            parameters: parameters
+
+        Returns:
+            Dict containing query results with column metadata and records
+        """
+        tx_id = ''
+        try:
+            # Begin read-only transaction
+            tx = self.data_client.begin_transaction(
+                resourceArn=self.cluster_arn,
+                secretArn=self.secret_arn,
+                database=self.database,
+            )
+
+            tx_id = tx['transactionId']
+
+            self.data_client.execute_statement(
+                resourceArn=self.cluster_arn,
+                secretArn=self.secret_arn,
+                database=self.database,
+                sql='SET TRANSACTION READ ONLY',
+                transactionId=tx_id,
+            )
+
+            execute_params = {
+                'resourceArn': self.cluster_arn,
+                'secretArn': self.secret_arn,
+                'database': self.database,
+                'sql': query,
+                'includeResultMetadata': True,
+                'transactionId': tx_id,
+            }
+
+            if parameters is not None:
+                execute_params['parameters'] = parameters
+
+            result = self.data_client.execute_statement(**execute_params)
+
+            self.data_client.commit_transaction(
+                resourceArn=self.cluster_arn,
+                secretArn=self.secret_arn,
+                transactionId=tx_id,
+            )
+            return result
+        except Exception as e:
+            if tx_id:
+                self.data_client.rollback_transaction(
+                    resourceArn=self.cluster_arn,
+                    secretArn=self.secret_arn,
+                    transactionId=tx_id,
+                )
+            raise e
+
+    async def close(self) -> None:
+        """Close the database connection asynchronously."""
+        # RDS Data API doesn't maintain persistent connections
+        pass
+
+    async def check_connection_health(self) -> bool:
+        """Check if the RDS Data API connection is healthy.
+
+        Returns:
+            bool: True if the connection is healthy, False otherwise
+        """
+        try:
+            result = await self.execute_query('SELECT 1')
+            return len(result.get('records', [])) > 0
+        except Exception as e:
+            logger.error(f'RDS Data API connection health check failed: {str(e)}')
+            return False

awslabs/postgres_mcp_server/server.py

@@ -16,8 +16,9 @@
 
 import argparse
 import asyncio
-import boto3
 import sys
+from awslabs.postgres_mcp_server.connection import DBConnectionSingleton
+from awslabs.postgres_mcp_server.connection.psycopg_pool_connection import PsycopgPoolConnection
 from awslabs.postgres_mcp_server.mutable_sql_detector import (
     check_sql_injection_risk,
     detect_mutating_keywords,
@@ -49,104 +50,6 @@ class DummyCtx:
         pass
 
 
-class DBConnection:
-    """Class that wraps DB connection client by RDS API."""
-
-    def __init__(self, cluster_arn, secret_arn, database, region, readonly, is_test=False):
-        """Initialize a new DB connection.
-
-        Args:
-            cluster_arn: The ARN of the RDS cluster
-            secret_arn: The ARN of the secret containing credentials
-            database: The name of the database to connect to
-            region: The AWS region where the RDS instance is located
-            readonly: Whether the connection should be read-only
-            is_test: Whether this is a test connection
-        """
-        self.cluster_arn = cluster_arn
-        self.secret_arn = secret_arn
-        self.database = database
-        self.readonly = readonly
-        if not is_test:
-            self.data_client = boto3.client('rds-data', region_name=region)
-
-    @property
-    def readonly_query(self):
-        """Get whether this connection is read-only.
-
-        Returns:
-            bool: True if the connection is read-only, False otherwise
-        """
-        return self.readonly
-
-
-class DBConnectionSingleton:
-    """Manages a single DBConnection instance across the application.
-
-    This singleton ensures that only one DBConnection is created and reused.
-    """
-
-    _instance = None
-
-    def __init__(self, resource_arn, secret_arn, database, region, readonly, is_test=False):
-        """Initialize a new DB connection singleton.
-
-        Args:
-            resource_arn: The ARN of the RDS resource
-            secret_arn: The ARN of the secret containing credentials
-            database: The name of the database to connect to
-            region: The AWS region where the RDS instance is located
-            readonly: Whether the connection should be read-only
-            is_test: Whether this is a test connection
-        """
-        if not all([resource_arn, secret_arn, database, region]):
-            raise ValueError(
-                'Missing required connection parameters. '
-                'Please provide resource_arn, secret_arn, database, and region.'
-            )
-        self._db_connection = DBConnection(
-            resource_arn, secret_arn, database, region, readonly, is_test
-        )
-
-    @classmethod
-    def initialize(cls, resource_arn, secret_arn, database, region, readonly, is_test=False):
-        """Initialize the singleton instance if it doesn't exist.
-
-        Args:
-            resource_arn: The ARN of the RDS resource
-            secret_arn: The ARN of the secret containing credentials
-            database: The name of the database to connect to
-            region: The AWS region where the RDS instance is located
-            readonly: Whether the connection should be read-only
-            is_test: Whether this is a test connection
-        """
-        if cls._instance is None:
-            cls._instance = cls(resource_arn, secret_arn, database, region, readonly, is_test)
-
-    @classmethod
-    def get(cls):
-        """Get the singleton instance.
-
-        Returns:
-            DBConnectionSingleton: The singleton instance
-
-        Raises:
-            RuntimeError: If the singleton has not been initialized
-        """
-        if cls._instance is None:
-            raise RuntimeError('DBConnectionSingleton is not initialized.')
-        return cls._instance
-
-    @property
-    def db_connection(self):
-        """Get the database connection.
-
-        Returns:
-            DBConnection: The database connection instance
-        """
-        return self._db_connection
-
-
 def extract_cell(cell: dict):
     """Extracts the scalar or array value from a single cell."""
     if cell.get('isNull'):
@@ -177,14 +80,14 @@ def parse_execute_response(response: dict) -> list[dict]:
 
 
 mcp = FastMCP(
-    'apg-mcp MCP server. This is the starting point for all solutions created',
+    'pg-mcp MCP server. This is the starting point for all solutions created',
     dependencies=[
         'loguru',
     ],
 )
 
 
-@mcp.tool(name='run_query', description='Run a SQL query using boto3 execute_statement')
+@mcp.tool(name='run_query', description='Run a SQL query against PostgreSQL')
 async def run_query(
     sql: Annotated[str, Field(description='The SQL query to run')],
     ctx: Context,
@@ -193,12 +96,12 @@ async def run_query(
         Optional[List[Dict[str, Any]]], Field(description='Parameters for the SQL query')
     ] = None,
 ) -> list[dict]:  # type: ignore
-    """Run a SQL query using boto3 execute_statement.
+    """Run a SQL query against PostgreSQL.
 
     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 if called by MCP server.
+        db_connection: DB connection object passed by unit test. It should be None if called by MCP server.
         query_parameters: Parameters for the SQL query
 
     Returns:
@@ -209,7 +112,17 @@ async def run_query(
     global write_query_prohibited_key
 
     if db_connection is None:
-        db_connection = DBConnectionSingleton.get().db_connection
+        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'}]
+
+    if db_connection is None:
+        raise AssertionError('db_connection should never be None')
 
     if db_connection.readonly_query:
         matches = detect_mutating_keywords(sql)
@@ -233,27 +146,10 @@ async def run_query(
     try:
         logger.info(f'run_query: readonly:{db_connection.readonly_query}, SQL:{sql}')
 
-        if db_connection.readonly_query:
-            response = await asyncio.to_thread(
-                execute_readonly_query, db_connection, sql, query_parameters
-            )
-        else:
-            execute_params = {
-                'resourceArn': db_connection.cluster_arn,
-                'secretArn': db_connection.secret_arn,
-                'database': db_connection.database,
-                'sql': sql,
-                'includeResultMetadata': True,
-            }
-
-            if query_parameters:
-                execute_params['parameters'] = query_parameters
-
-            response = await asyncio.to_thread(
-                db_connection.data_client.execute_statement, **execute_params
-            )
+        # Execute the query using the abstract connection interface
+        response = await db_connection.execute_query(sql, query_parameters)
 
-        logger.success('run_query successfully executed query:{}', sql)
+        logger.success(f'run_query successfully executed query:{sql}')
         return parse_execute_response(response)
     except ClientError as e:
         logger.exception(client_error_code_key)
@@ -270,7 +166,7 @@ async def run_query(
 
 @mcp.tool(
     name='get_table_schema',
-    description='Fetch table columns and comments from Postgres using RDS Data API',
+    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
@@ -305,68 +201,6 @@ async def get_table_schema(
     return await run_query(sql=sql, ctx=ctx, query_parameters=params)
 
 
-def execute_readonly_query(
-    db_connection: DBConnection, query: str, parameters: Optional[List[Dict[str, Any]]] = None
-) -> dict:
-    """Execute a query under readonly transaction.
-
-    Args:
-        db_connection: connection object
-        query: query to run
-        parameters: parameters
-
-    Returns:
-        List of dictionary that contains query response rows
-    """
-    tx_id = ''
-    try:
-        # Begin read-only transaction
-        tx = db_connection.data_client.begin_transaction(
-            resourceArn=db_connection.cluster_arn,
-            secretArn=db_connection.secret_arn,
-            database=db_connection.database,
-        )
-
-        tx_id = tx['transactionId']
-
-        db_connection.data_client.execute_statement(
-            resourceArn=db_connection.cluster_arn,
-            secretArn=db_connection.secret_arn,
-            database=db_connection.database,
-            sql='SET TRANSACTION READ ONLY',
-            transactionId=tx_id,
-        )
-
-        execute_params = {
-            'resourceArn': db_connection.cluster_arn,
-            'secretArn': db_connection.secret_arn,
-            'database': db_connection.database,
-            'sql': query,
-            'includeResultMetadata': True,
-            'transactionId': tx_id,
-        }
-
-        if parameters is not None:
-            execute_params['parameters'] = parameters
-
-        result = db_connection.data_client.execute_statement(**execute_params)
-
-        db_connection.data_client.commit_transaction(
-            resourceArn=db_connection.cluster_arn,
-            secretArn=db_connection.secret_arn,
-            transactionId=tx_id,
-        )
-        return result
-    except Exception as e:
-        if tx_id:
-            db_connection.data_client.rollback_transaction(
-                resourceArn=db_connection.cluster_arn,
-                secretArn=db_connection.secret_arn,
-                transactionId=tx_id,
-            )
-        raise e
-
-
 def main():
     """Main entry point for the MCP server application."""
     global client_error_code_key
@@ -375,51 +209,111 @@ def main():
     parser = argparse.ArgumentParser(
         description='An AWS Labs Model Context Protocol (MCP) server for postgres'
     )
-    parser.add_argument('--resource_arn', required=True, help='ARN of the RDS cluster')
+
+    # Connection method 1: RDS Data API
+    parser.add_argument('--resource_arn', help='ARN of the RDS cluster (for RDS Data 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)')
+
+    # Common parameters
     parser.add_argument(
         '--secret_arn',
         required=True,
         help='ARN of the Secrets Manager secret for database credentials',
     )
     parser.add_argument('--database', required=True, help='Database name')
-    parser.add_argument(
-        '--region', required=True, help='AWS region for RDS Data API (default: us-west-2)'
-    )
-    parser.add_argument(
-        '--readonly', required=True, help='Enforce NL to SQL to only allow readonly sql statement'
-    )
+    parser.add_argument('--region', required=True, help='AWS region')
+    parser.add_argument('--readonly', required=True, help='Enforce readonly SQL statements')
+
     args = parser.parse_args()
 
-    logger.info(
-        'Postgres MCP init with CLUSTER_ARN:{}, SECRET_ARN:{}, REGION:{}, DATABASE:{}, READONLY:{}',
-        args.resource_arn,
-        args.secret_arn,
-        args.region,
-        args.database,
-        args.readonly,
-    )
+    # 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'
+        )
 
-    try:
-        DBConnectionSingleton.initialize(
-            args.resource_arn, args.secret_arn, args.database, args.region, args.readonly
+    if args.resource_arn and args.hostname:
+        parser.error(
+            'Cannot specify both --resource_arn and --hostname. Choose one connection method.'
         )
-    except BotoCoreError:
-        logger.exception('Failed to RDS API client object for Postgres. Exit the MCP server')
+
+    # Convert args to dict for easier handling
+    connection_params = vars(args)
+
+    # Convert readonly string to boolean
+    connection_params['readonly'] = args.readonly.lower() == 'true'
+
+    # Log connection information
+    connection_target = args.resource_arn if args.resource_arn else f'{args.hostname}:{args.port}'
+
+    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}'
+        )
+    else:
+        logger.info(
+            f'Postgres MCP init with psycopg: CONNECTION_TARGET:{connection_target}, PORT:{args.port}, DATABASE:{args.database}, READONLY:{args.readonly}'
+        )
+
+    # Create the appropriate database connection based on the provided parameters
+    db_connection = None
+
+    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'],
+                )
+
+                # 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 RDS API connection
+    # Test database connection
     ctx = DummyCtx()
-    response = asyncio.run(run_query('SELECT 1', ctx))
+    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 RDS API db connection to Postgres. Exit the MCP server')
+        logger.error('Failed to validate database connection to Postgres. Exit the MCP server')
         sys.exit(1)
 
-    logger.success('Successfully validated RDS API db connection to Postgres')
+    logger.success('Successfully validated database connection to Postgres')
 
     logger.info('Starting Postgres MCP server')
     mcp.run()

{awslabs_postgres_mcp_server-1.0.3.dist-info → awslabs_postgres_mcp_server-1.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: awslabs.postgres-mcp-server
-Version: 1.0.3
+Version: 1.0.4
 Summary: An AWS Labs Model Context Protocol (MCP) server for postgres
 Project-URL: homepage, https://awslabs.github.io/mcp/
 Project-URL: docs, https://awslabs.github.io/mcp/servers/postgres-mcp-server/
@@ -25,6 +25,7 @@ Requires-Dist: boto3>=1.38.5
 Requires-Dist: botocore>=1.38.5
 Requires-Dist: loguru>=0.7.0
 Requires-Dist: mcp[cli]>=1.11.0
+Requires-Dist: psycopg[pool]>=3.1.12
 Requires-Dist: pydantic>=2.10.6
 Description-Content-Type: text/markdown
 
@@ -58,6 +59,8 @@ An AWS Labs Model Context Protocol (MCP) server for Aurora Postgres
 
 Configure the MCP server in your MCP client configuration (e.g., for Amazon Q Developer CLI, edit `~/.aws/amazonq/mcp.json`):
 
+### Option 1: Using RDS Data API Connection (for Aurora Postgres)
+
 ```json
 {
   "mcpServers": {
@@ -83,6 +86,35 @@ Configure the MCP server in your MCP client configuration (e.g., for Amazon Q De
 }
 ```
 
+### Option 2: Using Direct PostgreSQL(psycopg) Connection (for Aurora Postgres and RDS Postgres)
+
+```json
+{
+  "mcpServers": {
+    "awslabs.postgres-mcp-server": {
+      "command": "uvx",
+      "args": [
+        "awslabs.postgres-mcp-server@latest",
+        "--hostname", "[your data]",
+        "--secret_arn", "[your data]",
+        "--database", "[your data]",
+        "--region", "[your data]",
+        "--readonly", "True"
+      ],
+      "env": {
+        "AWS_PROFILE": "your-aws-profile",
+        "AWS_REGION": "us-east-1",
+        "FASTMCP_LOG_LEVEL": "ERROR"
+      },
+      "disabled": false,
+      "autoApprove": []
+    }
+  }
+}
+```
+
+Note: The `--port` parameter is optional and defaults to 5432 (the standard PostgreSQL port). You only need to specify it if your PostgreSQL instance uses a non-standard port.
+
 ### Build and install docker image locally on the same host of your LLM client
 
 1. 'git clone https://github.com/awslabs/mcp.git'
@@ -91,6 +123,8 @@ Configure the MCP server in your MCP client configuration (e.g., for Amazon Q De
 
 ### Add or update your LLM client's config with following:
 
+#### Option 1: Using RDS Data API Connection (for Aurora Postgres)
+
 ```json
 {
   "mcpServers": {
@@ -115,8 +149,46 @@ Configure the MCP server in your MCP client configuration (e.g., for Amazon Q De
 }
 ```
 
+#### Option 2: Using Direct PostgreSQL (psycopg) Connection (for Aurora Postgres and RDS Postgres)
+
+```
+{
+  "mcpServers": {
+    "awslabs.postgres-mcp-server": {
+      "command": "docker",
+      "args": [
+        "run",
+        "-i",
+        "--rm",
+        "-e", "AWS_ACCESS_KEY_ID=[your data]",
+        "-e", "AWS_SECRET_ACCESS_KEY=[your data]",
+        "-e", "AWS_REGION=[your data]",
+        "awslabs/postgres-mcp-server:latest",
+        "--hostname", "[your data]",
+        "--secret_arn", "[your data]",
+        "--database", "[your data]",
+        "--region", "[your data]",
+        "--readonly", "True"
+      ]
+    }
+  }
+}
+```
+
+Note: The `--port` parameter is optional and defaults to 5432 (the standard PostgreSQL port). You only need to specify it if your PostgreSQL instance uses a non-standard port.
+
 NOTE: By default, only read-only queries are allowed and it is controlled by --readonly parameter above. Set it to False if you also want to allow writable DML or DDL.
 
+## Connection Methods
+
+This MCP server supports two connection methods:
+
+1. **RDS Data API Connection** (using `--resource_arn`): Uses the AWS RDS Data API to connect to Aurora PostgreSQL. This method requires that your Aurora cluster has the Data API enabled.
+
+2. **Direct PostgreSQL Connection** (using `--hostname`): Uses psycopg to connect directly to any PostgreSQL database, including Aurora PostgreSQL, RDS PostgreSQL, or self-hosted PostgreSQL instances. This method provides better performance for frequent queries but requires direct network access to the database.
+
+Choose the connection method that best fits your environment and requirements.
+
 ### AWS Authentication
 
 The MCP server uses the AWS profile specified in the `AWS_PROFILE` environment variable. If not provided, it defaults to the "default" profile in your AWS configuration file.

awslabs_postgres_mcp_server-1.0.4.dist-info/RECORD

@@ -0,0 +1,15 @@
+awslabs/__init__.py,sha256=WuqxdDgUZylWNmVoPKiK7qGsTB_G4UmuXIrJ-VBwDew,731
+awslabs/postgres_mcp_server/__init__.py,sha256=Hmj3ZFqpAA8HzLCrR48AtfJthOFX7HH6_iNzE_E_2Q4,673
+awslabs/postgres_mcp_server/mutable_sql_detector.py,sha256=bivLNglukgfi0IWcQSfDcwxFae4Q8CSvb--yWW617l0,2717
+awslabs/postgres_mcp_server/server.py,sha256=wenII8rWhhUVVo-m2C2cPWVwS1S_Us3LUKmKQrr9ebw,11621
+awslabs/postgres_mcp_server/connection/__init__.py,sha256=9j_gZeLcWoXu_Fny1PkC3RsnDMOHr9M3OM9EWvlGR7o,850
+awslabs/postgres_mcp_server/connection/abstract_db_connection.py,sha256=BGdyRoF1fgjLoFSiWg-XdS0KJ1UKuNgV03gGcRXOG8U,2080
+awslabs/postgres_mcp_server/connection/db_connection_singleton.py,sha256=Fwm6uTwB-bhprx-UgiOttDFtczeMUfW5M79wGL0x1Go,4279
+awslabs/postgres_mcp_server/connection/psycopg_pool_connection.py,sha256=qHbRR8-P__v5_-_Q-AQdZVD_nTHj-2Y1p-j1RnKRmrs,12593
+awslabs/postgres_mcp_server/connection/rds_api_connection.py,sha256=1G4qfik5Nu-KbAxRLFNshQ5DJxFnJQI6eqLJPXp-Qn4,5419
+awslabs_postgres_mcp_server-1.0.4.dist-info/METADATA,sha256=aiAGagfCfoskIxv-rVx9mPqTt7rVbqhAf_Oq_DJUlcg,8410
+awslabs_postgres_mcp_server-1.0.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+awslabs_postgres_mcp_server-1.0.4.dist-info/entry_points.txt,sha256=qHCxq_MTANxTB8-mA7Wl6H71qWEMwmG-I6_koII4AXY,88
+awslabs_postgres_mcp_server-1.0.4.dist-info/licenses/LICENSE,sha256=CeipvOyAZxBGUsFoaFqwkx54aPnIKEtm9a5u2uXxEws,10142
+awslabs_postgres_mcp_server-1.0.4.dist-info/licenses/NOTICE,sha256=qMIIe3h7I1Y4-CKejn50wbSKXEZLWhYHdKaRwKdXN9M,95
+awslabs_postgres_mcp_server-1.0.4.dist-info/RECORD,,

awslabs_postgres_mcp_server-1.0.3.dist-info/RECORD

@@ -1,10 +0,0 @@
-awslabs/__init__.py,sha256=WuqxdDgUZylWNmVoPKiK7qGsTB_G4UmuXIrJ-VBwDew,731
-awslabs/postgres_mcp_server/__init__.py,sha256=OpiWDc6HV8EoZVGVKb6ebzcIm7TUbZeALyFkLdm7YWg,673
-awslabs/postgres_mcp_server/mutable_sql_detector.py,sha256=bivLNglukgfi0IWcQSfDcwxFae4Q8CSvb--yWW617l0,2717
-awslabs/postgres_mcp_server/server.py,sha256=Hw8qPrjlqoKoPatP2485RN_ntXfwb09OIPYOAKyRtrk,14500
-awslabs_postgres_mcp_server-1.0.3.dist-info/METADATA,sha256=N7XIGheT6KY3uZguuG-1VpdJchn37MMX5Trc6w7uI90,5952
-awslabs_postgres_mcp_server-1.0.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-awslabs_postgres_mcp_server-1.0.3.dist-info/entry_points.txt,sha256=qHCxq_MTANxTB8-mA7Wl6H71qWEMwmG-I6_koII4AXY,88
-awslabs_postgres_mcp_server-1.0.3.dist-info/licenses/LICENSE,sha256=CeipvOyAZxBGUsFoaFqwkx54aPnIKEtm9a5u2uXxEws,10142
-awslabs_postgres_mcp_server-1.0.3.dist-info/licenses/NOTICE,sha256=qMIIe3h7I1Y4-CKejn50wbSKXEZLWhYHdKaRwKdXN9M,95
-awslabs_postgres_mcp_server-1.0.3.dist-info/RECORD,,