From 00543bfff422e88c0f8d5722e3625d2bb148d436 Mon Sep 17 00:00:00 2001
From: William Patton <wllmpttn24@gmail.com>
Date: Tue, 2 Jan 2024 15:14:31 -0800
Subject: [PATCH] Custom BatchRequestError handling in pipeline.request_batch

We can filter out some more of the excess error traceback that isn't helpful to the readers.
---
 gunpowder/pipeline.py | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/gunpowder/pipeline.py b/gunpowder/pipeline.py
index f9976c26..3e4bba8f 100644
--- a/gunpowder/pipeline.py
+++ b/gunpowder/pipeline.py
@@ -1,5 +1,8 @@
-import logging
 from gunpowder.nodes import BatchProvider
+from gunpowder.nodes.batch_provider import BatchRequestError
+
+import logging
+import traceback
 
 logger = logging.getLogger(__name__)
 
@@ -21,13 +24,19 @@ def __str__(self):
 
 
 class PipelineRequestError(Exception):
-    def __init__(self, pipeline, request):
+    def __init__(self, pipeline, request, original_traceback=None):
         self.pipeline = pipeline
         self.request = request
+        self.original_traceback = original_traceback
 
     def __str__(self):
         return (
-            "Exception in pipeline:\n"
+            (
+                ("".join(self.original_traceback) )
+                if self.original_traceback is not None
+                else ""
+            )
+            + "Exception in pipeline:\n"
             f"{self.pipeline}\n"
             "while trying to process request\n"
             f"{self.request}"
@@ -123,6 +132,11 @@ def request_batch(self, request):
 
         try:
             return self.output.request_batch(request)
+        except BatchRequestError as e:
+            tb = traceback.format_exception(type(e), e, e.__traceback__)
+            if isinstance(e, BatchRequestError):
+                tb = tb[-1:]
+            raise PipelineRequestError(self, request, original_traceback=tb) from None
         except Exception as e:
             raise PipelineRequestError(self, request) from e