awslabs.redshift-mcp-server 0.0.1__py3-none-any.whl

awslabs/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+
+# This file is part of the awslabs namespace.
+# It is intentionally minimal to support PEP 420 namespace packages.

awslabs/redshift_mcp_server/__init__.py

@@ -0,0 +1,17 @@
+# 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.
+
+"""awslabs.redshift-mcp-server"""
+
+__version__ = '0.0.1'

awslabs/redshift_mcp_server/consts.py

@@ -0,0 +1,136 @@
+# 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.
+
+"""Redshift MCP Server constants."""
+
+# Defaults
+DEFAULT_AWS_REGION = 'us-east-1'
+DEFAULT_LOG_LEVEL = 'WARNING'
+
+# Timeouts (seconds), etc
+CLIENT_TIMEOUT = 60
+DATA_CLIENT_TIMEOUT = 60
+QUERY_TIMEOUT = 3600
+QUERY_POLL_INTERVAL = 2
+
+# Best practices
+
+CLIENT_BEST_PRACTICES = """
+## AWS Client Best Practices
+
+### Authentication
+
+- Default AWS credentials chain (IAM roles, ~/.aws/credentials, etc.).
+- AWS_PROFILE environment variable (if set).
+- AWS_REGION environment variable (if set).
+"""
+
+REDSHIFT_BEST_PRACTICES = """
+## Amazon Redshift Best Practices
+
+### Query Guidelines
+
+- Always specify the database and schema when referencing objects to avoid ambiguity.
+- Leverage distribution in WHERE and JOIN predicates and sort keys in ORDER BY for optimal query performance.
+- Use LIMIT clauses for exploratory queries to avoid large result sets.
+- Analyze table to update table statistics if it is not updated or too off before making a decision on the query structure.
+- Prefer explicitly specifying columns in SELECT over "*" for better performance.
+
+### Connection Guidelines
+
+- We are use the Redshift API and Redshift Data API.
+- Leverage IAM authentication when possible instead of secrets (database passwords).
+"""
+
+# SQL queries
+
+SVV_REDSHIFT_DATABASES_QUERY = """
+SELECT
+    database_name,
+    database_owner,
+    database_type,
+    database_acl,
+    database_options,
+    database_isolation_level
+FROM pg_catalog.svv_redshift_databases
+ORDER BY database_name;
+"""
+
+SVV_ALL_SCHEMAS_QUERY = """
+SELECT
+    database_name,
+    schema_name,
+    schema_owner,
+    schema_type,
+    schema_acl,
+    source_database,
+    schema_option
+FROM pg_catalog.svv_all_schemas
+WHERE database_name = {}
+ORDER BY schema_name;
+"""
+
+SVV_ALL_TABLES_QUERY = """
+SELECT
+    database_name,
+    schema_name,
+    table_name,
+    table_acl,
+    table_type,
+    remarks
+FROM pg_catalog.svv_all_tables
+WHERE database_name = {} AND schema_name = {}
+ORDER BY table_name;
+"""
+
+SVV_ALL_COLUMNS_QUERY = """
+SELECT
+    database_name,
+    schema_name,
+    table_name,
+    column_name,
+    ordinal_position,
+    column_default,
+    is_nullable,
+    data_type,
+    character_maximum_length,
+    numeric_precision,
+    numeric_scale,
+    remarks
+FROM pg_catalog.svv_all_columns
+WHERE database_name = {} AND schema_name = {} AND table_name = {}
+ORDER BY ordinal_position;
+"""
+
+# SQL guardrails
+
+# Single-lines comments.
+re_slc = r'--.*?$'
+
+
+def re_mlc(g: int) -> str:
+    """Multi-line comments, considering balanced recursion."""
+    return rf'(?P<mlc{g}>(?:\/\*)(?:[^\/\*]|\/[^\*]|\*[^\/]|(?P>mlc{g}))*(?:\*\/))'
+
+
+def re_sp(g: int) -> str:
+    """Whitespaces, comments, semicolons which can occur between words."""
+    return rf'({re_slc}|{re_mlc(g)}|\s|;)'
+
+
+# We consider `(END|COMMIT|ROLLBACK|ABORT) [WORK|TRANSACTION]` as a breaker for the `BEGIN READ ONLY; {sql}; END;`
+# guarding wrapper, having there might be variations of whitespaces and comments in the construct.
+SUSPICIOUS_QUERY_REGEXP = (
+    rf'(?im)^{re_sp(1)}*(END|COMMIT|ROLLBACK|ABORT)({re_sp(2)}+(WORK|TRANSACTION))?{re_sp(3)}*;'
+)

awslabs/redshift_mcp_server/models.py

@@ -0,0 +1,141 @@
+# 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.
+
+"""Redshift MCP Server Pydantic models."""
+
+from datetime import datetime
+from pydantic import BaseModel, Field
+from typing import Dict, Optional
+
+
+class RedshiftCluster(BaseModel):
+    """Information about a Redshift cluster or serverless workgroup."""
+
+    identifier: str = Field(..., description='Unique identifier for the cluster/workgroup')
+    type: str = Field(..., description='Type of cluster (provisioned or serverless)')
+    status: str = Field(..., description='Current status of the cluster')
+    database_name: str = Field(..., description='Default database name')
+    endpoint: Optional[str] = Field(None, description='Connection endpoint')
+    port: Optional[int] = Field(None, description='Connection port')
+    vpc_id: Optional[str] = Field(None, description='VPC ID where the cluster resides')
+    node_type: Optional[str] = Field(None, description='Node type (provisioned only)')
+    number_of_nodes: Optional[int] = Field(None, description='Number of nodes (provisioned only)')
+    creation_time: Optional[datetime] = Field(None, description='When the cluster was created')
+    master_username: Optional[str] = Field(None, description='Master username for the cluster')
+    publicly_accessible: Optional[bool] = Field(None, description='Whether publicly accessible')
+    encrypted: Optional[bool] = Field(None, description='Whether the cluster is encrypted')
+    tags: Optional[Dict[str, str]] = Field(
+        default_factory=dict, description='Tags associated with the cluster'
+    )
+
+
+class RedshiftDatabase(BaseModel):
+    """Information about a database in a Redshift cluster.
+
+    Based on the SVV_REDSHIFT_DATABASES system view.
+    """
+
+    database_name: str = Field(..., description='The name of the database')
+    database_owner: Optional[int] = Field(None, description='The database owner user ID')
+    database_type: Optional[str] = Field(
+        None, description='The type of database (local or shared)'
+    )
+    database_acl: Optional[str] = Field(
+        None, description='Access control information (for internal use)'
+    )
+    database_options: Optional[str] = Field(None, description='The properties of the database')
+    database_isolation_level: Optional[str] = Field(
+        None,
+        description='The isolation level of the database (Snapshot Isolation or Serializable)',
+    )
+
+
+class RedshiftSchema(BaseModel):
+    """Information about a schema in a Redshift database.
+
+    Based on the SVV_ALL_SCHEMAS system view.
+    """
+
+    database_name: str = Field(..., description='The name of the database where the schema exists')
+    schema_name: str = Field(..., description='The name of the schema')
+    schema_owner: Optional[int] = Field(None, description='The user ID of the schema owner')
+    schema_type: Optional[str] = Field(
+        None, description='The type of the schema (external, local, or shared)'
+    )
+    schema_acl: Optional[str] = Field(
+        None, description='The permissions for the specified user or user group for the schema'
+    )
+    source_database: Optional[str] = Field(
+        None, description='The name of the source database for external schema'
+    )
+    schema_option: Optional[str] = Field(
+        None, description='The options of the schema (external schema attribute)'
+    )
+
+
+class RedshiftTable(BaseModel):
+    """Information about a table in a Redshift database.
+
+    Based on the SVV_ALL_TABLES system view.
+    """
+
+    database_name: str = Field(..., description='The name of the database where the table exists')
+    schema_name: str = Field(..., description='The schema name for the table')
+    table_name: str = Field(..., description='The name of the table')
+    table_acl: Optional[str] = Field(
+        None, description='The permissions for the specified user or user group for the table'
+    )
+    table_type: Optional[str] = Field(
+        None,
+        description='The type of the table (views, base tables, external tables, shared tables)',
+    )
+    remarks: Optional[str] = Field(None, description='Remarks about the table')
+
+
+class RedshiftColumn(BaseModel):
+    """Information about a column in a Redshift table.
+
+    Based on the SVV_ALL_COLUMNS system view.
+    """
+
+    database_name: str = Field(..., description='The name of the database')
+    schema_name: str = Field(..., description='The name of the schema')
+    table_name: str = Field(..., description='The name of the table')
+    column_name: str = Field(..., description='The name of the column')
+    ordinal_position: Optional[int] = Field(
+        None, description='The position of the column in the table'
+    )
+    column_default: Optional[str] = Field(None, description='The default value of the column')
+    is_nullable: Optional[str] = Field(
+        None, description='Whether the column is nullable (yes or no)'
+    )
+    data_type: Optional[str] = Field(None, description='The data type of the column')
+    character_maximum_length: Optional[int] = Field(
+        None, description='The maximum number of characters in the column'
+    )
+    numeric_precision: Optional[int] = Field(None, description='The numeric precision')
+    numeric_scale: Optional[int] = Field(None, description='The numeric scale')
+    remarks: Optional[str] = Field(None, description='Remarks about the column')
+
+
+class QueryResult(BaseModel):
+    """Result of a SQL query execution."""
+
+    columns: list[str] = Field(..., description='List of column names in the result set')
+    rows: list[list] = Field(..., description='List of rows, where each row is a list of values')
+    row_count: int = Field(..., description='Number of rows returned')
+    execution_time_ms: Optional[int] = Field(
+        None, description='Query execution time in milliseconds'
+    )
+    query_id: str = Field(..., description='Unique identifier for the query execution')