HiveShardResolver.java

package org.hivedb.hibernate;


import org.hibernate.shards.ShardId;
import org.hibernate.shards.strategy.resolution.ShardResolutionStrategy;
import org.hibernate.shards.strategy.selection.ShardResolutionStrategyData;
import org.hibernate.shards.util.Lists;
import org.hivedb.Hive;
import org.hivedb.configuration.EntityConfig;
import org.hivedb.configuration.EntityHiveConfig;
import org.hivedb.configuration.EntityIndexConfig;
import org.hivedb.util.classgen.ReflectionTools;
import org.hivedb.util.functional.Transform;
import org.hivedb.util.functional.Unary;

import java.util.Collection;
import java.util.List;

public class HiveShardResolver implements ShardResolutionStrategy {
  private EntityHiveConfig hiveConfig;
  private Hive hive;

  public HiveShardResolver(EntityHiveConfig hiveConfig, Hive hive) {
    this.hiveConfig = hiveConfig;
    this.hive = hive;
  }

  @SuppressWarnings("unchecked")
  public List<ShardId> selectShardIdsFromShardResolutionStrategyData(ShardResolutionStrategyData data) {

    final Class clazz;
    try {
      clazz = Class.forName(data.getEntityName());
    } catch (ClassNotFoundException e) {
      throw new RuntimeException(e);
    }
    final Class<?> resolvedEntityInterface = new EntityResolver(hiveConfig).resolveToEntityOrRelatedEntity(clazz);
    if (resolvedEntityInterface != null) {
      EntityConfig config = hiveConfig.getEntityConfig(resolvedEntityInterface);
      Collection<Integer> ids = (config.isPartitioningResource())
        ? hive.directory().getNodeIdsOfPrimaryIndexKey(data.getId())
        : hive.directory().getNodeIdsOfResourceId(config.getResourceName(), data.getId());
      return Lists.newArrayList(Transform.map(nodeIdToShardIdConverter(), ids));
    } else {
      // Only for Hibernate entities that are not Hive entities, but are used for secondary indexes
      for (EntityConfig entityConfig : hiveConfig.getEntityConfigs()) {
        for (EntityIndexConfig entityIndexConfig : entityConfig.getEntityIndexConfigs()) {
          if (ReflectionTools.isComplexCollectionItemProperty(entityConfig.getRepresentedInterface(), entityIndexConfig.getPropertyName())
            && ReflectionTools.doesImplementOrExtend(clazz, ReflectionTools.getCollectionItemType(entityConfig.getRepresentedInterface(), entityIndexConfig.getPropertyName()))) {
            Collection<Integer> ids = hive.directory().getNodeIdsOfSecondaryIndexKey(
              entityConfig.getResourceName(),
              hive.getPartitionDimension().getResource(entityConfig.getResourceName()).getSecondaryIndex(entityIndexConfig.getIndexName()).getName(),
              data.getId());
            return Lists.newArrayList(Transform.map(nodeIdToShardIdConverter(), ids));
          }
        }
      }
    }
    throw new RuntimeException(String.format("Could not resolve class to a Hive entity nor a secondary index Hibernate entity: %s", clazz.getSimpleName()));
  }

  public static Unary<Integer, ShardId> nodeIdToShardIdConverter() {
    return new Unary<Integer, ShardId>() {
      public ShardId f(Integer id) {
        return new ShardId(id);
      }
    };
  }
}