@@ -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