From 8e85097999f5b35a480d99d288490f6fb6578f2e Mon Sep 17 00:00:00 2001
From: jisookim0513 <jisookim@Jisoos-MacBook-Pro.local>
Date: Thu, 19 Jun 2014 14:38:11 -0700
Subject: [PATCH] fix header so it passes the entire context

---
 .../io/druid/client/DirectDruidClient.java    | 29 ++++++++++++-------
 .../java/io/druid/server/QueryResource.java   |  6 ++--
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/server/src/main/java/io/druid/client/DirectDruidClient.java b/server/src/main/java/io/druid/client/DirectDruidClient.java
index 541ae1df62e..e69cea4e417 100644
--- 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.QueryToolChest;
 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.Closeable;
 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 class DirectDruidClient<T> implements QueryRunner<T>
                   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));
-                      }
+
+                  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
+                          )
+                      );
                     }
-                    catch (IOException e) {
-                    }
-                    ((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
index 353e7014e06..ce6e255cd0d 100644
--- 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 class QueryResource
           )
       ) {
 
-        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 class QueryResource
                 isSmile ? APPLICATION_JSON : APPLICATION_SMILE
             )
             .header("X-Druid-Query-Id", queryId)
-            .header("Missing-Segments", missingSegments)
+            .header("Context", headerContext)
             .build();
       }
     }