Skip to content

Commit 359156e

Browse files
authored
fix: use serialisation helper function (#808)
* fix: use serialisation helper function * fix: use helper to fetch token, only close if flight client exists
1 parent a4818c8 commit 359156e

File tree

1 file changed

+7
-4
lines changed

1 file changed

+7
-4
lines changed

graphdatascience/query_runner/gds_arrow_client.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)