@@ -127,7 +127,9 @@ def test_hive_schema_to_arrow_schema_preserves_column_names(self):
127
127
]
128
128
129
129
t_table_schema = ttypes.TTableSchema(columns)
130
- arrow_schema = ThriftDatabricksClient._hive_schema_to_arrow_schema(t_table_schema)
130
+ arrow_schema = ThriftDatabricksClient._hive_schema_to_arrow_schema(
131
+ t_table_schema
132
+ )
131
133
132
134
self.assertEqual(arrow_schema.field(0).name, "column 1")
133
135
self.assertEqual(arrow_schema.field(1).name, "column 2")
@@ -1147,7 +1149,9 @@ def test_execute_statement_calls_client_and_handle_execute_response(
1147
1149
thrift_backend._handle_execute_response = Mock()
1148
1150
cursor_mock = Mock()
1149
1151
1150
- result = thrift_backend.execute_command("foo", Mock(), 100, 200, Mock(), cursor_mock)
1152
+ result = thrift_backend.execute_command(
1153
+ "foo", Mock(), 100, 200, Mock(), cursor_mock
1154
+ )
1151
1155
# Verify the result is a ResultSet
1152
1156
self.assertIsInstance(result, ResultSet)
1153
1157
@@ -1419,8 +1423,12 @@ def test_create_arrow_table_raises_error_for_unsupported_type(self):
1419
1423
with self.assertRaises(OperationalError):
1420
1424
thrift_backend._create_arrow_table(t_row_set, Mock(), None, Mock())
1421
1425
1422
- @patch("databricks.sql.backend.thrift_backend.convert_arrow_based_set_to_arrow_table")
1423
- @patch("databricks.sql.backend.thrift_backend.convert_column_based_set_to_arrow_table")
1426
+ @patch(
1427
+ "databricks.sql.backend.thrift_backend.convert_arrow_based_set_to_arrow_table"
1428
+ )
1429
+ @patch(
1430
+ "databricks.sql.backend.thrift_backend.convert_column_based_set_to_arrow_table"
1431
+ )
1424
1432
def test_create_arrow_table_calls_correct_conversion_method(
1425
1433
self, convert_col_mock, convert_arrow_mock
1426
1434
):
@@ -1643,7 +1651,8 @@ def test_handle_execute_response_sets_active_op_handle(self):
1643
1651
"databricks.sql.thrift_api.TCLIService.TCLIService.Client.GetOperationStatus"
1644
1652
)
1645
1653
@patch(
1646
- "databricks.sql.backend.thrift_backend._retry_policy", new_callable=retry_policy_factory
1654
+ "databricks.sql.backend.thrift_backend._retry_policy",
1655
+ new_callable=retry_policy_factory,
1647
1656
)
1648
1657
def test_make_request_will_retry_GetOperationStatus(
1649
1658
self, mock_retry_policy, mock_GetOperationStatus, t_transport_class
@@ -1718,7 +1727,8 @@ def test_make_request_will_retry_GetOperationStatus(
1718
1727
"databricks.sql.thrift_api.TCLIService.TCLIService.Client.GetOperationStatus"
1719
1728
)
1720
1729
@patch(
1721
- "databricks.sql.backend.thrift_backend._retry_policy", new_callable=retry_policy_factory
1730
+ "databricks.sql.backend.thrift_backend._retry_policy",
1731
+ new_callable=retry_policy_factory,
1722
1732
)
1723
1733
def test_make_request_will_retry_GetOperationStatus_for_http_error(
1724
1734
self, mock_retry_policy, mock_gos
@@ -1795,7 +1805,8 @@ def test_make_request_wont_retry_if_error_code_not_429_or_503(
1795
1805
1796
1806
@patch("databricks.sql.auth.thrift_http_client.THttpClient")
1797
1807
@patch(
1798
- "databricks.sql.backend.thrift_backend._retry_policy", new_callable=retry_policy_factory
1808
+ "databricks.sql.backend.thrift_backend._retry_policy",
1809
+ new_callable=retry_policy_factory,
1799
1810
)
1800
1811
def test_make_request_will_retry_stop_after_attempts_count_if_retryable(
1801
1812
self, mock_retry_policy, t_transport_class
@@ -1975,7 +1986,12 @@ def test_retry_args_passthrough(self, mock_http_client):
1975
1986
@patch("thrift.transport.THttpClient.THttpClient")
1976
1987
def test_retry_args_bounding(self, mock_http_client):
1977
1988
retry_delay_test_args_and_expected_values = {}
1978
- for k, (_, _, min, max) in databricks.sql.backend.thrift_backend._retry_policy.items():
1989
+ for k, (
1990
+ _,
1991
+ _,
1992
+ min,
1993
+ max,
1994
+ ) in databricks.sql.backend.thrift_backend._retry_policy.items():
1979
1995
retry_delay_test_args_and_expected_values[k] = (
1980
1996
(min - 1, min),
1981
1997
(max + 1, max),
@@ -2171,7 +2187,9 @@ def test_protocol_v3_fails_if_initial_namespace_set(self, tcli_client_class):
2171
2187
)
2172
2188
2173
2189
@patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
2174
- @patch("databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response")
2190
+ @patch(
2191
+ "databricks.sql.backend.thrift_backend.ThriftDatabricksClient._handle_execute_response"
2192
+ )
2175
2193
def test_execute_command_sets_complex_type_fields_correctly(
2176
2194
self, mock_handle_execute_response, tcli_service_class
2177
2195
):
0 commit comments