preliminary reetries

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/auth/thrift_http_client.py

@@ -352,13 +352,21 @@ def _check_rest_response_for_error(
         """
         if status_code >= 400:
             error_message = f"REST HTTP request failed with status {status_code}"
+            error_code = None
 
             # Try to extract error details from JSON response
             if response_data:
                 try:
                     error_details = json.loads(response_data.decode("utf-8"))
-                    if isinstance(error_details, dict) and "message" in error_details:
-                        error_message = f"{error_message}: {error_details['message']}"
+                    if isinstance(error_details, dict):
+                        if "message" in error_details:
+                            error_message = (
+                                f"{error_message}: {error_details['message']}"
+                            )
+                        if "error_code" in error_details:
+                            error_code = error_details["error_code"]
+                        elif "errorCode" in error_details:
+                            error_code = error_details["errorCode"]
                     logger.error(
                         f"Request failed (status {status_code}): {error_details}"
                     )
@@ -369,6 +377,69 @@ def _check_rest_response_for_error(
             else:
                 logger.error(f"Request failed (status {status_code}): No response data")
 
-            from databricks.sql.exc import RequestError
+            from databricks.sql.exc import (
+                RequestError,
+                OperationalError,
+                DatabaseError,
+                SessionAlreadyClosedError,
+                CursorAlreadyClosedError,
+            )
 
-            raise RequestError(error_message)
+            # Map status codes to appropriate exceptions to match Thrift behavior
+            if status_code == 429:
+                # Rate limiting errors
+                retry_after = None
+                if self.headers and "Retry-After" in self.headers:
+                    retry_after = self.headers["Retry-After"]
+
+                rate_limit_msg = f"Maximum rate has been exceeded. Please reduce the rate of requests and try again"
+                if retry_after:
+                    rate_limit_msg += f" after {retry_after} seconds."
+                raise RequestError(rate_limit_msg)
+
+            elif status_code == 503:
+                # Service unavailable errors
+                raise OperationalError(
+                    "TEMPORARILY_UNAVAILABLE: Service temporarily unavailable"
+                )
+
+            elif status_code == 404:
+                # Not found errors - could be session or operation already closed
+                if error_message and "session" in error_message.lower():
+                    raise SessionAlreadyClosedError(
+                        "Session was closed by a prior request"
+                    )
+                elif error_message and (
+                    "operation" in error_message.lower()
+                    or "statement" in error_message.lower()
+                ):
+                    raise CursorAlreadyClosedError(
+                        "Operation was canceled by a prior request"
+                    )
+                else:
+                    raise RequestError(error_message)
+
+            elif status_code == 401:
+                # Authentication errors
+                raise OperationalError(
+                    "Authentication failed. Please check your credentials."
+                )
+
+            elif status_code == 403:
+                # Permission errors
+                raise OperationalError(
+                    "Permission denied. You do not have access to this resource."
+                )
+
+            elif status_code == 400:
+                # Bad request errors - often syntax errors
+                if error_message and "syntax" in error_message.lower():
+                    raise DatabaseError(
+                        f"Syntax error in SQL statement: {error_message}"
+                    )
+                else:
+                    raise RequestError(error_message)
+
+            else:
+                # Generic errors
+                raise RequestError(error_message)

src/databricks/sql/backend/sea/backend.py

@@ -130,14 +130,40 @@ def __init__(
         # Extract warehouse ID from http_path
         self.warehouse_id = self._extract_warehouse_id(http_path)
 
-        # Initialize ThriftHttpClient
+        # Extract retry policy parameters
+        retry_policy = kwargs.get("_retry_policy", None)
+        retry_stop_after_attempts_count = kwargs.get(
+            "_retry_stop_after_attempts_count", 30
+        )
+        retry_stop_after_attempts_duration = kwargs.get(
+            "_retry_stop_after_attempts_duration", 600
+        )
+        retry_delay_min = kwargs.get("_retry_delay_min", 1)
+        retry_delay_max = kwargs.get("_retry_delay_max", 60)
+        retry_delay_default = kwargs.get("_retry_delay_default", 5)
+        retry_dangerous_codes = kwargs.get("_retry_dangerous_codes", [])
+
+        # Create retry policy if not provided
+        if not retry_policy:
+            from databricks.sql.auth.retry import DatabricksRetryPolicy
+
+            retry_policy = DatabricksRetryPolicy(
+                delay_min=retry_delay_min,
+                delay_max=retry_delay_max,
+                stop_after_attempts_count=retry_stop_after_attempts_count,
+                stop_after_attempts_duration=retry_stop_after_attempts_duration,
+                delay_default=retry_delay_default,
+                force_dangerous_codes=retry_dangerous_codes,
+            )
+
+        # Initialize ThriftHttpClient with retry policy
         thrift_client = THttpClient(
             auth_provider=auth_provider,
             uri_or_host=f"https://{server_hostname}:{port}",
             path=http_path,
             ssl_options=ssl_options,
             max_connections=kwargs.get("max_connections", 1),
-            retry_policy=kwargs.get("_retry_stop_after_attempts_count", 30),
+            retry_policy=retry_policy,
         )
 
         # Set custom headers
@@ -475,48 +501,99 @@ def execute_command(
             result_compression=result_compression,
         )
 
-        response_data = self.http_client.post(
-            path=self.STATEMENT_PATH, data=request.to_dict()
-        )
-        response = ExecuteStatementResponse.from_dict(response_data)
-        statement_id = response.statement_id
-        if not statement_id:
-            raise ServerOperationError(
-                "Failed to execute command: No statement ID returned",
-                {
-                    "operation-id": None,
-                    "diagnostic-info": None,
-                },
+        try:
+            response_data = self.http_client.post(
+                path=self.STATEMENT_PATH, data=request.to_dict()
             )
+            response = ExecuteStatementResponse.from_dict(response_data)
+            statement_id = response.statement_id
+
+            if not statement_id:
+                raise ServerOperationError(
+                    "Failed to execute command: No statement ID returned",
+                    {
+                        "operation-id": None,
+                        "diagnostic-info": None,
+                    },
+                )
 
-        command_id = CommandId.from_sea_statement_id(statement_id)
+            command_id = CommandId.from_sea_statement_id(statement_id)
 
-        # Store the command ID in the cursor
-        cursor.active_command_id = command_id
+            # Store the command ID in the cursor
+            cursor.active_command_id = command_id
 
-        # If async operation, return and let the client poll for results
-        if async_op:
-            return None
+            # If async operation, return and let the client poll for results
+            if async_op:
+                return None
 
-        # For synchronous operation, wait for the statement to complete
-        status = response.status
-        state = status.state
+            # For synchronous operation, wait for the statement to complete
+            status = response.status
+            state = status.state
 
-        # Keep polling until we reach a terminal state
-        while state in [CommandState.PENDING, CommandState.RUNNING]:
-            time.sleep(0.5)  # add a small delay to avoid excessive API calls
-            state = self.get_query_state(command_id)
+            # Keep polling until we reach a terminal state
+            while state in [CommandState.PENDING, CommandState.RUNNING]:
+                time.sleep(0.5)  # add a small delay to avoid excessive API calls
+                state = self.get_query_state(command_id)
 
-        if state != CommandState.SUCCEEDED:
-            raise ServerOperationError(
-                f"Statement execution did not succeed: {status.error.message if status.error else 'Unknown error'}",
-                {
-                    "operation-id": command_id.to_sea_statement_id(),
-                    "diagnostic-info": None,
-                },
-            )
+            if state != CommandState.SUCCEEDED:
+                error_message = (
+                    status.error.message if status.error else "Unknown error"
+                )
+                error_code = status.error.error_code if status.error else None
+
+                # Map error codes to appropriate exceptions to match Thrift behavior
+                from databricks.sql.exc import (
+                    DatabaseError,
+                    ProgrammingError,
+                    OperationalError,
+                )
+
+                if (
+                    error_code == "SYNTAX_ERROR"
+                    or "syntax error" in error_message.lower()
+                ):
+                    raise DatabaseError(
+                        f"Syntax error in SQL statement: {error_message}"
+                    )
+                elif error_code == "TEMPORARILY_UNAVAILABLE":
+                    raise OperationalError(
+                        f"Service temporarily unavailable: {error_message}"
+                    )
+                elif error_code == "PERMISSION_DENIED":
+                    raise OperationalError(f"Permission denied: {error_message}")
+                else:
+                    raise ServerOperationError(
+                        f"Statement execution failed: {error_message}",
+                        {
+                            "operation-id": command_id.to_sea_statement_id(),
+                            "diagnostic-info": None,
+                        },
+                    )
 
-        return self.get_execution_result(command_id, cursor)
+            return self.get_execution_result(command_id, cursor)
+
+        except Exception as e:
+            # Map exceptions to match Thrift behavior
+            from databricks.sql.exc import DatabaseError, OperationalError, RequestError
+
+            if isinstance(e, (DatabaseError, OperationalError, RequestError)):
+                # Pass through these exceptions as they're already properly typed
+                raise
+            elif "syntax error" in str(e).lower():
+                # Syntax errors
+                raise DatabaseError(f"Syntax error in SQL statement: {str(e)}")
+            elif "permission denied" in str(e).lower():
+                # Permission errors
+                raise OperationalError(f"Permission denied: {str(e)}")
+            elif "database" in str(e).lower() and "not found" in str(e).lower():
+                # Database not found errors
+                raise DatabaseError(f"Database not found: {str(e)}")
+            elif "table" in str(e).lower() and "not found" in str(e).lower():
+                # Table not found errors
+                raise DatabaseError(f"Table not found: {str(e)}")
+            else:
+                # Generic operational errors
+                raise OperationalError(f"Error executing command: {str(e)}")
 
     def cancel_command(self, command_id: CommandId) -> None:
         """

src/databricks/sql/backend/sea/utils/http_client_adapter.py

@@ -9,6 +9,7 @@
 from typing import Dict, Optional, Any
 
 from databricks.sql.auth.thrift_http_client import THttpClient
+from databricks.sql.auth.retry import CommandType
 
 logger = logging.getLogger(__name__)
 
@@ -36,14 +37,49 @@ def __init__(
         """
         self.thrift_client = thrift_client
 
+    def _determine_command_type(
+        self, path: str, method: str, data: Optional[Dict[str, Any]] = None
+    ) -> CommandType:
+        """
+        Determine the CommandType based on the request path and method.
+
+        Args:
+            path: API endpoint path
+            method: HTTP method (GET, POST, DELETE)
+            data: Request payload data
+
+        Returns:
+            CommandType: The appropriate CommandType enum value
+        """
+        # Extract the base path component (e.g., "sessions", "statements")
+        path_parts = path.strip("/").split("/")
+        base_path = path_parts[0] if path_parts else ""
+
+        # Check for specific operations based on path and method
+        if "statements" in path:
+            if method == "POST" and any(part == "cancel" for part in path_parts):
+                return CommandType.CLOSE_OPERATION
+            elif method == "POST" and not any(part == "cancel" for part in path_parts):
+                return CommandType.EXECUTE_STATEMENT
+            elif method == "GET":
+                return CommandType.GET_OPERATION_STATUS
+            elif method == "DELETE":
+                return CommandType.CLOSE_OPERATION
+        elif "sessions" in path:
+            if method == "DELETE":
+                return CommandType.CLOSE_SESSION
+
+        # Default for any other operations
+        return CommandType.OTHER
+
     def get(
         self,
         path: str,
         params: Optional[Dict[str, Any]] = None,
         headers: Optional[Dict[str, str]] = None,
     ) -> Dict[str, Any]:
         """
-        Convenience method for GET requests.
+        Convenience method for GET requests with retry support.
 
         Args:
             path: API endpoint path
@@ -53,6 +89,10 @@ def get(
         Returns:
             Response data parsed from JSON
         """
+        command_type = self._determine_command_type(path, "GET")
+        self.thrift_client.set_retry_command_type(command_type)
+        self.thrift_client.startRetryTimer()
+
         return self.thrift_client.make_rest_request(
             "GET", path, params=params, headers=headers
         )
@@ -65,7 +105,7 @@ def post(
         headers: Optional[Dict[str, str]] = None,
     ) -> Dict[str, Any]:
         """
-        Convenience method for POST requests.
+        Convenience method for POST requests with retry support.
 
         Args:
             path: API endpoint path
@@ -76,6 +116,10 @@ def post(
         Returns:
             Response data parsed from JSON
         """
+        command_type = self._determine_command_type(path, "POST", data)
+        self.thrift_client.set_retry_command_type(command_type)
+        self.thrift_client.startRetryTimer()
+
         return self.thrift_client.make_rest_request(
             "POST", path, data=data, params=params, headers=headers
         )
@@ -88,7 +132,7 @@ def delete(
         headers: Optional[Dict[str, str]] = None,
     ) -> Dict[str, Any]:
         """
-        Convenience method for DELETE requests.
+        Convenience method for DELETE requests with retry support.
 
         Args:
             path: API endpoint path
@@ -99,6 +143,10 @@ def delete(
         Returns:
             Response data parsed from JSON
         """
+        command_type = self._determine_command_type(path, "DELETE", data)
+        self.thrift_client.set_retry_command_type(command_type)
+        self.thrift_client.startRetryTimer()
+
         return self.thrift_client.make_rest_request(
             "DELETE", path, data=data, params=params, headers=headers
         )