|
18 | 18 |
|
19 | 19 | import static org.apache.arrow.flight.FlightTestUtil.LOCALHOST; |
20 | 20 | import static org.apache.arrow.flight.Location.forGrpcInsecure; |
| 21 | +import static org.junit.jupiter.api.Assertions.assertArrayEquals; |
21 | 22 | import static org.junit.jupiter.api.Assertions.assertEquals; |
| 23 | +import static org.junit.jupiter.api.Assertions.assertFalse; |
22 | 24 | import static org.junit.jupiter.api.Assertions.assertNotNull; |
23 | 25 | import static org.junit.jupiter.api.Assertions.assertNull; |
| 26 | +import static org.junit.jupiter.api.Assertions.assertTrue; |
24 | 27 | import static org.junit.jupiter.api.Assertions.fail; |
25 | 28 |
|
26 | 29 | import java.io.IOException; |
27 | 30 | import java.util.Collections; |
28 | 31 | import java.util.List; |
29 | 32 | import java.util.concurrent.CompletableFuture; |
| 33 | +import java.util.concurrent.atomic.AtomicBoolean; |
30 | 34 | import java.util.function.BiConsumer; |
31 | 35 | import org.apache.arrow.flight.FlightClient.ClientStreamListener; |
32 | 36 | import org.apache.arrow.memory.BufferAllocator; |
@@ -87,6 +91,71 @@ public void doPutCustomCode() { |
87 | 91 | }); |
88 | 92 | } |
89 | 93 |
|
| 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 | + |
90 | 159 | /** Make sure uncaught exceptions in DoPut are intercepted. */ |
91 | 160 | @Test |
92 | 161 | public void doPutUncaught() { |
@@ -232,6 +301,22 @@ public ErrorRecorder onCallStarted( |
232 | 301 | } |
233 | 302 | } |
234 | 303 |
|
| 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 | + |
235 | 320 | /** A producer that throws the given exception on a call. */ |
236 | 321 | static class ErrorProducer extends NoOpFlightProducer { |
237 | 322 |
|
|
0 commit comments