From 9e52f5e896cce594e3125c5592bf9d3a822a759a Mon Sep 17 00:00:00 2001
From: Italo Garcia <italo.garcia@aiven.io>
Date: Thu, 13 Jun 2024 20:00:39 +0200
Subject: [PATCH] senders: Skips messages larger than 256K

GCP quotas page for the logging v2(https://cloud.google.com/logging/quotas) specifies
that a LogEntry should have at maximum 256KB in size, that is not a hard limit, but
an estimation.
In anycase we should check if any of our LogEntry objects are bigger than the quota
allows and not send those logs and instead let the user know what happened through a
custom message
---
 journalpump/senders/google_cloud_logging.py | 14 +++++++++++++-
 test/test_google_cloud_logging.py           | 18 ++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/journalpump/senders/google_cloud_logging.py b/journalpump/senders/google_cloud_logging.py
index 5c518c3..6dbe1ec 100644
--- a/journalpump/senders/google_cloud_logging.py
+++ b/journalpump/senders/google_cloud_logging.py
@@ -23,6 +23,11 @@ class GoogleCloudLoggingSender(LogSender):
         0: "EMERGENCY",
     }
 
+    # A bit on the safe side, not exactly 256KB but this
+    # is an approximation anyway
+    # according to https://cloud.google.com/logging/quotas
+    _LOG_ENTRY_QUOTA = 256_000
+
     def __init__(self, *, config, googleapiclient_request_builder=None, **kwargs):
         super().__init__(config=config, max_send_interval=config.get("max_send_interval", 0.3), **kwargs)
         credentials = None
@@ -61,7 +66,14 @@ def send_messages(self, *, messages, cursor):
 
         for message in messages:
             msg_str = message.decode("utf8")
-            msg = json.loads(msg_str)
+            # This might not measure exactly 256K but should be
+            # a good enough approximation to handle this error.
+            # In anycase, we are handling this error in the
+            # exception handler below also.
+            if len(message) > GoogleCloudLoggingSender._LOG_ENTRY_QUOTA:
+                msg = { "MESSAGE": "Log entry cannot be logged because its size is greater than GCP logging quota of 256K" }
+            else:
+                msg = json.loads(msg_str)
             timestamp = msg.pop("timestamp", None)
             journald_priority = msg.pop("PRIORITY", None)
 
diff --git a/test/test_google_cloud_logging.py b/test/test_google_cloud_logging.py
index 7614e34..36d02af 100644
--- a/test/test_google_cloud_logging.py
+++ b/test/test_google_cloud_logging.py
@@ -162,3 +162,21 @@ def test_correct_timestamp(self):
             cursor=None,
         )
         assert sender._sent_count == 1  # pylint: disable=protected-access
+
+    def test_big_logentry_is_skip(self):
+        """Check that message was not marked as sent if GoogleApi returns error"""
+        request_builder = self._generate_request_builder(
+            [{"jsonPayload": { "MESSAGE": "Log entry cannot be logged because its size is greater than GCP logging quota of 256K" }}],
+        )
+
+        sender = GoogleCloudLoggingSender(
+            name="googlecloudlogging",
+            reader=mock.Mock(),
+            stats=mock.Mock(),
+            field_filter=None,
+            config=self.CONFIG,
+            googleapiclient_request_builder=request_builder,
+        )
+        message = f'{{"MESSAGE": "{"A" * 257_000}"}}'
+        sender.send_messages(messages=[message.encode()], cursor=None)
+        assert sender._sent_count == 1