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

awslabs/s3_tables_mcp_server/models.py

@@ -0,0 +1,274 @@
+# 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 S3 Tables MCP Server models."""
+
+from awslabs.s3_tables_mcp_server.constants import (
+    TABLE_ARN_PATTERN,
+    TABLE_BUCKET_ARN_PATTERN,
+    TABLE_BUCKET_NAME_PATTERN,
+    TABLE_NAME_PATTERN,
+)
+from datetime import datetime
+from enum import Enum
+from pydantic import BaseModel, Field, model_validator
+from typing import List, Optional, Union
+
+
+# Enums
+class OpenTableFormat(str, Enum):
+    """Supported open table formats."""
+
+    ICEBERG = 'ICEBERG'
+
+
+class TableBucketType(str, Enum):
+    """Table bucket type."""
+
+    CUSTOMER = 'customer'
+    AWS = 'aws'
+
+
+class TableType(str, Enum):
+    """Table type."""
+
+    CUSTOMER = 'customer'
+    AWS = 'aws'
+
+
+class MaintenanceStatus(str, Enum):
+    """Maintenance status."""
+
+    ENABLED = 'enabled'
+    DISABLED = 'disabled'
+
+
+class JobStatus(str, Enum):
+    """Job status."""
+
+    NOT_YET_RUN = 'Not_Yet_Run'
+    SUCCESSFUL = 'Successful'
+    FAILED = 'Failed'
+    DISABLED = 'Disabled'
+
+
+class TableBucketMaintenanceType(str, Enum):
+    """Table bucket maintenance type."""
+
+    ICEBERG_UNREFERENCED_FILE_REMOVAL = 'icebergUnreferencedFileRemoval'
+
+
+class TableMaintenanceType(str, Enum):
+    """Table maintenance type."""
+
+    ICEBERG_COMPACTION = 'icebergCompaction'
+    ICEBERG_SNAPSHOT_MANAGEMENT = 'icebergSnapshotManagement'
+
+
+class TableMaintenanceJobType(str, Enum):
+    """Table maintenance job type."""
+
+    ICEBERG_COMPACTION = 'icebergCompaction'
+    ICEBERG_SNAPSHOT_MANAGEMENT = 'icebergSnapshotManagement'
+    ICEBERG_UNREFERENCED_FILE_REMOVAL = 'icebergUnreferencedFileRemoval'
+
+
+# Core Models
+class TableBucketSummary(BaseModel):
+    """Table bucket summary."""
+
+    arn: str = Field(pattern=TABLE_BUCKET_ARN_PATTERN)
+    name: str = Field(min_length=3, max_length=63, pattern=TABLE_BUCKET_NAME_PATTERN)
+    owner_account_id: str = Field(min_length=12, max_length=12, pattern=r'[0-9].*')
+    created_at: datetime
+    table_bucket_id: Optional[str] = None
+    type: Optional[TableBucketType] = None
+
+
+class TableBucket(BaseModel):
+    """Complete table bucket information."""
+
+    arn: str = Field(pattern=TABLE_BUCKET_ARN_PATTERN)
+    name: str = Field(min_length=3, max_length=63, pattern=TABLE_BUCKET_NAME_PATTERN)
+    owner_account_id: str = Field(min_length=12, max_length=12, pattern=r'[0-9].*')
+    created_at: datetime
+    table_bucket_id: Optional[str] = None
+    type: Optional[TableBucketType] = None
+
+
+class NamespaceSummary(BaseModel):
+    """Namespace summary."""
+
+    namespace: List[str]
+    created_at: datetime
+    created_by: str = Field(min_length=12, max_length=12, pattern=r'[0-9].*')
+    owner_account_id: str = Field(min_length=12, max_length=12, pattern=r'[0-9].*')
+    namespace_id: Optional[str] = None
+    table_bucket_id: Optional[str] = None
+
+
+class TableSummary(BaseModel):
+    """Table summary."""
+
+    namespace: List[str]
+    name: str = Field(min_length=1, max_length=255, pattern=TABLE_NAME_PATTERN)
+    type: TableType
+    table_arn: str = Field(pattern=TABLE_ARN_PATTERN)
+    created_at: datetime
+    modified_at: datetime
+    namespace_id: Optional[str] = None
+    table_bucket_id: Optional[str] = None
+
+
+class Table(BaseModel):
+    """Complete table information."""
+
+    name: str = Field(min_length=1, max_length=255, pattern=TABLE_NAME_PATTERN)
+    type: TableType
+    table_arn: str = Field(pattern=TABLE_ARN_PATTERN, alias='tableARN')
+    namespace: List[str]
+    namespace_id: Optional[str] = None
+    version_token: str = Field(min_length=1, max_length=2048)
+    metadata_location: Optional[str] = Field(None, min_length=1, max_length=2048)
+    warehouse_location: str = Field(min_length=1, max_length=2048)
+    created_at: datetime
+    created_by: str = Field(min_length=12, max_length=12, pattern=r'[0-9].*')
+    managed_by_service: Optional[str] = None
+    modified_at: datetime
+    modified_by: str = Field(min_length=12, max_length=12, pattern=r'[0-9].*')
+    owner_account_id: str = Field(min_length=12, max_length=12, pattern=r'[0-9].*')
+    format: OpenTableFormat
+    table_bucket_id: Optional[str] = None
+
+
+# Maintenance Models
+class IcebergCompactionSettings(BaseModel):
+    """Settings for Iceberg compaction."""
+
+    target_file_size_mb: Optional[int] = Field(None, ge=1, le=2147483647)
+
+
+class IcebergSnapshotManagementSettings(BaseModel):
+    """Settings for Iceberg snapshot management."""
+
+    min_snapshots_to_keep: Optional[int] = Field(None, ge=1, le=2147483647)
+    max_snapshot_age_hours: Optional[int] = Field(None, ge=1, le=2147483647)
+
+
+class TableMaintenanceJobStatusValue(BaseModel):
+    """Table maintenance job status value."""
+
+    status: JobStatus
+    last_run_timestamp: Optional[datetime] = None
+    failure_message: Optional[str] = None
+
+
+class TableMaintenanceConfigurationValue(BaseModel):
+    """Table maintenance configuration value."""
+
+    status: Optional[MaintenanceStatus] = None
+    settings: Optional[Union[IcebergCompactionSettings, IcebergSnapshotManagementSettings]] = None
+
+
+class IcebergUnreferencedFileRemovalSettings(BaseModel):
+    """Settings for unreferenced file removal."""
+
+    unreferenced_days: Optional[int] = Field(None, ge=1, le=2147483647)
+    non_current_days: Optional[int] = Field(None, ge=1, le=2147483647)
+
+
+class TableBucketMaintenanceSettings(BaseModel):
+    """Contains details about the maintenance settings for the table bucket."""
+
+    iceberg_unreferenced_file_removal: Optional[IcebergUnreferencedFileRemovalSettings] = Field(
+        None, description='Settings for unreferenced file removal.'
+    )
+
+    @model_validator(mode='after')
+    def validate_only_one_setting(self) -> 'TableBucketMaintenanceSettings':
+        """Validate that only one setting is specified."""
+        settings = [self.iceberg_unreferenced_file_removal]
+        if sum(1 for s in settings if s is not None) > 1:
+            raise ValueError('Only one maintenance setting can be specified')
+        return self
+
+
+class TableBucketMaintenanceConfigurationValue(BaseModel):
+    """Details about the values that define the maintenance configuration for a table bucket."""
+
+    settings: Optional[TableBucketMaintenanceSettings] = Field(
+        None, description='Contains details about the settings of the maintenance configuration.'
+    )
+    status: Optional[MaintenanceStatus] = Field(
+        None, description='The status of the maintenance configuration.'
+    )
+
+
+# Resource Models
+class TableBucketsResource(BaseModel):
+    """Resource containing all table buckets."""
+
+    table_buckets: List[TableBucketSummary]
+    total_count: int
+
+
+class NamespacesResource(BaseModel):
+    """Resource containing all namespaces."""
+
+    namespaces: List[NamespaceSummary]
+    total_count: int
+
+
+class TablesResource(BaseModel):
+    """Resource containing all tables."""
+
+    tables: List[TableSummary]
+    total_count: int
+
+
+# Error Models
+class S3TablesError(BaseModel):
+    """S3 Tables error response."""
+
+    error_code: str
+    error_message: str
+    request_id: Optional[str] = None
+    resource_name: Optional[str] = None
+
+
+# Schema Models
+class SchemaField(BaseModel):
+    """Iceberg schema field."""
+
+    name: str
+    type: str
+    required: Optional[bool] = None
+
+
+class IcebergSchema(BaseModel):
+    """Iceberg table schema."""
+
+    fields: List[SchemaField]
+
+
+class IcebergMetadata(BaseModel):
+    """Iceberg table metadata."""
+
+    table_schema: IcebergSchema = Field(alias='schema')
+
+
+class TableMetadata(BaseModel):
+    """Table metadata union."""
+
+    iceberg: Optional[IcebergMetadata] = None

awslabs/s3_tables_mcp_server/namespaces.py

@@ -0,0 +1,63 @@
+# 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.
+
+"""Namespace Management tools for S3 Tables MCP Server."""
+
+from .utils import get_s3tables_client, handle_exceptions
+from typing import Any, Dict, Optional
+
+
+@handle_exceptions
+async def create_namespace(
+    table_bucket_arn: str,
+    namespace: str,
+    region_name: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Create a new namespace."""
+    client = get_s3tables_client(region_name)
+    response = client.create_namespace(tableBucketARN=table_bucket_arn, namespace=[namespace])
+
+    return dict(response)
+
+
+@handle_exceptions
+async def delete_namespace(
+    table_bucket_arn: str,
+    namespace: str,
+    region_name: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Delete a namespace.
+
+    Permissions:
+    You must have the s3tables:DeleteNamespace permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+    response = client.delete_namespace(tableBucketARN=table_bucket_arn, namespace=namespace)
+    return dict(response)
+
+
+@handle_exceptions
+async def get_namespace(
+    table_bucket_arn: str, namespace: str, region_name: Optional[str] = None
+) -> Dict[str, Any]:
+    """Get details about a namespace.
+
+    Gets details about a namespace.
+
+    Permissions:
+    You must have the s3tables:GetNamespace permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+    response = client.get_namespace(tableBucketARN=table_bucket_arn, namespace=namespace)
+    return dict(response)

awslabs/s3_tables_mcp_server/resources.py

@@ -0,0 +1,231 @@
+# 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.
+
+"""MCP resource definitions for S3 Tables MCP Server."""
+
+import json
+from .models import (
+    NamespacesResource,
+    NamespaceSummary,
+    TableBucketsResource,
+    TableBucketSummary,
+    TablesResource,
+    TableSummary,
+)
+from .utils import get_s3tables_client
+from pydantic import BaseModel
+from typing import Any, Callable, Dict, List, Optional, Type, TypeVar
+
+
+T = TypeVar('T')
+ResourceT = TypeVar('ResourceT', bound=BaseModel)
+
+
+def create_error_response(error: Exception, resource_name: str) -> str:
+    """Create a standardized error response.
+
+    Args:
+        error: The exception that occurred
+        resource_name: The name of the resource being accessed
+
+    Returns:
+        A JSON string containing the error response
+    """
+    return json.dumps({'error': str(error), resource_name: [], 'total_count': 0})
+
+
+async def paginate_and_collect(
+    paginator: Any,
+    collection_key: str,
+    item_constructor: Callable[[Dict[str, Any]], T],
+    **pagination_args,
+) -> List[T]:
+    """Collect items from a paginated response.
+
+    Args:
+        paginator: The paginator to use
+        collection_key: The key in the response that contains the items
+        item_constructor: A function that constructs an item from a response
+        **pagination_args: Additional arguments to pass to the paginator
+
+    Returns:
+        A list of constructed items
+    """
+    items = []
+    for page in paginator.paginate(**pagination_args):
+        for item in page.get(collection_key, []):
+            items.append(item_constructor(item))
+    return items
+
+
+async def create_resource_response(
+    items: List[T], resource_class: Type[ResourceT], resource_name: str
+) -> str:
+    """Create a resource response.
+
+    Args:
+        items: The items to include in the resource
+        resource_class: The resource class to use
+        resource_name: The name of the resource
+
+    Returns:
+        A JSON string containing the resource response
+    """
+    try:
+        resource = resource_class(**{resource_name: items, 'total_count': len(items)})
+        return resource.model_dump_json()
+    except Exception as e:
+        return create_error_response(e, resource_name)
+
+
+async def list_table_buckets_resource(region_name: Optional[str] = None) -> str:
+    """List all S3 Tables buckets.
+
+    Lists table buckets for your account. Requires s3tables:ListTableBuckets permission.
+    The API supports pagination with continuationToken and filtering with prefix.
+
+    Returns:
+        A JSON string containing the list of table buckets and total count.
+    """
+    try:
+        client = get_s3tables_client(region_name)
+        paginator = client.get_paginator('list_table_buckets')
+
+        table_buckets = await paginate_and_collect(
+            paginator=paginator,
+            collection_key='tableBuckets',
+            item_constructor=lambda bucket: TableBucketSummary(
+                arn=bucket['arn'],
+                name=bucket['name'],
+                owner_account_id=bucket['ownerAccountId'],
+                created_at=bucket['createdAt'],
+                table_bucket_id=bucket.get('tableBucketId'),
+                type=bucket.get('type'),
+            ),
+        )
+
+        return await create_resource_response(
+            items=table_buckets, resource_class=TableBucketsResource, resource_name='table_buckets'
+        )
+
+    except Exception as e:
+        return create_error_response(e, 'table_buckets')
+
+
+async def get_table_buckets(region_name: Optional[str] = None) -> List[TableBucketSummary]:
+    """Get all table buckets as TableBucketSummary objects.
+
+    Returns:
+        A list of TableBucketSummary objects
+    """
+    response = await list_table_buckets_resource(region_name=region_name)
+    data = json.loads(response)
+    if 'error' in data:
+        raise Exception(data['error'])
+    return [TableBucketSummary(**bucket) for bucket in data['table_buckets']]
+
+
+async def list_namespaces_resource(region_name: Optional[str] = None) -> str:
+    """List all namespaces across all table buckets.
+
+    Lists the namespaces within table buckets. Requires s3tables:ListNamespaces permission.
+    The API supports pagination with continuationToken and filtering with prefix.
+
+    Returns:
+        A JSON string containing the list of namespaces and total count.
+    """
+    try:
+        client = get_s3tables_client(region_name)
+
+        # Get all table buckets
+        table_buckets = await get_table_buckets(region_name=region_name)
+
+        # Then get namespaces for each bucket
+        all_namespaces = []
+        for bucket in table_buckets:
+            try:
+                namespace_paginator = client.get_paginator('list_namespaces')
+                namespaces = await paginate_and_collect(
+                    paginator=namespace_paginator,
+                    collection_key='namespaces',
+                    item_constructor=lambda namespace: NamespaceSummary(
+                        namespace=namespace['namespace'],
+                        created_at=namespace['createdAt'],
+                        created_by=namespace['createdBy'],
+                        owner_account_id=namespace['ownerAccountId'],
+                        namespace_id=namespace.get('namespaceId'),
+                        table_bucket_id=namespace.get('tableBucketId'),
+                    ),
+                    tableBucketARN=bucket.arn,
+                )
+                all_namespaces.extend(namespaces)
+            except Exception as e:
+                return create_error_response(e, 'namespaces')
+
+        return await create_resource_response(
+            items=all_namespaces, resource_class=NamespacesResource, resource_name='namespaces'
+        )
+
+    except Exception as e:
+        return create_error_response(e, 'namespaces')
+
+
+async def list_tables_resource(region_name: Optional[str] = None) -> str:
+    """List all Iceberg tables across all table buckets and namespaces.
+
+    Lists tables in the given table bucket. Requires s3tables:ListTables permission.
+    The API supports pagination with continuationToken, filtering with prefix and namespace,
+    and limiting results with maxTables.
+
+    Returns:
+        A JSON string containing the list of tables and total count.
+    """
+    try:
+        client = get_s3tables_client(region_name)
+
+        # Get all table buckets
+        table_buckets = await get_table_buckets(region_name=region_name)
+
+        # Then get tables for each bucket
+        all_tables = []
+        for bucket in table_buckets:
+            try:
+                table_paginator = client.get_paginator('list_tables')
+
+                tables = await paginate_and_collect(
+                    paginator=table_paginator,
+                    collection_key='tables',
+                    item_constructor=lambda table: TableSummary(
+                        namespace=table['namespace'],
+                        name=table['name'],
+                        type=table['type'],
+                        table_arn=table['tableARN'],
+                        created_at=table['createdAt'],
+                        modified_at=table['modifiedAt'],
+                        namespace_id=table.get('namespaceId'),
+                        table_bucket_id=table.get('tableBucketId'),
+                    ),
+                    tableBucketARN=bucket.arn,
+                )
+
+                all_tables.extend(tables)
+            except Exception as e:
+                return create_error_response(e, 'tables')
+
+        return await create_resource_response(
+            items=all_tables, resource_class=TablesResource, resource_name='tables'
+        )
+
+    except Exception as e:
+        return create_error_response(e, 'tables')

awslabs/s3_tables_mcp_server/s3_operations.py

@@ -0,0 +1,43 @@
+# 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.
+
+"""S3 Operations tools for S3 Tables MCP Server."""
+
+from .utils import get_s3_client, handle_exceptions
+from typing import Any, Dict, Optional
+
+
+@handle_exceptions
+async def get_bucket_metadata_table_configuration(
+    bucket: str, region_name: Optional[str] = None
+) -> Dict[str, Any]:
+    """Get the metadata table configuration for an S3 bucket.
+
+    Gets the metadata table configuration for an S3 bucket. This configuration
+    determines how metadata is stored and managed for the bucket.
+
+    Permissions:
+    You must have the s3:GetBucketMetadataTableConfiguration permission to use this operation.
+
+    Args:
+        bucket: The name of the S3 bucket
+        region_name: Optional AWS region name. If not provided, uses AWS_REGION environment variable
+                    or defaults to 'us-east-1'.
+
+    Returns:
+        Dict containing the bucket metadata table configuration
+    """
+    client = get_s3_client(region_name)
+    response = client.get_bucket_metadata_table_configuration(Bucket=bucket)
+    return dict(response)