Skip to content

Commit 5a802b0

Browse files
committed
GH-448: Preserve error metadata when server middleware rejects calls
Forward exception trailers through the existing Flight-to-gRPC conversion so middleware rejection preserves custom text and binary error metadata. Verify the real client/server path and a rejection without custom metadata. Generated-by: Codex and DeepSeek
1 parent 617edd5 commit 5a802b0

2 files changed

Lines changed: 90 additions & 2 deletions

File tree

‎flight/flight-core/src/main/java/org/apache/arrow/flight/grpc/ServerInterceptorAdapter.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,11 @@ public <ReqT, RespT> Listener<ReqT> interceptCall(
100100
try {
101101
m = factory.factory.onCallStarted(info, headerAdapter, requestContextAdapter);
102102
} catch (FlightRuntimeException e) {
103-
// Cancel call
104-
call.close(StatusUtils.toGrpcStatus(e.status()), new Metadata());
103+
// Cancel call, but forward any custom error metadata attached to the Flight status
104+
// (GH-448).
105+
final Metadata trailers = Status.trailersFromThrowable(StatusUtils.toGrpcException(e));
106+
call.close(
107+
StatusUtils.toGrpcStatus(e.status()), trailers == null ? new Metadata() : trailers);
105108
return new Listener<ReqT>() {};
106109
}
107110
middleware.add(m);

‎flight/flight-core/src/test/java/org/apache/arrow/flight/TestServerMiddleware.java‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,15 +18,19 @@
1818

1919
import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST;
2020
import static org.apache.arrow.flight.Location.forGrpcInsecure;
21+
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
2122
import static org.junit.jupiter.api.Assertions.assertEquals;
23+
import static org.junit.jupiter.api.Assertions.assertFalse;
2224
import static org.junit.jupiter.api.Assertions.assertNotNull;
2325
import static org.junit.jupiter.api.Assertions.assertNull;
26+
import static org.junit.jupiter.api.Assertions.assertTrue;
2427
import static org.junit.jupiter.api.Assertions.fail;
2528

2629
import java.io.IOException;
2730
import java.util.Collections;
2831
import java.util.List;
2932
import java.util.concurrent.CompletableFuture;
33+
import java.util.concurrent.atomic.AtomicBoolean;
3034
import java.util.function.BiConsumer;
3135
import org.apache.arrow.flight.FlightClient.ClientStreamListener;
3236
import org.apache.arrow.memory.BufferAllocator;
@@ -87,6 +91,71 @@ public void doPutCustomCode() {
8791
});
8892
}
8993

94+
/**
95+
* Make sure custom error metadata attached to a middleware rejection reaches the client (GH-448).
96+
*/
97+
@Test
98+
public void middlewareRejectionPreservesMetadata() {
99+
final ErrorFlightMetadata metadata = new ErrorFlightMetadata();
100+
metadata.insert("x-foo", "foo");
101+
// Values with an embedded 0 and non-ASCII bytes must survive the round trip unchanged.
102+
final byte[] binaryValue = new byte[] {0, (byte) 0x80, 'a'};
103+
metadata.insert("x-bar-bin", binaryValue);
104+
final AtomicBoolean producerCalled = new AtomicBoolean(false);
105+
test(
106+
new NoOpFlightProducer() {
107+
@Override
108+
public FlightInfo getFlightInfo(CallContext context, FlightDescriptor descriptor) {
109+
producerCalled.set(true);
110+
throw new AssertionError(
111+
"Middleware rejection should prevent the producer from being invoked");
112+
}
113+
},
114+
Collections.singletonList(
115+
new ServerMiddlewarePair<>(
116+
FlightServerMiddleware.Key.of("m"),
117+
new RejectingFactory(
118+
CallStatus.UNAUTHENTICATED
119+
.withDescription("metadata rejection")
120+
.withMetadata(metadata)))),
121+
(allocator, client) -> {
122+
final CallStatus status =
123+
FlightTestUtil.assertCode(
124+
FlightStatusCode.UNAUTHENTICATED,
125+
() -> client.getInfo(FlightDescriptor.path("test")));
126+
assertEquals("metadata rejection", status.description());
127+
final ErrorFlightMetadata errorMetadata = status.metadata();
128+
// Assert presence first so a missing key fails clearly instead of throwing.
129+
assertTrue(errorMetadata.containsKey("x-foo"));
130+
assertTrue(errorMetadata.containsKey("x-bar-bin"));
131+
assertEquals("foo", errorMetadata.get("x-foo"));
132+
assertArrayEquals(binaryValue, errorMetadata.getByte("x-bar-bin"));
133+
assertFalse(producerCalled.get());
134+
});
135+
}
136+
137+
/** Control case: a rejection without metadata still propagates code and description. */
138+
@Test
139+
public void middlewareRejectionWithoutMetadata() {
140+
test(
141+
new NoOpFlightProducer(),
142+
Collections.singletonList(
143+
new ServerMiddlewarePair<>(
144+
FlightServerMiddleware.Key.of("m"),
145+
new RejectingFactory(CallStatus.UNAVAILABLE.withDescription("plain rejection")))),
146+
(allocator, client) -> {
147+
final CallStatus status =
148+
FlightTestUtil.assertCode(
149+
FlightStatusCode.UNAVAILABLE,
150+
() -> client.getInfo(FlightDescriptor.path("test")));
151+
assertEquals("plain rejection", status.description());
152+
// gRPC adds its own transport trailers (e.g. content-type), so only assert that no
153+
// custom error metadata came through.
154+
assertFalse(status.metadata().containsKey("x-foo"));
155+
assertFalse(status.metadata().containsKey("x-bar-bin"));
156+
});
157+
}
158+
90159
/** Make sure uncaught exceptions in DoPut are intercepted. */
91160
@Test
92161
public void doPutUncaught() {
@@ -232,6 +301,22 @@ public ErrorRecorder onCallStarted(
232301
}
233302
}
234303

304+
/** A middleware factory that rejects every call before the producer is invoked. */
305+
static class RejectingFactory implements FlightServerMiddleware.Factory<FlightServerMiddleware> {
306+
307+
private final CallStatus status;
308+
309+
RejectingFactory(CallStatus status) {
310+
this.status = status;
311+
}
312+
313+
@Override
314+
public FlightServerMiddleware onCallStarted(
315+
CallInfo info, CallHeaders incomingHeaders, RequestContext context) {
316+
throw status.toRuntimeException();
317+
}
318+
}
319+
235320
/** A producer that throws the given exception on a call. */
236321
static class ErrorProducer extends NoOpFlightProducer {
237322

0 commit comments

Comments
 (0)