[PECO-1137] Reintroduce protocol checking to Python test fw (#248)

* Put in some unit tests, will add e2e

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

* Added e2e test

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

* Linted

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

* re-bumped thrift files

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

* Changed structure to store protocol version as feature of connection

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

* Fixed parameters test

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

* Fixed comments

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

* Update src/databricks/sql/client.py

Co-authored-by: Jesse <jwhitehouse@airpost.net>
Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

* Fixed comments

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

* Removed extra indent

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

---------

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>
Co-authored-by: Jesse <jwhitehouse@airpost.net>
Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: 2 people

src/databricks/sql/client.py

@@ -13,6 +13,7 @@
     SessionAlreadyClosedError,
     CursorAlreadyClosedError,
 )
+from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.thrift_backend import ThriftBackend
 from databricks.sql.utils import (
     ExecuteResponse,
@@ -196,9 +197,11 @@ def read(self) -> Optional[OAuthToken]:
             **kwargs,
         )
 
-        self._session_handle = self.thrift_backend.open_session(
+        self._open_session_resp = self.thrift_backend.open_session(
             session_configuration, catalog, schema
         )
+        self._session_handle = self._open_session_resp.sessionHandle
+        self.protocol_version = self.get_protocol_version(self._open_session_resp)
         self.use_cloud_fetch = kwargs.get("use_cloud_fetch", True)
         self.open = True
         logger.info("Successfully opened session " + str(self.get_session_id_hex()))
@@ -225,6 +228,30 @@ def __del__(self):
     def get_session_id(self):
         return self.thrift_backend.handle_to_id(self._session_handle)
 
+    @staticmethod
+    def get_protocol_version(openSessionResp):
+        """
+        Since the sessionHandle will sometimes have a serverProtocolVersion, it takes
+        precedence over the serverProtocolVersion defined in the OpenSessionResponse.
+        """
+        if (
+            openSessionResp.sessionHandle
+            and hasattr(openSessionResp.sessionHandle, "serverProtocolVersion")
+            and openSessionResp.sessionHandle.serverProtocolVersion
+        ):
+            return openSessionResp.sessionHandle.serverProtocolVersion
+        return openSessionResp.serverProtocolVersion
+
+    @staticmethod
+    def server_parameterized_queries_enabled(protocolVersion):
+        if (
+            protocolVersion
+            and protocolVersion >= ttypes.TProtocolVersion.SPARK_CLI_SERVICE_PROTOCOL_V8
+        ):
+            return True
+        else:
+            return False
+
     def get_session_id_hex(self):
         return self.thrift_backend.handle_to_hex_id(self._session_handle)
 
@@ -501,6 +528,13 @@ def execute(
         """
         if parameters is None:
             parameters = []
+
+        elif not Connection.server_parameterized_queries_enabled(
+            self.connection.protocol_version
+        ):
+            raise NotSupportedError(
+                "Parameterized operations are not supported by this server. DBR 14.1 is required."
+            )
         else:
             parameters = named_parameters_to_tsparkparams(parameters)