Java Code Examples for org.apache.kudu.client.KuduSession#apply()

@Override
public void deleteRows(Block rowIds)
{
    Schema schema = table.getSchema();
    KuduSession session = clientSession.newSession();
    session.setFlushMode(FlushMode.AUTO_FLUSH_BACKGROUND);
    try {
        try {
            for (int i = 0; i < rowIds.getPositionCount(); i++) {
                int len = rowIds.getSliceLength(i);
                Slice slice = rowIds.getSlice(i, 0, len);
                PartialRow row = KeyEncoderAccessor.decodePrimaryKey(schema, slice.getBytes());
                Delete delete = table.newDelete();
                RowHelper.copyPrimaryKey(schema, row, delete.getRow());
                session.apply(delete);
            }
        }
        finally {
            session.close();
        }
    }
    catch (KuduException e) {
        throw new RuntimeException(e);
    }
}
 

@Override
public void createSchema(KuduClient client, String schemaName)
{
    if (DEFAULT_SCHEMA.equals(schemaName)) {
        throw new SchemaAlreadyExistsException(schemaName);
    }
    else {
        try {
            KuduTable schemasTable = getSchemasTable(client);
            KuduSession session = client.newSession();
            try {
                Upsert upsert = schemasTable.newUpsert();
                upsert.getRow().addString(0, schemaName);
                session.apply(upsert);
            }
            finally {
                session.close();
            }
        }
        catch (KuduException e) {
            throw new PrestoException(GENERIC_INTERNAL_ERROR, e);
        }
    }
}
 

private void createAndFillSchemasTable(KuduClient client)
        throws KuduException
{
    List<String> existingSchemaNames = listSchemaNamesFromTablets(client);
    ColumnSchema schemaColumnSchema = new ColumnSchema.ColumnSchemaBuilder("schema", Type.STRING)
            .key(true).build();
    Schema schema = new Schema(ImmutableList.of(schemaColumnSchema));
    CreateTableOptions options = new CreateTableOptions();
    options.addHashPartitions(ImmutableList.of(schemaColumnSchema.getName()), 2);
    KuduTable schemasTable = client.createTable(rawSchemasTableName, schema, options);
    KuduSession session = client.newSession();
    try {
        session.setFlushMode(SessionConfiguration.FlushMode.AUTO_FLUSH_BACKGROUND);
        for (String schemaName : existingSchemaNames) {
            Insert insert = schemasTable.newInsert();
            insert.getRow().addString(0, schemaName);
            session.apply(insert);
        }
    }
    finally {
        session.close();
    }
}
 

private void createAndFillSchemasTable() throws KuduException {
    List<String> existingSchemaNames = listSchemaNamesFromTablets();
    ColumnSchema tenantColumnSchema = new ColumnSchema.ColumnSchemaBuilder("tenant", Type.STRING)
            .key(true).build();
    ColumnSchema schemaColumnSchema = new ColumnSchema.ColumnSchemaBuilder("schema", Type.STRING)
            .key(true).build();
    Schema schema = new Schema(ImmutableList.of(tenantColumnSchema, schemaColumnSchema));
    CreateTableOptions options = new CreateTableOptions();
    options.setNumReplicas(1); // TODO config
    options.addHashPartitions(ImmutableList.of(tenantColumnSchema.getName()), 2);
    KuduTable schemasTable = client.createTable(rawSchemasTableName, schema, options);
    KuduSession session = client.newSession();
    session.setFlushMode(SessionConfiguration.FlushMode.AUTO_FLUSH_BACKGROUND);
    try {
        for (String schemaName : existingSchemaNames) {
            Insert insert = schemasTable.newInsert();
            fillSchemaRow(insert.getRow(), schemaName);
            session.apply(insert);
        }
    } finally {
        session.close();
    }
}
 

@Override
public void createSchema(String schemaName) {
    if (DEFAULT_SCHEMA.equals(schemaName)) {
        throw new SchemaAlreadyExistsException(schemaName);
    }
    else {
        try {
            KuduTable schemasTable = getSchemasTable();
            KuduSession session = client.newSession();
            try {
                Upsert upsert = schemasTable.newUpsert();
                fillSchemaRow(upsert.getRow(), schemaName);
                session.apply(upsert);
            }
            finally {
                session.close();
            }
        }
        catch (KuduException e) {
            throw new PrestoException(GENERIC_INTERNAL_ERROR, e);
        }
    }
}
 

@Override
public void dropSchema(String schemaName) {
    if (DEFAULT_SCHEMA.equals(schemaName)) {
        throw new PrestoException(GENERIC_USER_ERROR, "Deleting default schema not allowed.");
    }
    else {
        try {
            for (SchemaTableName table : listTables(schemaName)) {
                dropTable(table);
            }
            KuduTable schemasTable = getSchemasTable();
            KuduSession session = client.newSession();
            try {
                Delete delete = schemasTable.newDelete();
                fillSchemaRow(delete.getRow(), schemaName);
                session.apply(delete);
            }
            finally {
                session.close();
            }
        }
        catch (KuduException e) {
            throw new PrestoException(GENERIC_INTERNAL_ERROR, e);
        }
    }
}
 

@Override
public void deleteRows(Block rowIds) {
    Schema schema = table.getSchema();
    KuduSession session = clientSession.newSession();
    session.setFlushMode(FlushMode.AUTO_FLUSH_BACKGROUND);
    try {
        try {
            for (int i = 0; i < rowIds.getPositionCount(); i++) {
                int len = rowIds.getSliceLength(i);
                Slice slice = rowIds.getSlice(i, 0, len);
                PartialRow row = KeyEncoderAccessor.decodePrimaryKey(schema, slice.getBytes());
                Delete delete = table.newDelete();
                RowHelper.copyPrimaryKey(schema, row, delete.getRow());
                session.apply(delete);
            }
        } finally {
            session.close();
        }
    } catch (KuduException e) {
        throw new RuntimeException(e);
    }
}
 

@Override
public void applyRandomMutations(List<Row> planned) throws Exception {
  KuduConnection connection = getConnection();
  KuduSession session = connection.getSession();
  KuduTable table = connection.getTable(config.getString(TABLE_CONFIG_NAME));

  List<Operation> operations = extractOperations(planned, table);

  for (Operation operation : operations) {
    session.apply(operation);
  }

  // Wait until all operations have completed before checking for errors.
  while (session.hasPendingOperations()) {
    Thread.sleep(1);
  }

  // Fail fast on any error applying mutations
  if (session.countPendingErrors() > 0) {
    RowError firstError = session.getPendingErrors().getRowErrors()[0];
    String errorMessage = String.format("Kudu output error '%s' during operation '%s' at tablet server '%s'",
        firstError.getErrorStatus(), firstError.getOperation(), firstError.getTsUUID());

    throw new RuntimeException(errorMessage);
  }
}
 

@Override
public void dropSchema(KuduClient client, String schemaName)
{
    if (DEFAULT_SCHEMA.equals(schemaName)) {
        throw new PrestoException(GENERIC_USER_ERROR, "Deleting default schema not allowed.");
    }
    else {
        try {
            String prefix = getPrefixForTablesOfSchema(schemaName);
            for (String name : client.getTablesList(prefix).getTablesList()) {
                client.deleteTable(name);
            }

            KuduTable schemasTable = getSchemasTable(client);
            KuduSession session = client.newSession();
            try {
                Delete delete = schemasTable.newDelete();
                delete.getRow().addString(0, schemaName);
                session.apply(delete);
            }
            finally {
                session.close();
            }
        }
        catch (KuduException e) {
            throw new PrestoException(GENERIC_INTERNAL_ERROR, e);
        }
    }
}
 

public void truncateTable() throws Exception
{
  AbstractKuduPartitionScanner<UnitTestTablePojo,InputOperatorControlTuple> scannerForDeletingRows =
      unitTestStepwiseScanInputOperator.getScanner();
  List<KuduScanToken> scansForAllTablets = unitTestStepwiseScanInputOperator
      .getPartitioner().getKuduScanTokensForSelectAllColumns();
  ApexKuduConnection aCurrentConnection = scannerForDeletingRows.getConnectionPoolForThreads().get(0);
  KuduSession aSessionForDeletes = aCurrentConnection.getKuduClient().newSession();
  KuduTable currentTable = aCurrentConnection.getKuduTable();
  for ( KuduScanToken aTabletScanToken : scansForAllTablets) {
    KuduScanner aScanner = aTabletScanToken.intoScanner(aCurrentConnection.getKuduClient());
    while ( aScanner.hasMoreRows()) {
      RowResultIterator itrForRows = aScanner.nextRows();
      while ( itrForRows.hasNext()) {
        RowResult aRow = itrForRows.next();
        int intRowKey = aRow.getInt("introwkey");
        String stringRowKey = aRow.getString("stringrowkey");
        long timestampRowKey = aRow.getLong("timestamprowkey");
        Delete aDeleteOp = currentTable.newDelete();
        aDeleteOp.getRow().addInt("introwkey",intRowKey);
        aDeleteOp.getRow().addString("stringrowkey", stringRowKey);
        aDeleteOp.getRow().addLong("timestamprowkey",timestampRowKey);
        aSessionForDeletes.apply(aDeleteOp);
      }
    }
  }
  aSessionForDeletes.close();
  Thread.sleep(2000); // Sleep to allow for scans to complete
}
 

public void addTestDataRows(int numRowsInEachPartition) throws Exception
{
  int intRowKeyStepsize = Integer.MAX_VALUE / SPLIT_COUNT_FOR_INT_ROW_KEY;
  int splitBoundaryForIntRowKey = intRowKeyStepsize;
  int[] inputrowkeyPartitionEntries = new int[SPLIT_COUNT_FOR_INT_ROW_KEY + 1];
  // setting the int keys that will fall in the range of all partitions
  for ( int i = 0; i < SPLIT_COUNT_FOR_INT_ROW_KEY; i++) {
    inputrowkeyPartitionEntries[i] = splitBoundaryForIntRowKey + 3; // 3 to fall into the partition next to boundary
    splitBoundaryForIntRowKey += intRowKeyStepsize;
  }
  inputrowkeyPartitionEntries[SPLIT_COUNT_FOR_INT_ROW_KEY] = splitBoundaryForIntRowKey + 3;
  AbstractKuduPartitionScanner<UnitTestTablePojo,InputOperatorControlTuple> scannerForAddingRows =
      unitTestStepwiseScanInputOperator.getScanner();
  ApexKuduConnection aCurrentConnection = scannerForAddingRows.getConnectionPoolForThreads().get(0);
  KuduSession aSessionForInserts = aCurrentConnection.getKuduClient().newSession();
  KuduTable currentTable = aCurrentConnection.getKuduTable();
  long seedValueForTimestampRowKey = 0L; // constant to allow for data landing on first partition for unit tests
  for ( int i = 0; i <= SPLIT_COUNT_FOR_INT_ROW_KEY; i++) { // range key iterator
    int intRowKeyBaseValue = inputrowkeyPartitionEntries[i] + i;
    for ( int k = 0; k < 2; k++) { // hash key iterator . The table defines two hash partitions
      long timestampRowKeyValue = seedValueForTimestampRowKey + k; // to avoid spilling to another tablet
      String stringRowKeyValue = "" + timestampRowKeyValue + k; // to avoid spilling to another tablet randomly
      for ( int y = 0; y < numRowsInEachPartition; y++) {
        Upsert aNewRow = currentTable.newUpsert();
        PartialRow rowValue  = aNewRow.getRow();
        // Start assigning row keys below the current split boundary.
        rowValue.addInt("introwkey",intRowKeyBaseValue - y - 1);
        rowValue.addString("stringrowkey",stringRowKeyValue);
        rowValue.addLong("timestamprowkey",timestampRowKeyValue);
        rowValue.addLong("longdata",(seedValueForTimestampRowKey + y));
        rowValue.addString("stringdata", ("" + seedValueForTimestampRowKey + y));
        OperationResponse response = aSessionForInserts.apply(aNewRow);
      }
    }
  }
  List<OperationResponse> insertResponse = aSessionForInserts.flush();
  aSessionForInserts.close();
  Thread.sleep(2000); // Sleep to allow for scans to complete
}
 

@Before
public void init() throws Exception {
    testRunner = TestRunners.newTestRunner(SampleProcessor.class);
    testRunner.setValidateExpressionUsage(false);
    final String tableName = "table1";

    KuduClient client =  harness.getClient();
    List<ColumnSchema> columns = new ArrayList<>();
    columns.add(new ColumnSchema.ColumnSchemaBuilder("string", Type.STRING).key(true).build());
    columns.add(new ColumnSchema.ColumnSchemaBuilder("binary", Type.BINARY).build());
    columns.add(new ColumnSchema.ColumnSchemaBuilder("bool", Type.BOOL).build());
    columns.add(new ColumnSchema
            .ColumnSchemaBuilder("decimal", Type.DECIMAL)
            .typeAttributes(DecimalUtil.typeAttributes(DecimalUtil.MAX_DECIMAL64_PRECISION, 1))
            .build()
    );
    columns.add(new ColumnSchema.ColumnSchemaBuilder("double", Type.DOUBLE).build());
    columns.add(new ColumnSchema.ColumnSchemaBuilder("float", Type.FLOAT).build());
    columns.add(new ColumnSchema.ColumnSchemaBuilder("int8", Type.INT8).build());
    columns.add(new ColumnSchema.ColumnSchemaBuilder("int16", Type.INT16).build());
    columns.add(new ColumnSchema.ColumnSchemaBuilder("int32", Type.INT32).build());
    columns.add(new ColumnSchema.ColumnSchemaBuilder("int64", Type.INT64).build());
    columns.add(new ColumnSchema.ColumnSchemaBuilder("unixtime_micros", Type.UNIXTIME_MICROS).build());
    columns.add(new ColumnSchema.ColumnSchemaBuilder("varchar_3", Type.VARCHAR).typeAttributes(
            new ColumnTypeAttributes.ColumnTypeAttributesBuilder().length(3).build()
    ).build());
    Schema schema = new Schema(columns);

    CreateTableOptions opts = new CreateTableOptions().setRangePartitionColumns(Collections.singletonList("string"));
    client.createTable(tableName, schema, opts);

    KuduTable table = client.openTable(tableName);
    KuduSession session = client.newSession();

    Insert insert = table.newInsert();
    PartialRow row = insert.getRow();
    row.addString("string", "string1");
    row.addBinary("binary", "binary1".getBytes());
    row.addBoolean("bool",true);
    row.addDecimal("decimal", BigDecimal.valueOf(0.1));
    row.addDouble("double",0.2);
    row.addFloat("float",0.3f);
    row.addByte("int8", (byte) 1);
    row.addShort("int16", (short) 2);
    row.addInt("int32",3);
    row.addLong("int64",4L);
    row.addTimestamp("unixtime_micros", new Timestamp(nowMillis));
    row.addVarchar("varchar_3", "SFO");
    session.apply(insert);

    insert = table.newInsert();
    row = insert.getRow();
    row.addString("string", "string2");
    row.addBinary("binary", "binary2".getBytes());
    row.addBoolean("bool",false);
    row.addDecimal("decimal", BigDecimal.valueOf(0.1));
    row.addDouble("double",1.2);
    row.addFloat("float",1.3f);
    row.addByte("int8", (byte) 11);
    row.addShort("int16", (short) 12);
    row.addInt("int32",13);
    row.addLong("int64",14L);
    row.addTimestamp("unixtime_micros", new Timestamp(nowMillis+(1000L * 60 * 60 * 24 * 365))); //+ 1 year
    row.addVarchar("varchar_3", "SJC");
    session.apply(insert);

    session.close();

    kuduLookupService = new KuduLookupService();
    testRunner.addControllerService("kuduLookupService", kuduLookupService);
    testRunner.setProperty(kuduLookupService, KuduLookupService.KUDU_MASTERS, "testLocalHost:7051");
    testRunner.setProperty(kuduLookupService, KuduLookupService.KUDU_REPLICA_SELECTION, KuduLookupService.LEADER_ONLY);
    testRunner.setProperty(kuduLookupService, KuduLookupService.TABLE_NAME, tableName);
    kuduLookupService.kuduClient = client;
}