fix: sql.thrift_backend -> sql.backend.thrift_backend in tests and examples

Author: varun-edachali-dbx

examples/v3_retries_query_execute.py

@@ -26,7 +26,7 @@
 # which means all redirects will be followed. In this case, a redirect will count toward the
 # _retry_stop_after_attempts_count which means that by default the connector will not enter an endless retry loop.
 #
-# For complete information about configuring retries, see the docstring for databricks.sql.thrift_backend.ThriftBackend
+# For complete information about configuring retries, see the docstring for databricks.sql.backend.thrift_backend.ThriftBackend
 
 with sql.connect(
     server_hostname=os.getenv("DATABRICKS_SERVER_HOSTNAME"),

tests/unit/test_thrift_backend.py

@@ -133,7 +133,7 @@ def test_hive_schema_to_arrow_schema_preserves_column_names(self):
         self.assertEqual(arrow_schema.field(2).name, "column 2")
         self.assertEqual(arrow_schema.field(3).name, "")
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_bad_protocol_versions_are_rejected(self, tcli_service_client_cass):
         t_http_client_instance = tcli_service_client_cass.return_value
         bad_protocol_versions = [
@@ -163,7 +163,7 @@ def test_bad_protocol_versions_are_rejected(self, tcli_service_client_cass):
                 "expected server to use a protocol version", str(cm.exception)
             )
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_okay_protocol_versions_succeed(self, tcli_service_client_cass):
         t_http_client_instance = tcli_service_client_cass.return_value
         good_protocol_versions = [
@@ -642,7 +642,7 @@ def test_handle_execute_response_sets_compression_in_direct_results(
             )
             self.assertEqual(execute_response.lz4_compressed, lz4Compressed)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_handle_execute_response_checks_operation_state_in_polls(
         self, tcli_service_class
     ):
@@ -686,7 +686,7 @@ def test_handle_execute_response_checks_operation_state_in_polls(
                 if op_state_resp.errorMessage:
                     self.assertIn(op_state_resp.errorMessage, str(cm.exception))
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_get_status_uses_display_message_if_available(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
 
@@ -724,7 +724,7 @@ def test_get_status_uses_display_message_if_available(self, tcli_service_class):
         self.assertEqual(display_message, str(cm.exception))
         self.assertIn(diagnostic_info, str(cm.exception.message_with_context()))
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_direct_results_uses_display_message_if_available(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
 
@@ -825,7 +825,7 @@ def test_handle_execute_response_checks_direct_results_for_error_statuses(self):
                         thrift_backend._handle_execute_response(error_resp, Mock())
                     self.assertIn("this is a bad error", str(cm.exception))
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_handle_execute_response_can_handle_without_direct_results(
         self, tcli_service_class
     ):
@@ -917,7 +917,7 @@ def test_handle_execute_response_can_handle_with_direct_results(self):
                     ttypes.TOperationState.FINISHED_STATE,
                 )
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_use_arrow_schema_if_available(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         arrow_schema_mock = MagicMock(name="Arrow schema mock")
@@ -946,7 +946,7 @@ def test_use_arrow_schema_if_available(self, tcli_service_class):
 
         self.assertEqual(execute_response.arrow_schema_bytes, arrow_schema_mock)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         hive_schema_mock = MagicMock(name="Hive schema mock")
@@ -976,7 +976,7 @@ def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
     @patch(
         "databricks.sql.utils.ResultSetQueueFactory.build_queue", return_value=Mock()
     )
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_handle_execute_response_reads_has_more_rows_in_direct_results(
         self, tcli_service_class, build_queue
     ):
@@ -1020,7 +1020,7 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
     @patch(
         "databricks.sql.utils.ResultSetQueueFactory.build_queue", return_value=Mock()
     )
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_handle_execute_response_reads_has_more_rows_in_result_response(
         self, tcli_service_class, build_queue
     ):
@@ -1075,7 +1075,7 @@ def test_handle_execute_response_reads_has_more_rows_in_result_response(
 
                 self.assertEqual(has_more_rows, has_more_rows_resp)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_arrow_batches_row_count_are_respected(self, tcli_service_class):
         # make some semi-real arrow batches and check the number of rows is correct in the queue
         tcli_service_instance = tcli_service_class.return_value
@@ -1128,7 +1128,7 @@ def test_arrow_batches_row_count_are_respected(self, tcli_service_class):
 
         self.assertEqual(arrow_queue.n_valid_rows, 15 * 10)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_execute_statement_calls_client_and_handle_execute_response(
         self, tcli_service_class
     ):
@@ -1157,7 +1157,7 @@ def test_execute_statement_calls_client_and_handle_execute_response(
             response, cursor_mock
         )
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_get_catalogs_calls_client_and_handle_execute_response(
         self, tcli_service_class
     ):
@@ -1185,7 +1185,7 @@ def test_get_catalogs_calls_client_and_handle_execute_response(
             response, cursor_mock
         )
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_get_schemas_calls_client_and_handle_execute_response(
         self, tcli_service_class
     ):
@@ -1222,7 +1222,7 @@ def test_get_schemas_calls_client_and_handle_execute_response(
             response, cursor_mock
         )
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_get_tables_calls_client_and_handle_execute_response(
         self, tcli_service_class
     ):
@@ -1263,7 +1263,7 @@ def test_get_tables_calls_client_and_handle_execute_response(
             response, cursor_mock
         )
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_get_columns_calls_client_and_handle_execute_response(
         self, tcli_service_class
     ):
@@ -1304,7 +1304,7 @@ def test_get_columns_calls_client_and_handle_execute_response(
             response, cursor_mock
         )
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_open_session_user_provided_session_id_optional(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         tcli_service_instance.OpenSession.return_value = self.open_session_resp
@@ -1320,7 +1320,7 @@ def test_open_session_user_provided_session_id_optional(self, tcli_service_class
         thrift_backend.open_session({}, None, None)
         self.assertEqual(len(tcli_service_instance.OpenSession.call_args_list), 1)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_op_handle_respected_in_close_command(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         thrift_backend = ThriftDatabricksClient(
@@ -1337,7 +1337,7 @@ def test_op_handle_respected_in_close_command(self, tcli_service_class):
             self.operation_handle,
         )
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_session_handle_respected_in_close_session(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
         thrift_backend = ThriftDatabricksClient(
@@ -1354,7 +1354,7 @@ def test_session_handle_respected_in_close_session(self, tcli_service_class):
             self.session_handle,
         )
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_non_arrow_non_column_based_set_triggers_exception(
         self, tcli_service_class
     ):
@@ -1403,8 +1403,8 @@ def test_create_arrow_table_raises_error_for_unsupported_type(self):
         with self.assertRaises(OperationalError):
             thrift_backend._create_arrow_table(t_row_set, Mock(), None, Mock())
 
-    @patch("databricks.sql.thrift_backend.convert_arrow_based_set_to_arrow_table")
-    @patch("databricks.sql.thrift_backend.convert_column_based_set_to_arrow_table")
+    @patch("databricks.sql.backend.thrift_backend.convert_arrow_based_set_to_arrow_table")
+    @patch("databricks.sql.backend.thrift_backend.convert_column_based_set_to_arrow_table")
     def test_create_arrow_table_calls_correct_conversion_method(
         self, convert_col_mock, convert_arrow_mock
     ):
@@ -1597,7 +1597,7 @@ def test_convert_column_based_set_to_arrow_table_uses_types_from_col_set(self):
         self.assertEqual(arrow_table.column(2).to_pylist(), [1.15, 2.2, 3.3])
         self.assertEqual(arrow_table.column(3).to_pylist(), [b"\x11", b"\x22", b"\x33"])
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_cancel_command_uses_active_op_handle(self, tcli_service_class):
         tcli_service_instance = tcli_service_class.return_value
 
@@ -1627,7 +1627,7 @@ def test_handle_execute_response_sets_active_op_handle(self):
         "databricks.sql.thrift_api.TCLIService.TCLIService.Client.GetOperationStatus"
     )
     @patch(
-        "databricks.sql.thrift_backend._retry_policy", new_callable=retry_policy_factory
+        "databricks.sql.backend.thrift_backend._retry_policy", new_callable=retry_policy_factory
     )
     def test_make_request_will_retry_GetOperationStatus(
         self, mock_retry_policy, mock_GetOperationStatus, t_transport_class
@@ -1681,7 +1681,7 @@ def test_make_request_will_retry_GetOperationStatus(
         )
 
         with self.assertLogs(
-            "databricks.sql.thrift_backend", level=logging.WARNING
+            "databricks.sql.backend.thrift_backend", level=logging.WARNING
         ) as cm:
             with self.assertRaises(RequestError):
                 thrift_backend.make_request(client.GetOperationStatus, req)
@@ -1702,7 +1702,7 @@ def test_make_request_will_retry_GetOperationStatus(
         "databricks.sql.thrift_api.TCLIService.TCLIService.Client.GetOperationStatus"
     )
     @patch(
-        "databricks.sql.thrift_backend._retry_policy", new_callable=retry_policy_factory
+        "databricks.sql.backend.thrift_backend._retry_policy", new_callable=retry_policy_factory
     )
     def test_make_request_will_retry_GetOperationStatus_for_http_error(
         self, mock_retry_policy, mock_gos
@@ -1779,7 +1779,7 @@ def test_make_request_wont_retry_if_error_code_not_429_or_503(
 
     @patch("databricks.sql.auth.thrift_http_client.THttpClient")
     @patch(
-        "databricks.sql.thrift_backend._retry_policy", new_callable=retry_policy_factory
+        "databricks.sql.backend.thrift_backend._retry_policy", new_callable=retry_policy_factory
     )
     def test_make_request_will_retry_stop_after_attempts_count_if_retryable(
         self, mock_retry_policy, t_transport_class
@@ -1959,7 +1959,7 @@ def test_retry_args_passthrough(self, mock_http_client):
     @patch("thrift.transport.THttpClient.THttpClient")
     def test_retry_args_bounding(self, mock_http_client):
         retry_delay_test_args_and_expected_values = {}
-        for k, (_, _, min, max) in databricks.sql.thrift_backend._retry_policy.items():
+        for k, (_, _, min, max) in databricks.sql.backend.thrift_backend._retry_policy.items():
             retry_delay_test_args_and_expected_values[k] = (
                 (min - 1, min),
                 (max + 1, max),
@@ -1986,7 +1986,7 @@ def test_retry_args_bounding(self, mock_http_client):
             for arg, val in retry_delay_expected_vals.items():
                 self.assertEqual(getattr(backend, arg), val)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_configuration_passthrough(self, tcli_client_class):
         tcli_service_instance = tcli_client_class.return_value
         tcli_service_instance.OpenSession.return_value = self.open_session_resp
@@ -2011,7 +2011,7 @@ def test_configuration_passthrough(self, tcli_client_class):
         open_session_req = tcli_client_class.return_value.OpenSession.call_args[0][0]
         self.assertEqual(open_session_req.configuration, expected_config)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_cant_set_timestamp_as_string_to_true(self, tcli_client_class):
         tcli_service_instance = tcli_client_class.return_value
         tcli_service_instance.OpenSession.return_value = self.open_session_resp
@@ -2038,7 +2038,7 @@ def _construct_open_session_with_namespace(self, can_use_multiple_cats, cat, sch
             initialNamespace=ttypes.TNamespace(catalogName=cat, schemaName=schem),
         )
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_initial_namespace_passthrough_to_open_session(self, tcli_client_class):
         tcli_service_instance = tcli_client_class.return_value
 
@@ -2066,7 +2066,7 @@ def test_initial_namespace_passthrough_to_open_session(self, tcli_client_class):
                 self.assertEqual(open_session_req.initialNamespace.catalogName, cat)
                 self.assertEqual(open_session_req.initialNamespace.schemaName, schem)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_can_use_multiple_catalogs_is_set_in_open_session_req(
         self, tcli_client_class
     ):
@@ -2086,7 +2086,7 @@ def test_can_use_multiple_catalogs_is_set_in_open_session_req(
         open_session_req = tcli_client_class.return_value.OpenSession.call_args[0][0]
         self.assertTrue(open_session_req.canUseMultipleCatalogs)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_can_use_multiple_catalogs_is_false_fails_with_initial_catalog(
         self, tcli_client_class
     ):
@@ -2126,7 +2126,7 @@ def test_can_use_multiple_catalogs_is_false_fails_with_initial_catalog(
             )
             backend.open_session({}, cat, schem)
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
         tcli_service_instance = tcli_client_class.return_value
 
@@ -2154,8 +2154,8 @@ def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
             str(cm.exception),
         )
 
-    @patch("databricks.sql.thrift_backend.TCLIService.Client", autospec=True)
-    @patch("databricks.sql.thrift_backend.ThriftDatabricksClient._handle_execute_response")
+    @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
+    @patch("databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response")
     def test_execute_command_sets_complex_type_fields_correctly(
         self, mock_handle_execute_response, tcli_service_class
     ):