diff --git a/moto/dynamodb2/models.py b/moto/dynamodb2/models.py index 7525a43a9..7590ee1e1 100644 --- a/moto/dynamodb2/models.py +++ b/moto/dynamodb2/models.py @@ -3,6 +3,7 @@ from collections import defaultdict import datetime import decimal import json +import re from moto.compat import OrderedDict from moto.core import BaseBackend, BaseModel @@ -115,28 +116,31 @@ class Item(BaseModel): } def update(self, update_expression, expression_attribute_names, expression_attribute_values): - ACTION_VALUES = ['SET', 'set', 'REMOVE', 'remove'] - - action = None - for value in update_expression.split(): - if value in ACTION_VALUES: - # An action - action = value - continue - else: + # Update subexpressions are identifiable by the operator keyword, so split on that and + # get rid of the empty leading string. + parts = [p for p in re.split(r'\b(SET|REMOVE|ADD|DELETE)\b', update_expression) if p] + # make sure that we correctly found only operator/value pairs + assert len(parts) % 2 == 0, "Mismatched operators and values in update expression: '{}'".format(update_expression) + for action, valstr in zip(parts[:-1:2], parts[1::2]): + values = valstr.split(',') + for value in values: # A Real value - value = value.lstrip(":").rstrip(",") - for k, v in expression_attribute_names.items(): - value = value.replace(k, v) - if action == "REMOVE" or action == 'remove': - self.attrs.pop(value, None) - elif action == 'SET' or action == 'set': - key, value = value.split("=") - if value in expression_attribute_values: - self.attrs[key] = DynamoType( - expression_attribute_values[value]) + value = value.lstrip(":").rstrip(",").strip() + for k, v in expression_attribute_names.items(): + value = re.sub(r'{0}\b'.format(k), v, value) + + if action == "REMOVE": + self.attrs.pop(value, None) + elif action == 'SET': + key, value = value.split("=") + key = key.strip() + value = value.strip() + if value in expression_attribute_values: + self.attrs[key] = DynamoType(expression_attribute_values[value]) + else: + self.attrs[key] = DynamoType({"S": value}) else: - self.attrs[key] = DynamoType({"S": value}) + raise NotImplementedError('{} update action not yet supported'.format(action)) def update_with_attribute_updates(self, attribute_updates): for attribute_name, update_action in attribute_updates.items(): @@ -345,7 +349,6 @@ class Table(BaseModel): def query(self, hash_key, range_comparison, range_objs, limit, exclusive_start_key, scan_index_forward, index_name=None, **filter_kwargs): results = [] - if index_name: all_indexes = (self.global_indexes or []) + (self.indexes or []) indexes_by_name = dict((i['IndexName'], i) for i in all_indexes) diff --git a/moto/dynamodb2/responses.py b/moto/dynamodb2/responses.py index 1d9b70043..bf24e9964 100644 --- a/moto/dynamodb2/responses.py +++ b/moto/dynamodb2/responses.py @@ -316,24 +316,26 @@ class DynamoHandler(BaseResponse): else: index = table.schema - key_map = [column for _, column in sorted( - (k, v) for k, v in self.body['ExpressionAttributeNames'].items())] + reverse_attribute_lookup = dict((v, k) for k, v in + six.iteritems(self.body['ExpressionAttributeNames'])) if " AND " in key_condition_expression: expressions = key_condition_expression.split(" AND ", 1) - index_hash_key = [ - key for key in index if key['KeyType'] == 'HASH'][0] - hash_key_index_in_key_map = key_map.index( - index_hash_key['AttributeName']) + index_hash_key = [key for key in index if key['KeyType'] == 'HASH'][0] + hash_key_var = reverse_attribute_lookup.get(index_hash_key['AttributeName'], + index_hash_key['AttributeName']) + hash_key_regex = r'(^|[\s(]){0}\b'.format(hash_key_var) + i, hash_key_expression = next((i, e) for i, e in enumerate(expressions) + if re.search(hash_key_regex, e)) + hash_key_expression = hash_key_expression.strip('()') + expressions.pop(i) - hash_key_expression = expressions.pop( - hash_key_index_in_key_map).strip('()') - # TODO implement more than one range expression and OR - # operators + # TODO implement more than one range expression and OR operators range_key_expression = expressions[0].strip('()') range_key_expression_components = range_key_expression.split() range_comparison = range_key_expression_components[1] + if 'AND' in range_key_expression: range_comparison = 'BETWEEN' range_values = [