Re-allow use of dataset ID from keys as a fallback.

tseaver · tseaver · commit ae880b4b8eec · 2015-01-15T16:18:24.000-05:00
See: #552 (comment)
diff --git a/gcloud/datastore/api.py b/gcloud/datastore/api.py
@@ -103,7 +103,11 @@ def get(keys, missing=None, deferred=None, connection=None, dataset_id=None):
         return []
 
     connection = _require_connection(connection)
-    dataset_id = _require_dataset_id(dataset_id)
+    try:
+        dataset_id = _require_dataset_id(dataset_id)
+    except EnvironmentError:
+        dataset_id = keys[0].dataset_id
+
     if list(set([key.dataset_id for key in keys])) != [dataset_id]:
         raise ValueError('Keys do not match dataset ID')
 
@@ -157,7 +161,10 @@ def put(entities, connection=None, dataset_id=None):
         return
 
     connection = _require_connection(connection)
-    dataset_id = _require_dataset_id(dataset_id)
+    try:
+        dataset_id = _require_dataset_id(dataset_id)
+    except EnvironmentError:
+        dataset_id = entities[0].key.dataset_id
 
     current = Batch.current()
     in_batch = current is not None
@@ -192,7 +199,10 @@ def delete(keys, connection=None, dataset_id=None):
         return
 
     connection = _require_connection(connection)
-    dataset_id = _require_dataset_id(dataset_id)
+    try:
+        dataset_id = _require_dataset_id(dataset_id)
+    except EnvironmentError:
+        dataset_id = keys[0].dataset_id
 
     # We allow partial keys to attempt a delete, the backend will fail.
     current = Batch.current()
diff --git a/gcloud/datastore/test_api.py b/gcloud/datastore/test_api.py
@@ -148,30 +148,36 @@ def test_wo_connection(self):
         self.assertRaises(EnvironmentError,
                           self._callFUT, [key], dataset_id=DATASET_ID)
 
-    def test_wo_dataset_id(self):
+    def test_no_keys(self):
+        results = self._callFUT([])
+        self.assertEqual(results, [])
+
+    def test_miss(self):
         from gcloud.datastore.key import Key
         from gcloud.datastore.test_connection import _Connection
 
         DATASET_ID = 'DATASET'
         connection = _Connection()
         key = Key('Kind', 1234, dataset_id=DATASET_ID)
-        self.assertRaises(EnvironmentError,
-                          self._callFUT, [key], connection=connection)
-
-    def test_no_keys(self):
-        results = self._callFUT([])
+        results = self._callFUT([key], connection=connection,
+                                dataset_id=DATASET_ID)
         self.assertEqual(results, [])
 
-    def test_miss(self):
+    def test_miss_wo_dataset_id(self):
         from gcloud.datastore.key import Key
         from gcloud.datastore.test_connection import _Connection
 
         DATASET_ID = 'DATASET'
         connection = _Connection()
         key = Key('Kind', 1234, dataset_id=DATASET_ID)
-        results = self._callFUT([key], connection=connection,
-                                dataset_id=DATASET_ID)
+        results = self._callFUT([key], connection=connection)
         self.assertEqual(results, [])
+        expected = {
+            'dataset_id': DATASET_ID,
+            'key_pbs': [key.to_protobuf()],
+            'transaction_id': None,
+        }
+        self.assertEqual(connection._called_with, expected)
 
     def test_miss_w_missing(self):
         from gcloud.datastore import _datastore_v1_pb2 as datastore_pb
@@ -398,8 +404,18 @@ def test_no_dataset_id(self):
         entity.key = _Key(_DATASET)
 
         self.assertEqual(_implicit_environ.CONNECTION, None)
-        with self.assertRaises(EnvironmentError):
-            self._callFUT([entity], connection=connection)
+        result = self._callFUT([entity], connection=connection)
+
+        self.assertEqual(result, None)
+        self.assertEqual(len(connection._committed), 1)
+        dataset_id, mutation = connection._committed[0]
+        self.assertEqual(dataset_id, _DATASET)
+        upserts = list(mutation.upsert)
+        self.assertEqual(len(upserts), 1)
+        self.assertEqual(upserts[0].key, entity.key.to_protobuf())
+        properties = list(upserts[0].property)
+        self.assertEqual(properties[0].name, 'foo')
+        self.assertEqual(properties[0].value.string_value, u'bar')
 
     def test_no_entities(self):
         from gcloud.datastore import _implicit_environ
@@ -517,8 +533,14 @@ def test_no_dataset_id(self):
         key = _Key(_DATASET)
 
         self.assertEqual(_implicit_environ.CONNECTION, None)
-        with self.assertRaises(EnvironmentError):
-            self._callFUT([key], connection=connection)
+
+        result = self._callFUT([key], connection=connection)
+
+        self.assertEqual(result, None)
+        self.assertEqual(len(connection._committed), 1)
+        dataset_id, mutation = connection._committed[0]
+        self.assertEqual(dataset_id, _DATASET)
+        self.assertEqual(list(mutation.delete), [key.to_protobuf()])
 
     def test_no_keys(self):
         from gcloud.datastore import _implicit_environ
