dremio_arrow_flight_pandas_csv.py

username = "dremio"
password = "dremio123"
sqlquery = '''SELECT * FROM Business.Transportation."NYC Trips" limit 10000'''
hostname = "localhost"
flightport = 32010

from pyarrow import flight

class DremioClientAuthMiddlewareFactory(flight.ClientMiddlewareFactory):
    def __init__(self):
        self.call_credential = []

    def start_call(self, info):
        return DremioClientAuthMiddleware(self)

    def set_call_credential(self, call_credential):
        self.call_credential = call_credential


class DremioClientAuthMiddleware(flight.ClientMiddleware):
    def __init__(self, factory):
        self.factory = factory

    def received_headers(self, headers):
        auth_header_key = 'authorization'
        authorization_header = []
        for key in headers:
          if key.lower() == auth_header_key:
            authorization_header = headers.get(auth_header_key)
        self.factory.set_call_credential([
            b'authorization', authorization_header[0].encode("utf-8")])


scheme = "grpc+tcp"
connection_args = {}

initial_options = flight.FlightCallOptions(headers=[
    (b'routing-tag', b'test-routing-tag'),
    (b'routing-queue', b'Low Cost User Queries')
])
client_auth_middleware = DremioClientAuthMiddlewareFactory()
client = flight.FlightClient("{}://{}:{}".format(scheme, hostname, flightport),
  middleware=[client_auth_middleware], **connection_args)

bearer_token = client.authenticate_basic_token(username, password, initial_options)
print('[INFO] Authentication was successful')

flight_desc = flight.FlightDescriptor.for_command(sqlquery)
options = flight.FlightCallOptions(headers=[bearer_token])
schema = client.get_schema(flight_desc, options)
flight_info = client.get_flight_info(flight.FlightDescriptor.for_command(sqlquery),
    options)
reader = client.do_get(flight_info.endpoints[0].ticket, options)
print(reader.read_pandas().to_csv())