change names of ThriftBackend -> ThriftDatabricksClient in tests

Author: varun-edachali-dbx

tests/unit/test_client.py

@@ -15,7 +15,7 @@
     THandleIdentifier,
     TOperationType,
 )
-from databricks.sql.thrift_backend import ThriftBackend
+from databricks.sql.thrift_backend import ThriftDatabricksClient
 
 import databricks.sql
 import databricks.sql.client as client
@@ -27,10 +27,10 @@
 from tests.unit.test_arrow_queue import ArrowQueueSuite
 
 
-class ThriftBackendMockFactory:
+class ThriftDatabricksClientMockFactory:
     @classmethod
     def new(cls):
-        ThriftBackendMock = Mock(spec=ThriftBackend)
+        ThriftBackendMock = Mock(spec=ThriftDatabricksClient)
         ThriftBackendMock.return_value = ThriftBackendMock
 
         cls.apply_property_to_mock(ThriftBackendMock, staging_allowed_local_path=None)
@@ -80,7 +80,7 @@ class ClientTestSuite(unittest.TestCase):
         "access_token": "tok",
     }
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME, ThriftBackendMockFactory.new())
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME, ThriftDatabricksClientMockFactory.new())
     @patch("%s.client.ResultSet" % PACKAGE_NAME)
     def test_closing_connection_closes_commands(self, mock_result_set_class):
         # Test once with has_been_closed_server side, once without
@@ -97,7 +97,7 @@ def test_closing_connection_closes_commands(self, mock_result_set_class):
                 )
                 mock_result_set_class.return_value.close.assert_called_once_with()
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_cant_open_cursor_on_closed_connection(self, mock_client_class):
         connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
         self.assertTrue(connection.open)
@@ -107,7 +107,7 @@ def test_cant_open_cursor_on_closed_connection(self, mock_client_class):
             connection.cursor()
         self.assertIn("closed", str(cm.exception))
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     @patch("%s.client.Cursor" % PACKAGE_NAME)
     def test_arraysize_buffer_size_passthrough(
         self, mock_cursor_class, mock_client_class
@@ -124,7 +124,7 @@ def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
         mock_backend = Mock()
         result_set = client.ResultSet(
             connection=mock_connection,
-            thrift_backend=mock_backend,
+            backend=mock_backend,
             execute_response=Mock(),
         )
         # Setup session mock on the mock_connection
@@ -166,7 +166,7 @@ def test_executing_multiple_commands_uses_the_most_recent_command(
         mock_result_set_class.side_effect = mock_result_sets
 
         cursor = client.Cursor(
-            connection=Mock(), thrift_backend=ThriftBackendMockFactory.new()
+            connection=Mock(), backend=ThriftDatabricksClientMockFactory.new()
         )
         cursor.execute("SELECT 1;")
         cursor.execute("SELECT 1;")
@@ -215,7 +215,7 @@ def dict_product(self, dicts):
         """
         return (dict(zip(dicts.keys(), x)) for x in itertools.product(*dicts.values()))
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.client.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_get_schemas_parameters_passed_to_thrift_backend(self, mock_thrift_backend):
         req_args_combinations = self.dict_product(
             dict(
@@ -236,7 +236,7 @@ def test_get_schemas_parameters_passed_to_thrift_backend(self, mock_thrift_backe
                 for k, v in req_args.items():
                     self.assertEqual(v, call_args[k])
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.client.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_get_tables_parameters_passed_to_thrift_backend(self, mock_thrift_backend):
         req_args_combinations = self.dict_product(
             dict(
@@ -259,7 +259,7 @@ def test_get_tables_parameters_passed_to_thrift_backend(self, mock_thrift_backen
                 for k, v in req_args.items():
                     self.assertEqual(v, call_args[k])
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.client.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_get_columns_parameters_passed_to_thrift_backend(self, mock_thrift_backend):
         req_args_combinations = self.dict_product(
             dict(
@@ -310,7 +310,7 @@ def test_version_is_canonical(self):
         self.assertIsNotNone(re.match(canonical_version_re, version))
 
     def test_execute_parameter_passthrough(self):
-        mock_thrift_backend = ThriftBackendMockFactory.new()
+        mock_thrift_backend = ThriftDatabricksClientMockFactory.new()
         cursor = client.Cursor(Mock(), mock_thrift_backend)
 
         tests = [
@@ -334,16 +334,16 @@ def test_execute_parameter_passthrough(self):
                 expected_query,
             )
 
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.client.ThriftDatabricksClient" % PACKAGE_NAME)
     @patch("%s.client.ResultSet" % PACKAGE_NAME)
     def test_executemany_parameter_passhthrough_and_uses_last_result_set(
         self, mock_result_set_class, mock_thrift_backend
     ):
         # Create a new mock result set each time the class is instantiated
         mock_result_set_instances = [Mock(), Mock(), Mock()]
         mock_result_set_class.side_effect = mock_result_set_instances
-        mock_thrift_backend = ThriftBackendMockFactory.new()
-        cursor = client.Cursor(Mock(), mock_thrift_backend())
+        mock_backend = ThriftDatabricksClientMockFactory.new()
+        cursor = client.Cursor(Mock(), mock_backend())
 
         params = [{"x": None}, {"x": "foo1"}, {"x": "bar2"}]
         expected_queries = ["SELECT NULL", "SELECT 'foo1'", "SELECT 'bar2'"]
@@ -368,7 +368,7 @@ def test_executemany_parameter_passhthrough_and_uses_last_result_set(
             "last operation",
         )
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_commit_a_noop(self, mock_thrift_backend_class):
         c = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
         c.commit()
@@ -381,14 +381,14 @@ def test_setoutputsizes_a_noop(self):
         cursor = client.Cursor(Mock(), Mock())
         cursor.setoutputsize(1)
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_rollback_not_supported(self, mock_thrift_backend_class):
         c = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
         with self.assertRaises(NotSupportedError):
             c.rollback()
 
     @unittest.skip("JDW: skipping winter 2024 as we're about to rewrite this interface")
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.client.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_row_number_respected(self, mock_thrift_backend_class):
         def make_fake_row_slice(n_rows):
             mock_slice = Mock()
@@ -413,7 +413,7 @@ def make_fake_row_slice(n_rows):
         self.assertEqual(cursor.rownumber, 29)
 
     @unittest.skip("JDW: skipping winter 2024 as we're about to rewrite this interface")
-    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.client.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_disable_pandas_respected(self, mock_thrift_backend_class):
         mock_thrift_backend = mock_thrift_backend_class.return_value
         mock_table = Mock()
@@ -466,7 +466,7 @@ def test_column_name_api(self):
                 },
             )
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_cursor_keeps_connection_alive(self, mock_client_class):
         instance = mock_client_class.return_value
 
@@ -485,13 +485,13 @@ def test_cursor_keeps_connection_alive(self, mock_client_class):
 
     @patch("%s.utils.ExecuteResponse" % PACKAGE_NAME, autospec=True)
     @patch("%s.client.Cursor._handle_staging_operation" % PACKAGE_NAME)
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_staging_operation_response_is_handled(
         self, mock_client_class, mock_handle_staging_operation, mock_execute_response
     ):
         # If server sets ExecuteResponse.is_staging_operation True then _handle_staging_operation should be called
 
-        ThriftBackendMockFactory.apply_property_to_mock(
+        ThriftDatabricksClientMockFactory.apply_property_to_mock(
             mock_execute_response, is_staging_operation=True
         )
         mock_client_class.execute_command.return_value = mock_execute_response
@@ -504,7 +504,7 @@ def test_staging_operation_response_is_handled(
 
         mock_handle_staging_operation.call_count == 1
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME, ThriftBackendMockFactory.new())
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME, ThriftDatabricksClientMockFactory.new())
     def test_access_current_query_id(self):
         operation_id = "EE6A8778-21FC-438B-92D8-96AC51EE3821"
 

tests/unit/test_fetches.py

@@ -9,6 +9,7 @@
 
 import databricks.sql.client as client
 from databricks.sql.utils import ExecuteResponse, ArrowQueue
+from databricks.sql.thrift_backend import ThriftDatabricksClient
 
 
 @pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
@@ -39,7 +40,7 @@ def make_dummy_result_set_from_initial_results(initial_results):
         arrow_queue = ArrowQueue(arrow_table, len(initial_results), 0)
         rs = client.ResultSet(
             connection=Mock(),
-            thrift_backend=None,
+            backend=None,
             execute_response=ExecuteResponse(
                 status=None,
                 has_been_closed_server_side=True,
@@ -79,13 +80,13 @@ def fetch_results(
 
             return results, batch_index < len(batch_list)
 
-        mock_thrift_backend = Mock()
+        mock_thrift_backend = Mock(spec=ThriftDatabricksClient)
         mock_thrift_backend.fetch_results = fetch_results
         num_cols = len(batch_list[0][0]) if batch_list and batch_list[0] else 0
 
         rs = client.ResultSet(
             connection=Mock(),
-            thrift_backend=mock_thrift_backend,
+            backend=mock_thrift_backend,
             execute_response=ExecuteResponse(
                 status=None,
                 has_been_closed_server_side=False,

tests/unit/test_session.py

@@ -21,7 +21,7 @@ class SessionTestSuite(unittest.TestCase):
         "access_token": "tok",
     }
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_close_uses_the_correct_session_id(self, mock_client_class):
         instance = mock_client_class.return_value
 
@@ -36,7 +36,7 @@ def test_close_uses_the_correct_session_id(self, mock_client_class):
         close_session_id = instance.close_session.call_args[0][0].sessionId
         self.assertEqual(close_session_id, b"\x22")
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_auth_args(self, mock_client_class):
         # Test that the following auth args work:
         # token = foo,
@@ -63,15 +63,15 @@ def test_auth_args(self, mock_client_class):
             self.assertEqual(args["http_path"], http_path)
             connection.close()
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @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]
         self.assertIn(("foo", "bar"), call_args)
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_tls_arg_passthrough(self, mock_client_class):
         databricks.sql.connect(
             **self.DUMMY_CONNECTION_ARGS,
@@ -87,7 +87,7 @@ def test_tls_arg_passthrough(self, mock_client_class):
         self.assertEqual(kwargs["_tls_client_cert_key_file"], "trusted client cert")
         self.assertEqual(kwargs["_tls_client_cert_key_password"], "key password")
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_useragent_header(self, mock_client_class):
         databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
 
@@ -108,7 +108,7 @@ def test_useragent_header(self, mock_client_class):
         http_headers = mock_client_class.call_args[0][3]
         self.assertIn(user_agent_header_with_entry, http_headers)
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_context_manager_closes_connection(self, mock_client_class):
         instance = mock_client_class.return_value
 
@@ -123,7 +123,7 @@ def test_context_manager_closes_connection(self, mock_client_class):
         close_session_id = instance.close_session.call_args[0][0].sessionId
         self.assertEqual(close_session_id, b"\x22")
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_max_number_of_retries_passthrough(self, mock_client_class):
         databricks.sql.connect(
             _retry_stop_after_attempts_count=54, **self.DUMMY_CONNECTION_ARGS
@@ -133,12 +133,12 @@ def test_max_number_of_retries_passthrough(self, mock_client_class):
             mock_client_class.call_args[1]["_retry_stop_after_attempts_count"], 54
         )
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_socket_timeout_passthrough(self, mock_client_class):
         databricks.sql.connect(_socket_timeout=234, **self.DUMMY_CONNECTION_ARGS)
         self.assertEqual(mock_client_class.call_args[1]["_socket_timeout"], 234)
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_configuration_passthrough(self, mock_client_class):
         mock_session_config = Mock()
         databricks.sql.connect(
@@ -150,7 +150,7 @@ def test_configuration_passthrough(self, mock_client_class):
             mock_session_config,
         )
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_initial_namespace_passthrough(self, mock_client_class):
         mock_cat = Mock()
         mock_schem = Mock()
@@ -165,7 +165,7 @@ def test_initial_namespace_passthrough(self, mock_client_class):
             mock_client_class.return_value.open_session.call_args[0][2], mock_schem
         )
 
-    @patch("%s.session.ThriftBackend" % PACKAGE_NAME)
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_finalizer_closes_abandoned_connection(self, mock_client_class):
         instance = mock_client_class.return_value