From 2cea05bd19b93078c918912dec3df45dbcfe8e51 Mon Sep 17 00:00:00 2001
From: Paul Craciunoiu <paul@craciunoiu.net>
Date: Mon, 14 Dec 2015 15:42:17 -0700
Subject: [PATCH] More support for GSI

* Support query via multiple KeyConditionExpressions
* Support GSI updates via table.update()
---
 moto/dynamodb2/models.py    | 32 ++++++++++++++++++++++++++++++++
 moto/dynamodb2/responses.py | 31 +++++++++++++++++++++++++++----
 2 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/moto/dynamodb2/models.py b/moto/dynamodb2/models.py
index 66cc7a810..aab9f8397 100644
--- a/moto/dynamodb2/models.py
+++ b/moto/dynamodb2/models.py
@@ -404,6 +404,38 @@ class DynamoDBBackend(BaseBackend):
         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:
diff --git a/moto/dynamodb2/responses.py b/moto/dynamodb2/responses.py
index 7e703c896..20c45560d 100644
--- a/moto/dynamodb2/responses.py
+++ b/moto/dynamodb2/responses.py
@@ -123,8 +123,11 @@ class DynamoHandler(BaseResponse):
 
     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 @@ class DynamoHandler(BaseResponse):
         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: