From 06918fba201d1bc11049e0c289de6ab2bc2a910e Mon Sep 17 00:00:00 2001
From: Anatoli Kurtsevich <anatoli.kurtsevich@relational.ai>
Date: Thu, 19 Oct 2023 16:33:54 -0400
Subject: [PATCH] added retry mechanism to the retry flow

---
 CHANGELOG.md       |  4 ++++
 railib/__init__.py |  2 +-
 railib/rest.py     | 16 +++++++++-------
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b628405..3ec7da6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Changelog
 
+## v0.6.18
+
+* Added retry mechanism for the authentication flow.
+
 ## v0.6.17
 
 * Removed unnecessary dependencies of the SDK, bumped Pyarrow to v10.
diff --git a/railib/__init__.py b/railib/__init__.py
index c4de3da..f599a48 100644
--- a/railib/__init__.py
+++ b/railib/__init__.py
@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version_info__ = (0, 6, 17)
+__version_info__ = (0, 6, 18)
 __version__ = ".".join(map(str, __version_info__))
diff --git a/railib/rest.py b/railib/rest.py
index 77e71da..379142b 100644
--- a/railib/rest.py
+++ b/railib/rest.py
@@ -194,13 +194,15 @@ def _request_access_token(ctx: Context, url: str) -> AccessToken:
         data=data,
     )
     _print_request(req)
-    with urlopen(req) as rsp:
-        result = json.loads(rsp.read())
-        token = result.get(ACCESS_KEY_TOKEN_KEY, None)
-        if token is not None:
-            expires_in = result.get(EXPIRES_IN_KEY, None)
-            scope = result.get(SCOPE, None)
-            return AccessToken(token, scope, expires_in)
+    rsp = _urlopen_with_retry(req, ctx.retries)
+    result = json.loads(rsp.read())
+    token = result.get(ACCESS_KEY_TOKEN_KEY, None)
+
+    if token is not None:
+        expires_in = result.get(EXPIRES_IN_KEY, None)
+        scope = result.get(SCOPE, None)
+        return AccessToken(token, scope, expires_in)
+
     raise Exception("failed to get the access token")