diff --git a/tools/datafix/refresh_ids.py b/tools/datafix/refresh_ids.py index a749e9dee9e..45346ba42cd 100644 --- a/tools/datafix/refresh_ids.py +++ b/tools/datafix/refresh_ids.py @@ -15,6 +15,10 @@ MAX_BATCH_SIZE = 500 +class DryRunException(Exception): + """This exception is raised to cancel a transaction during dry runs""" + + def get_relevant_ids(verbose: bool) -> list[str]: relevant_ids = [] @@ -40,12 +44,15 @@ def get_relevant_ids(verbose: bool) -> list[str]: def reput_bugs(dryrun: bool, verbose: bool) -> None: """ Reput all bugs from a given source.""" + relevant_ids = [] # Uncomment below to load the state and skip the get_relevant_ids func - # relevant_ids = json.load(open('relevant_ids.json', 'r')) + # with open('relevant_ids.json', 'r') as f: + # relevant_ids = json.load(open('relevant_ids.json', 'r')) relevant_ids = get_relevant_ids(verbose) # Store the state incase we cancel halfway to avoid having to do the initial query again - json.dump(relevant_ids, open('relevant_ids.json', 'w')) + with open('relevant_ids.json', 'w') as f: + json.dump(relevant_ids, f) num_reputted = 0 time_start = time.perf_counter() @@ -68,7 +75,7 @@ def _reput_ndb(batch: int): if dryrun: print("Dry run mode. Preventing transaction from committing") - raise Exception("Dry run mode") # pylint: disable=broad-exception-raised + raise DryRunException print(f"Time elapsed: {(time.perf_counter() - time_start):.2f} seconds.") @@ -80,14 +87,13 @@ def _reput_ndb(batch: int): f"Reput {num_reputted} bugs... - {num_reputted/len(relevant_ids)*100:.2f}%" ) ndb.transaction(functools.partial(_reput_ndb, batch)) - except Exception as e: + except DryRunException: # Don't have the first batch's transaction-aborting exception stop # subsequent batches from being attempted. - if dryrun and e.args[0].startswith("Dry run mode"): - print("Dry run mode. Preventing transaction from committing") - else: - print([r for r in relevant_ids[batch:batch + MAX_BATCH_SIZE]]) - print(f"Exception {e} occurred. Continuing to next batch.") + print("Dry run mode. Preventing transaction from committing") + except Exception as e: + print([r for r in relevant_ids[batch:batch + MAX_BATCH_SIZE]]) + print(f"Exception {e} occurred. Continuing to next batch.") print("Reputted!")