add a few more options

Author: jettero

contrib/disk-queue-info.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python
 
+from __future__ import print_function
+
 import os
 import argparse
 import json
@@ -14,7 +16,9 @@ def get_args(*a):
     parser.add_argument('-e', '--evil-decode', action='store_true',
         help='attempt to decode all json found in queue (implies -p) and pretty-print them')
     parser.add_argument('-c', '--with-color', action='store_true',
-        help='attempt to pull in pygments and colorize the output of -e (implies -pe)')
+        help='attempt to pull in pygments and colorize the output of -e (implies -p and -e)')
+    parser.add_argument('-m', '--show-meta', action='store_true',
+        help='attempt to show meta-data as a comment in peek mode or a spurious _META_ field in evil mode')
     args = parser.parse_args(*a)
     if args.with_color:
         args.evil_decode = True
@@ -30,15 +34,21 @@ def evil_decode(docbytes):
     decoder = json.JSONDecoder()
     idx = WHITESPACE.match(docbytes, 0).end()
     while idx < len(docbytes):
-        obj, end = decoder.raw_decode(docbytes, idx)
-        yield obj
-        idx = WHITESPACE.match(docbytes, end).end()
+        try:
+            obj, end = decoder.raw_decode(docbytes, idx)
+            yield obj
+            idx = WHITESPACE.match(docbytes, end).end()
+        except ValueError as e:
+            print('docbytes:\n', docbytes)
+            raise
 
-def read_entries(dirname, evil=False, color=False):
+def read_entries(dirname, evil=False, color=False, meta=False):
     dq = DiskQueue(dirname)
-    for item in dq.iter_peek():
+    for item,meta in dq.iter_peek():
         if evil:
             for obj in evil_decode(item):
+                if meta:
+                    obj['_META_'] = meta
                 obj = json.dumps(obj, sort_keys=True, indent=2)
                 if color:
                     try:
@@ -49,6 +59,8 @@ def read_entries(dirname, evil=False, color=False):
                         pass
                 print(obj)
         else:
+            if meta:
+                print('# ', json.dumps(meta))
             print(item)
 
 def main(args):
@@ -57,11 +69,11 @@ def main(args):
             if not os.path.isdir(dirname):
                 raise Exception("directory does not exist")
             if args.peek:
-                read_entries(dirname, evil=args.evil_decode, color=args.with_color)
+                read_entries(dirname, evil=args.evil_decode, color=args.with_color, meta=args.show_meta)
             else:
                 show_info(dirname)
         except Exception as e:
-            print("# unable to read {}: {}".format(dirname, e))
+            print("# Exception while reading {}: {}".format(dirname, repr(e)))
 
 if __name__ == '__main__':
     try:

hubblestack/hec/dq.py

@@ -169,7 +169,7 @@ def iter_peek(self):
         ''' iterate and return all items in the disk queue (without removing any) '''
         for fname in self.files:
             with open(fname, 'rb') as fh:
-                yield self.decompress(fh.read())
+                yield self.decompress(fh.read()), self.read_meta(fname)
 
     def get(self):
         """ get the next item from the queue