/*
 * Druid - a distributed column store.
 * Copyright 2012 - 2015 Metamarkets Group Inc.
 *
 * Licensed 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.server;

import io.druid.client.CachingClusteredClient;
import io.druid.query.FinalizeResultsQueryRunner;
import io.druid.query.PostProcessingOperator;
import io.druid.query.Query;
import io.druid.query.QueryRunner;
import io.druid.query.QuerySegmentWalker;
import io.druid.query.QueryToolChest;
import io.druid.query.QueryToolChestWarehouse;
import io.druid.query.RetryQueryRunner;
import io.druid.query.RetryQueryRunnerConfig;
import io.druid.query.SegmentDescriptor;

import io.druid.query.UnionQueryRunner;
import java.util.Map;

import org.joda.time.Interval;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.inject.Inject;
import com.metamx.emitter.service.ServiceEmitter;

/**
 */
public class ClientQuerySegmentWalker implements QuerySegmentWalker
{
  private final ServiceEmitter emitter;
  private final CachingClusteredClient baseClient;
  private final QueryToolChestWarehouse warehouse;
  private final RetryQueryRunnerConfig retryConfig;
  private final ObjectMapper objectMapper;

  @Inject
  public ClientQuerySegmentWalker(
      ServiceEmitter emitter,
      CachingClusteredClient baseClient,
      QueryToolChestWarehouse warehouse,
      RetryQueryRunnerConfig retryConfig,
      ObjectMapper objectMapper
  )
  {
    this.emitter = emitter;
    this.baseClient = baseClient;
    this.warehouse = warehouse;
    this.retryConfig = retryConfig;
    this.objectMapper = objectMapper;
  }

  @Override
  public <T> QueryRunner<T> getQueryRunnerForIntervals(Query<T> query, Iterable<Interval> intervals)
  {
    return makeRunner(query);
  }

  @Override
  public <T> QueryRunner<T> getQueryRunnerForSegments(Query<T> query, Iterable<SegmentDescriptor> specs)
  {
    return makeRunner(query);
  }

  private <T> QueryRunner<T> makeRunner(final Query<T> query)
  {
    final QueryToolChest<T, Query<T>> toolChest = warehouse.getToolChest(query);
    final FinalizeResultsQueryRunner<T> baseRunner = new FinalizeResultsQueryRunner<T>(
        toolChest.postMergeQueryDecoration(
            toolChest.mergeResults(
                new UnionQueryRunner<T>(
                    toolChest.preMergeQueryDecoration(
                        new RetryQueryRunner<T>(
                            baseClient,
                            toolChest,
                            retryConfig,
                            objectMapper
                        )
                    ),
                    toolChest
                )
            )
        ),
        toolChest
    );

    final Map<String, Object> context = query.getContext();
    PostProcessingOperator<T> postProcessing = null;
    if (context != null) {
      postProcessing = objectMapper.convertValue(
          context.get("postProcessing"),
          new TypeReference<PostProcessingOperator<T>>()
          {
          }
      );
    }

    return postProcessing != null ?
           postProcessing.postProcess(baseRunner) : baseRunner;
  }
}

@@ -683,46 +683,377 @@
 ort 
-io.druid.client.CachingClusteredClient
+com.fasterxml.jackson.core.type.TypeReference;%0Aimport com.fasterxml.jackson.databind.ObjectMapper;%0Aimport com.google.common.base.Function;%0Aimport com.google.inject.Inject;%0Aimport com.metamx.emitter.service.ServiceEmitter;%0Aimport com.metamx.emitter.service.ServiceMetricEvent;%0Aimport io.druid.client.CachingClusteredClient;%0Aimport io.druid.query.CPUTimeMetricQueryRunner
 ;%0Aim
@@ -1464,17 +1464,16 @@
 riptor;%0A
-%0A
 import i
@@ -1508,31 +1508,8 @@
 er;%0A
-import java.util.Map;%0A%0A
 impo
@@ -1547,188 +1547,102 @@
 ort 
-com.fasterxml.jackson.core.type.TypeReference;%0Aimport com.fasterxml.jackson.databind.ObjectMapper;%0Aimport com.google.inject.Inject;%0Aimport com.metamx.emitter.service.ServiceEmitter
+javax.annotation.Nullable;%0Aimport java.util.Map;%0Aimport java.util.concurrent.atomic.AtomicLong
 ;%0A%0A/
@@ -2833,31 +2833,16 @@
   final 
-FinalizeResults
 QueryRun
@@ -2860,16 +2860,60 @@
 Runner =
+ CPUTimeMetricQueryRunner.safeBuild(%0A       
  new Fin
@@ -2932,32 +2932,36 @@
 QueryRunner%3CT%3E(%0A
+    
         toolChes
@@ -2992,32 +2992,36 @@
 on(%0A            
+   
+ 
 toolChest.mergeR
@@ -3028,16 +3028,20 @@
 esults(%0A
+    
         
@@ -3085,32 +3085,36 @@
                 
+  
+  
 toolChest.preMer
@@ -3124,32 +3124,36 @@
 ueryDecoration(%0A
+    
                 
@@ -3213,16 +3213,20 @@
         
+    
 baseClie
@@ -3225,24 +3225,28 @@
 baseClient,%0A
+    
             
@@ -3300,16 +3300,20 @@
         
+    
 retryCon
@@ -3313,24 +3313,28 @@
 etryConfig,%0A
+    
             
@@ -3378,34 +3378,41 @@
                 
-)%0A
+    )%0A   
                 
@@ -3407,35 +3407,40 @@
                 
+ 
 ),%0A
+    
                 
@@ -3461,34 +3461,42 @@
                 
-)%0A
+    )%0A    
             )%0A  
@@ -3493,35 +3493,43 @@
       )%0A        
+ 
-),%0A
+   ),%0A    
         toolChes
@@ -3526,24 +3526,362 @@
   toolChest%0A
+        ),%0A        new Function%3CQuery%3CT%3E, ServiceMetricEvent.Builder%3E()%0A        %7B%0A          @Nullable%0A          @Override%0A          public ServiceMetricEvent.Builder apply(Query%3CT%3E tQuery)%0A          %7B%0A            return toolChest.makeMetricBuilder(tQuery);%0A          %7D%0A        %7D,%0A        emitter,%0A        new AtomicLong(0L),%0A        true%0A
     );%0A%0A    

Add CPUTimeMetricQueryRunner to ClientQuerySegmentWalker