com.datastax.driver.core.Statement Java Examples

@Test
public void testClusterHintsPollerWhenNodeDown() throws UnknownHostException {
    ClusterHintsPoller clusterHintsPoller = new ClusterHintsPoller();
    Session mockSession = mock(Session.class);
    Cluster mockCluster = mock(Cluster.class);
    Metadata mockMetadata = mock(Metadata.class);
    when(mockCluster.getMetadata()).thenReturn(mockMetadata);
    when(mockCluster.getClusterName()).thenReturn("test-cluster");
    Host node1 = mock(Host.class);
    when(node1.getAddress()).thenReturn(InetAddress.getByName("127.0.0.1"));
    Host node2 = mock(Host.class);
    when(node2.getAddress()).thenReturn(InetAddress.getByName("127.0.0.2"));
    Host node3 = mock(Host.class);
    when(node3.getAddress()).thenReturn(InetAddress.getByName("127.0.0.3"));

    when(mockSession.getCluster()).thenReturn(mockCluster);
    // The first node queried is down
    when(mockSession.execute(any(Statement.class))).thenThrow(new NoHostAvailableException(ImmutableMap.<InetSocketAddress, Throwable>of()));

    when(mockMetadata.getAllHosts()).thenReturn(ImmutableSet.of(node1, node2, node3));
    HintsPollerResult actualResult = clusterHintsPoller.getOldestHintsInfo(mockSession);

    // Make sure HintsPollerResult fails
    assertFalse(actualResult.areAllHostsPolling(), "Result should show hosts failing");
    assertEquals(actualResult.getHostFailure(), ImmutableSet.of(InetAddress.getByName("127.0.0.1")), "Node 1 should return with host failure");
}
 

/**
 * This method returns a sorted map containing the columns that make up the (compound) primary key
 * in cassandra.
 *
 * <p>Each key in return value contains a name of a column in the cassandra key. The value for
 * each corresponding key in the map denotes the order in the Cassandra key.
 *
 * <p>Example, given the table:
 *
 * <p>create table mytable ( key_part_one int, key_part_two int, data boolean, PRIMARY
 * KEY(key_part_one, key_part_two) );
 *
 * <p>The method would return the following map:
 *
 * <p>"key_part_one", 0 "key_part_two", 1
 *
 * @param session the session to the Cassandra cluster.
 * @param keyspace they keyspace to query
 * @param table the table to query
 * @return see method description.
 */
static Map<String, Integer> primaryKeyOrder(Session session, String keyspace, String table) {
  HashMap<String, Integer> returnValue = new HashMap<>();
  Statement select = primarykeyCQL(keyspace, table);
  ResultSet rs = session.execute(select.toString());

  TreeMap<String, String> sortedMap = new TreeMap<>();
  for (Row row : rs) {
    String columnName = row.get(COLUMN_NAME_COLUMN, String.class);
    String kind = row.get(KIND_COLUMN, String.class);
    String position = row.get(POSITION_COLUMN, Integer.class).toString();
    if (kind.equals("clustering")) {
      sortedMap.put("clustering_" + position, columnName);
    } else {
      sortedMap.put(position, columnName);
    }
  }

  List<String> sortedKeyset = new ArrayList<>(sortedMap.keySet());
  for (int i = 0; i < sortedKeyset.size(); i++) {
    returnValue.put(sortedMap.get(sortedKeyset.get(i)), i);
  }

  return returnValue;
}
 

private Statement createDeleteAllMessagesStatement( Shard shard ) {

        Clause queueNameClause = QueryBuilder.eq(      COLUMN_QUEUE_NAME, shard.getQueueName() );
        Clause regionClause = QueryBuilder.eq(         COLUMN_REGION, shard.getRegion() );
        Clause shardIdClause = QueryBuilder.eq(        COLUMN_SHARD_ID, shard.getShardId() );

        DatabaseQueueMessage.Type dbqmType = Shard.Type.DEFAULT.equals( shard.getType() )
            ? DatabaseQueueMessage.Type.DEFAULT : DatabaseQueueMessage.Type.INFLIGHT;

        Statement deleteAll = QueryBuilder.delete().from( getTableName( dbqmType ))
            .where(queueNameClause)
            .and(regionClause)
            .and(shardIdClause);

        return deleteAll;
    }
 

public List<SizeEstimate> getSizeEstimates(String keyspaceName, String tableName)
{
    checkSizeEstimatesTableExist();
    Statement statement = select("partitions_count")
            .from(SYSTEM, SIZE_ESTIMATES)
            .where(eq("keyspace_name", keyspaceName))
            .and(eq("table_name", tableName));

    ResultSet result = executeWithSession(session -> session.execute(statement));
    ImmutableList.Builder<SizeEstimate> estimates = ImmutableList.builder();
    for (Row row : result.all()) {
        SizeEstimate estimate = new SizeEstimate(row.getLong("partitions_count"));
        estimates.add(estimate);
    }

    return estimates.build();
}
 

/**
 * Tunes CQL statement execution options (consistency level, fetch option and etc.).
 *
 * @param statement Statement.
 * @return Modified statement.
 */
private Statement tuneStatementExecutionOptions(Statement statement) {
    String qry = "";

    if (statement instanceof BoundStatement)
        qry = ((BoundStatement)statement).preparedStatement().getQueryString().trim().toLowerCase();
    else if (statement instanceof PreparedStatement)
        qry = ((PreparedStatement)statement).getQueryString().trim().toLowerCase();

    boolean readStatement = qry.startsWith("select");
    boolean writeStatement = statement instanceof Batch || statement instanceof BatchStatement ||
        qry.startsWith("insert") || qry.startsWith("delete") || qry.startsWith("update");

    if (readStatement && readConsistency != null)
        statement.setConsistencyLevel(readConsistency);

    if (writeStatement && writeConsistency != null)
        statement.setConsistencyLevel(writeConsistency);

    if (fetchSize != null)
        statement.setFetchSize(fetchSize);

    return statement;
}
 

/**
 * Executes a query sychronously, dynamically adjusting the fetch size down if necessary.
 */
public static ResultSet executeAdaptiveQuery(Session session, Statement statement, int fetchSize) {
    int remainingAdaptations = MAX_ADAPTATIONS;
    while (true) {
        try {
            statement.setFetchSize(fetchSize);
            ResultSet resultSet = session.execute(statement);
            return new AdaptiveResultSet(session, resultSet, remainingAdaptations);
        } catch (Throwable t) {
            if (isAdaptiveException(t) && --remainingAdaptations != 0 && fetchSize > MIN_FETCH_SIZE) {
                // Try again with half the fetch size
                fetchSize = Math.max(fetchSize / 2, MIN_FETCH_SIZE);
                _log.debug("Repeating previous query with fetch size {} due to {}", fetchSize, t.getMessage());
            } else {
                throw Throwables.propagate(t);
            }
        }
    }
}
 

public ResultSet update(Statement statement) throws Exception {
    if (statement instanceof BoundStatement) {
        BoundStatement boundStatement = (BoundStatement) statement;
        PreparedStatement preparedStatement = boundStatement.preparedStatement();
        String queryString = preparedStatement.getQueryString();
        if (!queryString.contains(" if ")) {
            throw new IllegalStateException("Unexpected update query: " + queryString);
        }
        if (!queryString.startsWith("update ") && !queryString.startsWith("insert ")) {
            throw new IllegalStateException("Unexpected update query: " + queryString);
        }
    }
    // do not use session.execute() because that calls getUninterruptibly() which can cause
    // central shutdown to timeout while waiting for executor service to shutdown
    return updateAsync(statement).get();
}
 

@Override
public void removeTransferLog(
        String queueName, String source, String dest, UUID messageId ) throws QakkaException {

    Statement query = QueryBuilder.select().all().from(TABLE_TRANSFER_LOG)
        .where(   QueryBuilder.eq( COLUMN_QUEUE_NAME, queueName ))
            .and( QueryBuilder.eq( COLUMN_DEST_REGION, dest ))
            .and( QueryBuilder.eq( COLUMN_MESSAGE_ID, messageId ));
    ResultSet rs = cassandraClient.getApplicationSession().execute( query );

    if ( rs.getAvailableWithoutFetching() == 0 ) {
        StringBuilder sb = new StringBuilder();
        sb.append( "Transfer log entry not found for queueName=" ).append( queueName );
        sb.append( " dest=" ).append( dest );
        sb.append( " messageId=" ).append( messageId );
        throw new QakkaException( sb.toString() );
    }

    Statement deleteQuery = QueryBuilder.delete().from(TABLE_TRANSFER_LOG)
            .where(   QueryBuilder.eq( COLUMN_QUEUE_NAME, queueName ))
                .and( QueryBuilder.eq( COLUMN_DEST_REGION, dest ))
            .and( QueryBuilder.eq( COLUMN_MESSAGE_ID, messageId ));
    cassandraClient.getApplicationSession().execute( deleteQuery );
}
 

@Override
protected List<Statement> prepareIndexUpdates(Device entity) {
   Device currentDevice = findById(entity.getId());
   if(currentDevice == null) {
      return prepareIndexInserts(entity.getId(), entity);
   }
   List<Statement> statements = new ArrayList<>();
   if(!StringUtils.equals(entity.getProtocolAddress(), currentDevice.getProtocolAddress())) {
      addProtocolAddressIndexDelete(statements, currentDevice);
      addProtocolAddressIndexInsert(statements, entity.getId(), entity);
   }
   if(!Objects.equal(entity.getPlace(), currentDevice.getPlace())) {
      addPlaceIndexDelete(statements, currentDevice);
      addPlaceIndexInsert(statements, entity.getId(), entity);
   }



   return statements;
}
 

@Test
public void handle_CorrectHandlerAndConditionalPutAndUpdateGiven_ShouldExecuteProperly() {
  // Arrange
  configureBehavior();
  mutations = prepareConditionalPuts();
  mutations.get(1).withCondition(new PutIfExists());
  when(session.execute(any(Statement.class))).thenReturn(results);
  when(results.wasApplied()).thenReturn(true);

  // Act Assert
  assertThatCode(
          () -> {
            batch.handle(mutations);
          })
      .doesNotThrowAnyException();

  // Assert
  verify(insert).prepare(mutations.get(0));
  verify(insert).bind(prepared, mutations.get(0));
  verify(update).prepare(mutations.get(1));
  verify(update).bind(prepared, mutations.get(1));
}
 

private void retryQuery(String id, Statement query, final long startTime, int retryCount,
    DriverException e) throws DriverException {
  if (retryCount >= maxWriteRetries) {
    logger.error("[{}]: Query aborted after {} retry: ", id, retryCount, e.getMessage());
    metricFailed.inc(((BatchStatement) query).size());
    commitTimer.update(System.nanoTime() - startTime, TimeUnit.NANOSECONDS);
    throw e;
  } else {
    logger.warn("[{}]: Query failed, retrying {} of {}: {} ", id, retryCount, maxWriteRetries,
        e.getMessage());

    try {
      Thread.sleep(1000 * (1 << retryCount));
    } catch (InterruptedException ie) {
      logger.debug("[{}]: Interrupted: {}", id, ie);
    }
    _executeQuery(id, query, startTime, retryCount++);
  }
}
 

@Override
public Stream<ModelEntity> streamDeviceModelByPlaceId(UUID placeId, boolean includeTombstoned) {
   if(placeId == null) {
      return Stream.empty();
   }

   try (Context timerContext = streamDeviceModelByPlaceIdTimer.time())
   {
      Statement associationQuery = findIdsByPlace.bind(placeId);

      Function<Row, UUID> entityIdTransform =
         row -> row.getUUID("devid");

      Function<ResultSet, ModelEntity> entityTransform =
         resultSet -> toModel(resultSet.one(), includeTombstoned);

      return listByAssociation(associationQuery, entityIdTransform, entityTransform, asyncTimeoutMs).stream();
   }
}
 

protected E findOneByStatement(Statement statement) {
    E object = null;
    if (statement != null) {
        statement.setConsistencyLevel(cluster.getDefaultReadConsistencyLevel());
        ResultSet resultSet = getSession().execute(statement);
        Result<E> result = getMapper().map(resultSet);
        if (result != null) {
            object = result.one();
        }
    }
    return object;
}
 

private Statement doInsert(UUID placeId, UUID recordingId, long expiration, long actualTtlInSeconds, String fieldName, String value, Long size) {
	Statement insert = QueryBuilder.insertInto(getTableSpace(), TABLE_NAME).using(QueryBuilder.ttl((int)actualTtlInSeconds))
			.values(COLUMNS, new Object[]{placeId, fieldName, expiration, value, recordingId, size})
			.setRetryPolicy(new LoggingRetryPolicy(DowngradingConsistencyRetryPolicy.INSTANCE))			
			;
	return insert;
}
 

private Statement doInsert(UUID placeId, UUID recordingId, String fieldName, String value, Long size) {
	Statement insert = QueryBuilder.insertInto(getTableSpace(), TABLE_NAME)
			.values(COLUMNS, new Object[]{placeId, fieldName, value, recordingId, size})
			.setRetryPolicy(new LoggingRetryPolicy(DowngradingConsistencyRetryPolicy.INSTANCE))			
			;
	return insert;
}
 

public static List<Statement> getDataCleanupStatements(String tableSpace) {
	List<Statement> stmts = new ArrayList<>();
	stmts.add(QueryBuilder.truncate(tableSpace, PlaceRecordingIndexV2Table.TABLE_NAME));
	stmts.add(QueryBuilder.truncate(tableSpace, PlaceRecordingIndexV2FavoriteTable.TABLE_NAME));
	stmts.add(QueryBuilder.truncate(tableSpace, PurgeRecordingV2Table.TABLE_NAME));
	stmts.add(QueryBuilder.truncate(tableSpace, RecordingV2Table.TABLE_NAME));
	stmts.add(QueryBuilder.truncate(tableSpace, RecordingV2FavoriteTable.TABLE_NAME));
	stmts.add(QueryBuilder.truncate(tableSpace, VideoMetadataV2Table.TABLE_NAME));
	stmts.add(QueryBuilder.truncate(tableSpace, VideoMetadataV2FavoriteTable.TABLE_NAME));		
	return stmts;
}
 

@Override
public ListenableFuture<?> insertIFrame(UUID recordingId, double tsInSeconds, long frameByteOffset, long frameByteSize, long ttlInSeconds) {
	Statement stmt = recordingTable.insertIFrame(recordingId, ttlInSeconds, tsInSeconds, frameByteOffset, toblob(frameByteSize));
	long startTime = System.nanoTime();
	ListenableFuture<?> result = session.executeAsync(stmt);
	result.addListener(() -> InsertFrameTimer.update(System.nanoTime() - startTime, TimeUnit.NANOSECONDS), MoreExecutors.directExecutor());
	return result;
}
 

public boolean run() throws Exception
{
    List<BoundStatement> stmts = new ArrayList<>();
    partitionCount = partitions.size();

    for (PartitionIterator iterator : partitions)
        while (iterator.hasNext())
            stmts.add(bindRow(iterator.next()));

    rowCount += stmts.size();

    // 65535 is max number of stmts per batch, so if we have more, we need to manually batch them
    for (int j = 0 ; j < stmts.size() ; j += 65535)
    {
        List<BoundStatement> substmts = stmts.subList(j, Math.min(j + stmts.size(), j + 65535));
        Statement stmt;
        if (stmts.size() == 1)
        {
            stmt = substmts.get(0);
        }
        else
        {
            BatchStatement batch = new BatchStatement(batchType);
            batch.setConsistencyLevel(JavaDriverClient.from(cl));
            batch.addAll(substmts);
            stmt = batch;
        }

        try
        {
            validate(client.getSession().execute(stmt));
        }
        catch (ClassCastException e)
        {
            e.printStackTrace();
        }
    }
    return true;
}
 

@Override
public RetryDecision onReadTimeout(
    Statement statement, ConsistencyLevel cl, int requiredResponses, int receivedResponses,
    boolean dataRetrieved, int nbRetry
) {
    if (nbRetry < numRetries) {
        final String stmt;

        if (statement instanceof BoundStatement) {
            final BoundStatement bound = (BoundStatement) statement;
            final List<Object> variables = new ArrayList<>();

            for (int i = 0; i < bound.preparedStatement().getVariables().size(); i++) {
                variables.add(bound.getObject(i));
            }

            stmt = String.format("%s (%s)", bound.preparedStatement().getQueryString(),
                variables.toString());
        } else {
            stmt = statement.toString();
        }

        final boolean tryNextHost = nbRetry % rotateHost == (rotateHost - 1);

        log.info("Request failing ({}/{}) retrying ({}) (decision: {})...", nbRetry, numRetries,
            stmt, tryNextHost ? "tryNextHost" : "retry");

        return tryNextHost ? RetryDecision.tryNextHost(cl) : RetryDecision.retry(cl);
    }

    return receivedResponses >= requiredResponses && !dataRetrieved ? RetryDecision.retry(cl)
        : RetryDecision.rethrow();
}
 

@Override
public RetryDecision onReadTimeout(Statement stmnt, ConsistencyLevel cl, int requiredResponses,
    int receivedResponses, boolean dataReceived, int rTime) {
  if (dataReceived) {
    return RetryDecision.ignore();
  } else if (rTime < readAttempts) {
    return receivedResponses >= requiredResponses ? RetryDecision.retry(cl)
        : RetryDecision.rethrow();
  } else {
    return RetryDecision.rethrow();
  }

}
 

public Statement insertField(UUID recordingId, long expiration, long actualTtlInSeconds, MetadataAttribute attribute, String value) {
	Statement insert = QueryBuilder.insertInto(getTableSpace(), TABLE_NAME).using(QueryBuilder.ttl((int)actualTtlInSeconds))
			.values(COLUMNS, new Object[]{recordingId, expiration, attribute.name().toLowerCase(), value})
			.setRetryPolicy(new LoggingRetryPolicy(DowngradingConsistencyRetryPolicy.INSTANCE))			
			;
	return insert;		
}
 

/**
 * @param ses Session.
 * @param stmt Statement.
 * @param ctrl Control.
 * @param log Logger.
 * @param clo Closure for loaded values.
 */
public LoadCacheCustomQueryWorker(CassandraSession ses, Statement stmt, PersistenceController ctrl,
                                  IgniteLogger log, IgniteBiInClosure<K, V> clo) {
    this.ses = ses;
    this.stmt = stmt;
    this.ctrl = ctrl;
    this.log = log;
    this.clo = clo;
}
 

public ListenableFuture<?> writeAsync(Statement statement) throws Exception {
    if (statement.getConsistencyLevel() == null && writeConsistencyLevel != null) {
        statement.setConsistencyLevel(writeConsistencyLevel);
    }
    return throttleWrite(() -> {
        // for now, need to record metrics in the same method because CassandraWriteMetrics
        // relies on some thread locals
        cassandraWriteMetrics.recordMetrics(statement);
        return wrappedSession.executeAsync(statement);
    });
}
 

@Override
public RetryDecision onRequestError(
    final Statement statement, final ConsistencyLevel cl, final DriverException e,
    final int nbRetry
) {
    return RetryDecision.rethrow();
}
 

@Test
public void handle_WasAppliedReturnedFalse_ShouldThrowRetriableExecutionException() {
  // Arrange
  configureBehavior();
  mutations = prepareConditionalPuts();
  when(session.execute(any(Statement.class))).thenReturn(results);
  when(results.wasApplied()).thenReturn(false);

  // Act Assert
  assertThatThrownBy(
          () -> {
            batch.handle(mutations);
          })
      .isInstanceOf(NoMutationException.class);
}
 

private static ResultSet executeQueryWithFreshSession(
    Statement statement,
    String contactPoint,
    Session session,
    com.datastax.driver.core.Cluster cluster) {
  if (session == null) {
    cluster = defaultBuilder().addContactPoint(contactPoint).build();
    session = cluster.connect();
  }
  ResultSet rs = session.execute(statement);
  cluster.close();
  ;
  return rs;
}
 

private void testGenerateUpdate(String table, String updateKeys, String updateMethod, List<Tuple<String, Object>> records, String expected) {
    Map<String, Object> recordContentMap = records.stream()
            .collect(Collectors.toMap(Tuple::getKey, Tuple::getValue));

    List<String> fieldNames = records.stream().map(Tuple::getKey).collect(Collectors.toList());

    when(schema.getFieldNames()).thenReturn(fieldNames);
    Statement actual = testSubject.generateUpdate(table, schema, updateKeys, updateMethod, recordContentMap);

    assertEquals(expected, actual.toString());
}
 

private Statement generateInsert(String cassandraTable, RecordSchema schema, Map<String, Object> recordContentMap) {
    Insert insertQuery;
    if (cassandraTable.contains(".")) {
        String[] keyspaceAndTable = cassandraTable.split("\\.");
        insertQuery = QueryBuilder.insertInto(keyspaceAndTable[0], keyspaceAndTable[1]);
    } else {
        insertQuery = QueryBuilder.insertInto(cassandraTable);
    }
    for (String fieldName : schema.getFieldNames()) {
        insertQuery.value(fieldName, recordContentMap.get(fieldName));
    }
    return insertQuery;
}
 

@Override
public RetryDecision onWriteTimeout(final Statement statement, final ConsistencyLevel cl, final WriteType writeType,
                                    final int requiredAcks, final int receivedAcks, final int nbRetry) {

  final RetryDecision decision = delegate.onWriteTimeout(statement, cl, writeType, requiredAcks, receivedAcks, nbRetry);
  for (final TagMetrics c : tagMetrics) {
    c.writeTimeouts.inc();
  }
  return decision;
}
 

@Override
protected List<Statement> prepareIndexInserts(String id, Hub entity) {
   List<Statement> indexInserts = new ArrayList<>();
   addMacAddrIndexInsert(indexInserts, id, entity);
   addAccountIndexInsert(indexInserts, id, entity);
   addPlaceIndexInsert(indexInserts, id, entity);
   return indexInserts;
}