delegate protocol version to SessionId

Author: varun-edachali-dbx

src/databricks/sql/ids.py

@@ -127,6 +127,22 @@ def to_hex_id(self) -> str:
         else:
             return str(self.guid)
 
+    def get_protocol_version(self):
+        """
+        Since the sessionHandle will sometimes have a serverProtocolVersion, it takes
+        precedence over the serverProtocolVersion defined in the OpenSessionResponse.
+        """
+        if self.backend_type != BackendType.THRIFT:
+            return None
+        session_handle = self.to_thrift_handle()
+        if (
+            session_handle
+            and hasattr(session_handle, "serverProtocolVersion")
+            and session_handle.serverProtocolVersion
+        ):
+            return session_handle.serverProtocolVersion
+        return None
+
 
 class CommandId:
     """

src/databricks/sql/session.py

@@ -96,20 +96,7 @@ def open(self):
 
     @staticmethod
     def get_protocol_version(session_id: SessionId):
-        """
-        Since the sessionHandle will sometimes have a serverProtocolVersion, it takes
-        precedence over the serverProtocolVersion defined in the OpenSessionResponse.
-        """
-        if session_id.backend_type != BackendType.THRIFT:
-            return None
-        session_handle = session_id.to_thrift_handle()
-        if (
-            session_handle
-            and hasattr(session_handle, "serverProtocolVersion")
-            and session_handle.serverProtocolVersion
-        ):
-            return session_handle.serverProtocolVersion
-        return None
+        return session_id.get_protocol_version()
 
     @staticmethod
     def server_parameterized_queries_enabled(protocolVersion):