From 4c8db866860b422538f2e92eaaec78de587d3403 Mon Sep 17 00:00:00 2001 From: openhands Date: Mon, 15 Sep 2025 12:38:44 +0000 Subject: [PATCH] AIP-308: Add username authentication support for Redshift MCP Modified to pass username to execute-statement and batch statement Added a username auth option that allows to pull a different username and use default temp creds mechanisms. --- src/redshift-mcp-server/README.md | 43 +++++++++++ .../awslabs/redshift_mcp_server/consts.py | 6 ++ .../awslabs/redshift_mcp_server/models.py | 14 ++++ .../awslabs/redshift_mcp_server/redshift.py | 73 ++++++++++++++++--- .../awslabs/redshift_mcp_server/server.py | 19 +++++ 5 files changed, 145 insertions(+), 10 deletions(-) diff --git a/src/redshift-mcp-server/README.md b/src/redshift-mcp-server/README.md index d3e3e64815..39cb1cc221 100644 --- a/src/redshift-mcp-server/README.md +++ b/src/redshift-mcp-server/README.md @@ -9,6 +9,7 @@ This MCP server provides tools to discover, explore, and query Amazon Redshift c - **Cluster Discovery**: Automatically discover both provisioned Redshift clusters and serverless workgroups - **Metadata Exploration**: Browse databases, schemas, tables, and columns - **Safe Query Execution**: Execute SQL queries in a READ ONLY mode (a safe READ WRITE support is planned to be implemnted in the future versions) +- **Multiple Authentication Methods**: Support for both IAM authentication and username authentication - **Multi-Cluster Support**: Work with multiple clusters and workgroups simultaneously ## Prerequisites @@ -27,6 +28,18 @@ This MCP server provides tools to discover, explore, and query Amazon Redshift c - Region specified in your AWS profile configuration 3. **Permissions**: Ensure your AWS credentials have the required permissions (see [Permissions](#permissions) section) +### Authentication Methods + +1. **IAM Authentication** (Default): Uses AWS IAM credentials to authenticate with Redshift + - Requires appropriate IAM permissions for Redshift Data API + - No additional configuration needed beyond AWS credentials + +2. **Username Authentication**: Uses database username to authenticate + - Set the following environment variables: + - `REDSHIFT_AUTH_TYPE=username` + - `REDSHIFT_USERNAME=` + - Or provide these parameters directly to the MCP tools + ## Installation | Cursor | VS Code | @@ -53,6 +66,36 @@ Configure the MCP server in your MCP client configuration (e.g., for Amazon Q De } ``` +#### Local Development Configuration + +For running from a local directory: + +```json +{ + "mcpServers": { + "awslabs.redshift-mcp-server": { + "command": "aws-vault", + "args": [ + "exec", + "your-aws-profile", + "--", + "uv", + "run", + "--directory", + "/path/to/redshift-mcp-server", + "awslabs.redshift-mcp-server" + ], + "env": { + "AWS_DEFAULT_REGION": "us-east-1", + "FASTMCP_LOG_LEVEL": "INFO", + "REDSHIFT_AUTH_TYPE": "username", + "REDSHIFT_USERNAME": "sagemaker_readonly" + } + } + } +} +``` + ### Windows Installation For Windows users, the MCP server configuration format is slightly different: diff --git a/src/redshift-mcp-server/awslabs/redshift_mcp_server/consts.py b/src/redshift-mcp-server/awslabs/redshift_mcp_server/consts.py index 155dacf47c..5d01fea291 100644 --- a/src/redshift-mcp-server/awslabs/redshift_mcp_server/consts.py +++ b/src/redshift-mcp-server/awslabs/redshift_mcp_server/consts.py @@ -14,6 +14,11 @@ """Redshift MCP Server constants.""" +# Authentication +AUTH_TYPE_IAM = "iam" +AUTH_TYPE_USERNAME = "username" +DEFAULT_AUTH_TYPE = AUTH_TYPE_IAM + # System CLIENT_CONNECT_TIMEOUT = 60 CLIENT_READ_TIMEOUT = 600 @@ -59,6 +64,7 @@ - We are use the Redshift API and Redshift Data API. - Leverage IAM authentication when possible instead of secrets (database passwords). +- When IAM authentication is not available, username authentication can be used as an alternative. """ # SQL queries diff --git a/src/redshift-mcp-server/awslabs/redshift_mcp_server/models.py b/src/redshift-mcp-server/awslabs/redshift_mcp_server/models.py index 377de8c7da..3bf0d11609 100644 --- a/src/redshift-mcp-server/awslabs/redshift_mcp_server/models.py +++ b/src/redshift-mcp-server/awslabs/redshift_mcp_server/models.py @@ -38,6 +38,20 @@ class RedshiftCluster(BaseModel): tags: Optional[Dict[str, str]] = Field( default_factory=dict, description='Tags associated with the cluster' ) + auth_type: Optional[str] = Field( + None, description='Authentication type (iam or username)' + ) + + +class RedshiftAuthConfig(BaseModel): + """Authentication configuration for Redshift connections.""" + + auth_type: str = Field( + 'iam', description='Authentication type (iam or username)' + ) + username: Optional[str] = Field( + None, description='Username for authentication' + ) class RedshiftDatabase(BaseModel): diff --git a/src/redshift-mcp-server/awslabs/redshift_mcp_server/redshift.py b/src/redshift-mcp-server/awslabs/redshift_mcp_server/redshift.py index cb857c8382..8bb3043e3e 100644 --- a/src/redshift-mcp-server/awslabs/redshift_mcp_server/redshift.py +++ b/src/redshift-mcp-server/awslabs/redshift_mcp_server/redshift.py @@ -20,6 +20,9 @@ import regex from awslabs.redshift_mcp_server import __version__ from awslabs.redshift_mcp_server.consts import ( + AUTH_TYPE_IAM, + AUTH_TYPE_USERNAME, + DEFAULT_AUTH_TYPE, CLIENT_CONNECT_TIMEOUT, CLIENT_READ_TIMEOUT, CLIENT_RETRIES, @@ -40,11 +43,14 @@ class RedshiftClientManager: """Manages AWS clients for Redshift operations.""" def __init__( - self, config: Config, aws_region: str | None = None, aws_profile: str | None = None + self, config: Config, aws_region: str | None = None, aws_profile: str | None = None, + auth_type: str | None = None, username: str | None = None ): """Initialize the client manager.""" self.aws_region = aws_region self.aws_profile = aws_profile + self.auth_type = auth_type + self.username = username self._redshift_client = None self._redshift_serverless_client = None self._redshift_data_client = None @@ -187,20 +193,59 @@ async def execute_statement( sqls = [f"SET application_name TO '{CLIENT_USER_AGENT_NAME}/{__version__}';"] + sqls logger.debug(f'Protected and versioned SQL: {" ".join(sqls)}') + + # For execute_statement, we need to use the actual SQL, not the transaction wrapper + actual_sql = sql # Execute the query using Data API + auth_type = os.environ.get("REDSHIFT_AUTH_TYPE", "iam") + username = os.environ.get("REDSHIFT_USERNAME") + + # Prepare common parameters + params = { + "Database": database_name, + "Sqls": sqls + } + + # Add authentication parameters based on auth type + if auth_type == "username" and username: + logger.debug(f"Using username authentication with username: {username}") + params["DbUser"] = username + elif auth_type == "iam" and username: + logger.debug(f"Using IAM authentication with database user: {username}") + params["DbUser"] = username + + # Add cluster identifier based on cluster type if cluster_info['type'] == 'provisioned': logger.debug(f'Using ClusterIdentifier for provisioned cluster: {cluster_identifier}') - response = data_client.batch_execute_statement( - ClusterIdentifier=cluster_identifier, Database=database_name, Sqls=sqls - ) + params["ClusterIdentifier"] = cluster_identifier elif cluster_info['type'] == 'serverless': logger.debug(f'Using WorkgroupName for serverless workgroup: {cluster_identifier}') - response = data_client.batch_execute_statement( - WorkgroupName=cluster_identifier, Database=database_name, Sqls=sqls - ) + params["WorkgroupName"] = cluster_identifier else: raise Exception(f'Unknown cluster type: {cluster_info["type"]}') + + # Execute the statement using execute_statement (works with sagemaker_readonly) + logger.debug(f"Auth type: {auth_type}, Username: {username}") + + if (auth_type == "iam" or auth_type == "username") and username: + # For IAM or username auth, use execute_statement with DbUser + logger.debug(f"Using execute_statement with {auth_type} authentication") + execute_params = { + "Database": params["Database"], + "Sql": actual_sql, # Use the actual SQL, not the transaction wrapper + "DbUser": params.get("DbUser") + } + if params.get("ClusterIdentifier"): + execute_params["ClusterIdentifier"] = params.get("ClusterIdentifier") + if params.get("WorkgroupName"): + execute_params["WorkgroupName"] = params.get("WorkgroupName") + + response = data_client.execute_statement(**execute_params) + else: + # Fallback to batch_execute_statement for other cases + logger.warning(f"Falling back to batch_execute_statement. Auth type: {auth_type}, Username: {username}") + response = data_client.batch_execute_statement(**params) query_id = response['Id'] logger.debug(f'Started query execution: {query_id}') @@ -229,9 +274,15 @@ async def execute_statement( raise Exception(f'Query timed out after {QUERY_TIMEOUT} seconds') # Get user query results - subquery2_id = status_response['SubStatements'][2]['Id'] - results_response = data_client.get_statement_result(Id=subquery2_id) - return results_response, subquery2_id + if 'SubStatements' in status_response: + # batch_execute_statement response format + subquery2_id = status_response['SubStatements'][2]['Id'] + results_response = data_client.get_statement_result(Id=subquery2_id) + return results_response, subquery2_id + else: + # execute_statement response format - use the main query ID + results_response = data_client.get_statement_result(Id=query_id) + return results_response, query_id async def discover_clusters() -> list[dict]: @@ -619,4 +670,6 @@ async def execute_query(cluster_identifier: str, database_name: str, sql: str) - ), aws_region=os.environ.get('AWS_REGION'), aws_profile=os.environ.get('AWS_PROFILE'), + auth_type=os.environ.get('REDSHIFT_AUTH_TYPE'), + username=os.environ.get('REDSHIFT_USERNAME'), ) diff --git a/src/redshift-mcp-server/awslabs/redshift_mcp_server/server.py b/src/redshift-mcp-server/awslabs/redshift_mcp_server/server.py index 941f8611d2..c6b9162aed 100644 --- a/src/redshift-mcp-server/awslabs/redshift_mcp_server/server.py +++ b/src/redshift-mcp-server/awslabs/redshift_mcp_server/server.py @@ -17,12 +17,16 @@ import os import sys from awslabs.redshift_mcp_server.consts import ( + AUTH_TYPE_IAM, + AUTH_TYPE_USERNAME, CLIENT_BEST_PRACTICES, + DEFAULT_AUTH_TYPE, DEFAULT_LOG_LEVEL, REDSHIFT_BEST_PRACTICES, ) from awslabs.redshift_mcp_server.models import ( QueryResult, + RedshiftAuthConfig, RedshiftCluster, RedshiftColumn, RedshiftDatabase, @@ -57,6 +61,8 @@ This MCP server provides comprehensive access to Amazon Redshift clusters and serverless workgroups. +It supports both IAM authentication and username authentication for connecting to Redshift clusters. + ## Available Tools ### list_clusters @@ -536,9 +542,19 @@ async def execute_query_tool( sql: str = Field( ..., description='The SQL statement to execute. Should be a single SQL statement.' ), + auth_type: str = Field( + DEFAULT_AUTH_TYPE, + description='Authentication type to use (iam or username). Default is iam.', + ), + username: str = Field( + None, + description='Username for authentication. Required if auth_type is username.', + ), ) -> QueryResult: """Execute a SQL query against a Redshift cluster or serverless workgroup. + This tool supports both IAM authentication and username authentication. + This tool uses the Redshift Data API to execute SQL queries and return results. It supports both provisioned clusters and serverless workgroups, and handles various data types in the result set. @@ -549,6 +565,7 @@ async def execute_query_tool( - The cluster must be available and accessible. - Required IAM permissions: redshift-data:ExecuteStatement, redshift-data:DescribeStatement, redshift-data:GetStatementResult. - The user must have appropriate permissions to execute queries in the specified database. + - For username authentication, provide the username parameter and set auth_type to "username". ## Parameters @@ -557,6 +574,8 @@ async def execute_query_tool( - database_name: The database name to execute the query against. IMPORTANT: Use a valid database name from the list_databases tool. - sql: The SQL statement to execute. Should be a single SQL statement. + - auth_type: Authentication type to use (iam or username). Default is iam. + - username: Username for authentication. Required if auth_type is username. ## Response Structure