Support passing X-Request-ID to Druid

CHANGELOG.md

@@ -10,6 +10,7 @@ Current
 
 ### Added:
 
+- [Support for X-Request-ID](https://github.com/yahoo/fili/pull/68)
 - [Documentation for topN](https://github.com/yahoo/fili/pull/43)
 
 
@@ -58,6 +59,8 @@ Current
 
 ### Fixed:
 
+- [Fixed typo emit -> omit in Utils.java omitField()](https://github.com/yahoo/fili/pull/68)
+
 - [Adds read locking to all attempts to read the Lucene index](https://github.com/yahoo/fili/pull/52)
   * Before, if Fili attempted to read from the Lucene indices (i.e. processing a query with filters)
     while loading dimension indices, the request would fail and we would get a

fili-core/src/main/java/com/yahoo/bard/webservice/logging/LogBlock.java

@@ -36,6 +36,18 @@ public void add(LogInfo phase) {
         body.put(phase.getName(), phase);
     }
 
+    /**
+     * Create a new LogBlock with the same body but updated uuid.
+     *
+     * @param uuid  New uuid.
+     * @return New LogBlock with updated uuid.
+     */
+    public LogBlock withUuid(String uuid) {
+        LogBlock newLogBlock = new LogBlock(uuid);
+        body.entrySet().forEach(entry -> newLogBlock.body.put(entry.getKey(), entry.getValue()));
+        return newLogBlock;
+    }
+
     /**
      * Add an entry of given class in this LogBlock that holds no information.
      * It is meant to be used only to define the order of {@link LogInfo} parts inside a LogBlock.

fili-core/src/main/java/com/yahoo/bard/webservice/logging/RequestLog.java

@@ -400,6 +400,19 @@ public static String getId() {
         return current.logId;
     }
 
+    /**
+     * Prepend an id prefix to generated druid query id.
+     *
+     * @param idPrefix  Prefix for queryId sent to druid
+     */
+    public static void addIdPrefix(String idPrefix) {
+        RequestLog current = RLOG.get();
+        String newId = idPrefix + getId();
+        current.info = current.info.withUuid(newId);
+        current.logId = newId;
+        MDC.put(ID_KEY, newId);
+    }
+
     /**
      * Overwrite current thread's request log context.
      * It first clears the request log of the current thread and then fills it with the contents found in ctx

fili-core/src/main/java/com/yahoo/bard/webservice/util/Utils.java

@@ -196,16 +196,16 @@ public static MultivaluedMap<String, String> headersToLowerCase(MultivaluedMap<S
      * This method is recursive.
      *
      * @param node The root of the tree of json nodes.
-     * @param fieldName The name of the node to be emitted.
+     * @param fieldName The name of the node to be omitted.
      * @param mapper  The object mapper that creates and empty node.
      */
-    public static void emitField(JsonNode node, String fieldName, ObjectMapper mapper) {
+    public static void omitField(JsonNode node, String fieldName, ObjectMapper mapper) {
         if (node.has("context")) {
             ((ObjectNode) node).replace(fieldName, mapper.createObjectNode());
         }
 
         for (JsonNode child : node) {
-            emitField(child, fieldName, mapper);
+            omitField(child, fieldName, mapper);
         }
     }
 

fili-core/src/main/java/com/yahoo/bard/webservice/web/filters/BardLoggingFilter.java

@@ -26,6 +26,7 @@
 import java.util.List;
 import java.util.Map.Entry;
 import java.util.Set;
+import java.util.regex.Pattern;
 
 import javax.annotation.Priority;
 import javax.inject.Singleton;
@@ -68,10 +69,12 @@ public class BardLoggingFilter implements ContainerRequestFilter, ContainerRespo
     private static final String PROPERTY_NANOS = LoggingFilter.class.getName() + ".nanos";
     private static final String PROPERTY_REQ_LEN = LoggingFilter.class.getName() + ".reqlen";
     private static final String PROPERTY_OUTPUT_STREAM = LoggingFilter.class.getName() + ".ostream";
+    private static final Pattern VALID_REQUEST_ID = Pattern.compile("[a-zA-Z0-9+/=-]+");
 
     public static final double MILLISECONDS_PER_NANOSECOND = 1000000.0;
     public static final String TOTAL_TIMER = "TotalTime";
     public static final String CLIENT_TOTAL_TIMER = "ClientRequestTotalTime";
+    public static final String X_REQUEST_ID_HEADER = "x-request-id";
 
     /**
      * Intercept the Container request to add length of request and a start timestamp.
@@ -83,6 +86,7 @@ public class BardLoggingFilter implements ContainerRequestFilter, ContainerRespo
     @Override
     public void filter(ContainerRequestContext request) throws IOException {
 
+        appendRequestId(request.getHeaders().getFirst(X_REQUEST_ID_HEADER));
         RequestLog.startTiming(TOTAL_TIMER);
         RequestLog.startTiming(this);
         RequestLog.record(new Preface(request));
@@ -108,6 +112,7 @@ public void filter(ContainerRequestContext request) throws IOException {
     @Override
     public void filter(ContainerRequestContext request, ContainerResponseContext response)
         throws IOException {
+        appendRequestId(request.getHeaders().getFirst(X_REQUEST_ID_HEADER));
         RequestLog.startTiming(this);
         StringBuilder debugMsgBuilder = new StringBuilder();
 
@@ -168,6 +173,7 @@ public void filter(ContainerRequestContext request, ContainerResponseContext res
      */
     @Override
     public void filter(ClientRequestContext request) throws IOException {
+        appendRequestId(request.getStringHeaders().getFirst(X_REQUEST_ID_HEADER));
         RequestLog.startTiming(CLIENT_TOTAL_TIMER);
         request.setProperty(PROPERTY_NANOS, System.nanoTime());
     }
@@ -198,6 +204,30 @@ public void filter(ClientRequestContext request, ClientResponseContext response)
         RequestLog.stopTiming(CLIENT_TOTAL_TIMER);
     }
 
+    /**
+     * Add a request id to pass to druid. Only added if x-request-id follows the format specified at:
+     * https://devcenter.heroku.com/articles/http-request-id
+     *
+     * @param requestId  Request id to add as queryId prefix to druid
+     */
+    private void appendRequestId(String requestId) {
+        if (isValidRequestId(requestId)) {
+            return; // Ignore according to https://devcenter.heroku.com/articles/http-request-id
+        }
+        RequestLog.addIdPrefix(requestId);
+    }
+
+    /**
+     * Validate whether or not a string is acceptable as an x-request-id header argument.
+     *
+     * @param requestId  Request id to validate
+     * @return True if valid, false otherwise
+     */
+    private boolean isValidRequestId(String requestId) {
+        return requestId == null || requestId.isEmpty() || requestId.length() > 200
+               || !VALID_REQUEST_ID.matcher(requestId).matches();
+    }
+
     /**
      * Render the URI as a string.
      *

fili-core/src/main/java/com/yahoo/bard/webservice/web/handlers/CacheRequestHandler.java

@@ -133,7 +133,7 @@ public boolean handleRequest(
      */
     protected String getKey(DruidAggregationQuery<?> druidQuery) throws JsonProcessingException {
         JsonNode root = mapper.valueToTree(druidQuery);
-        Utils.emitField(root, "context", mapper);
+        Utils.omitField(root, "context", mapper);
         return writer.writeValueAsString(root);
     }
 }

fili-core/src/main/java/com/yahoo/bard/webservice/web/handlers/CacheV2RequestHandler.java

@@ -156,7 +156,7 @@ public boolean handleRequest(
      */
     protected String getKey(DruidAggregationQuery<?> druidQuery) throws JsonProcessingException {
         JsonNode root = mapper.valueToTree(druidQuery);
-        Utils.emitField(root, "context", mapper);
+        Utils.omitField(root, "context", mapper);
         return writer.writeValueAsString(root);
     }
 }

fili-core/src/test/groovy/com/yahoo/bard/webservice/util/GroovyTestUtils.groovy

@@ -39,8 +39,8 @@ abstract class GroovyTestUtils extends Specification {
             JsonNode actualRoot = MAPPER.readTree(actual);
             JsonNode expectedRoot = MAPPER.readTree(expected);
 
-            Utils.emitField(actualRoot, "context", MAPPER);
-            Utils.emitField(expectedRoot, "context", MAPPER);
+            Utils.omitField(actualRoot, "context", MAPPER);
+            Utils.omitField(expectedRoot, "context", MAPPER);
 
             def actualProcessed = MAPPER.writer().withDefaultPrettyPrinter().writeValueAsString(actualRoot);
             def expectedProcessed = MAPPER.writer().withDefaultPrettyPrinter().writeValueAsString(expectedRoot);

fili-core/src/test/groovy/com/yahoo/bard/webservice/web/endpoints/BaseDataServletComponentSpec.groovy

@@ -21,6 +21,7 @@ import org.slf4j.LoggerFactory
 import spock.lang.Specification
 import spock.lang.Timeout
 
+import javax.ws.rs.core.MultivaluedHashMap
 import javax.ws.rs.core.MultivaluedMap
 import javax.ws.rs.core.Response
 
@@ -71,6 +72,13 @@ abstract class BaseDataServletComponentSpec extends Specification {
         populatePhysicalTableAvailability()
     }
 
+    /**
+     * Used to append headers to API request made by this spec.
+     */
+    MultivaluedHashMap<String, String> getAdditionalApiRequestHeaders() {
+        return [:]
+    }
+
     /**
      * Populates the interval availability of the physical tables.
      * <p>
@@ -206,7 +214,7 @@ abstract class BaseDataServletComponentSpec extends Specification {
         }
 
         // Make the call
-        Response response = httpCall.request().get()
+        Response response = httpCall.request().headers(getAdditionalApiRequestHeaders()).get()
         if (response.status != 200) {
             LOG.error( "***  *** Response status: ${response.status}: ${response.readEntity(String)}")
         }

fili-core/src/test/groovy/com/yahoo/bard/webservice/web/endpoints/RequestIdPrefixesDruidQueryIdSpec.groovy

@@ -0,0 +1,87 @@
+// Copyright 2016 Yahoo Inc.
+// Licensed under the terms of the Apache license. Please see LICENSE.md file distributed with this work for terms.
+package com.yahoo.bard.webservice.web.endpoints
+
+import com.yahoo.bard.webservice.data.dimension.BardDimensionField
+import com.yahoo.bard.webservice.data.dimension.DimensionDictionary
+import com.yahoo.bard.webservice.util.JsonSlurper
+import com.yahoo.bard.webservice.util.JsonSortStrategy
+import com.yahoo.bard.webservice.web.filters.BardLoggingFilter
+import org.apache.commons.lang.StringUtils
+
+import javax.ws.rs.core.MultivaluedHashMap
+
+
+class RequestIdPrefixesDruidQueryIdSpec extends BaseDataServletComponentSpec {
+    def prefixId = UUID.randomUUID().toString();
+
+    @Override
+    def setup() {
+        DimensionDictionary dimensionStore = jtb.configurationLoader.dimensionDictionary
+        dimensionStore.findByApiName("color").with {
+            addDimensionRow(BardDimensionField.makeDimensionRow(it, "Foo", "FooDesc"))
+            addDimensionRow(BardDimensionField.makeDimensionRow(it, "Bar", "BarDesc"))
+            addDimensionRow(BardDimensionField.makeDimensionRow(it, "Baz", "BazDesc"))
+        }
+    }
+
+    @Override
+    Class<?>[] getResourceClasses() {
+        [ DataServlet.class, BardLoggingFilter.class ]
+    }
+
+    @Override
+    String getTarget() {
+        return "data/shapes/week/color"
+    }
+
+    @Override
+    Map<String, List<String>> getQueryParams() {
+        [
+                "metrics": ["width","depth"],
+                "dateTime": ["2014-06-02/2014-06-30"],
+                "sort": ["width|desc","depth|asc"]
+        ]
+    }
+
+    @Override
+    boolean compareResult(String result, String expectedResult, JsonSortStrategy sortStrategy = JsonSortStrategy.SORT_MAPS) {
+        def parsedJson = new JsonSlurper(sortStrategy).parseText(result)
+        if (parsedJson.context != null) {
+            return parsedJson.context.queryId.startsWith(prefixId)
+        }
+        // Default to true-- the base spec runs an extra test which we consider a noop.
+        return true
+    }
+
+    @Override
+    String getExpectedDruidQuery() {
+        // Noop
+        """{}"""
+    }
+
+    @Override
+    String getFakeDruidResponse() {
+        // Noop
+        """[]"""
+    }
+
+    @Override
+    String getExpectedApiResponse() {
+        // Noop
+        """{}"""
+    }
+
+    @Override
+    MultivaluedHashMap<String, String> getAdditionalApiRequestHeaders() {
+        return ["x-request-id": prefixId]
+    }
+
+    def "verify invalid x-request-id values"() {
+        BardLoggingFilter filter = new BardLoggingFilter()
+        assert !filter.isValidRequestId('abcd$') // Invalid char
+        assert !filter.isValidRequestId(StringUtils.leftPad('a', 200, 'a')) // Too long
+        assert !filter.isValidRequestId('') // empty string
+        assert !filter.isValidRequestId(null) // null
+    }
+}