Convenience improvements to v3 retry logic (#199)

Signed-off-by: Jesse Whitehouse <[email protected]>

Author: Jesse

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## 2.9.x (Unreleased)
 
+- Other: Add `examples/v3_retries_query_execute.py` 
+- Other: suppress log message when `_enable_v3_retries` is not `True`
+- Other: make this connector backwards compatible with `urllib3>=1.0.0`
+
 ## 2.9.1 (2023-08-11)
 
 - Other: Explicitly pin urllib3 to ^2.0.0

examples/README.md

@@ -38,4 +38,5 @@ To run all of these examples you can clone the entire repository to your disk. O
 this example the string `ExamplePartnerTag` will be added to the the user agent on every request.
 - **`staging_ingestion.py`** shows how the connector handles Databricks' experimental staging ingestion commands `GET`, `PUT`, and `REMOVE`.
 - **`sqlalchemy.py`** shows a basic example of connecting to Databricks with [SQLAlchemy](https://www.sqlalchemy.org/). 
-- **`custom_cred_provider.py`** shows how to pass a custom credential provider to bypass connector authentication. Please install databricks-sdk prior to running this example.
+- **`custom_cred_provider.py`** shows how to pass a custom credential provider to bypass connector authentication. Please install databricks-sdk prior to running this example.
+- **`v3_retries_query_execute.py`** shows how to enable v3 retries in connector version 2.9.x including how to enable retries for non-default retry cases.

examples/v3_retries_query_execute.py

@@ -0,0 +1,35 @@
+from databricks import sql
+import os
+
+# Users of connector versions >= 2.9.0 and <= 3.0.0 can use the v3 retry behaviour by setting _enable_v3_retries=True
+# This flag will be deprecated in databricks-sql-connector~=3.0.0 as it will become the default.
+#
+# The new retry behaviour is defined in src/databricks/sql/auth/retry.py
+# 
+# The new retry behaviour allows users to force the connector to automatically retry requests that fail with codes
+# that are not retried by default (in most cases only codes 429 and 503 are retried by default). Additional HTTP
+# codes to retry are specified as a list passed to `_retry_dangerous_codes`.
+#
+# Note that, as implied in the name, doing this is *dangerous* and should not be configured in all usages.
+# With the default behaviour, ExecuteStatement Thrift commands are only retried for codes 429 and 503 because
+# we can be certain at run-time that the statement never reached Databricks compute. These codes are returned by
+# the SQL gateway / load balancer. So there is no risk that retrying the request would result in a doubled
+# (or tripled etc) command execution. These codes are always accompanied by a Retry-After header, which we honour.
+#
+# However, if your use-case emits idempotent queries such as SELECT statements, it can be helpful to retry 
+# for 502 (Bad Gateway) codes etc. In these cases, there is a possibility that the initial command _did_ reach
+# Databricks compute and retrying it could result in additional executions. Retrying under these conditions uses
+# an exponential back-off since a Retry-After header is not present.
+
+with sql.connect(server_hostname = os.getenv("DATABRICKS_SERVER_HOSTNAME"),
+                 http_path       = os.getenv("DATABRICKS_HTTP_PATH"),
+                 access_token    = os.getenv("DATABRICKS_TOKEN"),
+                 _enable_v3_retries = True,
+                 _retry_dangerous_codes=[502,400]) as connection:
+
+  with connection.cursor() as cursor:
+    cursor.execute("SELECT * FROM default.diamonds LIMIT 2")
+    result = cursor.fetchall()
+
+    for row in result:
+      print(row)

src/databricks/sql/thrift_backend.py

@@ -375,9 +375,10 @@ def attempt_request(attempt):
                 unsafe_logger.debug("Sending request: {}".format(request))
 
                 # These three lines are no-ops if the v3 retry policy is not in use
-                this_command_type = CommandType.get(this_method_name)
-                self._transport.set_retry_command_type(this_command_type)
-                self._transport.startRetryTimer()
+                if self.enable_v3_retries:
+                    this_command_type = CommandType.get(this_method_name)
+                    self._transport.set_retry_command_type(this_command_type)
+                    self._transport.startRetryTimer()
 
                 response = method(request)