Merge pull request #2553 from waprin/logging_gax_optional

Allows Explicitly Enabling/Disabling GAX for Logging/Pubsub

Author: waprin

logging/google/cloud/logging/client.py

@@ -71,11 +71,25 @@ class Client(JSONClient):
     :param http: An optional HTTP object to make requests. If not passed, an
                  ``http`` object is created that is bound to the
                  ``credentials`` for the current object.
+
+    :type use_gax: bool
+    :param use_gax: (Optional) Explicitly specifies whether
+                    to use the gRPC transport (via GAX) or HTTP. If unset,
+                    falls back to the ``GOOGLE_CLOUD_DISABLE_GRPC`` environment
+                    variable
     """
 
     _connection_class = Connection
     _logging_api = _sinks_api = _metrics_api = None
 
+    def __init__(self, project=None, credentials=None,
+                 http=None, use_gax=None):
+        super(Client, self).__init__(project, credentials, http)
+        if use_gax is None:
+            self._use_gax = _USE_GAX
+        else:
+            self._use_gax = use_gax
+
     @property
     def logging_api(self):
         """Helper for logging-related API calls.
@@ -85,7 +99,7 @@ def logging_api(self):
         https://cloud.google.com/logging/docs/api/ref_v2beta1/rest/v2beta1/projects.logs
         """
         if self._logging_api is None:
-            if _USE_GAX:
+            if self._use_gax:
                 generated = GeneratedLoggingAPI()
                 self._logging_api = GAXLoggingAPI(generated)
             else:

logging/unit_tests/test_client.py

@@ -39,14 +39,14 @@ def test_ctor(self):
         self.assertEqual(client.project, self.PROJECT)
 
     def test_logging_api_wo_gax(self):
-        from google.cloud.logging.connection import _LoggingAPI
-        from google.cloud.logging import client as MUT
         from google.cloud._testing import _Monkey
-        client = self._makeOne(self.PROJECT, credentials=_Credentials())
-        conn = client.connection = object()
+        from google.cloud.logging import client as MUT
+        from google.cloud.logging.connection import _LoggingAPI
 
         with _Monkey(MUT, _USE_GAX=False):
-            api = client.logging_api
+            client = self._makeOne(self.PROJECT, credentials=_Credentials())
+        conn = client.connection = object()
+        api = client.logging_api
 
         self.assertIsInstance(api, _LoggingAPI)
         self.assertIs(api._connection, conn)
@@ -85,6 +85,19 @@ def __init__(self, _wrapped):
         again = client.logging_api
         self.assertIs(again, api)
 
+    def test_no_gax_ctor(self):
+        from google.cloud._testing import _Monkey
+        from google.cloud.logging import client as MUT
+        from google.cloud.logging.connection import _LoggingAPI
+
+        creds = _Credentials()
+        with _Monkey(MUT, _USE_GAX=True):
+            client = self._makeOne(project=self.PROJECT, credentials=creds,
+                                   use_gax=False)
+
+        api = client.logging_api
+        self.assertIsInstance(api, _LoggingAPI)
+
     def test_sinks_api_wo_gax(self):
         from google.cloud.logging.connection import _SinksAPI
         from google.cloud.logging import client as MUT

pubsub/google/cloud/pubsub/client.py

@@ -63,7 +63,20 @@ class Client(JSONClient):
     :param http: An optional HTTP object to make requests. If not passed, an
                  ``http`` object is created that is bound to the
                  ``credentials`` for the current object.
+
+    :type use_gax: bool
+    :param use_gax: (Optional) Explicitly specifies whether
+                    to use the gRPC transport (via GAX) or HTTP. If unset,
+                    falls back to the ``GOOGLE_CLOUD_DISABLE_GRPC`` environment
+                    variable
     """
+    def __init__(self, project=None, credentials=None,
+                 http=None, use_gax=None):
+        super(Client, self).__init__(project, credentials, http)
+        if use_gax is None:
+            self._use_gax = _USE_GAX
+        else:
+            self._use_gax = use_gax
 
     _connection_class = Connection
     _publisher_api = _subscriber_api = _iam_policy_api = None
@@ -72,7 +85,7 @@ class Client(JSONClient):
     def publisher_api(self):
         """Helper for publisher-related API calls."""
         if self._publisher_api is None:
-            if _USE_GAX:
+            if self._use_gax:
                 generated = make_gax_publisher_api(self.connection)
                 self._publisher_api = GAXPublisherAPI(generated)
             else:
@@ -83,7 +96,7 @@ def publisher_api(self):
     def subscriber_api(self):
         """Helper for subscriber-related API calls."""
         if self._subscriber_api is None:
-            if _USE_GAX:
+            if self._use_gax:
                 generated = make_gax_subscriber_api(self.connection)
                 self._subscriber_api = GAXSubscriberAPI(generated)
             else:

pubsub/unit_tests/test_client.py

@@ -34,18 +34,33 @@ def test_publisher_api_wo_gax(self):
         from google.cloud.pubsub import client as MUT
         from google.cloud._testing import _Monkey
         creds = _Credentials()
-        client = self._makeOne(project=self.PROJECT, credentials=creds)
-        conn = client.connection = object()
 
         with _Monkey(MUT, _USE_GAX=False):
-            api = client.publisher_api
+            client = self._makeOne(project=self.PROJECT, credentials=creds)
+
+        conn = client.connection = object()
+        api = client.publisher_api
 
         self.assertIsInstance(api, _PublisherAPI)
         self.assertIs(api._connection, conn)
         # API instance is cached
         again = client.publisher_api
         self.assertIs(again, api)
 
+    def test_no_gax_ctor(self):
+        from google.cloud._testing import _Monkey
+        from google.cloud.pubsub.connection import _PublisherAPI
+        from google.cloud.pubsub import client as MUT
+
+        creds = _Credentials()
+        with _Monkey(MUT, _USE_GAX=True):
+            client = self._makeOne(project=self.PROJECT, credentials=creds,
+                                   use_gax=False)
+
+        self.assertFalse(client._use_gax)
+        api = client.publisher_api
+        self.assertIsInstance(api, _PublisherAPI)
+
     def test_publisher_api_w_gax(self):
         from google.cloud.pubsub import client as MUT
         from google.cloud._testing import _Monkey
@@ -84,11 +99,12 @@ def test_subscriber_api_wo_gax(self):
         from google.cloud.pubsub import client as MUT
         from google.cloud._testing import _Monkey
         creds = _Credentials()
-        client = self._makeOne(project=self.PROJECT, credentials=creds)
-        conn = client.connection = object()
 
         with _Monkey(MUT, _USE_GAX=False):
-            api = client.subscriber_api
+            client = self._makeOne(project=self.PROJECT, credentials=creds)
+
+        conn = client.connection = object()
+        api = client.subscriber_api
 
         self.assertIsInstance(api, _SubscriberAPI)
         self.assertIs(api._connection, conn)