I have this code t hat I'm using to execute queries in Dremio using pyarrow flight:
class DremioConnector:
env: str
auth_token: str
def __init__(self, env: str, auth_token: str):
self.env = env
self.auth_token = auth_token
def get_dremio_client(self):
dremio_url = get_url_from_template(DREMIO_URL, self.env)
return FlightClient(f"grpc+tls://{dremio_url}:*****", tls_root_certs=get_cert(),
middleware=[DremioClientAuthMiddlewareFactory(), CookieMiddlewareFactory()], **{}, )
def get_dremio_credentials(self):
...
...
return username, password
def __create_flight_call_options(self, username: str, password: str, client: FlightClient) -> FlightCallOptions:
headers: list[Any] = []
bearer_token = client.authenticate_basic_token(username, password, FlightCallOptions(headers=headers))
headers.append(bearer_token)
return FlightCallOptions(headers=headers)
def run_query(self, query: str, username: str, password: str, client: FlightClient) -> FlightStreamReader:
flight_desc = FlightDescriptor.for_command(query)
flight_info = client.get_flight_info(flight_desc, self.__create_flight_call_options(username, password, client))
reader = client.do_get(flight_info.endpoints[0].ticket, self.__create_flight_call_options(username, password, client))
return reader
connector = DremioConnector('env', "auth_token")
dremio_client = connector.get_dremio_client()
recently I've started getting this error: E pyarrow._flight.FlightInternalError: Could not finish writing before closing
when calling
bearer_token = client.authenticate_basic_token(username, password, FlightCallOptions(headers=headers))
Does anyone know what might be the issue? I have no idea why this error started popping up, since it was working before. Am I doing something wrong here?
Here is the code I use for connecting to Dremio using PyArrow, works pretty well for Dremio Cloud, hopefully this helps.
Simple PyArrow Client for Dremio