take TOpenSessionResp in get_protocol_version to maintain existing interface

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

Author: varun-edachali-dbx

src/databricks/sql/client.py

@@ -50,6 +50,7 @@
 from databricks.sql.backend.types import CommandId, BackendType, CommandState, SessionId
 
 from databricks.sql.thrift_api.TCLIService.ttypes import (
+    TOpenSessionResp,
     TSparkParameter,
     TOperationState,
 )
@@ -320,8 +321,16 @@ def protocol_version(self):
         return self.session.protocol_version
 
     @staticmethod
-    def get_protocol_version(session_id: SessionId):
+    def get_protocol_version(openSessionResp: TOpenSessionResp):
         """Delegate to Session class static method"""
+        properties = (
+            {"serverProtocolVersion": openSessionResp.serverProtocolVersion}
+            if openSessionResp.serverProtocolVersion
+            else {}
+        )
+        session_id = SessionId.from_thrift_handle(
+            openSessionResp.sessionHandle, properties
+        )
         return Session.get_protocol_version(session_id)
 
     @property

tests/unit/test_parameters.py

@@ -64,13 +64,7 @@ class TestSessionHandleChecks(object):
         ],
     )
     def test_get_protocol_version_fallback_behavior(self, test_input, expected):
-        properties = (
-            {"serverProtocolVersion": test_input.serverProtocolVersion}
-            if test_input.serverProtocolVersion
-            else {}
-        )
-        session_id = SessionId.from_thrift_handle(test_input.sessionHandle, properties)
-        assert Connection.get_protocol_version(session_id) == expected
+        assert Connection.get_protocol_version(test_input) == expected
 
     @pytest.mark.parametrize(
         "test_input,expected",