From fabe731480a82b8708498a8f8a9aa22547f2e9b1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jesse=20Tu=C4=9Flu?= Date: Wed, 8 Jul 2026 14:33:39 -0700 Subject: [PATCH 1/2] feat: ensure statusCode is included in router logs --- .../server/AsyncQueryForwardingServlet.java | 49 ++++++++++++---- .../AsyncQueryForwardingServletTest.java | 57 +++++++++++++++++-- 2 files changed, 91 insertions(+), 15 deletions(-) diff --git a/services/src/main/java/org/apache/druid/server/AsyncQueryForwardingServlet.java b/services/src/main/java/org/apache/druid/server/AsyncQueryForwardingServlet.java index 4054e6249b1e..3eadfbd65148 100644 --- a/services/src/main/java/org/apache/druid/server/AsyncQueryForwardingServlet.java +++ b/services/src/main/java/org/apache/druid/server/AsyncQueryForwardingServlet.java @@ -24,7 +24,6 @@ import com.fasterxml.jackson.dataformat.smile.SmileFactory; import com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes; import com.google.common.annotations.VisibleForTesting; -import com.google.common.collect.ImmutableMap; import com.google.inject.Inject; import com.google.inject.Provider; import org.apache.calcite.avatica.remote.ProtobufTranslation; @@ -387,7 +386,7 @@ private void handleQueryParseException( final String errorMessage = exceptionToReport.getMessage() == null ? "no error message" : exceptionToReport.getMessage(); - AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(request); + final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(request); if (isNativeQuery) { requestLogger.logNativeQuery( @@ -395,7 +394,18 @@ private void handleQueryParseException( null, DateTimes.nowUtc(), request.getRemoteAddr(), - new QueryStats(ImmutableMap.of("success", false, "exception", errorMessage, "identity", authenticationResult.getIdentity())) + new QueryStats( + Map.of( + "success", + false, + DruidMetrics.STATUS_CODE, + httpStatusCode, + "exception", + errorMessage, + "identity", + authenticationResult.getIdentity() + ) + ) ) ); } else { @@ -405,7 +415,18 @@ private void handleQueryParseException( null, DateTimes.nowUtc(), request.getRemoteAddr(), - new QueryStats(ImmutableMap.of("success", false, "exception", errorMessage, "identity", authenticationResult.getIdentity())) + new QueryStats( + Map.of( + "success", + false, + DruidMetrics.STATUS_CODE, + httpStatusCode, + "exception", + errorMessage, + "identity", + authenticationResult.getIdentity() + ) + ) ) ); } @@ -770,7 +791,7 @@ public void onComplete(Result result) failedQueryCount.incrementAndGet(); } - AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req); + final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req); // As router is simply a proxy, we don't make an effort to construct the error code from the exception ourselves. // We rely on broker to set this for us if the error occurs downstream. @@ -789,11 +810,13 @@ public void onComplete(Result result) DateTimes.nowUtc(), req.getRemoteAddr(), new QueryStats( - ImmutableMap.of( + Map.of( "query/time", TimeUnit.NANOSECONDS.toMillis(requestTimeNs), "success", success, + DruidMetrics.STATUS_CODE, + statusCode, "identity", authenticationResult.getIdentity() ) @@ -816,11 +839,13 @@ public void onComplete(Result result) DateTimes.nowUtc(), req.getRemoteAddr(), new QueryStats( - ImmutableMap.of( + Map.of( "query/time", TimeUnit.NANOSECONDS.toMillis(requestTimeNs), "success", success, + DruidMetrics.STATUS_CODE, + statusCode, "identity", authenticationResult.getIdentity() ) @@ -863,7 +888,7 @@ public void onFailure(Response response, Throwable failure) // We rely on broker to set this for us if the error occurs downstream. // Otherwise, if there's a router/client error, we log this as an unknown error. final int statusCode = determineStatusCode(false, response.getStatus()); - AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req); + final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req); emitQueryTime(requestTimeNs, false, sqlQueryId, queryId, statusCode, authenticationResult); //noinspection VariableNotUsedInsideIf @@ -878,9 +903,11 @@ public void onFailure(Response response, Throwable failure) DateTimes.nowUtc(), req.getRemoteAddr(), new QueryStats( - ImmutableMap.of( + Map.of( "success", false, + DruidMetrics.STATUS_CODE, + statusCode, "exception", errorMessage == null ? "no message" : errorMessage, "identity", @@ -910,9 +937,11 @@ public void onFailure(Response response, Throwable failure) DateTimes.nowUtc(), req.getRemoteAddr(), new QueryStats( - ImmutableMap.of( + Map.of( "success", false, + DruidMetrics.STATUS_CODE, + statusCode, "exception", errorMessage == null ? "no message" : errorMessage, "identity", diff --git a/services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java b/services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java index c3b50514e826..9903eb379dbe 100644 --- a/services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java +++ b/services/src/test/java/org/apache/druid/server/AsyncQueryForwardingServletTest.java @@ -73,6 +73,7 @@ import org.apache.druid.server.initialization.jetty.JettyServerInitUtils; import org.apache.druid.server.initialization.jetty.JettyServerInitializer; import org.apache.druid.server.log.NoopRequestLogger; +import org.apache.druid.server.log.RequestLogger; import org.apache.druid.server.metrics.NoopServiceEmitter; import org.apache.druid.server.router.QueryHostFinder; import org.apache.druid.server.router.RendezvousHashAvaticaConnectionBalancer; @@ -506,7 +507,7 @@ public QueryMetrics> makeMetrics() } @Test - public void testMetricsEmittedWithErrorStatusCodeButNoResultException() + public void testMetricsEmittedWithErrorStatusCodeButNoResultException() throws IOException { final TimeseriesQuery query = Druids.newTimeseriesQueryBuilder() .dataSource("foo") @@ -539,6 +540,7 @@ public Throwable getFailure() }; final StubServiceEmitter stubServiceEmitter = StubServiceEmitter.createStarted(); + final RequestLogger requestLogger = Mockito.mock(RequestLogger.class); final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet( new MapQueryToolChestWarehouse(ImmutableMap.of()), TestHelper.makeJsonMapper(), @@ -547,7 +549,7 @@ public Throwable getFailure() null, null, stubServiceEmitter, - NoopRequestLogger.instance(), + requestLogger, makeRouterTestOverrideEmittingFactory(), new AuthenticatorMapper(ImmutableMap.of()), new Properties(), @@ -566,6 +568,7 @@ public Throwable getFailure() 504, stubServiceEmitter.getMetricEvents("query/time").get(0).toMap().get(DruidMetrics.STATUS_CODE) ); + assertNativeQueryStatusCodeMatchesMetric(requestLogger, getEmittedStatusCode(stubServiceEmitter)); Assert.assertEquals("false", stubServiceEmitter.getMetricEvents("query/time").get(0).toMap().get("success")); Assert.assertEquals("testUser", stubServiceEmitter.getMetricEvents("query/time").get(0).toMap().get("identity")); } @@ -643,7 +646,7 @@ public Throwable getFailure() } @Test - public void testOnFailureWithExceptionAndUnassignedStatusCode() + public void testOnFailureWithExceptionAndUnassignedStatusCode() throws IOException { final TimeseriesQuery query = Druids.newTimeseriesQueryBuilder() .dataSource("foo") @@ -665,6 +668,7 @@ public void testOnFailureWithExceptionAndUnassignedStatusCode() Mockito.when(responseMock.getHeaders()).thenReturn(HttpFields.build()); final StubServiceEmitter stubServiceEmitter = StubServiceEmitter.createStarted(); + final RequestLogger requestLogger = Mockito.mock(RequestLogger.class); final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet( new MapQueryToolChestWarehouse(ImmutableMap.of()), TestHelper.makeJsonMapper(), @@ -673,7 +677,7 @@ public void testOnFailureWithExceptionAndUnassignedStatusCode() null, null, stubServiceEmitter, - NoopRequestLogger.instance(), + requestLogger, makeRouterTestOverrideEmittingFactory(), new AuthenticatorMapper(ImmutableMap.of()), new Properties(), @@ -693,6 +697,7 @@ public void testOnFailureWithExceptionAndUnassignedStatusCode() 500, // Should default to 500 when status is 0 stubServiceEmitter.getMetricEvents("query/time").get(0).toMap().get(DruidMetrics.STATUS_CODE) ); + assertNativeQueryStatusCodeMatchesMetric(requestLogger, getEmittedStatusCode(stubServiceEmitter)); Assert.assertEquals("false", stubServiceEmitter.getMetricEvents("query/time").get(0).toMap().get("success")); Assert.assertEquals("testUser", stubServiceEmitter.getMetricEvents("query/time").get(0).toMap().get("identity")); } @@ -913,6 +918,7 @@ public CompletableFuture abort(Throwable throwable) }; final Result result = new Result(proxyRequestMock, response); final StubServiceEmitter stubServiceEmitter = StubServiceEmitter.createStarted(); + final RequestLogger requestLogger = Mockito.mock(RequestLogger.class); final AsyncQueryForwardingServlet servlet = new AsyncQueryForwardingServlet( new MapQueryToolChestWarehouse(ImmutableMap.of()), jsonMapper, @@ -921,7 +927,7 @@ public CompletableFuture abort(Throwable throwable) null, null, stubServiceEmitter, - NoopRequestLogger.instance(), + requestLogger, new DefaultGenericQueryMetricsFactory(), new AuthenticatorMapper(ImmutableMap.of()), properties, @@ -967,12 +973,53 @@ protected void doService( stubServiceEmitter.getMetricEvents("query/time").get(0).toMap().get(DruidMetrics.STATUS_CODE) ); } + if (!isJDBCSql) { + assertStatusCodeMatchesMetric(requestLogger, isNativeSql, getEmittedStatusCode(stubServiceEmitter)); + } // This test is mostly about verifying that the servlet calls the right methods the right number of times. EasyMock.verify(hostFinder, requestMock); Assert.assertEquals(1, didService.get()); } + private static int getEmittedStatusCode(StubServiceEmitter stubServiceEmitter) + { + return (int) stubServiceEmitter.getMetricEvents("query/time").get(0).toMap().get(DruidMetrics.STATUS_CODE); + } + + private static void assertStatusCodeMatchesMetric( + RequestLogger requestLogger, + boolean isSqlQuery, + int metricStatusCode + ) throws IOException + { + if (isSqlQuery) { + final ArgumentCaptor requestLogLineCaptor = ArgumentCaptor.forClass(RequestLogLine.class); + Mockito.verify(requestLogger).logSqlQuery(requestLogLineCaptor.capture()); + assertStatusCode(requestLogLineCaptor.getValue(), metricStatusCode); + } else { + assertNativeQueryStatusCodeMatchesMetric(requestLogger, metricStatusCode); + } + } + + private static void assertNativeQueryStatusCodeMatchesMetric( + RequestLogger requestLogger, + int metricStatusCode + ) throws IOException + { + final ArgumentCaptor requestLogLineCaptor = ArgumentCaptor.forClass(RequestLogLine.class); + Mockito.verify(requestLogger).logNativeQuery(requestLogLineCaptor.capture()); + assertStatusCode(requestLogLineCaptor.getValue(), metricStatusCode); + } + + private static void assertStatusCode(RequestLogLine requestLogLine, int metricStatusCode) + { + Assert.assertEquals( + metricStatusCode, + requestLogLine.getQueryStats().getStats().get(DruidMetrics.STATUS_CODE) + ); + } + private static Server makeTestDeleteServer(int port, final CountDownLatch latch) { Server server = new Server(port); From a3a98908eef93d6d7e47429c3ac399850f98eaa5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jesse=20Tu=C4=9Flu?= Date: Thu, 9 Jul 2026 09:13:02 -0700 Subject: [PATCH 2/2] clean up --- .../server/AsyncQueryForwardingServlet.java | 126 ++++++++---------- 1 file changed, 58 insertions(+), 68 deletions(-) diff --git a/services/src/main/java/org/apache/druid/server/AsyncQueryForwardingServlet.java b/services/src/main/java/org/apache/druid/server/AsyncQueryForwardingServlet.java index 3eadfbd65148..7b8cd453b400 100644 --- a/services/src/main/java/org/apache/druid/server/AsyncQueryForwardingServlet.java +++ b/services/src/main/java/org/apache/druid/server/AsyncQueryForwardingServlet.java @@ -394,18 +394,7 @@ private void handleQueryParseException( null, DateTimes.nowUtc(), request.getRemoteAddr(), - new QueryStats( - Map.of( - "success", - false, - DruidMetrics.STATUS_CODE, - httpStatusCode, - "exception", - errorMessage, - "identity", - authenticationResult.getIdentity() - ) - ) + buildRequestLogQueryStats(false, httpStatusCode, authenticationResult.getIdentity(), null, errorMessage) ) ); } else { @@ -415,18 +404,7 @@ private void handleQueryParseException( null, DateTimes.nowUtc(), request.getRemoteAddr(), - new QueryStats( - Map.of( - "success", - false, - DruidMetrics.STATUS_CODE, - httpStatusCode, - "exception", - errorMessage, - "identity", - authenticationResult.getIdentity() - ) - ) + buildRequestLogQueryStats(false, httpStatusCode, authenticationResult.getIdentity(), null, errorMessage) ) ); } @@ -440,6 +418,38 @@ private void handleQueryParseException( ); } + /** + * Builds the {@link QueryStats} recorded in router request logs. Success paths pass {@code queryTimeMs}; + * failure paths pass {@code exceptionMessage}. Exactly one of the two is non-null at every call site. + */ + private static QueryStats buildRequestLogQueryStats( + boolean success, + int statusCode, + String identity, + @Nullable Long queryTimeMs, + @Nullable String exceptionMessage + ) + { + if (exceptionMessage != null) { + return new QueryStats( + Map.of( + "success", success, + DruidMetrics.STATUS_CODE, statusCode, + "exception", exceptionMessage, + "identity", identity + ) + ); + } + return new QueryStats( + Map.of( + "query/time", queryTimeMs, + "success", success, + DruidMetrics.STATUS_CODE, statusCode, + "identity", identity + ) + ); + } + protected void doService( HttpServletRequest request, HttpServletResponse response @@ -809,17 +819,12 @@ public void onComplete(Result result) sqlQuery.getContext(), DateTimes.nowUtc(), req.getRemoteAddr(), - new QueryStats( - Map.of( - "query/time", - TimeUnit.NANOSECONDS.toMillis(requestTimeNs), - "success", - success, - DruidMetrics.STATUS_CODE, - statusCode, - "identity", - authenticationResult.getIdentity() - ) + buildRequestLogQueryStats( + success, + statusCode, + authenticationResult.getIdentity(), + TimeUnit.NANOSECONDS.toMillis(requestTimeNs), + null ) ) ); @@ -838,17 +843,12 @@ public void onComplete(Result result) query, DateTimes.nowUtc(), req.getRemoteAddr(), - new QueryStats( - Map.of( - "query/time", - TimeUnit.NANOSECONDS.toMillis(requestTimeNs), - "success", - success, - DruidMetrics.STATUS_CODE, - statusCode, - "identity", - authenticationResult.getIdentity() - ) + buildRequestLogQueryStats( + success, + statusCode, + authenticationResult.getIdentity(), + TimeUnit.NANOSECONDS.toMillis(requestTimeNs), + null ) ) ); @@ -902,17 +902,12 @@ public void onFailure(Response response, Throwable failure) sqlQuery.getContext(), DateTimes.nowUtc(), req.getRemoteAddr(), - new QueryStats( - Map.of( - "success", - false, - DruidMetrics.STATUS_CODE, - statusCode, - "exception", - errorMessage == null ? "no message" : errorMessage, - "identity", - authenticationResult.getIdentity() - ) + buildRequestLogQueryStats( + false, + statusCode, + authenticationResult.getIdentity(), + null, + errorMessage == null ? "no message" : errorMessage ) ) ); @@ -936,17 +931,12 @@ public void onFailure(Response response, Throwable failure) query, DateTimes.nowUtc(), req.getRemoteAddr(), - new QueryStats( - Map.of( - "success", - false, - DruidMetrics.STATUS_CODE, - statusCode, - "exception", - errorMessage == null ? "no message" : errorMessage, - "identity", - authenticationResult.getIdentity() - ) + buildRequestLogQueryStats( + false, + statusCode, + authenticationResult.getIdentity(), + null, + errorMessage == null ? "no message" : errorMessage ) ) );