add minimal retry func

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

Author: varun-edachali-dbx

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

@@ -130,14 +130,35 @@ 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
@@ -394,7 +415,7 @@ def _results_message_to_execute_response(self, sea_response, command_id):
             description=description,
             has_been_closed_server_side=False,
             lz4_compressed=lz4_compressed,
-            is_staging_operation=False,
+            is_staging_operation=manifest_obj.is_volume_operation,
             arrow_schema_bytes=None,
             result_format=manifest_obj.format,
         )
@@ -475,48 +496,56 @@ 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:
+                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,
+                    },
+                )
 
-        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 RequestError, OperationalError
+            if isinstance(e, (ServerOperationError, RequestError)):
+                raise
+            else:
+                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,50 @@ 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[-1] if path_parts else ""
+        
+        # Check for specific operations based on path and method
+        if "statements" in path:
+            if method == "POST" and "cancel" in path:
+                return CommandType.CLOSE_OPERATION
+            elif method == "POST" and "cancel" not in path:
+                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 == "POST":
+                # Creating a new session
+                return CommandType.OTHER
+            elif 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 +90,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 +106,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 +117,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 +133,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 +144,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
         )