@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(); }