| 85 | |
| 86 | |
| 87 | class FlightServer(flight.FlightServerBase): |
| 88 | def __init__(self, delegate, **kwargs): |
| 89 | super().__init__(**kwargs) |
| 90 | if delegate: |
| 91 | self.delegate = flight.connect( |
| 92 | delegate, |
| 93 | middleware=(TracingClientMiddlewareFactory(),)) |
| 94 | else: |
| 95 | self.delegate = None |
| 96 | |
| 97 | def list_actions(self, context): |
| 98 | return [ |
| 99 | ("get-trace-id", "Get the trace context ID."), |
| 100 | ] |
| 101 | |
| 102 | def do_action(self, context, action): |
| 103 | trace_middleware = context.get_middleware("trace") |
| 104 | if trace_middleware: |
| 105 | TraceContext.set_trace_id(trace_middleware.trace_id) |
| 106 | if action.type == "get-trace-id": |
| 107 | if self.delegate: |
| 108 | for result in self.delegate.do_action(action): |
| 109 | yield result |
| 110 | else: |
| 111 | trace_id = TraceContext.current_trace_id().encode("utf-8") |
| 112 | print("Returning trace ID:", trace_id) |
| 113 | buf = pa.py_buffer(trace_id) |
| 114 | yield pa.flight.Result(buf) |
| 115 | else: |
| 116 | raise KeyError(f"Unknown action {action.type!r}") |
| 117 | |
| 118 | |
| 119 | def main(): |