fix header so it passes the entire context

jon-wei · Jun 19, 2014 · 8e85097 · 8e85097
1 parent 35e080b
commit 8e85097
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 14 deletions.
diff --git a/server/src/main/java/io/druid/client/DirectDruidClient.java b/server/src/main/java/io/druid/client/DirectDruidClient.java
@@ -54,6 +54,7 @@
 import io.druid.query.QueryToolChestWarehouse;
 import io.druid.query.QueryWatcher;
 import io.druid.query.Result;
+import io.druid.query.RetryQueryRunner;
 import io.druid.query.SegmentDescriptor;
 import io.druid.query.aggregation.MetricManipulatorFns;
 import org.jboss.netty.handler.codec.http.HttpChunk;
@@ -64,8 +65,8 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URL;
+import java.util.ArrayList;
 import java.util.Iterator;
-import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CancellationException;
@@ -160,18 +161,24 @@ public ClientResponse<AppendableByteArrayInputStream> handleResponse(HttpRespons
                   startTime = System.currentTimeMillis();
                   byteCount += response.getContent().readableBytes();
 
-                  if (!response.getHeader("Missing-Segments").equals("")) {
-                    LinkedList missingSegments = new LinkedList();
-                    try {
-                      missingSegments = objectMapper.readValue(response.getHeader("Missing-Segments"), LinkedList.class);
-                      for (int i = missingSegments.size(); i > 0; i--) {
-                        missingSegments.add(objectMapper.convertValue(missingSegments.remove(0), SegmentDescriptor.class));
-                      }
-                    }
-                    catch (IOException e) {
+
+                  List missingSegments = new ArrayList();
+                  try {
+                    Map<String, Object> headerContext = objectMapper.readValue(response.getHeader("Context"), Map.class);
+                    missingSegments = (List)headerContext.get(RetryQueryRunner.missingSegments);
+                    for (int i = missingSegments.size(); i > 0; i--) {
+                      missingSegments.add(
+                          objectMapper.convertValue(
+                              missingSegments.remove(0),
+                              SegmentDescriptor.class
+                          )
+                      );
                     }
-                    ((List) context.get("missingSegments")).addAll(missingSegments);
                   }
+                  catch (IOException e) {
+                    e.printStackTrace();
+                  }
+                  ((List) context.get("missingSegments")).addAll(missingSegments);
 
                   return super.handleResponse(response);
                 }

diff --git a/server/src/main/java/io/druid/server/QueryResource.java b/server/src/main/java/io/druid/server/QueryResource.java
@@ -166,9 +166,9 @@ public Object accumulate(Object accumulated, Object in)
           )
       ) {
 
-        String missingSegments = "";
+        String headerContext = "";
         if (!((List)context.get("missingSegments")).isEmpty()) {
-          missingSegments = jsonMapper.writeValueAsString(context.get("missingSegments"));
+          headerContext = jsonMapper.writeValueAsString(context);
         }
 
         long requestTime = System.currentTimeMillis() - start;
@@ -213,7 +213,7 @@ public void write(OutputStream outputStream) throws IOException, WebApplicationE
                 isSmile ? APPLICATION_JSON : APPLICATION_SMILE
             )
             .header("X-Druid-Query-Id", queryId)
-            .header("Missing-Segments", missingSegments)
+            .header("Context", headerContext)
             .build();
       }
     }