<|start_filename|>src/main/java/io/appform/dropwizard/sharding/caching/RelationalCache.java<|end_filename|>
/*
 * Copyright 2016 <NAME> <<EMAIL>>
 *
 * 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.appform.dropwizard.sharding.caching;


import org.hibernate.criterion.DetachedCriteria;

import java.util.List;
import java.util.function.Function;

/**
 * A simple cache interface which allows plugging in any caching framework or infrastructure to enable
 * write through caching
 */
public interface RelationalCache<V> {

    /**
     * Write through method that will be called if cache enabled {@link io.appform.dropwizard.sharding.dao.CacheableRelationalDao#save(String, Object)} is used
     * @param parentKey The key of the parent that the entity is related to
     * @param key The key that needs to be used to write this element to cache
     * @param entity Entity that needs to be written into cache
     */
    void put(String parentKey, Object key, V entity);

    /**
     * Write through method that will be called if cache enabled {@link io.appform.dropwizard.sharding.dao.CacheableRelationalDao#select(String, DetachedCriteria, int, int, Function)} is used
     * @param parentKey The key of the parent that the entity is related to
     * @param entities List of entities that needs to be written into cache
     */
    void put(String parentKey, List<V> entities);

    /**
     * Write through method that will be called if cache enabled {@link io.appform.dropwizard.sharding.dao.CacheableRelationalDao#select(String, DetachedCriteria, int, int, Function)} is used
     * @param parentKey The key of the parent that the entity is related to
     * @param entities List of entities that needs to be written into cache
     */
    void put(String parentKey, int first, int numResults, List<V> entities);

    /**
     * Read through exists method that will be called if cache enabled {@link io.appform.dropwizard.sharding.dao.CacheableRelationalDao#exists(String, Object)} is used
     * @param parentKey The key of the parent that the entity is related to
     * @param key The key that needs to be used to write this element to cache
     * @return Whether the entity exists or not.
     */
    boolean exists(String parentKey, Object key);

    /**
     * Read through method that will be called if a cache enabled {@link io.appform.dropwizard.sharding.dao.CacheableRelationalDao#get(String, Object)} is used
     * @param parentKey The key of the parent the entity is related to
     * @param key The key of the entity that needs to be read
     * @return entity Entity that was read through the cache
     */
    V get(String parentKey, Object key);

    /**
     * Read through method that will be called if a cache enabled {@link io.appform.dropwizard.sharding.dao.CacheableRelationalDao#select(String, DetachedCriteria, int, int, Function)} is used
     * @param parentKey The key of the parent the entity is related to
     * @return Entities that was read through the cache
     */
    List<V> select(String parentKey);

    /**
     * Read through method that will be called if a cache enabled {@link io.appform.dropwizard.sharding.dao.CacheableRelationalDao#select(String, DetachedCriteria, int, int, Function)} is used
     * @param parentKey The key of the parent the entity is related to
     * @return Entities that was read through the cache
     */
    List<V> select(String parentKey, int first, int numResults);

}


<|start_filename|>src/test/java/io/appform/dropwizard/sharding/sharding/BalancedShardManagerTest.java<|end_filename|>
/*
 * Copyright 2019 <NAME> <<EMAIL>>
 *  
 * 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.appform.dropwizard.sharding.sharding;

import org.junit.Test;

import static org.junit.Assert.assertEquals;

public class BalancedShardManagerTest {

   @Test(expected = IllegalArgumentException.class)
    public void testShardForBucket() throws Exception {
        new BalancedShardManager(5);
    }

   @Test(expected = IllegalArgumentException.class)
    public void testShardForOddBucket() throws Exception {
        new BalancedShardManager(9);
    }

    @Test(expected = IllegalArgumentException.class)
    public void testShardForEvenNon2PowerBucket() throws Exception {
        new BalancedShardManager(40);
    }

    @Test
    public void testShardFor64Bucket() throws Exception {
        BalancedShardManager shardManager = new BalancedShardManager(64);
        assertEquals(63, shardManager.shardForBucket(1023));
    }

    @Test
    public void testShardFor32Bucket() throws Exception {
        BalancedShardManager shardManager = new BalancedShardManager(32);
        assertEquals(31, shardManager.shardForBucket(1023));
    }

    @Test
    public void testShardFor31Bucket() throws Exception {
        BalancedShardManager shardManager = new BalancedShardManager(16);
        assertEquals(15, shardManager.shardForBucket(1023));
    }
}

<|start_filename|>src/test/java/io/appform/dropwizard/sharding/ShardInfoProviderTest.java<|end_filename|>
package io.appform.dropwizard.sharding;

import org.junit.Assert;
import org.junit.Test;

public class ShardInfoProviderTest {

    @Test
    public void testGetShardId() {
        ShardInfoProvider shardInfoProvider = new ShardInfoProvider("default");
        int shardId = shardInfoProvider.shardId("connectionpool-default-2");
        Assert.assertEquals(2, shardId);

        shardId = shardInfoProvider.shardId("connectionpool");
        Assert.assertEquals(-1, shardId);

        shardId = shardInfoProvider.shardId("connectionpool-test-1");
        Assert.assertEquals(1, shardId);

        shardId = shardInfoProvider.shardId("connectionpool-1");
        Assert.assertEquals(-1, shardId);

        shardId = shardInfoProvider.shardId("default-1");
        Assert.assertEquals(-1, shardId);
    }

    @Test
    public void testGetNamespace() {
        ShardInfoProvider shardInfoProvider = new ShardInfoProvider("default");
        String namespace = shardInfoProvider.namespace("connectionpool-default-2");
        Assert.assertEquals("default", namespace);

        namespace = shardInfoProvider.namespace("connectionpool");
        Assert.assertNull(namespace);

        namespace = shardInfoProvider.namespace("connectionpool-default");
        Assert.assertNull(namespace);

        namespace = shardInfoProvider.namespace("connectionpool-test-1");
        Assert.assertEquals("test", namespace);
    }

    @Test
    public void testGetShardName() {
        ShardInfoProvider shardInfoProvider = new ShardInfoProvider("default");
        String shardName = shardInfoProvider.shardName(1);
        Assert.assertEquals("connectionpool-default-1", shardName);
    }
}


<|start_filename|>src/main/java/io/appform/dropwizard/sharding/sharding/LegacyShardManager.java<|end_filename|>
/*
 * Copyright 2016 <NAME> <<EMAIL>>
 *
 * 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.appform.dropwizard.sharding.sharding;

import com.google.common.base.Preconditions;
import com.google.common.collect.Range;
import com.google.common.collect.RangeMap;
import com.google.common.collect.TreeRangeMap;
import lombok.Builder;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import lombok.val;

/**
 * Manages shard to bucket mapping.
 */
@ToString
@Slf4j
public class LegacyShardManager extends ShardManager {
    private static final int MIN_BUCKET = 0;
    private static final int MAX_BUCKET = 999;
    private final int numShards;

    private RangeMap<Integer, Integer> buckets = TreeRangeMap.create();

    public LegacyShardManager(int numShards) {
        this(numShards, new InMemoryLocalShardBlacklistingStore());
    }

    @Builder
    public LegacyShardManager(int numShards, ShardBlacklistingStore shardBlacklistingStore) {
        super(shardBlacklistingStore);
        Preconditions.checkArgument(
                numShards > 1 && ((numShards & (numShards - 1)) == 0),
                "Shard manager only support 2^n shards." +
                        " Also it is senseless to use anything other than 2^n shards for scale out.");
        this.numShards = numShards;
        final RangeMap<Integer, Integer> assignedBuckets = TreeRangeMap.create();
        int interval = MAX_BUCKET / numShards;
        log.trace("Interval: {}", interval);
        int shardCounter = 0;
        boolean endReached = false;
        boolean altAssignment = ((MAX_BUCKET - (interval * numShards)) > interval);
        for (int start = MIN_BUCKET; !endReached; start += interval, shardCounter++) {
            int remaining = MAX_BUCKET - start;
            log.trace("Remaining: {}. Interval: {}, 2x Interval: {}", remaining, interval, 2 * interval);
            int end = start + interval - 1;
            endReached = shardCounter == numShards - 1;
            end = endReached
                    ? MAX_BUCKET
                    : end;
            log.trace("Assigning {} elements, from {} to {} into shard {}", end - start + 1, start, end, shardCounter);
            //End is reached when remaining items (max -
            assignedBuckets.put(Range.closed(start, end), shardCounter);
        }
        Preconditions.checkArgument(assignedBuckets.asMapOfRanges().size() == numShards,
                                    "There is an issue in shard allocation. " +
                                            "Not all shards have been allocated to. Please contact devs.");
        buckets.putAll(assignedBuckets);
        log.info("Buckets to shard allocation: {}", buckets);
    }


    @Override
    public int numBuckets() {
        return 1000;
    }

    @Override
    public int numShards() {
        return numShards;
    }

    @Override
    protected int shardForBucketImpl(int bucketId) {
        Preconditions.checkArgument(bucketId >=MIN_BUCKET && bucketId <= MAX_BUCKET, "Bucket id can only be in the range of [1-1000] (inclusive)");
        val entry = buckets.getEntry(bucketId);
        if(null == entry) {
            throw new IllegalAccessError("Bucket not mapped to any shard");
        }
        return entry.getValue();
    }


}


<|start_filename|>src/main/java/io/appform/dropwizard/sharding/caching/LookupCache.java<|end_filename|>
/*
 * Copyright 2016 <NAME> <<EMAIL>>
 *
 * 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.appform.dropwizard.sharding.caching;


import io.appform.dropwizard.sharding.dao.CacheableLookupDao;

/**
 * A simple cache interface which allows plugging in any caching framework or infrastructure to enable
 * write through caching
 */
public interface LookupCache<V> {

    /**
     * Write through method that will be called if cache enabled {@link io.appform.dropwizard.sharding.dao.CacheableLookupDao#save(Object)} is used
     * @param key The key that needs to be used to write this element to cache
     * @param entity Entity that needs to be written into cache
     */
    void put(String key, V entity);

    /**
     * Read through exists method that will be called if cache enabled {@link io.appform.dropwizard.sharding.dao.CacheableLookupDao#exists(String)} is used
     * @param key The key of the entity that needs to be checked
     * @return whether the entity exists or not.
     */
    boolean exists(String key);

    /**
     * Read through method that will be called if a cache enabled {@link io.appform.dropwizard.sharding.dao.CacheableLookupDao#get(String)} is used
     * @param key The key of the entity that needs to be read
     * @return entity Entity that was read through the cache
     */
    V get(String key);
}


<|start_filename|>src/main/java/io/appform/dropwizard/sharding/sharding/ShardManager.java<|end_filename|>
/*
 * Copyright 2019 <NAME> <<EMAIL>>
 *
 * 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.appform.dropwizard.sharding.sharding;

import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.LoadingCache;
import io.appform.dropwizard.sharding.exceptions.ShardBlacklistedException;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.TimeUnit;


/**
 *
 */
@ToString
@Slf4j
public abstract class ShardManager {

    private final ShardBlacklistingStore shardBlacklistingStore;
    private LoadingCache<Integer, Boolean> blackListedShards;

    abstract public int numBuckets();

    abstract protected int numShards();

    abstract protected int shardForBucketImpl(int bucketId);

    protected ShardManager(ShardBlacklistingStore shardBlacklistingStore) {
        this.shardBlacklistingStore = shardBlacklistingStore;
        this.blackListedShards = Caffeine.newBuilder()
                .maximumSize(10_000)
                .expireAfterWrite(1, TimeUnit.MINUTES)
                .refreshAfterWrite(15, TimeUnit.SECONDS)
                .build(shardBlacklistingStore::blacklisted);
    }

    public int shardForBucket(int bucketId) {
        final int shard = shardForBucketImpl(bucketId);
        final Boolean isBlacklisted = blackListedShards.get(shard);
        if(null != isBlacklisted && isBlacklisted) {
            throw new ShardBlacklistedException(shard);
        }
        return shard;
    }

    public boolean isMappedToValidShard(int bucketId) {
        final int shard = shardForBucketImpl(bucketId);
        final Boolean isBlacklisted = blackListedShards.get(shard);
        if(null != isBlacklisted && isBlacklisted) {
            return false;
        }
        return true;
    }

    public void blacklistShard(int shardId) {
        if(shardId >=0 && shardId < numShards()) {
            shardBlacklistingStore.blacklist(shardId);
            blackListedShards.refresh(shardId);
        }
    }

    public boolean isBlacklisted(int shardId) {
        if (shardId >= 0
                && shardId < numShards()
                && shardBlacklistingStore != null) {
            return shardBlacklistingStore.blacklisted(shardId);
        }
        return false;
    }

    public void unblacklistShard(int shardId) {
        if(shardId >=0 && shardId < numShards()) {
            shardBlacklistingStore.unblacklist(shardId);
            blackListedShards.refresh(shardId);
        }
    }
}


<|start_filename|>src/main/java/io/appform/dropwizard/sharding/sharding/InMemoryLocalShardBlacklistingStore.java<|end_filename|>
/*
 * Copyright 2018 <NAME> <<EMAIL>>
 *
 * 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.appform.dropwizard.sharding.sharding;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 *
 */
public class InMemoryLocalShardBlacklistingStore implements ShardBlacklistingStore {

    private final ConcurrentMap<Integer, Boolean> blacklisted = new ConcurrentHashMap<>();

    @Override
    public void blacklist(int shardId) {
        blacklisted.put(shardId, true);
    }

    @Override
    public void unblacklist(int shardId) {
        blacklisted.remove(shardId);
    }

    @Override
    public boolean blacklisted(int shardId) {
        return blacklisted.getOrDefault(shardId, false);
    }

}


<|start_filename|>src/main/java/io/appform/dropwizard/sharding/sharding/ShardedTransaction.java<|end_filename|>
package io.appform.dropwizard.sharding.sharding;

import io.dropwizard.hibernate.HibernateBundle;

import java.lang.annotation.Retention;
import java.lang.annotation.Target;

import static java.lang.annotation.ElementType.METHOD;
import static java.lang.annotation.RetentionPolicy.RUNTIME;

/**
 * Used to mark a transactional operation in a derviced {@link io.dropwizard.hibernate.AbstractDAO}
 * for use by {@link io.appform.dropwizard.sharding.dao.WrapperDao}
 */
@Target(METHOD)
@Retention(RUNTIME)
public @interface ShardedTransaction {
    String value() default HibernateBundle.DEFAULT_NAME;
    boolean readOnly() default false;
}


<|start_filename|>src/main/java/io/appform/dropwizard/sharding/utils/TransactionHandler.java<|end_filename|>
/*
 * Copyright 2016 <NAME> <<EMAIL>>
 *
 * 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.appform.dropwizard.sharding.utils;

import lombok.Getter;
import org.hibernate.*;
import org.hibernate.context.internal.ManagedSessionContext;
import org.hibernate.resource.transaction.spi.TransactionStatus;

/**
 * A transaction handler utility class
 */
public class TransactionHandler {


    // Context variables
    @Getter
    private Session session;
    private final SessionFactory sessionFactory;
    private boolean readOnly;

    public TransactionHandler(SessionFactory sessionFactory, boolean readOnly) {
        this.sessionFactory = sessionFactory;
        this.readOnly = readOnly;
    }

    public void beforeStart() {

        session = sessionFactory.openSession();
        try {
            configureSession();
            ManagedSessionContext.bind(session);
            beginTransaction();
        } catch (Throwable th) {
            session.close();
            session = null;
            ManagedSessionContext.unbind(sessionFactory);
            throw th;
        }
    }

    public void afterEnd() {
        if (session == null) {
            return;
        }

        try {
            commitTransaction();
        } catch (Exception e) {
            rollbackTransaction();
            throw e;
        } finally {
            session.close();
            session = null;
            ManagedSessionContext.unbind(sessionFactory);
        }

    }

    public void onError() {
        if (session == null) {
            return;
        }

        try {
            rollbackTransaction();
        } finally {
            session.close();
            session = null;
            ManagedSessionContext.unbind(sessionFactory);
        }
    }

    private void configureSession() {
        session.setDefaultReadOnly(readOnly);
        session.setCacheMode(CacheMode.NORMAL);
        session.setHibernateFlushMode(FlushMode.AUTO);
    }

    private void beginTransaction() {
        session.beginTransaction();
    }

    private void rollbackTransaction() {
        final Transaction txn = session.getTransaction();
        if (txn != null && txn.getStatus() == TransactionStatus.ACTIVE) {
            txn.rollback();
        }
    }

    private void commitTransaction() {
        final Transaction txn = session.getTransaction();
        if (txn != null && txn.getStatus() == TransactionStatus.ACTIVE) {
            txn.commit();
        }
    }
}


<|start_filename|>src/main/java/io/appform/dropwizard/sharding/sharding/LookupKey.java<|end_filename|>
package io.appform.dropwizard.sharding.sharding;

import io.appform.dropwizard.sharding.dao.LookupDao;

import java.lang.annotation.Retention;
import java.lang.annotation.Target;

import static java.lang.annotation.ElementType.FIELD;
import static java.lang.annotation.ElementType.METHOD;
import static java.lang.annotation.RetentionPolicy.RUNTIME;

/**
 * Annotation to annotate top level lookup key to be used by {@link LookupDao}
 */
@Target({FIELD, METHOD})
@Retention(RUNTIME)
public @interface LookupKey {
}


<|start_filename|>src/main/java/io/appform/dropwizard/sharding/dao/UpdateOperationMeta.java<|end_filename|>
package io.appform.dropwizard.sharding.dao;

import lombok.Builder;
import lombok.Data;

import java.util.Map;

@Data
@Builder
public class UpdateOperationMeta {

    private String queryName;

    private Map<String, Object> params;

}


<|start_filename|>src/main/java/io/appform/dropwizard/sharding/config/BlacklistConfig.java<|end_filename|>
package io.appform.dropwizard.sharding.config;

import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;

@Data
@Builder
@NoArgsConstructor
@AllArgsConstructor
public class BlacklistConfig {

    private boolean skipNativeHealthcheck;

}


<|start_filename|>src/test/java/io/appform/dropwizard/sharding/dao/testdata/multi/MultiPackageTestEntity.java<|end_filename|>
package io.appform.dropwizard.sharding.dao.testdata.multi;

import io.appform.dropwizard.sharding.sharding.LookupKey;
import lombok.*;

import javax.persistence.*;

@Entity
@Table(name="audits")
@Data
@NoArgsConstructor
@AllArgsConstructor
@ToString
@Builder
public class MultiPackageTestEntity {
    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    @Column(name = "id")
    private long id;

    @LookupKey
    @Column(name = "lookup")
    private String lookup;

    @Column(name = "text")
    private String text;

}



<|start_filename|>src/test/java/io/appform/dropwizard/sharding/dao/LookupDaoTest.java<|end_filename|>
/*
 * Copyright 2016 <NAME> <<EMAIL>>
 *
 * 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.appform.dropwizard.sharding.dao;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import io.appform.dropwizard.sharding.dao.testdata.entities.Audit;
import io.appform.dropwizard.sharding.dao.testdata.entities.Phone;
import io.appform.dropwizard.sharding.dao.testdata.entities.TestEntity;
import io.appform.dropwizard.sharding.dao.testdata.entities.Transaction;
import io.appform.dropwizard.sharding.sharding.BalancedShardManager;
import io.appform.dropwizard.sharding.sharding.ShardManager;
import io.appform.dropwizard.sharding.sharding.impl.ConsistentHashBucketIdExtractor;
import io.appform.dropwizard.sharding.utils.ShardCalculator;
import lombok.val;
import org.hibernate.SessionFactory;
import org.hibernate.boot.registry.StandardServiceRegistry;
import org.hibernate.boot.registry.StandardServiceRegistryBuilder;
import org.hibernate.cfg.Configuration;
import org.hibernate.criterion.DetachedCriteria;
import org.hibernate.criterion.Restrictions;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import java.util.List;
import java.util.Optional;
import java.util.UUID;

import static org.junit.Assert.*;

public class LookupDaoTest {

    private List<SessionFactory> sessionFactories = Lists.newArrayList();
    private LookupDao<TestEntity> lookupDao;
    private LookupDao<Phone> phoneDao;
    private RelationalDao<Transaction> transactionDao;
    private RelationalDao<Audit> auditDao;

    private SessionFactory buildSessionFactory(String dbName) {
        Configuration configuration = new Configuration();
        configuration.setProperty("hibernate.dialect",
                                  "org.hibernate.dialect.H2Dialect");
        configuration.setProperty("hibernate.connection.driver_class",
                                  "org.h2.Driver");
        configuration.setProperty("hibernate.connection.url", "jdbc:h2:mem:" + dbName);
        configuration.setProperty("hibernate.hbm2ddl.auto", "create");
        configuration.setProperty("hibernate.current_session_context_class", "managed");
        configuration.setProperty("hibernate.show_sql", "true");
        configuration.setProperty("hibernate.format_sql", "true");
        configuration.addAnnotatedClass(TestEntity.class);
        configuration.addAnnotatedClass(Phone.class);
        configuration.addAnnotatedClass(Transaction.class);
        configuration.addAnnotatedClass(Audit.class);

        StandardServiceRegistry serviceRegistry
                = new StandardServiceRegistryBuilder().applySettings(
                configuration.getProperties())
                .build();
        return configuration.buildSessionFactory(serviceRegistry);
    }

    @Before
    public void before() {
        for (int i = 0; i < 2; i++) {
            sessionFactories.add(buildSessionFactory(String.format("db_%d", i)));
        }
        final ShardManager shardManager = new BalancedShardManager(sessionFactories.size());
        final ShardCalculator<String> shardCalculator = new ShardCalculator<>(shardManager,
                                                                              new ConsistentHashBucketIdExtractor<>(
                                                                                      shardManager));
        lookupDao = new LookupDao<>(sessionFactories, TestEntity.class, shardCalculator);
        phoneDao = new LookupDao<>(sessionFactories, Phone.class, shardCalculator);
        transactionDao = new RelationalDao<>(sessionFactories, Transaction.class, shardCalculator);
        auditDao = new RelationalDao<>(sessionFactories, Audit.class, shardCalculator);
    }

    @After
    public void after() {
        sessionFactories.forEach(SessionFactory::close);
    }

    @Test
    public void testSave() throws Exception {
        TestEntity testEntity = TestEntity.builder()
                .externalId("testId")
                .text("Some Text")
                .build();
        lookupDao.save(testEntity);

        assertEquals(true, lookupDao.exists("testId"));
        assertEquals(false, lookupDao.exists("testId1"));
        Optional<TestEntity> result = lookupDao.get("testId");
        assertEquals("Some Text",
                     result.get()
                             .getText());

        testEntity.setText("Some New Text");
        lookupDao.save(testEntity);
        result = lookupDao.get("testId");
        assertEquals("Some New Text",
                     result.get()
                             .getText());

        boolean updateStatus = lookupDao.update("testId", entity -> {
            if (entity.isPresent()) {
                TestEntity e = entity.get();
                e.setText("Updated text");
                return e;
            }
            return null;
        });

        assertTrue(updateStatus);
        result = lookupDao.get("testId");
        assertEquals("Updated text",
                     result.get()
                             .getText());

        updateStatus = lookupDao.update("testIdxxx", entity -> {
            if (entity.isPresent()) {
                TestEntity e = entity.get();
                e.setText("Updated text");
                return e;
            }
            return null;
        });

        assertFalse(updateStatus);
    }

    @Test
    public void testScatterGather() throws Exception {
        List<TestEntity> results = lookupDao.scatterGather(DetachedCriteria.forClass(TestEntity.class)
                                                                   .add(Restrictions.eq("externalId", "testId")));
        assertTrue(results.isEmpty());

        TestEntity testEntity = TestEntity.builder()
                .externalId("testId")
                .text("Some Text")
                .build();
        lookupDao.save(testEntity);
        results = lookupDao.scatterGather(DetachedCriteria.forClass(TestEntity.class)
                                                  .add(Restrictions.eq("externalId", "testId")));
        assertFalse(results.isEmpty());
        assertEquals("Some Text",
                     results.get(0)
                             .getText());
    }

    @Test
    public void testListGetQuery() throws Exception {
        List<String> lookupKeys = Lists.newArrayList();
        lookupKeys.add("testId1");
        List<TestEntity> results = lookupDao.get(lookupKeys);
        assertTrue(results.isEmpty());

        TestEntity testEntity1 = TestEntity.builder()
                .externalId("testId1")
                .text("Some Text 1")
                .build();
        lookupDao.save(testEntity1);
        results = lookupDao.get(lookupKeys);
        assertFalse(results.isEmpty());
        assertEquals(1, results.size());
        assertEquals("Some Text 1",
                     results.get(0)
                             .getText());

        TestEntity testEntity2 = TestEntity.builder()
                .externalId("testId2")
                .text("Some Text 2")
                .build();
        lookupDao.save(testEntity2);
        lookupKeys.add("testId2");
        results = lookupDao.get(lookupKeys);
        assertFalse(results.isEmpty());
        assertEquals(2, results.size());
    }

    @Test
    public void testUpdateUsingNamedQueryRowUpdated() throws Exception {
        val id = UUID.randomUUID().toString();
        val testEntity = TestEntity.builder()
                .externalId(id)
                .text(UUID.randomUUID().toString())
                .build();
        lookupDao.save(testEntity);

        val newText = UUID.randomUUID().toString();
        int rowsUpdated = lookupDao.updateUsingQuery(id, UpdateOperationMeta.builder()
                .queryName("testTextUpdateQuery")
                .params(ImmutableMap.of("externalId", id,
                        "text", newText))
                .build());
        assertEquals(1, rowsUpdated);

        val persistedEntity = lookupDao.get(id).orElse(null);
        assertNotNull(persistedEntity);
        assertEquals(newText, persistedEntity.getText());
    }

    @Test
    public void testUpdateUsingNamedQueryNoRowUpdated() throws Exception {
        val id = UUID.randomUUID().toString();
        val testEntity = TestEntity.builder()
                .externalId(id)
                .text(UUID.randomUUID().toString())
                .build();
        lookupDao.save(testEntity);

        val newText = UUID.randomUUID().toString();
        int rowsUpdated = lookupDao.updateUsingQuery(id, UpdateOperationMeta.builder()
                .queryName("testTextUpdateQuery")
                .params(ImmutableMap.of("externalId", UUID.randomUUID().toString(),
                        "text", newText))
                .build());
        assertEquals(0, rowsUpdated);

        val persistedEntity = lookupDao.get(id).orElse(null);
        assertNotNull(persistedEntity);
        assertEquals(testEntity.getText(), persistedEntity.getText());
    }

    @Test
    public void testSaveInParentBucket() throws Exception {
        final String phoneNumber = "9830968020";

        Phone phone = Phone.builder()
                .phone(phoneNumber)
                .build();

        Phone savedPhone = phoneDao.save(phone)
                .get();

        Transaction transaction = Transaction.builder()
                .transactionId("testTxn")
                .to("9830703153")
                .amount(100)
                .phone(savedPhone)
                .build();

        transactionDao.save(savedPhone.getPhone(), transaction)
                .get();
        {
            Transaction resultTx = transactionDao.get(phoneNumber, "testTxn")
                    .get();
            assertEquals(phoneNumber,
                         resultTx.getPhone()
                                 .getPhone());
            assertTrue(transactionDao.exists(phoneNumber, "testTxn"));
            assertFalse(transactionDao.exists(phoneNumber, "testTxn1"));
        }
        {
            Optional<Transaction> resultTx = transactionDao.get(phoneNumber, "testTxn1");
            assertFalse(resultTx.isPresent());
        }
        saveAudit(phoneNumber, "testTxn", "Started");
        saveAudit(phoneNumber, "testTxn", "Underway");
        saveAudit(phoneNumber, "testTxn", "Completed");

        assertEquals(3, auditDao.count(phoneNumber, DetachedCriteria.forClass(Audit.class)
                .add(Restrictions.eq("transaction.transactionId", "testTxn"))));

        List<Audit> audits = auditDao.select(phoneNumber, DetachedCriteria.forClass(Audit.class)
                .add(Restrictions.eq("transaction.transactionId", "testTxn")), 0, 10);
        assertEquals("Started",
                     audits.get(0)
                             .getText());

    }

    private void saveAudit(String phone, String transaction, String text) throws Exception {
        auditDao.save(phone, Audit.builder()
                .text(text)
                .transaction(Transaction.builder()
                                     .transactionId(transaction)
                                     .build())
                .build());
    }

    @Test
    public void testHierarchy() throws Exception {
        final String phoneNumber = "9986032019";
        saveHierarchy(phoneNumber);
        saveHierarchy("9986402019");

        List<Audit> audits = auditDao.select(phoneNumber, DetachedCriteria.forClass(Audit.class)
                .add(Restrictions.eq("transaction.transactionId", "newTxn-" + phoneNumber)), 0, 10);

        assertEquals(2, audits.size());

        List<Audit> allAudits = auditDao.scatterGather(DetachedCriteria.forClass(Audit.class), 0, 10);
        assertEquals(4, allAudits.size());
    }


    private void saveHierarchy(String phone) throws Exception {

        Transaction transaction = Transaction.builder()
                .transactionId("newTxn-" + phone)
                .amount(100)
                .to("9986402019")
                .build();

        Audit started = Audit.builder()
                .text("Started")
                .transaction(transaction)
                .build();

        Audit completed = Audit.builder()
                .text("Completed")
                .transaction(transaction)
                .build();

        transaction.setAudits(ImmutableList.of(started, completed));

        transactionDao.save(phone, transaction);
    }

    @Test
    public void deleteTest() throws Exception {
        TestEntity testEntity = TestEntity.builder()
                .externalId("testId")
                .text("Some Text")
                .build();
        lookupDao.save(testEntity);
        Assert.assertNotNull(lookupDao.get("testId")
                                     .orElse(null));
        Assert.assertTrue(lookupDao.delete("testId"));
        Assert.assertNull(lookupDao.get("testId")
                                  .orElse(null));
    }

    @Test
    public void testCount() throws Exception {
        DetachedCriteria criteria = DetachedCriteria.forClass(TestEntity.class)
                .add(Restrictions.eq("text", "TEST_TYPE"));

        Assert.assertEquals(
                0L,
                (long)lookupDao.count(criteria).stream().reduce(0L, Long::sum)
        );

        TestEntity testEntity = TestEntity.builder()
                .externalId("testId2")
                .text("TEST_TYPE")
                .build();
        lookupDao.save(testEntity);

        testEntity.setExternalId("testId3");
        lookupDao.save(testEntity);

        Assert.assertEquals(
                2L,
                (long)lookupDao.count(criteria).stream().reduce(0L, Long::sum)
        );


        lookupDao.delete("testId2");
        Assert.assertEquals(
                1L,
                (long)lookupDao.count(criteria).stream().reduce(0L, Long::sum)
        );

    }
}

<|start_filename|>src/test/java/io/appform/dropwizard/sharding/dao/CacheableLookupDaoTest.java<|end_filename|>
/*
 * Copyright 2016 <NAME> <<EMAIL>>
 *
 * 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.appform.dropwizard.sharding.dao;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import io.appform.dropwizard.sharding.caching.LookupCache;
import io.appform.dropwizard.sharding.caching.RelationalCache;
import io.appform.dropwizard.sharding.dao.testdata.entities.Audit;
import io.appform.dropwizard.sharding.dao.testdata.entities.Phone;
import io.appform.dropwizard.sharding.dao.testdata.entities.TestEntity;
import io.appform.dropwizard.sharding.dao.testdata.entities.Transaction;
import io.appform.dropwizard.sharding.sharding.BalancedShardManager;
import io.appform.dropwizard.sharding.sharding.ShardManager;
import io.appform.dropwizard.sharding.sharding.impl.ConsistentHashBucketIdExtractor;
import io.appform.dropwizard.sharding.utils.ShardCalculator;
import org.apache.commons.lang3.StringUtils;
import org.hibernate.SessionFactory;
import org.hibernate.boot.registry.StandardServiceRegistry;
import org.hibernate.boot.registry.StandardServiceRegistryBuilder;
import org.hibernate.cfg.Configuration;
import org.hibernate.criterion.DetachedCriteria;
import org.hibernate.criterion.Restrictions;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import static org.junit.Assert.*;

public class CacheableLookupDaoTest {

    private List<SessionFactory> sessionFactories = Lists.newArrayList();
    private CacheableLookupDao<TestEntity> lookupDao;
    private CacheableLookupDao<Phone> phoneDao;
    private CacheableRelationalDao<Transaction> transactionDao;
    private CacheableRelationalDao<Audit> auditDao;

    private SessionFactory buildSessionFactory(String dbName) {
        Configuration configuration = new Configuration();
        configuration.setProperty("hibernate.dialect",
                                  "org.hibernate.dialect.H2Dialect");
        configuration.setProperty("hibernate.connection.driver_class",
                                  "org.h2.Driver");
        configuration.setProperty("hibernate.connection.url", "jdbc:h2:mem:" + dbName);
        configuration.setProperty("hibernate.hbm2ddl.auto", "create");
        configuration.setProperty("hibernate.current_session_context_class", "managed");
        configuration.addAnnotatedClass(TestEntity.class);
        configuration.addAnnotatedClass(Phone.class);
        configuration.addAnnotatedClass(Transaction.class);
        configuration.addAnnotatedClass(Audit.class);

        StandardServiceRegistry serviceRegistry
                = new StandardServiceRegistryBuilder().applySettings(
                configuration.getProperties())
                .build();
        return configuration.buildSessionFactory(serviceRegistry);
    }


    @Before
    public void before() {
        for (int i = 0; i < 2; i++) {
            sessionFactories.add(buildSessionFactory(String.format("db_%d", i)));
        }
        final ShardManager shardManager = new BalancedShardManager(sessionFactories.size());
        lookupDao = new CacheableLookupDao<>(
                sessionFactories,
                TestEntity.class,
                new ShardCalculator<>(shardManager, new ConsistentHashBucketIdExtractor<>(shardManager)),
                new LookupCache<TestEntity>() {

                    private Map<String, TestEntity> cache = new HashMap<>();

                    @Override
                    public void put(String key, TestEntity entity) {
                        cache.put(key, entity);
                    }

                    @Override
                    public boolean exists(String key) {
                        return cache.containsKey(key);
                    }

                    @Override
                    public TestEntity get(String key) {
                        return cache.get(key);
                    }
                });
        phoneDao = new CacheableLookupDao<>(sessionFactories,
                                            Phone.class,
                                            new ShardCalculator<>(shardManager,
                                                                  new ConsistentHashBucketIdExtractor<>(shardManager)),
                                            new LookupCache<Phone>() {

                                                private Map<String, Phone> cache = new HashMap<>();

                                                @Override
                                                public void put(String key, Phone entity) {
                                                    cache.put(key, entity);
                                                }

                                                @Override
                                                public boolean exists(String key) {
                                                    return cache.containsKey(key);
                                                }

                                                @Override
                                                public Phone get(String key) {
                                                    return cache.get(key);
                                                }
                                            });
        transactionDao = new CacheableRelationalDao<>(sessionFactories,
                                                      Transaction.class,
                                                      new ShardCalculator<>(shardManager,
                                                                            new ConsistentHashBucketIdExtractor<>(
                                                                                    shardManager)),
                                                      new RelationalCache<Transaction>() {

                                                          private Map<String, Object> cache = new HashMap<>();

                                                          @Override
                                                          public void put(
                                                                  String parentKey,
                                                                  Object key,
                                                                  Transaction entity) {
                                                              cache.put(StringUtils.join(parentKey, key, ':'), entity);
                                                          }

                                                          @Override
                                                          public void put(
                                                                  String parentKey,
                                                                  List<Transaction> entities) {
                                                              cache.put(parentKey, entities);
                                                          }

                                                          @Override
                                                          public void put(
                                                                  String parentKey,
                                                                  int first,
                                                                  int numResults,
                                                                  List<Transaction> entities) {
                                                              cache.put(StringUtils.join(parentKey,
                                                                                         first,
                                                                                         numResults,
                                                                                         ':'), entities);
                                                          }

                                                          @Override
                                                          public boolean exists(String parentKey, Object key) {
                                                              return cache.containsKey(StringUtils.join(parentKey,
                                                                                                        key,
                                                                                                        ':'));
                                                          }

                                                          @Override
                                                          public Transaction get(String parentKey, Object key) {
                                                              return (Transaction) cache.get(StringUtils.join(parentKey,
                                                                                                              key,
                                                                                                              ':'));
                                                          }

                                                          @Override
                                                          public List<Transaction> select(String parentKey) {
                                                              return (List<Transaction>) cache.get(parentKey);
                                                          }

                                                          @Override
                                                          public List<Transaction> select(
                                                                  String parentKey,
                                                                  int first,
                                                                  int numResults) {
                                                              return (List<Transaction>) cache.get(StringUtils.join(
                                                                      parentKey,
                                                                      first,
                                                                      numResults,
                                                                      ':'));
                                                          }
                                                      });
        auditDao = new CacheableRelationalDao<>(sessionFactories,
                                                Audit.class,
                                                new ShardCalculator<>(shardManager,
                                                                      new ConsistentHashBucketIdExtractor<>(shardManager)),
                                                new RelationalCache<Audit>() {

                                                    private Map<String, Object> cache = new HashMap<>();

                                                    @Override
                                                    public void put(String parentKey, Object key, Audit entity) {
                                                        cache.put(StringUtils.join(parentKey, key, ':'), entity);
                                                    }

                                                    @Override
                                                    public void put(String parentKey, List<Audit> entities) {
                                                        cache.put(parentKey, entities);
                                                    }

                                                    @Override
                                                    public void put(
                                                            String parentKey,
                                                            int first,
                                                            int numResults,
                                                            List<Audit> entities) {
                                                        cache.put(StringUtils.join(parentKey, first, numResults, ':'),
                                                                  entities);
                                                    }

                                                    @Override
                                                    public boolean exists(String parentKey, Object key) {
                                                        return cache.containsKey(StringUtils.join(parentKey, key, ':'));
                                                    }

                                                    @Override
                                                    public Audit get(String parentKey, Object key) {
                                                        return (Audit) cache.get(StringUtils.join(parentKey, key, ':'));
                                                    }

                                                    @Override
                                                    public List<Audit> select(String parentKey) {
                                                        return (List<Audit>) cache.get(parentKey);
                                                    }

                                                    @Override
                                                    public List<Audit> select(
                                                            String parentKey,
                                                            int first,
                                                            int numResults) {
                                                        return (List<Audit>) cache.get(StringUtils.join(parentKey,
                                                                                                        first,
                                                                                                        numResults,
                                                                                                        ':'));
                                                    }
                                                });
    }

    @After
    public void after() {
        sessionFactories.forEach(SessionFactory::close);
    }

    @Test
    public void testSave() throws Exception {
        TestEntity testEntity = TestEntity.builder()
                .externalId("testId")
                .text("Some Text")
                .build();
        lookupDao.save(testEntity);

        assertEquals(true, lookupDao.exists("testId"));
        assertEquals(false, lookupDao.exists("testId1"));
        Optional<TestEntity> result = lookupDao.get("testId");
        assertEquals("Some Text",
                     result.get()
                             .getText());

        testEntity.setText("Some New Text");
        lookupDao.save(testEntity);
        result = lookupDao.get("testId");
        assertEquals("Some New Text",
                     result.get()
                             .getText());

        boolean updateStatus = lookupDao.update("testId", entity -> {
            if (entity.isPresent()) {
                TestEntity e = entity.get();
                e.setText("Updated text");
                return e;
            }
            return null;
        });

        assertTrue(updateStatus);
        result = lookupDao.get("testId");
        assertEquals("Updated text",
                     result.get()
                             .getText());

        updateStatus = lookupDao.update("testIdxxx", entity -> {
            if (entity.isPresent()) {
                TestEntity e = entity.get();
                e.setText("Updated text");
                return e;
            }
            return null;
        });

        assertFalse(updateStatus);
    }

    @Test
    public void testScatterGather() throws Exception {
        List<TestEntity> results = lookupDao.scatterGather(DetachedCriteria.forClass(TestEntity.class)
                                                                   .add(Restrictions.eq("externalId", "testId")));
        assertTrue(results.isEmpty());

        TestEntity testEntity = TestEntity.builder()
                .externalId("testId")
                .text("Some Text")
                .build();
        lookupDao.save(testEntity);
        results = lookupDao.scatterGather(DetachedCriteria.forClass(TestEntity.class)
                                                  .add(Restrictions.eq("externalId", "testId")));
        assertFalse(results.isEmpty());
        assertEquals("Some Text",
                     results.get(0)
                             .getText());
    }

    @Test
    public void testSaveInParentBucket() throws Exception {
        final String phoneNumber = "9830968020";

        Phone phone = Phone.builder()
                .phone(phoneNumber)
                .build();

        Phone savedPhone = phoneDao.save(phone)
                .get();

        Transaction transaction = Transaction.builder()
                .transactionId("testTxn")
                .to("9830703153")
                .amount(100)
                .phone(savedPhone)
                .build();

        transactionDao.save(savedPhone.getPhone(), transaction)
                .get();
        {
            Transaction resultTx = transactionDao.get(phoneNumber, "testTxn")
                    .get();
            assertEquals(phoneNumber,
                         resultTx.getPhone()
                                 .getPhone());
            assertTrue(transactionDao.exists(phoneNumber, "testTxn"));
            assertFalse(transactionDao.exists(phoneNumber, "testTxn1"));
        }
        {
            Optional<Transaction> resultTx = transactionDao.get(phoneNumber, "testTxn1");
            assertFalse(resultTx.isPresent());
        }
        saveAudit(phoneNumber, "testTxn", "Started");
        saveAudit(phoneNumber, "testTxn", "Underway");
        saveAudit(phoneNumber, "testTxn", "Completed");

        assertEquals(3, auditDao.count(phoneNumber, DetachedCriteria.forClass(Audit.class)
                .add(Restrictions.eq("transaction.transactionId", "testTxn"))));

        List<Audit> audits = auditDao.select(phoneNumber, DetachedCriteria.forClass(Audit.class)
                .add(Restrictions.eq("transaction.transactionId", "testTxn")), 0, 10);
        assertEquals("Started",
                     audits.get(0)
                             .getText());

    }

    private void saveAudit(String phone, String transaction, String text) throws Exception {
        auditDao.save(phone, Audit.builder()
                .text(text)
                .transaction(Transaction.builder()
                                     .transactionId(transaction)
                                     .build())
                .build());
    }

    @Test
    public void testHierarchy() throws Exception {
        final String phoneNumber = "9986032019";
        saveHierarchy(phoneNumber);
        saveHierarchy("9986402019");

        List<Audit> audits = auditDao.select(phoneNumber, DetachedCriteria.forClass(Audit.class)
                .add(Restrictions.eq("transaction.transactionId", "newTxn-" + phoneNumber)), 0, 10);

        assertEquals(2, audits.size());

        List<Audit> allAudits = auditDao.scatterGather(DetachedCriteria.forClass(Audit.class), 0, 10);
        assertEquals(4, allAudits.size());
    }


    private void saveHierarchy(String phone) throws Exception {

        Transaction transaction = Transaction.builder()
                .transactionId("newTxn-" + phone)
                .amount(100)
                .to("9986402019")
                .build();

        Audit started = Audit.builder()
                .text("Started")
                .transaction(transaction)
                .build();

        Audit completed = Audit.builder()
                .text("Completed")
                .transaction(transaction)
                .build();

        transaction.setAudits(ImmutableList.of(started, completed));

        transactionDao.save(phone, transaction);
    }
}

<|start_filename|>src/main/java/io/appform/dropwizard/sharding/utils/ShardCalculator.java<|end_filename|>
/*
 * Copyright 2016 <NAME> <<EMAIL>>
 *
 * 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.appform.dropwizard.sharding.utils;

import io.appform.dropwizard.sharding.sharding.BucketIdExtractor;
import io.appform.dropwizard.sharding.sharding.ShardManager;
import lombok.extern.slf4j.Slf4j;

/**
 * Utility class for calculating shards.
 */
@Slf4j
public class ShardCalculator<T> {

    private final ShardManager shardManager;
    private final BucketIdExtractor<T> extractor;

    public ShardCalculator(ShardManager shardManager, BucketIdExtractor<T> extractor) {
        this.shardManager = shardManager;
        this.extractor = extractor;
    }

    public int shardId(T key) {
        int bucketId = extractor.bucketId(key);
        return shardManager.shardForBucket(bucketId);
    }

    public boolean isOnValidShard(T key) {
        int bucketId = extractor.bucketId(key);
        return shardManager.isMappedToValidShard(bucketId);
    }
}


<|start_filename|>src/main/java/io/appform/dropwizard/sharding/BalancedDBShardingBundle.java<|end_filename|>
/*
 * Copyright 2019 <NAME> <<EMAIL>>
 *
 * 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.appform.dropwizard.sharding;

import io.appform.dropwizard.sharding.sharding.BalancedShardManager;
import io.appform.dropwizard.sharding.sharding.ShardBlacklistingStore;
import io.appform.dropwizard.sharding.sharding.ShardManager;
import io.dropwizard.Configuration;
import lombok.extern.slf4j.Slf4j;

import java.util.List;

/**
 * A dropwizard bundle that provides sharding over normal RDBMS.
 */
@Slf4j
public abstract class BalancedDBShardingBundle<T extends Configuration> extends DBShardingBundleBase<T> {

    public BalancedDBShardingBundle(
            String dbNamespace,
            Class<?> entity,
            Class<?>... entities) {
        super(dbNamespace, entity, entities);
    }

    public BalancedDBShardingBundle(String dbNamespace, List<String> classPathPrefixList) {
        super(dbNamespace, classPathPrefixList);
    }

    public BalancedDBShardingBundle(Class<?> entity, Class<?>... entities) {
        super(entity, entities);
    }

    public BalancedDBShardingBundle(String... classPathPrefixes) {
        super(classPathPrefixes);
    }

    @Override
    protected ShardManager createShardManager(int numShards, ShardBlacklistingStore shardBlacklistingStore) {
        return new BalancedShardManager(numShards, shardBlacklistingStore);
    }

}


<|start_filename|>src/main/java/io/appform/dropwizard/sharding/healthcheck/ShardHealthCheckMeta.java<|end_filename|>
package io.appform.dropwizard.sharding.healthcheck;

import com.codahale.metrics.health.HealthCheck;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;

@Data
@Builder
@NoArgsConstructor
@AllArgsConstructor
public class ShardHealthCheckMeta {

    private HealthCheck healthCheck;
    private int shardId;

}


<|start_filename|>src/test/java/io/appform/dropwizard/sharding/dao/RelationalDaoTest.java<|end_filename|>
/*
 * Copyright 2016 <NAME> <<EMAIL>>
 *
 * 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.appform.dropwizard.sharding.dao;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import io.appform.dropwizard.sharding.dao.testdata.entities.RelationalEntity;
import io.appform.dropwizard.sharding.sharding.BalancedShardManager;
import io.appform.dropwizard.sharding.sharding.ShardManager;
import io.appform.dropwizard.sharding.sharding.impl.ConsistentHashBucketIdExtractor;
import io.appform.dropwizard.sharding.utils.ShardCalculator;
import lombok.val;
import org.hibernate.SessionFactory;
import org.hibernate.boot.registry.StandardServiceRegistry;
import org.hibernate.boot.registry.StandardServiceRegistryBuilder;
import org.hibernate.cfg.Configuration;
import org.hibernate.criterion.DetachedCriteria;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import java.util.List;
import java.util.UUID;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;

public class RelationalDaoTest {

    private List<SessionFactory> sessionFactories = Lists.newArrayList();
    private RelationalDao<RelationalEntity> relationalDao;

    private SessionFactory buildSessionFactory(String dbName) {
        Configuration configuration = new Configuration();
        configuration.setProperty("hibernate.dialect",
                                  "org.hibernate.dialect.H2Dialect");
        configuration.setProperty("hibernate.connection.driver_class",
                                  "org.h2.Driver");
        configuration.setProperty("hibernate.connection.url", "jdbc:h2:mem:" + dbName);
        configuration.setProperty("hibernate.hbm2ddl.auto", "create");
        configuration.setProperty("hibernate.current_session_context_class", "managed");
        configuration.addAnnotatedClass(RelationalEntity.class);

        StandardServiceRegistry serviceRegistry
                = new StandardServiceRegistryBuilder().applySettings(
                configuration.getProperties())
                .build();
        return configuration.buildSessionFactory(serviceRegistry);
    }

    @Before
    public void before() {
        for (int i = 0; i < 16; i++) {
            sessionFactories.add(buildSessionFactory(String.format("db_%d", i)));
        }
        final ShardManager shardManager = new BalancedShardManager(sessionFactories.size());
        relationalDao = new RelationalDao<>(sessionFactories,
                                            RelationalEntity.class,
                                            new ShardCalculator<>(shardManager,
                                                                  new ConsistentHashBucketIdExtractor<>(shardManager)));
    }

    @After
    public void after() {
        sessionFactories.forEach(SessionFactory::close);
    }

    @Test
    public void testBulkSave() throws Exception {
        String key = "testPhone";
        RelationalEntity entityOne = RelationalEntity.builder()
                .key("1")
                .value("abcd")
                .build();
        RelationalEntity entityTwo = RelationalEntity.builder()
                .key("2")
                .value("abcd")
                .build();
        relationalDao.saveAll(key, Lists.newArrayList(entityOne, entityTwo));
        List<RelationalEntity> entities = relationalDao.select(key,
                                                               DetachedCriteria.forClass(RelationalEntity.class),
                                                               0,
                                                               10);
        assertEquals(2, entities.size());

    }

    @Test
    public void testUpdateUsingQuery() throws Exception {
        val relationalKey = UUID.randomUUID().toString();

        val entityOne = RelationalEntity.builder()
                .key("1")
                .keyTwo("1")
                .value(UUID.randomUUID().toString())
                .build();
        relationalDao.save(relationalKey, entityOne);

        val entityTwo = RelationalEntity.builder()
                .key("2")
                .keyTwo("2")
                .value(UUID.randomUUID().toString())
                .build();
        relationalDao.save(relationalKey, entityTwo);

        val entityThree = RelationalEntity.builder()
                .key("3")
                .keyTwo("2")
                .value(UUID.randomUUID().toString())
                .build();
        relationalDao.save(relationalKey, entityThree);


        val newValue = UUID.randomUUID().toString();
        int rowsUpdated = relationalDao.updateUsingQuery(relationalKey,
                UpdateOperationMeta.builder()
                        .queryName("testUpdateUsingKeyTwo")
                        .params(ImmutableMap.of("keyTwo", "2",
                                "value", newValue))
                        .build()
        );
        assertEquals(2, rowsUpdated);

        val persistedEntityTwo = relationalDao.get(relationalKey, "2").orElse(null);
        assertNotNull(persistedEntityTwo);
        assertEquals(newValue, persistedEntityTwo.getValue());

        val persistedEntityThree = relationalDao.get(relationalKey, "3").orElse(null);
        assertNotNull(persistedEntityThree);
        assertEquals(newValue, persistedEntityThree.getValue());


    }

    @Test
    public void testUpdateUsingQueryNoRowUpdated() throws Exception {
        val relationalKey = UUID.randomUUID().toString();

        val entityOne = RelationalEntity.builder()
                .key("1")
                .keyTwo("1")
                .value(UUID.randomUUID().toString())
                .build();
        relationalDao.save(relationalKey, entityOne);

        val entityTwo = RelationalEntity.builder()
                .key("2")
                .keyTwo("2")
                .value(UUID.randomUUID().toString())
                .build();
        relationalDao.save(relationalKey, entityTwo);

        val entityThree = RelationalEntity.builder()
                .key("3")
                .keyTwo("2")
                .value(UUID.randomUUID().toString())
                .build();
        relationalDao.save(relationalKey, entityThree);


        val newValue = UUID.randomUUID().toString();
        int rowsUpdated = relationalDao.updateUsingQuery(relationalKey,
                UpdateOperationMeta.builder()
                        .queryName("testUpdateUsingKeyTwo")
                        .params(ImmutableMap.of("keyTwo", UUID.randomUUID().toString(),
                                "value", newValue))
                        .build()
        );
        assertEquals(0, rowsUpdated);


        val persistedEntityOne = relationalDao.get(relationalKey, "1").orElse(null);
        assertNotNull(persistedEntityOne);
        assertEquals(entityOne.getValue(), persistedEntityOne.getValue());

        val persistedEntityTwo = relationalDao.get(relationalKey, "2").orElse(null);
        assertNotNull(persistedEntityTwo);
        assertEquals(entityTwo.getValue(), persistedEntityTwo.getValue());

        val persistedEntityThree = relationalDao.get(relationalKey, "3").orElse(null);
        assertNotNull(persistedEntityThree);
        assertEquals(entityThree.getValue(), persistedEntityThree.getValue());
    }
}

<|start_filename|>src/test/java/io/appform/dropwizard/sharding/sharding/ShardManagerTest.java<|end_filename|>
/*
 * Copyright 2016 <NAME> <<EMAIL>>
 *
 * 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.appform.dropwizard.sharding.sharding;

import io.appform.dropwizard.sharding.exceptions.ShardBlacklistedException;
import org.junit.Assert;
import org.junit.Test;

import static org.junit.Assert.assertEquals;

public class ShardManagerTest {

   @Test
    public void testShardForBucket() throws Exception {
        ShardManager shardManager = new LegacyShardManager(8);
        assertEquals(0, shardManager.shardForBucket(100));
    }

    @Test
    public void testBlacklisting() throws Exception {
        ShardManager shardManager = new LegacyShardManager(8);
        assertEquals(0, shardManager.shardForBucket(100));
        shardManager.blacklistShard(0);
        Thread.sleep(5000); //Sleeping for 5 seconds
        try {
            shardManager.shardForBucket(100);
            Assert.fail("Should have errored out");
        } catch (ShardBlacklistedException e) {

        }
        shardManager.unblacklistShard(0);
        Thread.sleep(5000); //Sleeping for 5 seconds
        assertEquals(0, shardManager.shardForBucket(100));
    }
}

<|start_filename|>src/test/java/io/appform/dropwizard/sharding/dao/testdata/entities/RelationalEntity.java<|end_filename|>
package io.appform.dropwizard.sharding.dao.testdata.entities;

import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;

import javax.persistence.Column;
import javax.persistence.Entity;
import javax.persistence.Id;
import javax.persistence.NamedQueries;
import javax.persistence.NamedQuery;
import javax.persistence.Table;

@Entity
@Data
@NoArgsConstructor
@AllArgsConstructor
@Builder
@Table(name = "relations")
@NamedQueries({
        @NamedQuery(name = "testUpdateUsingKeyTwo", query = "update RelationalEntity set value = :value where keyTwo =:keyTwo")})
public class RelationalEntity {

    @Id
    @Column(name = "key", nullable = false, unique = true)
    private String key;

    private String keyTwo;

    private String value;

}


<|start_filename|>src/main/java/io/appform/dropwizard/sharding/healthcheck/BlacklistingAwareHealthCheck.java<|end_filename|>

package io.appform.dropwizard.sharding.healthcheck;

import com.codahale.metrics.health.HealthCheck;
import io.appform.dropwizard.sharding.config.BlacklistConfig;
import io.appform.dropwizard.sharding.sharding.ShardManager;
import lombok.extern.slf4j.Slf4j;


@Slf4j
public class BlacklistingAwareHealthCheck extends HealthCheck {

    private final int shardId;
    private final HealthCheck baseHealthCheck;
    private final ShardManager shardManager;
    private final BlacklistConfig blacklistConfig;

    public BlacklistingAwareHealthCheck(final int shardId,
                                        final HealthCheck baseHealthCheck,
                                        final ShardManager shardManager,
                                        final BlacklistConfig blacklistConfig) {
        this.shardId = shardId;
        this.baseHealthCheck = baseHealthCheck;
        this.shardManager = shardManager;
        this.blacklistConfig = blacklistConfig;
    }

    @Override
    protected Result check() {
        if (blacklistConfig.isSkipNativeHealthcheck()) {
            return Result.healthy();
        }

        if (shardManager.isBlacklisted(shardId)) {
            log.info("returning healthy since shard is blacklisted [{}]", shardId);
            return Result.healthy();
        }

        final Result result = baseHealthCheck.execute();
        log.debug("DBSharding HealthCheck for shardId: {}, Status: {}, {}", shardId,
                result.isHealthy(), result.getMessage());
        return result;
    }
}

<|start_filename|>src/test/java/io/appform/dropwizard/sharding/BalancedDbShardingBundleWithMultipleClassPath.java<|end_filename|>
/*
 * Copyright 2019 <NAME> <<EMAIL>>
 *
 * 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.appform.dropwizard.sharding;

import io.appform.dropwizard.sharding.config.ShardedHibernateFactory;
import io.appform.dropwizard.sharding.dao.LookupDao;
import io.appform.dropwizard.sharding.dao.testdata.entities.TestEntity;
import io.appform.dropwizard.sharding.dao.testdata.multi.MultiPackageTestEntity;
import org.junit.Assert;
import org.junit.Test;

import java.util.Optional;

public class BalancedDbShardingBundleWithMultipleClassPath extends DBShardingBundleTestBase {



    @Override
    protected DBShardingBundleBase<TestConfig> getBundle() {
        return new BalancedDBShardingBundle<TestConfig>("io.appform.dropwizard.sharding.dao.testdata.entities", "io.appform.dropwizard.sharding.dao.testdata.multi") {
            @Override
            protected ShardedHibernateFactory getConfig(TestConfig config) {
                return testConfig.getShards();
            }
        };
    }


    @Test
    public void testMultiPackage() throws Exception {

        DBShardingBundleBase<TestConfig> bundle = getBundle();

        bundle.initialize(bootstrap);
        bundle.initBundles(bootstrap);
        bundle.runBundles(testConfig, environment);
        bundle.run(testConfig, environment);
        LookupDao<MultiPackageTestEntity> lookupDao = bundle.createParentObjectDao(MultiPackageTestEntity.class);

        MultiPackageTestEntity multiPackageTestEntity = MultiPackageTestEntity.builder()
                .text("Testing multi package scanning")
                .lookup("123")
                .build();

        Optional<MultiPackageTestEntity> saveMultiPackageTestEntity = lookupDao.save(multiPackageTestEntity);
        Assert.assertEquals(multiPackageTestEntity.getText(), saveMultiPackageTestEntity.get().getText());

        Optional<MultiPackageTestEntity> fetchedMultiPackageTestEntity = lookupDao.get(multiPackageTestEntity.getLookup());
        Assert.assertEquals(saveMultiPackageTestEntity.get().getText(), fetchedMultiPackageTestEntity.get().getText());

        LookupDao<TestEntity> testEntityLookupDao = bundle.createParentObjectDao(TestEntity.class);

        TestEntity testEntity = TestEntity.builder()
                .externalId("E123")
                .text("Test Second Package")
                .build();
        Optional<TestEntity> savedTestEntity = testEntityLookupDao.save(testEntity);
        Assert.assertEquals(testEntity.getText(), savedTestEntity.get().getText());

        Optional<TestEntity> fetchedTestEntity = testEntityLookupDao.get(testEntity.getExternalId());
        Assert.assertEquals(savedTestEntity.get().getText(), fetchedTestEntity.get().getText());


    }
}
