diff --git a/src/databricks/sql/backend/sea/backend.py b/src/databricks/sql/backend/sea/backend.py index 76903ccd..ff0a259c 100644 --- a/src/databricks/sql/backend/sea/backend.py +++ b/src/databricks/sql/backend/sea/backend.py @@ -11,6 +11,7 @@ ResultCompression, WaitTimeout, ) +from databricks.sql.thrift_api.TCLIService import ttypes if TYPE_CHECKING: from databricks.sql.client import Cursor @@ -402,7 +403,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, ) -> Union["ResultSet", None]: @@ -436,9 +437,9 @@ 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, + type=param.type, ) ) diff --git a/tests/unit/test_sea_backend.py b/tests/unit/test_sea_backend.py index f30c92ed..a2c1aeb0 100644 --- a/tests/unit/test_sea_backend.py +++ b/tests/unit/test_sea_backend.py @@ -13,6 +13,7 @@ _filter_session_configuration, ) from databricks.sql.backend.types import SessionId, CommandId, CommandState, BackendType +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 ( @@ -354,7 +355,7 @@ def test_command_execution_advanced( "status": {"state": "SUCCEEDED"}, } mock_http_client._make_request.return_value = execute_response - param = {"name": "param1", "value": "value1", "type": "STRING"} + param = ttypes.TSparkParameter(name="param1", value="value1", type="STRING") with patch.object(sea_client, "get_execution_result"): sea_client.execute_command(