/** Test inserts using KeyRange query */
@Test
public void testKeyRangeWrites() throws Exception {
Random random = new Random();
VtGate vtgate = VtGate.connect("localhost:" + testEnv.port, 0);
vtgate.begin();
String sql =
"insert into vtgate_test "
+ "(id, name, keyspace_id, age) "
+ "values (:id, :name, :keyspace_id, :age)";
int count = 20;
// Insert 20 rows per shard
for (String shardName : testEnv.shardKidMap.keySet()) {
List<KeyspaceId> kids = testEnv.getKeyspaceIds(shardName);
KeyspaceId minKid = Collections.min(kids);
KeyspaceId maxKid = Collections.max(kids);
KeyRange kr = new KeyRange(minKid, maxKid);
for (int i = 0; i < count; i++) {
KeyspaceId kid = kids.get(i % kids.size());
Query query =
new QueryBuilder(sql, testEnv.keyspace, "master")
.addBindVar(
BindVariable.forULong(
"id", UnsignedLong.valueOf("" + Math.abs(random.nextInt()))))
.addBindVar(BindVariable.forString("name", ("name_" + i)))
.addBindVar(BindVariable.forULong("keyspace_id", (UnsignedLong) kid.getId()))
.addBindVar(BindVariable.forNull("age"))
.addKeyRange(kr)
.build();
vtgate.execute(query);
}
}
vtgate.commit();
vtgate.close();
// Check 40 rows exist in total
vtgate = VtGate.connect("localhost:" + testEnv.port, 0);
String selectSql = "select * from vtgate_test";
Query allRowsQuery =
new QueryBuilder(selectSql, testEnv.keyspace, "master")
.setKeyspaceIds(testEnv.getAllKeyspaceIds())
.build();
Cursor cursor = vtgate.execute(allRowsQuery);
Assert.assertEquals(count * 2, cursor.getRowsAffected());
// Check 20 rows exist per shard
for (String shardName : testEnv.shardKidMap.keySet()) {
Query shardRows =
new QueryBuilder(selectSql, testEnv.keyspace, "master")
.setKeyspaceIds(testEnv.getKeyspaceIds(shardName))
.build();
cursor = vtgate.execute(shardRows);
Assert.assertEquals(count, cursor.getRowsAffected());
}
vtgate.close();
}
/** Test reads using Keyrange query */
@Test
public void testKeyRangeReads() throws Exception {
int rowsPerShard = 10;
// insert rows in each shard using ExecuteKeyspaceIds
for (String shardName : testEnv.shardKidMap.keySet()) {
Util.insertRowsInShard(testEnv, shardName, rowsPerShard);
}
VtGate vtgate = VtGate.connect("localhost:" + testEnv.port, 0);
String selectSql = "select * from vtgate_test";
// Check ALL KeyRange query returns rows from both shards
Query allRangeQuery =
new QueryBuilder(selectSql, testEnv.keyspace, "master").addKeyRange(KeyRange.ALL).build();
Cursor cursor = vtgate.execute(allRangeQuery);
Assert.assertEquals(rowsPerShard * 2, cursor.getRowsAffected());
// Check KeyRange query limited to a single shard returns 10 rows each
for (String shardName : testEnv.shardKidMap.keySet()) {
List<KeyspaceId> shardKids = testEnv.getKeyspaceIds(shardName);
KeyspaceId minKid = Collections.min(shardKids);
KeyspaceId maxKid = Collections.max(shardKids);
KeyRange shardKeyRange = new KeyRange(minKid, maxKid);
Query shardRangeQuery =
new QueryBuilder(selectSql, testEnv.keyspace, "master")
.addKeyRange(shardKeyRange)
.build();
cursor = vtgate.execute(shardRangeQuery);
Assert.assertEquals(rowsPerShard, cursor.getRowsAffected());
}
// Now make a cross-shard KeyRange and check all rows are returned
Iterator<String> shardNameIter = testEnv.shardKidMap.keySet().iterator();
KeyspaceId kidShard1 = testEnv.getKeyspaceIds(shardNameIter.next()).get(2);
KeyspaceId kidShard2 = testEnv.getKeyspaceIds(shardNameIter.next()).get(2);
KeyRange crossShardKeyrange;
if (kidShard1.compareTo(kidShard2) < 0) {
crossShardKeyrange = new KeyRange(kidShard1, kidShard2);
} else {
crossShardKeyrange = new KeyRange(kidShard2, kidShard1);
}
Query shardRangeQuery =
new QueryBuilder(selectSql, testEnv.keyspace, "master")
.addKeyRange(crossShardKeyrange)
.build();
cursor = vtgate.execute(shardRangeQuery);
Assert.assertEquals(rowsPerShard * 2, cursor.getRowsAffected());
vtgate.close();
}
/** Test queries are routed to the right shard based on based on keyspace ids */
@Test
public void testQueryRouting() throws Exception {
for (String shardName : testEnv.shardKidMap.keySet()) {
Util.insertRowsInShard(testEnv, shardName, 10);
}
VtGate vtgate = VtGate.connect("localhost:" + testEnv.port, 0);
String allRowsSql = "select * from vtgate_test";
for (String shardName : testEnv.shardKidMap.keySet()) {
Query shardRows =
new QueryBuilder(allRowsSql, testEnv.keyspace, "master")
.setKeyspaceIds(testEnv.getKeyspaceIds(shardName))
.build();
Cursor cursor = vtgate.execute(shardRows);
Assert.assertEquals(10, cursor.getRowsAffected());
}
vtgate.close();
}