check types fix

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

src/databricks/sql/backend/types.py

@@ -282,7 +282,6 @@ def __init__(
         operation_type: Optional[int] = None,
         has_result_set: bool = False,
         modified_row_count: Optional[int] = None,
-        statement_type: Optional[StatementType] = None,
     ):
         """
         Initialize a CommandId.
@@ -302,7 +301,7 @@ def __init__(
         self.operation_type = operation_type
         self.has_result_set = has_result_set
         self.modified_row_count = modified_row_count
-        self._statement_type = statement_type
+        self._statement_type = StatementType.NONE
 
     def __str__(self) -> str:
         """
@@ -421,7 +420,7 @@ def set_statement_type(self, statement_type: StatementType):
         self._statement_type = statement_type
 
     @property
-    def statement_type(self) -> Optional[StatementType]:
+    def statement_type(self) -> StatementType:
         """
         Get the statement type for this command.
         """

tests/unit/test_client.py

@@ -189,6 +189,7 @@ def test_arraysize_buffer_size_passthrough(
     def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
         mock_connection = Mock()
         mock_backend = Mock()
+        mock_results = Mock()
         mock_backend.fetch_results.return_value = (Mock(), False, 0)
 
         result_set = ThriftResultSet(
@@ -223,7 +224,9 @@ def test_closing_result_set_hard_closes_commands(self):
 
         mock_thrift_backend.fetch_results.return_value = (Mock(), False, 0)
         result_set = ThriftResultSet(
-            mock_connection, mock_results_response, mock_thrift_backend, session_id_hex=Mock())
+            mock_connection, mock_results_response, mock_thrift_backend, session_id_hex=Mock()
+        )
+        result_set.results = mock_results
 
         result_set.close()