@@ -140,7 +140,8 @@ def request_token(self) -> Optional[str]:
140140 a token from the server and returns it.
141141 """
142142 if self ._auth :
143- self ._flight_client .authenticate_basic_token (self ._auth [0 ], self ._auth [1 ])
143+ client = self ._client ()
144+ client .authenticate_basic_token (self ._auth [0 ], self ._auth [1 ])
144145 return self ._auth_middleware .token ()
145146 else :
146147 return "IGNORED"
@@ -654,8 +655,9 @@ def _do_get(
654655
655656 ticket = flight .Ticket (json .dumps (payload ).encode ("utf-8" ))
656657
658+ client = self ._client ()
657659 try :
658- get = self . _flight_client .do_get (ticket )
660+ get = client .do_get (ticket )
659661 arrow_table = get .read_all ()
660662 except Exception as e :
661663 self .handle_flight_error (e )
@@ -683,10 +685,11 @@ def __exit__(
683685 exception_value : Optional [BaseException ],
684686 traceback : Optional [TracebackType ],
685687 ) -> None :
686- self ._flight_client . close ()
688+ self .close ()
687689
688690 def close (self ) -> None :
689- self ._flight_client .close ()
691+ if self ._flight_client :
692+ self ._flight_client .close ()
690693
691694 def _versioned_action_type (self , action_type : str ) -> str :
692695 return self ._arrow_endpoint_version .prefix () + action_type
0 commit comments