From 05c2e7a508c240a05af28696c9c683a21fcccf39 Mon Sep 17 00:00:00 2001
From: Hugo Rivera Calzadillas <roguh@roguh.com>
Date: Thu, 3 Oct 2024 23:21:52 -0600
Subject: [PATCH] feat: Add the max_retries parameter (#598)

This is passed to requests.adapters.HTTPAdapter to allow for more retries
---
 src/keycloak/connection.py        | 10 ++++++++--
 src/keycloak/keycloak_admin.py    |  6 ++++++
 src/keycloak/keycloak_openid.py   |  6 ++++++
 src/keycloak/openid_connection.py |  3 +++
 4 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/src/keycloak/connection.py b/src/keycloak/connection.py
index 9d55c37..f7b8ffd 100644
--- a/src/keycloak/connection.py
+++ b/src/keycloak/connection.py
@@ -53,9 +53,13 @@ class ConnectionManager(object):
         Either a path to an SSL certificate file, or two-tuple of
         (certificate file, key file).
     :type cert: Union[str,Tuple[str,str]]
+    :param max_retries: The total number of times to retry HTTP requests.
+    :type max_retries: int
     """
 
-    def __init__(self, base_url, headers={}, timeout=60, verify=True, proxies=None, cert=None):
+    def __init__(
+        self, base_url, headers={}, timeout=60, verify=True, proxies=None, cert=None, max_retries=1
+    ):
         """Init method.
 
         :param base_url: The server URL.
@@ -73,6 +77,8 @@ class ConnectionManager(object):
             Either a path to an SSL certificate file, or two-tuple of
             (certificate file, key file).
         :type cert: Union[str,Tuple[str,str]]
+        :param max_retries: The total number of times to retry HTTP requests.
+        :type max_retries: int
         """
         self.base_url = base_url
         self.headers = headers
@@ -85,7 +91,7 @@ class ConnectionManager(object):
         # retry once to reset connection with Keycloak after  tomcat's ConnectionTimeout
         # see https://github.com/marcospereirampj/python-keycloak/issues/36
         for protocol in ("https://", "http://"):
-            adapter = HTTPAdapter(max_retries=1)
+            adapter = HTTPAdapter(max_retries=max_retries)
             # adds POST to retry whitelist
             allowed_methods = set(adapter.max_retries.allowed_methods)
             allowed_methods.add("POST")
diff --git a/src/keycloak/keycloak_admin.py b/src/keycloak/keycloak_admin.py
index 7af64ab..2c61b08 100644
--- a/src/keycloak/keycloak_admin.py
+++ b/src/keycloak/keycloak_admin.py
@@ -77,6 +77,8 @@ class KeycloakAdmin:
         Either a path to an SSL certificate file, or two-tuple of
         (certificate file, key file).
     :type cert: Union[str,Tuple[str,str]]
+    :param max_retries: The total number of times to retry HTTP requests.
+    :type max_retries: int
     :param connection: A KeycloakOpenIDConnection as an alternative to individual params.
     :type connection: KeycloakOpenIDConnection
     """
@@ -99,6 +101,7 @@ class KeycloakAdmin:
         user_realm_name=None,
         timeout=60,
         cert=None,
+        max_retries=1,
         connection: Optional[KeycloakOpenIDConnection] = None,
     ):
         """Init method.
@@ -134,6 +137,8 @@ class KeycloakAdmin:
         :param cert: An SSL certificate used by the requested host to authenticate the client.
             Either a path to an SSL certificate file, or two-tuple of (certificate file, key file).
         :type cert: Union[str,Tuple[str,str]]
+        :param max_retries: The total number of times to retry HTTP requests.
+        :type max_retries: int
         :param connection: An OpenID Connection as an alternative to individual params.
         :type connection: KeycloakOpenIDConnection
         """
@@ -152,6 +157,7 @@ class KeycloakAdmin:
             custom_headers=custom_headers,
             timeout=timeout,
             cert=cert,
+            max_retries=max_retries,
         )
 
     @property
diff --git a/src/keycloak/keycloak_openid.py b/src/keycloak/keycloak_openid.py
index a4453eb..47e6114 100644
--- a/src/keycloak/keycloak_openid.py
+++ b/src/keycloak/keycloak_openid.py
@@ -77,6 +77,8 @@ class KeycloakOpenID:
     :param cert: An SSL certificate used by the requested host to authenticate the client.
         Either a path to an SSL certificate file, or two-tuple of
         (certificate file, key file).
+    :param max_retries: The total number of times to retry HTTP requests.
+    :type max_retries: int
     """
 
     def __init__(
@@ -90,6 +92,7 @@ class KeycloakOpenID:
         proxies=None,
         timeout=60,
         cert=None,
+        max_retries=1,
     ):
         """Init method.
 
@@ -114,6 +117,8 @@ class KeycloakOpenID:
             Either a path to an SSL certificate file, or two-tuple of
             (certificate file, key file).
         :type cert: Union[str,Tuple[str,str]]
+        :param max_retries: The total number of times to retry HTTP requests.
+        :type max_retries: int
         """
         self.client_id = client_id
         self.client_secret_key = client_secret_key
@@ -126,6 +131,7 @@ class KeycloakOpenID:
             verify=verify,
             proxies=proxies,
             cert=cert,
+            max_retries=max_retries,
         )
 
         self.authorization = Authorization()
diff --git a/src/keycloak/openid_connection.py b/src/keycloak/openid_connection.py
index 73ab9d8..2a2e077 100644
--- a/src/keycloak/openid_connection.py
+++ b/src/keycloak/openid_connection.py
@@ -73,6 +73,7 @@ class KeycloakOpenIDConnection(ConnectionManager):
         user_realm_name=None,
         timeout=60,
         cert=None,
+        max_retries=1,
     ):
         """Init method.
 
@@ -108,6 +109,8 @@ class KeycloakOpenIDConnection(ConnectionManager):
             Either a path to an SSL certificate file, or two-tuple of
             (certificate file, key file).
         :type cert: Union[str,Tuple[str,str]]
+        :param max_retries: The total number of times to retry HTTP requests.
+        :type max_retries: int
         """
         # token is renewed when it hits 90% of its lifetime. This is to account for any possible
         # clock skew.