Ensures StackDriver worker thread is alive when send logs

temlogger/helpers.py

@@ -1,8 +1,6 @@
 import base64
 import tempfile
 
-
-from google.cloud.logging import Client as LoggingClient
 from importlib import import_module
 
 
@@ -43,6 +41,8 @@ def import_string_list(dotted_path_list=[]):
 
 
 def load_google_client(base64_data, scopes=[]):
+    from google.cloud.logging import Client as LoggingClient
+
     if not base64_data:
         return ''
 

temlogger/providers/stackdriver.py

@@ -1,3 +1,4 @@
+from google.cloud.logging.handlers import CloudLoggingHandler
 from .base import FormatterBase
 
 
@@ -6,3 +7,17 @@ class StackDriverFormatter(FormatterBase):
     def format(self, record):
         message = super().format(record)
         return message
+
+
+class StackDriverLoggingHandler(CloudLoggingHandler):
+
+    def ensure_transport_worker_is_running(self):
+        """
+        Ensures that transport worker is running when any log is done
+        """
+        if not self.transport.worker.is_alive:
+            self.transport.worker.start()
+
+    def emit(self, record):
+        self.ensure_transport_worker_is_running()
+        super().emit(record)

temlogger/temlogger.py

@@ -181,6 +181,7 @@ def get_logger_stackdriver(self, name, event_handlers=[]):
         """
         import google.cloud.logging
         from .providers.stackdriver import StackDriverFormatter
+        from .providers.stackdriver import StackDriverLoggingHandler
 
         app_name = config.get_app_name()
         logging_environment = config.get_environment()
@@ -197,7 +198,7 @@ def get_logger_stackdriver(self, name, event_handlers=[]):
         logger.setLevel(config.get_log_level_parsed())
         logger.logging_provider = LoggingProvider.STACK_DRIVER
 
-        handler = client.get_default_handler()
+        handler = StackDriverLoggingHandler(client)
 
         # Setup logger explicitly with this handler
         handler.setFormatter(StackDriverFormatter(

temlogger/tests/test_get_loggers.py

@@ -162,6 +162,33 @@ def test_stackdriver_with_credentials_base64_as_environment(self):
         temlogger.config.set_provider('stackdriver')
         temlogger.getLogger('stackdriver-base64')
 
+    @mock.patch("google.cloud.logging.handlers.CloudLoggingHandler.emit")
+    def test_stackdriver_worker_thread_is_alive_when_send_logs(self, mock_emit):
+        """
+        This problem happen when temlogger is used with celery.
+        The fork process made by celery, don't start the worker
+        """
+        valid_cred = encode_file_as_base64(VALID_GOOGLE_CREDENTIALS)
+
+        os.environ['TEMLOGGER_GOOGLE_CREDENTIALS_BASE64'] = valid_cred
+        temlogger.config.set_provider('stackdriver')
+        temlogger.config.set_log_level('DEBUG')
+
+        logger = temlogger.getLogger('stackdriver-worker-stopped-log')
+        handler = logger.handlers[0]
+
+        handler.transport.worker.stop()
+        self.assertFalse(handler.transport.worker.is_alive)
+
+        logger.info('Activate thread on log info')
+        self.assertTrue(handler.transport.worker.is_alive)
+
+        handler.transport.worker.stop()
+        self.assertFalse(handler.transport.worker.is_alive)
+
+        logger.debug('Activate thread on log debug')
+        self.assertTrue(handler.transport.worker.is_alive)
+
 
 class TestConsoleLogger(unittest.TestCase):