Skip to content

Commit

Permalink
fix: use serialisation helper function (#808)
Browse files Browse the repository at this point in the history
* fix: use serialisation helper function

* fix: use helper to fetch token, only close if flight client exists
  • Loading branch information
dhrudevalia authored Jan 2, 2025
1 parent a4818c8 commit 359156e
Showing 1 changed file with 7 additions and 4 deletions.
11 changes: 7 additions & 4 deletions graphdatascience/query_runner/gds_arrow_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,8 @@ def request_token(self) -> Optional[str]:
a token from the server and returns it.
"""
if self._auth:
self._flight_client.authenticate_basic_token(self._auth[0], self._auth[1])
client = self._client()
client.authenticate_basic_token(self._auth[0], self._auth[1])
return self._auth_middleware.token()
else:
return "IGNORED"
Expand Down Expand Up @@ -654,8 +655,9 @@ def _do_get(

ticket = flight.Ticket(json.dumps(payload).encode("utf-8"))

client = self._client()
try:
get = self._flight_client.do_get(ticket)
get = client.do_get(ticket)
arrow_table = get.read_all()
except Exception as e:
self.handle_flight_error(e)
Expand Down Expand Up @@ -683,10 +685,11 @@ def __exit__(
exception_value: Optional[BaseException],
traceback: Optional[TracebackType],
) -> None:
self._flight_client.close()
self.close()

def close(self) -> None:
self._flight_client.close()
if self._flight_client:
self._flight_client.close()

def _versioned_action_type(self, action_type: str) -> str:
return self._arrow_endpoint_version.prefix() + action_type
Expand Down

0 comments on commit 359156e

Please sign in to comment.