def flush(self):
"""Transmits the buffered logs to Cloud Pub/Sub."""
self.acquire()
try:
while self.buffer:
body = {
'messages': [{
'data':
compat_urlsafe_b64encode(self.format(r))
} for r in self.buffer[:MAX_BATCH_SIZE]]
}
self._publish_body(self._client, body, self._project_id,
self._topic)
self.buffer = self.buffer[MAX_BATCH_SIZE:]
except RecoverableError:
# Cloud Pub/Sub API didn't receive the logs, most
# likely because of intermittent errors. This handler
# ignores this case and keep the logs in its buffer in
# a hope of future success.
pass
finally:
# Cut off the buffer at the _buf_hard_limit.
# The default value of -1 means unlimited buffer.
if self._buf_hard_limit != -1:
self.buffer = self.buffer[:self._buf_hard_limit]
self.release()
def setUp(self):
self.mocked_client = mock.MagicMock()
self.topic = 'projects/test-project/topics/test-topic'
self.projects = self.mocked_client.projects.return_value
self.topics = self.projects.topics.return_value
self.topics_publish = self.topics.publish.return_value
self.log_msg = 'Test message'
self.expected_payload = compat_urlsafe_b64encode(self.log_msg)
self.expected_body = {'messages': [{'data': self.expected_payload}]}
self.r = logging.LogRecord('test', logging.INFO, None, 0, self.log_msg,
[], None)
def setUp(self):
self.mocked_client = mock.MagicMock()
self.topic = 'projects/test-project/topics/test-topic'
self.publish_body = mock.MagicMock()
self.handler = pubsub_logging.PubsubHandler(
topic=self.topic,
client=self.mocked_client,
retry=self.RETRY,
capacity=self.BATCH_NUM,
publish_body=self.publish_body)
self.log_msg = 'Test message'
self.expected_payload = compat_urlsafe_b64encode(self.log_msg)
self.expected_body = {'messages': [{'data': self.expected_payload}]}
self.r = logging.LogRecord('test', logging.INFO, None, 0, self.log_msg,
[], None)
def test_formated_message(self):
"""Tests if a formatter assigned to the handler is used."""
self.log_msg = 'Test message'
self.expected_payload = compat_urlsafe_b64encode(
'Test Formatter - test - CRITICAL - ' + self.log_msg)
self.expected_body = {'messages': [{'data': self.expected_payload}]}
self.handler = pubsub_logging.AsyncPubsubHandler(
topic=self.topic,
client=self.mocked_client,
retry=self.RETRY,
worker_num=1,
publish_body=self.mock_publish_body)
self.handler.setFormatter(
logging.Formatter(
'Test Formatter - %(name)s - %(levelname)s - %(message)s'))
r = logging.LogRecord('test', logging.CRITICAL, None, 0, self.log_msg,
[], None)
self.handler.emit(r)
self.handler.close()
def send_loop(client, q, topic, retry, logger, format_func,
publish_body): # pragma: NO COVER
"""Process loop for indefinitely sending logs to Cloud Pub/Sub.
Args:
client: Pub/Sub client. If it's None, a new client will be created.
q: mp.JoinableQueue instance to get the message from.
topic: Cloud Pub/Sub topic name to send the logs.
retry: How many times to retry upon Cloud Pub/Sub API failure.
logger: A logger for informing failures within this function.
format_func: A callable for formatting the logs.
publish_body: A callable for sending the logs.
"""
if client is None:
client = get_pubsub_client()
while True:
try:
logs = q.get()
except Empty:
continue
try:
body = {
'messages': [{
'data': compat_urlsafe_b64encode(format_func(r))
} for r in logs]
}
publish_body(client, body, topic, retry)
except errors.RecoverableError as e:
# Records the exception and puts the logs back to the deque
# and prints the exception to stderr.
q.put(logs)
logger.exception(e)
except Exception as e:
logger.exception(e)
logger.warn('There was a non recoverable error, exiting.')
return
q.task_done()
def test_compat_urlsafe_b64encode(self):
v = 'test'
expected = 'dGVzdA=='
result = compat_urlsafe_b64encode(v)
self.assertEqual(expected, result)