Merge pull request #32 from mcriscenzo/fix/drop_and_keep_key_filtering

Fix - Fetch - Key filtering

Author: tomreitz

lightbeam/__main__.py

@@ -127,7 +127,7 @@ def main(argv=None):
         logger=logger,
         selector=args.selector or "*",
         exclude=args.exclude or "",
-        keep_keys=args.keep_keys or "",
+        keep_keys=args.keep_keys or "*",
         drop_keys=args.drop_keys or "",
         query=args.query or "{}",
         params=args.params,

lightbeam/fetch.py

@@ -89,8 +89,9 @@ async def get_endpoint_records(self, endpoint, limit, offset, file_handle=None):
                                 payload_keys = list(values[0].keys())
                                 final_keys = util.apply_selections(payload_keys, self.lightbeam.keep_keys, self.lightbeam.drop_keys)
                                 do_key_filtering = len(payload_keys) != len(final_keys)
+
                                 for v in values:
-                                    if do_key_filtering: row = {k: v[k] for k in final_keys}
+                                    if do_key_filtering: row = {k: v.get(k, None) for k in final_keys} #v.get() to account for missing keys
                                     else: row = v
                                     if file_handle: file_handle.write(json.dumps(row)+"\n")
                                     else: self.lightbeam.results.append(row)

lightbeam/lightbeam.py

@@ -50,7 +50,7 @@ class Lightbeam:
     MAX_STATUS_REASONS_TO_DISPLAY = 10
     DATA_FILE_EXTENSIONS = ['json', 'jsonl', 'ndjson']
 
-    def __init__(self, config_file, logger=None, selector="*", exclude="", keep_keys="", drop_keys="", query="{}", params="", wipe=False, force=False, older_than="", newer_than="", resend_status_codes="", results_file=""):
+    def __init__(self, config_file, logger=None, selector="*", exclude="", keep_keys="*", drop_keys="", query="{}", params="", wipe=False, force=False, older_than="", newer_than="", resend_status_codes="", results_file=""):
         self.config_file = config_file
         self.logger = logger
         self.errors = 0

lightbeam/util.py

@@ -1,6 +1,7 @@
 import re
 import json
 import itertools
+import copy
 
 # Strips newlines from a string
 # Replace single-quotes with backticks
@@ -59,7 +60,8 @@ def apply_selections(keys, keep, drop):
         for payload_key, keep_key in list(itertools.product(keys, keep)):
             if (keys_match(payload_key, keep_key)):
                 final_keys.append(payload_key)
-    else: final_keys = keys
+    # copy rather than direct pointer assignment, so this function doesn't modify the payload_keys variable in parent code
+    else: final_keys = keys.copy()
     # remove from `final_keys` keys that match `drop`
     if drop and drop != [""]:
         for payload_key, drop_key in list(itertools.product(keys, drop)):