com.datastax.driver.core.SimpleStatement Java Examples

@Override
public void transactionMarker() throws Exception {
    BatchStatement batchStatement = new BatchStatement();
    batchStatement.add(new SimpleStatement(
            "INSERT INTO test.users (id,  fname, lname) VALUES (100, 'f100', 'l100')"));
    batchStatement.add(new SimpleStatement(
            "INSERT INTO test.users (id,  fname, lname) VALUES (101, 'f101', 'l101')"));
    PreparedStatement preparedStatement =
            session.prepare("INSERT INTO test.users (id,  fname, lname) VALUES (?, ?, ?)");
    for (int i = 200; i < 210; i++) {
        BoundStatement boundStatement = new BoundStatement(preparedStatement);
        boundStatement.bind(i, "f" + i, "l" + i);
        batchStatement.add(boundStatement);
    }
    batchStatement.add(new SimpleStatement(
            "INSERT INTO test.users (id,  fname, lname) VALUES (300, 'f300', 'l300')"));
    session.execute(batchStatement);
}
 

private void primeAndExecuteQueries(String[] primed, String[] queries) throws Exception {
  SuccessResult result = getSampleSuccessResult();
  for (String primeQuery : primed) {
    server.prime(when(primeQuery).then(result));
  }

  try (com.datastax.driver.core.Cluster driverCluster =
      defaultBuilder(server.getCluster())
          .withRetryPolicy(FallthroughRetryPolicy.INSTANCE)
          .build()) {
    Session session = driverCluster.connect();
    server.getCluster().clearLogs();
    for (String executeQuery : queries) {
      SimpleStatement stmt = new SimpleStatement(executeQuery);
      stmt.setDefaultTimestamp(100);
      session.execute(stmt);
    }
  }
}
 

@Override
public void transactionMarker() throws Exception {
    BatchStatement batchStatement = new BatchStatement();
    batchStatement.add(new SimpleStatement(
            "INSERT INTO test.users (id,  fname, lname) VALUES (100, 'f100', 'l100')"));
    batchStatement.add(new SimpleStatement(
            "INSERT INTO test.users (id,  fname, lname) VALUES (101, 'f101', 'l101')"));
    PreparedStatement preparedStatement =
            session.prepare("INSERT INTO test.users (id,  fname, lname) VALUES (?, ?, ?)");
    for (int i = 200; i < 210; i++) {
        BoundStatement boundStatement = new BoundStatement(preparedStatement);
        boundStatement.bind(i, "f" + i, "l" + i);
        batchStatement.add(boundStatement);
    }
    batchStatement.add(new SimpleStatement(
            "INSERT INTO test.users (id,  fname, lname) VALUES (300, 'f300', 'l300')"));
    session.executeAsync(batchStatement).get();
}
 

@Override
public Cluster deleteCluster(String clusterName) {
  getRepairSchedulesForCluster(clusterName).forEach(schedule -> deleteRepairSchedule(schedule.getId()));
  session.executeAsync(deleteRepairRunByClusterPrepStmt.bind(clusterName));

  getEventSubscriptions(clusterName)
      .stream()
      .filter(subscription -> subscription.getId().isPresent())
      .forEach(subscription -> deleteEventSubscription(subscription.getId().get()));

  Statement stmt = new SimpleStatement(SELECT_REPAIR_UNIT);
  stmt.setIdempotent(true);
  ResultSet results = session.execute(stmt);
  for (Row row : results) {
    if (row.getString("cluster_name").equals(clusterName)) {
      UUID id = row.getUUID("id");
      session.executeAsync(deleteRepairUnitPrepStmt.bind(id));
    }
  }
  Cluster cluster = getCluster(clusterName);
  session.execute(deleteClusterPrepStmt.bind(clusterName));
  return cluster;
}
 

@Override
public Collection<Cluster> getClusters() {
  // cache the clusters list for ten seconds
  if (System.currentTimeMillis() - clustersCacheAge.get() > TimeUnit.SECONDS.toMillis(10)) {
    clustersCacheAge.set(System.currentTimeMillis());
    Collection<Cluster> clusters = Lists.<Cluster>newArrayList();
    for (Row row : session.execute(new SimpleStatement(SELECT_CLUSTER).setIdempotent(Boolean.TRUE))) {
      try {
        clusters.add(parseCluster(row));
      } catch (IOException ex) {
        LOG.error("Failed parsing cluster {}", row.getString("name"), ex);
      }
    }
    clustersCache.set(Collections.unmodifiableCollection(clusters));
  }
  return clustersCache.get();
}
 

private ArgumentMatcher<Statement> getHostStatementMatcher(final Host host, final String query)
        throws Exception {
    return new ArgumentMatcher<Statement>() {
        @Override
        public boolean matches(Object argument) {
            SelectedHostStatement statement = (SelectedHostStatement) argument;

            return ((SimpleStatement)statement.getStatement()).getQueryString().equals(query) &&
                    Objects.equals(statement.getHostCordinator().getAddress(), host.getAddress());
        }

        @Override
        public void describeTo(Description description) {
            description.appendText(format("query:%s host:%s", query, host.getAddress().toString()));
        }
    };
}
 

public int[] executeBatch() throws SQLException
  {
  	int[] returnCounts= new int[batchQueries.size()];
  	List<ResultSetFuture> futures = new ArrayList<ResultSetFuture>();
  	if (logger.isTraceEnabled() || this.connection.debugMode) logger.debug("CQL statements: "+ batchQueries.size());
  	for(String q:batchQueries){
  		if (logger.isTraceEnabled() || this.connection.debugMode) logger.debug("CQL: "+ q);
  		SimpleStatement stmt = new SimpleStatement(q);
  		stmt.setConsistencyLevel(this.connection.defaultConsistencyLevel);
  		ResultSetFuture resultSetFuture = this.connection.getSession().executeAsync(stmt);
  		futures.add(resultSetFuture);
  	}

  	int i=0;
for (ResultSetFuture future : futures){
	future.getUninterruptibly();
	returnCounts[i]=1;
	i++;
}
      
      return returnCounts;
  }
 

private Optional<String> asString(Statement statement) {
    if (statement instanceof PreparedStatement) {
        PreparedStatement preparedStatement = (PreparedStatement) statement;
        return Optional.of(preparedStatement.getQueryString());
    }
    if (statement instanceof SimpleStatement) {
        SimpleStatement simpleStatement = (SimpleStatement) statement;
        return Optional.of(simpleStatement.getQueryString());
    }
    return Optional.empty();
}
 

@Override
public List<UUID> getLeaders() {
  return session.execute(new SimpleStatement(SELECT_LEADERS))
      .all()
      .stream()
      .map(leader -> leader.getUUID("leader_id"))
      .collect(Collectors.toList());
}
 

/**
 * Execute a range query
 * @param destTable
 * @param range
 */
public void executeRangeQuery(final String destTable, final String format, final Hyperrectangle range) {
	System.out.println("# Execute range query in range " + range);

	final Stopwatch stopwatch = Stopwatch.createStarted();
	final TupleBuilder tupleBuilder = TupleBuilderFactory.getBuilderForFormat(format);

	long readRecords = 0;
	long resultRecords = 0;

	final SimpleStatement statement = new SimpleStatement("SELECT * FROM " + destTable);
	statement.setFetchSize(2000); // Use 2000 tuples per page

	final ResultSet rs = session.execute(statement);

	for (final Row row : rs) {

		// Request next page
	    if (rs.getAvailableWithoutFetching() == 100 && !rs.isFullyFetched()) {
	        rs.fetchMoreResults(); // this is asynchronous
	    }

	    readRecords++;

	    final long id = row.getLong(0);
	    final String text = row.getString(1);

	    final Tuple tuple = tupleBuilder.buildTuple(text, Long.toString(id));

	    if(tuple.getBoundingBox().intersects(range)) {
	    	resultRecords++;
	    }
	}

	System.out.println("# Read records " + readRecords + " result records " + resultRecords);
	System.out.println("# Execution time in sec " + stopwatch.elapsed(TimeUnit.SECONDS));
}
 

/**
 * This executes the query to retrieve result from database.
 * It then converts each row into tuple and emit that into output port.
 */
@Override
public void emitTuples()
{
  String query = queryToRetrieveData();
  logger.debug("select statement: {}", query);

  SimpleStatement stmt = new SimpleStatement(query);
  stmt.setFetchSize(fetchSize);
  try {
    if (nextPageState != null) {
      stmt.setPagingState(nextPageState);
    }
    ResultSet result = store.getSession().execute(stmt);
    nextPageState = result.getExecutionInfo().getPagingState();

    if (!result.isExhausted()) {
      for (Row row : result) {
        T tuple = getTuple(row);
        emit(tuple);
      }
    } else {
      // No rows available wait for some time before retrying so as to not continuously slam the database
      Thread.sleep(waitForDataTimeout);
    }
  } catch (Exception ex) {
    store.disconnect();
    DTThrowable.rethrow(ex);
  }
}
 

@Test
public void testCreateExitSpan() throws Throwable {
    interceptor.beforeMethod(objectInstance, method, new Object[] {new SimpleStatement("SELECT * FROM test")}, null, null);
    interceptor.afterMethod(objectInstance, method, new Object[] {new SimpleStatement("SELECT * FROM test")}, null, null);

    assertThat(segmentStorage.getTraceSegments().size(), is(1));
    TraceSegment segment = segmentStorage.getTraceSegments().get(0);
    assertThat(SegmentHelper.getSpans(segment).size(), is(1));
    AbstractTracingSpan span = SegmentHelper.getSpans(segment).get(0);
    SpanAssert.assertLayer(span, SpanLayer.DB);
    assertThat(span.getOperationName(), is(Constants.CASSANDRA_OP_PREFIX));
    SpanAssert.assertTag(span, 0, Constants.CASSANDRA_DB_TYPE);
    SpanAssert.assertTag(span, 1, "test");
    SpanAssert.assertTag(span, 2, "SELECT * FROM test");
}
 

@Override
public List<ODataEntry> readTable(String tableName) throws ODataServiceFault {
    Statement statement = new SimpleStatement("Select * from " + this.keyspace + "." + tableName);
    ResultSet resultSet = this.session.execute(statement);
    Iterator<Row> iterator = resultSet.iterator();
    List<ODataEntry> entryList = new ArrayList<>();
    ColumnDefinitions columnDefinitions = resultSet.getColumnDefinitions();
    while (iterator.hasNext()) {
        ODataEntry dataEntry = createDataEntryFromRow(tableName, iterator.next(), columnDefinitions);
        entryList.add(dataEntry);
    }
    return entryList;
}
 

@Override
public Mono<Result> check() {
    // execute a simple query to check if cassandra is responding
    // idea from: https://stackoverflow.com/questions/10246287
    return queryExecutor.execute(new SimpleStatement(SAMPLE_QUERY))
        .map(resultSet -> Result.healthy(COMPONENT_NAME))
        .onErrorResume(e -> Mono.just(Result.unhealthy(COMPONENT_NAME, "Error checking Cassandra backend", e)));
}
 

private void printStatement(Statement statement) {
    if (printStatements) {
        if (statement instanceof BoundStatement) {
            BoundStatement boundStatement = (BoundStatement) statement;
            System.out.println("Executing: " + boundStatement.preparedStatement().getQueryString());
        } else if (statement instanceof SimpleStatement) {
            SimpleStatement simpleStatement = (SimpleStatement) statement;
            System.out.println("Executing: " + simpleStatement.getQueryString());
        } else {
            System.out.println("Executing: " + statement);
        }
    }
}
 

public ResultSet read(String query) throws Exception {
    if (!query.startsWith("select ")) {
        throw new IllegalStateException("Unexpected read query: " + query);
    }
    // do not use session.execute() because that calls getUninterruptibly() which can cause
    // central shutdown to timeout while waiting for executor service to shutdown
    return readAsync(new SimpleStatement(query)).get();
}
 

@Override
@SuppressWarnings("unchecked")
public boolean matches(Object obj) {
    if (obj != null && obj instanceof SimpleStatement) {
        SimpleStatement receivedStmt = (SimpleStatement) obj;
        Object[] expectedValues = Whitebox.getInternalState(expectedStmt, "values");
        Object[] receivedValues = Whitebox.getInternalState(receivedStmt, "values");

        return matchValues(expectedValues, receivedValues);
    }
    return false;
}
 

@Override
public Optional<RepairUnit> getRepairUnit(RepairUnit.Builder params) {
  // brute force again
  RepairUnit repairUnit = null;
  Statement stmt = new SimpleStatement(SELECT_REPAIR_UNIT);
  stmt.setIdempotent(Boolean.TRUE);
  ResultSet results = session.execute(stmt);
  for (Row repairUnitRow : results) {
    if (repairUnitRow.getString("cluster_name").equals(params.clusterName)
        && repairUnitRow.getString("keyspace_name").equals(params.keyspaceName)
        && repairUnitRow.getSet("column_families", String.class).equals(params.columnFamilies)
        && repairUnitRow.getBool("incremental_repair") == params.incrementalRepair
        && repairUnitRow.getSet("nodes", String.class).equals(params.nodes)
        && repairUnitRow.getSet("datacenters", String.class).equals(params.datacenters)
        && repairUnitRow
            .getSet("blacklisted_tables", String.class)
            .equals(params.blacklistedTables)
        && repairUnitRow.getInt("repair_thread_count") == params.repairThreadCount) {

      repairUnit = RepairUnit.builder()
          .clusterName(repairUnitRow.getString("cluster_name"))
          .keyspaceName(repairUnitRow.getString("keyspace_name"))
          .columnFamilies(repairUnitRow.getSet("column_families", String.class))
          .incrementalRepair(repairUnitRow.getBool("incremental_repair"))
          .nodes(repairUnitRow.getSet("nodes", String.class))
          .datacenters(repairUnitRow.getSet("datacenters", String.class))
          .blacklistedTables(repairUnitRow.getSet("blacklisted_tables", String.class))
          .repairThreadCount(repairUnitRow.getInt("repair_thread_count"))
          .build(repairUnitRow.getUUID("id"));
      // exit the loop once we find a match
      break;
    }
  }

  return Optional.ofNullable(repairUnit);
}
 

public static BatchStatement makeNativeBatchStatement(List<String> queries, List<List> values) {

    BatchStatement statement = new BatchStatement();
    Iterator<List> valuesIterator = values.iterator();
    for (String query : queries) {
      List value = valuesIterator.next();
      statement.add(new SimpleStatement(query, value.toArray(new Object[value.size()])));
    }
    return statement;
  }
 

@Test
public void testShouldReturnClientTimeout() throws Exception {
  server.prime(when(query));

  thrown.expect(OperationTimedOutException.class);
  query(new SimpleStatement(query).setReadTimeoutMillis(1000));
}
 

@Override
public Collection<RepairSchedule> getAllRepairSchedules() {
  Collection<RepairSchedule> schedules = Lists.<RepairSchedule>newArrayList();
  Statement stmt = new SimpleStatement(SELECT_REPAIR_SCHEDULE);
  stmt.setIdempotent(Boolean.TRUE);
  ResultSet scheduleResults = session.execute(stmt);
  for (Row scheduleRow : scheduleResults) {
    schedules.add(createRepairScheduleFromRow(scheduleRow));
  }

  return schedules;
}
 

private List<Row> loadMigrations(String tablePrefix) {
    Session session = cassandra.getCluster().connect(CassandraJUnitRule.TEST_KEYSPACE);
    if (tablePrefix == null || tablePrefix.isEmpty()) {
        return session.execute(
                new SimpleStatement("SELECT * FROM schema_migration;")).all();
    }
    return session.execute(
            new SimpleStatement(String.format("SELECT * FROM %s_schema_migration;", tablePrefix))).all();
}
 

public static ResultSet makeNativeQueryWithPositionalParams(
    String query, String contactPoint, Object... values) {
  SimpleStatement statement = new SimpleStatement(query, values);
  return executeQueryWithFreshSession(statement, contactPoint);
}
 

@Override
public Object runPreQuery(InternalParamCollection params, int queryLevel) throws DataServiceFault {
    ResultSet rs = null;
     /*
        There is no point of creating prepared statements for dynamic queries
     */
    if (isDynamicQuery(params)) {
        Object[] result = this.processDynamicQuery(this.getQuery(), params);
        String dynamicCql = (String) result[0];
        int currentParamCount = (Integer) result[1];
        String processedSQL = this.createProcessedQuery(dynamicCql, params, currentParamCount);
        if (DispatchStatus.isBatchRequest() && this.isNativeBatchRequestsSupported()) {
            /* handle batch requests */
            if (DispatchStatus.isFirstBatchRequest()) {
                this.batchStatement.set(new BatchStatement());
            }
            SimpleStatement simpleStatement = new SimpleStatement(processedSQL);
            this.batchStatement.get().add(simpleStatement);
            if (DispatchStatus.isLastBatchRequest()) {
                this.getSession().execute(this.batchStatement.get());
            }
        } else {
            SimpleStatement statement = new SimpleStatement(processedSQL, this.bindParams(params));
            rs = this.getSession().execute(statement);
        }
    } else {
        this.checkAndCreateStatement();
        if (DispatchStatus.isBatchRequest() && this.isNativeBatchRequestsSupported()) {
            /* handle batch requests */
            if (DispatchStatus.isFirstBatchRequest()) {
                this.batchStatement.set(new BatchStatement());
            }
            this.batchStatement.get().add(this.getStatement().bind(this.bindParams(params)));
            if (DispatchStatus.isLastBatchRequest()) {
                this.getSession().execute(this.batchStatement.get());
            }
        } else {
            rs = this.getSession().execute(this.getStatement().bind(this.bindParams(params)));
        }
    }
    return rs;
}
 

public StatementMatcher(SimpleStatement expectedStmt) {
    this.expectedStmt = expectedStmt;
}
 

@Test
public void testEqualsInForDeepRecordReader() {

    // Static stubs
    PowerMockito.mockStatic(CassandraClientProvider.class);

    // Stubbing
    when(config.getPageSize()).thenReturn(PAGE_SIZE_CONSTANT);
    when(config.getTable()).thenReturn(TABLE_NAME_CONSTANT);
    when(config.getInputColumns()).thenReturn(INPUT_COLUMNS_CONSTANT);
    when(config.fetchTableMetadata()).thenReturn(tableMetadata);
    when(config.getEqualsInValue()).thenReturn(equalsInValue);
    when(config.getPartitionerClassName()).thenReturn("org.apache.cassandra.dht.Murmur3Partitioner");

    when(tokenRange.getReplicas()).thenReturn(Arrays.asList(LOCALHOST_CONSTANT));
    when(tokenRange.getStartTokenAsComparable()).thenReturn(-8000000000000000000L, -7000000000000000000L,
            2200000000000000000L, 2300000000000000000L, 2600000000000000000L);
    when(tokenRange.getEndTokenAsComparable()).thenReturn(-7000000000000000000L, 2200000000000000000L,
            2300000000000000000L, 2600000000000000000L, -8000000000000000000L);

    when(CassandraClientProvider.trySessionForLocation(any(String.class), any(CassandraDeepJobConfig.class),
            any(Boolean.class))).thenReturn(Pair.create(session, LOCALHOST_CONSTANT));

    when(tableMetadata.getPartitionKey()).thenReturn(Arrays.asList(columnMetadata, columnMetadata));
    when(tableMetadata.getClusteringColumns()).thenReturn(new ArrayList<ColumnMetadata>());

    when(columnMetadata.getName()).thenReturn("col1", "col2");
    when(columnMetadata.getType()).thenReturn(DataType.bigint());

    when(equalsInValue.getEqualsList()).thenReturn(Arrays.asList(Pair.create("col1", (Serializable) 1L)));
    when(equalsInValue.getInField()).thenReturn("col2");
    when(equalsInValue.getInValues()).thenReturn(
            Arrays.asList((Serializable) 1L, (Serializable) 2L, (Serializable) 3L, (Serializable) 4L,
                    (Serializable) 5L));

    Object[] values = { 1L, Arrays.asList(1L, 4L) };
    SimpleStatement stmt = new SimpleStatement(
            "SELECT \"col1\",\"col2\",\"col3\" FROM \"TABLENAME\" WHERE col1 = ? AND col2 IN ?", values);
    when(session.execute(any(Statement.class))).thenReturn(resultSet);

    DeepRecordReader recordReader = new DeepRecordReader(config, tokenRange);

    // TODO How do we check two statements are the same?
    verify(session, times(1)).execute(Matchers.argThat(new StatementMatcher(stmt)));
}
 

@Override
public Observable<ResultSet> execute(String query, Scheduler scheduler) {
    return scheduleStatement(new SimpleStatement(query), scheduler);
}
 

@Override
public Observable<ResultSet> execute(String query) {
    return scheduleStatement(new SimpleStatement(query), Schedulers.computation());
}
 

Observable<ResultSet> createTemporaryTable(String tempTableName) {
    return Observable.just(tempTableName)
            .map(t -> new SimpleStatement(String.format(TempStatement.CREATE_TABLE.getStatement(), t)))
            .flatMap(st -> rxSession.execute(st));
}
 

public void dropCassandraTable(String tableName) {
  String drop_stmt = String.format("DROP TABLE IF EXISTS %s;", tableName);
  getCassandraClient().execute(new SimpleStatement(drop_stmt).setReadTimeoutMillis(60000));
  LOG.info("Dropped Cassandra table " + tableName + " using query: [" + drop_stmt + "]");
}