import java.io.File;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.hibernate.SessionFactory;
import org.hibernate.cfg.AnnotationConfiguration;
import org.hibernate.cfg.Configuration;
import org.hibernate.shards.ShardId;
import org.hibernate.shards.ShardedConfiguration;
import org.hibernate.shards.cfg.ConfigurationToShardConfigurationAdapter;
import org.hibernate.shards.cfg.ShardConfiguration;
import org.hibernate.shards.strategy.ShardStrategy;
import org.hibernate.shards.strategy.ShardStrategyFactory;
import org.hibernate.shards.strategy.ShardStrategyImpl;
import org.hibernate.shards.strategy.access.SequentialShardAccessStrategy;
import org.hibernate.shards.strategy.resolution.AllShardsShardResolutionStrategy;

public class StandardSessionFactoryHelper {

    private final File TS_ENTERTAINMENT_FILE =
        new File("src/main/resources/news/shard.sports_entertainment.hibernate.cfg.xml");
    
    private final File BUSINESS_SCIENCES_WORLD_FILE = 
        new File("src/main/resources/news/shard.business_sciences_world.cfg.xml");

    public SessionFactory createSessionFactory() {

        AnnotationConfiguration prototypeConfig = new AnnotationConfiguration()
                .configure(SPORTS_ENTERTAINMENT_FILE)
                .addAnnotatedClass(Article.class);

        ShardedConfiguration shardedConfig = new ShardedConfiguration(
                prototypeConfig,
                buildPhysicalShardConfigurations(),
                buildShardStrategyFactory(),
                buildVirtualShardMap());

        return shardedConfig.buildShardedSessionFactory();
    }

    private List buildPhysicalShardConfigurations() {
        return
            Arrays.asList(
                buildShardConfig(SPORTS_ENTERTAINMENT_FILE),
                buildShardConfig(BUSINESS_SCIENCES_WORLD_FILE)
            );
    }

    private Map buildVirtualShardMap() {
        final int sportsEntertainmentPhysicalShardId = 0;
        final int businessSciencesWorldPhysicalShardId = 1;

        Map virtualShardMap = new HashMap();

        virtualShardMap.put( 
                getVirtualShardId(Category.ENTERTAINMENT), 
                                  sportsEntertainmentPhysicalShardId);        
        virtualShardMap.put( 
                getVirtualShardId(Category.SPORTS),             
                                  sportsEntertainmentPhysicalShardId);

        virtualShardMap.put( 
                getVirtualShardId(Category.BUSINESS),   
                                  businessSciencesWorldPhysicalShardId);        
        virtualShardMap.put( 
                getVirtualShardId(Category.HEALTH),
                                  businessSciencesWorldPhysicalShardId);
        virtualShardMap.put( 
                getVirtualShardId(Category.SCIENCE),    
                                  businessSciencesWorldPhysicalShardId);
        virtualShardMap.put( 
                getVirtualShardId(Category.TECHNOLOGY), 
                                  businessSciencesWorldPhysicalShardId);
        virtualShardMap.put( 
                getVirtualShardId(Category.WORLD    
                                  businessSciencesWorldPhysicalShardId);

        return Collections.unmodifiableMap(virtualShardMap);
    }

    private Integer getVirtualShardId(Category category) {
        return new CategoryShardSelector().getShardIdForCategory(category).getId();
    }

    private ShardStrategyFactory buildShardStrategyFactory() {
        return
            new ShardStrategyFactory() {
                public ShardStrategy newShardStrategy(List shardIds) {
                    return new ShardStrategyImpl(
                            new CategoryShardSelector(),
                            new AllShardsShardResolutionStrategy(shardIds),
                            new SequentialShardAccessStrategy());
                }
            };
    }

    private ShardConfiguration buildShardConfig(File configFile) {
        return
            new ConfigurationToShardConfigurationAdapter(
                    new Configuration().configure(configFile) );
    }

}	

In our example the application is sharding based on the category of an article. The list of categories is fairly stable so we can use an enum to enumerate the categories. The mapping of categories and virtual shard IDs is maintained in the SessionFactoryHelper.

public enum Category {
    BUSINESS, ENTERTAINMENT, HEALTH, SCIENCE, SPORTS, TECHNOLOGY, WORLD
}

This is the main model/entity class for the application.

import java.math.BigInteger;
import java.util.Date;

import javax.persistence.Column;
import javax.persistence.Entity;
import javax.persistence.EnumType;
import javax.persistence.Enumerated;
import javax.persistence.GeneratedValue;
import javax.persistence.Id;
import javax.persistence.Table;

import org.hibernate.annotations.GenericGenerator;

@Entity
@Table(name = "articles")
public class Article {

    @Id
    @GeneratedValue( generator = "ArticleIdGenerator")
    @GenericGenerator( name = "ArticleIdGenerator",
                       strategy = "org.hibernate.shards.id.ShardedUUIDGenerator")
    @Column( name = "article_id" )
    private BigInteger articleId;

    @Enumerated(EnumType.STRING)
    @Column(name = "category")
    private Category category;

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

    @Column(name = "submit_time")
    private Date submitTime;

    // Getters and setters
}

This class is used by Hibernate Shards to figure out which shard to put an object on. For our case it's a pretty simple mapping from an Article class's category value to a shard ID.

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import org.hibernate.shards.ShardId;
import org.hibernate.shards.strategy.selection.ShardSelectionStrategy;

public class ContinentBasedShardSelector implements ShardSelectionStrategy {

    private final Map continentShardMap;

    public ContinentBasedShardSelector() {
        Map tempMap = new HashMap();
        tempMap.put(Continent.AFRICA,        new ShardId(0));
        tempMap.put(Continent.ANTARCTICA,    new ShardId(1));
        tempMap.put(Continent.ASIA,          new ShardId(2));
        tempMap.put(Continent.AUSTRALIA,     new ShardId(3));
        tempMap.put(Continent.EUROPE,        new ShardId(4));
        tempMap.put(Continent.NORTH_AMERICA, new ShardId(5));
        tempMap.put(Continent.SOUTH_AMERICA, new ShardId(6));

        continentShardMap = Collections.unmodifiableMap(tempMap);
    }

    @Override
    public ShardId selectShardIdForNewObject(Object obj) {
        if(obj instanceof WeatherReport) {
            return continentShardMap.get(((WeatherReport)obj).getContinent());
        }
        throw new IllegalArgumentException(
                "Can only shard " + WeatherReport.class.getName());
    }

}

This is a "driver" application to test our sharding setup. It saves an article for each category and then verifies that each article went to the correct physical database shard.

import static com.amazonaws.rds.samples.news.Category.BUSINESS;
import static com.amazonaws.rds.samples.news.Category.TECHNOLOGY;
import static com.amazonaws.rds.samples.news.Category.ENTERTAINMENT;
import static com.amazonaws.rds.samples.news.Category.SCIENCE;
import static com.amazonaws.rds.samples.news.Category.SPORTS;
import static com.amazonaws.rds.samples.news.Category.HEALTH;
import static com.amazonaws.rds.samples.news.Category.WORLD;
import static java.util.Arrays.asList;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.Date;
import java.util.List;

import org.hibernate.Session;
import org.hibernate.SessionFactory;
import org.hibernate.Transaction;
import org.junit.Assert;
import org.junit.Test;

public class ShardingWithRdsApplicationTest {

    private static final String JDBC_URL_SPORTS_ENTERTAINMENT = 
        "jdbc:mysql://my_sports_entertainment_db_url/news";
    private static final String JDBC_URL_BUSINESS_SCIENCES_WORLD = 
        "jdbc:mysql://my_business_sciences_world_db_url/news";
    
    private final String DATABASE_USERNAME = "my_username";
    private final String DATABASE_PASSWORD = "my_password";

    @Test
    public void testArticlesAreSavedToCorrectShards() throws Exception {

        StandardSessionFactoryHelper sessionFactoryHelper = 
            new StandardSessionFactoryHelper();
        SessionFactory sessionFactory = sessionFactoryHelper.createSessionFactory();

        Session session = sessionFactory.openSession();
        Transaction tx = session.beginTransaction();

        for (Category category : Category.values()) {
            session.save(createArticle(category));
        }

        tx.commit();

        session.close();

        verifyDatabaseContainsOnlyGivenCategories(
                asList(
BUSINESS, TECHNOLOGY, ENTERTAINMENT, SPORTS, WORLD),
                       "JDBC_URL_SPORTS_ENTERTAINMENT);

        verifyDatabaseContainsOnlyGivenCategories(
                asList(HEALTH, SCIENCE),
                       "JDBC_URL_BUSINESS_SCIENCES_WORLD);
    }

    private Article createArticle(Category category) {
        Article article = new Article();
        article.setCategory(category);
        article.setTitle("Article about " + category);
        article.setSubmitTime(new Date());
        return article;
    }

    private void verifyDatabaseContainsOnlyGivenCategories(
            List expectedCategories, String databaseUrl)
    throws Exception {

        Class.forName("com.mysql.jdbc.Driver");

        Connection connection = null;
        Statement statement = null;
        ResultSet rs = null;

        try {
            connection = DriverManager.getConnection(databaseUrl, 
                                                     DATABASE_USERNAME, 
                                                     DATABASE_PASSWORD);

            statement = connection.createStatement();
            rs = statement.executeQuery("select * from articles");

            while (rs.next()) {
                Category categoryInDatabase = 
                    Category.valueOf(rs.getString("categories"));
                Assert.assertTrue(expectedCategories.contains(categoryInDatabase));
            }
        }
        finally {
            if (rs != null) { rs.close(); }
            if (statement != null) { statement.close(); }
            if (connection != null) { connection.close(); }
        }
    }

}