databricks-sql-connector 2.8.0__tar.gz → 2.9.1.dev1__tar.gz

{databricks_sql_connector-2.8.0 → databricks_sql_connector-2.9.1.dev1}/CHANGELOG.md

@@ -1,10 +1,20 @@
 # Release History
 
-## 2.8.x (Unreleased)
+## 2.9.x (Unreleased)
+
+## 2.9.1 (2023-08-11)
+
+- Other: Explicitly pin urllib3 to ^2.0.0
+
+## 2.9.0 (2023-08-10)
+
+- Replace retry handling with DatabricksRetryPolicy. This is disabled by default. To enable, set `enable_v3_retries=True` when creating `databricks.sql.client`
+- Other: Fix typo in README quick start example
+- Other: Add autospec to Client mocks and tidy up `make_request`
 
 ## 2.8.0 (2023-07-21)
 
-- Add support for Cloud Fetch (#146, #151, #154)
+- Add support for Cloud Fetch. Disabled by default. Set `use_cloud_fetch=True` when building `databricks.sql.client` to enable it (#146, #151, #154)
 - SQLAlchemy has_table function now honours schema= argument and adds catalog= argument (#174)
 - SQLAlchemy set non_native_boolean_check_constraint False as it's not supported by Databricks (#120)
 - Fix: Revised SQLAlchemy dialect and examples for compatibility with SQLAlchemy==1.3.x (#173)

{databricks_sql_connector-2.8.0 → databricks_sql_connector-2.9.1.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: databricks-sql-connector
-Version: 2.8.0
+Version: 2.9.1.dev1
 Summary: Databricks SQL Connector for Python
 License: Apache-2.0
 Author: Databricks
@@ -25,6 +25,7 @@ Requires-Dist: pyarrow (>=6.0.0) ; python_version >= "3.7" and python_version <
 Requires-Dist: requests (>=2.18.1,<3.0.0)
 Requires-Dist: sqlalchemy (>=1.3.24,<2.0.0)
 Requires-Dist: thrift (>=0.16.0,<0.17.0)
+Requires-Dist: urllib3 (>=2.0.0,<3.0.0)
 Project-URL: Bug Tracker, https://github.com/databricks/databricks-sql-python/issues
 Project-URL: Homepage, https://github.com/databricks/databricks-sql-python
 Description-Content-Type: text/markdown
@@ -70,7 +71,7 @@ from databricks import sql
 
 host = os.getenv("DATABRICKS_HOST")
 http_path = os.getenv("DATABRICKS_HTTP_PATH")
-access_token = os.getenv("DATABRICKS_ACCESS_TOKEN")
+access_token = os.getenv("DATABRICKS_TOKEN")
 
 connection = sql.connect(
   server_hostname=host,

{databricks_sql_connector-2.8.0 → databricks_sql_connector-2.9.1.dev1}/README.md

@@ -39,7 +39,7 @@ from databricks import sql
 
 host = os.getenv("DATABRICKS_HOST")
 http_path = os.getenv("DATABRICKS_HTTP_PATH")
-access_token = os.getenv("DATABRICKS_ACCESS_TOKEN")
+access_token = os.getenv("DATABRICKS_TOKEN")
 
 connection = sql.connect(
   server_hostname=host,

{databricks_sql_connector-2.8.0 → databricks_sql_connector-2.9.1.dev1}/pyproject.toml

@@ -1,35 +1,35 @@
 [tool.poetry]
 name = "databricks-sql-connector"
-version = "2.8.0"
+version = "2.9.1dev1"
 description = "Databricks SQL Connector for Python"
 authors = ["Databricks <databricks-sql-connector-maintainers@databricks.com>"]
 license = "Apache-2.0"
 readme = "README.md"
-packages = [{include = "databricks", from = "src"}]
+packages = [{ include = "databricks", from = "src" }]
 include = ["CHANGELOG.md"]
 
 [tool.poetry.dependencies]
 python = "^3.7.1"
 thrift = "^0.16.0"
 pandas = [
-    {version = ">=1.2.5,<1.4.0", python = ">=3.7,<3.8"},
-    {version =">=1.2.5,<3.0.0", python = ">=3.8"}
+    { version = ">=1.2.5,<1.4.0", python = ">=3.7,<3.8" },
+    { version = ">=1.2.5,<3.0.0", python = ">=3.8" },
 ]
-
 pyarrow = [
-    {version = ">=6.0.0", python = ">=3.7,<3.11"},
-    {version = ">=10.0.1", python = ">=3.11"}
+    { version = ">=6.0.0", python = ">=3.7,<3.11" },
+    { version = ">=10.0.1", python = ">=3.11" },
 ]
 lz4 = "^4.0.2"
-requests="^2.18.1"
-oauthlib="^3.1.0"
+requests = "^2.18.1"
+oauthlib = "^3.1.0"
 numpy = [
-    {version = ">=1.16.6", python = ">=3.7,<3.11"},
-    {version = ">=1.23.4", python = ">=3.11"}
+    { version = ">=1.16.6", python = ">=3.7,<3.11" },
+    { version = ">=1.23.4", python = ">=3.11" },
 ]
 sqlalchemy = "^1.3.24"
 openpyxl = "^3.0.10"
 alembic = "^1.0.11"
+urllib3 = "^2.0.0"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.1.2"
@@ -60,9 +60,5 @@ exclude = '/(\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|\.svn|_build|buck
 minversion = "6.0"
 log_cli = "false"
 log_cli_level = "INFO"
-testpaths = [
-    "tests"
-]
-env_files = [
-    "test.env"
-]
+testpaths = ["tests"]
+env_files = ["test.env"]

{databricks_sql_connector-2.8.0 → databricks_sql_connector-2.9.1.dev1}/src/databricks/sql/__init__.py

@@ -28,7 +28,7 @@ DATETIME = DBAPITypeObject("timestamp")
 DATE = DBAPITypeObject("date")
 ROWID = DBAPITypeObject()
 
-__version__ = "2.8.0"
+__version__ = "2.9.1dev1"
 USER_AGENT_NAME = "PyDatabricksSqlConnector"
 
 # These two functions are pyhive legacy

databricks_sql_connector-2.9.1.dev1/src/databricks/sql/auth/retry.py

@@ -0,0 +1,410 @@
+import logging
+import time
+import typing
+from enum import Enum
+from typing import List, Optional, Tuple, Union
+
+from urllib3 import BaseHTTPResponse  # type: ignore
+from urllib3 import Retry
+from urllib3.util.retry import RequestHistory
+
+from databricks.sql.exc import (
+    CursorAlreadyClosedError,
+    MaxRetryDurationError,
+    NonRecoverableNetworkError,
+    OperationalError,
+    SessionAlreadyClosedError,
+    UnsafeToRetryError,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class CommandType(Enum):
+    EXECUTE_STATEMENT = "ExecuteStatement"
+    CLOSE_SESSION = "CloseSession"
+    CLOSE_OPERATION = "CloseOperation"
+    OTHER = "Other"
+
+    @classmethod
+    def get(cls, value: str):
+        value_name_map = {i.value: i.name for i in cls}
+        valid_command = value_name_map.get(value, False)
+        if valid_command:
+            return getattr(cls, str(valid_command))
+        else:
+            return cls.OTHER
+
+
+class DatabricksRetryPolicy(Retry):
+    """
+    Implements our v3 retry policy by extending urllib3's robust default retry behaviour.
+
+    Retry logic varies based on the overall wall-clock request time and Thrift CommandType
+    being issued. ThriftBackend starts a timer and sets the current CommandType prior to
+    initiating a network request. See `self.should_retry()` for details about what we do
+    and do not retry.
+
+    :param delay_min:
+        Float of seconds for the minimum delay between retries. This is an alias for urllib3's
+        `backoff_factor`.
+
+    :param delay_max:
+        Float of seconds for the maximum delay between retries. This is an alias for urllib3's
+        `backoff_max`
+
+    :param stop_after_attempts_count:
+        Integer maximum number of attempts that will be retried. This is an alias for urllib3's
+        `total`.
+
+    :param stop_after_attempts_duration:
+        Float of maximum number of seconds within which a request may be retried starting from
+        the beginning of the first request.
+
+    :param delay_default:
+        Float of seconds the connector will wait between sucessive GetOperationStatus
+        requests. This parameter is not used to retry failed network requests. We include
+        it in this class to keep all retry behaviour encapsulated in this file.
+
+    :param force_dangerous_codes:
+        List of integer HTTP status codes that the connector will retry, even for dangerous
+        commands like ExecuteStatement. This is passed to urllib3 by extending its status_forcelist
+
+    :param urllib3_kwargs:
+        Dictionary of arguments that are passed to Retry.__init__. Any setting of Retry() that
+        Databricks does not override or extend may be modified here.
+    """
+
+    def __init__(
+        self,
+        delay_min: float,
+        delay_max: float,
+        stop_after_attempts_count: int,
+        stop_after_attempts_duration: float,
+        delay_default: float,
+        force_dangerous_codes: List[int],
+        urllib3_kwargs: dict = {},
+    ):
+        # These values do not change from one command to the next
+        self.delay_max = delay_max
+        self.delay_min = delay_min
+        self.stop_after_attempts_count = stop_after_attempts_count
+        self.stop_after_attempts_duration = stop_after_attempts_duration
+        self._delay_default = delay_default
+        self.force_dangerous_codes = force_dangerous_codes
+
+        # the urllib3 kwargs are a mix of configuration (some of which we override)
+        # and counters like `total` or `connect` which may change between successive retries
+        # we only care about urllib3 kwargs that we alias, override, or add to in some way
+
+        # the length of _history increases as retries are performed
+        _history: Optional[Tuple[RequestHistory, ...]] = urllib3_kwargs.get("history")
+
+        if not _history:
+            # no attempts were made so we can retry the current command as many times as specified
+            # by the user
+            _attempts_remaining = self.stop_after_attempts_count
+        else:
+            # at least one of our attempts has been consumed, and urllib3 will have set a total
+            # `total` is a counter that begins equal to self.stop_after_attempts_count and is
+            # decremented after each unsuccessful request. When `total` is zero, urllib3 raises a
+            # MaxRetryError
+            _total: int = urllib3_kwargs.pop("total")
+            _attempts_remaining = _total
+
+        _urllib_kwargs_we_care_about = dict(
+            total=_attempts_remaining,
+            respect_retry_after_header=True,
+            backoff_factor=self.delay_min,
+            backoff_max=self.delay_max,
+            allowed_methods=["POST"],
+            status_forcelist=[429, 503, *self.force_dangerous_codes],
+        )
+
+        urllib3_kwargs.update(**_urllib_kwargs_we_care_about)
+
+        super().__init__(
+            **urllib3_kwargs,  # type: ignore
+        )
+
+    @classmethod
+    def __private_init__(
+        cls, retry_start_time: float, command_type: Optional[CommandType], **init_kwargs
+    ):
+        """
+        Returns a new instance of DatabricksRetryPolicy with the _retry_start_time and _command_type
+        properties already set. This method should only be called by DatabricksRetryPolicy itself between
+        successive Retry attempts.
+
+        :param retry_start_time:
+            Float unix timestamp. Used to monitor the overall request duration across successive
+            retries. Never set this value directly. Use self.start_retry_timer() instead. Users
+            never set this value. It is set by ThriftBackend immediately before issuing a network
+            request.
+
+        :param command_type:
+            CommandType of the current request being retried. Used to modify retry behaviour based
+            on the type of Thrift command being issued. See self.should_retry() for details. Users
+            never set this value directly. It is set by ThriftBackend immediately before issuing
+            a network request.
+
+        :param init_kwargs:
+            A dictionary of parameters that will be passed to __init__ in the new object
+        """
+
+        new_object = cls(**init_kwargs)
+        new_object._retry_start_time = retry_start_time
+        new_object.command_type = command_type
+        return new_object
+
+    def new(self, **urllib3_incremented_counters: typing.Any) -> Retry:
+        """This method is responsible for passing the entire Retry state to its next iteration.
+
+        urllib3 calls Retry.new() between successive requests as part of its `.increment()` method
+        as shown below:
+
+        ```python
+            new_retry = self.new(
+                total=total,
+                connect=connect,
+                read=read,
+                redirect=redirect,
+                status=status_count,
+                other=other,
+                history=history,
+        )
+        ```
+
+        The arguments it passes to `.new()` (total, connect, read, etc.) are those modified by `.increment()`.
+
+        Since self.__init__ has a different signature than Retry.__init__ , we implement our own `self.new()`
+        to pipe our Databricks-specific state while preserving the super-class's behaviour.
+
+        """
+
+        # These arguments will match the function signature for self.__init__
+        databricks_init_params = dict(
+            delay_min=self.delay_min,
+            delay_max=self.delay_max,
+            stop_after_attempts_count=self.stop_after_attempts_count,
+            stop_after_attempts_duration=self.stop_after_attempts_duration,
+            delay_default=self.delay_default,
+            force_dangerous_codes=self.force_dangerous_codes,
+            urllib3_kwargs={},
+        )
+
+        # Gather urllib3's current retry state _before_ increment was called
+        # These arguments match the function signature for Retry.__init__
+        # Note: if we update urllib3 we may need to add/remove arguments from this dict
+        urllib3_init_params = dict(
+            total=self.total,
+            connect=self.connect,
+            read=self.read,
+            redirect=self.redirect,
+            status=self.status,
+            other=self.other,
+            allowed_methods=self.allowed_methods,
+            status_forcelist=self.status_forcelist,
+            backoff_factor=self.backoff_factor,  # type: ignore
+            backoff_max=self.backoff_max,  # type: ignore
+            raise_on_redirect=self.raise_on_redirect,
+            raise_on_status=self.raise_on_status,
+            history=self.history,
+            remove_headers_on_redirect=self.remove_headers_on_redirect,
+            respect_retry_after_header=self.respect_retry_after_header,
+            backoff_jitter=self.backoff_jitter,  # type: ignore
+        )
+
+        # Update urllib3's current state to reflect the incremented counters
+        urllib3_init_params.update(**urllib3_incremented_counters)
+
+        # Include urllib3's current state in our __init__ params
+        databricks_init_params["urllib3_kwargs"].update(**urllib3_init_params)  # type: ignore
+
+        return type(self).__private_init__(
+            retry_start_time=self._retry_start_time,
+            command_type=self.command_type,
+            **databricks_init_params,
+        )
+
+    @property
+    def command_type(self) -> Optional[CommandType]:
+        return self._command_type
+
+    @command_type.setter
+    def command_type(self, value: CommandType) -> None:
+        self._command_type = value
+
+    @property
+    def delay_default(self) -> float:
+        """Time in seconds the connector will wait between requests polling a GetOperationStatus Request
+
+        This property is never read by urllib3 for the purpose of retries. It's stored in this class
+        to keep all retry logic in one place.
+
+        This property is only set by __init__ and cannot be modified afterward.
+        """
+        return self._delay_default
+
+    def start_retry_timer(self) -> None:
+        """Timer is used to monitor the overall time across successive requests
+
+        Should only be called by ThriftBackend before sending a Thrift command"""
+        self._retry_start_time = time.time()
+
+    def check_timer_duration(self) -> float:
+        """Return time in seconds since the timer was started"""
+
+        if self._retry_start_time is None:
+            raise OperationalError(
+                "Cannot check retry timer. Timer was not started for this request."
+            )
+        else:
+            return time.time() - self._retry_start_time
+
+    def check_proposed_wait(self, proposed_wait: Union[int, float]) -> None:
+        """Raise an exception if the proposed wait would exceed the configured max_attempts_duration"""
+
+        proposed_overall_time = self.check_timer_duration() + proposed_wait
+        if proposed_overall_time > self.stop_after_attempts_duration:
+            raise MaxRetryDurationError(
+                f"Retry request would exceed Retry policy max retry duration of {self.stop_after_attempts_duration} seconds"
+            )
+
+    def sleep_for_retry(self, response: BaseHTTPResponse) -> bool:  # type: ignore
+        """Sleeps for the duration specified in the response Retry-After header, if present
+
+        A MaxRetryDurationError will be raised if doing so would exceed self.max_attempts_duration
+
+        This method is only called by urllib3 internals.
+        """
+        retry_after = self.get_retry_after(response)
+        if retry_after:
+            self.check_proposed_wait(retry_after)
+            time.sleep(retry_after)
+            return True
+
+        return False
+
+    def get_backoff_time(self) -> float:
+        """Calls urllib3's built-in get_backoff_time.
+
+        Never returns a value larger than self.delay_max
+        A MaxRetryDurationError will be raised if the calculated backoff would exceed self.max_attempts_duration
+
+        Note: within urllib3, a backoff is only calculated in cases where a Retry-After header is not present
+            in the previous unsuccessful request and `self.respect_retry_after_header` is True (which is always true)
+        """
+
+        proposed_backoff = super().get_backoff_time()
+        proposed_backoff = min(proposed_backoff, self.delay_max)
+        self.check_proposed_wait(proposed_backoff)
+
+        return proposed_backoff
+
+    def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
+        """This method encapsulates the connector's approach to retries.
+
+        We always retry a request unless one of these conditions is met:
+
+            1. The request received a 200 (Success) status code
+               Because the request succeeded .
+            2. The request received a 501 (Not Implemented) status code
+               Because this request can never succeed.
+            3. The request received a 404 (Not Found) code and the request CommandType
+               was CloseSession or CloseOperation. This code indicates that the session
+               or cursor was already closed. Further retries will always return the same
+               code.
+            4. The request CommandType was ExecuteStatement and the HTTP code does not
+               appear in the default status_forcelist or force_dangerous_codes list. By
+               default, this means ExecuteStatement is only retried for codes 429 and 503.
+               This limit prevents automatically retrying non-idempotent commands that could
+               be destructive.
+
+
+        Q: What about OSErrors and Redirects?
+        A: urllib3 automatically retries in both scenarios
+
+        Returns True if the request should be retried. Returns False or raises an exception
+        if a retry would violate the configured policy.
+        """
+
+        # Request succeeded. Don't retry.
+        if status_code == 200:
+            return False, "200 codes are not retried"
+
+        # Request failed and server said NotImplemented. This isn't recoverable. Don't retry.
+        if status_code == 501:
+            raise NonRecoverableNetworkError("Received code 501 from server.")
+
+        # Request failed and this method is not retryable. We only retry POST requests.
+        if not self._is_method_retryable(method):  # type: ignore
+            return False, "Only POST requests are retried"
+
+        # Request failed with 404 because CloseSession returns 404 if you repeat the request.
+        if (
+            status_code == 404
+            and self.command_type == CommandType.CLOSE_SESSION
+            and len(self.history) > 0
+        ):
+            raise SessionAlreadyClosedError(
+                "CloseSession received 404 code from Databricks. Session is already closed."
+            )
+
+        # Request failed with 404 because CloseOperation returns 404 if you repeat the request.
+        if (
+            status_code == 404
+            and self.command_type == CommandType.CLOSE_OPERATION
+            and len(self.history) > 0
+        ):
+            raise CursorAlreadyClosedError(
+                "CloseOperation received 404 code from Databricks. Cursor is already closed."
+            )
+
+        # Request failed, was an ExecuteStatement and the command may have reached the server
+        if (
+            self.command_type == CommandType.EXECUTE_STATEMENT
+            and status_code not in self.status_forcelist
+            and status_code not in self.force_dangerous_codes
+        ):
+            raise UnsafeToRetryError(
+                "ExecuteStatement command can only be retried for codes 429 and 503"
+            )
+
+        # Request failed with a dangerous code, was an ExecuteStatement, but user forced retries for this
+        # dangerous code. Note that these lines _are not required_ to make these requests retry. They would
+        # retry automatically. This code is included only so that we can log the exact reason for the retry.
+        # This gives users signal that their _retry_dangerous_codes setting actually did something.
+        if (
+            self.command_type == CommandType.EXECUTE_STATEMENT
+            and status_code in self.force_dangerous_codes
+        ):
+            return (
+                True,
+                f"Request failed with dangerous code {status_code} that is one of the configured _retry_dangerous_codes.",
+            )
+
+        # None of the above conditions applied. Eagerly retry.
+        logger.debug(
+            f"This request should be retried: {self.command_type and self.command_type.value}"
+        )
+        return (
+            True,
+            "Failed requests are retried by default per configured DatabricksRetryPolicy",
+        )
+
+    def is_retry(
+        self, method: str, status_code: int, has_retry_after: bool = False
+    ) -> bool:
+        """
+        Called by urllib3 when determining whether or not to retry
+
+        Logs a debug message if the request will be retried
+        """
+
+        should_retry, msg = self.should_retry(method, status_code)
+
+        if should_retry:
+            logger.debug(msg)
+
+        return should_retry

{databricks_sql_connector-2.8.0 → databricks_sql_connector-2.9.1.dev1}/src/databricks/sql/auth/thrift_http_client.py

@@ -15,6 +15,8 @@ from io import BytesIO
 
 from urllib3 import HTTPConnectionPool, HTTPSConnectionPool, ProxyManager
 
+from databricks.sql.auth.retry import CommandType, DatabricksRetryPolicy
+
 
 class THttpClient(thrift.transport.THttpClient.THttpClient):
     def __init__(
@@ -28,6 +30,7 @@ class THttpClient(thrift.transport.THttpClient.THttpClient):
         key_file=None,
         ssl_context=None,
         max_connections: int = 1,
+        retry_policy: Union[DatabricksRetryPolicy, int] = 0,
     ):
         if port is not None:
             warnings.warn(
@@ -81,6 +84,10 @@ class THttpClient(thrift.transport.THttpClient.THttpClient):
 
         self.max_connections = max_connections
 
+        # If retry_policy == 0 then urllib3 will not retry automatically
+        # this falls back to the pre-v3 behaviour where thrift_backend.py handles retry logic
+        self.retry_policy = retry_policy
+
         self.__wbuf = BytesIO()
         self.__resp: Union[None, HTTPResponse] = None
         self.__timeout = None
@@ -92,6 +99,13 @@ class THttpClient(thrift.transport.THttpClient.THttpClient):
         self._headers = headers
         super().setCustomHeaders(headers)
 
+    def startRetryTimer(self):
+        """Notify DatabricksRetryPolicy of the request start time
+
+        This is used to enforce the retry_stop_after_attempts_duration
+        """
+        self.retry_policy and self.retry_policy.start_retry_timer()
+
     def open(self):
 
         # self.__pool replaces the self.__http used by the original THttpClient
@@ -167,6 +181,7 @@ class THttpClient(thrift.transport.THttpClient.THttpClient):
             headers=headers,
             preload_content=False,
             timeout=self.__timeout,
+            retries=self.retry_policy,
         )
 
         # Get reply to flush the request
@@ -188,3 +203,12 @@ class THttpClient(thrift.transport.THttpClient.THttpClient):
         )
         cr = base64.b64encode(ap.encode()).strip()
         return "Basic " + six.ensure_str(cr)
+
+    def set_retry_command_type(self, value: CommandType):
+        """Pass the provided CommandType to the retry policy"""
+        if isinstance(self.retry_policy, DatabricksRetryPolicy):
+            self.retry_policy.command_type = value
+        else:
+            logger.warning(
+                "DatabricksRetryPolicy is currently bypassed. The CommandType cannot be set."
+            )

{databricks_sql_connector-2.8.0 → databricks_sql_connector-2.9.1.dev1}/src/databricks/sql/client.py

@@ -8,7 +8,11 @@ import os
 
 from databricks.sql import __version__
 from databricks.sql import *
-from databricks.sql.exc import OperationalError
+from databricks.sql.exc import (
+    OperationalError,
+    SessionAlreadyClosedError,
+    CursorAlreadyClosedError,
+)
 from databricks.sql.thrift_backend import ThriftBackend
 from databricks.sql.utils import ExecuteResponse, ParamEscaper, inject_parameters
 from databricks.sql.types import Row
@@ -257,6 +261,9 @@ class Connection:
 
         try:
             self.thrift_backend.close_session(self._session_handle)
+        except RequestError as e:
+            if isinstance(e.args[1], SessionAlreadyClosedError):
+                logger.info("Session was closed by a prior request")
         except DatabaseError as e:
             if "Invalid SessionHandle" in str(e):
                 logger.warning(
@@ -958,6 +965,9 @@ class ResultSet:
                 and self.connection.open
             ):
                 self.thrift_backend.close_command(self.command_id)
+        except RequestError as e:
+            if isinstance(e.args[1], CursorAlreadyClosedError):
+                logger.info("Operation was canceled by a prior request")
         finally:
             self.has_been_closed_server_side = True
             self.op_state = self.thrift_backend.CLOSED_OP_STATE

{databricks_sql_connector-2.8.0 → databricks_sql_connector-2.9.1.dev1}/src/databricks/sql/exc.py

@@ -93,3 +93,25 @@ class RequestError(OperationalError):
     """
 
     pass
+
+
+class MaxRetryDurationError(RequestError):
+    """Thrown if the next HTTP request retry would exceed the configured
+    stop_after_attempts_duration
+    """
+
+
+class NonRecoverableNetworkError(RequestError):
+    """Thrown if an HTTP code 501 is received"""
+
+
+class UnsafeToRetryError(RequestError):
+    """Thrown if ExecuteStatement request receives a code other than 200, 429, or 503"""
+
+
+class SessionAlreadyClosedError(RequestError):
+    """Thrown if CloseSession receives a code 404. ThriftBackend should gracefully proceed as this is expected."""
+
+
+class CursorAlreadyClosedError(RequestError):
+    """Thrown if CancelOperation receives a code 404. ThriftBackend should gracefully proceed as this is expected."""

{databricks_sql_connector-2.8.0 → databricks_sql_connector-2.9.1.dev1}/src/databricks/sql/thrift_backend.py

@@ -17,9 +17,11 @@ import thrift.transport.TTransport
 import urllib3.exceptions
 
 import databricks.sql.auth.thrift_http_client
+from databricks.sql.auth.thrift_http_client import CommandType
 from databricks.sql.auth.authenticators import AuthProvider
 from databricks.sql.thrift_api.TCLIService import TCLIService, ttypes
 from databricks.sql import *
+from databricks.sql.exc import MaxRetryDurationError
 from databricks.sql.thrift_api.TCLIService.TCLIService import (
     Client as TCLIServiceClient,
 )
@@ -70,6 +72,12 @@ class ThriftBackend:
     CLOSED_OP_STATE = ttypes.TOperationState.CLOSED_STATE
     ERROR_OP_STATE = ttypes.TOperationState.ERROR_STATE
 
+    _retry_delay_min: float
+    _retry_delay_max: float
+    _retry_stop_after_attempts_count: int
+    _retry_stop_after_attempts_duration: float
+    _retry_delay_default: float
+
     def __init__(
         self,
         server_hostname: str,
@@ -113,9 +121,15 @@ class ThriftBackend:
         #
         # _retry_stop_after_attempts_count
         #  The maximum number of times we should retry retryable requests (defaults to 24)
+        # _retry_dangerous_codes
+        #  An iterable of integer HTTP status codes. ExecuteStatement commands will be retried if these codes are received.
+        #  (defaults to [])
         # _socket_timeout
         #  The timeout in seconds for socket send, recv and connect operations. Should be a positive float or integer.
         #  (defaults to 900)
+        # _enable_v3_retries
+        # Whether to use the DatabricksRetryPolicy implemented in urllib3
+        # (defaults to False)
         # max_download_threads
         #  Number of threads for handling cloud fetch downloads. Defaults to 10
 
@@ -166,10 +180,28 @@ class ThriftBackend:
 
         self._auth_provider = auth_provider
 
+        # Connector version 3 retry approach
+        self.enable_v3_retries = kwargs.get("_enable_v3_retries", False)
+        self.force_dangerous_codes = kwargs.get("_retry_dangerous_codes", [])
+
+        additional_transport_args = {}
+        if self.enable_v3_retries:
+            self.retry_policy = databricks.sql.auth.thrift_http_client.DatabricksRetryPolicy(
+                delay_min=self._retry_delay_min,
+                delay_max=self._retry_delay_max,
+                stop_after_attempts_count=self._retry_stop_after_attempts_count,
+                stop_after_attempts_duration=self._retry_stop_after_attempts_duration,
+                delay_default=self._retry_delay_default,
+                force_dangerous_codes=self.force_dangerous_codes,
+            )
+
+            additional_transport_args["retry_policy"] = self.retry_policy
+
         self._transport = databricks.sql.auth.thrift_http_client.THttpClient(
             auth_provider=self._auth_provider,
             uri_or_host=uri,
             ssl_context=ssl_context,
+            **additional_transport_args,  # type: ignore
         )
 
         timeout = kwargs.get("_socket_timeout", DEFAULT_SOCKET_TIMEOUT)
@@ -188,6 +220,7 @@ class ThriftBackend:
 
         self._request_lock = threading.RLock()
 
+    # TODO: Move this bounding logic into DatabricksRetryPolicy for v3 (PECO-918)
     def _initialize_retry_args(self, kwargs):
         # Configure retries & timing: use user-settings or defaults, and bound
         # by policy. Log.warn when given param gets restricted.
@@ -335,15 +368,17 @@ class ThriftBackend:
 
             error, error_message, retry_delay = None, None, None
             try:
-                # The MagicMocks in our unit tests have a `name` property instead of `__name__`.
-                logger.debug(
-                    "Sending request: {}(<REDACTED>)".format(
-                        getattr(
-                            method, "__name__", getattr(method, "name", "UnknownMethod")
-                        )
-                    )
-                )
+
+                this_method_name = getattr(method, "__name__")
+
+                logger.debug("Sending request: {}(<REDACTED>)".format(this_method_name))
                 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()
+
                 response = method(request)
 
                 # Calling `close()` here releases the active HTTP connection back to the pool
@@ -359,9 +394,16 @@ class ThriftBackend:
             except urllib3.exceptions.HTTPError as err:
                 # retry on timeout. Happens a lot in Azure and it is safe as data has not been sent to server yet
 
+                # TODO: don't use exception handling for GOS polling...
+
                 gos_name = TCLIServiceClient.GetOperationStatus.__name__
                 if method.__name__ == gos_name:
-                    retry_delay = bound_retry_delay(attempt, self._retry_delay_default)
+                    delay_default = (
+                        self.enable_v3_retries
+                        and self.retry_policy.delay_default
+                        or self._retry_delay_default
+                    )
+                    retry_delay = bound_retry_delay(attempt, delay_default)
                     logger.info(
                         f"GetOperationStatus failed with HTTP error and will be retried: {str(err)}"
                     )

{databricks_sql_connector-2.8.0 → databricks_sql_connector-2.9.1.dev1}/src/databricks/sqlalchemy/dialect/__init__.py

@@ -74,7 +74,7 @@ class DatabricksDialect(default.DefaultDialect):
     driver: str = "databricks-sql-python"
     default_schema_name: str = "default"
 
-    preparer = DatabricksIdentifierPreparer
+    preparer = DatabricksIdentifierPreparer  # type: ignore
     type_compiler = DatabricksTypeCompiler
     ddl_compiler = DatabricksDDLCompiler
     supports_statement_cache: bool = True