awslabs.redshift-mcp-server 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl
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/redshift_mcp_server/__init__.py +1 -1
- awslabs/redshift_mcp_server/consts.py +21 -12
- awslabs/redshift_mcp_server/redshift.py +270 -132
- awslabs/redshift_mcp_server/server.py +17 -6
- {awslabs_redshift_mcp_server-0.0.6.dist-info → awslabs_redshift_mcp_server-0.0.8.dist-info}/METADATA +13 -9
- awslabs_redshift_mcp_server-0.0.8.dist-info/RECORD +12 -0
- awslabs_redshift_mcp_server-0.0.6.dist-info/RECORD +0 -12
- {awslabs_redshift_mcp_server-0.0.6.dist-info → awslabs_redshift_mcp_server-0.0.8.dist-info}/WHEEL +0 -0
- {awslabs_redshift_mcp_server-0.0.6.dist-info → awslabs_redshift_mcp_server-0.0.8.dist-info}/entry_points.txt +0 -0
- {awslabs_redshift_mcp_server-0.0.6.dist-info → awslabs_redshift_mcp_server-0.0.8.dist-info}/licenses/LICENSE +0 -0
- {awslabs_redshift_mcp_server-0.0.6.dist-info → awslabs_redshift_mcp_server-0.0.8.dist-info}/licenses/NOTICE +0 -0
|
@@ -14,26 +14,35 @@
|
|
|
14
14
|
|
|
15
15
|
"""Redshift MCP Server constants."""
|
|
16
16
|
|
|
17
|
-
#
|
|
18
|
-
|
|
17
|
+
# System
|
|
18
|
+
CLIENT_CONNECT_TIMEOUT = 60
|
|
19
|
+
CLIENT_READ_TIMEOUT = 600
|
|
20
|
+
CLIENT_RETRIES = {'max_attempts': 5, 'mode': 'adaptive'}
|
|
21
|
+
CLIENT_USER_AGENT_NAME = 'awslabs/mcp/redshift-mcp-server'
|
|
19
22
|
DEFAULT_LOG_LEVEL = 'WARNING'
|
|
20
|
-
|
|
21
|
-
# Timeouts (seconds), etc
|
|
22
|
-
CLIENT_TIMEOUT = 60
|
|
23
|
-
DATA_CLIENT_TIMEOUT = 60
|
|
24
23
|
QUERY_TIMEOUT = 3600
|
|
25
|
-
QUERY_POLL_INTERVAL =
|
|
24
|
+
QUERY_POLL_INTERVAL = 1
|
|
25
|
+
SESSION_KEEPALIVE = 600
|
|
26
26
|
|
|
27
27
|
# Best practices
|
|
28
28
|
|
|
29
29
|
CLIENT_BEST_PRACTICES = """
|
|
30
30
|
## AWS Client Best Practices
|
|
31
31
|
|
|
32
|
-
### Authentication
|
|
32
|
+
### Authentication and Configuration
|
|
33
33
|
|
|
34
34
|
- Default AWS credentials chain (IAM roles, ~/.aws/credentials, etc.).
|
|
35
35
|
- AWS_PROFILE environment variable (if set).
|
|
36
|
-
-
|
|
36
|
+
- Region configuration (in order of precedence):
|
|
37
|
+
- AWS_REGION environment variable (highest priority)
|
|
38
|
+
- AWS_DEFAULT_REGION environment variable
|
|
39
|
+
- Region specified in AWS profile configuration
|
|
40
|
+
|
|
41
|
+
### Error Handling
|
|
42
|
+
|
|
43
|
+
- Always print out AWS client errors in full to help diagnose configuration issues.
|
|
44
|
+
- For region-related errors, suggest checking AWS_REGION, AWS_DEFAULT_REGION, or AWS profile configuration.
|
|
45
|
+
- For credential errors, suggest verifying AWS credentials setup and permissions.
|
|
37
46
|
"""
|
|
38
47
|
|
|
39
48
|
REDSHIFT_BEST_PRACTICES = """
|
|
@@ -77,7 +86,7 @@ SELECT
|
|
|
77
86
|
source_database,
|
|
78
87
|
schema_option
|
|
79
88
|
FROM pg_catalog.svv_all_schemas
|
|
80
|
-
WHERE database_name =
|
|
89
|
+
WHERE database_name = :database_name
|
|
81
90
|
ORDER BY schema_name;
|
|
82
91
|
"""
|
|
83
92
|
|
|
@@ -90,7 +99,7 @@ SELECT
|
|
|
90
99
|
table_type,
|
|
91
100
|
remarks
|
|
92
101
|
FROM pg_catalog.svv_all_tables
|
|
93
|
-
WHERE database_name =
|
|
102
|
+
WHERE database_name = :database_name AND schema_name = :schema_name
|
|
94
103
|
ORDER BY table_name;
|
|
95
104
|
"""
|
|
96
105
|
|
|
@@ -109,7 +118,7 @@ SELECT
|
|
|
109
118
|
numeric_scale,
|
|
110
119
|
remarks
|
|
111
120
|
FROM pg_catalog.svv_all_columns
|
|
112
|
-
WHERE database_name =
|
|
121
|
+
WHERE database_name = :database_name AND schema_name = :schema_name AND table_name = :table_name
|
|
113
122
|
ORDER BY ordinal_position;
|
|
114
123
|
"""
|
|
115
124
|
|
|
@@ -18,12 +18,16 @@ import asyncio
|
|
|
18
18
|
import boto3
|
|
19
19
|
import os
|
|
20
20
|
import regex
|
|
21
|
+
import time
|
|
21
22
|
from awslabs.redshift_mcp_server import __version__
|
|
22
23
|
from awslabs.redshift_mcp_server.consts import (
|
|
23
|
-
|
|
24
|
-
|
|
24
|
+
CLIENT_CONNECT_TIMEOUT,
|
|
25
|
+
CLIENT_READ_TIMEOUT,
|
|
26
|
+
CLIENT_RETRIES,
|
|
27
|
+
CLIENT_USER_AGENT_NAME,
|
|
25
28
|
QUERY_POLL_INTERVAL,
|
|
26
29
|
QUERY_TIMEOUT,
|
|
30
|
+
SESSION_KEEPALIVE,
|
|
27
31
|
SUSPICIOUS_QUERY_REGEXP,
|
|
28
32
|
SVV_ALL_COLUMNS_QUERY,
|
|
29
33
|
SVV_ALL_SCHEMAS_QUERY,
|
|
@@ -37,7 +41,9 @@ from loguru import logger
|
|
|
37
41
|
class RedshiftClientManager:
|
|
38
42
|
"""Manages AWS clients for Redshift operations."""
|
|
39
43
|
|
|
40
|
-
def __init__(
|
|
44
|
+
def __init__(
|
|
45
|
+
self, config: Config, aws_region: str | None = None, aws_profile: str | None = None
|
|
46
|
+
):
|
|
41
47
|
"""Initialize the client manager."""
|
|
42
48
|
self.aws_region = aws_region
|
|
43
49
|
self.aws_profile = aws_profile
|
|
@@ -50,15 +56,12 @@ class RedshiftClientManager:
|
|
|
50
56
|
"""Get or create the Redshift client for provisioned clusters."""
|
|
51
57
|
if self._redshift_client is None:
|
|
52
58
|
try:
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
|
|
59
|
-
'redshift', config=self._config, region_name=self.aws_region
|
|
60
|
-
)
|
|
61
|
-
logger.info('Created Redshift client with default credentials')
|
|
59
|
+
# Session works with None values - uses default credentials/region chain
|
|
60
|
+
session = boto3.Session(profile_name=self.aws_profile, region_name=self.aws_region)
|
|
61
|
+
self._redshift_client = session.client('redshift', config=self._config)
|
|
62
|
+
logger.info(
|
|
63
|
+
f'Created Redshift client with profile: {self.aws_profile or "default"}, region: {self.aws_region or "default"}'
|
|
64
|
+
)
|
|
62
65
|
except Exception as e:
|
|
63
66
|
logger.error(f'Error creating Redshift client: {str(e)}')
|
|
64
67
|
raise
|
|
@@ -69,19 +72,14 @@ class RedshiftClientManager:
|
|
|
69
72
|
"""Get or create the Redshift Serverless client."""
|
|
70
73
|
if self._redshift_serverless_client is None:
|
|
71
74
|
try:
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
else:
|
|
81
|
-
self._redshift_serverless_client = boto3.client(
|
|
82
|
-
'redshift-serverless', config=self._config, region_name=self.aws_region
|
|
83
|
-
)
|
|
84
|
-
logger.info('Created Redshift Serverless client with default credentials')
|
|
75
|
+
# Session works with None values - uses default credentials/region chain
|
|
76
|
+
session = boto3.Session(profile_name=self.aws_profile, region_name=self.aws_region)
|
|
77
|
+
self._redshift_serverless_client = session.client(
|
|
78
|
+
'redshift-serverless', config=self._config
|
|
79
|
+
)
|
|
80
|
+
logger.info(
|
|
81
|
+
f'Created Redshift Serverless client with profile: {self.aws_profile or "default"}, region: {self.aws_region or "default"}'
|
|
82
|
+
)
|
|
85
83
|
except Exception as e:
|
|
86
84
|
logger.error(f'Error creating Redshift Serverless client: {str(e)}')
|
|
87
85
|
raise
|
|
@@ -92,19 +90,12 @@ class RedshiftClientManager:
|
|
|
92
90
|
"""Get or create the Redshift Data API client."""
|
|
93
91
|
if self._redshift_data_client is None:
|
|
94
92
|
try:
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
f'Created Redshift Data API client with profile: {self.aws_profile}'
|
|
102
|
-
)
|
|
103
|
-
else:
|
|
104
|
-
self._redshift_data_client = boto3.client(
|
|
105
|
-
'redshift-data', config=self._config, region_name=self.aws_region
|
|
106
|
-
)
|
|
107
|
-
logger.info('Created Redshift Data API client with default credentials')
|
|
93
|
+
# Session works with None values - uses default credentials/region chain
|
|
94
|
+
session = boto3.Session(profile_name=self.aws_profile, region_name=self.aws_region)
|
|
95
|
+
self._redshift_data_client = session.client('redshift-data', config=self._config)
|
|
96
|
+
logger.info(
|
|
97
|
+
f'Created Redshift Data API client with profile: {self.aws_profile or "default"}, region: {self.aws_region or "default"}'
|
|
98
|
+
)
|
|
108
99
|
except Exception as e:
|
|
109
100
|
logger.error(f'Error creating Redshift Data API client: {str(e)}')
|
|
110
101
|
raise
|
|
@@ -112,61 +103,124 @@ class RedshiftClientManager:
|
|
|
112
103
|
return self._redshift_data_client
|
|
113
104
|
|
|
114
105
|
|
|
115
|
-
|
|
116
|
-
"""
|
|
106
|
+
class RedshiftSessionManager:
|
|
107
|
+
"""Manages Redshift Data API sessions for connection reuse."""
|
|
108
|
+
|
|
109
|
+
def __init__(self, session_keepalive: int, app_name: str):
|
|
110
|
+
"""Initialize the session manager.
|
|
111
|
+
|
|
112
|
+
Args:
|
|
113
|
+
session_keepalive: Session keepalive timeout in seconds.
|
|
114
|
+
app_name: Application name to set in sessions.
|
|
115
|
+
"""
|
|
116
|
+
self._sessions = {} # {cluster:database -> session_info}
|
|
117
|
+
self._session_keepalive = session_keepalive
|
|
118
|
+
self._app_name = app_name
|
|
119
|
+
|
|
120
|
+
async def session(
|
|
121
|
+
self, cluster_identifier: str, database_name: str, cluster_info: dict
|
|
122
|
+
) -> str:
|
|
123
|
+
"""Get or create a session for the given cluster and database.
|
|
124
|
+
|
|
125
|
+
Args:
|
|
126
|
+
cluster_identifier: The cluster identifier to get session for.
|
|
127
|
+
database_name: The database name to get session for.
|
|
128
|
+
cluster_info: Cluster information dictionary from discover_clusters.
|
|
129
|
+
|
|
130
|
+
Returns:
|
|
131
|
+
Session ID for use in ExecuteStatement calls.
|
|
132
|
+
"""
|
|
133
|
+
# Check existing session
|
|
134
|
+
session_key = f'{cluster_identifier}:{database_name}'
|
|
135
|
+
if session_key in self._sessions:
|
|
136
|
+
session_info = self._sessions[session_key]
|
|
137
|
+
if not self._is_session_expired(session_info):
|
|
138
|
+
logger.debug(f'Reusing existing session: {session_info["session_id"]}')
|
|
139
|
+
return session_info['session_id']
|
|
140
|
+
else:
|
|
141
|
+
logger.debug(f'Session expired, removing: {session_info["session_id"]}')
|
|
142
|
+
del self._sessions[session_key]
|
|
143
|
+
|
|
144
|
+
# Create new session with application name
|
|
145
|
+
session_id = await self._create_session_with_app_name(
|
|
146
|
+
cluster_identifier, database_name, cluster_info
|
|
147
|
+
)
|
|
117
148
|
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
"""
|
|
121
|
-
if value is None:
|
|
122
|
-
return 'NULL'
|
|
149
|
+
# Store session
|
|
150
|
+
self._sessions[session_key] = {'session_id': session_id, 'created_at': time.time()}
|
|
123
151
|
|
|
124
|
-
|
|
125
|
-
|
|
126
|
-
return "'" + repr('"' + value)[2:]
|
|
152
|
+
logger.info(f'Created new session: {session_id} for {cluster_identifier}:{database_name}')
|
|
153
|
+
return session_id
|
|
127
154
|
|
|
155
|
+
async def _create_session_with_app_name(
|
|
156
|
+
self, cluster_identifier: str, database_name: str, cluster_info: dict
|
|
157
|
+
) -> str:
|
|
158
|
+
"""Create a new session by executing SET application_name.
|
|
128
159
|
|
|
129
|
-
|
|
130
|
-
|
|
160
|
+
Args:
|
|
161
|
+
cluster_identifier: The cluster identifier.
|
|
162
|
+
database_name: The database name.
|
|
163
|
+
cluster_info: Cluster information dictionary.
|
|
131
164
|
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
|
|
165
|
+
Returns:
|
|
166
|
+
Session ID from the ExecuteStatement response.
|
|
167
|
+
"""
|
|
168
|
+
# Set application name to create session
|
|
169
|
+
app_name_sql = f"SET application_name TO '{self._app_name}';"
|
|
170
|
+
|
|
171
|
+
# Execute statement to create session
|
|
172
|
+
statement_id = await _execute_statement(
|
|
173
|
+
cluster_info=cluster_info,
|
|
174
|
+
cluster_identifier=cluster_identifier,
|
|
175
|
+
database_name=database_name,
|
|
176
|
+
sql=app_name_sql,
|
|
177
|
+
session_keepalive=self._session_keepalive,
|
|
178
|
+
)
|
|
135
179
|
|
|
136
|
-
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
180
|
+
# Get session ID from the response
|
|
181
|
+
data_client = client_manager.redshift_data_client()
|
|
182
|
+
status_response = data_client.describe_statement(Id=statement_id)
|
|
183
|
+
session_id = status_response['SessionId']
|
|
140
184
|
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
allow_read_write: Indicates if read-write mode should be activated.
|
|
185
|
+
logger.debug(f'Created session with application name: {session_id}')
|
|
186
|
+
return session_id
|
|
144
187
|
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
|
|
148
|
-
|
|
149
|
-
|
|
150
|
-
else:
|
|
151
|
-
# Check if SQL contains suspicious patterns trying to break the transaction context
|
|
152
|
-
if regex.compile(SUSPICIOUS_QUERY_REGEXP).search(sql):
|
|
153
|
-
logger.error(f'SQL contains suspicious pattern, execution rejected: {sql}')
|
|
154
|
-
raise Exception(f'SQL contains suspicious pattern, execution rejected: {sql}')
|
|
188
|
+
def _is_session_expired(self, session_info: dict) -> bool:
|
|
189
|
+
"""Check if a session has expired based on keepalive timeout.
|
|
190
|
+
|
|
191
|
+
Args:
|
|
192
|
+
session_info: Session information dictionary.
|
|
155
193
|
|
|
156
|
-
|
|
194
|
+
Returns:
|
|
195
|
+
True if session is expired, False otherwise.
|
|
196
|
+
"""
|
|
197
|
+
return (time.time() - session_info['created_at']) > self._session_keepalive
|
|
157
198
|
|
|
158
199
|
|
|
159
|
-
async def
|
|
160
|
-
cluster_identifier: str,
|
|
200
|
+
async def _execute_protected_statement(
|
|
201
|
+
cluster_identifier: str,
|
|
202
|
+
database_name: str,
|
|
203
|
+
sql: str,
|
|
204
|
+
parameters: list[dict] | None = None,
|
|
205
|
+
allow_read_write: bool = False,
|
|
161
206
|
) -> tuple[dict, str]:
|
|
162
|
-
"""Execute a SQL statement against a Redshift cluster
|
|
207
|
+
"""Execute a SQL statement against a Redshift cluster in a protected fashion.
|
|
208
|
+
|
|
209
|
+
The SQL is protected by wrapping it in a transaction block with READ ONLY or READ WRITE mode
|
|
210
|
+
based on allow_read_write flag. Transaction breaker protection is implemented
|
|
211
|
+
to prevent unauthorized modifications.
|
|
163
212
|
|
|
164
|
-
|
|
213
|
+
The SQL execution takes the form:
|
|
214
|
+
1. Get or create session (with SET application_name)
|
|
215
|
+
2. BEGIN [READ ONLY|READ WRITE];
|
|
216
|
+
3. <user sql>
|
|
217
|
+
4. END;
|
|
165
218
|
|
|
166
219
|
Args:
|
|
167
220
|
cluster_identifier: The cluster identifier to query.
|
|
168
221
|
database_name: The database to execute the query against.
|
|
169
222
|
sql: The SQL statement to execute.
|
|
223
|
+
parameters: Optional list of parameter dictionaries with 'name' and 'value' keys.
|
|
170
224
|
allow_read_write: Indicates if read-write mode should be activated.
|
|
171
225
|
|
|
172
226
|
Returns:
|
|
@@ -177,9 +231,7 @@ async def execute_statement(
|
|
|
177
231
|
Raises:
|
|
178
232
|
Exception: If cluster not found, query fails, or times out.
|
|
179
233
|
"""
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
# First, check if this is a provisioned cluster or serverless workgroup
|
|
234
|
+
# Get cluster info
|
|
183
235
|
clusters = await discover_clusters()
|
|
184
236
|
cluster_info = None
|
|
185
237
|
for cluster in clusters:
|
|
@@ -192,54 +244,131 @@ async def execute_statement(
|
|
|
192
244
|
f'Cluster {cluster_identifier} not found. Please use list_clusters to get valid cluster identifiers.'
|
|
193
245
|
)
|
|
194
246
|
|
|
195
|
-
#
|
|
196
|
-
|
|
197
|
-
logger.debug(f'Protected SQL: {" ".join(protected_sqls)}')
|
|
247
|
+
# Get session (creates if needed, sets app name automatically)
|
|
248
|
+
session_id = await session_manager.session(cluster_identifier, database_name, cluster_info)
|
|
198
249
|
|
|
199
|
-
#
|
|
200
|
-
if
|
|
201
|
-
|
|
202
|
-
|
|
203
|
-
|
|
204
|
-
|
|
205
|
-
|
|
206
|
-
|
|
207
|
-
|
|
208
|
-
|
|
209
|
-
|
|
210
|
-
|
|
211
|
-
|
|
250
|
+
# Check for suspicious patterns in read-only mode
|
|
251
|
+
if not allow_read_write:
|
|
252
|
+
if regex.compile(SUSPICIOUS_QUERY_REGEXP).search(sql):
|
|
253
|
+
logger.error(f'SQL contains suspicious pattern, execution rejected: {sql}')
|
|
254
|
+
raise Exception(f'SQL contains suspicious pattern, execution rejected: {sql}')
|
|
255
|
+
|
|
256
|
+
# Execute BEGIN statement
|
|
257
|
+
begin_sql = 'BEGIN READ WRITE;' if allow_read_write else 'BEGIN READ ONLY;'
|
|
258
|
+
await _execute_statement(
|
|
259
|
+
cluster_info=cluster_info,
|
|
260
|
+
cluster_identifier=cluster_identifier,
|
|
261
|
+
database_name=database_name,
|
|
262
|
+
sql=begin_sql,
|
|
263
|
+
session_id=session_id,
|
|
264
|
+
)
|
|
265
|
+
|
|
266
|
+
# Execute user SQL with parameters
|
|
267
|
+
user_query_id = await _execute_statement(
|
|
268
|
+
cluster_info=cluster_info,
|
|
269
|
+
cluster_identifier=cluster_identifier,
|
|
270
|
+
database_name=database_name,
|
|
271
|
+
sql=sql,
|
|
272
|
+
parameters=parameters,
|
|
273
|
+
session_id=session_id,
|
|
274
|
+
)
|
|
275
|
+
|
|
276
|
+
# Execute END statement to close transaction
|
|
277
|
+
await _execute_statement(
|
|
278
|
+
cluster_info=cluster_info,
|
|
279
|
+
cluster_identifier=cluster_identifier,
|
|
280
|
+
database_name=database_name,
|
|
281
|
+
sql='END;',
|
|
282
|
+
session_id=session_id,
|
|
283
|
+
)
|
|
284
|
+
|
|
285
|
+
# Get results from user query
|
|
286
|
+
data_client = client_manager.redshift_data_client()
|
|
287
|
+
results_response = data_client.get_statement_result(Id=user_query_id)
|
|
288
|
+
return results_response, user_query_id
|
|
212
289
|
|
|
213
|
-
query_id = response['Id']
|
|
214
|
-
logger.debug(f'Started query execution: {query_id}')
|
|
215
290
|
|
|
216
|
-
|
|
291
|
+
async def _execute_statement(
|
|
292
|
+
cluster_info: dict,
|
|
293
|
+
cluster_identifier: str,
|
|
294
|
+
database_name: str,
|
|
295
|
+
sql: str,
|
|
296
|
+
parameters: list[dict] | None = None,
|
|
297
|
+
session_id: str | None = None,
|
|
298
|
+
session_keepalive: int | None = None,
|
|
299
|
+
query_poll_interval: float = QUERY_POLL_INTERVAL,
|
|
300
|
+
query_timeout: float = QUERY_TIMEOUT,
|
|
301
|
+
) -> str:
|
|
302
|
+
"""Execute a single statement with optional session support and parameters.
|
|
303
|
+
|
|
304
|
+
Args:
|
|
305
|
+
cluster_info: Cluster information dictionary.
|
|
306
|
+
cluster_identifier: The cluster identifier.
|
|
307
|
+
database_name: The database name.
|
|
308
|
+
sql: The SQL statement to execute.
|
|
309
|
+
parameters: Optional list of parameter dictionaries with 'name' and 'value' keys.
|
|
310
|
+
session_id: Optional session ID to use.
|
|
311
|
+
session_keepalive: Optional session keepalive seconds (only used when session_id is None).
|
|
312
|
+
query_poll_interval: Polling interval in seconds for checking query status.
|
|
313
|
+
query_timeout: Maximum time in seconds to wait for query completion.
|
|
314
|
+
|
|
315
|
+
Returns:
|
|
316
|
+
Statement ID from the ExecuteStatement response.
|
|
317
|
+
"""
|
|
318
|
+
data_client = client_manager.redshift_data_client()
|
|
319
|
+
|
|
320
|
+
# Build request parameters
|
|
321
|
+
request_params: dict[str, str | int | list[dict]] = {'Sql': sql}
|
|
322
|
+
|
|
323
|
+
# Add database and cluster/workgroup identifier only if not using session
|
|
324
|
+
if not session_id:
|
|
325
|
+
request_params['Database'] = database_name
|
|
326
|
+
if cluster_info['type'] == 'provisioned':
|
|
327
|
+
request_params['ClusterIdentifier'] = cluster_identifier
|
|
328
|
+
elif cluster_info['type'] == 'serverless':
|
|
329
|
+
request_params['WorkgroupName'] = cluster_identifier
|
|
330
|
+
else:
|
|
331
|
+
raise Exception(f'Unknown cluster type: {cluster_info["type"]}')
|
|
332
|
+
|
|
333
|
+
# Add parameters if provided
|
|
334
|
+
if parameters:
|
|
335
|
+
request_params['Parameters'] = parameters
|
|
336
|
+
|
|
337
|
+
# Add session ID if provided, otherwise add session keepalive
|
|
338
|
+
if session_id:
|
|
339
|
+
request_params['SessionId'] = session_id
|
|
340
|
+
elif session_keepalive is not None:
|
|
341
|
+
request_params['SessionKeepAliveSeconds'] = session_keepalive
|
|
342
|
+
|
|
343
|
+
response = data_client.execute_statement(**request_params)
|
|
344
|
+
statement_id = response['Id']
|
|
345
|
+
|
|
346
|
+
logger.debug(
|
|
347
|
+
f'Executed statement: {statement_id}' + (f' in session {session_id}' if session_id else '')
|
|
348
|
+
)
|
|
349
|
+
|
|
350
|
+
# Wait for statement completion
|
|
217
351
|
wait_time = 0
|
|
218
|
-
|
|
219
|
-
|
|
220
|
-
status_response = data_client.describe_statement(Id=query_id)
|
|
352
|
+
while wait_time < query_timeout:
|
|
353
|
+
status_response = data_client.describe_statement(Id=statement_id)
|
|
221
354
|
status = status_response['Status']
|
|
222
355
|
|
|
223
356
|
if status == 'FINISHED':
|
|
224
|
-
logger.debug(f'
|
|
357
|
+
logger.debug(f'Statement completed: {statement_id}')
|
|
225
358
|
break
|
|
226
359
|
elif status in ['FAILED', 'ABORTED']:
|
|
227
360
|
error_msg = status_response.get('Error', 'Unknown error')
|
|
228
|
-
logger.error(f'
|
|
229
|
-
raise Exception(f'
|
|
361
|
+
logger.error(f'Statement failed: {error_msg}')
|
|
362
|
+
raise Exception(f'Statement failed: {error_msg}')
|
|
230
363
|
|
|
231
|
-
|
|
232
|
-
|
|
233
|
-
wait_time += QUERY_POLL_INTERVAL
|
|
364
|
+
await asyncio.sleep(query_poll_interval)
|
|
365
|
+
wait_time += query_poll_interval
|
|
234
366
|
|
|
235
|
-
if wait_time >=
|
|
236
|
-
logger.error(f'
|
|
237
|
-
raise Exception(f'
|
|
367
|
+
if wait_time >= query_timeout:
|
|
368
|
+
logger.error(f'Statement timed out: {statement_id}')
|
|
369
|
+
raise Exception(f'Statement timed out after {wait_time} seconds')
|
|
238
370
|
|
|
239
|
-
|
|
240
|
-
subquery1_id = status_response['SubStatements'][1]['Id']
|
|
241
|
-
results_response = data_client.get_statement_result(Id=subquery1_id)
|
|
242
|
-
return results_response, subquery1_id
|
|
371
|
+
return statement_id
|
|
243
372
|
|
|
244
373
|
|
|
245
374
|
async def discover_clusters() -> list[dict]:
|
|
@@ -342,7 +471,7 @@ async def discover_databases(cluster_identifier: str, database_name: str = 'dev'
|
|
|
342
471
|
logger.info(f'Discovering databases in cluster {cluster_identifier}')
|
|
343
472
|
|
|
344
473
|
# Execute the query using the common function
|
|
345
|
-
results_response, _ = await
|
|
474
|
+
results_response, _ = await _execute_protected_statement(
|
|
346
475
|
cluster_identifier=cluster_identifier,
|
|
347
476
|
database_name=database_name,
|
|
348
477
|
sql=SVV_REDSHIFT_DATABASES_QUERY,
|
|
@@ -387,10 +516,11 @@ async def discover_schemas(cluster_identifier: str, schema_database_name: str) -
|
|
|
387
516
|
)
|
|
388
517
|
|
|
389
518
|
# Execute the query using the common function
|
|
390
|
-
results_response, _ = await
|
|
519
|
+
results_response, _ = await _execute_protected_statement(
|
|
391
520
|
cluster_identifier=cluster_identifier,
|
|
392
521
|
database_name=schema_database_name,
|
|
393
|
-
sql=SVV_ALL_SCHEMAS_QUERY
|
|
522
|
+
sql=SVV_ALL_SCHEMAS_QUERY,
|
|
523
|
+
parameters=[{'name': 'database_name', 'value': schema_database_name}],
|
|
394
524
|
)
|
|
395
525
|
|
|
396
526
|
schemas = []
|
|
@@ -440,12 +570,14 @@ async def discover_tables(
|
|
|
440
570
|
)
|
|
441
571
|
|
|
442
572
|
# Execute the query using the common function
|
|
443
|
-
results_response, _ = await
|
|
573
|
+
results_response, _ = await _execute_protected_statement(
|
|
444
574
|
cluster_identifier=cluster_identifier,
|
|
445
575
|
database_name=table_database_name,
|
|
446
|
-
sql=SVV_ALL_TABLES_QUERY
|
|
447
|
-
|
|
448
|
-
|
|
576
|
+
sql=SVV_ALL_TABLES_QUERY,
|
|
577
|
+
parameters=[
|
|
578
|
+
{'name': 'database_name', 'value': table_database_name},
|
|
579
|
+
{'name': 'schema_name', 'value': table_schema_name},
|
|
580
|
+
],
|
|
449
581
|
)
|
|
450
582
|
|
|
451
583
|
tables = []
|
|
@@ -498,14 +630,15 @@ async def discover_columns(
|
|
|
498
630
|
)
|
|
499
631
|
|
|
500
632
|
# Execute the query using the common function
|
|
501
|
-
results_response, _ = await
|
|
633
|
+
results_response, _ = await _execute_protected_statement(
|
|
502
634
|
cluster_identifier=cluster_identifier,
|
|
503
635
|
database_name=column_database_name,
|
|
504
|
-
sql=SVV_ALL_COLUMNS_QUERY
|
|
505
|
-
|
|
506
|
-
|
|
507
|
-
|
|
508
|
-
|
|
636
|
+
sql=SVV_ALL_COLUMNS_QUERY,
|
|
637
|
+
parameters=[
|
|
638
|
+
{'name': 'database_name', 'value': column_database_name},
|
|
639
|
+
{'name': 'schema_name', 'value': column_schema_name},
|
|
640
|
+
{'name': 'table_name', 'value': column_table_name},
|
|
641
|
+
],
|
|
509
642
|
)
|
|
510
643
|
|
|
511
644
|
columns = []
|
|
@@ -562,7 +695,7 @@ async def execute_query(cluster_identifier: str, database_name: str, sql: str) -
|
|
|
562
695
|
start_time = time.time()
|
|
563
696
|
|
|
564
697
|
# Execute the query using the common function
|
|
565
|
-
results_response, query_id = await
|
|
698
|
+
results_response, query_id = await _execute_protected_statement(
|
|
566
699
|
cluster_identifier=cluster_identifier, database_name=database_name, sql=sql
|
|
567
700
|
)
|
|
568
701
|
|
|
@@ -620,11 +753,16 @@ async def execute_query(cluster_identifier: str, database_name: str, sql: str) -
|
|
|
620
753
|
# Global client manager instance
|
|
621
754
|
client_manager = RedshiftClientManager(
|
|
622
755
|
config=Config(
|
|
623
|
-
connect_timeout=
|
|
624
|
-
read_timeout=
|
|
625
|
-
retries=
|
|
626
|
-
user_agent_extra=f'
|
|
756
|
+
connect_timeout=CLIENT_CONNECT_TIMEOUT,
|
|
757
|
+
read_timeout=CLIENT_READ_TIMEOUT,
|
|
758
|
+
retries=CLIENT_RETRIES,
|
|
759
|
+
user_agent_extra=f'{CLIENT_USER_AGENT_NAME}/{__version__}',
|
|
627
760
|
),
|
|
628
|
-
aws_region=os.environ.get('AWS_REGION'
|
|
761
|
+
aws_region=os.environ.get('AWS_REGION'),
|
|
629
762
|
aws_profile=os.environ.get('AWS_PROFILE'),
|
|
630
763
|
)
|
|
764
|
+
|
|
765
|
+
# Global session manager instance
|
|
766
|
+
session_manager = RedshiftSessionManager(
|
|
767
|
+
session_keepalive=SESSION_KEEPALIVE, app_name=f'{CLIENT_USER_AGENT_NAME}/{__version__}'
|
|
768
|
+
)
|
|
@@ -85,7 +85,7 @@ This tool uses the Redshift Data API to run queries and return results.
|
|
|
85
85
|
|
|
86
86
|
## Getting Started
|
|
87
87
|
|
|
88
|
-
1. Ensure your AWS credentials are configured (
|
|
88
|
+
1. Ensure your AWS configuration and credentials are configured (environment variables or profile configuration file).
|
|
89
89
|
2. Use the list_clusters tool to discover available Redshift instances.
|
|
90
90
|
3. Note the cluster identifiers for use with other tools (coming in future milestones).
|
|
91
91
|
|
|
@@ -219,7 +219,9 @@ async def list_databases_tool(
|
|
|
219
219
|
"""
|
|
220
220
|
try:
|
|
221
221
|
logger.info(f'Discovering databases on cluster: {cluster_identifier}')
|
|
222
|
-
databases_data = await discover_databases(
|
|
222
|
+
databases_data = await discover_databases(
|
|
223
|
+
cluster_identifier=cluster_identifier, database_name=database_name
|
|
224
|
+
)
|
|
223
225
|
|
|
224
226
|
# Convert to RedshiftDatabase models
|
|
225
227
|
databases = []
|
|
@@ -302,7 +304,9 @@ async def list_schemas_tool(
|
|
|
302
304
|
logger.info(
|
|
303
305
|
f'Discovering schemas in database {schema_database_name} on cluster {cluster_identifier}'
|
|
304
306
|
)
|
|
305
|
-
schemas_data = await discover_schemas(
|
|
307
|
+
schemas_data = await discover_schemas(
|
|
308
|
+
cluster_identifier=cluster_identifier, schema_database_name=schema_database_name
|
|
309
|
+
)
|
|
306
310
|
|
|
307
311
|
# Convert to RedshiftSchema models
|
|
308
312
|
schemas = []
|
|
@@ -394,7 +398,9 @@ async def list_tables_tool(
|
|
|
394
398
|
f'Discovering tables in schema {table_schema_name} in database {table_database_name} on cluster {cluster_identifier}'
|
|
395
399
|
)
|
|
396
400
|
tables_data = await discover_tables(
|
|
397
|
-
cluster_identifier,
|
|
401
|
+
cluster_identifier=cluster_identifier,
|
|
402
|
+
table_database_name=table_database_name,
|
|
403
|
+
table_schema_name=table_schema_name,
|
|
398
404
|
)
|
|
399
405
|
|
|
400
406
|
# Convert to RedshiftTable models
|
|
@@ -500,7 +506,10 @@ async def list_columns_tool(
|
|
|
500
506
|
f'Discovering columns in table {column_table_name} in schema {column_schema_name} in database {column_database_name} on cluster {cluster_identifier}'
|
|
501
507
|
)
|
|
502
508
|
columns_data = await discover_columns(
|
|
503
|
-
cluster_identifier,
|
|
509
|
+
cluster_identifier=cluster_identifier,
|
|
510
|
+
column_database_name=column_database_name,
|
|
511
|
+
column_schema_name=column_schema_name,
|
|
512
|
+
column_table_name=column_table_name,
|
|
504
513
|
)
|
|
505
514
|
|
|
506
515
|
# Convert to RedshiftColumn models
|
|
@@ -594,7 +603,9 @@ async def execute_query_tool(
|
|
|
594
603
|
"""
|
|
595
604
|
try:
|
|
596
605
|
logger.info(f'Executing query on cluster {cluster_identifier} in database {database_name}')
|
|
597
|
-
query_result_data = await execute_query(
|
|
606
|
+
query_result_data = await execute_query(
|
|
607
|
+
cluster_identifier=cluster_identifier, database_name=database_name, sql=sql
|
|
608
|
+
)
|
|
598
609
|
|
|
599
610
|
# Convert to QueryResult model
|
|
600
611
|
query_result = QueryResult(**query_result_data)
|
{awslabs_redshift_mcp_server-0.0.6.dist-info → awslabs_redshift_mcp_server-0.0.8.dist-info}/METADATA
RENAMED
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: awslabs.redshift-mcp-server
|
|
3
|
-
Version: 0.0.
|
|
3
|
+
Version: 0.0.8
|
|
4
4
|
Summary: An AWS Labs Model Context Protocol (MCP) server for Redshift
|
|
5
5
|
Project-URL: homepage, https://awslabs.github.io/mcp/
|
|
6
6
|
Project-URL: docs, https://awslabs.github.io/mcp/servers/redshift-mcp-server/
|
|
@@ -52,7 +52,11 @@ This MCP server provides tools to discover, explore, and query Amazon Redshift c
|
|
|
52
52
|
### AWS Client Requirements
|
|
53
53
|
|
|
54
54
|
1. **Credentials**: Configure AWS credentials via AWS CLI, or environment variables
|
|
55
|
-
2. **
|
|
55
|
+
2. **Region**: Configure AWS region using one of the following (in order of precedence):
|
|
56
|
+
- `AWS_REGION` environment variable (highest priority)
|
|
57
|
+
- `AWS_DEFAULT_REGION` environment variable
|
|
58
|
+
- Region specified in your AWS profile configuration
|
|
59
|
+
3. **Permissions**: Ensure your AWS credentials have the required permissions (see [Permissions](#permissions) section)
|
|
56
60
|
|
|
57
61
|
## Installation
|
|
58
62
|
|
|
@@ -70,7 +74,7 @@ Configure the MCP server in your MCP client configuration (e.g., for Amazon Q De
|
|
|
70
74
|
"args": ["awslabs.redshift-mcp-server@latest"],
|
|
71
75
|
"env": {
|
|
72
76
|
"AWS_PROFILE": "default",
|
|
73
|
-
"
|
|
77
|
+
"AWS_DEFAULT_REGION": "us-east-1",
|
|
74
78
|
"FASTMCP_LOG_LEVEL": "INFO"
|
|
75
79
|
},
|
|
76
80
|
"disabled": false,
|
|
@@ -79,6 +83,7 @@ Configure the MCP server in your MCP client configuration (e.g., for Amazon Q De
|
|
|
79
83
|
}
|
|
80
84
|
}
|
|
81
85
|
```
|
|
86
|
+
|
|
82
87
|
### Windows Installation
|
|
83
88
|
|
|
84
89
|
For Windows users, the MCP server configuration format is slightly different:
|
|
@@ -99,16 +104,15 @@ For Windows users, the MCP server configuration format is slightly different:
|
|
|
99
104
|
"awslabs.redshift-mcp-server.exe"
|
|
100
105
|
],
|
|
101
106
|
"env": {
|
|
102
|
-
"FASTMCP_LOG_LEVEL": "ERROR",
|
|
103
107
|
"AWS_PROFILE": "your-aws-profile",
|
|
104
|
-
"
|
|
108
|
+
"AWS_DEFAULT_REGION": "us-east-1",
|
|
109
|
+
"FASTMCP_LOG_LEVEL": "ERROR"
|
|
105
110
|
}
|
|
106
111
|
}
|
|
107
112
|
}
|
|
108
113
|
}
|
|
109
114
|
```
|
|
110
115
|
|
|
111
|
-
|
|
112
116
|
or docker after a successful `docker build -t awslabs/redshift-mcp-server:latest .`:
|
|
113
117
|
|
|
114
118
|
```json
|
|
@@ -122,7 +126,7 @@ or docker after a successful `docker build -t awslabs/redshift-mcp-server:latest
|
|
|
122
126
|
"--interactive",
|
|
123
127
|
"--env", "AWS_ACCESS_KEY_ID=[your data]",
|
|
124
128
|
"--env", "AWS_SECRET_ACCESS_KEY=[your data]",
|
|
125
|
-
"--env", "
|
|
129
|
+
"--env", "AWS_DEFAULT_REGION=[your data]",
|
|
126
130
|
"awslabs/redshift-mcp-server:latest"
|
|
127
131
|
]
|
|
128
132
|
}
|
|
@@ -132,7 +136,8 @@ or docker after a successful `docker build -t awslabs/redshift-mcp-server:latest
|
|
|
132
136
|
|
|
133
137
|
### Environment Variables
|
|
134
138
|
|
|
135
|
-
- `AWS_REGION`: AWS region to use (
|
|
139
|
+
- `AWS_REGION`: AWS region to use (overrides all other region settings)
|
|
140
|
+
- `AWS_DEFAULT_REGION`: Default AWS region (used if AWS_REGION not set and no region in profile)
|
|
136
141
|
- `AWS_PROFILE`: AWS profile to use (optional, uses default if not specified)
|
|
137
142
|
- `FASTMCP_LOG_LEVEL`: Logging level (`DEBUG`, `INFO`, `WARNING`, `ERROR`)
|
|
138
143
|
- `LOG_FILE`: Path to log file (optional, logs to stdout if not specified)
|
|
@@ -444,7 +449,6 @@ Your AWS credentials need the following IAM permissions:
|
|
|
444
449
|
"redshift-serverless:ListWorkgroups",
|
|
445
450
|
"redshift-serverless:GetWorkgroup",
|
|
446
451
|
"redshift-data:ExecuteStatement",
|
|
447
|
-
"redshift-data:BatchExecuteStatement",
|
|
448
452
|
"redshift-data:DescribeStatement",
|
|
449
453
|
"redshift-data:GetStatementResult"
|
|
450
454
|
],
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
awslabs/__init__.py,sha256=WuqxdDgUZylWNmVoPKiK7qGsTB_G4UmuXIrJ-VBwDew,731
|
|
2
|
+
awslabs/redshift_mcp_server/__init__.py,sha256=h53d7Qk7HjCGqRe4pn0vhr4rBIlvYwav_vEZYkB4RsQ,673
|
|
3
|
+
awslabs/redshift_mcp_server/consts.py,sha256=Ou07Q81KMsSw1AYKrU4sawCjbeD2tovQCnYVc5CPWyk,4302
|
|
4
|
+
awslabs/redshift_mcp_server/models.py,sha256=p6oKcVz4xfaqQzXjJrZK9YlfuUnzMDCphXbTK1LT1k4,6437
|
|
5
|
+
awslabs/redshift_mcp_server/redshift.py,sha256=ljmLLXXUpAN-wkpxu9rNNOCqYMCCLvGXWNVKDzOTMfA,29821
|
|
6
|
+
awslabs/redshift_mcp_server/server.py,sha256=msHWX_g8g82kM705wgWjkeK0vVEou_Cz1emEI7VuDPM,26503
|
|
7
|
+
awslabs_redshift_mcp_server-0.0.8.dist-info/METADATA,sha256=tGwiCpH4UWcUxiWCEHw-EiYdl5Wkv6ZMKmy18sD9Ak4,16523
|
|
8
|
+
awslabs_redshift_mcp_server-0.0.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
|
|
9
|
+
awslabs_redshift_mcp_server-0.0.8.dist-info/entry_points.txt,sha256=o2G-onpmq80KMTQw2OrY1G07GmwPaurcJcIqfvMR9Sw,88
|
|
10
|
+
awslabs_redshift_mcp_server-0.0.8.dist-info/licenses/LICENSE,sha256=CeipvOyAZxBGUsFoaFqwkx54aPnIKEtm9a5u2uXxEws,10142
|
|
11
|
+
awslabs_redshift_mcp_server-0.0.8.dist-info/licenses/NOTICE,sha256=iIvfV8gFGERQ7xLtxV8bD_Lsrj9KOIPpvk49qp5-K0c,95
|
|
12
|
+
awslabs_redshift_mcp_server-0.0.8.dist-info/RECORD,,
|
|
@@ -1,12 +0,0 @@
|
|
|
1
|
-
awslabs/__init__.py,sha256=WuqxdDgUZylWNmVoPKiK7qGsTB_G4UmuXIrJ-VBwDew,731
|
|
2
|
-
awslabs/redshift_mcp_server/__init__.py,sha256=dSiLO85AiX9WDWinzOaxhO-qdzqjs6VvfOnDmD18-ec,673
|
|
3
|
-
awslabs/redshift_mcp_server/consts.py,sha256=_r7BEcNEjW8bEOxMwcb6ScCv3w9aXR0vRmd7HzAokyI,3683
|
|
4
|
-
awslabs/redshift_mcp_server/models.py,sha256=p6oKcVz4xfaqQzXjJrZK9YlfuUnzMDCphXbTK1LT1k4,6437
|
|
5
|
-
awslabs/redshift_mcp_server/redshift.py,sha256=FfzCPi5njzUvehNURHqAaHZr9Lz60Z9bAFXveGHLZR4,24696
|
|
6
|
-
awslabs/redshift_mcp_server/server.py,sha256=UNlBqgYS8KLhT9fFKKIPbszH1MtFTY5xCWbED_h-CJ8,26100
|
|
7
|
-
awslabs_redshift_mcp_server-0.0.6.dist-info/METADATA,sha256=gIGrr8hnNi5DpDW3Ye-tLUtSTqnPsmjPupyt82rN4Mc,16184
|
|
8
|
-
awslabs_redshift_mcp_server-0.0.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
|
|
9
|
-
awslabs_redshift_mcp_server-0.0.6.dist-info/entry_points.txt,sha256=o2G-onpmq80KMTQw2OrY1G07GmwPaurcJcIqfvMR9Sw,88
|
|
10
|
-
awslabs_redshift_mcp_server-0.0.6.dist-info/licenses/LICENSE,sha256=CeipvOyAZxBGUsFoaFqwkx54aPnIKEtm9a5u2uXxEws,10142
|
|
11
|
-
awslabs_redshift_mcp_server-0.0.6.dist-info/licenses/NOTICE,sha256=iIvfV8gFGERQ7xLtxV8bD_Lsrj9KOIPpvk49qp5-K0c,95
|
|
12
|
-
awslabs_redshift_mcp_server-0.0.6.dist-info/RECORD,,
|
{awslabs_redshift_mcp_server-0.0.6.dist-info → awslabs_redshift_mcp_server-0.0.8.dist-info}/WHEEL
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|