From 0c5f447950a7338f4374d789f567209841bd9d8c Mon Sep 17 00:00:00 2001
From: Adrien Perrin <adrien.perrin@nersc.no>
Date: Wed, 7 Oct 2020 13:37:52 +0000
Subject: [PATCH] add GELF handler for Celery logging

---
 Dockerfile                   |  1 +
 geospaas_processing/tasks.py | 13 +++++++++++++
 tests/test_tasks.py          | 19 +++++++++++++++++++
 3 files changed, 33 insertions(+)

diff --git a/Dockerfile b/Dockerfile
index 9bb95693..c61d64ff 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -4,6 +4,7 @@ FROM ${BASE_IMAGE} as base
 RUN pip install --no-cache-dir \
     celery==4.4 \
     django-celery-results==1.2 \
+    graypy==2.1.0 \
     paramiko==2.7.1 \
     redis==3.5 \
     scp==0.13.2
diff --git a/geospaas_processing/tasks.py b/geospaas_processing/tasks.py
index 2ae1f165..af3d8a3b 100644
--- a/geospaas_processing/tasks.py
+++ b/geospaas_processing/tasks.py
@@ -14,6 +14,9 @@
 import shutil
 
 import celery
+import celery.utils.log
+import celery.signals
+import graypy.handler
 import scp
 from django.db import connection
 
@@ -30,6 +33,16 @@
 app.config_from_object('django.conf:settings', namespace='CELERY')
 
 
+@celery.signals.after_setup_logger.connect
+def setup_logger(logger, *args, **kwargs):  # pylint: disable=unused-argument
+    """Set up a GELF handler for Celery tasks if the necessary environment variables are set"""
+    logging_host = os.getenv('GEOSPAAS_PROCESSING_LOGGING_HOST')
+    logging_port = os.getenv('GEOSPAAS_PROCESSING_LOGGING_PORT')
+    if logging_host and logging_port:
+        gelf_handler = graypy.handler.GELFTCPHandler(logging_host, logging_port, facility=__name__)
+        logger.addHandler(gelf_handler)
+
+
 class FaultTolerantTask(celery.Task):  # pylint: disable=abstract-method
     """
     Workaround for https://github.com/celery/django-celery/issues/121.
diff --git a/tests/test_tasks.py b/tests/test_tasks.py
index 2fed68f0..a95523e3 100644
--- a/tests/test_tasks.py
+++ b/tests/test_tasks.py
@@ -7,6 +7,7 @@
 import unittest.mock as mock
 
 import celery
+import graypy.handler
 import scp
 
 import geospaas_processing.downloaders as downloaders
@@ -54,6 +55,24 @@ def test_retry_if_locked(self):
         mock_task.retry.assert_called()
 
 
+class SignalsTestCase(unittest.TestCase):
+    """Unit tests for Celery signal functions"""
+
+    def test_setup_logger(self):
+        """
+        The setup_logger() functions must add a GELF handler if the
+        right environment variables are defined
+        """
+        with mock.patch('os.getenv', return_value='test'):
+            logger = logging.Logger('test_logger')
+            tasks.setup_logger(logger)
+            handler = logger.handlers[0]
+        self.assertIsInstance(handler, graypy.handler.GELFTCPHandler)
+        self.assertEqual(handler.host, 'test')
+        self.assertEqual(handler.port, 'test')
+        self.assertEqual(handler.facility, tasks.__name__)
+
+
 class DownloadTestCase(unittest.TestCase):
     """Tests for the download() task"""