Datastore: Add client_options to constructors for manual clients. (#9055)

* Add `client_options` to datastore

Author: emar-kar

packages/google-cloud-datastore/google/cloud/datastore/client.py

@@ -15,6 +15,7 @@
 
 import os
 
+import google.api_core.client_options
 from google.cloud._helpers import _LocalStack
 from google.cloud._helpers import _determine_default_project as _base_default_project
 from google.cloud.client import ClientWithProject
@@ -201,6 +202,11 @@ class Client(ClientWithProject):
                         you only need to set this if you're developing your
                         own library or partner tool.
 
+    :type client_options: :class:`~google.api_core.client_options.ClientOptions`
+                          or :class:`dict`
+    :param client_options: (Optional) Client options used to set user options on the
+                           client. API Endpoint should be set through client_options.
+
     :type _http: :class:`~requests.Session`
     :param _http: (Optional) HTTP object to make requests. Can be any object
                   that defines ``request()`` with the same interface as
@@ -228,6 +234,7 @@ def __init__(
         namespace=None,
         credentials=None,
         client_info=_CLIENT_INFO,
+        client_options=None,
         _http=None,
         _use_grpc=None,
     ):
@@ -236,6 +243,7 @@ def __init__(
         )
         self.namespace = namespace
         self._client_info = client_info
+        self._client_options = client_options
         self._batch_stack = _LocalStack()
         self._datastore_api_internal = None
         if _use_grpc is None:
@@ -246,7 +254,15 @@ def __init__(
             host = os.environ[GCD_HOST]
             self._base_url = "http://" + host
         except KeyError:
-            self._base_url = _DATASTORE_BASE_URL
+            api_endpoint = _DATASTORE_BASE_URL
+            if client_options:
+                if type(client_options) == dict:
+                    client_options = google.api_core.client_options.from_dict(
+                        client_options
+                    )
+                if client_options.api_endpoint:
+                    api_endpoint = client_options.api_endpoint
+            self._base_url = api_endpoint
 
     @staticmethod
     def _determine_default(project):

packages/google-cloud-datastore/setup.py

@@ -30,7 +30,7 @@
 release_status = "Development Status :: 5 - Production/Stable"
 dependencies = [
     "google-api-core[grpc] >= 1.14.0, < 2.0.0dev",
-    "google-cloud-core >= 1.0.0, < 2.0dev",
+    "google-cloud-core >= 1.0.3, < 2.0dev",
 ]
 extras = {}
 

packages/google-cloud-datastore/tests/unit/test_client.py

@@ -128,6 +128,7 @@ def _make_one(
         namespace=None,
         credentials=None,
         client_info=None,
+        client_options=None,
         _http=None,
         _use_grpc=None,
     ):
@@ -136,6 +137,7 @@ def _make_one(
             namespace=namespace,
             credentials=credentials,
             client_info=client_info,
+            client_options=client_options,
             _http=_http,
             _use_grpc=_use_grpc,
         )
@@ -172,6 +174,7 @@ def test_constructor_w_implicit_inputs(self):
         self.assertIs(client._credentials, creds)
         self.assertIs(client._client_info, _CLIENT_INFO)
         self.assertIsNone(client._http_internal)
+        self.assertIsNone(client._client_options)
         self.assertEqual(client.base_url, _DATASTORE_BASE_URL)
 
         self.assertIsNone(client.current_batch)
@@ -181,18 +184,20 @@ def test_constructor_w_implicit_inputs(self):
         _determine_default_project.assert_called_once_with(None)
 
     def test_constructor_w_explicit_inputs(self):
-        from google.cloud.datastore.client import _DATASTORE_BASE_URL
+        from google.api_core.client_options import ClientOptions
 
         other = "other"
         namespace = "namespace"
         creds = _make_credentials()
         client_info = mock.Mock()
+        client_options = ClientOptions("endpoint")
         http = object()
         client = self._make_one(
             project=other,
             namespace=namespace,
             credentials=creds,
             client_info=client_info,
+            client_options=client_options,
             _http=http,
         )
         self.assertEqual(client.project, other)
@@ -201,8 +206,8 @@ def test_constructor_w_explicit_inputs(self):
         self.assertIs(client._client_info, client_info)
         self.assertIs(client._http_internal, http)
         self.assertIsNone(client.current_batch)
+        self.assertIs(client._base_url, "endpoint")
         self.assertEqual(list(client._batch_stack), [])
-        self.assertEqual(client.base_url, _DATASTORE_BASE_URL)
 
     def test_constructor_use_grpc_default(self):
         import google.cloud.datastore.client as MUT
@@ -243,12 +248,39 @@ def test_constructor_gcd_host(self):
             self.assertEqual(client.base_url, "http://" + host)
 
     def test_base_url_property(self):
+        from google.cloud.datastore.client import _DATASTORE_BASE_URL
+        from google.api_core.client_options import ClientOptions
+
         alternate_url = "https://alias.example.com/"
         project = "PROJECT"
         creds = _make_credentials()
         http = object()
+        client_options = ClientOptions()
 
-        client = self._make_one(project=project, credentials=creds, _http=http)
+        client = self._make_one(
+            project=project,
+            credentials=creds,
+            _http=http,
+            client_options=client_options,
+        )
+        self.assertEqual(client.base_url, _DATASTORE_BASE_URL)
+        client.base_url = alternate_url
+        self.assertEqual(client.base_url, alternate_url)
+
+    def test_base_url_property_w_client_options(self):
+        alternate_url = "https://alias.example.com/"
+        project = "PROJECT"
+        creds = _make_credentials()
+        http = object()
+        client_options = {"api_endpoint": "endpoint"}
+
+        client = self._make_one(
+            project=project,
+            credentials=creds,
+            _http=http,
+            client_options=client_options,
+        )
+        self.assertEqual(client.base_url, "endpoint")
         client.base_url = alternate_url
         self.assertEqual(client.base_url, alternate_url)