@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 selects using ExecuteKeyspaceIds */
@Test
public void testExecuteKeyspaceIds() throws Exception {
VtGate vtgate = VtGate.connect("localhost:" + testEnv.port, 0);
// Ensure empty table
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(CursorImpl.class, cursor.getClass());
Assert.assertEquals(0, cursor.getRowsAffected());
Assert.assertEquals(0, cursor.getLastRowId());
Assert.assertFalse(cursor.hasNext());
vtgate.close();
// Insert 10 rows
Util.insertRows(testEnv, 1000, 10);
vtgate = VtGate.connect("localhost:" + testEnv.port, 0);
cursor = vtgate.execute(allRowsQuery);
Assert.assertEquals(10, cursor.getRowsAffected());
Assert.assertEquals(0, cursor.getLastRowId());
Assert.assertTrue(cursor.hasNext());
// Fetch all rows from the first shard
KeyspaceId firstKid = testEnv.getAllKeyspaceIds().get(0);
Query query =
new QueryBuilder(selectSql, testEnv.keyspace, "master").addKeyspaceId(firstKid).build();
cursor = vtgate.execute(query);
// Check field types and values
Row row = cursor.next();
Cell idCell = row.next();
Assert.assertEquals("id", idCell.getName());
Assert.assertEquals(UnsignedLong.class, idCell.getType());
UnsignedLong id = row.getULong(idCell.getName());
Cell nameCell = row.next();
Assert.assertEquals("name", nameCell.getName());
Assert.assertEquals(byte[].class, nameCell.getType());
Assert.assertTrue(
Arrays.equals(("name_" + id.toString()).getBytes(), row.getBytes(nameCell.getName())));
Cell ageCell = row.next();
Assert.assertEquals("age", ageCell.getName());
Assert.assertEquals(Integer.class, ageCell.getType());
Assert.assertEquals(Integer.valueOf(2 * id.intValue()), row.getInt(ageCell.getName()));
vtgate.close();
}