databricks-sql-python/examples/v3_retries_query_execute.py at main · databricks/databricks-sql-python

45 lines (41 loc) · 2.57 KB

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.

# This new retry behaviour allows you to configure the maximum number of redirects that the connector will follow.

# Just set `_retry_max_redirects` to the integer number of redirects you want to allow. The default is None,

# which means all redirects will be followed. In this case, a redirect will count toward the

# _retry_stop_after_attempts_count which means that by default the connector will not enter an endless retry loop.

# For complete information about configuring retries, see the docstring for databricks.sql.thrift_backend.ThriftBackend

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],

_retry_max_redirects=2,

) as connection:

with connection.cursor() as cursor:

cursor.execute("SELECT * FROM default.diamonds LIMIT 2")

result = cursor.fetchall()

for row in result:

print(row)

Provide feedback