diff --git a/swh/journal/cli.py b/swh/journal/cli.py --- a/swh/journal/cli.py +++ b/swh/journal/cli.py @@ -7,6 +7,7 @@ import logging import mmap import os +import signal import time import click @@ -24,6 +25,12 @@ from swh.journal.backfill import JournalBackfiller +def clean_exit(signal, frame): + """Raise a proper SystemExit exception when a signal gets intercepted. This + allows finally blocks to run.""" + raise SystemExit(0) + + @click.group(name='journal', context_settings=CONTEXT_SETTINGS) @click.option('--config-file', '-C', default=None, type=click.Path(exists=True, dir_okay=False,), @@ -48,6 +55,9 @@ ctx.ensure_object(dict) + signal.signal(signal.SIGTERM, clean_exit) + signal.signal(signal.SIGINT, clean_exit) + log_level = ctx.obj.get('log_level', logging.INFO) logging.root.setLevel(log_level) logging.getLogger('kafka').setLevel(logging.INFO) @@ -106,10 +116,6 @@ # Log at most once per minute. logger.info('Processed %d messages.' % nb_messages) last_log_time = time.time() - except KeyboardInterrupt: - ctx.exit(0) - else: - print('Done.') finally: client.close() @@ -225,10 +231,8 @@ # Log at most once per minute. logger.info('Processed %d messages.' % nb_messages) last_log_time = time.time() - except KeyboardInterrupt: - ctx.exit(0) - else: - print('Done.') + finally: + client.close() def main():