Merge pull request #1257 from ISISScientificComputing/1256_manage_db_connection_lifetime

Manage DB connection lifetime during Message processing & show FQDN in logs instead of hostname

Author: Dimitar Tasev

queue_processors/queue_processor/handle_message.py

@@ -14,6 +14,7 @@
 """
 import logging
 import socket
+from contextlib import contextmanager
 from typing import Optional
 from django.db import transaction
 from django.utils import timezone
@@ -42,9 +43,35 @@ def __init__(self, queue_listener):
 
         self._logger = logging.getLogger("handle_queue_message")
 
+        self.database = None
+        self.data_model = None
+
+    def connect(self):
+        """
+        Starts a connection to the database
+        """
         self.database = db_access.start_database()
         self.data_model = self.database.data_model
 
+    def disconnect(self):
+        """
+        Disconnects from the database
+        """
+        self.database.disconnect()
+        self.database = None
+        self.data_model = None
+
+    @contextmanager
+    def connected(self):
+        """
+        Context manager for the connection state to the DB
+        """
+        self.connect()
+        try:
+            yield
+        finally:
+            self.disconnect()
+
     def data_ready(self, message: Message):
         """
         Called when destination queue was data_ready.
@@ -250,7 +277,7 @@ def _common_reduction_run_update(reduction_run, status, message):
         reduction_run.finished = timezone.now()
         reduction_run.message = message.message
         reduction_run.reduction_log = message.reduction_log
-        message.admin_log += "Running on host: %s" % socket.gethostname()
+        message.admin_log += "Running on host: %s" % socket.getfqdn()
         reduction_run.admin_log = message.admin_log
 
     @staticmethod

queue_processors/queue_processor/queue_listener.py

@@ -75,14 +75,15 @@ def on_message(self, headers, message):
             # has to submit an acknowledgement for receiving the message
             # (otherwise I think that it is not removed from the queue but I am not sure about that)
             self.client.ack(headers["message-id"], headers["subscription"])
-            try:
-                if destination == '/queue/DataReady':
-                    self.message_handler.data_ready(message)
-                else:
-                    self.logger.error("Received a message on an unknown topic '%s'", destination)
-            except Exception as exp:  # pylint:disable=broad-except
-                self.logger.error("Unhandled exception encountered: %s %s\n\n%s",
-                                  type(exp).__name__, exp, traceback.format_exc())
+            with self.message_handler.connected():
+                try:
+                    if destination == '/queue/DataReady':
+                        self.message_handler.data_ready(message)
+                    else:
+                        self.logger.error("Received a message on an unknown topic '%s'", destination)
+                except Exception as exp:  # pylint:disable=broad-except
+                    self.logger.error("Unhandled exception encountered: %s %s\n\n%s",
+                                      type(exp).__name__, exp, traceback.format_exc())
 
 
 def setup_connection(consumer_name) -> Tuple[QueueClient, QueueListener]:

queue_processors/queue_processor/tests/test_handle_message.py

@@ -76,6 +76,7 @@ def setUp(self):
         })
         with patch("logging.getLogger") as patched_logger:
             self.handler = HandleMessage(self.mocked_client)
+            self.handler.connect()
             self.mocked_logger = patched_logger.return_value
 
         db_handle = model.database.access.start_database()
@@ -91,6 +92,8 @@ def setUp(self):
         self.reduction_run.save()
 
     def tearDown(self) -> None:
+        if self.handler.database is not None:
+            self.handler.disconnect()
         self.experiment.delete()
         self.instrument.delete()
         self.reduction_run.delete()
@@ -392,6 +395,21 @@ def test_create_run_records_valid_rb_number(self):
             reduction_run, _, _ = self.handler.create_run_records(self.msg)
             assert reduction_run.experiment.reference_number == self.msg.rb_number
 
+    def test_connected(self):
+        """
+        Test the connected context manager properly starts/clears the DB connection
+        """
+        # Disconnect first to remove state from TestCase.setUp
+        self.handler.disconnect()
+
+        with self.handler.connected():
+            assert self.handler.database is not None
+            assert self.handler.data_model is not None
+
+        # at the end of the context there should be a disconnect
+        assert self.handler.database is None
+        assert self.handler.data_model is None
+
 
 if __name__ == '__main__':
     main()

queue_processors/queue_processor/tests/test_queue_listener.py

@@ -55,7 +55,7 @@ class TestQueueListener(TestCase):
     """
     def setUp(self):
         self.mocked_client = mock.Mock(spec=QueueClient)
-        self.mocked_handler = mock.Mock(spec=HandleMessage)
+        self.mocked_handler = mock.MagicMock(spec=HandleMessage)
         self.headers = self._get_header()
 
         with patch("queue_processors.queue_processor.queue_listener"
@@ -96,6 +96,7 @@ def test_on_message_can_receive_a_prepopulated_message(self):
         self.mocked_logger.info.assert_called_once()
         self.mocked_client.ack.assert_called_once_with(self.headers["message-id"], self.headers["subscription"])
         self.mocked_handler.data_ready.assert_called_once()
+        self.mocked_handler.connected.assert_called_once()
         self.assertIsInstance(self.mocked_handler.data_ready.call_args[0][0], Message)
 
     def test_on_message_sends_acknowledgement(self):
@@ -106,6 +107,7 @@ def test_on_message_sends_acknowledgement(self):
         self.mocked_logger.info.assert_called_once()
         self.mocked_client.ack.assert_called_once_with(self.headers["message-id"], self.headers["subscription"])
         self.mocked_handler.data_ready.assert_called_once()
+        self.mocked_handler.connected.assert_called_once()
         self.assertIsInstance(self.mocked_handler.data_ready.call_args[0][0], Message)
 
     def test_on_message_handler_catches_exceptions(self):

utils/clients/queue_client.py

@@ -94,7 +94,7 @@ def subscribe_queues(self, queue_list, consumer_name, listener):
         for queue in queue_list:
             # prefetchSize limits the processing to 1 message at a time
             self._connection.subscribe(destination=queue,
-                                       id=socket.gethostname(),
+                                       id=socket.getfqdn(),
                                        ack="client-individual",
                                        header={'activemq.prefetchSize': '1'})
             self._logger.info("[%s] Subscribing to %s", consumer_name, queue)