@@ -400,7 +400,7 @@ def test_get_status_uses_display_message_if_available(self, tcli_service_class):
400400
401401 thrift_backend = ThriftBackend ("foobar" , 443 , "path" , [], auth_provider = AuthProvider ())
402402 with self .assertRaises (DatabaseError ) as cm :
403- thrift_backend .execute_command (Mock (), Mock (), 100 , 100 , Mock (), Mock ())
403+ thrift_backend .execute_command (Mock (), Mock (), 100 , 100 , Mock (), Mock (), None )
404404
405405 self .assertEqual (display_message , str (cm .exception ))
406406 self .assertIn (diagnostic_info , str (cm .exception .message_with_context ()))
@@ -430,7 +430,7 @@ def test_direct_results_uses_display_message_if_available(self, tcli_service_cla
430430
431431 thrift_backend = ThriftBackend ("foobar" , 443 , "path" , [], auth_provider = AuthProvider ())
432432 with self .assertRaises (DatabaseError ) as cm :
433- thrift_backend .execute_command (Mock (), Mock (), 100 , 100 , Mock (), Mock ())
433+ thrift_backend .execute_command (Mock (), Mock (), 100 , 100 , Mock (), Mock (), None )
434434
435435 self .assertEqual (display_message , str (cm .exception ))
436436 self .assertIn (diagnostic_info , str (cm .exception .message_with_context ()))
@@ -718,7 +718,7 @@ def test_execute_statement_calls_client_and_handle_execute_response(self, tcli_s
718718 thrift_backend ._handle_execute_response = Mock ()
719719 cursor_mock = Mock ()
720720
721- thrift_backend .execute_command ("foo" , Mock (), 100 , 200 , Mock (), cursor_mock )
721+ thrift_backend .execute_command ("foo" , Mock (), 100 , 200 , Mock (), cursor_mock , None )
722722 # Check call to client
723723 req = tcli_service_instance .ExecuteStatement .call_args [0 ][0 ]
724724 get_direct_results = ttypes .TSparkGetDirectResults (maxRows = 100 , maxBytes = 200 )
@@ -877,7 +877,7 @@ def test_non_arrow_non_column_based_set_triggers_exception(self, tcli_service_cl
877877 thrift_backend = self ._make_fake_thrift_backend ()
878878
879879 with self .assertRaises (OperationalError ) as cm :
880- thrift_backend .execute_command ("foo" , Mock (), 100 , 100 , Mock (), Mock ())
880+ thrift_backend .execute_command ("foo" , Mock (), 100 , 100 , Mock (), Mock (), None )
881881 self .assertIn ("Expected results to be in Arrow or column based format" , str (cm .exception ))
882882
883883 def test_create_arrow_table_raises_error_for_unsupported_type (self ):
@@ -1447,7 +1447,7 @@ def test_execute_command_sets_complex_type_fields_correctly(self, mock_handle_ex
14471447 complex_arg_types ["_use_arrow_native_decimals" ] = decimals
14481448
14491449 thrift_backend = ThriftBackend ("foobar" , 443 , "path" , [], auth_provider = AuthProvider (), ** complex_arg_types )
1450- thrift_backend .execute_command (Mock (), Mock (), 100 , 100 , Mock (), Mock ())
1450+ thrift_backend .execute_command (Mock (), Mock (), 100 , 100 , Mock (), Mock (), None )
14511451 t_execute_statement_req = tcli_service_instance .ExecuteStatement .call_args [0 ][0 ]
14521452 # If the value is unset, the native type should default to True
14531453 self .assertEqual (t_execute_statement_req .useArrowNativeTypes .timestampAsArrow ,
0 commit comments