feat: include context on batch log errors

google/cloud/logging_v2/logger.py

@@ -15,6 +15,7 @@
 """Define API Loggers."""
 
 import collections
+import re
 
 from google.cloud.logging_v2._helpers import _add_defaults_to_filter
 from google.cloud.logging_v2.entries import LogEntry
@@ -25,6 +26,9 @@
 from google.cloud.logging_v2.handlers._monitored_resources import detect_resource
 from google.cloud.logging_v2._instrumentation import _add_instrumentation
 
+from google.api_core.exceptions import InvalidArgument
+from google.rpc.error_details_pb2 import DebugInfo
+
 import google.protobuf.message
 
 _GLOBAL_RESOURCE = Resource(type="global", labels={})
@@ -457,6 +461,37 @@ def commit(self, *, client=None):
             kwargs["labels"] = self.logger.labels
 
         entries = [entry.to_api_repr() for entry in self.entries]
-
-        client.logging_api.write_entries(entries, **kwargs)
+        try:
+            client.logging_api.write_entries(entries, partial_success=True, **kwargs)
+        except InvalidArgument as e:
+            # InvalidArgument is often sent when a log is too large
+            # attempt to attach extra contex on which log caused error
+            self._append_context_to_error(e)
+            raise e
         del self.entries[:]
+
+    def _append_context_to_error(self, err):
+        """
+        Attempts to Modify `write_entries` exception messages to contain
+        context on which log in the batch caused the error.
+
+        Best-effort basis. If another exception occurs while processing the
+        input exception, the input will be left unmodified
+
+        Args:
+            err (~google.api_core.exceptions.InvalidArgument):
+                The original exception object
+        """
+        try:
+            # find debug info proto if in details
+            debug_info = next(x for x in err.details if isinstance(x, DebugInfo))
+            # parse out the index of the faulty entry
+            error_idx = re.search("(?<=key: )[0-9]+", debug_info.detail).group(0)
+            # find the faulty entry object
+            found_entry = self.entries[int(error_idx)]
+            str_entry = str(found_entry.to_api_repr())
+            # modify error message to contain extra context
+            err.message = f"{err.message}: {str_entry:.2000}..."
+        except Exception:
+            # if parsing fails, abort changes and leave err unmodified
+            pass

tests/unit/test_logger.py

@@ -1697,6 +1697,79 @@ def test_context_mgr_failure(self):
         self.assertEqual(list(batch.entries), UNSENT)
         self.assertIsNone(api._write_entries_called_with)
 
+    def test_append_context_to_error(self):
+        """
+        If an InvalidArgument exception contains info on the log that threw it,
+        we should be able to add it to the exceptiom message. If not, the
+        exception should be unchanged
+        """
+        from google.api_core.exceptions import InvalidArgument
+        from google.rpc.error_details_pb2 import DebugInfo
+        from google.cloud.logging import TextEntry
+
+        logger = _Logger()
+        client = _Client(project=self.PROJECT)
+        batch = self._make_one(logger, client=client)
+        test_entries = [TextEntry(payload=str(i)) for i in range(11)]
+        batch.entries = test_entries
+        starting_message = "test"
+        # test that properly formatted exceptions add log details
+        for idx, entry in enumerate(test_entries):
+            api_entry = entry.to_api_repr()
+            err = InvalidArgument(
+                starting_message, details=["padding", DebugInfo(detail=f"key: {idx}")]
+            )
+            batch._append_context_to_error(err)
+            self.assertEqual(err.message, f"{starting_message}: {str(api_entry)}...")
+            self.assertIn(str(idx), str(entry))
+        # test with missing debug info
+        err = InvalidArgument(starting_message, details=[])
+        batch._append_context_to_error(err)
+        self.assertEqual(
+            err.message, starting_message, "message should have been unchanged"
+        )
+        # test with missing key
+        err = InvalidArgument(
+            starting_message, details=["padding", DebugInfo(detail=f"no k e y here")]
+        )
+        batch._append_context_to_error(err)
+        self.assertEqual(
+            err.message, starting_message, "message should have been unchanged"
+        )
+        # test with key out of range
+        err = InvalidArgument(
+            starting_message, details=["padding", DebugInfo(detail=f"key: 100")]
+        )
+        batch._append_context_to_error(err)
+        self.assertEqual(
+            err.message, starting_message, "message should have been unchanged"
+        )
+
+    def test_batch_error_gets_context(self):
+        """
+        Simulate an InvalidArgument sent as part of a batch commit, to ensure
+        _append_context_to_error is thrown
+        """
+        from google.api_core.exceptions import InvalidArgument
+        from google.rpc.error_details_pb2 import DebugInfo
+        from google.cloud.logging import TextEntry
+
+        logger = _Logger()
+        client = _Client(project=self.PROJECT)
+        starting_message = "hello"
+        exception = InvalidArgument(
+            starting_message, details=[DebugInfo(detail=f"key: 1")]
+        )
+        client.logging_api = _DummyLoggingExceptionAPI(exception)
+        batch = self._make_one(logger, client=client)
+        test_entries = [TextEntry(payload=str(i)) for i in range(11)]
+        batch.entries = test_entries
+        with self.assertRaises(InvalidArgument) as e:
+            batch.commit()
+            expected_log = test_entries[1]
+            api_entry = expected_log.to_api_repr()
+            self.assertEqual(e.message, f"{starting_message}: {str(api_entry)}...")
+
 
 class _Logger(object):
 
@@ -1725,6 +1798,25 @@ def logger_delete(self, logger_name):
         self._logger_delete_called_with = logger_name
 
 
+class _DummyLoggingExceptionAPI(object):
+    def __init__(self, exception):
+        self.exception = exception
+
+    def write_entries(
+        self,
+        entries,
+        *,
+        logger_name=None,
+        resource=None,
+        labels=None,
+        partial_success=False,
+    ):
+        raise self.exception
+
+    def logger_delete(self, logger_name):
+        raise self.exception
+
+
 class _Client(object):
     def __init__(self, project, connection=None):
         self.project = project