generalise open session, fix session tests to consider positional args

Author: varun-edachali-dbx

src/databricks/sql/session.py

@@ -84,27 +84,25 @@ def __init__(
             **kwargs,
         )
 
-        self._handle = None
         self.protocol_version = None
 
     def open(self) -> None:
-        self._open_session_resp = self.thrift_backend.open_session(
+        self._session_id = self.backend.open_session(
             self.session_configuration, self.catalog, self.schema
         )
-        self._handle = self._open_session_resp.sessionHandle
-        self.protocol_version = self.get_protocol_version(self._open_session_resp)
+        self.protocol_version = self.get_protocol_version(self._session_id)
         self.is_open = True
         logger.info("Successfully opened session " + str(self.get_id_hex()))
 
     @staticmethod
-    def get_protocol_version(openSessionResp):
+    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 sessionId.backend_type != BackendType.THRIFT:
+        if session_id.backend_type != BackendType.THRIFT:
             return None
-        session_handle = sessionId.to_thrift_handle()
+        session_handle = session_id.to_thrift_handle()
         if (
             session_handle
             and hasattr(session_handle, "serverProtocolVersion")

tests/unit/test_session.py

@@ -146,18 +146,18 @@ def test_socket_timeout_passthrough(self, mock_client_class):
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_configuration_passthrough(self, mock_client_class):
         mock_session_config = Mock()
-
+        
         # Create a mock SessionId that will be returned by open_session
         mock_session_id = SessionId(BackendType.THRIFT, b"\x22", b"\x33")
         mock_client_class.return_value.open_session.return_value = mock_session_id
-
+        
         databricks.sql.connect(
             session_configuration=mock_session_config, **self.DUMMY_CONNECTION_ARGS
         )
 
-        # Check that open_session was called with the correct session_configuration
-        call_kwargs = mock_client_class.return_value.open_session.call_args[1]
-        self.assertEqual(call_kwargs["session_configuration"], mock_session_config)
+        # Check that open_session was called with the correct session_configuration as first positional argument
+        call_args = mock_client_class.return_value.open_session.call_args[0]
+        self.assertEqual(call_args[0], mock_session_config)
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_initial_namespace_passthrough(self, mock_client_class):
@@ -171,11 +171,11 @@ def test_initial_namespace_passthrough(self, mock_client_class):
         databricks.sql.connect(
             **self.DUMMY_CONNECTION_ARGS, catalog=mock_cat, schema=mock_schem
         )
-
-        # Check that open_session was called with the correct catalog and schema
-        call_kwargs = mock_client_class.return_value.open_session.call_args[1]
-        self.assertEqual(call_kwargs["catalog"], mock_cat)
-        self.assertEqual(call_kwargs["schema"], mock_schem)
+        
+        # Check that open_session was called with the correct catalog and schema as positional arguments
+        call_args = mock_client_class.return_value.open_session.call_args[0]
+        self.assertEqual(call_args[1], mock_cat)  # catalog is second positional argument
+        self.assertEqual(call_args[2], mock_schem)  # schema is third positional argument
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_finalizer_closes_abandoned_connection(self, mock_client_class):