ShardedAvroBaseTest.java

package avrobase.shard;

import avrobase.AvroBase;
import avrobase.AvroBaseException;
import avrobase.AvroFormat;
import avrobase.Row;
import avrobase.mysql.KeyStrategy;
import avrobase.mysql.MysqlAB;
import bagcheck.User;
import com.jolbox.bonecp.BoneCPConfig;
import com.jolbox.bonecp.BoneCPDataSource;
import org.apache.avro.Schema;
import org.apache.avro.util.Utf8;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Test;
import thepusher.Pusher;
import thepusher.PusherBase;

import javax.sql.DataSource;
import java.nio.ByteBuffer;
import java.security.SecureRandom;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Random;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import static junit.framework.Assert.assertEquals;

/**
 * Test the system
 * <p/>
 * User: sam
 * Date: 10/10/10
 * Time: 11:03 AM
 */
@SuppressWarnings({"unchecked"})
public class ShardedAvroBaseTest {
  private static Random random = new SecureRandom();
  private static final KeyStrategy<String> KEYTX = new KeyStrategy<String>() {

    @Override
    public byte[] toBytes(String key) {
      return key.getBytes();
    }

    @Override
    public String fromBytes(byte[] row) {
      return new String(row);
    }

    @Override
    public String fromString(String key) {
      return key;
    }

    @Override
    public String toString(String row) {
      return row;
    }

    @Override
    public String newKey() {
      return String.valueOf(random.nextLong());
    }
  };
  private static final Comparator<String> STRING_COMPARATOR = new Comparator<String>() {
    public int compare(String s, String s1) {
      return s.compareTo(s1);
    }
  };

  private static ShardableMysqlAB mab1;
  private static ShardableMysqlAB mab2;
  private static ShardableMysqlAB mab3;
  private static ShardableMysqlAB mab4;
  private static AvroBase[] AVRO_BASES;

  @BeforeClass
  public static void setup() {
    try {
      Class.forName("com.mysql.jdbc.Driver");
    } catch (ClassNotFoundException e) {
      throw new AssertionError("Could not find JDBC driver: " + e);
    }

    final BoneCPConfig config = new BoneCPConfig();
    config.setJdbcUrl("jdbc:mysql://localhost:3306/shardedtest");
    config.setMaxConnectionsPerPartition(10);
    config.setMinConnectionsPerPartition(2);
    config.setPartitionCount(1);
    config.setLazyInit(true);
    config.setUsername("sam");
    config.setPassword("");
    BoneCPDataSource ds = new BoneCPDataSource(config);
    ExecutorService es = Executors.newCachedThreadPool();

    String MYSQL_SCHEMAS = "avro_schemas";
    mab1 = new ShardableMysqlAB(es, ds, "user", "profile", MYSQL_SCHEMAS, User.SCHEMA$, AvroFormat.JSON, KEYTX);
    mab2 = new ShardableMysqlAB(es, ds, "user2", "profile", MYSQL_SCHEMAS, User.SCHEMA$, AvroFormat.JSON, KEYTX);
    mab3 = new ShardableMysqlAB(es, ds, "user3", "profile", MYSQL_SCHEMAS, User.SCHEMA$, AvroFormat.JSON, KEYTX);
    mab4 = new ShardableMysqlAB(es, ds, "user4", "profile", MYSQL_SCHEMAS, User.SCHEMA$, AvroFormat.JSON, KEYTX);
    AVRO_BASES = new AvroBase[] {mab1, mab2, mab3, mab4};

    //noinspection unchecked
    for (AvroBase<User, String> ab : AVRO_BASES) {
      for (Row<User, String> userRow : ab.scan(null, null)) {
        ab.delete(userRow.row);
      }
    }
  }

  public static class ShardableMysqlAB extends MysqlAB<User, String> implements ShardableAvroBase<User, String> {
    public ShardableMysqlAB(ExecutorService es, DataSource datasource, java.lang.String table, java.lang.String family, java.lang.String schemaTable, Schema schema, AvroFormat storageFormat, KeyStrategy<String> keytx) throws AvroBaseException {
      super(es, datasource, table, family, schemaTable, schema, storageFormat, keytx);
    }

    @Override
    public byte[] representation() {
      return new byte[0];
    }

    @Override
    public void init(byte[] representation) {
    }

    @Override
    public Iterable<String> scanKeys(String start, String end) throws AvroBaseException {
      final byte[] startRow = start == null ? null : keytx.toBytes(start);
      final byte[] stopRow = end == null ? null : keytx.toBytes(end);
      StringBuilder statement = new StringBuilder("SELECT row FROM ");
      statement.append(mysqlTableName);
      if (startRow != null) {
        statement.append(" WHERE row >= ?");
      }
      if (stopRow != null) {
        if (startRow == null) {
          statement.append(" WHERE row < ?");
        } else {
          statement.append(" AND row < ?");
        }
      }
      return new Query<Iterable<String>>(datasource, statement.toString()) {
        public void setup(PreparedStatement ps) throws AvroBaseException, SQLException {
          int i = 1;
          if (startRow != null) {
            ps.setBytes(i++, startRow);
          }
          if (stopRow != null) {
            ps.setBytes(i, stopRow);
          }
        }

        public Iterable<String> execute(final ResultSet rs) throws AvroBaseException, SQLException {
          // TODO: Can't stream this yet due to database cursors
          List<String> rows = new ArrayList<String>();
          while (rs.next()) {
            byte[] row = rs.getBytes(1);
              rows.add(keytx.fromBytes(row));
          }
          return rows;
        }
      }.query();

    }
  }

  @Test
  public void create() {
    Pusher<SC> pusher = PusherBase.create(SC.class, Inject.class);
    pusher.bindInstance(SC.KEY_COMPARATOR, STRING_COMPARATOR);
    pusher.bindClass(SC.STRATEGY, ShardingStrategy.Partition.class);
    @SuppressWarnings({"unchecked"}) ShardedAvroBase<User, String> sab = pusher.create(ShardedAvroBase.class);
    sab.addShard(mab1, 1.0, false);

    User user = getUser();
    mab1.put("test", user);
    Row<User, String> test = mab1.get("test");
    assertEquals(user, test.value);
    mab1.delete("test");
  }

  @Test
  public void shard() {
    Pusher<SC> pusher = PusherBase.create(SC.class, Inject.class);
    pusher.bindInstance(SC.KEY_COMPARATOR, STRING_COMPARATOR);
    pusher.bindClass(SC.STRATEGY, ShardingStrategy.Partition.class);
    @SuppressWarnings({"unchecked"}) ShardedAvroBase<User, String> sab = pusher.create(ShardedAvroBase.class);
    sab.addShard(mab1, 1.0, false);

    int T = 1000;
    for (int i = 0; i < T; i++) {
      User user = getUser();
      String row = KEYTX.newKey();
      user.firstName = new Utf8(user.firstName.toString() + row);
      sab.put(row, user);
    }

    sab.addShard(mab2, 3.0, true);

    int count = 0;
    for (String tRow : mab1.scanKeys(null, null)) {
      count++;
    }
    assertEquals(T/4, count);
    for (String tRow : mab2.scanKeys(null, null)) {
      count++;
    }
    assertEquals(T, count);

    sab.addShard(mab3, 1.0, true);

    count = 0;
    for (String tRow : mab1.scanKeys(null, null)) {
      count++;
    }
    assertEquals(T/5, count);
    for (String tRow : mab2.scanKeys(null, null)) {
      count++;
    }
    assertEquals(T/5*4, count);
    for (String tRow : mab3.scanKeys(null, null)) {
      count++;
    }
    assertEquals(T, count);

    sab.addShard(mab4, 5.0, true);

    for (int i = 0; i < T; i++) {
      User user = getUser();
      String row = KEYTX.newKey();
      user.firstName = new Utf8(user.firstName.toString() + row);
      sab.put(row, user);
    }

    count = 0;
    for (String tRow : mab1.scanKeys(null, null)) {
      count++;
    }
    for (String tRow : mab2.scanKeys(null, null)) {
      count++;
    }
    for (String tRow : mab3.scanKeys(null, null)) {
      count++;
    }
    for (String tRow : mab4.scanKeys(null, null)) {
      count++;
    }
    assertEquals(T*2, count);
  }

  @AfterClass
  public static void teardown() {
    for (AvroBase<User, String> ab : AVRO_BASES) {
      for (Row<User, String> userRow : ab.scan(null, null)) {
        ab.delete(userRow.row);
      }
    }
  }

  private User getUser() {
    User user = new User();
    user.email = $("spullara@yahoo.com");
    user.firstName = $("Sam");
    user.lastName = $("Pullara");
    user.image = $("");
    user.password = ByteBuffer.allocate(0);
    return user;
  }

  Utf8 $(String s) {
    return new Utf8(s);
  }

}