Function main

python/examples/flight/middleware.py:119–171

()

def main():
    parser = argparse.ArgumentParser()

    subparsers = parser.add_subparsers(dest="command")
    client = subparsers.add_parser("client", help="Run the client.")
    client.add_argument("server")
    client.add_argument("--request-id", default=None)

    server = subparsers.add_parser("server", help="Run the server.")
    server.add_argument(
        "--listen",
        required=True,
        help="The location to listen on (example: grpc://localhost:5050)",
    )
    server.add_argument(
        "--delegate",
        required=False,
        default=None,
        help=("A location to delegate to. That is, this server will "
              "simply call the given server for the response. Demonstrates "
              "propagation of the trace ID between servers."),
    )
    server.add_argument(
        "--otel",
        action="store_true",
        help="Use OpenTelemetry instrumentation."
    )

    args = parser.parse_args()
    if not getattr(args, "command"):
        parser.print_help()
        return 1

    if args.command == "server":
        middleware = {"trace": TracingServerMiddlewareFactory()}
        if args.otel:
            middleware["otel"] = flight.TracingServerMiddlewareFactory()
        server = FlightServer(
            args.delegate,
            location=args.listen,
            middleware=middleware)
        server.serve()
    elif args.command == "client":
        client = flight.connect(
            args.server,
            middleware=(TracingClientMiddlewareFactory(),))
        if args.request_id:
            TraceContext.set_trace_id(args.request_id)
        else:
            TraceContext.set_trace_id("client-chosen-id")

        for result in client.do_action(flight.Action("get-trace-id", b"")):
            print(result.body.to_pybytes())


if __name__ == "__main__":