package com.bazaarvoice.emodb.sor.core;

import com.bazaarvoice.emodb.sor.condition.Conditions;
import com.bazaarvoice.emodb.sor.db.MigrationScanResult;
import com.bazaarvoice.emodb.sor.db.MigratorReaderDAO;
import com.bazaarvoice.emodb.sor.db.MigratorWriterDAO;
import com.bazaarvoice.emodb.sor.db.ScanRange;
import com.bazaarvoice.emodb.table.db.StashTableDAO;
import com.bazaarvoice.emodb.table.db.astyanax.FullConsistencyTimeProvider;
import com.bazaarvoice.emodb.table.db.astyanax.PlacementCache;
import com.bazaarvoice.emodb.table.db.consistency.HintsConsistencyTimeProvider;
import com.bazaarvoice.emodb.table.db.stash.StashTokenRange;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterators;
import com.google.common.util.concurrent.RateLimiter;
import com.google.inject.Inject;

import java.util.Iterator;
import java.util.List;

import static com.google.common.base.Preconditions.checkNotNull;

public class DefaultMigratorTools implements MigratorTools {

    private final MigratorReaderDAO _migratorReaderDao;
    private final MigratorWriterDAO _migratorWriterDao;
    private final FullConsistencyTimeProvider _fullConsistencyTimeProvider;
    private final PlacementCache _placementCache;
    private final StashTableDAO _stashTableDao;

    @Inject
    public DefaultMigratorTools(MigratorReaderDAO migratorReaderDAO, MigratorWriterDAO migratorWriterDAO,
                                HintsConsistencyTimeProvider fullConsistencyTimeProvider, PlacementCache placementCache,
                                StashTableDAO stashTableDAO) {
        _migratorReaderDao = checkNotNull(migratorReaderDAO, "migratorReaderDao");
        _migratorWriterDao = checkNotNull(migratorWriterDAO, "migratorWriterDao");
        _fullConsistencyTimeProvider = checkNotNull(fullConsistencyTimeProvider, "fullConsistencyTimeProvider");
        _placementCache = checkNotNull(placementCache, "placementCache");
        _stashTableDao = checkNotNull(stashTableDAO, "stashTableDao");
    }

    @Override
    public void writeRows(String placement, Iterator<MigrationScanResult> results, RateLimiter rateLimiter) {
        checkNotNull(placement, "placement");
        checkNotNull(results, "rows");
        _migratorWriterDao.writeRows(placement, results, rateLimiter);
    }

    @Override
    public Iterator<MigrationScanResult> readRows(String placement, ScanRange scanRange) {
        checkNotNull(placement, "placement");
        checkNotNull(scanRange, "scanRange");

        // Since the range may wrap from high to low end of the token range we need to unwrap it
        List<ScanRange> unwrappedRanges = scanRange.unwrapped();

        Iterator<StashTokenRange> tokenRanges = Iterators.concat(
                Iterators.transform(
                        unwrappedRanges.iterator(),
                        unwrappedRange -> _stashTableDao.getStashTokenRangesFromSnapshot(placement, placement, unwrappedRange.getFrom(), unwrappedRange.getTo())));

        return Iterators.concat(
                Iterators.transform(tokenRanges, tokenRange ->
                        _migratorReaderDao.readRows(placement, ScanRange.create(tokenRange.getFrom(), tokenRange.getTo()))
                )
        );
    }

    @Override
    public long getFullConsistencyTimestamp(String placement) {
        return _fullConsistencyTimeProvider.getMaxTimeStamp(_placementCache.get(placement).getKeyspace().getClusterName());
    }
}

@@ -67,32 +67,52 @@
 odb.
-sor.condition.Conditions
+common.dropwizard.guice.SystemTablePlacement
 ;%0Aim
@@ -660,55 +660,8 @@
 ge;%0A
-import com.google.common.collect.ImmutableSet;%0A
 impo
@@ -1252,16 +1252,64 @@
 ableDao;
+%0A    private final String _systemTablePlacement;
 %0A%0A    @I
@@ -1600,16 +1600,67 @@
 TableDAO
+, @SystemTablePlacement String systemTablePlacement
 ) %7B%0A    
@@ -2075,24 +2075,116 @@
 TableDao%22);%0A
+        _systemTablePlacement = checkNotNull(systemTablePlacement, %22systemTablePlacement%22);%0A
     %7D%0A%0A    @
@@ -2608,32 +2608,32 @@
 , %22placement%22);%0A
-
         checkNot
@@ -2659,24 +2659,265 @@
 anRange%22);%0A%0A
+        // System tables will not show up in the stash token ranges, so we need to scan the entire thing%0A        if (placement.equals(_systemTablePlacement)) %7B%0A            return _migratorReaderDao.readRows(placement, scanRange);%0A        %7D%0A%0A
         // S

Make sys placement be migrated without stash token range index (#208)