Examples with ResultSetFuture com.datastax.driver.core.ResultSetFuture used on opensource projects

Example 46 with ResultSetFuture

use of com.datastax.driver.core.ResultSetFuture in project nifi by apache.

the class QueryCassandra method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    FlowFile fileToProcess = null;
    if (context.hasIncomingConnection()) {
        fileToProcess = session.get();
        // we know that we should run only if we have a FlowFile.
        if (fileToProcess == null && context.hasNonLoopConnection()) {
            return;
        }
    }
    final ComponentLog logger = getLogger();
    final String selectQuery = context.getProperty(CQL_SELECT_QUERY).evaluateAttributeExpressions(fileToProcess).getValue();
    final long queryTimeout = context.getProperty(QUERY_TIMEOUT).evaluateAttributeExpressions(fileToProcess).asTimePeriod(TimeUnit.MILLISECONDS);
    final String outputFormat = context.getProperty(OUTPUT_FORMAT).getValue();
    final Charset charset = Charset.forName(context.getProperty(CHARSET).evaluateAttributeExpressions(fileToProcess).getValue());
    final StopWatch stopWatch = new StopWatch(true);
    if (fileToProcess == null) {
        fileToProcess = session.create();
    }
    try {
        // The documentation for the driver recommends the session remain open the entire time the processor is running
        // and states that it is thread-safe. This is why connectionSession is not in a try-with-resources.
        final Session connectionSession = cassandraSession.get();
        final ResultSetFuture queryFuture = connectionSession.executeAsync(selectQuery);
        final AtomicLong nrOfRows = new AtomicLong(0L);
        fileToProcess = session.write(fileToProcess, new OutputStreamCallback() {

            @Override
            public void process(final OutputStream out) throws IOException {
                try {
                    logger.debug("Executing CQL query {}", new Object[] { selectQuery });
                    final ResultSet resultSet;
                    if (queryTimeout > 0) {
                        resultSet = queryFuture.getUninterruptibly(queryTimeout, TimeUnit.MILLISECONDS);
                        if (AVRO_FORMAT.equals(outputFormat)) {
                            nrOfRows.set(convertToAvroStream(resultSet, out, queryTimeout, TimeUnit.MILLISECONDS));
                        } else if (JSON_FORMAT.equals(outputFormat)) {
                            nrOfRows.set(convertToJsonStream(resultSet, out, charset, queryTimeout, TimeUnit.MILLISECONDS));
                        }
                    } else {
                        resultSet = queryFuture.getUninterruptibly();
                        if (AVRO_FORMAT.equals(outputFormat)) {
                            nrOfRows.set(convertToAvroStream(resultSet, out, 0, null));
                        } else if (JSON_FORMAT.equals(outputFormat)) {
                            nrOfRows.set(convertToJsonStream(resultSet, out, charset, 0, null));
                        }
                    }
                } catch (final TimeoutException | InterruptedException | ExecutionException e) {
                    throw new ProcessException(e);
                }
            }
        });
        // set attribute how many rows were selected
        fileToProcess = session.putAttribute(fileToProcess, RESULT_ROW_COUNT, String.valueOf(nrOfRows.get()));
        // set mime.type based on output format
        fileToProcess = session.putAttribute(fileToProcess, CoreAttributes.MIME_TYPE.key(), JSON_FORMAT.equals(outputFormat) ? "application/json" : "application/avro-binary");
        logger.info("{} contains {} Avro records; transferring to 'success'", new Object[] { fileToProcess, nrOfRows.get() });
        session.getProvenanceReporter().modifyContent(fileToProcess, "Retrieved " + nrOfRows.get() + " rows", stopWatch.getElapsed(TimeUnit.MILLISECONDS));
        session.transfer(fileToProcess, REL_SUCCESS);
    } catch (final NoHostAvailableException nhae) {
        getLogger().error("No host in the Cassandra cluster can be contacted successfully to execute this query", nhae);
        // Log up to 10 error messages. Otherwise if a 1000-node cluster was specified but there was no connectivity,
        // a thousand error messages would be logged. However we would like information from Cassandra itself, so
        // cap the error limit at 10, format the messages, and don't include the stack trace (it is displayed by the
        // logger message above).
        getLogger().error(nhae.getCustomMessage(10, true, false));
        fileToProcess = session.penalize(fileToProcess);
        session.transfer(fileToProcess, REL_RETRY);
    } catch (final QueryExecutionException qee) {
        logger.error("Cannot execute the query with the requested consistency level successfully", qee);
        fileToProcess = session.penalize(fileToProcess);
        session.transfer(fileToProcess, REL_RETRY);
    } catch (final QueryValidationException qve) {
        if (context.hasIncomingConnection()) {
            logger.error("The CQL query {} is invalid due to syntax error, authorization issue, or another " + "validation problem; routing {} to failure", new Object[] { selectQuery, fileToProcess }, qve);
            fileToProcess = session.penalize(fileToProcess);
            session.transfer(fileToProcess, REL_FAILURE);
        } else {
            // This can happen if any exceptions occur while setting up the connection, statement, etc.
            logger.error("The CQL query {} is invalid due to syntax error, authorization issue, or another " + "validation problem", new Object[] { selectQuery }, qve);
            session.remove(fileToProcess);
            context.yield();
        }
    } catch (final ProcessException e) {
        if (context.hasIncomingConnection()) {
            logger.error("Unable to execute CQL select query {} for {} due to {}; routing to failure", new Object[] { selectQuery, fileToProcess, e });
            fileToProcess = session.penalize(fileToProcess);
            session.transfer(fileToProcess, REL_FAILURE);
        } else {
            logger.error("Unable to execute CQL select query {} due to {}", new Object[] { selectQuery, e });
            session.remove(fileToProcess);
            context.yield();
        }
    }
}

Example 47 with ResultSetFuture

use of com.datastax.driver.core.ResultSetFuture in project atlasdb by palantir.

the class CqlKeyValueService method getLatestTimestampsInternal.

private Map<Cell, Long> getLatestTimestampsInternal(final TableReference tableRef, Map<Cell, Long> timestampByCell) throws Exception {
    int fetchBatchCount = config.fetchBatchCount();
    Iterable<List<Cell>> partitions = Iterables.partition(timestampByCell.keySet(), fetchBatchCount);
    int numPartitions = (timestampByCell.size() / fetchBatchCount) + (timestampByCell.size() % fetchBatchCount > 0 ? 1 : 0);
    List<Future<Map<Cell, Long>>> futures = Lists.newArrayListWithCapacity(numPartitions);
    String loadOnlyTsQuery = "SELECT " + fieldNameProvider.row() + ", " + fieldNameProvider.column() + ", " + fieldNameProvider.timestamp() + " FROM " + getFullTableName(tableRef) + " WHERE " + fieldNameProvider.row() + " = ?" + " AND " + fieldNameProvider.column() + " = ?" + " LIMIT 1";
    if (timestampByCell.size() > fetchBatchCount) {
        log.warn("Re-batching in getLatestTimestamps a call to {} that attempted to multiget {} cells; " + "this may indicate overly-large batching on a higher level.\n{}", tableRef, timestampByCell.size(), CassandraKeyValueServices.getFilteredStackTrace("com.palantir"));
    }
    for (final List<Cell> partition : partitions) {
        futures.add(executor.submit(AnnotatedCallable.wrapWithThreadName(AnnotationType.PREPEND, "Atlas CQL getLatestTimestamps partition of " + partition.size(), () -> {
            PreparedStatement preparedStatement = getPreparedStatement(tableRef, loadOnlyTsQuery, session);
            preparedStatement.setConsistencyLevel(readConsistency);
            List<ResultSetFuture> resultSetFutures = Lists.newArrayListWithExpectedSize(partition.size());
            for (Cell c : partition) {
                BoundStatement boundStatement = preparedStatement.bind();
                boundStatement.setBytes(fieldNameProvider.row(), ByteBuffer.wrap(c.getRowName()));
                boundStatement.setBytes(fieldNameProvider.column(), ByteBuffer.wrap(c.getColumnName()));
                resultSetFutures.add(session.executeAsync(boundStatement));
            }
            Map<Cell, Long> res = Maps.newHashMapWithExpectedSize(partition.size());
            for (ResultSetFuture resultSetFuture : resultSetFutures) {
                ResultSet resultSet = resultSetFuture.getUninterruptibly();
                for (Row row : resultSet.all()) {
                    res.put(Cell.create(getRowName(row), getColName(row)), getTs(row));
                }
                cqlKeyValueServices.logTracedQuery(loadOnlyTsQuery, resultSet, session, cqlStatementCache.normalQuery);
            }
            return res;
        })));
    }
    Map<Cell, Long> res = Maps.newHashMapWithExpectedSize(timestampByCell.size());
    for (Future<Map<Cell, Long>> f : futures) {
        try {
            res.putAll(f.get());
        } catch (InterruptedException e) {
            throw Throwables.throwUncheckedException(e);
        } catch (ExecutionException e) {
            Throwables.throwIfInstance(e, Error.class);
            throw Throwables.throwUncheckedException(e.getCause());
        }
    }
    return res;
}

Example 48 with ResultSetFuture

use of com.datastax.driver.core.ResultSetFuture in project atlasdb by palantir.

the class CqlKeyValueService method multiPut.

@Override
public void multiPut(Map<TableReference, ? extends Map<Cell, byte[]>> valuesByTable, long timestamp) throws KeyAlreadyExistsException {
    Map<ResultSetFuture, TableReference> resultSetFutures = Maps.newHashMap();
    for (Entry<TableReference, ? extends Map<Cell, byte[]>> e : valuesByTable.entrySet()) {
        final TableReference table = e.getKey();
        // We sort here because some key value stores are more efficient if you store adjacent keys together.
        NavigableMap<Cell, byte[]> sortedMap = ImmutableSortedMap.copyOf(e.getValue());
        Iterable<List<Entry<Cell, byte[]>>> partitions = IterablePartitioner.partitionByCountAndBytes(sortedMap.entrySet(), getMultiPutBatchCount(), getMultiPutBatchSizeBytes(), table, CqlKeyValueServices.MULTIPUT_ENTRY_SIZING_FUNCTION);
        for (final List<Entry<Cell, byte[]>> p : partitions) {
            List<Entry<Cell, Value>> partition = Lists.transform(p, input -> Maps.immutableEntry(input.getKey(), Value.create(input.getValue(), timestamp)));
            resultSetFutures.put(getPutPartitionResultSetFuture(table, partition, TransactionType.NONE), table);
        }
    }
    for (Entry<ResultSetFuture, TableReference> result : resultSetFutures.entrySet()) {
        ResultSet resultSet;
        try {
            resultSet = result.getKey().getUninterruptibly();
            resultSet.all();
        } catch (Throwable t) {
            throw Throwables.throwUncheckedException(t);
        }
        cqlKeyValueServices.logTracedQuery(getPutQuery(result.getValue(), CassandraConstants.NO_TTL), resultSet, session, cqlStatementCache.normalQuery);
    }
}

Example 49 with ResultSetFuture

use of com.datastax.driver.core.ResultSetFuture in project teiid by teiid.

the class TestNativeCassandra method testNativeQuery.

@Test
public void testNativeQuery() throws Exception {
    CassandraExecutionFactory cef = new CassandraExecutionFactory();
    cef.setSupportsDirectQueryProcedure(true);
    String input = "call proc('a', 1)";
    TransformationMetadata metadata = RealMetadataFactory.fromDDL("create foreign procedure proc (in x string, in y integer) options (\"teiid_rel:native-query\" 'delete from $1 where $2')", "x", "y");
    TranslationUtility util = new TranslationUtility(metadata);
    Command command = util.parseCommand(input);
    ExecutionContext ec = Mockito.mock(ExecutionContext.class);
    RuntimeMetadata rm = Mockito.mock(RuntimeMetadata.class);
    CassandraConnection connection = Mockito.mock(CassandraConnection.class);
    ResultSetFuture rsf = Mockito.mock(ResultSetFuture.class);
    Mockito.stub(connection.executeQuery("delete from 'a' where 1")).toReturn(rsf);
    Execution execution = cef.createExecution(command, ec, rm, connection);
    execution.execute();
    Mockito.verify(connection).executeQuery("delete from 'a' where 1");
}

Example 50 with ResultSetFuture

use of com.datastax.driver.core.ResultSetFuture in project teiid by teiid.

the class TestUpdates method testBatchedUpdate.

@Test
public void testBatchedUpdate() throws TranslatorException {
    CassandraExecutionFactory cef = new CassandraExecutionFactory();
    String input = "insert into pm1.g1 (e1) values ('a')";
    TranslationUtility util = FakeTranslationFactory.getInstance().getExampleTranslationUtility();
    Command command = util.parseCommand(input);
    Command command1 = util.parseCommand("update pm1.g1 set e1 = 'b'");
    command = new BatchedUpdates(Arrays.asList(command, command1));
    ExecutionContext ec = Mockito.mock(ExecutionContext.class);
    RuntimeMetadata rm = Mockito.mock(RuntimeMetadata.class);
    CassandraConnection connection = Mockito.mock(CassandraConnection.class);
    ResultSetFuture rsf = Mockito.mock(ResultSetFuture.class);
    Mockito.stub(rsf.isDone()).toReturn(true);
    Mockito.stub(connection.executeBatch(Arrays.asList("INSERT INTO g1 (e1) VALUES ('a')", "UPDATE g1 SET e1 = 'b'"))).toReturn(rsf);
    UpdateExecution execution = (UpdateExecution) cef.createExecution(command, ec, rm, connection);
    execution.execute();
    assertArrayEquals(new int[] { 2 }, execution.getUpdateCounts());
    Mockito.verify(connection).executeBatch(Arrays.asList("INSERT INTO g1 (e1) VALUES ('a')", "UPDATE g1 SET e1 = 'b'"));
}