@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 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
public void testSplitQueryInvalidTable() throws Exception {
VtGate vtgate = VtGate.connect("localhost:" + testEnv.port, 0);
try {
vtgate.splitQuery("test_keyspace", "select id from invalid_table", 1);
Assert.fail("failed to raise connection exception");
} catch (ConnectionException e) {
Assert.assertTrue(
e.getMessage().contains("query validation error: can't find table in schema"));
} finally {
vtgate.close();
}
}
/** Test ALL keyrange fetches rows from all shards */
@Test
public void testAllKeyRange() throws Exception {
// Insert 10 rows across the shards
Util.insertRows(testEnv, 1000, 10);
VtGate vtgate = VtGate.connect("localhost:" + testEnv.port, 0);
String selectSql = "select * from vtgate_test";
Query allRowsQuery =
new QueryBuilder(selectSql, testEnv.keyspace, "master").addKeyRange(KeyRange.ALL).build();
Cursor cursor = vtgate.execute(allRowsQuery);
// Verify all rows returned
Assert.assertEquals(10, cursor.getRowsAffected());
vtgate.close();
}
/** Test DMLs are not allowed outside a transaction */
@Test
public void testDMLOutsideTransaction() throws ConnectionException {
VtGate vtgate = VtGate.connect("localhost:" + testEnv.port, 0);
String deleteSql = "delete from vtgate_test";
try {
vtgate.execute(
new QueryBuilder(deleteSql, testEnv.keyspace, "master")
.addKeyspaceId(testEnv.getAllKeyspaceIds().get(0))
.build());
Assert.fail("did not raise DatabaseException");
} catch (DatabaseException e) {
Assert.assertTrue(e.getMessage().contains("not_in_tx"));
} finally {
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();
}
/** 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();
}
@Override
public void close() throws IOException {
if (vtgate != null) {
try {
vtgate.close();
vtgate = null;
} catch (ConnectionException e) {
throw new RuntimeException(e);
}
}
}
@Test
public void testSplitQueryMultipleSplitsPerShard() throws Exception {
int rowCount = 30;
Util.insertRows(testEnv, 1, 30);
List<String> expectedSqls =
Lists.newArrayList(
"select id, keyspace_id from vtgate_test where id < 10",
"select id, keyspace_id from vtgate_test where id < 11",
"select id, keyspace_id from vtgate_test where id >= 10 and id < 19",
"select id, keyspace_id from vtgate_test where id >= 11 and id < 19",
"select id, keyspace_id from vtgate_test where id >= 19",
"select id, keyspace_id from vtgate_test where id >= 19");
Util.waitForTablet("rdonly", rowCount, 3, testEnv);
VtGate vtgate = VtGate.connect("localhost:" + testEnv.port, 0);
int splitCount = 6;
Map<Query, Long> queries =
vtgate.splitQuery("test_keyspace", "select id,keyspace_id from vtgate_test", splitCount);
vtgate.close();
// Verify 6 splits, 3 per shard
Assert.assertEquals(splitCount, queries.size());
Set<String> shardsInSplits = new HashSet<>();
for (Query q : queries.keySet()) {
String sql = q.getSql();
Assert.assertTrue(expectedSqls.contains(sql));
expectedSqls.remove(sql);
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()));
Assert.assertTrue(expectedSqls.size() == 0);
}
@Test
public void testDateFieldTypes() throws Exception {
DateTime dt = DateTime.now().minusDays(2).withMillisOfSecond(0);
Util.insertRows(testEnv, 10, 1, dt);
VtGate vtgate = VtGate.connect("localhost:" + testEnv.port, 0);
Query allRowsQuery =
new QueryBuilder("select * from vtgate_test", testEnv.keyspace, "master")
.setKeyspaceIds(testEnv.getAllKeyspaceIds())
.build();
Row row = vtgate.execute(allRowsQuery).next();
Assert.assertTrue(dt.equals(row.getDateTime("timestamp_col")));
Assert.assertTrue(dt.equals(row.getDateTime("datetime_col")));
Assert.assertTrue(
dt.withHourOfDay(0)
.withMinuteOfHour(0)
.withSecondOfMinute(0)
.equals(row.getDateTime("date_col")));
Assert.assertTrue(
dt.withYear(1970).withMonthOfYear(1).withDayOfMonth(1).equals(row.getDateTime("time_col")));
vtgate.close();
}
/** Fetch connection parameters from Configuraiton and open VtGate connection. */
@Override
public void initialize(InputSplit split, TaskAttemptContext context)
throws IOException, InterruptedException {
this.split = (VitessInputSplit) split;
conf = new VitessConf(context.getConfiguration());
try {
Class<? extends RpcClientFactory> rpcFactoryClass =
(Class<? extends RpcClientFactory>) Class.forName(conf.getRpcFactoryClass());
vtgate = VtGate.connect(conf.getHosts(), conf.getTimeoutMs(), rpcFactoryClass.newInstance());
} catch (ConnectionException
| ClassNotFoundException
| InstantiationException
| IllegalAccessException e) {
throw new RuntimeException(e);
}
}
/**
* Fetches the next row. If this is the first invocation for the split, execute the streaming
* query. Subsequent calls just advance the iterator.
*/
@Override
public boolean nextKeyValue() throws IOException, InterruptedException {
if (cursor == null) {
try {
cursor = vtgate.execute(split.getQuery());
} catch (DatabaseException | ConnectionException e) {
throw new RuntimeException(e);
}
}
if (!cursor.hasNext()) {
return false;
}
Row row = cursor.next();
rowWritable = new RowWritable(row);
rowsProcessed++;
return true;
}
/** 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();
}