Inject QueryMetrics factories via PolyBind (#4336)

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

@@ -22,9 +22,11 @@ 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.LazySingleton;
 import io.druid.guice.annotations.Json;
 import io.druid.jackson.DefaultObjectMapper;
 
+@LazySingleton
 public class DefaultGenericQueryMetricsFactory implements GenericQueryMetricsFactory
 {
   private static final GenericQueryMetricsFactory INSTANCE =

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

@@ -1,38 +0,0 @@
-/*
- * 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

@@ -23,7 +23,15 @@ 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.
+ * Implementations could be injected using
+ *
+ * PolyBind
+ *    .optionBinder(binder, Key.get(GenericQueryMetricsFactory.class))
+ *    .addBinding("myCustomGenericQueryMetricsFactory")
+ *    .to(MyCustomGenericQueryMetricsFactory.class);
+ *
+ * And then setting property:
+ * druid.query.generic.queryMetricsFactory=myCustomGenericQueryMetricsFactory
  */
 public interface GenericQueryMetricsFactory
 {

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

@@ -22,9 +22,11 @@ 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.LazySingleton;
 import io.druid.guice.annotations.Json;
 import io.druid.jackson.DefaultObjectMapper;
 
+@LazySingleton
 public class DefaultGroupByQueryMetricsFactory implements GroupByQueryMetricsFactory
 {
   private static final GroupByQueryMetricsFactory INSTANCE =

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

@@ -66,9 +66,6 @@ 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;
@@ -129,16 +126,6 @@ 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

@@ -20,7 +20,15 @@
 package io.druid.query.groupby;
 
 /**
- * Implementations of this interface could be injected using {@link GroupByQueryConfig#queryMetricsFactory} option.
+ * Implementations could be injected using
+ *
+ * PolyBind
+ *    .optionBinder(binder, Key.get(GroupByQueryMetricsFactory.class))
+ *    .addBinding("myCustomGroupByQueryMetricsFactory")
+ *    .to(MyCustomGroupByQueryMetricsFactory.class);
+ *
+ * And then setting property:
+ * druid.query.groupBy.queryMetricsFactory=myCustomGroupByQueryMetricsFactory
  */
 public interface GroupByQueryMetricsFactory
 {

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

@@ -22,9 +22,11 @@ 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.LazySingleton;
 import io.druid.guice.annotations.Json;
 import io.druid.jackson.DefaultObjectMapper;
 
+@LazySingleton
 public class DefaultTimeseriesQueryMetricsFactory implements TimeseriesQueryMetricsFactory
 {
   private static final TimeseriesQueryMetricsFactory INSTANCE =

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

@@ -1,38 +0,0 @@
-/*
- * 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

@@ -20,7 +20,15 @@
 package io.druid.query.timeseries;
 
 /**
- * Implementations of this interface could be injected using {@link TimeseriesQueryConfig#queryMetricsFactory} option.
+ * Implementations could be injected using
+ *
+ * PolyBind
+ *    .optionBinder(binder, Key.get(TimeseriesQueryMetricsFactory.class))
+ *    .addBinding("myCustomTimeseriesQueryMetricsFactory")
+ *    .to(MyCustomTimeseriesQueryMetricsFactory.class);
+ *
+ * And then setting property:
+ * druid.query.timeseries.queryMetricsFactory=myCustomTimeseriesQueryMetricsFactory
  */
 public interface TimeseriesQueryMetricsFactory
 {

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

@@ -22,9 +22,11 @@ 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.LazySingleton;
 import io.druid.guice.annotations.Json;
 import io.druid.jackson.DefaultObjectMapper;
 
+@LazySingleton
 public class DefaultTopNQueryMetricsFactory implements TopNQueryMetricsFactory
 {
   private static final TopNQueryMetricsFactory INSTANCE = new DefaultTopNQueryMetricsFactory(new DefaultObjectMapper());

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

@@ -31,21 +31,8 @@ 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

@@ -20,7 +20,15 @@
 package io.druid.query.topn;
 
 /**
- * Implementations of this interface could be injected using {@link TopNQueryConfig#queryMetricsFactory} option.
+ * Implementations could be injected using
+ *
+ * PolyBind
+ *    .optionBinder(binder, Key.get(TopNQueryMetricsFactory.class))
+ *    .addBinding("myCustomTopNQueryMetricsFactory")
+ *    .to(MyCustomTopNQueryMetricsFactory.class);
+ *
+ * And then setting property:
+ * druid.query.topN.queryMetricsFactory=myCustomTopNQueryMetricsFactory
  */
 public interface TopNQueryMetricsFactory
 {

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

@@ -21,11 +21,10 @@ package io.druid.guice;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.inject.Binder;
-import com.google.inject.Injector;
+import com.google.inject.Key;
 import com.google.inject.Module;
-import com.google.inject.Provides;
 import com.google.inject.multibindings.MapBinder;
-import io.druid.query.GenericQueryConfig;
+import io.druid.query.DefaultGenericQueryMetricsFactory;
 import io.druid.query.GenericQueryMetricsFactory;
 import io.druid.query.MapQueryToolChestWarehouse;
 import io.druid.query.Query;
@@ -33,6 +32,7 @@ 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;
@@ -48,10 +48,11 @@ 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;
@@ -63,6 +64,11 @@ import java.util.Map;
  */
 public class QueryToolChestModule implements Module
 {
+  public static final String GENERIC_QUERY_METRICS_FACTORY_PROPERTY = "druid.query.generic.queryMetricsFactory";
+  public static final String GROUPBY_QUERY_METRICS_FACTORY_PROPERTY = "druid.query.groupBy.queryMetricsFactory";
+  public static final String TIMESERIES_QUERY_METRICS_FACTORY_PROPERTY = "druid.query.timeseries.queryMetricsFactory";
+  public static final String TOPN_QUERY_METRICS_FACTORY_PROPERTY = "druid.query.topN.queryMetricsFactory";
+
   public final Map<Class<? extends Query>, Class<? extends QueryToolChest>> mappings =
       ImmutableMap.<Class<? extends Query>, Class<? extends QueryToolChest>>builder()
                   .put(TimeseriesQuery.class, TimeseriesQueryQueryToolChest.class)
@@ -87,36 +93,54 @@ public class QueryToolChestModule implements Module
 
     binder.bind(QueryToolChestWarehouse.class).to(MapQueryToolChestWarehouse.class);
 
-    JsonConfigProvider.bind(binder, "druid.query.generic", GenericQueryConfig.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
+    PolyBind.createChoice(
-  public GenericQueryMetricsFactory getGenericQueryMetricsFactory(Injector injector, GenericQueryConfig config)
+        binder,
-  {
+        GENERIC_QUERY_METRICS_FACTORY_PROPERTY,
-    return injector.getInstance(config.getQueryMetricsFactory());
+        Key.get(GenericQueryMetricsFactory.class),
-  }
+        Key.get(DefaultGenericQueryMetricsFactory.class)
+    );
+    PolyBind
+        .optionBinder(binder, Key.get(GenericQueryMetricsFactory.class))
+        .addBinding("default")
+        .to(DefaultGenericQueryMetricsFactory.class);
 
-  @Provides
+    PolyBind.createChoice(
-  public GroupByQueryMetricsFactory getGroupByQueryMetricsFactory(Injector injector, GroupByQueryConfig config)
+        binder,
-  {
+        GROUPBY_QUERY_METRICS_FACTORY_PROPERTY,
-    return injector.getInstance(config.getQueryMetricsFactory());
+        Key.get(GroupByQueryMetricsFactory.class),
-  }
+        Key.get(DefaultGroupByQueryMetricsFactory.class)
+    );
+    PolyBind
+        .optionBinder(binder, Key.get(GroupByQueryMetricsFactory.class))
+        .addBinding("default")
+        .to(DefaultGroupByQueryMetricsFactory.class);
 
-  @Provides
+    PolyBind.createChoice(
-  public TimeseriesQueryMetricsFactory getTimeseriesQueryMetricsFactory(Injector injector, TimeseriesQueryConfig config)
+        binder,
-  {
+        TIMESERIES_QUERY_METRICS_FACTORY_PROPERTY,
-    return injector.getInstance(config.getQueryMetricsFactory());
+        Key.get(TimeseriesQueryMetricsFactory.class),
-  }
+        Key.get(DefaultTimeseriesQueryMetricsFactory.class)
+    );
+    PolyBind
+        .optionBinder(binder, Key.get(TimeseriesQueryMetricsFactory.class))
+        .addBinding("default")
+        .to(DefaultTimeseriesQueryMetricsFactory.class);
 
-  @Provides
+    PolyBind.createChoice(
-  public TopNQueryMetricsFactory getTopNQueryMetricsFactory(Injector injector, TopNQueryConfig config)
+        binder,
-  {
+        TOPN_QUERY_METRICS_FACTORY_PROPERTY,
-    return injector.getInstance(config.getQueryMetricsFactory());
+        Key.get(TopNQueryMetricsFactory.class),
+        Key.get(DefaultTopNQueryMetricsFactory.class)
+    );
+    PolyBind
+        .optionBinder(binder, Key.get(TopNQueryMetricsFactory.class))
+        .addBinding("default")
+        .to(DefaultTopNQueryMetricsFactory.class);
   }
 }