From bb75d24f573c02c9c3bef722fa32a0bc5351ef99 Mon Sep 17 00:00:00 2001
From: Chris Wang <chris.wang@pointsbet.com>
Date: Thu, 17 Nov 2022 14:53:35 +1100
Subject: [PATCH] use batch update instead of update incidents one by one; add
 progress bay; retry smaller batch when batch fails

---
 .../mass_update_incidents.py                  | 50 ++++++++++++-------
 mass_update_incidents/requirements.txt        |  3 +-
 2 files changed, 35 insertions(+), 18 deletions(-)

diff --git a/mass_update_incidents/mass_update_incidents.py b/mass_update_incidents/mass_update_incidents.py
index fd731c8..01006e2 100755
--- a/mass_update_incidents/mass_update_incidents.py
+++ b/mass_update_incidents/mass_update_incidents.py
@@ -10,6 +10,7 @@
 import pprint
 
 import pdpyras
+from tqdm import tqdm
 import time
 
 # Default parameters:
@@ -55,24 +56,39 @@ def mass_update_incidents(args):
         PARAMETERS['date_range'] = 'all'
         print("Getting incidents of all time")
     print("Parameters: "+str(PARAMETERS))
-    try:
-        print("Please be patient as this can take a while for large volumes "
-            "of incidents.")
-        for incident in session.list_all('incidents', params=PARAMETERS):
-            print("* Incident {}: {}".format(incident['id'], args.action))
-            if args.dry_run:
-                continue
+    print("Please be patient as this can take a while for large volumes "
+        "of incidents.")
+
+    incidents = session.list_all('incidents', params=PARAMETERS)
+    jsonPayload = []
+    for incident in incidents:
+        jsonPayload.append({
+            'type': 'incident_reference',
+            'id': incident['id'],
+            'status': '{0}d'.format(args.action), # acknowledged or resolved
+        })
+
+    if args.dry_run:
+        print(jsonPayload)
+    else:
+        update(session, jsonPayload)
+
+
+def update(session, jsonPayload, batchSize=100):
+    url =  "https://api.pagerduty.com/incidents"
+    chunks=[jsonPayload[i:i + batchSize] for i in range(0, len(jsonPayload), batchSize)]
+    for chunk in tqdm(chunks):
+        try:
+            session.rput(url, json=chunk)
             time.sleep(0.25)
-            self_url =  f"https://api.pagerduty.com/incidents/{incident['id']}"
-            session.rput(self_url, json={
-                'type': 'incident_reference',
-                'id': incident['id'],
-                'status': '{0}d'.format(args.action), # acknowledged or resolved
-            })
-    except pdpyras.PDClientError as e:
-        if e.response is not None:
-            print(e.response.text)
-        raise e
+        except Exception as e:
+            if (batchSize == 1):
+                print("fail update following item")
+                print(chunk)
+                raise e
+            else:
+                print("Current chunk failed, trying smaller batch size")
+                update(session, chunk, batchSize/10)
 
 def main(argv=None):
     ap = argparse.ArgumentParser(description="Mass ack or resolve incidents "
diff --git a/mass_update_incidents/requirements.txt b/mass_update_incidents/requirements.txt
index 11eb542..a0859ce 100644
--- a/mass_update_incidents/requirements.txt
+++ b/mass_update_incidents/requirements.txt
@@ -1 +1,2 @@
-pdpyras >= 2.0.0
+pdpyras >= 2.1.0
+tqdm >= 4.64.0
\ No newline at end of file