fix merge artifacts

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -605,7 +605,7 @@ def open_session(self, session_configuration, catalog, schema) -> SessionId:
             session_id = SessionId.from_thrift_handle(
                 response.sessionHandle, properties
             )
-            self._session_id_hex = session_id.hex_guid
+            self._session_id_hex = session_id.guid_hex
             return session_id
         except:
             self._transport.close()

src/databricks/sql/backend/types.py

@@ -161,7 +161,7 @@ def __str__(self) -> str:
                 if isinstance(self.secret, bytes)
                 else str(self.secret)
             )
-            return f"{self.hex_guid}|{secret_hex}"
+            return f"{self.guid_hex}|{secret_hex}"
         return str(self.guid)
 
     @classmethod
@@ -240,7 +240,7 @@ def to_sea_session_id(self):
         return self.guid
 
     @property
-    def hex_guid(self) -> str:
+    def guid_hex(self) -> str:
         """
         Get a hexadecimal string representation of the session ID.
 

src/databricks/sql/session.py

@@ -131,7 +131,7 @@ def open(self):
 
     @staticmethod
     def get_protocol_version(session_id: SessionId):
-        return session_id.get_protocol_version()
+        return session_id.protocol_version
 
     @staticmethod
     def server_parameterized_queries_enabled(protocolVersion):

tests/unit/test_client.py

@@ -98,13 +98,9 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
             mock_thrift_client_class: Mock for ThriftBackend class
         """
 
+        # Test once with has_been_closed_server side, once without
         for closed in (True, False):
             with self.subTest(closed=closed):
-                # Set initial state based on whether the command is already closed
-                initial_state = (
-                    CommandState.CLOSED if closed else CommandState.SUCCEEDED
-                )
-
                 # Mock the execute response with controlled state
                 mock_execute_response = Mock(spec=ExecuteResponse)
 
@@ -114,11 +110,14 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
                 )
                 mock_execute_response.has_been_closed_server_side = closed
                 mock_execute_response.is_staging_operation = False
-                mock_execute_response.command_id = Mock(spec=CommandId)
+                mock_execute_response.description = []
 
-                # Mock the backend that will be used
+                # Mock the backend that will be used by the real ThriftResultSet
                 mock_backend = Mock(spec=ThriftDatabricksClient)
                 mock_backend.staging_allowed_local_path = None
+                mock_backend.fetch_results.return_value = (Mock(), False)
+
+                # Configure the decorator's mock to return our specific mock_backend
                 mock_thrift_client_class.return_value = mock_backend
 
                 # Create connection and cursor
@@ -137,16 +136,22 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
                 # Execute a command - this should set cursor.active_result_set to our real result set
                 cursor.execute("SELECT 1")
 
+                # Verify that cursor.execute() set up the result set correctly
+                self.assertIsInstance(cursor.active_result_set, ThriftResultSet)
+                self.assertEqual(
+                    cursor.active_result_set.has_been_closed_server_side, closed
+                )
+
                 # Close the connection - this should trigger the real close chain:
                 # connection.close() -> cursor.close() -> result_set.close()
                 connection.close()
 
                 # Verify the REAL close logic worked through the chain:
                 # 1. has_been_closed_server_side should always be True after close()
-                assert real_result_set.has_been_closed_server_side is True
+                self.assertTrue(real_result_set.has_been_closed_server_side)
 
-                # 2. op_state should always be CLOSED after close()
-                assert real_result_set.op_state == CommandState.CLOSED
+                # 2. status should always be CLOSED after close()
+                self.assertEqual(real_result_set.status, CommandState.CLOSED)
 
                 # 3. Backend close_command should be called appropriately
                 if not closed:
@@ -183,6 +188,7 @@ def test_arraysize_buffer_size_passthrough(
     def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
         mock_connection = Mock()
         mock_backend = Mock()
+        mock_backend.fetch_results.return_value = (Mock(), False)
 
         result_set = ThriftResultSet(
             connection=mock_connection,
@@ -209,6 +215,7 @@ def test_closing_result_set_hard_closes_commands(self):
         mock_session.open = True
         type(mock_connection).session = PropertyMock(return_value=mock_session)
 
+        mock_thrift_backend.fetch_results.return_value = (Mock(), False)
         result_set = ThriftResultSet(
             mock_connection, mock_results_response, mock_thrift_backend
         )

tests/unit/test_session.py

@@ -62,18 +62,18 @@ def test_auth_args(self, mock_client_class):
 
         for args in connection_args:
             connection = databricks.sql.connect(**args)
-            host, port, http_path, *_ = mock_client_class.call_args[0]
-            assert args["server_hostname"] == host
-            assert args["http_path"] == http_path
+            call_kwargs = mock_client_class.call_args[1]
+            assert args["server_hostname"] == call_kwargs["server_hostname"]
+            assert args["http_path"] == call_kwargs["http_path"]
             connection.close()
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_http_header_passthrough(self, mock_client_class):
         http_headers = [("foo", "bar")]
         databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS, http_headers=http_headers)
 
-        call_args = mock_client_class.call_args[0][3]
-        assert ("foo", "bar") in call_args
+        call_kwargs = mock_client_class.call_args[1]
+        assert ("foo", "bar") in call_kwargs["http_headers"]
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_tls_arg_passthrough(self, mock_client_class):
@@ -95,7 +95,8 @@ def test_tls_arg_passthrough(self, mock_client_class):
     def test_useragent_header(self, mock_client_class):
         databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
 
-        http_headers = mock_client_class.call_args[0][3]
+        call_kwargs = mock_client_class.call_args[1]
+        http_headers = call_kwargs["http_headers"]
         user_agent_header = (
             "User-Agent",
             "{}/{}".format(databricks.sql.USER_AGENT_NAME, databricks.sql.__version__),
@@ -109,7 +110,8 @@ def test_useragent_header(self, mock_client_class):
                 databricks.sql.USER_AGENT_NAME, databricks.sql.__version__, "foobar"
             ),
         )
-        http_headers = mock_client_class.call_args[0][3]
+        call_kwargs = mock_client_class.call_args[1]
+        http_headers = call_kwargs["http_headers"]
         assert user_agent_header_with_entry in http_headers
 
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)