Spaces:
Running
Running
| # Licensed to the Apache Software Foundation (ASF) under one | |
| # or more contributor license agreements. See the NOTICE file | |
| # distributed with this work for additional information | |
| # regarding copyright ownership. The ASF licenses this file | |
| # to you under the Apache License, Version 2.0 (the | |
| # "License"); you may not use this file except in compliance | |
| # with the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, | |
| # software distributed under the License is distributed on an | |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
| # KIND, either express or implied. See the License for the | |
| # specific language governing permissions and limitations | |
| # under the License. | |
| """An example Flight CLI client.""" | |
| import argparse | |
| import sys | |
| import pyarrow | |
| import pyarrow.flight | |
| import pyarrow.csv as csv | |
| def list_flights(args, client, connection_args={}): | |
| print('Flights\n=======') | |
| for flight in client.list_flights(): | |
| descriptor = flight.descriptor | |
| if descriptor.descriptor_type == pyarrow.flight.DescriptorType.PATH: | |
| print("Path:", descriptor.path) | |
| elif descriptor.descriptor_type == pyarrow.flight.DescriptorType.CMD: | |
| print("Command:", descriptor.command) | |
| else: | |
| print("Unknown descriptor type") | |
| print("Total records:", end=" ") | |
| if flight.total_records >= 0: | |
| print(flight.total_records) | |
| else: | |
| print("Unknown") | |
| print("Total bytes:", end=" ") | |
| if flight.total_bytes >= 0: | |
| print(flight.total_bytes) | |
| else: | |
| print("Unknown") | |
| print("Number of endpoints:", len(flight.endpoints)) | |
| print("Schema:") | |
| print(flight.schema) | |
| print('---') | |
| print('\nActions\n=======') | |
| for action in client.list_actions(): | |
| print("Type:", action.type) | |
| print("Description:", action.description) | |
| print('---') | |
| def do_action(args, client, connection_args={}): | |
| try: | |
| buf = pyarrow.allocate_buffer(0) | |
| action = pyarrow.flight.Action(args.action_type, buf) | |
| print('Running action', args.action_type) | |
| for result in client.do_action(action): | |
| print("Got result", result.body.to_pybytes()) | |
| except pyarrow.lib.ArrowIOError as e: | |
| print("Error calling action:", e) | |
| def push_data(args, client, connection_args={}): | |
| print('File Name:', args.file) | |
| my_table = csv.read_csv(args.file) | |
| print('Table rows=', str(len(my_table))) | |
| df = my_table.to_pandas() | |
| print(df.head()) | |
| writer, _ = client.do_put( | |
| pyarrow.flight.FlightDescriptor.for_path(args.file), my_table.schema) | |
| writer.write_table(my_table) | |
| writer.close() | |
| def get_flight(args, client, connection_args={}): | |
| if args.path: | |
| descriptor = pyarrow.flight.FlightDescriptor.for_path(*args.path) | |
| else: | |
| descriptor = pyarrow.flight.FlightDescriptor.for_command(args.command) | |
| info = client.get_flight_info(descriptor) | |
| for endpoint in info.endpoints: | |
| print('Ticket:', endpoint.ticket) | |
| for location in endpoint.locations: | |
| print(location) | |
| get_client = pyarrow.flight.FlightClient(location, | |
| **connection_args) | |
| reader = get_client.do_get(endpoint.ticket) | |
| df = reader.read_pandas() | |
| print(df) | |
| def _add_common_arguments(parser): | |
| parser.add_argument('--tls', action='store_true', | |
| help='Enable transport-level security') | |
| parser.add_argument('--tls-roots', default=None, | |
| help='Path to trusted TLS certificate(s)') | |
| parser.add_argument("--mtls", nargs=2, default=None, | |
| metavar=('CERTFILE', 'KEYFILE'), | |
| help="Enable transport-level security") | |
| parser.add_argument('host', type=str, | |
| help="Address or hostname to connect to") | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| subcommands = parser.add_subparsers() | |
| cmd_list = subcommands.add_parser('list') | |
| cmd_list.set_defaults(action='list') | |
| _add_common_arguments(cmd_list) | |
| cmd_list.add_argument('-l', '--list', action='store_true', | |
| help="Print more details.") | |
| cmd_do = subcommands.add_parser('do') | |
| cmd_do.set_defaults(action='do') | |
| _add_common_arguments(cmd_do) | |
| cmd_do.add_argument('action_type', type=str, | |
| help="The action type to run.") | |
| cmd_put = subcommands.add_parser('put') | |
| cmd_put.set_defaults(action='put') | |
| _add_common_arguments(cmd_put) | |
| cmd_put.add_argument('file', type=str, | |
| help="CSV file to upload.") | |
| cmd_get = subcommands.add_parser('get') | |
| cmd_get.set_defaults(action='get') | |
| _add_common_arguments(cmd_get) | |
| cmd_get_descriptor = cmd_get.add_mutually_exclusive_group(required=True) | |
| cmd_get_descriptor.add_argument('-p', '--path', type=str, action='append', | |
| help="The path for the descriptor.") | |
| cmd_get_descriptor.add_argument('-c', '--command', type=str, | |
| help="The command for the descriptor.") | |
| args = parser.parse_args() | |
| if not hasattr(args, 'action'): | |
| parser.print_help() | |
| sys.exit(1) | |
| commands = { | |
| 'list': list_flights, | |
| 'do': do_action, | |
| 'get': get_flight, | |
| 'put': push_data, | |
| } | |
| host, port = args.host.split(':') | |
| port = int(port) | |
| scheme = "grpc+tcp" | |
| connection_args = {} | |
| if args.tls: | |
| scheme = "grpc+tls" | |
| if args.tls_roots: | |
| with open(args.tls_roots, "rb") as root_certs: | |
| connection_args["tls_root_certs"] = root_certs.read() | |
| if args.mtls: | |
| with open(args.mtls[0], "rb") as cert_file: | |
| tls_cert_chain = cert_file.read() | |
| with open(args.mtls[1], "rb") as key_file: | |
| tls_private_key = key_file.read() | |
| connection_args["cert_chain"] = tls_cert_chain | |
| connection_args["private_key"] = tls_private_key | |
| client = pyarrow.flight.FlightClient(f"{scheme}://{host}:{port}", | |
| **connection_args) | |
| while True: | |
| try: | |
| action = pyarrow.flight.Action("healthcheck", b"") | |
| options = pyarrow.flight.FlightCallOptions(timeout=1) | |
| list(client.do_action(action, options=options)) | |
| break | |
| except pyarrow.ArrowIOError as e: | |
| if "Deadline" in str(e): | |
| print("Server is not ready, waiting...") | |
| commands[args.action](args, client, connection_args) | |
| if __name__ == '__main__': | |
| main() | |