getmoto / py-partiql-parser

Python Parser for PartiQL
MIT License
7 stars 2 forks source link

Case sensitive DynamoDB table name causes PartiQL parser to raise KeyError #23

Open jliu0812 opened 1 week ago

jliu0812 commented 1 week ago

When a DynamoDB table has mixed cases, the parser throws a KeyError message.

How to reproduce: Run the following code:

import boto3
from moto import mock_aws

table_name = 'CaseSensitiveTableName'

with mock_aws() as moto:
    # Initialize table
    dynamodb = boto3.client('dynamodb', region_name='us-west-2')
    table_obj = boto3.resource('dynamodb', region_name='us-west-2').Table(table_name)
    dynamodb.create_table(
        AttributeDefinitions=[
            {
                'AttributeName': 'pk',
                'AttributeType': 'S'
            },
            {
                'AttributeName': 'sk',
                'AttributeType': 'S'
            },
        ],
        TableName=table_name,
        KeySchema=[
            {
                'AttributeName': 'pk',
                'KeyType': 'HASH'
            },
            {
                'AttributeName': 'sk',
                'KeyType': 'RANGE'
            }
        ],
        BillingMode="PAY_PER_REQUEST",
    )
    dynamodb.get_waiter("table_exists").wait(TableName=table_name)

    # Put some data
    data = [
        {
            "pk": "test",
            "sk": "test1"
        },
        {
            "pk": "test2",
            "sk": "test2"
        }
    ]

    with table_obj.batch_writer() as writer:
        for item in data:
            writer.put_item(
                Item=item
            )

    # Query
    stmt = f"""SELECT * 
                      FROM "{table_name}" 
                   """

    response = dynamodb.execute_statement(
        Statement=stmt,
        ConsistentRead=False
    )

    print(response)

Traceback log:

...
File ".../build/venv/lib/python3.11/site-packages/py_partiql_parser/_internal/parser.py", line 108, in _parse_select
source_data = self.documents[list(from_parser.clauses.values())[0].lower()]
                  ~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 'casesensitivetablename'
jliu0812 commented 1 week ago

Testing locally, by changing the following line (https://github.com/getmoto/py-partiql-parser/blob/main/py_partiql_parser/_internal/parser.py#L108)

from

source_data = self.documents[list(from_parser.clauses.values())[0].lower()]

to

source_data = self.documents[list(from_parser.clauses.values())[0]]

Fixes the problem.

Expected result:

{'Items': [{'pk': {'S': 'test'}, 'sk': {'S': 'test1'}}, {'pk': {'S': 'test2'}, 'sk': {'S': 'test2'}}], 'ResponseMetadata': {'RequestId': 'ycN6ilzvV30CTS3W1us16BLMKf9iphbpg6sYPZVrdQ9uOoBJA4sJ', 'HTTPStatusCode': 200, 'HTTPHeaders': {'server': 'amazon.com', 'date': 'Thu, 05 Sep 2024 16:42:24 GMT', 'x-amzn-requestid': 'ycN6ilzvV30CTS3W1us16BLMKf9iphbpg6sYPZVrdQ9uOoBJA4sJ', 'x-amz-crc32': '2096729874'}, 'RetryAttempts': 0}}
bblommers commented 1 week ago

Hi @jliu0812, I think that was a holdover from when we were using the same approach for DynamoDB and S3, and in S3 the 'tablename' (i.e. s3object) is case-insensitive.

This is now fixed with 0.5.6, just released. Can you upgrade and verify it now works for you?