Put in some unit tests, will add e2e

Author: nithinkdb

src/databricks/sql/client.py

@@ -225,6 +225,21 @@ def __del__(self):
     def get_session_id(self):
         return self.thrift_backend.handle_to_id(self._session_handle)
 
+<<<<<<< HEAD
+=======
+    def get_session_protocol_version(self):
+        return self.thrift_backend.extract_protocol_version_from_handle(
+            self._session_handle
+        )
+    
+    @staticmethod
+    def server_parameterized_queries_enabled(protocolVersion):
+        if protocolVersion and protocolVersion >= ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V8:
+            return True
+        else:
+            return False
+
+>>>>>>> cd93869 (Put in some unit tests, will add e2e)
     def get_session_id_hex(self):
         return self.thrift_backend.handle_to_hex_id(self._session_handle)
 
@@ -501,6 +516,13 @@ def execute(
         """
         if parameters is None:
             parameters = []
+<<<<<<< HEAD
+=======
+        elif (not self.server_parameterized_queries_enabled(self.connection._session_handle)):
+            raise Error(
+                "Parameterized operations are not supported by this server. DBR 14.2 is required."
+            )
+>>>>>>> cd93869 (Put in some unit tests, will add e2e)
         else:
             parameters = named_parameters_to_tsparkparams(parameters)
 

src/databricks/sql/thrift_backend.py

@@ -506,7 +506,7 @@ def _check_protocol_version(self, t_open_session_resp):
                 "SPARK_CLI_SERVICE_PROTOCOL_V2, "
                 "instead got: {}".format(protocol_version)
             )
-
+    
     def _check_initial_namespace(self, catalog, schema, response):
         if not (catalog or schema):
             return
@@ -569,7 +569,7 @@ def open_session(self, session_configuration, catalog, schema):
             response = self.make_request(self._client.OpenSession, open_session_req)
             self._check_initial_namespace(catalog, schema, response)
             self._check_protocol_version(response)
-            return response.sessionHandle
+            return self.get_session_handle_from_resp(response)
         except:
             self._transport.close()
             raise
@@ -1014,6 +1014,13 @@ def cancel_command(self, active_op_handle):
         req = ttypes.TCancelOperationReq(active_op_handle)
         self.make_request(self._client.CancelOperation, req)
 
+    @staticmethod
+    def get_session_handle_from_resp(t_open_session_resp):
+        sessionHandle = t_open_session_resp.sessionHandle
+        if sessionHandle.serverProtocolVersion == None:
+            sessionHandle.serverProtocolVersion = t_open_session_resp.serverProtocolVersion
+        return sessionHandle
+    
     @staticmethod
     def handle_to_id(session_handle):
         return session_handle.sessionId.guid

tests/e2e/common/parameterized_query_tests.py

@@ -29,6 +29,12 @@ def _get_one_result(self, query: str, parameters: Union[Dict, List[Dict]]) -> Tu
             with conn.cursor() as cursor:
                 cursor.execute(query, parameters=parameters)
                 return cursor.fetchone()
+            
+    def _test_protocol_checking(self) -> Tuple:
+        with self.connection() as conn:
+            with conn.cursor() as cursor:
+                cursor.execute(query, parameters=parameters)
+                return cursor.fetchone()
 
     def _quantize(self, input: Union[float, int], place_value=2) -> Decimal:
 

tests/unit/test_parameters.py

@@ -9,14 +9,29 @@
 from databricks.sql.thrift_api.TCLIService.ttypes import (
     TSparkParameter,
     TSparkParameterValue,
+    TSessionHandle,
+    TOpenSessionResp
 )
 from databricks.sql.utils import DbSqlParameter, DbSqlType
 import pytest
 
-from decimal import Decimal
+from databricks.sql.thrift_backend import ThriftBackend
+from databricks.sql.thrift_api.TCLIService import ttypes
 
+from decimal import Decimal
+from databricks.sql.client import Connection
 from typing import List
 
+class TestSessionHandleChecks(object):
+    def test_get_session_handle(self):
+        assert ThriftBackend.get_session_handle_from_resp(TOpenSessionResp(serverProtocolVersion=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V7, sessionHandle=TSessionHandle(1, None) )) == TSessionHandle(1, ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V7)
+        assert ThriftBackend.get_session_handle_from_resp(TOpenSessionResp(serverProtocolVersion=None, sessionHandle=TSessionHandle(1, None) )) == TSessionHandle(1, None)
+        assert ThriftBackend.get_session_handle_from_resp(TOpenSessionResp(serverProtocolVersion=ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V8, sessionHandle=TSessionHandle(1, None) )) == TSessionHandle(1, ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V8)
+    
+    def test_parameters_enabled(self):
+        assert Connection.server_parameterized_queries_enabled(None) == False
+        assert Connection.server_parameterized_queries_enabled(ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V7) == False
+        assert Connection.server_parameterized_queries_enabled(ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V8) == True
 
 class TestTSparkParameterConversion(object):
     @pytest.mark.parametrize(