fix typing

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

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -584,7 +584,7 @@ def open_session(self, session_configuration, catalog, schema) -> SessionId:
             self._check_initial_namespace(catalog, schema, response)
             self._check_protocol_version(response)
             if response.sessionHandle is None:
-                return None
+                raise DatabaseError("Server returned null session handle")
             info = (
                 {"serverProtocolVersion": response.serverProtocolVersion}
                 if response.serverProtocolVersion
@@ -875,66 +875,6 @@ def get_execution_result(self, command_id: CommandId, cursor):
             arrow_schema_bytes=schema_bytes,
         )
 
-    def get_execution_result(self, op_handle, cursor):
-
-        assert op_handle is not None
-
-        req = ttypes.TFetchResultsReq(
-            operationHandle=ttypes.TOperationHandle(
-                op_handle.operationId,
-                op_handle.operationType,
-                False,
-                op_handle.modifiedRowCount,
-            ),
-            maxRows=cursor.arraysize,
-            maxBytes=cursor.buffer_size_bytes,
-            orientation=ttypes.TFetchOrientation.FETCH_NEXT,
-            includeResultSetMetadata=True,
-        )
-
-        resp = self.make_request(self._client.FetchResults, req)
-
-        t_result_set_metadata_resp = resp.resultSetMetadata
-
-        lz4_compressed = t_result_set_metadata_resp.lz4Compressed
-        is_staging_operation = t_result_set_metadata_resp.isStagingOperation
-        has_more_rows = resp.hasMoreRows
-        description = self._hive_schema_to_description(
-            t_result_set_metadata_resp.schema
-        )
-
-        if pyarrow:
-            schema_bytes = (
-                t_result_set_metadata_resp.arrowSchema
-                or self._hive_schema_to_arrow_schema(t_result_set_metadata_resp.schema)
-                .serialize()
-                .to_pybytes()
-            )
-        else:
-            schema_bytes = None
-
-        queue = ResultSetQueueFactory.build_queue(
-            row_set_type=resp.resultSetMetadata.resultFormat,
-            t_row_set=resp.results,
-            arrow_schema_bytes=schema_bytes,
-            max_download_threads=self.max_download_threads,
-            lz4_compressed=lz4_compressed,
-            description=description,
-            ssl_options=self._ssl_options,
-        )
-
-        return ExecuteResponse(
-            arrow_queue=queue,
-            status=resp.status,
-            has_been_closed_server_side=False,
-            has_more_rows=has_more_rows,
-            lz4_compressed=lz4_compressed,
-            is_staging_operation=is_staging_operation,
-            command_handle=op_handle,
-            description=description,
-            arrow_schema_bytes=schema_bytes,
-        )
-
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
         if initial_operation_status_resp:
             self._check_command_not_in_error_or_closed_state(
@@ -1165,10 +1105,6 @@ def _handle_execute_response_async(self, resp, cursor):
         cursor.active_command_id = command_id
         self._check_direct_results_for_error(resp.directResults)
 
-    def _handle_execute_response_async(self, resp, cursor):
-        cursor.active_op_handle = resp.operationHandle
-        self._check_direct_results_for_error(resp.directResults)
-
     def fetch_results(
         self,
         command_id: CommandId,

src/databricks/sql/backend/types.py

@@ -18,7 +18,7 @@ def guid_to_hex_id(guid: bytes) -> str:
     try:
         this_uuid = uuid.UUID(bytes=guid)
     except Exception as e:
-        logger.debug(f"Unable to convert bytes to UUID: {guid} -- {str(e)}")
+        logger.debug(f"Unable to convert bytes to UUID: {guid!r} -- {str(e)}")
         return str(guid)
     return str(this_uuid)
 

src/databricks/sql/client.py

@@ -855,6 +855,8 @@ def get_query_state(self) -> "TOperationState":
         :return:
         """
         self._check_not_closed()
+        if self.active_command_id is None:
+            raise Error("No active command to get state for")
         return self.backend.get_query_state(self.active_command_id)
 
     def is_query_pending(self):
@@ -870,19 +872,6 @@ def is_query_pending(self):
             ttypes.TOperationState.PENDING_STATE,
         ]
 
-    def is_query_pending(self):
-        """
-        Checks whether the async executing query is in pending state or not
-
-        :return:
-        """
-        operation_state = self.get_query_state()
-
-        return not operation_state or operation_state in [
-            ttypes.TOperationState.RUNNING_STATE,
-            ttypes.TOperationState.PENDING_STATE,
-        ]
-
     def get_async_execution_result(self):
         """
 

tests/unit/test_thrift_backend.py

@@ -52,6 +52,7 @@ class ThriftBackendTestSuite(unittest.TestCase):
     open_session_resp = ttypes.TOpenSessionResp(
         status=okay_status,
         serverProtocolVersion=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V4,
+        sessionHandle=session_handle,
     )
 
     metadata_resp = ttypes.TGetResultSetMetadataResp(
@@ -177,7 +178,9 @@ def test_okay_protocol_versions_succeed(self, tcli_service_client_cass):
 
         for protocol_version in good_protocol_versions:
             t_http_client_instance.OpenSession.return_value = ttypes.TOpenSessionResp(
-                status=self.okay_status, serverProtocolVersion=protocol_version
+                status=self.okay_status,
+                serverProtocolVersion=protocol_version,
+                sessionHandle=self.session_handle,
             )
 
             thrift_backend = self._make_fake_thrift_backend()
@@ -2061,6 +2064,7 @@ def _construct_open_session_with_namespace(self, can_use_multiple_cats, cat, sch
             serverProtocolVersion=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V4,
             canUseMultipleCatalogs=can_use_multiple_cats,
             initialNamespace=ttypes.TNamespace(catalogName=cat, schemaName=schem),
+            sessionHandle=self.session_handle,
         )
 
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
@@ -2160,6 +2164,7 @@ def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
             serverProtocolVersion=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V3,
             canUseMultipleCatalogs=True,
             initialNamespace=ttypes.TNamespace(catalogName="cat", schemaName="schem"),
+            sessionHandle=self.session_handle,
         )
 
         backend = ThriftDatabricksClient(