From db4f0881a244c80161d4fdb948c6542c983dfce5 Mon Sep 17 00:00:00 2001
From: "Joshua M. Boniface" <joshua@boniface.me>
Date: Mon, 29 Jan 2024 09:35:14 -0500
Subject: [PATCH] Improve error handling and retries

1. Use the actual response code from the server on error, or 504 on
timeouts instead of 500.
2. Retry GET requests 3 times and only error if the last fails
---
 client-cli/pvc/lib/common.py | 35 ++++++++++++++++++++++++++---------
 1 file changed, 26 insertions(+), 9 deletions(-)

diff --git a/client-cli/pvc/lib/common.py b/client-cli/pvc/lib/common.py
index b2f1f79f..2ec38f72 100644
--- a/client-cli/pvc/lib/common.py
+++ b/client-cli/pvc/lib/common.py
@@ -140,15 +140,31 @@ def call_api(
     # Determine the request type and hit the API
     disable_warnings()
     try:
+        response = None
         if operation == "get":
-            response = requests.get(
-                uri,
-                timeout=timeout,
-                headers=headers,
-                params=params,
-                data=data,
-                verify=config["verify_ssl"],
-            )
+            retry_on_code = [429, 500, 502, 503, 504]
+            for i in range(3):
+                failed = False
+                try:
+                    response = requests.get(
+                        uri,
+                        timeout=timeout,
+                        headers=headers,
+                        params=params,
+                        data=data,
+                        verify=config["verify_ssl"],
+                    )
+                    if response.status_code in retry_on_code:
+                        failed = True
+                        continue
+                except requests.exceptions.ConnectionError:
+                    failed = True
+                    pass
+            if failed:
+                error = f"Code {response.status_code}" if response else "Timeout"
+                raise requests.exceptions.ConnectionError(
+                    f"Failed to connect after 3 tries ({error})"
+                )
         if operation == "post":
             response = requests.post(
                 uri,
@@ -189,7 +205,8 @@ def call_api(
             )
     except Exception as e:
         message = "Failed to connect to the API: {}".format(e)
-        response = ErrorResponse({"message": message}, 500)
+        code = response.status_code if response else 504
+        response = ErrorResponse({"message": message}, code)
 
     # Display debug output
     if config["debug"]: