From dd99e2c3723819225e41c4a38cbe92bbf3f9d03c Mon Sep 17 00:00:00 2001
From: Tres Seaver <tseaver@palladion.com>
Date: Fri, 10 Jul 2015 13:40:38 -0400
Subject: [PATCH] Move dataset ID / connection discovery from
 '_implicit_environ' to 'client'.

---
 gcloud/datastore/client.py      | 66 +++++++++++++++++++++++++++++++--
 gcloud/datastore/test_client.py |  2 +-
 2 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/gcloud/datastore/client.py b/gcloud/datastore/client.py
index cb7c5a043ab5..ec267400ac7d 100644
--- a/gcloud/datastore/client.py
+++ b/gcloud/datastore/client.py
@@ -13,20 +13,80 @@
 # limitations under the License.
 """Convenience wrapper for invoking APIs/factories w/ a dataset ID."""
 
+import os
+
 from gcloud._helpers import _LocalStack
+from gcloud._helpers import _app_engine_id
+from gcloud._helpers import _compute_engine_id
 from gcloud.datastore import helpers
+from gcloud.datastore.connection import Connection
 from gcloud.datastore.batch import Batch
 from gcloud.datastore.entity import Entity
 from gcloud.datastore.key import Key
 from gcloud.datastore.query import Query
 from gcloud.datastore.transaction import Transaction
-from gcloud.datastore._implicit_environ import _determine_default_dataset_id
-from gcloud.datastore._implicit_environ import get_connection
 
 
 _MAX_LOOPS = 128
 """Maximum number of iterations to wait for deferred keys."""
 
+_DATASET_ENV_VAR_NAME = 'GCLOUD_DATASET_ID'
+"""Environment variable defining default dataset ID."""
+
+_GCD_DATASET_ENV_VAR_NAME = 'DATASTORE_DATASET'
+"""Environment variable defining default dataset ID under GCD."""
+
+
+def _get_production_dataset_id():
+    """Gets the production application ID if it can be inferred."""
+    return os.getenv(_DATASET_ENV_VAR_NAME)
+
+
+def _get_gcd_dataset_id():
+    """Gets the GCD application ID if it can be inferred."""
+    return os.getenv(_GCD_DATASET_ENV_VAR_NAME)
+
+
+def _determine_default_dataset_id(dataset_id=None):
+    """Determine default dataset ID explicitly or implicitly as fall-back.
+
+    In implicit case, supports four environments. In order of precedence, the
+    implicit environments are:
+
+    * GCLOUD_DATASET_ID environment variable
+    * DATASTORE_DATASET environment variable (for ``gcd`` testing)
+    * Google App Engine application ID
+    * Google Compute Engine project ID (from metadata server)
+
+    :type dataset_id: string
+    :param dataset_id: Optional. The dataset ID to use as default.
+
+    :rtype: string or ``NoneType``
+    :returns: Default dataset ID if it can be determined.
+    """
+    if dataset_id is None:
+        dataset_id = _get_production_dataset_id()
+
+    if dataset_id is None:
+        dataset_id = _get_gcd_dataset_id()
+
+    if dataset_id is None:
+        dataset_id = _app_engine_id()
+
+    if dataset_id is None:
+        dataset_id = _compute_engine_id()
+
+    return dataset_id
+
+
+def _get_connection():
+    """Shortcut method to establish a connection to the Cloud Datastore.
+
+    :rtype: :class:`gcloud.datastore.connection.Connection`
+    :returns: A connection defined with the proper credentials.
+    """
+    return Connection.from_environment()
+
 
 def _extended_lookup(connection, dataset_id, key_pbs,
                      missing=None, deferred=None,
@@ -126,7 +186,7 @@ def __init__(self, dataset_id=None, namespace=None, connection=None):
             raise EnvironmentError('Dataset ID could not be inferred.')
         self.dataset_id = dataset_id
         if connection is None:
-            connection = get_connection()
+            connection = _get_connection()
         self.connection = connection
         self._batch_stack = _LocalStack()
         self.namespace = namespace
diff --git a/gcloud/datastore/test_client.py b/gcloud/datastore/test_client.py
index 611dad5f8845..aeee85f7968b 100644
--- a/gcloud/datastore/test_client.py
+++ b/gcloud/datastore/test_client.py
@@ -58,7 +58,7 @@ def test_ctor_w_implicit_inputs(self):
         klass = self._getTargetClass()
         with _Monkey(_MUT,
                      _determine_default_dataset_id=lambda x: x or OTHER,
-                     get_connection=lambda: conn):
+                     _get_connection=lambda: conn):
             client = klass()
         self.assertEqual(client.dataset_id, OTHER)
         self.assertEqual(client.namespace, None)