Repository: googleapis/google-cloud-python

Issue:
Pub/Sub should send streaming pull requests over unary RPCs



Discussion:

Resolving patch:
diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/helper_threads.py b/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/helper_threads.py
--- a/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/helper_threads.py
+++ b/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/helper_threads.py
@@ -111,6 +111,6 @@ def __call__(self):
             try:
                 self._callback(items)
             except Exception as exc:
-                _LOGGER.error('%s: %s', exc.__class__.__name__, exc)
+                _LOGGER.exception('Error in queue callback worker: %s', exc)
 
         _LOGGER.debug('Exiting the QueueCallbackWorker.')
diff --git a/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py b/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py
--- a/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py
+++ b/pubsub/google/cloud/pubsub_v1/subscriber/_protocol/streaming_pull_manager.py
@@ -14,13 +14,15 @@
 
 from __future__ import division
 
+import collections
 import functools
 import logging
 import threading
 
-from google.api_core import exceptions
 import grpc
+import six
 
+from google.api_core import exceptions
 from google.cloud.pubsub_v1 import types
 from google.cloud.pubsub_v1.subscriber._protocol import bidi
 from google.cloud.pubsub_v1.subscriber._protocol import dispatcher
@@ -84,6 +86,10 @@ class StreamingPullManager(object):
             scheduler will be used.
     """
 
+    _UNARY_REQUESTS = True
+    """If set to True, this class will make requests over a separate unary
+    RPC instead of over the streaming RPC."""
+
     def __init__(self, client, subscription, flow_control=types.FlowControl(),
                  scheduler=None):
         self._client = client
@@ -220,9 +226,41 @@ def maybe_resume_consumer(self):
         else:
             _LOGGER.debug('Did not resume, current load is %s', self.load)
 
+    def _send_unary_request(self, request):
+        """Send a request using a separate unary request instead of over the
+        stream.
+
+        Args:
+            request (types.StreamingPullRequest): The stream request to be
+                mapped into unary requests.
+        """
+        if request.ack_ids:
+            self._client.acknowledge(
+                subscription=self._subscription,
+                ack_ids=list(request.ack_ids))
+
+        if request.modify_deadline_ack_ids:
+            # Send ack_ids with the same deadline seconds together.
+            deadline_to_ack_ids = collections.defaultdict(list)
+
+            for n, ack_id in enumerate(request.modify_deadline_ack_ids):
+                deadline = request.modify_deadline_seconds[n]
+                deadline_to_ack_ids[deadline].append(ack_id)
+
+            for deadline, ack_ids in six.iteritems(deadline_to_ack_ids):
+                self._client.modify_ack_deadline(
+                    subscription=self._subscription,
+                    ack_ids=ack_ids,
+                    ack_deadline_seconds=deadline)
+
+        _LOGGER.debug('Sent request(s) over unary RPC.')
+
     def send(self, request):
         """Queue a request to be sent to the RPC."""
-        self._rpc.send(request)
+        if self._UNARY_REQUESTS:
+            self._send_unary_request(request)
+        else:
+            self._rpc.send(request)
 
     def open(self, callback):
         """Begin consuming messages.


Test patch:
diff --git a/pubsub/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py b/pubsub/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py
--- a/pubsub/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py
+++ b/pubsub/tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py
@@ -174,8 +174,41 @@ def test_resume_not_paused():
     manager._consumer.resume.assert_not_called()
 
 
-def test_send():
+def test_send_unary():
     manager = make_manager()
+    manager._UNARY_REQUESTS = True
+
+    manager.send(types.StreamingPullRequest(
+        ack_ids=['ack_id1', 'ack_id2'],
+        modify_deadline_ack_ids=['ack_id3', 'ack_id4', 'ack_id5'],
+        modify_deadline_seconds=[10, 20, 20]))
+
+    manager._client.acknowledge.assert_called_once_with(
+        subscription=manager._subscription, ack_ids=['ack_id1', 'ack_id2'])
+
+    manager._client.modify_ack_deadline.assert_has_calls([
+        mock.call(
+            subscription=manager._subscription,
+            ack_ids=['ack_id3'], ack_deadline_seconds=10),
+        mock.call(
+            subscription=manager._subscription,
+            ack_ids=['ack_id4', 'ack_id5'], ack_deadline_seconds=20),
+        ], any_order=True)
+
+
+def test_send_unary_empty():
+    manager = make_manager()
+    manager._UNARY_REQUESTS = True
+
+    manager.send(types.StreamingPullRequest())
+
+    manager._client.acknowledge.assert_not_called()
+    manager._client.modify_ack_deadline.assert_not_called()
+
+
+def test_send_streaming():
+    manager = make_manager()
+    manager._UNARY_REQUESTS = False
     manager._rpc = mock.create_autospec(bidi.BidiRpc, instance=True)
 
     manager.send(mock.sentinel.request)