diff --git a/tools/run_tests/python_utils/upload_rbe_results.py b/tools/run_tests/python_utils/upload_rbe_results.py index 36fb5f2537b..f5bf3ccc098 100755 --- a/tools/run_tests/python_utils/upload_rbe_results.py +++ b/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 = [] @@ -189,8 +216,9 @@ if __name__ == "__main__": elif 'tests' not in action['testAction']['testSuite']: continue else: - test_cases = action['testAction']['testSuite']['tests'][0][ - 'testSuite']['tests'] + test_cases = [] + for tests_item in action['testAction']['testSuite']['tests']: + test_cases += tests_item['testSuite']['tests'] for test_case in test_cases: if any(s in test_case['testCase'] for s in ['errors', 'failures']): result = 'FAILED' @@ -247,6 +275,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.')