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

awslabs/s3_tables_mcp_server/table_buckets.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.
+
+"""Table Bucket Management tools for S3 Tables MCP Server."""
+
+from .models import (
+    TableBucketMaintenanceConfigurationValue,
+    TableBucketMaintenanceType,
+)
+from .utils import get_s3tables_client, handle_exceptions
+from typing import Any, Dict, Optional
+
+
+@handle_exceptions
+async def create_table_bucket(
+    name: str,
+    region_name: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Create a new S3 Tables bucket."""
+    client = get_s3tables_client(region_name)
+
+    # Prepare parameters for create_table_bucket
+    params = {'name': name}
+
+    response = client.create_table_bucket(**params)
+    return dict(response)
+
+
+@handle_exceptions
+async def delete_table_bucket(
+    table_bucket_arn: str, region_name: Optional[str] = None
+) -> Dict[str, Any]:
+    """Delete a table bucket.
+
+    Permissions:
+    You must have the s3tables:DeleteTableBucket permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+    response = client.delete_table_bucket(tableBucketARN=table_bucket_arn)
+    return dict(response)
+
+
+@handle_exceptions
+async def put_table_bucket_maintenance_configuration(
+    table_bucket_arn: str,
+    maintenance_type: TableBucketMaintenanceType,
+    value: TableBucketMaintenanceConfigurationValue,
+    region_name: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Create or replace a maintenance configuration for a table bucket.
+
+    Permissions:
+    You must have the s3tables:PutTableBucketMaintenanceConfiguration permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+    response = client.put_table_bucket_maintenance_configuration(
+        tableBucketARN=table_bucket_arn,
+        type=maintenance_type.value,
+        value=value.model_dump(by_alias=True, exclude_none=True),
+    )
+    return dict(response)
+
+
+@handle_exceptions
+async def get_table_bucket(
+    table_bucket_arn: str, region_name: Optional[str] = None
+) -> Dict[str, Any]:
+    """Get details about a table bucket.
+
+    Gets details on a table bucket.
+
+    Permissions:
+    You must have the s3tables:GetTableBucket permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+    response = client.get_table_bucket(tableBucketARN=table_bucket_arn)
+    return dict(response)
+
+
+@handle_exceptions
+async def get_table_bucket_maintenance_configuration(
+    table_bucket_arn: str, region_name: Optional[str] = None
+) -> Dict[str, Any]:
+    """Get details about a maintenance configuration for a table bucket.
+
+    Gets details about a maintenance configuration for a given table bucket.
+
+    Permissions:
+    You must have the s3tables:GetTableBucketMaintenanceConfiguration permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+    response = client.get_table_bucket_maintenance_configuration(tableBucketARN=table_bucket_arn)
+    return dict(response)
+
+
+@handle_exceptions
+async def get_table_bucket_policy(
+    table_bucket_arn: str, region_name: Optional[str] = None
+) -> Dict[str, Any]:
+    """Get details about a table bucket policy.
+
+    Gets details about a table bucket policy.
+
+    Permissions:
+    You must have the s3tables:GetTableBucketPolicy permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+    response = client.get_table_bucket_policy(tableBucketARN=table_bucket_arn)
+    return dict(response)
+
+
+@handle_exceptions
+async def delete_table_bucket_policy(
+    table_bucket_arn: str, region_name: Optional[str] = None
+) -> Dict[str, Any]:
+    """Delete a table bucket policy.
+
+    Deletes a table bucket policy.
+
+    Permissions:
+    You must have the s3tables:DeleteTableBucketPolicy permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+    response = client.delete_table_bucket_policy(tableBucketARN=table_bucket_arn)
+    return dict(response)

awslabs/s3_tables_mcp_server/tables.py

@@ -0,0 +1,307 @@
+# 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.
+
+"""Table Management tools for S3 Tables MCP Server."""
+
+from .models import (
+    OpenTableFormat,
+    TableMaintenanceConfigurationValue,
+    TableMaintenanceType,
+    TableMetadata,
+)
+from .utils import get_s3tables_client, handle_exceptions
+from typing import Any, Dict, Optional
+
+
+@handle_exceptions
+async def create_table(
+    table_bucket_arn: str,
+    namespace: str,
+    name: str,
+    format: OpenTableFormat = OpenTableFormat.ICEBERG,
+    metadata: Optional[TableMetadata] = None,
+    region_name: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Create a new table associated with the given namespace in a table bucket.
+
+    Permissions:
+    You must have the s3tables:CreateTable permission to use this operation.
+    If using metadata parameter, you must have the s3tables:PutTableData permission.
+    """
+    client = get_s3tables_client(region_name)
+
+    # Prepare parameters for create_table
+    params: Dict[str, Any] = {
+        'tableBucketARN': table_bucket_arn,
+        'namespace': namespace,
+        'name': name,
+        'format': format.value,
+    }
+
+    # Add metadata if provided
+    if metadata:
+        params['metadata'] = metadata.model_dump(by_alias=True, exclude_none=True)
+
+    response = client.create_table(**params)
+    return dict(response)
+
+
+@handle_exceptions
+async def delete_table(
+    table_bucket_arn: str,
+    namespace: str,
+    name: str,
+    version_token: Optional[str] = None,
+    region_name: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Delete a table.
+
+    Permissions:
+    You must have the s3tables:DeleteTable permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+
+    # Prepare parameters for delete_table
+    params: Dict[str, Any] = {
+        'tableBucketARN': table_bucket_arn,
+        'namespace': namespace,
+        'name': name,
+    }
+
+    # Add version token if provided
+    if version_token:
+        params['versionToken'] = version_token
+
+    response = client.delete_table(**params)
+    return dict(response)
+
+
+@handle_exceptions
+async def get_table(
+    table_bucket_arn: str, namespace: str, name: str, region_name: Optional[str] = None
+) -> Dict[str, Any]:
+    """Get details about a table.
+
+    Gets details about a table.
+
+    Permissions:
+    You must have the s3tables:GetTable permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+    response = client.get_table(tableBucketARN=table_bucket_arn, namespace=namespace, name=name)
+    return dict(response)
+
+
+@handle_exceptions
+async def delete_table_policy(
+    table_bucket_arn: str, namespace: str, name: str, region_name: Optional[str] = None
+) -> Dict[str, Any]:
+    """Delete a table policy.
+
+    Deletes a table policy.
+
+    Permissions:
+    You must have the s3tables:DeleteTablePolicy permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+    response = client.delete_table_policy(
+        tableBucketARN=table_bucket_arn, namespace=namespace, name=name
+    )
+    return dict(response)
+
+
+@handle_exceptions
+async def get_table_maintenance_configuration(
+    table_bucket_arn: str, namespace: str, name: str, region_name: Optional[str] = None
+) -> Dict[str, Any]:
+    """Get details about the maintenance configuration of a table.
+
+    Gets details about the maintenance configuration of a table. For more information, see S3 Tables maintenance in the Amazon Simple Storage Service User Guide.
+
+    Permissions:
+    You must have the s3tables:GetTableMaintenanceConfiguration permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+    response = client.get_table_maintenance_configuration(
+        tableBucketARN=table_bucket_arn, namespace=namespace, name=name
+    )
+    return dict(response)
+
+
+@handle_exceptions
+async def get_table_maintenance_job_status(
+    table_bucket_arn: str, namespace: str, name: str, region_name: Optional[str] = None
+) -> Dict[str, Any]:
+    """Get the status of a maintenance job for a table.
+
+    Gets the status of a maintenance job for a table. For more information, see S3 Tables maintenance in the Amazon Simple Storage Service User Guide.
+
+    Permissions:
+    You must have the s3tables:GetTableMaintenanceJobStatus permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+    response = client.get_table_maintenance_job_status(
+        tableBucketARN=table_bucket_arn, namespace=namespace, name=name
+    )
+    return dict(response)
+
+
+@handle_exceptions
+async def get_table_metadata_location(
+    table_bucket_arn: str, namespace: str, name: str, region_name: Optional[str] = None
+) -> Dict[str, Any]:
+    """Get the location of the table metadata.
+
+    Gets the location of the table metadata.
+
+    Permissions:
+    You must have the s3tables:GetTableMetadataLocation permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+    response = client.get_table_metadata_location(
+        tableBucketARN=table_bucket_arn, namespace=namespace, name=name
+    )
+    return dict(response)
+
+
+@handle_exceptions
+async def get_table_policy(
+    table_bucket_arn: str, namespace: str, name: str, region_name: Optional[str] = None
+) -> Dict[str, Any]:
+    """Get details about a table policy.
+
+    Gets details about a table policy. For more information, see Viewing a table policy in the Amazon Simple Storage Service User Guide.
+
+    Permissions:
+    You must have the s3tables:GetTablePolicy permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+    response = client.get_table_policy(
+        tableBucketARN=table_bucket_arn, namespace=namespace, name=name
+    )
+    return dict(response)
+
+
+@handle_exceptions
+async def put_table_maintenance_configuration(
+    table_bucket_arn: str,
+    namespace: str,
+    name: str,
+    maintenance_type: TableMaintenanceType,
+    value: TableMaintenanceConfigurationValue,
+    region_name: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Create or replace a maintenance configuration for a table.
+
+    Creates a new maintenance configuration or replaces an existing maintenance configuration for a table.
+    For more information, see S3 Tables maintenance in the Amazon Simple Storage Service User Guide.
+
+    Permissions:
+    You must have the s3tables:PutTableMaintenanceConfiguration permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+
+    # Prepare parameters for put_table_maintenance_configuration
+    params: Dict[str, Any] = {
+        'tableBucketARN': table_bucket_arn,
+        'namespace': namespace,
+        'name': name,
+        'type': maintenance_type.value,
+        'value': value.model_dump(by_alias=True, exclude_none=True),
+    }
+
+    response = client.put_table_maintenance_configuration(**params)
+    return dict(response)
+
+
+@handle_exceptions
+async def rename_table(
+    table_bucket_arn: str,
+    namespace: str,
+    name: str,
+    new_name: Optional[str] = None,
+    new_namespace_name: Optional[str] = None,
+    version_token: Optional[str] = None,
+    region_name: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Rename a table or a namespace.
+
+    Renames a table or a namespace. For more information, see S3 Tables in the Amazon Simple Storage Service User Guide.
+
+    Args:
+        table_bucket_arn: The Amazon Resource Name (ARN) of the table bucket.
+        namespace: The namespace associated with the table.
+        name: The current name of the table.
+        new_name: Optional new name for the table.
+        new_namespace_name: Optional new name for the namespace.
+        version_token: Optional version token of the table.
+        region_name: Optional AWS region name.
+
+    Returns:
+        Dict containing the response from the rename operation.
+
+    Permissions:
+    You must have the s3tables:RenameTable permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+
+    # Prepare parameters for rename_table
+    params: Dict[str, Any] = {
+        'tableBucketARN': table_bucket_arn,
+        'namespace': namespace,
+        'name': name,
+    }
+
+    # Add optional parameters if provided
+    if new_name:
+        params['newName'] = new_name
+    if new_namespace_name:
+        params['newNamespaceName'] = new_namespace_name
+    if version_token:
+        params['versionToken'] = version_token
+
+    response = client.rename_table(**params)
+    return dict(response)
+
+
+@handle_exceptions
+async def update_table_metadata_location(
+    table_bucket_arn: str,
+    namespace: str,
+    name: str,
+    metadata_location: str,
+    version_token: str,
+    region_name: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Update the metadata location for a table.
+
+    Updates the metadata location for a table. The metadata location of a table must be an S3 URI that begins with the table's warehouse location.
+    The metadata location for an Apache Iceberg table must end with .metadata.json, or if the metadata file is Gzip-compressed, .metadata.json.gz.
+
+    Permissions:
+    You must have the s3tables:UpdateTableMetadataLocation permission to use this operation.
+    """
+    client = get_s3tables_client(region_name)
+
+    # Prepare parameters for update_table_metadata_location
+    params: Dict[str, Any] = {
+        'tableBucketARN': table_bucket_arn,
+        'namespace': namespace,
+        'name': name,
+        'metadataLocation': metadata_location,
+        'versionToken': version_token,
+    }
+
+    response = client.update_table_metadata_location(**params)
+    return dict(response)

awslabs/s3_tables_mcp_server/utils.py

@@ -0,0 +1,139 @@
+# 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.
+
+"""Common utilities and helpers for S3 Tables MCP Server."""
+
+import boto3
+import os
+from . import __version__
+from botocore.client import BaseClient
+from botocore.config import Config
+from functools import wraps
+from typing import Optional
+
+
+_user_agent_mode = 'ro'  # Default to read-only
+
+
+def set_user_agent_mode(allow_write: bool):
+    """Set the user agent mode to 'rw' (read-write) or 'ro' (read-only)."""
+    global _user_agent_mode
+    _user_agent_mode = 'rw' if allow_write else 'ro'
+
+
+def _user_agent_extra():
+    return f'awslabs/mcp/s3-tables-mcp-server/{__version__}/{_user_agent_mode}'
+
+
+def handle_exceptions(func):
+    """Decorator to handle exceptions consistently across tools."""
+
+    @wraps(func)
+    async def wrapper(*args, **kwargs):
+        try:
+            return await func(*args, **kwargs)
+        except Exception as e:
+            return {'error': str(e), 'tool': func.__name__}
+
+    return wrapper
+
+
+def get_s3tables_client(region_name: Optional[str] = None) -> BaseClient:
+    """Create a boto3 S3 Tables client.
+
+    Args:
+        region_name: Optional AWS region name. If not provided, uses AWS_REGION environment variable
+                    or defaults to 'us-east-1'.
+
+    Returns:
+        boto3.client: Configured S3 Tables client
+    """
+    region = region_name or os.getenv('AWS_REGION') or 'us-east-1'
+    config = Config(user_agent_extra=_user_agent_extra())
+    session = boto3.Session()
+    return session.client('s3tables', region_name=region, config=config)
+
+
+def get_s3_client(region_name: Optional[str] = None) -> BaseClient:
+    """Create a boto3 S3 client.
+
+    Args:
+        region_name: Optional AWS region name. If not provided, uses AWS_REGION environment variable
+                    or defaults to 'us-east-1'.
+
+    Returns:
+        boto3.client: Configured S3 client
+    """
+    region = region_name or os.getenv('AWS_REGION') or 'us-east-1'
+    config = Config(user_agent_extra=_user_agent_extra())
+    session = boto3.Session()
+    return session.client('s3', region_name=region, config=config)
+
+
+def get_sts_client(region_name: Optional[str] = None) -> BaseClient:
+    """Create a boto3 STS client.
+
+    Args:
+        region_name: Optional AWS region name. If not provided, uses AWS_REGION environment variable
+                    or defaults to 'us-east-1'.
+
+    Returns:
+        boto3.client: Configured STS client
+    """
+    region = region_name or os.getenv('AWS_REGION') or 'us-east-1'
+    config = Config(user_agent_extra=_user_agent_extra())
+    session = boto3.Session()
+    return session.client('sts', region_name=region, config=config)
+
+
+def get_athena_client(region_name: Optional[str] = None) -> BaseClient:
+    """Create a boto3 Athena client.
+
+    Args:
+        region_name: Optional AWS region name. If not provided, uses AWS_REGION environment variable
+                    or defaults to 'us-east-1'.
+
+    Returns:
+        boto3.client: Configured Athena client
+    """
+    region = region_name or os.getenv('AWS_REGION') or 'us-east-1'
+    config = Config(user_agent_extra=_user_agent_extra())
+    session = boto3.Session()
+    return session.client('athena', region_name=region, config=config)
+
+
+def pyiceberg_load_catalog(
+    catalog_name: str,
+    warehouse: str,
+    uri: str,
+    region: str,
+    rest_signing_name: str = 's3tables',
+    rest_sigv4_enabled: str = 'true',
+):
+    """Load a PyIceberg catalog with the given parameters."""
+    from pyiceberg.catalog import load_catalog
+
+    catalog = load_catalog(
+        catalog_name,
+        **{
+            'type': 'rest',
+            'warehouse': warehouse,
+            'uri': uri,
+            'rest.sigv4-enabled': rest_sigv4_enabled,
+            'rest.signing-name': rest_signing_name,
+            'rest.signing-region': region,
+        },
+    )
+    catalog._session.headers['User-Agent'] = _user_agent_extra()  # type: ignore[attr-defined]
+    return catalog