Merge pull request #1134 from brianrower/feature/add-operation-with-update-expression

Add support for the ADD and DELETE operations in dynamodb update expressions
This commit is contained in:
Jack Danger 2017-09-13 15:24:55 -07:00 committed by GitHub
commit 2f997a7349
3 changed files with 220 additions and 1 deletions

View File

@ -57,7 +57,7 @@ class DynamoType(object):
@property
def cast_value(self):
if self.type == 'N':
if self.is_number():
try:
return int(self.value)
except ValueError:
@ -76,6 +76,15 @@ class DynamoType(object):
comparison_func = get_comparison_func(range_comparison)
return comparison_func(self.cast_value, *range_values)
def is_number(self):
return self.type == 'N'
def is_set(self):
return self.type == 'SS' or self.type == 'NS' or self.type == 'BS'
def same_type(self, other):
return self.type == other.type
class Item(BaseModel):
@ -140,6 +149,55 @@ class Item(BaseModel):
self.attrs[key] = DynamoType(expression_attribute_values[value])
else:
self.attrs[key] = DynamoType({"S": value})
elif action == 'ADD':
key, value = value.split(" ", 1)
key = key.strip()
value_str = value.strip()
if value_str in expression_attribute_values:
dyn_value = DynamoType(expression_attribute_values[value])
else:
raise TypeError
# Handle adding numbers - value gets added to existing value,
# or added to 0 if it doesn't exist yet
if dyn_value.is_number():
existing = self.attrs.get(key, DynamoType({"N": '0'}))
if not existing.same_type(dyn_value):
raise TypeError()
self.attrs[key] = DynamoType({"N": str(
decimal.Decimal(existing.value) +
decimal.Decimal(dyn_value.value)
)})
# Handle adding sets - value is added to the set, or set is
# created with only this value if it doesn't exist yet
# New value must be of same set type as previous value
elif dyn_value.is_set():
existing = self.attrs.get(key, DynamoType({dyn_value.type: {}}))
if not existing.same_type(dyn_value):
raise TypeError()
new_set = set(existing.value).union(dyn_value.value)
self.attrs[key] = DynamoType({existing.type: list(new_set)})
else: # Number and Sets are the only supported types for ADD
raise TypeError
elif action == 'DELETE':
key, value = value.split(" ", 1)
key = key.strip()
value_str = value.strip()
if value_str in expression_attribute_values:
dyn_value = DynamoType(expression_attribute_values[value])
else:
raise TypeError
if not dyn_value.is_set():
raise TypeError
existing = self.attrs.get(key, None)
if existing:
if not existing.same_type(dyn_value):
raise TypeError
new_set = set(existing.value).difference(dyn_value.value)
self.attrs[key] = DynamoType({existing.type: list(new_set)})
else:
raise NotImplementedError('{} update action not yet supported'.format(action))

View File

@ -497,6 +497,9 @@ class DynamoHandler(BaseResponse):
except ValueError:
er = 'com.amazonaws.dynamodb.v20111205#ConditionalCheckFailedException'
return self.error(er)
except TypeError:
er = 'com.amazonaws.dynamodb.v20111205#ValidationException'
return self.error(er)
item_dict = item.to_json()
item_dict['ConsumedCapacityUnits'] = 0.5

View File

@ -5,6 +5,7 @@ from decimal import Decimal
import boto
import boto3
from boto3.dynamodb.conditions import Key
from botocore.exceptions import ClientError
import sure # noqa
from freezegun import freeze_time
from moto import mock_dynamodb2, mock_dynamodb2_deprecated
@ -1190,6 +1191,14 @@ def _create_table_with_range_key():
'AttributeName': 'subject',
'AttributeType': 'S'
},
{
'AttributeName': 'username',
'AttributeType': 'S'
},
{
'AttributeName': 'created',
'AttributeType': 'N'
}
],
ProvisionedThroughput={
'ReadCapacityUnits': 5,
@ -1392,6 +1401,155 @@ def test_update_item_with_expression():
'subject': '123',
})
@mock_dynamodb2
def test_update_item_add_with_expression():
table = _create_table_with_range_key()
item_key = {'forum_name': 'the-key', 'subject': '123'}
current_item = {
'forum_name': 'the-key',
'subject': '123',
'str_set': {'item1', 'item2', 'item3'},
'num_set': {1, 2, 3},
'num_val': 6
}
# Put an entry in the DB to play with
table.put_item(Item=current_item)
# Update item to add a string value to a string set
table.update_item(
Key=item_key,
UpdateExpression='ADD str_set :v',
ExpressionAttributeValues={
':v': {'item4'}
}
)
current_item['str_set'] = current_item['str_set'].union({'item4'})
dict(table.get_item(Key=item_key)['Item']).should.equal(current_item)
# Update item to add a num value to a num set
table.update_item(
Key=item_key,
UpdateExpression='ADD num_set :v',
ExpressionAttributeValues={
':v': {6}
}
)
current_item['num_set'] = current_item['num_set'].union({6})
dict(table.get_item(Key=item_key)['Item']).should.equal(current_item)
# Update item to add a value to a number value
table.update_item(
Key=item_key,
UpdateExpression='ADD num_val :v',
ExpressionAttributeValues={
':v': 20
}
)
current_item['num_val'] = current_item['num_val'] + 20
dict(table.get_item(Key=item_key)['Item']).should.equal(current_item)
# Attempt to add a number value to a string set, should raise Client Error
table.update_item.when.called_with(
Key=item_key,
UpdateExpression='ADD str_set :v',
ExpressionAttributeValues={
':v': 20
}
).should.have.raised(ClientError)
dict(table.get_item(Key=item_key)['Item']).should.equal(current_item)
# Attempt to add a number set to the string set, should raise a ClientError
table.update_item.when.called_with(
Key=item_key,
UpdateExpression='ADD str_set :v',
ExpressionAttributeValues={
':v': { 20 }
}
).should.have.raised(ClientError)
dict(table.get_item(Key=item_key)['Item']).should.equal(current_item)
# Attempt to update with a bad expression
table.update_item.when.called_with(
Key=item_key,
UpdateExpression='ADD str_set bad_value'
).should.have.raised(ClientError)
# Attempt to add a string value instead of a string set
table.update_item.when.called_with(
Key=item_key,
UpdateExpression='ADD str_set :v',
ExpressionAttributeValues={
':v': 'new_string'
}
).should.have.raised(ClientError)
@mock_dynamodb2
def test_update_item_delete_with_expression():
table = _create_table_with_range_key()
item_key = {'forum_name': 'the-key', 'subject': '123'}
current_item = {
'forum_name': 'the-key',
'subject': '123',
'str_set': {'item1', 'item2', 'item3'},
'num_set': {1, 2, 3},
'num_val': 6
}
# Put an entry in the DB to play with
table.put_item(Item=current_item)
# Update item to delete a string value from a string set
table.update_item(
Key=item_key,
UpdateExpression='DELETE str_set :v',
ExpressionAttributeValues={
':v': {'item2'}
}
)
current_item['str_set'] = current_item['str_set'].difference({'item2'})
dict(table.get_item(Key=item_key)['Item']).should.equal(current_item)
# Update item to delete a num value from a num set
table.update_item(
Key=item_key,
UpdateExpression='DELETE num_set :v',
ExpressionAttributeValues={
':v': {2}
}
)
current_item['num_set'] = current_item['num_set'].difference({2})
dict(table.get_item(Key=item_key)['Item']).should.equal(current_item)
# Try to delete on a number, this should fail
table.update_item.when.called_with(
Key=item_key,
UpdateExpression='DELETE num_val :v',
ExpressionAttributeValues={
':v': 20
}
).should.have.raised(ClientError)
dict(table.get_item(Key=item_key)['Item']).should.equal(current_item)
# Try to delete a string set from a number set
table.update_item.when.called_with(
Key=item_key,
UpdateExpression='DELETE num_set :v',
ExpressionAttributeValues={
':v': {'del_str'}
}
).should.have.raised(ClientError)
dict(table.get_item(Key=item_key)['Item']).should.equal(current_item)
# Attempt to update with a bad expression
table.update_item.when.called_with(
Key=item_key,
UpdateExpression='DELETE num_val badvalue'
).should.have.raised(ClientError)
@mock_dynamodb2
def test_boto3_query_gsi_range_comparison():