Repository: mercadona/rele
Pull request #143 (2019-12-21 20:06:32)

Configure timeout from .publish()
### :tophat: What?

Configure timeout from .publish()

### :thinking: Why?

More granular timeout configuration

### :link: Related issue

Fixes #137 



Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index a99b266..1c24dc9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,10 @@
 Changelog
 =========
 
+`0.9.1` (2020-1-2)
+* Ack messages when data not json serializable (#141)
+* Use ThreadScheduler instead of ThreadPoolExecutor (#145)
+
 `0.9.0` (2019-12-20)
 * Flask support via middleware (#127)
 * Add message attributes to metrics log (#128)
diff --git a/docs/api/settings.rst b/docs/api/settings.rst
index 31c8d75..73b8f64 100644
--- a/docs/api/settings.rst
+++ b/docs/api/settings.rst
@@ -123,6 +123,8 @@ Ack deadline for all subscribers in seconds.
     passes, the message is no longer considered outstanding, and Cloud Pub/Sub will attempt
     to redeliver the message.*
 
+.. _settings_publisher_timeout:
+
 ``PUBLISHER_TIMEOUT``
 ---------------------
 
@@ -136,7 +138,7 @@ Timeout that the publishing result will wait on the future to publish successful
 <https://googleapis.dev/python/pubsub/1.1.0/publisher/api/futures.html?highlight=result#google.cloud.pubsub_v1.publisher.futures.Future.result>`_
 
 ``THREADS_PER_SUBSCRIPTION``
-------------------
+----------------------------
 
 **Optional**
 
diff --git a/rele/__init__.py b/rele/__init__.py
index 9f2820e..71852af 100644
--- a/rele/__init__.py
+++ b/rele/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.9.0"
+__version__ = "0.9.1"
 default_app_config = "rele.apps.ReleConfig"
 
 from .client import Publisher, Subscriber  # noqa
diff --git a/rele/client.py b/rele/client.py
index 53a02d1..6d2725c 100644
--- a/rele/client.py
+++ b/rele/client.py
@@ -104,7 +104,7 @@ class Publisher:
     :param gc_project_id: string Google Cloud Project ID.
     :param credentials: string Google Cloud Credentials.
     :param encoder: A valid `json.encoder.JSONEncoder subclass <https://docs.python.org/3/library/json.html#json.JSONEncoder>`_  # noqa
-    :param timeout: float
+    :param timeout: float, default :ref:`settings_publisher_timeout`
     """
 
     def __init__(self, gc_project_id, credentials, encoder, timeout):
@@ -116,7 +116,7 @@ class Publisher:
         else:
             self._client = pubsub_v1.PublisherClient(credentials=credentials)
 
-    def publish(self, topic, data, blocking=False, **attrs):
+    def publish(self, topic, data, blocking=False, timeout=None, **attrs):
         """Publishes message to Google PubSub topic.
 
         Usage::
@@ -145,6 +145,7 @@ class Publisher:
         :param topic: string topic to publish the data.
         :param data: dict with the content of the message.
         :param blocking: boolean
+        :param timeout: float, default None fallsback to :ref:`settings_publisher_timeout`
         :param attrs: Extra parameters to be published.
         :return: `Future <https://googleapis.github.io/google-cloud-python/latest/pubsub/subscriber/api/futures.html>`_  # noqa
         """
@@ -157,6 +158,6 @@ class Publisher:
         if not blocking:
             return future
 
-        future.result(timeout=self._timeout)
+        future.result(timeout=timeout or self._timeout)
         run_middleware_hook("post_publish", topic)
         return future
diff --git a/rele/subscription.py b/rele/subscription.py
index 846618d..73f6b03 100644
--- a/rele/subscription.py
+++ b/rele/subscription.py
@@ -1,7 +1,6 @@
 import json
 import logging
 import time
-from concurrent import futures
 
 from .middleware import run_middleware_hook
 
@@ -64,7 +63,19 @@ class Callback:
         run_middleware_hook("pre_process_message", self._subscription, message)
         start_time = time.time()
 
-        data = json.loads(message.data.decode("utf-8"))
+        try:
+            data = json.loads(message.data.decode("utf-8"))
+        except json.JSONDecodeError as e:
+            message.ack()
+            run_middleware_hook(
+                "post_process_message_failure",
+                self._subscription,
+                e,
+                start_time,
+                message,
+            )
+            run_middleware_hook("post_process_message")
+            return
 
         try:
             res = self._subscription(data, **dict(message.attributes))
diff --git a/rele/worker.py b/rele/worker.py
index 08ea6fa..6b1e901 100644
--- a/rele/worker.py
+++ b/rele/worker.py
@@ -3,6 +3,8 @@ import sys
 import time
 from concurrent import futures
 
+from google.cloud.pubsub_v1.subscriber.scheduler import ThreadScheduler
+
 from .client import Subscriber
 from .middleware import run_middleware_hook
 from .subscription import Callback
@@ -56,9 +58,10 @@ class Worker:
             executor_kwargs = {
                 "thread_name_prefix": "ThreadPoolExecutor-ThreadScheduler"
             }
-            scheduler = futures.ThreadPoolExecutor(
+            executor = futures.ThreadPoolExecutor(
                 max_workers=self.threads_per_subscription, **executor_kwargs
             )
+            scheduler = ThreadScheduler(executor=executor)
             self._futures.append(
                 self._subscriber.consume(
                     subscription_name=subscription.name,


Test patch:
diff --git a/tests/test_client.py b/tests/test_client.py
index 954ff06..069f57a 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -63,6 +63,21 @@ class TestPublisher:
         )
         mock_future.assert_called_once_with(timeout=100)
 
+    @patch.object(concurrent.futures.Future, "result")
+    def test_publishes_data_with_client_timeout_when_blocking_and_timeout_specified(
+        self, mock_future, publisher
+    ):
+        publisher._timeout = 100.0
+        publisher.publish(
+            topic="order-cancelled", data={"foo": "bar"}, blocking=True, timeout=50
+        )
+
+        publisher._client.publish.return_value = mock_future
+        publisher._client.publish.assert_called_with(
+            ANY, b'{"foo": "bar"}', published_at=ANY
+        )
+        mock_future.assert_called_once_with(timeout=50)
+
 
 class TestSubscriber:
     @patch.object(SubscriberClient, "create_subscription")
diff --git a/tests/test_subscription.py b/tests/test_subscription.py
index ec33e67..d3e07d5 100644
--- a/tests/test_subscription.py
+++ b/tests/test_subscription.py
@@ -91,14 +91,33 @@ class TestCallback:
             message_id="1",
         )
 
-        return pubsub_v1.subscriber.message.Message(rele_message, "ack-id", MagicMock())
+        message = pubsub_v1.subscriber.message.Message(
+            rele_message, "ack-id", MagicMock()
+        )
+        message.ack = MagicMock(autospec=True)
+        return message
 
     @pytest.fixture
     def message_wrapper_empty(self):
         rele_message = pubsub_v1.types.PubsubMessage(
             data=b"", attributes={"lang": "es"}, message_id="1"
         )
-        return pubsub_v1.subscriber.message.Message(rele_message, "ack-id", MagicMock())
+        message = pubsub_v1.subscriber.message.Message(
+            rele_message, "ack-id", MagicMock()
+        )
+        message.ack = MagicMock(autospec=True)
+        return message
+
+    @pytest.fixture
+    def message_wrapper_invalid_json(self):
+        rele_message = pubsub_v1.types.PubsubMessage(
+            data=b"foobar", attributes={}, message_id="1"
+        )
+        message = pubsub_v1.subscriber.message.Message(
+            rele_message, "ack-id", MagicMock()
+        )
+        message.ack = MagicMock(autospec=True)
+        return message
 
     def test_log_start_processing_when_callback_called(
         self, caplog, message_wrapper, published_at
@@ -129,6 +148,7 @@ class TestCallback:
             res = callback(message_wrapper)
 
         assert res == 123
+        message_wrapper.ack.assert_called_once()
         assert len(caplog.records) == 3
         message_wrapper_log = caplog.records[1]
         assert message_wrapper_log.message == (
@@ -168,6 +188,7 @@ class TestCallback:
         res = callback(message_wrapper)
 
         assert res is None
+        message_wrapper.ack.assert_not_called()
         failed_log = caplog.records[-1]
         assert failed_log.message == (
             "Exception raised while processing "
@@ -187,6 +208,33 @@ class TestCallback:
         }
         assert failed_log.subscription_message == message_wrapper
 
+    def test_log_acks_called_message_when_not_json_serializable(
+        self, caplog, message_wrapper_invalid_json, published_at
+    ):
+        callback = Callback(sub_stub)
+        res = callback(message_wrapper_invalid_json)
+
+        assert res is None
+        message_wrapper_invalid_json.ack.assert_called_once()
+        failed_log = caplog.records[-1]
+        assert failed_log.message == (
+            "Exception raised while processing "
+            "message for rele-some-cool-topic - "
+            "sub_stub: JSONDecodeError"
+        )
+        assert failed_log.metrics == {
+            "name": "subscriptions",
+            "data": {
+                "agent": "rele",
+                "topic": "some-cool-topic",
+                "status": "failed",
+                "subscription": "rele-some-cool-topic",
+                "duration_seconds": pytest.approx(0.5, abs=0.5),
+                "attributes": {},
+            },
+        }
+        assert failed_log.subscription_message == message_wrapper_invalid_json
+
     def test_published_time_as_message_attribute(self, message_wrapper, caplog):
         callback = Callback(sub_published_time_type)
         callback(message_wrapper)
diff --git a/tests/test_worker.py b/tests/test_worker.py
index f8fb506..9d3d41b 100644
--- a/tests/test_worker.py
+++ b/tests/test_worker.py
@@ -3,6 +3,7 @@ from unittest.mock import ANY, patch
 import os
 
 import pytest
+from google.cloud.pubsub_v1.subscriber.scheduler import ThreadScheduler
 
 from rele import Subscriber, Worker, sub
 from rele.middleware import register_middleware
@@ -47,7 +48,8 @@ class TestWorker:
             subscription_name="rele-some-cool-topic", callback=ANY, scheduler=ANY
         )
         scheduler = mock_consume.call_args_list[0][1]["scheduler"]
-        assert isinstance(scheduler, futures.ThreadPoolExecutor)
+        assert isinstance(scheduler, ThreadScheduler)
+        assert isinstance(scheduler._executor, futures.ThreadPoolExecutor)
 
     def test_setup_creates_subscription_when_topic_given(
         self, mock_create_subscription, worker
@@ -71,7 +73,8 @@ class TestWorker:
             subscription_name="rele-some-cool-topic", callback=ANY, scheduler=ANY
         )
         scheduler = mock_consume.call_args_list[0][1]["scheduler"]
-        assert isinstance(scheduler, futures.ThreadPoolExecutor)
+        assert isinstance(scheduler, ThreadScheduler)
+        assert isinstance(scheduler._executor, futures.ThreadPoolExecutor)
         mock_wait_forever.assert_called_once()
 
     @patch.object(Worker, "_wait_forever")