@Test
public void testSplitQuery() throws Exception {
// Insert 20 rows per shard
for (String shardName : testEnv.shardKidMap.keySet()) {
Util.insertRowsInShard(testEnv, shardName, 20);
}
Util.waitForTablet("rdonly", 40, 3, testEnv);
VtGate vtgate = VtGate.connect("localhost:" + testEnv.port, 0);
Map<Query, Long> queries =
vtgate.splitQuery("test_keyspace", "select id,keyspace_id from vtgate_test", 1);
vtgate.close();
// Verify 2 splits, one per shard
Assert.assertEquals(2, queries.size());
Set<String> shardsInSplits = new HashSet<>();
for (Query q : queries.keySet()) {
Assert.assertEquals("select id,keyspace_id from vtgate_test", q.getSql());
Assert.assertEquals("test_keyspace", q.getKeyspace());
Assert.assertEquals("rdonly", q.getTabletType());
Assert.assertEquals(0, q.getBindVars().size());
Assert.assertEquals(null, q.getKeyspaceIds());
String start = Hex.encodeHexString(q.getKeyRanges().get(0).get("Start"));
String end = Hex.encodeHexString(q.getKeyRanges().get(0).get("End"));
shardsInSplits.add(start + "-" + end);
}
// Verify the keyrange queries in splits cover the entire keyspace
Assert.assertTrue(shardsInSplits.containsAll(testEnv.shardKidMap.keySet()));
}
@Test
public void testBatchExecuteKeyspaceIds() throws Exception {
int rowsPerShard = 5;
for (String shardName : testEnv.shardKidMap.keySet()) {
Util.insertRowsInShard(testEnv, shardName, rowsPerShard);
}
VtGate vtgate = VtGate.connect("localhost:" + testEnv.port, 0);
BatchQuery query =
new BatchQueryBuilder(testEnv.keyspace, "master")
.addSqlAndBindVars("select * from vtgate_test where id = 3", null)
.addSqlAndBindVars(
"select * from vtgate_test where id = :id",
Lists.newArrayList(BindVariable.forULong("id", UnsignedLong.valueOf("4"))))
.withKeyspaceIds(testEnv.getAllKeyspaceIds())
.build();
List<Long> expected = Lists.newArrayList(3L, 3L, 4L, 4L);
List<Cursor> cursors = vtgate.execute(query);
List<Long> actual = new ArrayList<>();
for (Cursor cursor : cursors) {
for (Row row : cursor) {
actual.add(row.getULong("id").longValue());
}
}
Assert.assertTrue(expected.equals(actual));
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();
}