/*
 * Druid - a distributed column store.
 * Copyright (C) 2014  Metamarkets Group Inc.
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
 */

package io.druid.client;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.metamx.common.guava.Sequence;
import com.metamx.common.guava.Sequences;
import io.druid.client.cache.Cache;
import io.druid.client.cache.CacheConfig;
import io.druid.client.cache.MapCache;
import io.druid.query.CacheStrategy;
import io.druid.query.Query;
import io.druid.query.QueryRunner;
import io.druid.query.QueryToolChest;
import io.druid.query.SegmentDescriptor;

import java.util.ArrayList;

public class CachePopulatingQueryRunner<T> implements QueryRunner<T>
{

  private final String segmentIdentifier;
  private final SegmentDescriptor segmentDescriptor;
  private final QueryRunner<T> base;
  private final QueryToolChest toolChest;
  private final Cache cache;
  private final ObjectMapper mapper;
  private final CacheConfig cacheConfig;

  public CachePopulatingQueryRunner(
      String segmentIdentifier,
      SegmentDescriptor segmentDescriptor, ObjectMapper mapper,
      Cache cache,
      QueryToolChest toolchest,
      QueryRunner<T> base,
      CacheConfig cacheConfig
  )
  {
    this.base = base;
    this.segmentIdentifier = segmentIdentifier;
    this.segmentDescriptor = segmentDescriptor;
    this.toolChest = toolchest;
    this.cache = cache;
    this.mapper = mapper;
    this.cacheConfig = cacheConfig;
  }

  @Override
  public Sequence<T> run(Query<T> query)
  {

    final CacheStrategy strategy = toolChest.getCacheStrategy(query);

    final boolean populateCache = Boolean.parseBoolean(query.getContextValue(CacheConfig.POPULATE_CACHE, "true"))
                                  && strategy != null
                                  && cacheConfig.isPopulateCache()
                                  // historical only populates distributed cache since the cache lookups are done at broker.
                                  && !(cache instanceof MapCache) ;
    Sequence<T> results = base.run(query);
    if (populateCache) {
      Cache.NamedKey key = CacheUtil.computeSegmentCacheKey(
          segmentIdentifier,
          segmentDescriptor,
          strategy.computeCacheKey(query)
      );
      CacheUtil.populate(
          cache,
          mapper,
          key,
          Sequences.toList(Sequences.map(results, strategy.prepareForCache()), new ArrayList())
      );
    }
    return results;

  }
}

@@ -882,24 +882,108 @@
 jectMapper;%0A
+import com.google.common.collect.Lists;%0Aimport com.metamx.common.guava.Accumulator;%0A
 import com.m
@@ -2782,11 +2782,37 @@
 che)
- 
 ;%0A
+    if (populateCache) %7B%0A  
     
@@ -2854,33 +2854,8 @@
 y);%0A
-    if (populateCache) %7B%0A
     
@@ -3016,24 +3016,105 @@
 y)%0A      );%0A
+      ArrayList%3CT%3E resultAsList = Sequences.toList(results, new ArrayList%3CT%3E());%0A
       CacheU
@@ -3191,46 +3191,36 @@
     
-Sequences.toList(Sequences.map
+Lists.transform
 (result
-s
+AsList
 , st
@@ -3248,65 +3248,115 @@
 e())
-, new Array
+%0A      );%0A      return Sequences.simple(resultAs
 List
-()
 )
+;
 %0A    
-  );
+%7D else %7B
 %0A   
- %7D%0A
    
- 
 return 
-results;%0A
+base.run(query);%0A    %7D
 %0A  %7D%0A
+%0A
 %7D%0A

Fix cache population simple way