From eff498d91dc951b48166995af1fdb5b4958d198c Mon Sep 17 00:00:00 2001
From: Tres Seaver <tseaver@palladion.com>
Date: Mon, 27 Oct 2014 17:21:46 -0400
Subject: [PATCH] Implement 'allocate_ids' RPC on Connection.

Fixes #272.
---
 gcloud/datastore/connection.py      | 20 +++++++++
 gcloud/datastore/test_connection.py | 66 +++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/gcloud/datastore/connection.py b/gcloud/datastore/connection.py
index 71efc538ec22..0139fc231faf 100644
--- a/gcloud/datastore/connection.py
+++ b/gcloud/datastore/connection.py
@@ -342,6 +342,26 @@ def rollback(self, dataset_id):
         self._rpc(dataset_id, 'rollback', request,
                   datastore_pb.RollbackResponse)
 
+    def allocate_ids(self, dataset_id, key_pbs):
+        """Obtain backend-generated IDs for a set of keys.
+
+        :type dataset_id: string
+        :param dataset_id: The dataset to which the transaction belongs.
+
+        :type key_pbs: list of :class:`gcloud.datastore.datastore_v1_pb2.Key`
+        :param key_pbs: The keys for which the backend should allocate IDs.
+
+        :rtype: list of :class:`gcloud.datastore.datastore_v1_pb2.Key`
+        :returns: An equal number of keys,  with IDs filled in by the backend.
+        """
+        request = datastore_pb.AllocateIdsRequest()
+        for key_pb in key_pbs:
+            request.key.add().CopyFrom(key_pb)
+        # Nothing to do with this response, so just execute the method.
+        response = self._rpc(dataset_id, 'allocateIds', request,
+                             datastore_pb.AllocateIdsResponse)
+        return list(response.key)
+
     #
     #   Entity-related helper methods.
     #
diff --git a/gcloud/datastore/test_connection.py b/gcloud/datastore/test_connection.py
index ce56dcb5f7c4..53db7d65a8a9 100644
--- a/gcloud/datastore/test_connection.py
+++ b/gcloud/datastore/test_connection.py
@@ -560,6 +560,72 @@ def id(self):
         request.ParseFromString(cw['body'])
         self.assertEqual(request.transaction, TRANSACTION)
 
+    def test_allocate_ids_empty(self):
+        from gcloud.datastore.connection import datastore_pb
+
+        DATASET_ID = 'DATASET'
+        rsp_pb = datastore_pb.AllocateIdsResponse()
+        conn = self._makeOne()
+        URI = '/'.join([
+            conn.API_BASE_URL,
+            'datastore',
+            conn.API_VERSION,
+            'datasets',
+            DATASET_ID,
+            'allocateIds',
+        ])
+        http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
+        self.assertEqual(conn.allocate_ids(DATASET_ID, []), [])
+        cw = http._called_with
+        self.assertEqual(cw['uri'], URI)
+        self.assertEqual(cw['method'], 'POST')
+        self.assertEqual(cw['headers']['Content-Type'],
+                         'application/x-protobuf')
+        self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT)
+        rq_class = datastore_pb.AllocateIdsRequest
+        request = rq_class()
+        request.ParseFromString(cw['body'])
+        self.assertEqual(list(request.key), [])
+
+    def test_allocate_ids_non_empty(self):
+        from gcloud.datastore.connection import datastore_pb
+        from gcloud.datastore.key import Key
+
+        DATASET_ID = 'DATASET'
+        before_key_pbs = [
+            Key(path=[{'kind': 'Kind'}]).to_protobuf(),
+            Key(path=[{'kind': 'Kind'}]).to_protobuf(),
+            ]
+        after_key_pbs = [
+            Key(path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf(),
+            Key(path=[{'kind': 'Kind', 'id': 2345}]).to_protobuf(),
+            ]
+        rsp_pb = datastore_pb.AllocateIdsResponse()
+        rsp_pb.key.add().CopyFrom(after_key_pbs[0])
+        rsp_pb.key.add().CopyFrom(after_key_pbs[1])
+        conn = self._makeOne()
+        URI = '/'.join([
+            conn.API_BASE_URL,
+            'datastore',
+            conn.API_VERSION,
+            'datasets',
+            DATASET_ID,
+            'allocateIds',
+        ])
+        http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
+        self.assertEqual(conn.allocate_ids(DATASET_ID, before_key_pbs),
+                         after_key_pbs)
+        cw = http._called_with
+        self.assertEqual(cw['uri'], URI)
+        self.assertEqual(cw['method'], 'POST')
+        self.assertEqual(cw['headers']['Content-Type'],
+                         'application/x-protobuf')
+        self.assertEqual(cw['headers']['User-Agent'], conn.USER_AGENT)
+        rq_class = datastore_pb.AllocateIdsRequest
+        request = rq_class()
+        request.ParseFromString(cw['body'])
+        self.assertEqual(list(request.key), before_key_pbs)
+
     def test_save_entity_wo_transaction_w_upsert(self):
         from gcloud.datastore.connection import datastore_pb
         from gcloud.datastore.key import Key