code: /*
   * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.integrations.destination.cassandra;
  
  import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.now;
  
  import com.datastax.oss.driver.api.core.CqlSession;
  import com.datastax.oss.driver.api.core.cql.BatchStatement;
  import com.datastax.oss.driver.api.core.cql.BatchType;
  import com.datastax.oss.driver.api.core.cql.BoundStatement;
  import com.datastax.oss.driver.api.core.cql.PreparedStatement;
  import com.datastax.oss.driver.api.core.metadata.TokenMap;
  import com.datastax.oss.driver.api.core.type.DataTypes;
  import com.datastax.oss.driver.api.core.uuid.Uuids;
  import com.datastax.oss.driver.api.querybuilder.QueryBuilder;
  import com.datastax.oss.driver.api.querybuilder.SchemaBuilder;
  import io.airbyte.integrations.base.JavaBaseConstants;
  import java.io.Closeable;
  import java.time.Instant;
  import java.util.List;
  import java.util.UUID;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.Future;
  import java.util.stream.Collectors;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  class CassandraCqlProvider implements Closeable {
  
    private static final Logger LOGGER = LoggerFactory.getLogger(CassandraCqlProvider.class);
  
    private static final int N_THREADS = Runtime.getRuntime().availableProcessors();
  
    private final ExecutorService executorService;
  
    private final CqlSession cqlSession;
  
    private final CassandraConfig cassandraConfig;
  
    private final String columnId;
  
    private final String columnData;
  
    private final String columnTimestamp;
  
    public CassandraCqlProvider(CassandraConfig cassandraConfig) {
      this.cassandraConfig = cassandraConfig;
      this.cqlSession = SessionManager.initSession(cassandraConfig);
      var nameTransformer = new CassandraNameTransformer(cassandraConfig);
      this.columnId = nameTransformer.outputColumn(JavaBaseConstants.COLUMN_NAME_AB_ID);
      this.columnData = nameTransformer.outputColumn(JavaBaseConstants.COLUMN_NAME_DATA);
      this.columnTimestamp = nameTransformer.outputColumn(JavaBaseConstants.COLUMN_NAME_EMITTED_AT);
      this.executorService = Executors.newFixedThreadPool(N_THREADS);
    }
  
    public void createKeySpaceIfNotExists(String keyspace, int replicationFactor) {
      var query = SchemaBuilder.createKeyspace(keyspace)
          .ifNotExists()
          .withSimpleStrategy(replicationFactor)
          .build();
      cqlSession.execute(query);
    }
  
    public void createTableIfNotExists(String keyspace, String tableName) {
      var query = SchemaBuilder.createTable(keyspace, tableName)
          .ifNotExists()
          .withPartitionKey(columnId, DataTypes.UUID)
          .withColumn(columnData, DataTypes.TEXT)
          .withColumn(columnTimestamp, DataTypes.TIMESTAMP)
          .build();
      cqlSession.execute(query);
    }
  
    public void dropTableIfExists(String keyspace, String tableName) {
      var query = SchemaBuilder.dropTable(keyspace, tableName)
          .ifExists()
          .build();
      cqlSession.execute(query);
    }
  
    public void insert(String keyspace, String tableName, String jsonData) {
      var query = QueryBuilder.insertInto(keyspace, tableName)
          .value(columnId, QueryBuilder.literal(Uuids.random()))
          .value(columnData, QueryBuilder.literal(jsonData))
          .value(columnTimestamp, QueryBuilder.toTimestamp(now()))
          .build();
      cqlSession.execute(query);
    }
  
    public void truncate(String keyspace, String tableName) {
      var query = QueryBuilder.truncate(keyspace, tableName).build();
      cqlSession.execute(query);
    }
  
    public List<CassandraRecord> select(String keyspace, String tableName) {
      var query = QueryBuilder.selectFrom(keyspace, tableName)
          .columns(columnId, columnData, columnTimestamp)
          .build();
      return cqlSession.execute(query)
          .map(result -> new CassandraRecord(
              result.get(columnId, UUID.class),
              result.get(columnData, String.class),
              result.get(columnTimestamp, Instant.class)))
          .all();
    }
  
    public List<Tuple<String, List<String>>> retrieveMetadata() {
      return cqlSession.getMetadata().getKeyspaces().values().stream()
          .map(keyspace -> Tuple.of(keyspace.getName().toString(), keyspace.getTables().values()
              .stream()
              .map(table -> table.getName().toString())
              .collect(Collectors.toList())))
          .collect(Collectors.toList());
    }
  
    public void copy(String keyspace, String sourceTable, String destinationTable) {
      var select = String.format("SELECT * FROM %s.%s WHERE token(%s) > ? AND token(%s) <= ?",
          keyspace, sourceTable, columnId, columnId);
  
      var selectStatement = cqlSession.prepare(select);
  
      var insert = String.format("INSERT INTO %s.%s (%s, %s, %s) VALUES (?, ?, ?)",
          keyspace, destinationTable, columnId, columnData, columnTimestamp);
  
      var insertStatement = cqlSession.prepare(insert);
  
      // perform full table scan in parallel using token ranges
      // optimal for copying large amounts of data
      cqlSession.getMetadata().getTokenMap()
          .map(TokenMap::getTokenRanges)
          .orElseThrow(IllegalStateException::new)
          .stream()
          .flatMap(range -> range.unwrap().stream())
          .map(range -> selectStatement.bind(range.getStart(), range.getEnd()))
          // explore datastax 4.x async api as an alternative for async processing
          .map(selectBoundStatement -> executorService.submit(() -> batchInsert(selectBoundStatement, insertStatement)))
          .forEach(this::awaitThread);
  
    }
  
    private void batchInsert(BoundStatement select, PreparedStatement insert) {
      // unlogged removes the log record for increased insert speed
      var batchStatement = BatchStatement.builder(BatchType.UNLOGGED);
  
      cqlSession.execute(select).all().stream()
          .map(r -> CassandraRecord.of(
              r.get(columnId, UUID.class),
              r.get(columnData, String.class),
              r.get(columnTimestamp, Instant.class)))
          .map(r -> insert.bind(r.getId(), r.getData(), r.getTimestamp()))
          .forEach(batchStatement::addStatement);
  
      cqlSession.execute(batchStatement.build());
    }
  
    private void awaitThread(Future<?> future) {
      try {
        future.get();
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        LOGGER.error("Interrupted thread while copying data with reason: ", e);
      } catch (ExecutionException e) {
        LOGGER.error("Error while copying data with reason: ", e);
      }
    }
  
    @Override
    public void close() {
      // wait for tasks completion and terminate executor gracefully
      executorService.shutdown();
      // close cassandra session for the given config
      SessionManager.closeSession(cassandraConfig);
    }
  
  }
  