Add query index support for boto3 via dynamodb2

README.md

@@ -64,7 +64,7 @@ It gets even better! Moto isn't just S3. Here's the status of the other AWS serv
 | Data Pipeline         | @mock_datapipeline| basic endpoints done             |
 |------------------------------------------------------------------------------|
 | DynamoDB              | @mock_dynamodb   | core endpoints done               |
-| DynamoDB2             | @mock_dynamodb2  | core endpoints done - no indexes  |
+| DynamoDB2             | @mock_dynamodb2  | core endpoints + partial indexes  |
 |------------------------------------------------------------------------------|
 | EC2                   | @mock_ec2        | core endpoints done               |
 |     - AMI             |                  | core endpoints done               |

moto/dynamodb2/models.py

@@ -278,20 +278,63 @@ def delete_item(self, hash_key, range_key):
         except KeyError:
             return None
 
-    def query(self, hash_key, range_comparison, range_objs):
+    def query(self, hash_key, range_comparison, range_objs, index_name=None):
         results = []
         last_page = True  # Once pagination is implemented, change this
 
-        possible_results = [item for item in list(self.all_items()) if isinstance(item, Item) and item.hash_key == hash_key]
+        if index_name:
+            all_indexes = (self.global_indexes or []) + (self.indexes or [])
+            indexes_by_name = dict((i['IndexName'], i) for i in all_indexes)
+            if index_name not in indexes_by_name:
+                raise ValueError('Invalid index: %s for table: %s. Available indexes are: %s' % (
+                    index_name, self.name, ', '.join(indexes_by_name.keys())
+                ))
+
+            index = indexes_by_name[index_name]
+            try:
+                index_hash_key = [key for key in index['KeySchema'] if key['KeyType'] == 'HASH'][0]
+            except IndexError:
+                raise ValueError('Missing Hash Key. KeySchema: %s' % index['KeySchema'])
+
+            possible_results = []
+            for item in self.all_items():
+                if not isinstance(item, Item):
+                    continue
+                item_hash_key = item.attrs.get(index_hash_key['AttributeName'])
+                if item_hash_key and item_hash_key == hash_key:
+                    possible_results.append(item)
+        else:
+            possible_results = [item for item in list(self.all_items()) if isinstance(item, Item) and item.hash_key == hash_key]
+
+        if index_name:
+            try:
+                index_range_key = [key for key in index['KeySchema'] if key['KeyType'] == 'RANGE'][0]
+            except IndexError:
+                index_range_key = None
+
         if range_comparison:
-            for result in possible_results:
-                if result.range_key.compare(range_comparison, range_objs):
-                    results.append(result)
+            if index_name and not index_range_key:
+                raise ValueError('Range Key comparison but no range key found for index: %s' % index_name)
+
+            elif index_name:
+                for result in possible_results:
+                    if result.attrs.get(index_range_key['AttributeName']).compare(range_comparison, range_objs):
+                        results.append(result)
+            else:
+                for result in possible_results:
+                    if result.range_key.compare(range_comparison, range_objs):
+                        results.append(result)
         else:
             # If we're not filtering on range key, return all values
             results = possible_results
 
-        results.sort(key=lambda item: item.range_key)
+        if index_name:
+
+            if index_range_key:
+                results.sort(key=lambda item: item.attrs[index_range_key['AttributeName']].value
+                                                if item.attrs.get(index_range_key['AttributeName']) else None)
+        else:
+            results.sort(key=lambda item: item.range_key)
         return results, last_page
 
     def all_items(self):
@@ -361,6 +404,38 @@ def update_table_throughput(self, name, throughput):
         table.throughput = throughput
         return table
 
+    def update_table_global_indexes(self, name, global_index_updates):
+        table = self.tables[name]
+        gsis_by_name = dict((i['IndexName'], i) for i in table.global_indexes)
+        for gsi_update in global_index_updates:
+            gsi_to_create = gsi_update.get('Create')
+            gsi_to_update = gsi_update.get('Update')
+            gsi_to_delete = gsi_update.get('Delete')
+
+            if gsi_to_delete:
+                index_name = gsi_to_delete['IndexName']
+                if index_name not in gsis_by_name:
+                    raise ValueError('Global Secondary Index does not exist, but tried to delete: %s' %
+                                     gsi_to_delete['IndexName'])
+
+                del gsis_by_name[index_name]
+
+            if gsi_to_update:
+                index_name = gsi_to_update['IndexName']
+                if index_name not in gsis_by_name:
+                    raise ValueError('Global Secondary Index does not exist, but tried to update: %s' %
+                                     gsi_to_update['IndexName'])
+                gsis_by_name[index_name].update(gsi_to_update)
+
+            if gsi_to_create:
+                if gsi_to_create['IndexName'] in gsis_by_name:
+                    raise ValueError('Global Secondary Index already exists: %s' % gsi_to_create['IndexName'])
+
+                gsis_by_name[gsi_to_create['IndexName']] = gsi_to_create
+
+        table.global_indexes = gsis_by_name.values()
+        return table
+
     def put_item(self, table_name, item_attrs, expected=None, overwrite=False):
         table = self.tables.get(table_name)
         if not table:
@@ -400,15 +475,15 @@ def get_item(self, table_name, keys):
         hash_key, range_key = self.get_keys_value(table, keys)
         return table.get_item(hash_key, range_key)
 
-    def query(self, table_name, hash_key_dict, range_comparison, range_value_dicts):
+    def query(self, table_name, hash_key_dict, range_comparison, range_value_dicts, index_name=None):
         table = self.tables.get(table_name)
         if not table:
             return None, None
 
         hash_key = DynamoType(hash_key_dict)
         range_values = [DynamoType(range_value) for range_value in range_value_dicts]
 
-        return table.query(hash_key, range_comparison, range_values)
+        return table.query(hash_key, range_comparison, range_values, index_name)
 
     def scan(self, table_name, filters):
         table = self.tables.get(table_name)

moto/dynamodb2/responses.py

@@ -123,8 +123,11 @@ def delete_table(self):
 
     def update_table(self):
         name = self.body['TableName']
-        throughput = self.body["ProvisionedThroughput"]
-        table = dynamodb_backend2.update_table_throughput(name, throughput)
+        if 'GlobalSecondaryIndexUpdates' in self.body:
+            table = dynamodb_backend2.update_table_global_indexes(name, self.body['GlobalSecondaryIndexUpdates'])
+        if 'ProvisionedThroughput' in self.body:
+            throughput = self.body["ProvisionedThroughput"]
+            table = dynamodb_backend2.update_table_throughput(name, throughput)
         return dynamo_json_dump(table.describe)
 
     def describe_table(self):
@@ -239,11 +242,31 @@ def query(self):
         if key_condition_expression:
             value_alias_map = self.body['ExpressionAttributeValues']
 
+            table = dynamodb_backend2.get_table(name)
+            index_name = self.body.get('IndexName')
+            if index_name:
+                all_indexes = (table.global_indexes or []) + (table.indexes or [])
+                indexes_by_name = dict((i['IndexName'], i) for i in all_indexes)
+                if index_name not in indexes_by_name:
+                    raise ValueError('Invalid index: %s for table: %s. Available indexes are: %s' % (
+                        index_name, name, ', '.join(indexes_by_name.keys())
+                    ))
+
+                index = indexes_by_name[index_name]['KeySchema']
+            else:
+                index = table.schema
+
+            key_map = [column for _, column in sorted((k, v) for k, v in self.body['ExpressionAttributeNames'].items())]
+
             if " AND " in key_condition_expression:
                 expressions = key_condition_expression.split(" AND ", 1)
-                hash_key_expression = expressions[0]
+
+                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'])
+
+                hash_key_expression = expressions.pop(hash_key_index_in_key_map).strip('()')
                 # TODO implement more than one range expression and OR operators
-                range_key_expression = expressions[1].replace(")", "")
+                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:
@@ -291,7 +314,8 @@ def query(self):
                             range_comparison = None
                             range_values = []
 
-        items, last_page = dynamodb_backend2.query(name, hash_key, range_comparison, range_values)
+        index_name = self.body.get('IndexName')
+        items, last_page = dynamodb_backend2.query(name, hash_key, range_comparison, range_values, index_name=index_name)
         if items is None:
             er = 'com.amazonaws.dynamodb.v20111205#ResourceNotFoundException'
             return self.error(er)