awslabs.s3-tables-mcp-server 0.0.1__py3-none-any.whl

awslabs/__init__.py

@@ -0,0 +1,15 @@
+# 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 Labs MCP package."""

awslabs/s3_tables_mcp_server/__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.
+
+# This file is part of the awslabs namespace.
+# It is intentionally minimal to support PEP 420 namespace packages.
+
+__version__ = '0.0.0'

awslabs/s3_tables_mcp_server/constants.py

@@ -0,0 +1,167 @@
+# 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.
+
+"""Constants used throughout the S3 Tables MCP Server.
+
+This module contains all the constant values used across the S3 Tables MCP Server,
+including version information, regex patterns for validation, and field definitions
+for Pydantic models.
+"""
+
+from pydantic import Field
+
+
+# Patterns
+TABLE_BUCKET_NAME_PATTERN = r'[a-z0-9][a-z0-9-]{1,61}[a-z0-9]'
+"""
+Regex pattern for validating S3 bucket names.
+Valid bucket names must:
+- Be between 3 and 63 characters long
+- Start and end with a letter or number
+- Contain only lowercase letters, numbers, and hyphens
+- Not contain consecutive hyphens
+"""
+
+TABLE_BUCKET_ARN_PATTERN = (
+    r'arn:aws[-a-z0-9]*:[a-z0-9]+:[-a-z0-9]*:[0-9]{12}:bucket/[a-z0-9_-]{3,63}'
+)
+"""
+Regex pattern for validating S3 bucket ARNs.
+Format: arn:aws[-a-z0-9]*:[a-z0-9]+:[-a-z0-9]*:[0-9]{12}:bucket/[bucket-name]
+Example: arn:aws:s3:::my-bucket
+"""
+
+TABLE_NAME_PATTERN = r'[0-9a-z_]*'
+"""
+Regex pattern for validating table names.
+Valid table names must:
+- Contain only lowercase letters, numbers, and underscores
+- Have a maximum length of 255 characters
+"""
+
+TABLE_ARN_PATTERN = (
+    r'arn:aws[-a-z0-9]*:[a-z0-9]+:[-a-z0-9]*:[0-9]{12}:bucket/[a-z0-9_-]{3,63}/table/[0-9a-f-]{36}'
+)
+"""
+Regex pattern for validating table ARNs.
+Format: arn:aws[-a-z0-9]*:[a-z0-9]+:[-a-z0-9]*:[0-9]{12}:bucket/[bucket-name]/table/[uuid]
+Example: arn:aws:s3:::my-bucket/table/123e4567-e89b-12d3-a456-426614174000
+"""
+
+# Field Definitions
+TABLE_BUCKET_ARN_FIELD = Field(
+    ...,
+    description='Table bucket ARN',
+    pattern=TABLE_BUCKET_ARN_PATTERN,
+    min_length=1,
+    max_length=2048,
+)
+"""
+Pydantic field for table bucket ARN validation.
+Required field that must match the TABLE_BUCKET_ARN_PATTERN.
+"""
+
+TABLE_ARN_FIELD = Field(..., description='Table ARN', pattern=TABLE_ARN_PATTERN)
+"""
+Pydantic field for table ARN validation.
+Required field that must match the TABLE_ARN_PATTERN.
+"""
+
+NAMESPACE_NAME_FIELD = Field(
+    ...,
+    description='The name of the namespace. Must be 1-255 characters long and contain only alphanumeric characters, underscores, and hyphens.',
+    min_length=1,
+    max_length=255,
+    pattern=r'^[a-zA-Z0-9_-]+$',
+)
+"""
+Pydantic field for namespace name validation.
+Required field that must:
+- Be 1-255 characters long
+- Contain only alphanumeric characters, underscores, and hyphens
+"""
+
+TABLE_NAME_FIELD = Field(
+    ...,
+    description='The name of the table. Must be 1-255 characters long and contain only alphanumeric characters, underscores, and hyphens.',
+    min_length=1,
+    max_length=255,
+    pattern=TABLE_NAME_PATTERN,
+)
+"""
+Pydantic field for table name validation.
+Required field that must:
+- Be 1-255 characters long
+- Contain only alphanumeric characters, underscores, and hyphens
+- Match the TABLE_NAME_PATTERN
+"""
+
+REGION_NAME_FIELD = Field(
+    None,
+    description='The AWS region name where the operation should be performed.',
+    min_length=1,
+    max_length=64,
+)
+"""
+Pydantic field for AWS region name.
+Optional field that can be used to specify the AWS region for operations.
+Example values: 'us-east-1', 'eu-west-1', 'ap-southeast-2'
+"""
+
+# Query-specific fields
+QUERY_FIELD = Field(
+    default=None,
+    description='Optional SQL query. If not provided, will execute SELECT * FROM table. Must be a read operation.',
+    min_length=1,
+    max_length=10000,
+)
+"""
+Pydantic field for SQL query validation.
+Optional field that must be a valid read operation.
+"""
+
+OUTPUT_LOCATION_FIELD = Field(
+    default=None,
+    description='Optional S3 location for query results. If not provided, will use default Athena results bucket.',
+    pattern=r'^s3://[a-z0-9-]+/[a-z0-9-./]*$',
+    min_length=1,
+    max_length=2048,
+)
+"""
+Pydantic field for output location validation.
+Optional field that must be a valid S3 URI.
+"""
+
+WORKGROUP_FIELD = Field(
+    default='primary',
+    description='Athena workgroup to use for query execution.',
+    pattern=r'^[a-zA-Z0-9_-]+$',
+    min_length=1,
+    max_length=128,
+)
+"""
+Pydantic field for workgroup validation.
+Optional field that must contain only letters, numbers, hyphens, and underscores.
+Defaults to 'primary'.
+"""
+
+S3_URL_FIELD = Field(
+    ...,
+    description='The S3 URL of the file to preview (format: s3://bucket-name/key)',
+    min_length=1,
+)
+"""
+Pydantic field for S3 URL validation.
+Required field that must be a valid S3 URI.
+"""

awslabs/s3_tables_mcp_server/database.py

@@ -0,0 +1,140 @@
+# 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 query operations for S3 Tables MCP Server.
+
+This module provides functions for executing queries against S3 Tables using Athena.
+It handles query execution, result retrieval, and proper formatting of responses.
+"""
+
+import sqlparse
+from .engines.pyiceberg import PyIcebergConfig, PyIcebergEngine
+from typing import Any, Dict
+
+
+WRITE_OPERATIONS = {
+    'ADD',
+    'ALTER',
+    'ANALYZE',
+    'BEGIN',
+    'COMMIT',
+    'COPY',
+    'CREATE',
+    'DELETE',
+    'DROP',
+    'EXPORT',
+    'GRANT',
+    'IMPORT',
+    'INSERT',
+    'LOAD',
+    'LOCK',
+    'MERGE',
+    'MSCK',
+    'REDUCE',
+    'REFRESH',
+    'REPLACE',
+    'RESET',
+    'REVOKE',
+    'ROLLBACK',
+    'SET',
+    'START',
+    'TRUNCATE',
+    'UNCACHE',
+    'UNLOCK',
+    'UPDATE',
+    'UPSERT',
+    'VACUUM',
+    'VALUES',
+    'WRITE',
+}
+
+READ_OPERATIONS = {
+    'DESC',
+    'DESCRIBE',
+    'EXPLAIN',
+    'LIST',
+    'SELECT',
+    'SHOW',
+    'USE',
+}
+
+# Disallowed destructive operations for write
+DESTRUCTIVE_OPERATIONS = {'DELETE', 'DROP', 'MERGE', 'REPLACE', 'TRUNCATE', 'VACUUM'}
+
+
+def _get_query_operations(query: str) -> set:
+    """Extract all top-level SQL operations from the query as a set."""
+    parsed = sqlparse.parse(query)
+    operations = set()
+    for stmt in parsed:
+        tokens = [token.value.upper() for token in stmt.tokens if not token.is_whitespace]
+        for token in tokens:
+            if token.isalpha():
+                operations.add(token)
+    return operations
+
+
+async def query_database_resource(
+    warehouse: str,
+    region: str,
+    namespace: str,
+    query: str,
+    uri: str = 'https://s3tables.us-west-2.amazonaws.com/iceberg',
+    catalog_name: str = 's3tablescatalog',
+    rest_signing_name: str = 's3tables',
+    rest_sigv4_enabled: str = 'true',
+) -> Dict[str, Any]:
+    """Execute a read-only query against a database using PyIceberg."""
+    operations = _get_query_operations(query)
+    disallowed = operations & WRITE_OPERATIONS
+    if disallowed:
+        raise ValueError(f'Write operations are not allowed in read-only queries: {disallowed}')
+    config = PyIcebergConfig(
+        warehouse=warehouse,
+        uri=uri,
+        region=region,
+        namespace=namespace,
+        catalog_name=catalog_name,
+        rest_signing_name=rest_signing_name,
+        rest_sigv4_enabled=rest_sigv4_enabled,
+    )
+    engine = PyIcebergEngine(config)
+    result = engine.execute_query(query)
+    return result
+
+
+async def append_rows_to_table_resource(
+    warehouse: str,
+    region: str,
+    namespace: str,
+    table_name: str,
+    rows: list,
+    uri: str = 'https://s3tables.us-west-2.amazonaws.com/iceberg',
+    catalog_name: str = 's3tablescatalog',
+    rest_signing_name: str = 's3tables',
+    rest_sigv4_enabled: str = 'true',
+) -> Dict[str, Any]:
+    """Append rows to an Iceberg table using PyIceberg."""
+    config = PyIcebergConfig(
+        warehouse=warehouse,
+        uri=uri,
+        region=region,
+        namespace=namespace,
+        catalog_name=catalog_name,
+        rest_signing_name=rest_signing_name,
+        rest_sigv4_enabled=rest_sigv4_enabled,
+    )
+    engine = PyIcebergEngine(config)
+    engine.append_rows(table_name, rows)
+    return {'status': 'success', 'rows_appended': len(rows)}

awslabs/s3_tables_mcp_server/engines/__init__.py

@@ -0,0 +1,13 @@
+# 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/s3_tables_mcp_server/engines/pyiceberg.py

@@ -0,0 +1,239 @@
+# 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.
+
+"""Engine for interacting with Iceberg tables using pyiceberg and daft (read-only)."""
+
+import pyarrow as pa
+from ..utils import pyiceberg_load_catalog
+from daft import Catalog as DaftCatalog
+from daft.session import Session
+from datetime import date, datetime, time
+from decimal import Decimal
+from pydantic import BaseModel
+from pyiceberg.types import (
+    BinaryType,
+    BooleanType,
+    DateType,
+    DecimalType,
+    DoubleType,
+    FixedType,
+    FloatType,
+    IntegerType,
+    ListType,
+    LongType,
+    MapType,
+    StringType,
+    StructType,
+    TimestampType,
+    TimestamptzType,
+    TimeType,
+    UUIDType,
+)
+
+# pyiceberg and daft imports
+from typing import Any, Dict, Optional
+
+
+class PyIcebergConfig(BaseModel):
+    """Configuration for PyIceberg/Daft connection."""
+
+    warehouse: str  # e.g. 'arn:aws:s3tables:us-west-2:484907528679:bucket/customer-data-bucket'
+    uri: str  # e.g. 'https://s3tables.us-west-2.amazonaws.com/iceberg'
+    region: str  # e.g. 'us-west-2'
+    namespace: str  # e.g. 'retail_data'
+    catalog_name: str = 's3tablescatalog'  # default
+    rest_signing_name: str = 's3tables'
+    rest_sigv4_enabled: str = 'true'
+
+
+def convert_value_for_append(value, iceberg_type):
+    """Convert a value to the appropriate type for appending to an Iceberg table column.
+
+    Args:
+        value: The value to convert. Can be of various types (str, int, float, etc.).
+        iceberg_type: The Iceberg type to convert the value to.
+
+    Returns:
+        The value converted to the appropriate type for the Iceberg column, or None if value is None.
+
+    Raises:
+        NotImplementedError: If the iceberg_type is a complex type (ListType, MapType, StructType).
+        ValueError: If the conversion is unsupported or fails.
+    """
+    if value is None:
+        return None
+    # Already correct type
+    if isinstance(iceberg_type, BooleanType) and isinstance(value, bool):
+        return value
+    if isinstance(iceberg_type, (IntegerType, LongType)) and isinstance(value, int):
+        return value
+    if isinstance(iceberg_type, (FloatType, DoubleType)) and isinstance(value, float):
+        return value
+    if isinstance(iceberg_type, DecimalType) and isinstance(value, Decimal):
+        return value
+    if isinstance(iceberg_type, DateType) and isinstance(value, date):
+        return value
+    if isinstance(iceberg_type, TimeType) and isinstance(value, time):
+        return value
+    if isinstance(iceberg_type, (TimestampType, TimestamptzType)) and isinstance(value, datetime):
+        return value
+    if isinstance(iceberg_type, StringType) and isinstance(value, str):
+        return value
+    # Convert from string
+    if isinstance(value, str):
+        if isinstance(iceberg_type, BooleanType):
+            return value.lower() in ('true', '1', 'yes')
+        if isinstance(iceberg_type, (IntegerType, LongType)):
+            return int(value)
+        if isinstance(iceberg_type, (FloatType, DoubleType)):
+            return float(value)
+        if isinstance(iceberg_type, DecimalType):
+            return Decimal(value)
+        if isinstance(iceberg_type, DateType):
+            return date.fromisoformat(value)
+        if isinstance(iceberg_type, TimeType):
+            return time.fromisoformat(value)
+        if isinstance(iceberg_type, (TimestampType, TimestamptzType)):
+            return datetime.fromisoformat(value)
+        if isinstance(iceberg_type, StringType):
+            return value
+        if isinstance(iceberg_type, UUIDType):
+            import uuid
+
+            return uuid.UUID(value)
+        if isinstance(iceberg_type, (BinaryType, FixedType)):
+            return bytes.fromhex(value)
+    # Convert from number
+    if isinstance(value, (int, float)):
+        if isinstance(iceberg_type, (IntegerType, LongType)):
+            return int(value)
+        if isinstance(iceberg_type, (FloatType, DoubleType)):
+            return float(value)
+        if isinstance(iceberg_type, DecimalType):
+            return Decimal(str(value))
+        if isinstance(iceberg_type, StringType):
+            return str(value)
+    if isinstance(iceberg_type, (ListType, MapType, StructType)):
+        raise NotImplementedError(f'Complex type {iceberg_type} not supported in append_rows')
+    raise ValueError(f'Unsupported conversion from {type(value)} to {iceberg_type}')
+
+
+class PyIcebergEngine:
+    """Engine for read-only queries on Iceberg tables using pyiceberg and daft."""
+
+    def __init__(self, config: PyIcebergConfig):
+        """Initialize the PyIcebergEngine with the given configuration.
+
+        Args:
+            config: PyIcebergConfig object containing connection parameters.
+        """
+        self.config = config
+        self._catalog: Optional[Any] = None
+        self._session: Optional[Session] = None
+        self._initialize_connection()
+
+    def _initialize_connection(self):
+        try:
+            self._catalog = pyiceberg_load_catalog(
+                self.config.catalog_name,
+                self.config.warehouse,
+                self.config.uri,
+                self.config.region,
+                self.config.rest_signing_name,
+                self.config.rest_sigv4_enabled,
+            )
+            self._session = Session()
+            self._session.attach(DaftCatalog.from_iceberg(self._catalog))
+            self._session.set_namespace(self.config.namespace)
+        except Exception as e:
+            raise ConnectionError(f'Failed to initialize PyIceberg connection: {str(e)}')
+
+    def execute_query(self, query: str) -> Dict[str, Any]:
+        """Execute a SQL query against the Iceberg catalog using Daft.
+
+        Args:
+            query: SQL query to execute
+
+        Returns:
+            Dict containing:
+                - columns: List of column names
+                - rows: List of rows, where each row is a list of values
+        """
+        if not self._session:
+            raise ConnectionError('No active session for PyIceberg/Daft')
+        try:
+            result = self._session.sql(query)
+            if result is None:
+                raise Exception('Query execution returned None result')
+            df = result.collect()
+            columns = df.column_names
+            rows = df.to_pylist()
+            return {
+                'columns': columns,
+                'rows': [list(row.values()) for row in rows],
+            }
+        except Exception as e:
+            raise Exception(f'Error executing query: {str(e)}')
+
+    def test_connection(self) -> bool:
+        """Test the connection by listing namespaces."""
+        if not self._session:
+            return False
+        try:
+            _ = self._session.list_namespaces()
+            return True
+        except Exception:
+            return False
+
+    def append_rows(self, table_name: str, rows: list[dict]) -> None:
+        """Append rows to an Iceberg table using pyiceberg.
+
+        Args:
+            table_name: The name of the table (e.g., 'namespace.tablename' or just 'tablename' if namespace is set)
+            rows: List of dictionaries, each representing a row to append
+
+        Raises:
+            Exception: If appending fails
+        """
+        if not self._catalog:
+            raise ConnectionError('No active catalog for PyIceberg')
+        try:
+            # If table_name does not contain a dot, prepend the namespace
+            if '.' not in table_name:
+                full_table_name = f'{self.config.namespace}.{table_name}'
+            else:
+                full_table_name = table_name
+            table = self._catalog.load_table(full_table_name)
+            iceberg_schema = table.schema()
+            converted_rows = []
+            for row in rows:
+                converted_row = {}
+                for field in iceberg_schema.fields:
+                    field_name = field.name
+                    field_type = field.field_type
+                    value = row.get(field_name)
+                    if field.required and value is None:
+                        raise ValueError(f'Required field {field_name} is missing or None')
+                    try:
+                        converted_row[field_name] = convert_value_for_append(value, field_type)
+                    except (ValueError, TypeError) as e:
+                        raise ValueError(
+                            f'Error converting value for field {field_name}: {str(e)}'
+                        )
+                converted_rows.append(converted_row)
+            schema = iceberg_schema.as_arrow()
+            pa_table = pa.Table.from_pylist(converted_rows, schema=schema)
+            table.append(pa_table)
+        except Exception as e:
+            raise Exception(f'Error appending rows: {str(e)}')