Merge branch 'sea-migration' into sea-optimise-success

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -82,6 +82,7 @@ def execute_command(
         parameters: List[ttypes.TSparkParameter],
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
+        row_limit: Optional[int] = None,
     ) -> Union[ResultSet, None]:
         """
         Executes a SQL command or query within the specified session.
@@ -100,6 +101,7 @@ def execute_command(
             parameters: List of parameters to bind to the query
             async_op: Whether to execute the command asynchronously
             enforce_embedded_schema_correctness: Whether to enforce schema correctness
+            row_limit: Maximum number of rows in the response.
 
         Returns:
             If async_op is False, returns a ResultSet object containing the

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

@@ -14,6 +14,7 @@
     WaitTimeout,
     MetadataCommands,
 )
+from databricks.sql.thrift_api.TCLIService import ttypes
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
@@ -424,7 +425,7 @@ def execute_command(
         lz4_compression: bool,
         cursor: Cursor,
         use_cloud_fetch: bool,
-        parameters: List[Dict[str, Any]],
+        parameters: List[ttypes.TSparkParameter],
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
         row_limit: Optional[int] = None,
@@ -459,9 +460,11 @@ def execute_command(
             for param in parameters:
                 sea_parameters.append(
                     StatementParameter(
-                        name=param["name"],
-                        value=param["value"],
-                        type=param["type"] if "type" in param else None,
+                        name=param.name,
+                        value=(
+                            param.value.stringValue if param.value is not None else None
+                        ),
+                        type=param.type,
                     )
                 )
 

src/databricks/sql/backend/thrift_backend.py

@@ -239,10 +239,6 @@ def __init__(
     def max_download_threads(self) -> int:
         return self._max_download_threads
 
-    @property
-    def max_download_threads(self) -> int:
-        return self._max_download_threads
-
     # 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

tests/unit/test_sea_backend.py

@@ -13,6 +13,8 @@
     _filter_session_configuration,
 )
 from databricks.sql.backend.types import SessionId, CommandId, CommandState, BackendType
+from databricks.sql.parameters.native import IntegerParameter, TDbsqlParameter
+from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import SSLOptions
 from databricks.sql.auth.authenticators import AuthProvider
 from databricks.sql.exc import (
@@ -352,7 +354,8 @@ def test_command_execution_advanced(
             "status": {"state": "SUCCEEDED"},
         }
         mock_http_client._make_request.return_value = execute_response
-        param = {"name": "param1", "value": "value1", "type": "STRING"}
+        dbsql_param = IntegerParameter(name="param1", value=1)
+        param = dbsql_param.as_tspark_param(named=True)
 
         with patch.object(sea_client, "_response_to_result_set"):
             sea_client.execute_command(
@@ -371,8 +374,8 @@ def test_command_execution_advanced(
             assert "parameters" in kwargs["data"]
             assert len(kwargs["data"]["parameters"]) == 1
             assert kwargs["data"]["parameters"][0]["name"] == "param1"
-            assert kwargs["data"]["parameters"][0]["value"] == "value1"
-            assert kwargs["data"]["parameters"][0]["type"] == "STRING"
+            assert kwargs["data"]["parameters"][0]["value"] == "1"
+            assert kwargs["data"]["parameters"][0]["type"] == "INT"
 
         # Test execution failure
         mock_http_client.reset_mock()