예제 #1
0
  @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()));
  }
예제 #2
0
 @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();
 }
예제 #3
0
  /** 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();
  }
예제 #4
0
 @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();
   }
 }
예제 #5
0
 /** 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();
 }
예제 #6
0
 /** 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();
   }
 }
예제 #7
0
  /** 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();
  }
예제 #8
0
  /** 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();
  }
예제 #9
0
 @Override
 public void close() throws IOException {
   if (vtgate != null) {
     try {
       vtgate.close();
       vtgate = null;
     } catch (ConnectionException e) {
       throw new RuntimeException(e);
     }
   }
 }
예제 #10
0
  @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);
  }
예제 #11
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();
  }
예제 #12
0
 /** 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);
   }
 }
예제 #13
0
 /**
  * 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;
 }
예제 #14
0
  /** 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();
  }