content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package com.datastax.oss.driver.core.cql;
  
  import static org.assertj.core.api.Assertions.assertThat;
  
  import com.datastax.oss.driver.api.core.CqlSession;
  import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
  import com.datastax.oss.driver.api.core.cql.AsyncResultSet;
  import com.datastax.oss.driver.api.core.cql.BatchStatement;
  import com.datastax.oss.driver.api.core.cql.BatchStatementBuilder;
  import com.datastax.oss.driver.api.core.cql.DefaultBatchType;
  import com.datastax.oss.driver.api.core.cql.PreparedStatement;
  import com.datastax.oss.driver.api.core.cql.Row;
  import com.datastax.oss.driver.api.core.cql.SimpleStatement;
  import com.datastax.oss.driver.api.testinfra.ccm.CcmRule;
  import com.datastax.oss.driver.api.testinfra.ccm.SchemaChangeSynchronizer;
  import com.datastax.oss.driver.api.testinfra.session.SessionRule;
  import com.datastax.oss.driver.api.testinfra.session.SessionUtils;
  import com.datastax.oss.driver.categories.ParallelizableTests;
  import java.util.Iterator;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.CompletionStage;
  import java.util.function.Function;
  import org.junit.BeforeClass;
  import org.junit.ClassRule;
  import org.junit.Test;
  import org.junit.experimental.categories.Category;
  import org.junit.rules.RuleChain;
  import org.junit.rules.TestRule;
  
  @Category(ParallelizableTests.class)
  public class AsyncResultSetIT {
  
    private static final int PAGE_SIZE = 100;
    private static final int ROWS_PER_PARTITION = 1000;
    private static final String PARTITION_KEY1 = "part";
    private static final String PARTITION_KEY2 = "part2";
  
    private static final CcmRule CCM_RULE = CcmRule.getInstance();
  
    private static final SessionRule<CqlSession> SESSION_RULE =
        SessionRule.builder(CCM_RULE)
            .withConfigLoader(
                SessionUtils.configLoaderBuilder()
                    .withInt(DefaultDriverOption.REQUEST_PAGE_SIZE, PAGE_SIZE)
                    .build())
            .build();
  
    @ClassRule
    public static final TestRule CHAIN = RuleChain.outerRule(CCM_RULE).around(SESSION_RULE);
  
    @BeforeClass
    public static void setupSchema() {
      // create table and load data across two partitions so we can test paging across tokens.
      SchemaChangeSynchronizer.withLock(
          () -> {
            SESSION_RULE
                .session()
                .execute(
                    SimpleStatement.builder(
                            "CREATE TABLE IF NOT EXISTS test (k0 text, k1 int, v int, PRIMARY KEY(k0, k1))")
                        .setExecutionProfile(SESSION_RULE.slowProfile())
                        .build());
          });
  
      PreparedStatement prepared =
          SESSION_RULE.session().prepare("INSERT INTO test (k0, k1, v) VALUES (?, ?, ?)");
  
      BatchStatementBuilder batchPart1 = BatchStatement.builder(DefaultBatchType.UNLOGGED);
      BatchStatementBuilder batchPart2 = BatchStatement.builder(DefaultBatchType.UNLOGGED);
      for (int i = 0; i < ROWS_PER_PARTITION; i++) {
        batchPart1.addStatement(prepared.bind(PARTITION_KEY1, i, i));
        batchPart2.addStatement(
            prepared.bind(PARTITION_KEY2, i + ROWS_PER_PARTITION, i + ROWS_PER_PARTITION));
      }
  
      SESSION_RULE
          .session()
          .execute(batchPart1.setExecutionProfile(SESSION_RULE.slowProfile()).build());
      SESSION_RULE
          .session()
          .execute(batchPart2.setExecutionProfile(SESSION_RULE.slowProfile()).build());
    }
  
    @Test
    public void should_only_iterate_over_rows_in_current_page() throws Exception {
      // very basic test that just ensures that iterating over an AsyncResultSet only visits the first
      // page.
      CompletionStage<AsyncResultSet> result =
          SESSION_RULE
              .session()
              .executeAsync(
                  SimpleStatement.builder("SELECT * FROM test where k0 = ?")
                      .addPositionalValue(PARTITION_KEY1)
                      .build());
  
      AsyncResultSet rs = result.toCompletableFuture().get();
  
      // Should only receive rows in page.
      assertThat(rs.remaining()).isEqualTo(PAGE_SIZE);
      assertThat(rs.hasMorePages()).isTrue();
  
      Iterator<Row> rowIt = rs.currentPage().iterator();
      for (int i = 0; i < PAGE_SIZE; i++) {
        Row row = rowIt.next();
        assertThat(row.getString("k0")).isEqualTo(PARTITION_KEY1);
        assertThat(row.getInt("k1")).isEqualTo(i);
        assertThat(row.getInt("v")).isEqualTo(i);
      }
    }
  
    @Test
    public void should_iterate_over_all_pages_asynchronously_single_partition() throws Exception {
      // Validates async paging behavior over single partition.
      CompletionStage<PageStatistics> result =
          SESSION_RULE
              .session()
              .executeAsync(
                  SimpleStatement.builder("SELECT * FROM test where k0 = ?")
                      .addPositionalValue(PARTITION_KEY1)
                      .build())
              .thenCompose(new AsyncResultSetConsumingFunction());
  
      PageStatistics stats = result.toCompletableFuture().get();
  
      assertThat(stats.rows).isEqualTo(ROWS_PER_PARTITION);
      assertThat(stats.pages).isEqualTo((int) Math.ceil(ROWS_PER_PARTITION / (double) PAGE_SIZE));
    }
  
    @Test
    public void should_iterate_over_all_pages_asynchronously_cross_partition() throws Exception {
      // Validates async paging behavior over a range query.
      CompletionStage<PageStatistics> result =
          SESSION_RULE
              .session()
              .executeAsync("SELECT * FROM test")
              .thenCompose(new AsyncResultSetConsumingFunction());
  
      PageStatistics stats = result.toCompletableFuture().get();
  
      assertThat(stats.rows).isEqualTo(ROWS_PER_PARTITION * 2);
      assertThat(stats.pages).isEqualTo((int) Math.ceil(ROWS_PER_PARTITION * 2 / (double) PAGE_SIZE));
    }
  
    private static class PageStatistics {
      int rows;
      int pages;
  
      PageStatistics(int rows, int pages) {
        this.rows = rows;
        this.pages = pages;
      }
    }
  
    private static class AsyncResultSetConsumingFunction
        implements Function<AsyncResultSet, CompletionStage<PageStatistics>> {
  
      // number of rows paged before exercising this function.
      private final int rowsSoFar;
      // number of pages encountered before exercising this function.
      private final int pagesSoFar;
  
      AsyncResultSetConsumingFunction() {
        this(0, 0);
      }
  
      AsyncResultSetConsumingFunction(int rowsSoFar, int pagesSoFar) {
        this.rowsSoFar = rowsSoFar;
        this.pagesSoFar = pagesSoFar;
      }
  
      @Override
      public CompletionStage<PageStatistics> apply(AsyncResultSet result) {
        int consumedRows = rowsSoFar;
  
        // Only count page if it has rows.
        int pages = result.remaining() == 0 ? pagesSoFar : pagesSoFar + 1;
  
        // iterate over page and ensure data is in order.
        for (Row row : result.currentPage()) {
          int v = row.getInt("v");
          if (v != consumedRows) {
            CompletableFuture<PageStatistics> next = new CompletableFuture<>();
            next.completeExceptionally(
                new Exception(String.format("Expected v == %d, got %d.", consumedRows, v)));
            return next;
          }
          consumedRows++;
        }
  
        if (result.hasMorePages()) {
          return result
              .fetchNextPage()
              .thenComposeAsync(new AsyncResultSetConsumingFunction(consumedRows, pages));
        } else {
          CompletableFuture<PageStatistics> next = new CompletableFuture<>();
          next.complete(new PageStatistics(consumedRows, pages));
          return next;
        }
      }
    }
  }
  