HiveShardSelector.java

package org.hivedb.hibernate;

import org.hibernate.shards.ShardId;
import org.hibernate.shards.strategy.selection.ShardSelectionStrategy;
import org.hivedb.Hive;
import org.hivedb.HiveLockableException;
import org.hivedb.HiveRuntimeException;
import org.hivedb.configuration.EntityConfig;
import org.hivedb.configuration.EntityHiveConfig;
import org.hivedb.util.classgen.ReflectionTools;
import org.hivedb.util.functional.Atom;
import org.hivedb.util.functional.Transform;
import org.hivedb.util.functional.Unary;

import java.util.Collection;

public class HiveShardSelector implements ShardSelectionStrategy {
  private EntityHiveConfig hiveConfig;
  private Hive hive;

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

  // The Hive HAS to be responsible for shard allocation
  public ShardId selectShardIdForNewObject(Object entity) {
    EntityConfig config = hiveConfig.getEntityConfig(resolveEntityConfigClass(entity.getClass()));

    if (!hive.directory().doesPrimaryIndexKeyExist(config.getPrimaryIndexKey(entity)))
      try {
        hive.directory().insertPrimaryIndexKey(config.getPrimaryIndexKey(entity));
      } catch (HiveLockableException e) {
        throw new HiveRuntimeException(e.getMessage(), e);
      }

    Collection<Integer> nodeIds =
      hive.directory().getNodeIdsOfPrimaryIndexKey(config.getPrimaryIndexKey(entity));

    return Atom.getFirstOrThrow(Transform.map(HiveShardResolver.nodeIdToShardIdConverter(), nodeIds));
  }

  @SuppressWarnings("unchecked")
  private Class<?> resolveEntityConfigClass(Class<?> clazz) {
    return ReflectionTools.whichIsImplemented(
      clazz,
      Transform.map(new Unary<EntityConfig, Class>() {
        public Class f(EntityConfig entityConfig) {
          return entityConfig.getRepresentedInterface();
        }
      },
        hiveConfig.getEntityConfigs()));
  }
}