From 1e636b708255a93771f0117985b370358ab04456 Mon Sep 17 00:00:00 2001
From: Jesse Whitehouse <jesse.whitehouse@databricks.com>
Date: Mon, 7 Aug 2023 19:38:33 -0400
Subject: [PATCH] WIP: I'm working on making new() a bit easier to understand

This is hard to debug when it breaks

Signed-off-by: Jesse Whitehouse <jesse.whitehouse@databricks.com>
---
 src/databricks/sql/auth/thrift_http_client.py | 60 ++++++++++++-------
 1 file changed, 38 insertions(+), 22 deletions(-)

diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py
index 922788c6..47982d9f 100644
--- a/src/databricks/sql/auth/thrift_http_client.py
+++ b/src/databricks/sql/auth/thrift_http_client.py
@@ -56,18 +56,22 @@ class DatabricksRetryPolicy(Retry):
 
     Retry.allowed_methods is set to POST
     Retry.status_forcelist is set to codes 429 and 503
+
+    kwargs is used by urllib3 to pass state between successive Retry iterations. 
     """
 
     def __init__(
         self,
-        *args,
         delay_min: float,
         delay_max: float,
         stop_after_attempts_count: int,
         stop_after_attempts_duration: float,
         delay_default: float,
         force_dangerous_codes: List[int],
-        **kwargs,
+        _retry_start_time: Optional[float] = None,
+        _command_type: Optional[CommandType] = None,
+        _attempts_remaining: Optional[Union[bool, int,  None]] = None,
+        extra_urllib3_kwargs={},
     ):
 
         self.backoff_factor: float = delay_min
@@ -77,25 +81,29 @@ def __init__(
         self.force_dangerous_codes = force_dangerous_codes
         self.delay_max = delay_max
 
-        self._retry_start_time = kwargs.pop("_retry_start_time", None)
-        self.command_type = kwargs.pop("_command_type", CommandType.OTHER)
+        self._retry_start_time = _retry_start_time
+        self.command_type = _command_type
 
-        _total = min(
-            stop_after_attempts_count, kwargs.pop("total", stop_after_attempts_count)
-        )
+        
+        # attempts remaining
+        self._attempts_remaining = min(_attempts_remaining or stop_after_attempts_count, stop_after_attempts_count)
 
         super().__init__(
-            total=_total,
+            total=self._attempts_remaining,
             respect_retry_after_header=True,
             backoff_max=self.delay_max,
             allowed_methods=["POST"],
             status_forcelist=[429, 503, *self.force_dangerous_codes],
-            *args,
-            **kwargs,  # type: ignore
+            **extra_urllib3_kwargs,  # type: ignore
         )
 
-    def new(self, **kw: typing.Any) -> Retry:
+    def new(self, **kw_from_urllib3: typing.Any) -> Retry:
         """urllib3 calls Retry.new() between successive requests
+
+        All the contents of `kw` are what urllib3 passes to new(), which in turn calls __init__
+        Since we don't override the arguments that urllib3 passes internally, we need to map them
+        to the new function signature of DatabricksRetryPolicy
+
         We need to override this method to include mapping of our extra keyword arguments:
             delay_min
             delay_max
@@ -103,7 +111,22 @@ def new(self, **kw: typing.Any) -> Retry:
             stop-after_attemps_duration
             delay_default
         """
-        params = dict(
+
+        # These are the arguments needed to implement our custom behaviour
+        databricks_params = dict(
+            delay_min=self.backoff_factor,
+            delay_max=self.backoff_max,  # type: ignore
+            stop_after_attempts_count=self.total,
+            stop_after_attempts_duration=self.stop_after_attempts_duration,
+            delay_default=self.delay_default,
+            force_dangerous_codes=self.force_dangerous_codes,
+            _retry_start_time=self._retry_start_time,
+            _command_type=self.command_type,
+            _attempts_remaining = self._attempts_remaining
+        )
+
+        # These are arguments that urllib3 needs to pass through successive Retry iterations
+        urllib3_params = dict(
             connect=self.connect,
             read=self.read,
             redirect=self.redirect,
@@ -115,18 +138,11 @@ def new(self, **kw: typing.Any) -> Retry:
             history=self.history,
             remove_headers_on_redirect=self.remove_headers_on_redirect,
             backoff_jitter=self.backoff_jitter,  # type: ignore
-            delay_min=self.backoff_factor,
-            delay_max=self.backoff_max,  # type: ignore
-            stop_after_attempts_count=self.total,
-            stop_after_attempts_duration=self.stop_after_attempts_duration,
-            delay_default=self.delay_default,
-            force_dangerous_codes=self.force_dangerous_codes,
-            _retry_start_time=self._retry_start_time,
-            _command_type=self.command_type,
         )
 
-        params.update(kw)
-        return type(self)(**params)  # type: ignore[arg-type]
+        kw_from_urllib3.update(**urllib3_params)
+
+        return type(self)(**databricks_params, extra_urllib3_kwargs=urllib3_params)  # type: ignore[arg-type]
 
     @property
     def command_type(self) -> CommandType: