handle SIGPIPE errors when piping output to another process (#73)

* handle SIGPIPE errors when piping output to another process

* fix SIGINT handling within threaded queue

* attempt to fix CI tests

* remove test print

* another attempt at CI tests fix

* fix handler to store cursor in handle_results

* use queue.join() unless we're running on python2

* remove SIGPIPE handling

* - fix logging stacks
- fix printed exception on broken pipe due to stdout being closed prematurely

* just catch broken pipe so we don't close stdout/err on tests

* remove py2 workaround from Worker

* - apparently queue.join() doesn't handle SIGINT on Windows even on py3
- print on exit to not leave progress bar on current line

* uncomment flush func

Author: timabrmsn

src/code42cli/cmds/alerts/extraction.py

@@ -37,7 +37,7 @@ def extract(sdk, profile, output_logger, args):
             args: Command line args used to build up alert query filters.
     """
     store = AlertCursorStore(profile.name) if args.incremental else None
-    handlers = create_handlers(output_logger, store, event_key=u"alerts", sdk=sdk)
+    handlers = create_handlers(sdk, AlertExtractor, output_logger, store)
     extractor = AlertExtractor(sdk, handlers)
     if args.advanced_query:
         exit_if_advanced_query_used_with_other_search_args(args)

src/code42cli/cmds/search_shared/extraction.py

@@ -33,7 +33,8 @@ def verify_begin_date_requirements(args, cursor_store):
         exit(1)
 
 
-def create_handlers(output_logger, cursor_store, event_key, sdk=None):
+def create_handlers(sdk, extractor_class, output_logger, cursor_store):
+    extractor = extractor_class(sdk, ExtractionHandlers())
     handlers = ExtractionHandlers()
     handlers.TOTAL_EVENTS = 0
 
@@ -52,19 +53,21 @@ def handle_error(exception):
         handlers.get_cursor_position = cursor_store.get_stored_cursor_timestamp
 
     @warn_interrupt(
-        warning=u"Cancelling operation cleanly to keep checkpoint data accurate. One moment..."
+        warning=u"Attempting to cancel cleanly to keep checkpoint data accurate. One moment..."
     )
     def handle_response(response):
         response_dict = json.loads(response.text)
-        events = response_dict.get(event_key)
-        if event_key == u"alerts":
+        events = response_dict.get(extractor._key)
+        if extractor._key == u"alerts":
             try:
                 events = get_alert_details(sdk, events)
             except Exception as ex:
                 handlers.handle_error(ex)
         handlers.TOTAL_EVENTS += len(events)
         for event in events:
             output_logger.info(event)
+        last_event_timestamp = extractor._get_timestamp_from_item(event)
+        handlers.record_cursor_position(last_event_timestamp)
 
     handlers.handle_response = handle_response
     return handlers

src/code42cli/cmds/securitydata/extraction.py

@@ -28,7 +28,7 @@ def extract(sdk, profile, output_logger, args):
             args: Command line args used to build up file event query filters.
     """
     store = FileEventCursorStore(profile.name) if args.incremental else None
-    handlers = create_handlers(output_logger, store, event_key=u"fileEvents")
+    handlers = create_handlers(sdk, FileEventExtractor, output_logger, store)
     extractor = FileEventExtractor(sdk, handlers)
     if args.advanced_query:
         exit_if_advanced_query_used_with_other_search_args(args)

src/code42cli/logger.py

@@ -6,6 +6,8 @@
 from code42cli.compat import str
 from code42cli.util import get_user_project_path, is_interactive, color_text_red
 
+# prevent loggers from printing stacks to stderr if a pipe is broken
+logging.raiseExceptions = False
 
 logger_deps_lock = Lock()
 ERROR_LOG_FILE_NAME = u"code42_errors.log"

src/code42cli/main.py

@@ -15,10 +15,12 @@
 from code42cli.cmds.profile import ProfileSubcommandLoader
 from code42cli.commands import Command, SubcommandLoader
 from code42cli.invoker import CommandInvoker
+from code42cli.util import flush_stds_out_err_without_printing_error
 
 
 # Handle KeyboardInterrupts by just exiting instead of printing out a stack
 def exit_on_interrupt(signal, frame):
+    print()
     sys.exit(1)
 
 
@@ -118,7 +120,10 @@ def _create_legal_hold_loader(self):
 def main():
     top = Command(u"", u"", subcommand_loader=MainSubcommandLoader(u""))
     invoker = CommandInvoker(top)
-    invoker.run(sys.argv[1:])
+    try:
+        invoker.run(sys.argv[1:])
+    finally:
+        flush_stds_out_err_without_printing_error()
 
 
 if __name__ == u"__main__":

src/code42cli/util.py

@@ -50,6 +50,22 @@ def is_interactive():
     return sys.stdin.isatty()
 
 
+def flush_stds_out_err_without_printing_error():
+    """Workaround for bug in python3 that causes exception to be printed on broken pipe: 
+    https://bugs.python.org/issue11380
+    """
+    try:
+        sys.stdout.flush()
+    except BrokenPipeError:
+        try:
+            sys.stdout.close()
+        except BrokenPipeError:
+            try:
+                sys.stderr.flush()
+            except BrokenPipeError:
+                sys.stderr.close()
+
+
 def get_url_parts(url_str):
     parts = url_str.split(u":")
     port = None
@@ -129,19 +145,20 @@ def my_important_func():
 
     def __init__(self, warning="Cancelling operation cleanly, one moment... "):
         self.warning = warning
-        self.old_handler = None
+        self.old_int_handler = None
         self.interrupted = False
         self.exit_instructions = "Hit CTRL-C again to force quit."
 
     def __enter__(self):
-        self.old_handler = getsignal(SIGINT)
+        self.old_int_handler = getsignal(SIGINT)
         signal(SIGINT, self._handle_interrupts)
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if self.interrupted:
             exit(1)
-        signal(SIGINT, self.old_handler)
+        signal(SIGINT, self.old_int_handler)
+
         return False
 
     def _handle_interrupts(self, sig, frame):

src/code42cli/worker.py

@@ -1,4 +1,5 @@
 from threading import Thread, Lock
+from time import sleep
 
 from py42.exceptions import Py42HTTPError, Py42ForbiddenError
 
@@ -81,7 +82,8 @@ def stats(self):
     def wait(self):
         """Wait for the tasks in the queue to complete. This should usually be called before 
         program termination."""
-        self._queue.join()
+        while not self._queue.empty():
+            sleep(0.5)
 
     def _process_queue(self):
         while True: