From 60caa641f34ef4e9065f1efb868a62916b8a3454 Mon Sep 17 00:00:00 2001
From: Gian Merlino <gianmerlino@gmail.com>
Date: Thu, 20 Apr 2017 01:47:50 +0900
Subject: [PATCH] Restore backwards compatibility of Query. (#4185)

---
 processing/src/main/java/io/druid/query/BaseQuery.java      | 6 ------
 processing/src/main/java/io/druid/query/Query.java          | 2 --
 processing/src/main/java/io/druid/query/QueryContexts.java  | 6 ++++++
 .../src/test/java/io/druid/query/QueryContextsTest.java     | 4 ++--
 server/src/main/java/io/druid/server/QueryResource.java     | 3 ++-
 5 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/processing/src/main/java/io/druid/query/BaseQuery.java b/processing/src/main/java/io/druid/query/BaseQuery.java
index edb1ca5bf32..0186e00075a 100644
--- a/processing/src/main/java/io/druid/query/BaseQuery.java
+++ b/processing/src/main/java/io/druid/query/BaseQuery.java
@@ -176,12 +176,6 @@ public abstract class BaseQuery<T extends Comparable<T>> implements Query<T>
     return withOverriddenContext(ImmutableMap.<String, Object>of(QUERYID, id));
   }
 
-  @Override
-  public Query<T> withDefaultTimeout(long defaultTimeout)
-  {
-    return withOverriddenContext(ImmutableMap.of(QueryContexts.DEFAULT_TIMEOUT_KEY, defaultTimeout));
-  }
-
   @Override
   public boolean equals(Object o)
   {
diff --git a/processing/src/main/java/io/druid/query/Query.java b/processing/src/main/java/io/druid/query/Query.java
index cfbf6f2d340..9ad178161ea 100644
--- a/processing/src/main/java/io/druid/query/Query.java
+++ b/processing/src/main/java/io/druid/query/Query.java
@@ -99,6 +99,4 @@ public interface Query<T>
   String getId();
 
   Query<T> withDataSource(DataSource dataSource);
-
-  Query<T> withDefaultTimeout(long defaultTimeout);
 }
diff --git a/processing/src/main/java/io/druid/query/QueryContexts.java b/processing/src/main/java/io/druid/query/QueryContexts.java
index b59c6bc2a2b..5a0409e430d 100644
--- a/processing/src/main/java/io/druid/query/QueryContexts.java
+++ b/processing/src/main/java/io/druid/query/QueryContexts.java
@@ -20,6 +20,7 @@
 package io.druid.query;
 
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
 import io.druid.java.util.common.ISE;
 
 public class QueryContexts
@@ -114,6 +115,11 @@ public class QueryContexts
     return timeout;
   }
 
+  public static <T> Query<T> withDefaultTimeout(Query<T> query, long defaultTimeout)
+  {
+    return query.withOverriddenContext(ImmutableMap.of(QueryContexts.DEFAULT_TIMEOUT_KEY, defaultTimeout));
+  }
+
   static <T> long getDefaultTimeout(Query<T> query)
   {
     final long defaultTimeout = parseLong(query, DEFAULT_TIMEOUT_KEY, DEFAULT_TIMEOUT_MILLIS);
diff --git a/processing/src/test/java/io/druid/query/QueryContextsTest.java b/processing/src/test/java/io/druid/query/QueryContextsTest.java
index c656f077110..79a7ebf992d 100644
--- a/processing/src/test/java/io/druid/query/QueryContextsTest.java
+++ b/processing/src/test/java/io/druid/query/QueryContextsTest.java
@@ -106,7 +106,7 @@ public class QueryContextsTest
     );
     Assert.assertEquals(300_000, QueryContexts.getTimeout(query));
 
-    query = query.withDefaultTimeout(60_000);
+    query = QueryContexts.withDefaultTimeout(query, 60_000);
     Assert.assertEquals(60_000, QueryContexts.getTimeout(query));
   }
 
@@ -121,7 +121,7 @@ public class QueryContextsTest
     );
     Assert.assertEquals(1000, QueryContexts.getTimeout(query));
 
-    query = query.withDefaultTimeout(1_000_000);
+    query = QueryContexts.withDefaultTimeout(query, 1_000_000);
     Assert.assertEquals(1000, QueryContexts.getTimeout(query));
   }
 }
diff --git a/server/src/main/java/io/druid/server/QueryResource.java b/server/src/main/java/io/druid/server/QueryResource.java
index 554cc34de83..45fcd98dff1 100644
--- a/server/src/main/java/io/druid/server/QueryResource.java
+++ b/server/src/main/java/io/druid/server/QueryResource.java
@@ -40,6 +40,7 @@ import io.druid.java.util.common.guava.Yielders;
 import io.druid.query.DruidMetrics;
 import io.druid.query.GenericQueryMetricsFactory;
 import io.druid.query.Query;
+import io.druid.query.QueryContexts;
 import io.druid.query.QueryInterruptedException;
 import io.druid.query.QueryMetrics;
 import io.druid.query.QuerySegmentWalker;
@@ -190,7 +191,7 @@ public class QueryResource implements QueryCountStatsProvider
         queryId = UUID.randomUUID().toString();
         query = query.withId(queryId);
       }
-      query = query.withDefaultTimeout(config.getDefaultQueryTimeout());
+      query = QueryContexts.withDefaultTimeout(query, config.getDefaultQueryTimeout());
 
       toolChest = warehouse.getToolChest(query);