From e40446c700da36411e588f51666a204c6f8b3a92 Mon Sep 17 00:00:00 2001
From: Holly Gong <39108850+hogo6002@users.noreply.github.com>
Date: Wed, 16 Oct 2024 16:36:54 +1100
Subject: [PATCH] feat(datafix): Allow reputing for each bug ID (#2736)

Add a `--bugs` flag to the `reput_bug` script to allow reputing only
specific bugs.
---
 tools/datafix/reput_bugs.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/tools/datafix/reput_bugs.py b/tools/datafix/reput_bugs.py
index 73e6acb6743..470bf175a96 100644
--- a/tools/datafix/reput_bugs.py
+++ b/tools/datafix/reput_bugs.py
@@ -13,14 +13,17 @@
 MAX_BATCH_SIZE = 500
 
 
-def reput_bugs(dryrun: bool, source: str) -> None:
+def reput_bugs(dryrun: bool, source: str, ids: list) -> None:
   """ Reput all bugs from a given source."""
   query = osv.Bug.query()
-  query = query.filter(osv.Bug.source == source)
-  print(f"Running query {query.filters} "
-        f"on {query.kind}...")
+  if ids:
+    result = [ndb.Key(query.kind, id) for id in ids[0]]
+  else:
+    query = query.filter(osv.Bug.source == source)
+    print(f"Running query {query.filters} "
+          f"on {query.kind}...")
+    result = list(query.fetch(keys_only=True))
 
-  result = list(query.fetch(keys_only=True))
   result.sort(key=lambda r: r.id())
   # result = [r for r in result if not r.id()[0].isnumeric()]
   print(f"Retrieved {len(result)} bugs to examine for reputting")
@@ -78,11 +81,17 @@ def main() -> None:
       dest="project",
       default="oss-vdb-test",
       help="GCP project to operate on")
+  parser.add_argument(
+      "--bugs",
+      action="append",
+      nargs="+",
+      required=False,
+      help=f"The bug IDs to operate on ({MAX_BATCH_SIZE} at most)")
   args = parser.parse_args()
 
   client = ndb.Client(project=args.project)
   with client.context():
-    reput_bugs(args.dryrun, args.source)
+    reput_bugs(args.dryrun, args.source, args.bugs)
 
 
 if __name__ == "__main__":