org.apache.kudu.client.KuduSession Java Examples

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

private LinkedList<OperationResponse> queueInsert(MockPutKudu putKudu, KuduSession session, boolean sync, ResultCode... results) throws Exception {
    LinkedList<OperationResponse> responses = new LinkedList<>();
    for (ResultCode result : results) {
        boolean ok = result == OK;
        Tuple<Insert, OperationResponse> tuple = insert(ok);
        putKudu.queue(tuple.getKey());

        if (result == EXCEPTION) {
            when(session.apply(tuple.getKey())).thenThrow(mock(KuduException.class));
            // Stop processing the rest of the records on the first exception
            break;
        } else {
            responses.add(tuple.getValue());
            if (sync) {
                when(session.apply(tuple.getKey())).thenReturn(ok ? null : tuple.getValue());

                // In AUTO_FLUSH_SYNC mode, PutKudu immediately knows when an operation has failed.
                // In that case, it does not process the rest of the records in the FlowFile.
                if (result == FAIL) break;
            }
        }
    }
    return responses;
}
 

private KuduSession openKuduSession(List<ConfigIssue> issues) {
  KuduSession session = kuduClient.newSession();
  try {
    session.setExternalConsistencyMode(ExternalConsistencyMode.valueOf(configBean.consistencyMode.name()));
  } catch (IllegalArgumentException ex) {
    issues.add(
        getContext().createConfigIssue(
            Groups.KUDU.name(),
            KuduConfigBean.CONF_PREFIX + CONSISTENCY_MODE,
            Errors.KUDU_02
        )
    );
  }
  session.setMutationBufferSpace(configBean.mutationBufferSpace);
  session.setFlushMode(SessionConfiguration.FlushMode.MANUAL_FLUSH);
  return session;
}
 

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

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

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

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

public KuduTableClient(KuduClient client, KuduTable kuduTable, KuduSession kuduSession,int batchSize) {
    this.client = client;
    this.kuduTable = kuduTable;
    this.metaTable = getMetaTable(kuduTable);
    this.kuduSession = kuduSession;
    this.batchSize = batchSize;
}
 

protected void flushKuduSession(final KuduSession kuduSession, boolean close, final List<RowError> rowErrors) throws KuduException {
    final List<OperationResponse> responses = close ? kuduSession.close() : kuduSession.flush();

    if (kuduSession.getFlushMode() == SessionConfiguration.FlushMode.AUTO_FLUSH_BACKGROUND) {
        rowErrors.addAll(Arrays.asList(kuduSession.getPendingErrors().getRowErrors()));
    } else {
        responses.stream()
                .filter(OperationResponse::hasRowError)
                .map(OperationResponse::getRowError)
                .forEach(rowErrors::add);
    }
}
 

@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
}
 

public ApexKuduConnection buildMockWiring(AbstractKuduInputOperator abstractKuduInputOperator,
    int numScanTokens) throws Exception
{
  ApexKuduConnection mockedConnectionHandle = PowerMockito.mock(ApexKuduConnection.class);
  ApexKuduConnection.ApexKuduConnectionBuilder mockedConnectionHandleBuilder = PowerMockito.mock(
      ApexKuduConnection.ApexKuduConnectionBuilder.class);
  KuduClient mockedClient = PowerMockito.mock(KuduClient.class);
  KuduSession mockedKuduSession = PowerMockito.mock(KuduSession.class);
  KuduTable mockedKuduTable = PowerMockito.mock(KuduTable.class);
  KuduScanToken.KuduScanTokenBuilder mockedScanTokenBuilder = PowerMockito.mock(
      KuduScanToken.KuduScanTokenBuilder.class);
  List<KuduScanToken> mockedScanTokens = new ArrayList<>();
  int scanTokensToBuild = numScanTokens;
  for (int i = 0; i < scanTokensToBuild; i++) {
    mockedScanTokens.add(PowerMockito.mock(KuduScanToken.class));
  }
  PowerMockito.mockStatic(KryoCloneUtils.class);
  when(KryoCloneUtils.cloneObject(abstractKuduInputOperator)).thenReturn(abstractKuduInputOperator);
  //wire the mocks
  when(abstractKuduInputOperator.getApexKuduConnectionInfo()).thenReturn(mockedConnectionHandleBuilder);
  when(mockedConnectionHandle.getKuduClient()).thenReturn(mockedClient);
  when(mockedClient.newSession()).thenReturn(mockedKuduSession);
  when(mockedConnectionHandle.getKuduTable()).thenReturn(mockedKuduTable);
  when(mockedConnectionHandle.getKuduSession()).thenReturn(mockedKuduSession);
  when(mockedConnectionHandle.getBuilderForThisConnection()).thenReturn(mockedConnectionHandleBuilder);
  when(mockedClient.openTable(tableName)).thenReturn(mockedKuduTable);
  when(mockedConnectionHandleBuilder.build()).thenReturn(mockedConnectionHandle);
  when(mockedKuduTable.getSchema()).thenReturn(schemaForUnitTests);
  when(mockedClient.newScanTokenBuilder(mockedKuduTable)).thenReturn(mockedScanTokenBuilder);
  when(mockedScanTokenBuilder.build()).thenReturn(mockedScanTokens);
  return mockedConnectionHandle;
}
 

@Override
protected KuduSession createKuduSession(final KuduClient client) {
    return session;
}
 

public MockPutKudu(KuduSession session) {
    this.session = session;
    this.insertQueue = new LinkedList<>();
}
 

public MockPutKudu() {
    this(mock(KuduSession.class));
}
 

public KuduSession newSession()
{
    return client.newSession();
}
 

protected KuduSession createKuduSession(final KuduClient client) {
    final KuduSession kuduSession = client.newSession();
    kuduSession.setMutationBufferSpace(batchSize);
    kuduSession.setFlushMode(flushMode);
    return kuduSession;
}
 

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

public KuduSession getSession() {
  return client.newSession();
}
 

public void setKuduSession(KuduSession kuduSession)
{
  this.kuduSession = kuduSession;
}
 

public KuduSession getKuduSession()
{
  return kuduSession;
}
 

KuduSession getSession() {
  return session;
}
 

private KuduSession obtainSession() {
    return client.newSession();
}
 

private KuduSession obtainSession() {
    KuduSession session = client.newSession();
    session.setFlushMode(writerConfig.getFlushMode());
    return session;
}
 

public KuduSession getKuduSession() {
    return kuduSession;
}
 

@Override
public KuduSession newSession() {
    return client.newSession();
}