awslabs.s3-tables-mcp-server 0.0.3__py3-none-any.whl → 0.0.5__py3-none-any.whl

awslabs/s3_tables_mcp_server/__init__.py

@@ -15,4 +15,4 @@
 # This file is part of the awslabs namespace.
 # It is intentionally minimal to support PEP 420 namespace packages.
 
-__version__ = '0.0.3'
+__version__ = '0.0.5'

awslabs/s3_tables_mcp_server/file_processor/csv.py

@@ -18,13 +18,8 @@ This module provides functionality for processing and analyzing uploaded files,
 particularly focusing on CSV file handling and import capabilities.
 """
 
-import io
-import os
 import pyarrow.csv as pc
-from ..utils import get_s3_client, pyiceberg_load_catalog
-from pyiceberg.exceptions import NoSuchTableError
-from typing import Dict
-from urllib.parse import urlparse
+from .utils import import_file_to_table
 
 
 async def import_csv_to_table(
@@ -37,87 +32,19 @@ async def import_csv_to_table(
     catalog_name: str = 's3tablescatalog',
     rest_signing_name: str = 's3tables',
     rest_sigv4_enabled: str = 'true',
-) -> Dict:
-    """Import data from a CSV file into an S3 table.
-
-    This function reads data from a CSV file stored in S3 and imports it into an existing S3 table.
-    If the table doesn't exist, it will be created using the schema inferred from the CSV file.
-
-    Args:
-        warehouse: Warehouse string for Iceberg catalog
-        region: AWS region for S3Tables/Iceberg REST endpoint
-        namespace: The namespace containing the table
-        table_name: The name of the table to import data into
-        s3_url: The S3 URL of the CSV file (format: s3://bucket-name/key)
-        uri: REST URI for Iceberg catalog
-        catalog_name: Catalog name
-        rest_signing_name: REST signing name
-        rest_sigv4_enabled: Enable SigV4 signing
-
-    Returns:
-        A dictionary containing:
-        - status: 'success' or 'error'
-        - message: Success message or error details
-        - rows_processed: Number of rows processed (on success)
-        - file_processed: Name of the processed file
-        - table_created: Boolean indicating if a new table was created (on success)
-    """
-    # Parse S3 URL
-    parsed = urlparse(s3_url)
-    bucket = parsed.netloc
-    key = parsed.path.lstrip('/')
-
-    try:
-        # Load Iceberg catalog
-        catalog = pyiceberg_load_catalog(
-            catalog_name,
-            warehouse,
-            uri,
-            region,
-            rest_signing_name,
-            rest_sigv4_enabled,
-        )
-
-        # Get S3 client and read the CSV file to infer schema
-        s3_client = get_s3_client()
-        response = s3_client.get_object(Bucket=bucket, Key=key)
-        csv_data = response['Body'].read()
-
-        # Read CSV file into PyArrow Table to infer schema
-        # Convert bytes to file-like object for PyArrow
-        csv_buffer = io.BytesIO(csv_data)
-        csv_table = pc.read_csv(csv_buffer)
-        csv_schema = csv_table.schema
-
-        table_created = False
-        try:
-            # Try to load existing table
-            table = catalog.load_table(f'{namespace}.{table_name}')
-        except NoSuchTableError:
-            # Table doesn't exist, create it using the CSV schema
-            try:
-                table = catalog.create_table(
-                    identifier=f'{namespace}.{table_name}',
-                    schema=csv_schema,
-                )
-                table_created = True
-            except Exception as create_error:
-                return {
-                    'status': 'error',
-                    'error': f'Failed to create table: {str(create_error)}',
-                }
-
-        # Append data to Iceberg table
-        table.append(csv_table)
-
-        return {
-            'status': 'success',
-            'message': f'Successfully imported {csv_table.num_rows} rows{" and created new table" if table_created else ""}',
-            'rows_processed': csv_table.num_rows,
-            'file_processed': os.path.basename(key),
-            'table_created': table_created,
-            'table_uuid': table.metadata.table_uuid,
-        }
-
-    except Exception as e:
-        return {'status': 'error', 'error': str(e)}
+    preserve_case: bool = False,
+):
+    """Import a CSV file into an S3 table using PyArrow."""
+    return await import_file_to_table(
+        warehouse=warehouse,
+        region=region,
+        namespace=namespace,
+        table_name=table_name,
+        s3_url=s3_url,
+        uri=uri,
+        create_pyarrow_table=pc.read_csv,
+        catalog_name=catalog_name,
+        rest_signing_name=rest_signing_name,
+        rest_sigv4_enabled=rest_sigv4_enabled,
+        preserve_case=preserve_case,
+    )

awslabs/s3_tables_mcp_server/file_processor/parquet.py

@@ -13,10 +13,7 @@
 # limitations under the License.
 
 import pyarrow.parquet as pq
-from awslabs.s3_tables_mcp_server.utils import get_s3_client, pyiceberg_load_catalog
-from io import BytesIO
-from pyiceberg.exceptions import NoSuchTableError
-from typing import Dict
+from .utils import import_file_to_table
 
 
 async def import_parquet_to_table(
@@ -29,88 +26,19 @@ async def import_parquet_to_table(
     catalog_name: str = 's3tablescatalog',
     rest_signing_name: str = 's3tables',
     rest_sigv4_enabled: str = 'true',
-) -> Dict:
-    """Import data from a Parquet file into an S3 table.
-
-    This function reads data from a Parquet file stored in S3 and imports it into an existing Iceberg table.
-    If the table doesn't exist, it will be created using the schema from the Parquet file.
-
-    Args:
-        warehouse: Warehouse string for Iceberg catalog
-        region: AWS region for S3Tables/Iceberg REST endpoint
-        namespace: The namespace containing the table
-        table_name: The name of the table to import data into
-        s3_url: The S3 URL of the Parquet file
-        uri: REST URI for Iceberg catalog
-        catalog_name: Catalog name
-        rest_signing_name: REST signing name
-        rest_sigv4_enabled: Enable SigV4 signing
-
-    Returns:
-        A dictionary containing:
-        - status: 'success' or 'error'
-        - message: Success message or error details
-        - rows_processed: Number of rows processed (on success)
-        - file_processed: Name of the processed file
-        - table_created: Boolean indicating if a new table was created (on success)
-    """
-    import os
-    from urllib.parse import urlparse
-
-    # Parse S3 URL
-    parsed = urlparse(s3_url)
-    bucket = parsed.netloc
-    key = parsed.path.lstrip('/')
-
-    try:
-        # Load Iceberg catalog
-        catalog = pyiceberg_load_catalog(
-            catalog_name,
-            warehouse,
-            uri,
-            region,
-            rest_signing_name,
-            rest_sigv4_enabled,
-        )
-
-        # Get S3 client and read the Parquet file first to get the schema
-        s3_client = get_s3_client()
-        response = s3_client.get_object(Bucket=bucket, Key=key)
-        parquet_data = BytesIO(response['Body'].read())
-
-        # Read Parquet file into PyArrow Table
-        parquet_table = pq.read_table(parquet_data)
-        parquet_schema = parquet_table.schema
-
-        table_created = False
-        try:
-            # Try to load existing table
-            table = catalog.load_table(f'{namespace}.{table_name}')
-        except NoSuchTableError:
-            # Table doesn't exist, create it using the Parquet schema
-            try:
-                table = catalog.create_table(
-                    identifier=f'{namespace}.{table_name}',
-                    schema=parquet_schema,
-                )
-                table_created = True
-            except Exception as create_error:
-                return {
-                    'status': 'error',
-                    'error': f'Failed to create table: {str(create_error)}',
-                }
-
-        # Append data to Iceberg table
-        table.append(parquet_table)
-
-        return {
-            'status': 'success',
-            'message': f'Successfully imported {parquet_table.num_rows} rows{" and created new table" if table_created else ""}',
-            'rows_processed': parquet_table.num_rows,
-            'file_processed': os.path.basename(key),
-            'table_created': table_created,
-            'table_uuid': table.metadata.table_uuid,
-        }
-
-    except Exception as e:
-        return {'status': 'error', 'error': str(e)}
+    preserve_case: bool = False,
+):
+    """Import a Parquet file into an S3 table using PyArrow."""
+    return await import_file_to_table(
+        warehouse=warehouse,
+        region=region,
+        namespace=namespace,
+        table_name=table_name,
+        s3_url=s3_url,
+        uri=uri,
+        create_pyarrow_table=pq.read_table,
+        catalog_name=catalog_name,
+        rest_signing_name=rest_signing_name,
+        rest_sigv4_enabled=rest_sigv4_enabled,
+        preserve_case=preserve_case,
+    )

awslabs/s3_tables_mcp_server/file_processor/utils.py

@@ -0,0 +1,157 @@
+# 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 file processing utilities.
+
+This module provides utility functions for file processing operations,
+particularly focusing on column name conversion and schema transformation.
+"""
+
+import os
+import pyarrow as pa
+from ..utils import get_s3_client, pyiceberg_load_catalog
+from io import BytesIO
+from pydantic.alias_generators import to_snake
+from pyiceberg.exceptions import NoSuchTableError
+from typing import Any, Callable, Dict
+from urllib.parse import urlparse
+
+
+def convert_column_names_to_snake_case(schema: pa.Schema) -> pa.Schema:
+    """Convert column names in PyArrow schema to snake_case.
+
+    Args:
+        schema: PyArrow schema with original column names
+
+    Returns:
+        PyArrow schema with converted column names
+
+    Raises:
+        ValueError: If duplicate column names exist after conversion
+    """
+    # Extract original column names
+    original_names = schema.names
+
+    # Convert each column name to snake_case
+    converted_names = [to_snake(name) for name in original_names]
+
+    # Check for duplicates after conversion using set and len
+    if len(set(converted_names)) != len(converted_names):
+        raise ValueError(
+            f'Duplicate column names after case conversion. '
+            f'Original names: {original_names}. Converted names: {converted_names}'
+        )
+
+    # Create new schema with converted column names
+    new_fields = []
+    for i, field in enumerate(schema):
+        new_field = pa.field(
+            converted_names[i], field.type, nullable=field.nullable, metadata=field.metadata
+        )
+        new_fields.append(new_field)
+
+    return pa.schema(new_fields, metadata=schema.metadata)
+
+
+async def import_file_to_table(
+    warehouse: str,
+    region: str,
+    namespace: str,
+    table_name: str,
+    s3_url: str,
+    uri: str,
+    create_pyarrow_table: Callable[[Any], pa.Table],
+    catalog_name: str = 's3tablescatalog',
+    rest_signing_name: str = 's3tables',
+    rest_sigv4_enabled: str = 'true',
+    preserve_case: bool = False,
+) -> Dict:
+    """Import data from a file (CSV, Parquet, etc.) into an S3 table using a provided PyArrow table creation function."""
+    # Parse S3 URL
+    parsed = urlparse(s3_url)
+    bucket = parsed.netloc
+    key = parsed.path.lstrip('/')
+
+    try:
+        # Load Iceberg catalog
+        catalog = pyiceberg_load_catalog(
+            catalog_name,
+            warehouse,
+            uri,
+            region,
+            rest_signing_name,
+            rest_sigv4_enabled,
+        )
+
+        # Get S3 client and read the file
+        s3_client = get_s3_client()
+        response = s3_client.get_object(Bucket=bucket, Key=key)
+        file_bytes = response['Body'].read()
+
+        # Create PyArrow Table and Schema (file-like interface)
+        file_like = BytesIO(file_bytes)
+        pyarrow_table = create_pyarrow_table(file_like)
+        pyarrow_schema = pyarrow_table.schema
+
+        # Convert column names to snake_case unless preserve_case is True
+        columns_converted = False
+        if not preserve_case:
+            try:
+                pyarrow_schema = convert_column_names_to_snake_case(pyarrow_schema)
+                pyarrow_table = pyarrow_table.rename_columns(pyarrow_schema.names)
+                columns_converted = True
+            except Exception as conv_err:
+                return {
+                    'status': 'error',
+                    'error': f'Column name conversion failed: {str(conv_err)}',
+                }
+
+        table_created = False
+        try:
+            # Try to load existing table
+            table = catalog.load_table(f'{namespace}.{table_name}')
+        except NoSuchTableError:
+            # Table doesn't exist, create it using the schema
+            try:
+                table = catalog.create_table(
+                    identifier=f'{namespace}.{table_name}',
+                    schema=pyarrow_schema,
+                )
+                table_created = True
+            except Exception as create_error:
+                return {
+                    'status': 'error',
+                    'error': f'Failed to create table: {str(create_error)}',
+                }
+
+        # Append data to Iceberg table
+        table.append(pyarrow_table)
+
+        # Build message with warnings if applicable
+        message = f'Successfully imported {pyarrow_table.num_rows} rows{" and created new table" if table_created else ""}'
+        if columns_converted:
+            message += '. WARNING: Column names were converted to snake_case format. To preserve the original case, set preserve_case to True.'
+
+        return {
+            'status': 'success',
+            'message': message,
+            'rows_processed': pyarrow_table.num_rows,
+            'file_processed': os.path.basename(key),
+            'table_created': table_created,
+            'table_uuid': table.metadata.table_uuid,
+            'columns': pyarrow_schema.names,
+        }
+
+    except Exception as e:
+        return {'status': 'error', 'error': str(e)}

awslabs/s3_tables_mcp_server/server.py

@@ -589,6 +589,9 @@ async def import_csv_to_table(
         str, Field('s3tables', description='REST signing name')
     ] = 's3tables',
     rest_sigv4_enabled: Annotated[str, Field('true', description='Enable SigV4 signing')] = 'true',
+    preserve_case: Annotated[
+        bool, Field(..., description='Preserve case of column names')
+    ] = False,
 ) -> dict:
     """Import data from a CSV file into an S3 table.
 
@@ -596,6 +599,7 @@ async def import_csv_to_table(
     If the table doesn't exist, it will be created with a schema inferred from the CSV file.
     If the table exists, the CSV file schema must be compatible with the table's schema.
     The tool will validate the schema before attempting to import the data.
+    If preserve_case is True, the column names will not be converted to snake_case. Otherwise, the column names will be converted to snake_case.
 
     Returns error dictionary with status and error message if:
         - URL is not a valid S3 URL
@@ -615,6 +619,7 @@ async def import_csv_to_table(
         catalog_name: 's3tablescatalog'
         rest_signing_name: 's3tables'
         rest_sigv4_enabled: 'true'
+        preserve_case: False
 
     Permissions:
     You must have:
@@ -634,6 +639,7 @@ async def import_csv_to_table(
         catalog_name=catalog_name,
         rest_signing_name=rest_signing_name,
         rest_sigv4_enabled=rest_sigv4_enabled,
+        preserve_case=preserve_case,
     )
 
 
@@ -656,6 +662,9 @@ async def import_parquet_to_table(
         str, Field('s3tables', description='REST signing name')
     ] = 's3tables',
     rest_sigv4_enabled: Annotated[str, Field('true', description='Enable SigV4 signing')] = 'true',
+    preserve_case: Annotated[
+        bool, Field(..., description='Preserve case of column names')
+    ] = False,
 ) -> dict:
     """Import data from a Parquet file into an S3 table.
 
@@ -663,6 +672,7 @@ async def import_parquet_to_table(
     If the table doesn't exist, it will be created with a schema inferred from the Parquet file.
     If the table exists, the Parquet file schema must be compatible with the table's schema.
     The tool will validate the schema before attempting to import the data.
+    If preserve_case is True, the column names will not be converted to snake_case. Otherwise, the column names will be converted to snake_case.
 
     Returns error dictionary with status and error message if:
         - URL is not a valid S3 URL
@@ -688,6 +698,7 @@ async def import_parquet_to_table(
         catalog_name: 's3tablescatalog'
         rest_signing_name: 's3tables'
         rest_sigv4_enabled: 'true'
+        preserve_case: False
 
     Permissions:
     You must have:
@@ -708,6 +719,7 @@ async def import_parquet_to_table(
         catalog_name=catalog_name,
         rest_signing_name=rest_signing_name,
         rest_sigv4_enabled=rest_sigv4_enabled,
+        preserve_case=preserve_case,
     )
 
 

{awslabs_s3_tables_mcp_server-0.0.3.dist-info → awslabs_s3_tables_mcp_server-0.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: awslabs.s3-tables-mcp-server
-Version: 0.0.3
+Version: 0.0.5
 Summary: An AWS Labs Model Context Protocol (MCP) server for awslabs.s3-tables-mcp-server
 Project-URL: homepage, https://awslabs.github.io/mcp/
 Project-URL: docs, https://awslabs.github.io/mcp/servers/s3-tables-mcp-server/

{awslabs_s3_tables_mcp_server-0.0.3.dist-info → awslabs_s3_tables_mcp_server-0.0.5.dist-info}/RECORD

@@ -1,23 +1,24 @@
 awslabs/__init__.py,sha256=BHp8_uaBohApyLlmVWvYVe5bSrH59FvLJ5cNTigMV_8,644
-awslabs/s3_tables_mcp_server/__init__.py,sha256=T0dJ6usOanvF7UeP7NAt_YhLunvmWkSXvnjOvbL467Y,754
+awslabs/s3_tables_mcp_server/__init__.py,sha256=M4Nr7HtAc0wvjrfj0MUcaHdA1ljUTxNFYPdSQ__i6EQ,754
 awslabs/s3_tables_mcp_server/constants.py,sha256=qCWY9A9PAQXdVz-anO26zbQ72Dp79nGM7xeLR062a_o,4971
 awslabs/s3_tables_mcp_server/database.py,sha256=YorxcSx-9typfQ5W_LzwNPZkP47u__QSLJlp0fBsZLg,3851
 awslabs/s3_tables_mcp_server/models.py,sha256=zWTFJLBhIZRLEgOCTyNcGvbItxqYbFJKH6se1EzXDjY,8097
 awslabs/s3_tables_mcp_server/namespaces.py,sha256=KZqxJiEnlpxkqvbfygezbr0szwyDP2O0J6osyiPUzwg,2071
 awslabs/s3_tables_mcp_server/resources.py,sha256=PXZo0sTVn34tXJ4mlw_OS90p12SNoLZs4Re0gV815wk,8281
 awslabs/s3_tables_mcp_server/s3_operations.py,sha256=Zq3oe-uHuKbW87b_WQyM-6HZ0_ikbgiagb2SVesltdg,1656
-awslabs/s3_tables_mcp_server/server.py,sha256=cvXDTZuK1sGpYfjLbF6iLGe49BSA0yx4rSp73UEBcvE,32008
+awslabs/s3_tables_mcp_server/server.py,sha256=vNVDpX7axJAmrLaBKoNGt8RG6_8n3iYEHH-aBmbBEbc,32660
 awslabs/s3_tables_mcp_server/table_buckets.py,sha256=JHmpB_P9h0Hz5Uis25_GPTD1G-mIODVwjaswwIGyCS4,4471
 awslabs/s3_tables_mcp_server/tables.py,sha256=ITnRDHHrtRWLsRhff4TP4B7gGT_jRXy994oxK3x10a4,10143
 awslabs/s3_tables_mcp_server/utils.py,sha256=SReyS3KsdikI9ycL5RsvtVI7MiRnA1W9bTiXGKf1lHc,4517
 awslabs/s3_tables_mcp_server/engines/__init__.py,sha256=O4wlFva3THWmjfaXfJAwi29mxJSKIhM0jcebVfd3S5U,615
 awslabs/s3_tables_mcp_server/engines/pyiceberg.py,sha256=wzkySQZgx7L8Fn2Oqk8Yz4V-hQZDUempd8q0IwvCY_4,5784
 awslabs/s3_tables_mcp_server/file_processor/__init__.py,sha256=8PeggFRY3ZKBdxcFPEqSSHkSJBZ57eOs-z0fqkMHn9E,978
-awslabs/s3_tables_mcp_server/file_processor/csv.py,sha256=Sngc5mfJDLxQaINBUJLBn5OLc842rv9FqqcJ1upK6iw,4406
-awslabs/s3_tables_mcp_server/file_processor/parquet.py,sha256=Lr7mtqsK9jqlWokQv74dgdEgYmNKlCJ869yNNMrm69o,4189
-awslabs_s3_tables_mcp_server-0.0.3.dist-info/METADATA,sha256=DYOFGTR6IgR7l2ZciRbZWF4yjQ5FSTL8ilpBdmKMHFY,11511
-awslabs_s3_tables_mcp_server-0.0.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-awslabs_s3_tables_mcp_server-0.0.3.dist-info/entry_points.txt,sha256=WRA45Bi2dVY5hskxkka_e7BAGRqG1KiW3ImTBnHSyLs,90
-awslabs_s3_tables_mcp_server-0.0.3.dist-info/licenses/LICENSE,sha256=CeipvOyAZxBGUsFoaFqwkx54aPnIKEtm9a5u2uXxEws,10142
-awslabs_s3_tables_mcp_server-0.0.3.dist-info/licenses/NOTICE,sha256=jm-1A_8i-wl7KYs2Ynj2A29vXWJCMKLHmGfy4P_B51Y,96
-awslabs_s3_tables_mcp_server-0.0.3.dist-info/RECORD,,
+awslabs/s3_tables_mcp_server/file_processor/csv.py,sha256=1JfFFKbnK1qvDvpOKWoXIv9zca8MRlPXk2Phd7-GgIU,1650
+awslabs/s3_tables_mcp_server/file_processor/parquet.py,sha256=TA9HmVPAnwOc3hv3zsL7J7sRd8gMINGFp5_r1eX7uUE,1458
+awslabs/s3_tables_mcp_server/file_processor/utils.py,sha256=MkNE0c1tlR_bs5ZeCTqAuRCLQWOc3mJqnJxATbGwoEc,5613
+awslabs_s3_tables_mcp_server-0.0.5.dist-info/METADATA,sha256=fLAINA7r57U-LwJj_rSbCAioG2Q57F7Xeef4r0LlgEI,11511
+awslabs_s3_tables_mcp_server-0.0.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+awslabs_s3_tables_mcp_server-0.0.5.dist-info/entry_points.txt,sha256=WRA45Bi2dVY5hskxkka_e7BAGRqG1KiW3ImTBnHSyLs,90
+awslabs_s3_tables_mcp_server-0.0.5.dist-info/licenses/LICENSE,sha256=CeipvOyAZxBGUsFoaFqwkx54aPnIKEtm9a5u2uXxEws,10142
+awslabs_s3_tables_mcp_server-0.0.5.dist-info/licenses/NOTICE,sha256=jm-1A_8i-wl7KYs2Ynj2A29vXWJCMKLHmGfy4P_B51Y,96
+awslabs_s3_tables_mcp_server-0.0.5.dist-info/RECORD,,