Implement Key.compare_to_proto to check pb keys against existing.

Addresses sixth part of googleapis#451.

gcloud/datastore/key.py

@@ -153,6 +153,92 @@ def completed_key(self, id_or_name):
         new_key._flat_path += (id_or_name,)
         return new_key
 
+    def _validate_protobuf_dataset_id(self, protobuf):
+        """Checks that dataset ID on protobuf matches current one.
+
+        The value of the dataset ID may have changed from unprefixed
+        (e.g. 'foo') to prefixed (e.g. 's~foo' or 'e~foo').
+
+        :type protobuf: :class:`gcloud.datastore.datastore_v1_pb2.Key`
+        :param protobuf: A protobuf representation of the key. Expected to be
+                         returned after a datastore operation.
+
+        :rtype: :class:`str`
+        """
+        proto_dataset_id = protobuf.partition_id.dataset_id
+        if proto_dataset_id == self.dataset_id:
+            return
+
+        # Since they don't match, we check to see if `proto_dataset_id` has a
+        # prefix.
+        unprefixed = None
+        prefix = proto_dataset_id[:2]
+        if prefix in ('s~', 'e~'):
+            unprefixed = proto_dataset_id[2:]
+
+        if unprefixed != self.dataset_id:
+            raise ValueError('Dataset ID on protobuf does not match.',
+                             proto_dataset_id, self.dataset_id)
+
+    def compare_to_proto(self, protobuf):
+        """Checks current key against a protobuf; updates if partial.
+
+        If the current key is partial, returns a new key that has been
+        completed otherwise returns the current key.
+
+        The value of the dataset ID may have changed from implicit (i.e. None,
+        with the ID implied from the dataset.Dataset object associated with the
+        Entity/Key), but if it was implicit before, we leave it as implicit.
+
+        :type protobuf: :class:`gcloud.datastore.datastore_v1_pb2.Key`
+        :param protobuf: A protobuf representation of the key. Expected to be
+                         returned after a datastore operation.
+
+        :rtype: :class:`gcloud.datastore.key.Key`
+        :returns: The current key if not partial.
+        :raises: `ValueError` if the namespace or dataset ID of `protobuf`
+                 don't match the current values or if the path from `protobuf`
+                 doesn't match.
+        """
+        if self.namespace is None:
+            if protobuf.partition_id.HasField('namespace'):
+                raise ValueError('Namespace unset on key but set on protobuf.')
+        elif protobuf.partition_id.namespace != self.namespace:
+            raise ValueError('Namespace on protobuf does not match.',
+                             protobuf.partition_id.namespace, self.namespace)
+
+        # Check that dataset IDs match if not implicit.
+        if self.dataset_id is not None:
+            self._validate_protobuf_dataset_id(protobuf)
+
+        path = []
+        for element in protobuf.path_element:
+            key_part = {}
+            for descriptor, value in element._fields.items():
+                key_part[descriptor.name] = value
+            path.append(key_part)
+
+        if path == self.path:
+            return self
+
+        if not self.is_partial:
+            raise ValueError('Proto path does not match completed key.',
+                             path, self.path)
+
+        last_part = path[-1]
+        id_or_name = None
+        if 'id' in last_part:
+            id_or_name = last_part.pop('id')
+        elif 'name' in last_part:
+            id_or_name = last_part.pop('name')
+
+        # We have edited path by popping from the last part, so check again.
+        if path != self.path:
+            raise ValueError('Proto path does not match partial key.',
+                             path, self.path)
+
+        return self.complete_key(id_or_name)
+
     def to_protobuf(self):
         """Return a protobuf corresponding to the key.
 

gcloud/datastore/test_key.py

@@ -86,6 +86,87 @@ def test_completed_key_on_complete(self):
         key = self._makeOne('KIND', 1234)
         self.assertRaises(ValueError, key.completed_key, 5678)
 
+    def test_compare_to_proto_incomplete_w_id(self):
+        _ID = 1234
+        key = self._makeOne('KIND')
+        pb = key.to_protobuf()
+        pb.path_element[0].id = _ID
+        new_key = key.compare_to_proto(pb)
+        self.assertFalse(new_key is key)
+        self.assertEqual(new_key.id, _ID)
+        self.assertEqual(new_key.name, None)
+
+    def test_compare_to_proto_incomplete_w_name(self):
+        _NAME = 'NAME'
+        key = self._makeOne('KIND')
+        pb = key.to_protobuf()
+        pb.path_element[0].name = _NAME
+        new_key = key.compare_to_proto(pb)
+        self.assertFalse(new_key is key)
+        self.assertEqual(new_key.id, None)
+        self.assertEqual(new_key.name, _NAME)
+
+    def test_compare_to_proto_incomplete_w_incomplete(self):
+        key = self._makeOne('KIND')
+        pb = key.to_protobuf()
+        new_key = key.compare_to_proto(pb)
+        self.assertTrue(new_key is key)
+
+    def test_compare_to_proto_incomplete_w_bad_path(self):
+        key = self._makeOne('KIND1', 1234, 'KIND2')
+        pb = key.to_protobuf()
+        pb.path_element[0].kind = 'NO_KIND'
+        self.assertRaises(ValueError, key.compare_to_proto, pb)
+
+    def test_compare_to_proto_complete_w_id(self):
+        key = self._makeOne('KIND', 1234)
+        pb = key.to_protobuf()
+        pb.path_element[0].id = 5678
+        self.assertRaises(ValueError, key.compare_to_proto, pb)
+
+    def test_compare_to_proto_complete_w_name(self):
+        key = self._makeOne('KIND', 1234)
+        pb = key.to_protobuf()
+        pb.path_element[0].name = 'NAME'
+        self.assertRaises(ValueError, key.compare_to_proto, pb)
+
+    def test_compare_to_proto_complete_w_incomplete(self):
+        key = self._makeOne('KIND', 1234)
+        pb = key.to_protobuf()
+        pb.path_element[0].ClearField('id')
+        self.assertRaises(ValueError, key.compare_to_proto, pb)
+
+    def test_compare_to_proto_complete_diff_dataset(self):
+        key = self._makeOne('KIND', 1234, dataset_id='DATASET')
+        pb = key.to_protobuf()
+        pb.partition_id.dataset_id = 's~' + key.dataset_id
+        new_key = key.compare_to_proto(pb)
+        self.assertTrue(new_key is key)
+
+    def test_compare_to_proto_complete_bad_dataset(self):
+        key = self._makeOne('KIND', 1234, dataset_id='DATASET')
+        pb = key.to_protobuf()
+        pb.partition_id.dataset_id = 'BAD_PRE~' + key.dataset_id
+        self.assertRaises(ValueError, key.compare_to_proto, pb)
+
+    def test_compare_to_proto_complete_valid_namespace(self):
+        key = self._makeOne('KIND', 1234, namespace='NAMESPACE')
+        pb = key.to_protobuf()
+        new_key = key.compare_to_proto(pb)
+        self.assertTrue(new_key is key)
+
+    def test_compare_to_proto_complete_namespace_unset_on_pb(self):
+        key = self._makeOne('KIND', 1234, namespace='NAMESPACE')
+        pb = key.to_protobuf()
+        pb.partition_id.ClearField('namespace')
+        self.assertRaises(ValueError, key.compare_to_proto, pb)
+
+    def test_compare_to_proto_complete_namespace_unset_on_key(self):
+        key = self._makeOne('KIND', 1234)
+        pb = key.to_protobuf()
+        pb.partition_id.namespace = 'NAMESPACE'
+        self.assertRaises(ValueError, key.compare_to_proto, pb)
+
     def test_to_protobuf_defaults(self):
         from gcloud.datastore.datastore_v1_pb2 import Key as KeyPB
         _KIND = 'KIND'