easier debugging for upload_rbe_results.py

tools/run_tests/python_utils/upload_rbe_results.py

@@ -134,15 +134,42 @@ def _get_resultstore_data(api_key, invocation_id):
 
 if __name__ == "__main__":
     # Arguments are necessary if running in a non-Kokoro environment.
-    argp = argparse.ArgumentParser(description='Upload RBE results.')
-    argp.add_argument('--api_key', default='', type=str)
-    argp.add_argument('--invocation_id', default='', type=str)
+    argp = argparse.ArgumentParser(
+        description=
+        'Fetches results for given RBE invocation and uploads them to BigQuery table.'
+    )
+    argp.add_argument('--api_key',
+                      default='',
+                      type=str,
+                      help='The API key to read from ResultStore API')
+    argp.add_argument('--invocation_id',
+                      default='',
+                      type=str,
+                      help='UUID of bazel invocation to fetch.')
+    argp.add_argument('--bq_dump_file',
+                      default=None,
+                      type=str,
+                      help='Dump JSON data to file just before uploading')
+    argp.add_argument('--resultstore_dump_file',
+                      default=None,
+                      type=str,
+                      help='Dump JSON data as received from ResultStore API')
+    argp.add_argument('--skip_upload',
+                      default=False,
+                      action='store_const',
+                      const=True,
+                      help='Skip uploading to bigquery')
     args = argp.parse_args()
 
     api_key = args.api_key or _get_api_key()
     invocation_id = args.invocation_id or _get_invocation_id()
     resultstore_actions = _get_resultstore_data(api_key, invocation_id)
 
+    if args.resultstore_dump_file:
+        with open(args.resultstore_dump_file, 'w') as f:
+            json.dump(resultstore_actions, f, indent=4, sort_keys=True)
+        print('Dumped resultstore data to file %s' % args.resultstore_dump_file)
+
     # google.devtools.resultstore.v2.Action schema:
     # https://github.com/googleapis/googleapis/blob/master/google/devtools/resultstore/v2/action.proto
     bq_rows = []
@@ -247,6 +274,14 @@ if __name__ == "__main__":
                     }
                 })
 
-    # BigQuery sometimes fails with large uploads, so batch 1,000 rows at a time.
-    for i in range((len(bq_rows) / 1000) + 1):
-        _upload_results_to_bq(bq_rows[i * 1000:(i + 1) * 1000])
+    if args.bq_dump_file:
+        with open(args.bq_dump_file, 'w') as f:
+            json.dump(bq_rows, f, indent=4, sort_keys=True)
+        print('Dumped BQ data to file %s' % args.bq_dump_file)
+
+    if not args.skip_upload:
+        # BigQuery sometimes fails with large uploads, so batch 1,000 rows at a time.
+        for i in range((len(bq_rows) / 1000) + 1):
+            _upload_results_to_bq(bq_rows[i * 1000:(i + 1) * 1000])
+    else:
+        print('Skipped upload to bigquery.')