From ca42fdf4c2585e79cabfe4f9dedc8613249ec5f3 Mon Sep 17 00:00:00 2001
From: Tres Seaver <tseaver@palladion.com>
Date: Fri, 10 Jul 2015 13:24:30 -0400
Subject: [PATCH] Update batch/transaction/query to hold client.

Update batch/transaction to push themselves onto client's stack, rather
than batches.

Updated API-invoking methods to take 'client', rather than 'connection',
falling back to the client held by self.

See #944.
---
 gcloud/datastore/query.py            |  4 ++--
 gcloud/datastore/test_transaction.py | 13 ++++++++-----
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/gcloud/datastore/query.py b/gcloud/datastore/query.py
index 4adfece8520d..1ae2c43dfa24 100644
--- a/gcloud/datastore/query.py
+++ b/gcloud/datastore/query.py
@@ -102,7 +102,7 @@ def dataset_id(self):
 
         :rtype: str
         """
-        return self._dataset_id
+        return self._dataset_id or self._client.dataset_id
 
     @property
     def namespace(self):
@@ -111,7 +111,7 @@ def namespace(self):
         :rtype: string or None
         :returns: the namespace assigned to this query
         """
-        return self._namespace
+        return self._namespace or self._client.namespace
 
     @namespace.setter
     def namespace(self, value):
diff --git a/gcloud/datastore/test_transaction.py b/gcloud/datastore/test_transaction.py
index 7d0f822a2e3a..fe8d38d657ea 100644
--- a/gcloud/datastore/test_transaction.py
+++ b/gcloud/datastore/test_transaction.py
@@ -26,6 +26,7 @@ def _makeOne(self, client):
 
     def test_ctor(self):
         from gcloud.datastore._datastore_v1_pb2 import Mutation
+        from gcloud.datastore.test_batch import _Client
 
         _DATASET = 'DATASET'
         connection = _Connection()
@@ -40,6 +41,7 @@ def test_ctor(self):
 
     def test_current(self):
         from gcloud.datastore.test_client import _NoCommitBatch
+        from gcloud.datastore.test_batch import _Client
         _DATASET = 'DATASET'
         connection = _Connection()
         client = _Client(_DATASET, connection)
@@ -65,6 +67,7 @@ def test_current(self):
         self.assertTrue(xact2.current() is None)
 
     def test_begin(self):
+        from gcloud.datastore.test_batch import _Client
         _DATASET = 'DATASET'
         connection = _Connection(234)
         client = _Client(_DATASET, connection)
@@ -74,6 +77,7 @@ def test_begin(self):
         self.assertEqual(connection._begun, _DATASET)
 
     def test_begin_tombstoned(self):
+        from gcloud.datastore.test_batch import _Client
         _DATASET = 'DATASET'
         connection = _Connection(234)
         client = _Client(_DATASET, connection)
@@ -88,6 +92,7 @@ def test_begin_tombstoned(self):
         self.assertRaises(ValueError, xact.begin)
 
     def test_rollback(self):
+        from gcloud.datastore.test_batch import _Client
         _DATASET = 'DATASET'
         connection = _Connection(234)
         client = _Client(_DATASET, connection)
@@ -98,6 +103,7 @@ def test_rollback(self):
         self.assertEqual(connection._rolled_back, (_DATASET, 234))
 
     def test_commit_no_auto_ids(self):
+        from gcloud.datastore.test_batch import _Client
         _DATASET = 'DATASET'
         connection = _Connection(234)
         client = _Client(_DATASET, connection)
@@ -109,6 +115,7 @@ def test_commit_no_auto_ids(self):
         self.assertEqual(xact.id, None)
 
     def test_commit_w_auto_ids(self):
+        from gcloud.datastore.test_batch import _Client
         _DATASET = 'DATASET'
         _KIND = 'KIND'
         _ID = 123
@@ -127,6 +134,7 @@ def test_commit_w_auto_ids(self):
         self.assertEqual(entity.key.path, [{'kind': _KIND, 'id': _ID}])
 
     def test_context_manager_no_raise(self):
+        from gcloud.datastore.test_batch import _Client
         _DATASET = 'DATASET'
         connection = _Connection(234)
         client = _Client(_DATASET, connection)
@@ -217,8 +225,3 @@ def _push_batch(self, batch):
 
     def _pop_batch(self):
         return self._batches.pop(0)
-
-    @property
-    def current_batch(self):
-        if self._batches:
-            return self._batches[0]