Make QueryMetrics factories configurable (#4268)

* Ensure QueryMetrics factories accept Json ObjectMapper; Make QueryMetrics factories configurable

* Update QueryMetrics Javadocs

* Add javadocs to QueryMetrics factories

* Move queryMetricsFactory defaults to getter methods of config classes

processing/src/main/java/io/druid/query/DefaultGenericQueryMetricsFactory.java

@@ -22,6 +22,7 @@ package io.druid.query;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.inject.Inject;
+import io.druid.guice.annotations.Json;
 import io.druid.jackson.DefaultObjectMapper;
 
 public class DefaultGenericQueryMetricsFactory implements GenericQueryMetricsFactory
@@ -43,7 +44,7 @@ public class DefaultGenericQueryMetricsFactory implements GenericQueryMetricsFac
   private final ObjectMapper jsonMapper;
 
   @Inject
-  public DefaultGenericQueryMetricsFactory(ObjectMapper jsonMapper)
+  public DefaultGenericQueryMetricsFactory(@Json ObjectMapper jsonMapper)
   {
     this.jsonMapper = jsonMapper;
   }

processing/src/main/java/io/druid/query/GenericQueryConfig.java

@@ -0,0 +1,38 @@
+/*
+ * Licensed to Metamarkets Group Inc. (Metamarkets) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. Metamarkets licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package io.druid.query;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public class GenericQueryConfig
+{
+  @JsonProperty
+  private Class<? extends GenericQueryMetricsFactory> queryMetricsFactory;
+
+  public Class<? extends GenericQueryMetricsFactory> getQueryMetricsFactory()
+  {
+    return queryMetricsFactory != null ? queryMetricsFactory : DefaultGenericQueryMetricsFactory.class;
+  }
+
+  public void setQueryMetricsFactory(Class<? extends GenericQueryMetricsFactory> queryMetricsFactory)
+  {
+    this.queryMetricsFactory = queryMetricsFactory;
+  }
+}

processing/src/main/java/io/druid/query/GenericQueryMetricsFactory.java

@@ -22,6 +22,8 @@ package io.druid.query;
 /**
  * This factory is used for DI of custom {@link QueryMetrics} implementations for all query types, which don't (yet)
  * need to emit custom dimensions and/or metrics, i. e. they are good with the generic {@link QueryMetrics} interface.
+ *
+ * Implementations could be injected using {@link GenericQueryConfig#queryMetricsFactory} option.
  */
 public interface GenericQueryMetricsFactory
 {

processing/src/main/java/io/druid/query/QueryMetrics.java

@@ -122,9 +122,13 @@ import org.joda.time.Interval;
  *  5. Inject and use SearchQueryMetricsFactory instead of {@link GenericQueryMetricsFactory} in {@link
  *  io.druid.query.search.SearchQueryQueryToolChest}.
  *
- *  6. Specify `binder.bind(SearchQueryMetricsFactory.class).to(DefaultSearchQueryMetricsFactory.class)` in
+ *  6. Establish injection of SearchQueryMetricsFactory using config and provider method in QueryToolChestModule
- *  QueryToolChestModule (if the query type belongs to the core druid-processing, e. g. SearchQuery) or in some
+ *  (see how it is done in QueryToolChestModule for existing query types with custom metrics, e. g. {@link
- *  extension-specific Guice module otherwise, if the query type is defined in an extension, e. g. ScanQuery.
+ *  io.druid.query.topn.TopNQueryMetricsFactory}), if the query type belongs to the core druid-processing, e. g.
+ *  SearchQuery. If the query type defined in an extension, you can specify
+ *  `binder.bind(ScanQueryMetricsFactory.class).to(DefaultScanQueryMetricsFactory.class)` in the extension's
+ *  Guice module, if the query type is defined in an extension, e. g. ScanQuery. Or establish similar configuration,
+ *  as for the core query types.
  *
  * This complex procedure is needed to ensure custom {@link GenericQueryMetricsFactory} specified by users still works
  * for the query type when query type decides to create their custom QueryMetrics subclass.

processing/src/main/java/io/druid/query/groupby/DefaultGroupByQueryMetricsFactory.java

@@ -22,6 +22,7 @@ package io.druid.query.groupby;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.inject.Inject;
+import io.druid.guice.annotations.Json;
 import io.druid.jackson.DefaultObjectMapper;
 
 public class DefaultGroupByQueryMetricsFactory implements GroupByQueryMetricsFactory
@@ -43,7 +44,7 @@ public class DefaultGroupByQueryMetricsFactory implements GroupByQueryMetricsFac
   private final ObjectMapper jsonMapper;
 
   @Inject
-  public DefaultGroupByQueryMetricsFactory(ObjectMapper jsonMapper)
+  public DefaultGroupByQueryMetricsFactory(@Json ObjectMapper jsonMapper)
   {
     this.jsonMapper = jsonMapper;
   }

processing/src/main/java/io/druid/query/groupby/GroupByQueryConfig.java

@@ -66,6 +66,9 @@ public class GroupByQueryConfig
   // Max on-disk temporary storage, per-query; when exceeded, the query fails
   private long maxOnDiskStorage = 0L;
 
+  @JsonProperty
+  private Class<? extends GroupByQueryMetricsFactory> queryMetricsFactory;
+
   public String getDefaultStrategy()
   {
     return defaultStrategy;
@@ -126,6 +129,16 @@ public class GroupByQueryConfig
     return maxOnDiskStorage;
   }
 
+  public Class<? extends GroupByQueryMetricsFactory> getQueryMetricsFactory()
+  {
+    return queryMetricsFactory != null ? queryMetricsFactory : DefaultGroupByQueryMetricsFactory.class;
+  }
+
+  public void setQueryMetricsFactory(Class<? extends GroupByQueryMetricsFactory> queryMetricsFactory)
+  {
+    this.queryMetricsFactory = queryMetricsFactory;
+  }
+
   public GroupByQueryConfig withOverrides(final GroupByQuery query)
   {
     final GroupByQueryConfig newConfig = new GroupByQueryConfig();

processing/src/main/java/io/druid/query/groupby/GroupByQueryMetricsFactory.java

@@ -19,6 +19,9 @@
 
 package io.druid.query.groupby;
 
+/**
+ * Implementations of this interface could be injected using {@link GroupByQueryConfig#queryMetricsFactory} option.
+ */
 public interface GroupByQueryMetricsFactory
 {
 

processing/src/main/java/io/druid/query/timeseries/DefaultTimeseriesQueryMetricsFactory.java

@@ -22,6 +22,7 @@ package io.druid.query.timeseries;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.inject.Inject;
+import io.druid.guice.annotations.Json;
 import io.druid.jackson.DefaultObjectMapper;
 
 public class DefaultTimeseriesQueryMetricsFactory implements TimeseriesQueryMetricsFactory
@@ -42,7 +43,7 @@ public class DefaultTimeseriesQueryMetricsFactory implements TimeseriesQueryMetr
   private final ObjectMapper jsonMapper;
 
   @Inject
-  public DefaultTimeseriesQueryMetricsFactory(ObjectMapper jsonMapper)
+  public DefaultTimeseriesQueryMetricsFactory(@Json ObjectMapper jsonMapper)
   {
     this.jsonMapper = jsonMapper;
   }

processing/src/main/java/io/druid/query/timeseries/TimeseriesQueryConfig.java

@@ -0,0 +1,38 @@
+/*
+ * Licensed to Metamarkets Group Inc. (Metamarkets) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. Metamarkets licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package io.druid.query.timeseries;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public class TimeseriesQueryConfig
+{
+  @JsonProperty
+  private Class<? extends TimeseriesQueryMetricsFactory> queryMetricsFactory;
+
+  public Class<? extends TimeseriesQueryMetricsFactory> getQueryMetricsFactory()
+  {
+    return queryMetricsFactory != null ? queryMetricsFactory : DefaultTimeseriesQueryMetricsFactory.class;
+  }
+
+  public void setQueryMetricsFactory(Class<? extends TimeseriesQueryMetricsFactory> queryMetricsFactory)
+  {
+    this.queryMetricsFactory = queryMetricsFactory;
+  }
+}

processing/src/main/java/io/druid/query/timeseries/TimeseriesQueryMetricsFactory.java

@@ -19,6 +19,9 @@
 
 package io.druid.query.timeseries;
 
+/**
+ * Implementations of this interface could be injected using {@link TimeseriesQueryConfig#queryMetricsFactory} option.
+ */
 public interface TimeseriesQueryMetricsFactory
 {
 

processing/src/main/java/io/druid/query/topn/DefaultTopNQueryMetricsFactory.java

@@ -22,6 +22,7 @@ package io.druid.query.topn;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.inject.Inject;
+import io.druid.guice.annotations.Json;
 import io.druid.jackson.DefaultObjectMapper;
 
 public class DefaultTopNQueryMetricsFactory implements TopNQueryMetricsFactory
@@ -41,7 +42,7 @@ public class DefaultTopNQueryMetricsFactory implements TopNQueryMetricsFactory
   private final ObjectMapper jsonMapper;
 
   @Inject
-  public DefaultTopNQueryMetricsFactory(ObjectMapper jsonMapper)
+  public DefaultTopNQueryMetricsFactory(@Json ObjectMapper jsonMapper)
   {
     this.jsonMapper = jsonMapper;
   }

processing/src/main/java/io/druid/query/topn/TopNQueryConfig.java

@@ -31,8 +31,21 @@ public class TopNQueryConfig
   @Min(1)
   private int minTopNThreshold = 1000;
 
+  @JsonProperty
+  private Class<? extends TopNQueryMetricsFactory> queryMetricsFactory;
+
   public int getMinTopNThreshold()
   {
     return minTopNThreshold;
   }
+
+  public Class<? extends TopNQueryMetricsFactory> getQueryMetricsFactory()
+  {
+    return queryMetricsFactory != null ? queryMetricsFactory : DefaultTopNQueryMetricsFactory.class;
+  }
+
+  public void setQueryMetricsFactory(Class<? extends TopNQueryMetricsFactory> queryMetricsFactory)
+  {
+    this.queryMetricsFactory = queryMetricsFactory;
+  }
 }

processing/src/main/java/io/druid/query/topn/TopNQueryMetricsFactory.java

@@ -19,6 +19,9 @@
 
 package io.druid.query.topn;
 
+/**
+ * Implementations of this interface could be injected using {@link TopNQueryConfig#queryMetricsFactory} option.
+ */
 public interface TopNQueryMetricsFactory
 {
 

server/src/main/java/io/druid/guice/QueryToolChestModule.java

@@ -21,17 +21,18 @@ package io.druid.guice;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.inject.Binder;
+import com.google.inject.Injector;
 import com.google.inject.Module;
+import com.google.inject.Provides;
 import com.google.inject.multibindings.MapBinder;
-import io.druid.query.DefaultGenericQueryMetricsFactory;
+import io.druid.query.GenericQueryConfig;
+import io.druid.query.GenericQueryMetricsFactory;
 import io.druid.query.MapQueryToolChestWarehouse;
 import io.druid.query.Query;
-import io.druid.query.GenericQueryMetricsFactory;
 import io.druid.query.QueryToolChest;
 import io.druid.query.QueryToolChestWarehouse;
 import io.druid.query.datasourcemetadata.DataSourceMetadataQuery;
 import io.druid.query.datasourcemetadata.DataSourceQueryQueryToolChest;
-import io.druid.query.groupby.DefaultGroupByQueryMetricsFactory;
 import io.druid.query.groupby.GroupByQuery;
 import io.druid.query.groupby.GroupByQueryConfig;
 import io.druid.query.groupby.GroupByQueryMetricsFactory;
@@ -47,11 +48,10 @@ import io.druid.query.select.SelectQueryConfig;
 import io.druid.query.select.SelectQueryQueryToolChest;
 import io.druid.query.timeboundary.TimeBoundaryQuery;
 import io.druid.query.timeboundary.TimeBoundaryQueryQueryToolChest;
-import io.druid.query.timeseries.DefaultTimeseriesQueryMetricsFactory;
 import io.druid.query.timeseries.TimeseriesQuery;
+import io.druid.query.timeseries.TimeseriesQueryConfig;
 import io.druid.query.timeseries.TimeseriesQueryMetricsFactory;
 import io.druid.query.timeseries.TimeseriesQueryQueryToolChest;
-import io.druid.query.topn.DefaultTopNQueryMetricsFactory;
 import io.druid.query.topn.TopNQuery;
 import io.druid.query.topn.TopNQueryConfig;
 import io.druid.query.topn.TopNQueryMetricsFactory;
@@ -87,15 +87,36 @@ public class QueryToolChestModule implements Module
 
     binder.bind(QueryToolChestWarehouse.class).to(MapQueryToolChestWarehouse.class);
 
-    binder.bind(GenericQueryMetricsFactory.class).to(DefaultGenericQueryMetricsFactory.class);
+    JsonConfigProvider.bind(binder, "druid.query.generic", GenericQueryConfig.class);
-    binder.bind(TopNQueryMetricsFactory.class).to(DefaultTopNQueryMetricsFactory.class);
-    binder.bind(GroupByQueryMetricsFactory.class).to(DefaultGroupByQueryMetricsFactory.class);
-    binder.bind(TimeseriesQueryMetricsFactory.class).to(DefaultTimeseriesQueryMetricsFactory.class);
-
     JsonConfigProvider.bind(binder, "druid.query.groupBy", GroupByQueryConfig.class);
     JsonConfigProvider.bind(binder, "druid.query.search", SearchQueryConfig.class);
+    JsonConfigProvider.bind(binder, "druid.query.timeseries", TimeseriesQueryConfig.class);
     JsonConfigProvider.bind(binder, "druid.query.topN", TopNQueryConfig.class);
     JsonConfigProvider.bind(binder, "druid.query.segmentMetadata", SegmentMetadataQueryConfig.class);
     JsonConfigProvider.bind(binder, "druid.query.select", SelectQueryConfig.class);
   }
+
+  @Provides
+  public GenericQueryMetricsFactory getGenericQueryMetricsFactory(Injector injector, GenericQueryConfig config)
+  {
+    return injector.getInstance(config.getQueryMetricsFactory());
+  }
+
+  @Provides
+  public GroupByQueryMetricsFactory getGroupByQueryMetricsFactory(Injector injector, GroupByQueryConfig config)
+  {
+    return injector.getInstance(config.getQueryMetricsFactory());
+  }
+
+  @Provides
+  public TimeseriesQueryMetricsFactory getTimeseriesQueryMetricsFactory(Injector injector, TimeseriesQueryConfig config)
+  {
+    return injector.getInstance(config.getQueryMetricsFactory());
+  }
+
+  @Provides
+  public TopNQueryMetricsFactory getTopNQueryMetricsFactory(Injector injector, TopNQueryConfig config)
+  {
+    return injector.getInstance(config.getQueryMetricsFactory());
+  }
 }