Merge pull request #1301 from dhermes/datastore-datastore-file-shim

Replacing datastore pb uses with "datastore" shim.

Author: dhermes

gcloud/datastore/_datastore_pb2.py

@@ -0,0 +1,38 @@
+# Copyright 2015 Google Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Datastore shim to emulate v1beta3 module structure.
+
+This module intended to pair with datastore.proto.
+"""
+
+from gcloud.datastore import _datastore_v1_pb2
+
+
+LookupRequest = _datastore_v1_pb2.LookupRequest
+LookupResponse = _datastore_v1_pb2.LookupResponse
+RunQueryRequest = _datastore_v1_pb2.RunQueryRequest
+RunQueryResponse = _datastore_v1_pb2.RunQueryResponse
+BeginTransactionRequest = _datastore_v1_pb2.BeginTransactionRequest
+BeginTransactionResponse = _datastore_v1_pb2.BeginTransactionResponse
+RollbackRequest = _datastore_v1_pb2.RollbackRequest
+RollbackResponse = _datastore_v1_pb2.RollbackResponse
+CommitRequest = _datastore_v1_pb2.CommitRequest
+CommitResponse = _datastore_v1_pb2.CommitResponse
+AllocateIdsRequest = _datastore_v1_pb2.AllocateIdsRequest
+AllocateIdsResponse = _datastore_v1_pb2.AllocateIdsResponse
+Mutation = _datastore_v1_pb2.Mutation
+MutationResult = _datastore_v1_pb2.MutationResult
+ReadOptions = _datastore_v1_pb2.ReadOptions
+Property = _datastore_v1_pb2.Property  # Not present in v1beta3

gcloud/datastore/batch.py

@@ -23,7 +23,7 @@
 
 from gcloud.datastore import helpers
 from gcloud.datastore.key import _dataset_ids_equal
-from gcloud.datastore import _datastore_v1_pb2 as datastore_pb
+from gcloud.datastore import _datastore_pb2
 
 
 class Batch(object):
@@ -64,7 +64,7 @@ class Batch(object):
 
     def __init__(self, client):
         self._client = client
-        self._mutation = datastore_pb.Mutation()
+        self._mutation = _datastore_pb2.Mutation()
         self._partial_key_entities = []
 
     def current(self):
@@ -109,7 +109,7 @@ def mutation(self):
         This getter returns the Mutation protobuf that
         has been built-up so far.
 
-        :rtype: :class:`gcloud.datastore._datastore_v1_pb2.Mutation`
+        :rtype: :class:`gcloud.datastore._datastore_pb2.Mutation`
         :returns: The Mutation protobuf to be sent in the commit request.
         """
         return self._mutation
@@ -223,7 +223,7 @@ def _assign_entity_to_pb(entity_pb, entity):
 
     Helper method for ``Batch.put``.
 
-    :type entity_pb: :class:`gcloud.datastore._datastore_v1_pb2.Entity`
+    :type entity_pb: :class:`gcloud.datastore._entity_pb2.Entity`
     :param entity_pb: The entity owned by a mutation.
 
     :type entity: :class:`gcloud.datastore.entity.Entity`

gcloud/datastore/connection.py

@@ -19,7 +19,7 @@
 from gcloud import connection
 from gcloud.environment_vars import GCD_HOST
 from gcloud.exceptions import make_exception
-from gcloud.datastore import _datastore_v1_pb2 as datastore_pb
+from gcloud.datastore import _datastore_pb2
 from gcloud.datastore import _entity_pb2
 
 
@@ -188,12 +188,12 @@ def lookup(self, dataset_id, key_pbs,
                   ``deferred`` is a list of
                   :class:`gcloud.datastore._entity_pb2.Key`.
         """
-        lookup_request = datastore_pb.LookupRequest()
+        lookup_request = _datastore_pb2.LookupRequest()
         _set_read_options(lookup_request, eventual, transaction_id)
         _add_keys_to_request(lookup_request.key, key_pbs)
 
         lookup_response = self._rpc(dataset_id, 'lookup', lookup_request,
-                                    datastore_pb.LookupResponse)
+                                    _datastore_pb2.LookupResponse)
 
         results = [result.entity for result in lookup_response.found]
         missing = [result.entity for result in lookup_response.missing]
@@ -255,15 +255,15 @@ def run_query(self, dataset_id, query_pb, namespace=None,
                                the given transaction.  Incompatible with
                                ``eventual==True``.
         """
-        request = datastore_pb.RunQueryRequest()
+        request = _datastore_pb2.RunQueryRequest()
         _set_read_options(request, eventual, transaction_id)
 
         if namespace:
             request.partition_id.namespace = namespace
 
         request.query.CopyFrom(query_pb)
         response = self._rpc(dataset_id, 'runQuery', request,
-                             datastore_pb.RunQueryResponse)
+                             _datastore_pb2.RunQueryResponse)
         return (
             [e.entity for e in response.batch.entity_result],
             response.batch.end_cursor,  # Assume response always has cursor.
@@ -279,14 +279,14 @@ def begin_transaction(self, dataset_id):
         :type dataset_id: string
         :param dataset_id: The ID dataset to which the transaction applies.
 
-        :rtype: :class:`._datastore_v1_pb2.BeginTransactionResponse`
+        :rtype: :class:`._datastore_pb2.BeginTransactionResponse`
         :returns': the result protobuf for the begin transaction request.
         """
-        request = datastore_pb.BeginTransactionRequest()
+        request = _datastore_pb2.BeginTransactionRequest()
         request.isolation_level = (
-            datastore_pb.BeginTransactionRequest.SERIALIZABLE)
+            _datastore_pb2.BeginTransactionRequest.SERIALIZABLE)
         response = self._rpc(dataset_id, 'beginTransaction', request,
-                             datastore_pb.BeginTransactionResponse)
+                             _datastore_pb2.BeginTransactionResponse)
         return response.transaction
 
     def commit(self, dataset_id, mutation_pb, transaction_id):
@@ -297,28 +297,28 @@ def commit(self, dataset_id, mutation_pb, transaction_id):
         :type dataset_id: string
         :param dataset_id: The ID dataset to which the transaction applies.
 
-        :type mutation_pb: :class:`datastore_pb.Mutation`.
+        :type mutation_pb: :class:`._datastore_pb2.Mutation`.
         :param mutation_pb: The protobuf for the mutations being saved.
 
         :type transaction_id: string or None
         :param transaction_id: The transaction ID returned from
                                :meth:`begin_transaction`.  Non-transactional
                                batches must pass ``None``.
 
-        :rtype: :class:`gcloud.datastore._datastore_v1_pb2.MutationResult`.
+        :rtype: :class:`gcloud.datastore._datastore_pb2.MutationResult`.
         :returns': the result protobuf for the mutation.
         """
-        request = datastore_pb.CommitRequest()
+        request = _datastore_pb2.CommitRequest()
 
         if transaction_id:
-            request.mode = datastore_pb.CommitRequest.TRANSACTIONAL
+            request.mode = _datastore_pb2.CommitRequest.TRANSACTIONAL
             request.transaction = transaction_id
         else:
-            request.mode = datastore_pb.CommitRequest.NON_TRANSACTIONAL
+            request.mode = _datastore_pb2.CommitRequest.NON_TRANSACTIONAL
 
         request.mutation.CopyFrom(mutation_pb)
         response = self._rpc(dataset_id, 'commit', request,
-                             datastore_pb.CommitResponse)
+                             _datastore_pb2.CommitResponse)
         return response.mutation_result
 
     def rollback(self, dataset_id, transaction_id):
@@ -334,11 +334,11 @@ def rollback(self, dataset_id, transaction_id):
         :param transaction_id: The transaction ID returned from
                                :meth:`begin_transaction`.
         """
-        request = datastore_pb.RollbackRequest()
+        request = _datastore_pb2.RollbackRequest()
         request.transaction = transaction_id
         # Nothing to do with this response, so just execute the method.
         self._rpc(dataset_id, 'rollback', request,
-                  datastore_pb.RollbackResponse)
+                  _datastore_pb2.RollbackResponse)
 
     def allocate_ids(self, dataset_id, key_pbs):
         """Obtain backend-generated IDs for a set of keys.
@@ -355,11 +355,11 @@ def allocate_ids(self, dataset_id, key_pbs):
         :rtype: list of :class:`gcloud.datastore._entity_pb2.Key`
         :returns: An equal number of keys,  with IDs filled in by the backend.
         """
-        request = datastore_pb.AllocateIdsRequest()
+        request = _datastore_pb2.AllocateIdsRequest()
         _add_keys_to_request(request.key, key_pbs)
         # Nothing to do with this response, so just execute the method.
         response = self._rpc(dataset_id, 'allocateIds', request,
-                             datastore_pb.AllocateIdsResponse)
+                             _datastore_pb2.AllocateIdsResponse)
         return list(response.key)
 
 
@@ -376,7 +376,7 @@ def _set_read_options(request, eventual, transaction_id):
 
     opts = request.read_options
     if eventual:
-        opts.read_consistency = datastore_pb.ReadOptions.EVENTUAL
+        opts.read_consistency = _datastore_pb2.ReadOptions.EVENTUAL
     elif transaction_id:
         opts.transaction = transaction_id
 

gcloud/datastore/helpers.py

@@ -263,7 +263,7 @@ def _get_value_from_property_pb(property_pb):
     Some work is done to coerce the return value into a more useful type
     (particularly in the case of a timestamp value, or a key value).
 
-    :type property_pb: :class:`gcloud.datastore._datastore_v1_pb2.Property`
+    :type property_pb: :class:`gcloud.datastore._datastore_pb2.Property`
     :param property_pb: The Property Protobuf.
 
     :returns: The value provided by the Protobuf.

gcloud/datastore/test_batch.py

@@ -26,7 +26,7 @@ def _makeOne(self, client):
         return self._getTargetClass()(client)
 
     def test_ctor(self):
-        from gcloud.datastore._datastore_v1_pb2 import Mutation
+        from gcloud.datastore._datastore_pb2 import Mutation
         _DATASET = 'DATASET'
         _NAMESPACE = 'NAMESPACE'
         connection = _Connection()