awslabs.mysql-mcp-server 1.0.7__tar.gz → 1.0.9__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/PKG-INFO +44 -1
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/README.md +42 -0
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/awslabs/mysql_mcp_server/__init__.py +1 -1
- awslabs_mysql_mcp_server-1.0.9/awslabs/mysql_mcp_server/connection/__init__.py +18 -0
- awslabs_mysql_mcp_server-1.0.9/awslabs/mysql_mcp_server/connection/abstract_db_connection.py +68 -0
- awslabs_mysql_mcp_server-1.0.9/awslabs/mysql_mcp_server/connection/asyncmy_pool_connection.py +267 -0
- awslabs_mysql_mcp_server-1.0.9/awslabs/mysql_mcp_server/connection/db_connection_singleton.py +142 -0
- awslabs_mysql_mcp_server-1.0.9/awslabs/mysql_mcp_server/connection/rds_data_api_connection.py +93 -0
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/awslabs/mysql_mcp_server/server.py +82 -145
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/pyproject.toml +2 -1
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/tests/conftest.py +32 -1
- awslabs_mysql_mcp_server-1.0.9/tests/test_abstract_db_connection.py +69 -0
- awslabs_mysql_mcp_server-1.0.9/tests/test_asyncmy_pool_connection.py +518 -0
- awslabs_mysql_mcp_server-1.0.9/tests/test_db_connection_singleton.py +135 -0
- awslabs_mysql_mcp_server-1.0.9/tests/test_rds_data_api_connection.py +156 -0
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/tests/test_server.py +202 -10
- awslabs_mysql_mcp_server-1.0.9/uv-requirements.txt +27 -0
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/uv.lock +40 -1
- awslabs_mysql_mcp_server-1.0.7/uv-requirements.txt +0 -24
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/.gitignore +0 -0
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/.python-version +0 -0
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/CHANGELOG.md +0 -0
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/Dockerfile +0 -0
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/LICENSE +0 -0
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/NOTICE +0 -0
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/awslabs/__init__.py +0 -0
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/awslabs/mysql_mcp_server/mutable_sql_detector.py +0 -0
- {awslabs_mysql_mcp_server-1.0.7 → awslabs_mysql_mcp_server-1.0.9}/docker-healthcheck.sh +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: awslabs.mysql-mcp-server
|
|
3
|
-
Version: 1.0.
|
|
3
|
+
Version: 1.0.9
|
|
4
4
|
Summary: An AWS Labs Model Context Protocol (MCP) server for mysql
|
|
5
5
|
Project-URL: homepage, https://awslabs.github.io/mcp/
|
|
6
6
|
Project-URL: docs, https://awslabs.github.io/mcp/servers/mysql-mcp-server/
|
|
@@ -21,6 +21,7 @@ Classifier: Programming Language :: Python :: 3.11
|
|
|
21
21
|
Classifier: Programming Language :: Python :: 3.12
|
|
22
22
|
Classifier: Programming Language :: Python :: 3.13
|
|
23
23
|
Requires-Python: >=3.10
|
|
24
|
+
Requires-Dist: asyncmy>=0.2.10
|
|
24
25
|
Requires-Dist: boto3>=1.38.14
|
|
25
26
|
Requires-Dist: botocore>=1.38.14
|
|
26
27
|
Requires-Dist: loguru>=0.7.0
|
|
@@ -58,6 +59,18 @@ An AWS Labs Model Context Protocol (MCP) server for Aurora MySQL
|
|
|
58
59
|
|
|
59
60
|
Configure the MCP server in your MCP client configuration (e.g., for Amazon Q Developer CLI, edit `~/.aws/amazonq/mcp.json`):
|
|
60
61
|
|
|
62
|
+
## Connection Methods
|
|
63
|
+
|
|
64
|
+
This MCP server supports two connection methods:
|
|
65
|
+
|
|
66
|
+
1. **RDS Data API Connection** (using `--resource_arn`): Uses the AWS RDS Data API to connect to Aurora MySQL. This method requires that your Aurora cluster has the Data API enabled.
|
|
67
|
+
|
|
68
|
+
2. **Direct MySQL Connection** (using `--hostname`): Uses asyncmy to connect directly to any MySQL database, including Aurora MySQL, RDS MySQL, RDS MariaDB, or self-hosted MySQL/MariaDB instances.
|
|
69
|
+
|
|
70
|
+
Choose the connection method that best fits your environment and requirements.
|
|
71
|
+
|
|
72
|
+
### Option 1: Using RDS Data API Connection (for Aurora MySQL)
|
|
73
|
+
|
|
61
74
|
```json
|
|
62
75
|
{
|
|
63
76
|
"mcpServers": {
|
|
@@ -82,6 +95,36 @@ Configure the MCP server in your MCP client configuration (e.g., for Amazon Q De
|
|
|
82
95
|
}
|
|
83
96
|
}
|
|
84
97
|
```
|
|
98
|
+
|
|
99
|
+
### Option 2: Using Direct MySQL Connection (for Aurora MySQL, RDS MySQL, and RDS MariaDB)
|
|
100
|
+
|
|
101
|
+
```json
|
|
102
|
+
{
|
|
103
|
+
"mcpServers": {
|
|
104
|
+
"awslabs.mysql-mcp-server": {
|
|
105
|
+
"command": "uvx",
|
|
106
|
+
"args": [
|
|
107
|
+
"awslabs.mysql-mcp-server@latest",
|
|
108
|
+
"--hostname", "[your data]",
|
|
109
|
+
"--secret_arn", "[your data]",
|
|
110
|
+
"--database", "[your data]",
|
|
111
|
+
"--region", "[your data]",
|
|
112
|
+
"--readonly", "True"
|
|
113
|
+
],
|
|
114
|
+
"env": {
|
|
115
|
+
"AWS_PROFILE": "your-aws-profile",
|
|
116
|
+
"AWS_REGION": "us-east-1",
|
|
117
|
+
"FASTMCP_LOG_LEVEL": "ERROR"
|
|
118
|
+
},
|
|
119
|
+
"disabled": false,
|
|
120
|
+
"autoApprove": []
|
|
121
|
+
}
|
|
122
|
+
}
|
|
123
|
+
}
|
|
124
|
+
```
|
|
125
|
+
|
|
126
|
+
Note: The `--port` parameter is optional and defaults to 3306 (the standard MySQL port). You only need to specify it if your MySQL instance uses a non-default port.
|
|
127
|
+
|
|
85
128
|
### Windows Installation
|
|
86
129
|
|
|
87
130
|
For Windows users, the MCP server configuration format is slightly different:
|
|
@@ -28,6 +28,18 @@ An AWS Labs Model Context Protocol (MCP) server for Aurora MySQL
|
|
|
28
28
|
|
|
29
29
|
Configure the MCP server in your MCP client configuration (e.g., for Amazon Q Developer CLI, edit `~/.aws/amazonq/mcp.json`):
|
|
30
30
|
|
|
31
|
+
## Connection Methods
|
|
32
|
+
|
|
33
|
+
This MCP server supports two connection methods:
|
|
34
|
+
|
|
35
|
+
1. **RDS Data API Connection** (using `--resource_arn`): Uses the AWS RDS Data API to connect to Aurora MySQL. This method requires that your Aurora cluster has the Data API enabled.
|
|
36
|
+
|
|
37
|
+
2. **Direct MySQL Connection** (using `--hostname`): Uses asyncmy to connect directly to any MySQL database, including Aurora MySQL, RDS MySQL, RDS MariaDB, or self-hosted MySQL/MariaDB instances.
|
|
38
|
+
|
|
39
|
+
Choose the connection method that best fits your environment and requirements.
|
|
40
|
+
|
|
41
|
+
### Option 1: Using RDS Data API Connection (for Aurora MySQL)
|
|
42
|
+
|
|
31
43
|
```json
|
|
32
44
|
{
|
|
33
45
|
"mcpServers": {
|
|
@@ -52,6 +64,36 @@ Configure the MCP server in your MCP client configuration (e.g., for Amazon Q De
|
|
|
52
64
|
}
|
|
53
65
|
}
|
|
54
66
|
```
|
|
67
|
+
|
|
68
|
+
### Option 2: Using Direct MySQL Connection (for Aurora MySQL, RDS MySQL, and RDS MariaDB)
|
|
69
|
+
|
|
70
|
+
```json
|
|
71
|
+
{
|
|
72
|
+
"mcpServers": {
|
|
73
|
+
"awslabs.mysql-mcp-server": {
|
|
74
|
+
"command": "uvx",
|
|
75
|
+
"args": [
|
|
76
|
+
"awslabs.mysql-mcp-server@latest",
|
|
77
|
+
"--hostname", "[your data]",
|
|
78
|
+
"--secret_arn", "[your data]",
|
|
79
|
+
"--database", "[your data]",
|
|
80
|
+
"--region", "[your data]",
|
|
81
|
+
"--readonly", "True"
|
|
82
|
+
],
|
|
83
|
+
"env": {
|
|
84
|
+
"AWS_PROFILE": "your-aws-profile",
|
|
85
|
+
"AWS_REGION": "us-east-1",
|
|
86
|
+
"FASTMCP_LOG_LEVEL": "ERROR"
|
|
87
|
+
},
|
|
88
|
+
"disabled": false,
|
|
89
|
+
"autoApprove": []
|
|
90
|
+
}
|
|
91
|
+
}
|
|
92
|
+
}
|
|
93
|
+
```
|
|
94
|
+
|
|
95
|
+
Note: The `--port` parameter is optional and defaults to 3306 (the standard MySQL port). You only need to specify it if your MySQL instance uses a non-default port.
|
|
96
|
+
|
|
55
97
|
### Windows Installation
|
|
56
98
|
|
|
57
99
|
For Windows users, the MCP server configuration format is slightly different:
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
|
|
15
|
+
"""aws.mysql-mcp-server.connection"""
|
|
16
|
+
|
|
17
|
+
from awslabs.mysql_mcp_server.connection.db_connection_singleton import DBConnectionSingleton
|
|
18
|
+
from awslabs.mysql_mcp_server.connection.abstract_db_connection import AbstractDBConnection
|
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
|
|
15
|
+
"""Abstract database connection interface for MySQL MCP Server."""
|
|
16
|
+
|
|
17
|
+
from abc import ABC, abstractmethod
|
|
18
|
+
from typing import Any, Dict, List, Optional
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class AbstractDBConnection(ABC):
|
|
22
|
+
"""Abstract base class for database connections."""
|
|
23
|
+
|
|
24
|
+
def __init__(self, readonly: bool):
|
|
25
|
+
"""Initialize the database connection.
|
|
26
|
+
|
|
27
|
+
Args:
|
|
28
|
+
readonly: Whether the connection should be read-only
|
|
29
|
+
"""
|
|
30
|
+
self._readonly = readonly
|
|
31
|
+
|
|
32
|
+
@property
|
|
33
|
+
def readonly_query(self) -> bool:
|
|
34
|
+
"""Get whether this connection is read-only.
|
|
35
|
+
|
|
36
|
+
Returns:
|
|
37
|
+
bool: True if the connection is read-only, False otherwise
|
|
38
|
+
"""
|
|
39
|
+
return self._readonly
|
|
40
|
+
|
|
41
|
+
@abstractmethod
|
|
42
|
+
async def execute_query(
|
|
43
|
+
self, sql: str, parameters: Optional[List[Dict[str, Any]]] = None
|
|
44
|
+
) -> Dict[str, Any]:
|
|
45
|
+
"""Execute a SQL query.
|
|
46
|
+
|
|
47
|
+
Args:
|
|
48
|
+
sql: The SQL query to execute
|
|
49
|
+
parameters: Optional parameters for the query
|
|
50
|
+
|
|
51
|
+
Returns:
|
|
52
|
+
Dict containing query results with column metadata and records
|
|
53
|
+
"""
|
|
54
|
+
pass
|
|
55
|
+
|
|
56
|
+
@abstractmethod
|
|
57
|
+
async def close(self) -> None:
|
|
58
|
+
"""Close the database connection."""
|
|
59
|
+
pass
|
|
60
|
+
|
|
61
|
+
@abstractmethod
|
|
62
|
+
async def check_connection_health(self) -> bool:
|
|
63
|
+
"""Check if the database connection is healthy.
|
|
64
|
+
|
|
65
|
+
Returns:
|
|
66
|
+
bool: True if the connection is healthy, False otherwise
|
|
67
|
+
"""
|
|
68
|
+
pass
|
|
@@ -0,0 +1,267 @@
|
|
|
1
|
+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
|
|
15
|
+
"""Python connector for MySQL and MariaDB MCP Server.
|
|
16
|
+
|
|
17
|
+
This connector provides direct connection to MySQL/MariaDB databases using asyncmy.
|
|
18
|
+
It supports both Aurora MySQL and RDS Mysql/MariaDB instances via direct connection.
|
|
19
|
+
"""
|
|
20
|
+
|
|
21
|
+
import boto3
|
|
22
|
+
import json
|
|
23
|
+
from asyncmy import Pool, create_pool
|
|
24
|
+
from awslabs.mysql_mcp_server.connection.abstract_db_connection import AbstractDBConnection
|
|
25
|
+
from loguru import logger
|
|
26
|
+
from typing import Any, Dict, List, Optional, Tuple
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
class AsyncmyPoolConnection(AbstractDBConnection):
|
|
30
|
+
"""Class that wraps DB connection using asyncmy connection pool.
|
|
31
|
+
|
|
32
|
+
This class can connect directly to any MySQL and MariaDB database, including:
|
|
33
|
+
- Aurora MySQL (using the cluster endpoint)
|
|
34
|
+
- RDS MySQL and RDS MariaDB (using the instance endpoint)
|
|
35
|
+
- Self-hosted MySQL and MariaDB
|
|
36
|
+
|
|
37
|
+
It uses AWS Secrets Manager (secret_arn and region) for authentication.
|
|
38
|
+
"""
|
|
39
|
+
|
|
40
|
+
def __init__(
|
|
41
|
+
self,
|
|
42
|
+
hostname: str,
|
|
43
|
+
port: int,
|
|
44
|
+
database: str,
|
|
45
|
+
readonly: bool,
|
|
46
|
+
secret_arn: str,
|
|
47
|
+
region: str,
|
|
48
|
+
min_size: int = 1,
|
|
49
|
+
max_size: int = 10,
|
|
50
|
+
):
|
|
51
|
+
"""Initialize a new DB connection pool.
|
|
52
|
+
|
|
53
|
+
Args:
|
|
54
|
+
hostname: Database host (Aurora cluster endpoint or RDS instance endpoint)
|
|
55
|
+
port: Database port (default 3306)
|
|
56
|
+
database: Database name
|
|
57
|
+
readonly: Whether connections should be read-only
|
|
58
|
+
secret_arn: ARN of the secret containing credentials
|
|
59
|
+
region: AWS region for Secrets Manager
|
|
60
|
+
min_size: Minimum number of connections in the pool
|
|
61
|
+
max_size: Maximum number of connections in the pool
|
|
62
|
+
"""
|
|
63
|
+
super().__init__(readonly)
|
|
64
|
+
self.hostname = hostname
|
|
65
|
+
self.port = port
|
|
66
|
+
self.min_size = min_size
|
|
67
|
+
self.max_size = max_size
|
|
68
|
+
self.pool: Optional[Pool] = None
|
|
69
|
+
self.database = database
|
|
70
|
+
|
|
71
|
+
# Get credentials from Secrets Manager
|
|
72
|
+
logger.info(f'Retrieving credentials from Secrets Manager: {secret_arn}')
|
|
73
|
+
self.user, self.password = _get_credentials_from_secret(secret_arn, region)
|
|
74
|
+
logger.info(f'Successfully retrieved credentials for user: {self.user}')
|
|
75
|
+
|
|
76
|
+
async def initialize_pool(self):
|
|
77
|
+
"""Initialize the connection pool."""
|
|
78
|
+
if self.pool is None:
|
|
79
|
+
logger.info(
|
|
80
|
+
f'Initializing connection pool with min_size={self.min_size}, max_size={self.max_size}'
|
|
81
|
+
)
|
|
82
|
+
|
|
83
|
+
self.pool = await create_pool(
|
|
84
|
+
minsize=self.min_size,
|
|
85
|
+
maxsize=self.max_size,
|
|
86
|
+
host=self.hostname,
|
|
87
|
+
port=self.port,
|
|
88
|
+
user=self.user,
|
|
89
|
+
password=self.password,
|
|
90
|
+
db=self.database,
|
|
91
|
+
autocommit=True,
|
|
92
|
+
)
|
|
93
|
+
|
|
94
|
+
logger.info('Connection pool initialized successfully')
|
|
95
|
+
|
|
96
|
+
if self._readonly:
|
|
97
|
+
await self._set_all_connections_readonly()
|
|
98
|
+
|
|
99
|
+
async def _set_all_connections_readonly(self):
|
|
100
|
+
"""Set all connections in the pool to read-only mode."""
|
|
101
|
+
if self.pool is None:
|
|
102
|
+
logger.warning('Connection pool is not initialized, cannot set read-only mode')
|
|
103
|
+
return
|
|
104
|
+
|
|
105
|
+
try:
|
|
106
|
+
async with self.pool.acquire() as conn:
|
|
107
|
+
async with conn.cursor() as cursor:
|
|
108
|
+
await cursor.execute('SET SESSION TRANSACTION READ ONLY;')
|
|
109
|
+
logger.info('Successfully set connection to read-only mode')
|
|
110
|
+
except Exception as e:
|
|
111
|
+
logger.warning(f'Failed to set connections to read-only mode: {str(e)}')
|
|
112
|
+
logger.warning('Continuing without setting read-only mode')
|
|
113
|
+
|
|
114
|
+
async def _get_connection(self):
|
|
115
|
+
"""Get a database connection from the pool."""
|
|
116
|
+
if self.pool is None:
|
|
117
|
+
await self.initialize_pool()
|
|
118
|
+
|
|
119
|
+
if self.pool is None:
|
|
120
|
+
raise ValueError('Failed to initialize connection pool')
|
|
121
|
+
|
|
122
|
+
return self.pool.acquire()
|
|
123
|
+
|
|
124
|
+
async def execute_query(
|
|
125
|
+
self, sql: str, parameters: Optional[List[Dict[str, Any]]] = None
|
|
126
|
+
) -> Dict[str, Any]:
|
|
127
|
+
"""Execute a SQL query using async connection."""
|
|
128
|
+
try:
|
|
129
|
+
async with await self._get_connection() as conn:
|
|
130
|
+
async with conn.cursor() as cursor:
|
|
131
|
+
if self._readonly:
|
|
132
|
+
await cursor.execute('SET TRANSACTION READ ONLY')
|
|
133
|
+
# Execute the query
|
|
134
|
+
if parameters:
|
|
135
|
+
params = list(_convert_parameters(self, parameters).values())
|
|
136
|
+
await cursor.execute(sql, params)
|
|
137
|
+
else:
|
|
138
|
+
await cursor.execute(sql)
|
|
139
|
+
|
|
140
|
+
# Check if there are results to fetch by examining the cursor's description
|
|
141
|
+
if cursor.description:
|
|
142
|
+
# Get column names
|
|
143
|
+
columns = [desc[0] for desc in cursor.description]
|
|
144
|
+
|
|
145
|
+
# Fetch all rows
|
|
146
|
+
rows = await cursor.fetchall()
|
|
147
|
+
|
|
148
|
+
# Structure the response to match the interface contract required by server.py
|
|
149
|
+
column_metadata = [{'label': col} for col in columns]
|
|
150
|
+
records = []
|
|
151
|
+
|
|
152
|
+
# Convert each row to the expected format
|
|
153
|
+
for row in rows:
|
|
154
|
+
record = []
|
|
155
|
+
for value in row:
|
|
156
|
+
if value is None:
|
|
157
|
+
record.append({'isNull': True})
|
|
158
|
+
elif isinstance(value, str):
|
|
159
|
+
record.append({'stringValue': value})
|
|
160
|
+
elif isinstance(value, int):
|
|
161
|
+
record.append({'longValue': value})
|
|
162
|
+
elif isinstance(value, float):
|
|
163
|
+
record.append({'doubleValue': value})
|
|
164
|
+
elif isinstance(value, bool):
|
|
165
|
+
record.append({'booleanValue': value})
|
|
166
|
+
elif isinstance(value, bytes):
|
|
167
|
+
record.append({'blobValue': value})
|
|
168
|
+
else:
|
|
169
|
+
# Convert other types to string
|
|
170
|
+
record.append({'stringValue': str(value)})
|
|
171
|
+
records.append(record)
|
|
172
|
+
|
|
173
|
+
return {'columnMetadata': column_metadata, 'records': records}
|
|
174
|
+
else:
|
|
175
|
+
# No results (e.g., for INSERT, UPDATE, etc.)
|
|
176
|
+
return {'columnMetadata': [], 'records': []}
|
|
177
|
+
|
|
178
|
+
except Exception as e:
|
|
179
|
+
logger.error(f'Database connection error: {str(e)}')
|
|
180
|
+
raise e
|
|
181
|
+
|
|
182
|
+
async def close(self) -> None:
|
|
183
|
+
"""Close all connections in the pool."""
|
|
184
|
+
if self.pool is not None:
|
|
185
|
+
logger.info('Closing connection pool')
|
|
186
|
+
await self.pool.close()
|
|
187
|
+
self.pool = None
|
|
188
|
+
logger.info('Connection pool closed successfully')
|
|
189
|
+
|
|
190
|
+
async def check_connection_health(self) -> bool:
|
|
191
|
+
"""Check if the connection is healthy."""
|
|
192
|
+
try:
|
|
193
|
+
result = await self.execute_query('SELECT 1')
|
|
194
|
+
return len(result.get('records', [])) > 0
|
|
195
|
+
except Exception as e:
|
|
196
|
+
logger.error(f'Connection health check failed: {str(e)}')
|
|
197
|
+
return False
|
|
198
|
+
|
|
199
|
+
|
|
200
|
+
def _convert_parameters(self, parameters: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
201
|
+
"""Transform structured parameter format to psycopg's native parameter format."""
|
|
202
|
+
result = {}
|
|
203
|
+
for param in parameters:
|
|
204
|
+
name = param.get('name')
|
|
205
|
+
value = param.get('value', {})
|
|
206
|
+
|
|
207
|
+
# Extract the value based on its type
|
|
208
|
+
if 'stringValue' in value:
|
|
209
|
+
result[name] = value['stringValue']
|
|
210
|
+
elif 'longValue' in value:
|
|
211
|
+
result[name] = value['longValue']
|
|
212
|
+
elif 'doubleValue' in value:
|
|
213
|
+
result[name] = value['doubleValue']
|
|
214
|
+
elif 'booleanValue' in value:
|
|
215
|
+
result[name] = value['booleanValue']
|
|
216
|
+
elif 'blobValue' in value:
|
|
217
|
+
result[name] = value['blobValue']
|
|
218
|
+
elif 'isNull' in value and value['isNull']:
|
|
219
|
+
result[name] = None
|
|
220
|
+
|
|
221
|
+
return result
|
|
222
|
+
|
|
223
|
+
|
|
224
|
+
def _get_credentials_from_secret(secret_arn: str, region: str) -> Tuple[str, str]:
|
|
225
|
+
"""Get database credentials from AWS Secrets Manager."""
|
|
226
|
+
try:
|
|
227
|
+
# Create a Secrets Manager client
|
|
228
|
+
logger.info(f'Creating Secrets Manager client in region {region}')
|
|
229
|
+
session = boto3.Session()
|
|
230
|
+
client = session.client(service_name='secretsmanager', region_name=region)
|
|
231
|
+
|
|
232
|
+
# Get the secret value
|
|
233
|
+
logger.info(f'Retrieving secret value for {secret_arn}')
|
|
234
|
+
get_secret_value_response = client.get_secret_value(SecretId=secret_arn)
|
|
235
|
+
logger.info('Successfully retrieved secret value')
|
|
236
|
+
|
|
237
|
+
# Parse the secret string
|
|
238
|
+
if 'SecretString' in get_secret_value_response:
|
|
239
|
+
secret = json.loads(get_secret_value_response['SecretString'])
|
|
240
|
+
logger.info(f'Secret keys: {", ".join(secret.keys())}')
|
|
241
|
+
|
|
242
|
+
# Extract username and password
|
|
243
|
+
username = secret.get('username') or secret.get('user') or secret.get('Username')
|
|
244
|
+
password = secret.get('password') or secret.get('Password')
|
|
245
|
+
|
|
246
|
+
if not username:
|
|
247
|
+
logger.error(
|
|
248
|
+
f'Username not found in secret. Available keys: {", ".join(secret.keys())}'
|
|
249
|
+
)
|
|
250
|
+
raise ValueError(
|
|
251
|
+
f'Secret does not contain username. Available keys: {", ".join(secret.keys())}'
|
|
252
|
+
)
|
|
253
|
+
|
|
254
|
+
if not password:
|
|
255
|
+
logger.error('Password not found in secret')
|
|
256
|
+
raise ValueError(
|
|
257
|
+
f'Secret does not contain password. Available keys: {", ".join(secret.keys())}'
|
|
258
|
+
)
|
|
259
|
+
|
|
260
|
+
logger.info(f'Successfully extracted credentials for user: {username}')
|
|
261
|
+
return username, password
|
|
262
|
+
else:
|
|
263
|
+
logger.error('Secret does not contain a SecretString')
|
|
264
|
+
raise ValueError('Secret does not contain a SecretString')
|
|
265
|
+
except Exception as e:
|
|
266
|
+
logger.error(f'Error retrieving secret: {str(e)}')
|
|
267
|
+
raise ValueError(f'Failed to retrieve credentials from Secrets Manager: {str(e)}')
|
|
@@ -0,0 +1,142 @@
|
|
|
1
|
+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
|
|
15
|
+
"""Database connection singleton for MySQL and MariaDB MCP Server."""
|
|
16
|
+
|
|
17
|
+
from awslabs.mysql_mcp_server.connection.asyncmy_pool_connection import AsyncmyPoolConnection
|
|
18
|
+
from awslabs.mysql_mcp_server.connection.rds_data_api_connection import RDSDataAPIConnection
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class DBConnectionSingleton:
|
|
22
|
+
"""Manages a single database connection instance across the application."""
|
|
23
|
+
|
|
24
|
+
_instance = None
|
|
25
|
+
|
|
26
|
+
def __init__(
|
|
27
|
+
self,
|
|
28
|
+
secret_arn: str,
|
|
29
|
+
database: str,
|
|
30
|
+
region: str,
|
|
31
|
+
readonly: bool = True,
|
|
32
|
+
is_test: bool = False,
|
|
33
|
+
resource_arn: str | None = None,
|
|
34
|
+
hostname: str | None = None,
|
|
35
|
+
port: int | None = None,
|
|
36
|
+
):
|
|
37
|
+
"""Initialize a new DB connection singleton using one of the two connection types.
|
|
38
|
+
|
|
39
|
+
1. RDS Data API Connection by specifying resource ARN
|
|
40
|
+
2. Direct MySQL or MariaDB connection by specifying hostname and port
|
|
41
|
+
|
|
42
|
+
Args:
|
|
43
|
+
secret_arn: The ARN of the secret containing credentials
|
|
44
|
+
database: The name of the database to connect to
|
|
45
|
+
region: The AWS region where the RDS instance is located
|
|
46
|
+
readonly: Whether the connection should be read-only (default: True)
|
|
47
|
+
resource_arn: The ARN of the RDS cluster (for using RDS Data API)
|
|
48
|
+
hostname: Database hostname (for using direct MySQL connection)
|
|
49
|
+
port: Database port (for using direct MySQL connection)
|
|
50
|
+
is_test: Whether this is a test connection (default: False)
|
|
51
|
+
"""
|
|
52
|
+
if resource_arn:
|
|
53
|
+
if not all([resource_arn, secret_arn, database, region]):
|
|
54
|
+
raise ValueError(
|
|
55
|
+
'Missing required connection parameters for RDS Data API. '
|
|
56
|
+
'Please provide resource_arn, secret_arn, database, and region.'
|
|
57
|
+
)
|
|
58
|
+
|
|
59
|
+
self._db_connection = RDSDataAPIConnection(
|
|
60
|
+
cluster_arn=resource_arn,
|
|
61
|
+
secret_arn=secret_arn,
|
|
62
|
+
database=database,
|
|
63
|
+
region=region,
|
|
64
|
+
readonly=readonly,
|
|
65
|
+
is_test=is_test,
|
|
66
|
+
)
|
|
67
|
+
else:
|
|
68
|
+
# Direct connection to MySQL/MariaDB
|
|
69
|
+
if not all([hostname, port, secret_arn, database, region]):
|
|
70
|
+
raise ValueError(
|
|
71
|
+
'Missing required connection parameters for direct MySQL connection. '
|
|
72
|
+
'Please provide hostname, port, secret_arn, database, and region.'
|
|
73
|
+
)
|
|
74
|
+
assert hostname is not None
|
|
75
|
+
assert port is not None
|
|
76
|
+
|
|
77
|
+
self._db_connection = AsyncmyPoolConnection(
|
|
78
|
+
hostname=hostname,
|
|
79
|
+
port=port,
|
|
80
|
+
secret_arn=secret_arn,
|
|
81
|
+
database=database,
|
|
82
|
+
region=region,
|
|
83
|
+
readonly=readonly,
|
|
84
|
+
)
|
|
85
|
+
|
|
86
|
+
@classmethod
|
|
87
|
+
def initialize(
|
|
88
|
+
cls,
|
|
89
|
+
secret_arn: str,
|
|
90
|
+
database: str,
|
|
91
|
+
region: str,
|
|
92
|
+
readonly: bool = True,
|
|
93
|
+
is_test: bool = False,
|
|
94
|
+
resource_arn: str | None = None,
|
|
95
|
+
hostname: str | None = None,
|
|
96
|
+
port: int | None = None,
|
|
97
|
+
):
|
|
98
|
+
"""Initialize the singleton instance if it doesn't exist.
|
|
99
|
+
|
|
100
|
+
Args:
|
|
101
|
+
resource_arn: The ARN of the RDS cluster (for using RDS Data API)
|
|
102
|
+
hostname: Database hostname (for using direct MySQL/MariaDB connection)
|
|
103
|
+
port: Database port (for using direct MySQL/MariaDB connection)
|
|
104
|
+
secret_arn: The ARN of the secret containing credentials
|
|
105
|
+
database: The name of the database to connect to
|
|
106
|
+
region: The AWS region where the RDS instance is located
|
|
107
|
+
readonly: Whether the connection should be read-only (default: True)
|
|
108
|
+
is_test: Whether this is a test connection (default: False)
|
|
109
|
+
"""
|
|
110
|
+
if cls._instance is None:
|
|
111
|
+
cls._instance = cls(
|
|
112
|
+
secret_arn=secret_arn,
|
|
113
|
+
database=database,
|
|
114
|
+
region=region,
|
|
115
|
+
readonly=readonly,
|
|
116
|
+
resource_arn=resource_arn,
|
|
117
|
+
hostname=hostname,
|
|
118
|
+
port=port,
|
|
119
|
+
is_test=is_test,
|
|
120
|
+
)
|
|
121
|
+
|
|
122
|
+
@classmethod
|
|
123
|
+
def get(cls):
|
|
124
|
+
"""Get the singleton instance.
|
|
125
|
+
|
|
126
|
+
Returns:
|
|
127
|
+
DBConnectionSingleton: The singleton instance
|
|
128
|
+
Raises:
|
|
129
|
+
RuntimeError: If the singleton has not been initialized
|
|
130
|
+
"""
|
|
131
|
+
if cls._instance is None:
|
|
132
|
+
raise RuntimeError('DBConnectionSingleton is not initialized.')
|
|
133
|
+
return cls._instance
|
|
134
|
+
|
|
135
|
+
@property
|
|
136
|
+
def db_connection(self):
|
|
137
|
+
"""Get the database connection.
|
|
138
|
+
|
|
139
|
+
Returns:
|
|
140
|
+
DBConnection: The database connection instance
|
|
141
|
+
"""
|
|
142
|
+
return self._db_connection
|