catch signals

mongoprof.py

@@ -3,18 +3,30 @@ __version__ = '0.1.0'
 
 import time
 import datetime
+import signal
 import argparse
 
 from pymongo import Connection
 from termcolor import colored
 
+quit = False
 
 def watch(dbname, refresh):
+    global quit
     db = getattr(Connection('localhost'), dbname)
     db.set_profiling_level(2)
     last_ts = datetime.datetime.utcnow()
     exclude_name = '{0}.system.profile'.format(dbname)
-    while True:
+
+    def ctrl_c(signal, frame):
+        global quit
+        print('returning profiling level to 0...')
+        db.set_profiling_level(0)
+        db.system.profile.drop()
+        quit = True
+    signal.signal(signal.SIGINT, ctrl_c)
+
+    while not quit:
         for e in db.system.profile.find({'ns': {'$ne': exclude_name},
                                          'ts': {'$gt': last_ts}}):
             output = []
@@ -35,7 +47,7 @@ def watch(dbname, refresh):
             else:
                 output.append(colored('unknown operation: {op}'.format(**e),
                                       'red'))
-                print e
+                print(e)
 
             if 'nscanned' in e:
                 output.append(colored('scanned {nscanned}'.format(**e),
@@ -48,12 +60,9 @@ def watch(dbname, refresh):
             if e.get('scanAndOrder'):
                 output.append(colored('scanAndOrder', 'red'))
             output.append(colored('{millis}ms'.format(**e), 'green'))
-            print ' '.join(output)
+            print(' '.join(output))
             last_ts = e['ts']
         time.sleep(refresh)
-    db.set_profiling_level(0)
-    db.system.profile.drop()
-
 
 def main():
     parser = argparse.ArgumentParser(description='watch mongo queries')