com.datastax.oss.driver.api.core.CqlSession Java Examples

@Override
protected void before() throws Throwable {

	dependency.before();

	try (CqlSession session = CqlSession.builder().addContactPoint(new InetSocketAddress(getHost(), getPort()))
			.withLocalDatacenter("datacenter1").build()) {

		if (requiredVersion != null) {

			Version cassandraReleaseVersion = CassandraVersion.getReleaseVersion(session);

			if (cassandraReleaseVersion.isLessThan(requiredVersion)) {
				throw new AssumptionViolatedException(
						String.format("Cassandra at %s:%s runs in Version %s but we require at least %s", getHost(), getPort(),
								cassandraReleaseVersion, requiredVersion));
			}
		}

		session.execute(String.format("CREATE KEYSPACE IF NOT EXISTS %s \n"
				+ "WITH REPLICATION = { 'class' : 'SimpleStrategy', 'replication_factor' : 1 };", keyspaceName));
	}
}
 

private Server createServer() {
    // distributed locks
    LockManager lockManager = new LockManager();
    JanusGraphFactory janusGraphFactory = new JanusGraphFactory(serverConfig);
    HadoopGraphFactory hadoopGraphFactory = new HadoopGraphFactory(serverConfig);

    Integer storagePort = serverConfig.getProperty(ConfigKey.STORAGE_PORT);
    String storageHostname = serverConfig.getProperty(ConfigKey.STORAGE_HOSTNAME);
    // CQL cluster used by KeyspaceManager to fetch all existing keyspaces
    CqlSession cqlSession = CqlSession.builder()
            .addContactPoint(new InetSocketAddress(storageHostname, storagePort))
            .withLocalDatacenter("datacenter1")
            .build();

    sessionFactory = new SessionFactory(lockManager, janusGraphFactory, hadoopGraphFactory, serverConfig);
    keyspaceManager = new KeyspaceManager(cqlSession, janusGraphFactory, sessionFactory);

    OpenRequest requestOpener = new ServerOpenRequest(sessionFactory);

    io.grpc.Server serverRPC = ServerBuilder.forPort(grpcPort)
            .addService(new SessionService(requestOpener))
            .addService(new KeyspaceService(keyspaceManager))
            .build();

    return ServerFactory.createServer(serverRPC);
}
 

protected long insert(CqlSession session) {

        try {
            session.execute(insertInto(KEYSPACE, TABLE)
                                .value(ID, literal("1"))
                                .value(DATE, literal(timeBased()))
                                .value(NAME, literal("stefan1"))
                                .build());

            session.execute(insertInto(TestEntity2.KEYSPACE_2, TestEntity2.TABLE_2)
                                .value(ID, literal("1"))
                                .value(DATE, literal(timeBased()))
                                .value(NAME, literal("stefan1"))
                                .build());

            Thread.sleep(2000);
        } catch (InterruptedException ex) {
            logger.error("Exception while sleeping!");
        }

        return System.currentTimeMillis();
    }
 

@Override
public void sendToNext(Document doc) {
  if (isRemembering()) {
    CqlSession session = getCassandra().getSession();
    PreparedStatement preparedQuery = getCassandra().getPreparedQuery(UPDATE_HASH_U);
    BoundStatement bind = preparedQuery.bind(doc.getHash(), doc.getId(), doc.getSourceScannerName());
    session.execute(bind);
  }
  superSendToNext(doc);
}
 

@Test
public void test(final MockTracer tracer) {
  try (final CqlSession session = createSession()) {
    createKeyspace(session);
  }

  assertEquals(1, tracer.finishedSpans().size());
}
 

private static void useKeyspace(CqlSession session, String keyspace) {
    try {
        session.execute("USE " + keyspace + ";");
    } catch (InvalidQueryException e) {
        session.execute("CREATE KEYSPACE IF NOT EXISTS " + keyspace + " WITH REPLICATION = { 'class' : 'SimpleStrategy', 'replication_factor' : 3 };");
        session.execute("USE " + keyspace + ";");
    }
}
 

@PostConstruct
public void initialize() {
    String cassandraHosts = env.getProperty("ea.cassandra.hosts","localhost:9042");
    String cassandraKeyspaceName = env.getProperty("ea.cassandra.keyspace","\"ElasticActors\"");
    Integer cassandraPort = env.getProperty("ea.cassandra.port", Integer.class, 9042);

    Set<String> hostSet = StringUtils.commaDelimitedListToSet(cassandraHosts);

    String[] contactPoints = new String[hostSet.size()];
    int i=0;
    for (String host : hostSet) {
        if(host.contains(":")) {
            contactPoints[i] = host.substring(0,host.indexOf(":"));
        } else {
            contactPoints[i] = host;
        }
        i+=1;
    }

    List<InetSocketAddress> contactPointAddresses = Arrays.stream(contactPoints)
            .map(host -> new InetSocketAddress(host, cassandraPort))
            .collect(Collectors.toList());

    DriverConfigLoader driverConfigLoader = DriverConfigLoader.programmaticBuilder()
            .withDuration(DefaultDriverOption.HEARTBEAT_INTERVAL, Duration.ofSeconds(60))
            .withString(DefaultDriverOption.REQUEST_CONSISTENCY, ConsistencyLevel.QUORUM.name())
            .withString(DefaultDriverOption.RETRY_POLICY_CLASS, "DefaultRetryPolicy")
            .withString(DefaultDriverOption.RECONNECTION_POLICY_CLASS, "ConstantReconnectionPolicy")
            .withDuration(DefaultDriverOption.RECONNECTION_BASE_DELAY, Duration.ofSeconds(env.getProperty("ea.cassandra.retryDownedHostsDelayInSeconds",Integer.class,1)))
            .withString(DefaultDriverOption.LOAD_BALANCING_POLICY_CLASS, "DcInferringLoadBalancingPolicy")
            .build();


    cassandraSession = CqlSession.builder()
            .addContactPoints(contactPointAddresses)
            .withConfigLoader(driverConfigLoader)
            .withKeyspace(cassandraKeyspaceName)
            .build();
}
 

@Override
public void initialize(GenericApplicationContext context) {
	Supplier<CassandraDataAutoConfiguration> configurationSupplier = () -> new CassandraDataAutoConfiguration(context.getBean(CqlSession.class));

	context.registerBean(CassandraCustomConversions.class, () -> configurationSupplier.get().cassandraCustomConversions());
	context.registerBean(CassandraMappingContext.class, () -> getCassandraMappingContext(context, configurationSupplier));
	context.registerBean(CassandraConverter.class, () -> configurationSupplier.get().cassandraConverter(context.getBean(CassandraMappingContext.class), context.getBean(CassandraCustomConversions.class)));
	context.registerBean(SessionFactoryFactoryBean.class, () -> getCassandraSessionFactoryBean(context, configurationSupplier));
	context.registerBean(CassandraTemplate.class, () -> getCassandraTemplate(context, configurationSupplier));
}
 

@Override
public void initialize(GenericApplicationContext context) {
	CassandraReactiveDataAutoConfiguration configuration = new CassandraReactiveDataAutoConfiguration();

	context.registerBean(ReactiveSession.class, () -> configuration.reactiveCassandraSession(context.getBean(CqlSession.class)));
	context.registerBean(ReactiveSessionFactory.class, () -> configuration.reactiveCassandraSessionFactory(context.getBean(ReactiveSession.class)));
	context.registerBean(ReactiveCassandraTemplate.class, () -> configuration.reactiveCassandraTemplate(context.getBean(ReactiveSession.class), context.getBean(CassandraConverter.class)));
}
 

@Override
public void initialize(GenericApplicationContext context) {
	CassandraAutoConfiguration configuration = new CassandraAutoConfiguration();
	context.registerBean(CqlSession.class, () -> configuration.cassandraSession(context.getBean(CqlSessionBuilder.class)));
	context.registerBean(CqlSessionBuilder.class, () -> configuration.cassandraSessionBuilder(properties, context.getBean(DriverConfigLoader.class), context.getBeanProvider(CqlSessionBuilderCustomizer.class)));
	context.registerBean(DriverConfigLoader.class, () -> configuration.cassandraDriverConfigLoader(properties, context.getBeanProvider(DriverConfigLoaderBuilderCustomizer.class)));
}
 

@Test
public void should_execute_statement_with_timestamp_option() throws Exception {
  //Given
  String statement1 = "INSERT INTO zeppelin.ts(key,val) VALUES('k','v1');";
  String statement2 = "@timestamp=15\n" +
          "INSERT INTO zeppelin.ts(key,val) VALUES('k','v2');";

  CqlSession session = EmbeddedCassandraServerHelper.getSession();
  // Insert v1 with current timestamp
  interpreter.interpret(statement1, intrContext);
  System.out.println("going to read data from zeppelin.ts;");
  session.execute("SELECT val FROM zeppelin.ts LIMIT 1")
          .forEach(x -> System.out.println("row " + x ));

  Thread.sleep(1);

  //When
  // Insert v2 with past timestamp
  interpreter.interpret(statement2, intrContext);
  System.out.println("going to read data from zeppelin.ts;");
  session.execute("SELECT val FROM zeppelin.ts LIMIT 1")
          .forEach(x -> System.out.println("row " + x ));
  final String actual = session.execute("SELECT val FROM zeppelin.ts LIMIT 1").one()
          .getString("val");

  //Then
  assertThat(actual).isEqualTo("v1");
}
 

@Override
public void activate() {
  super.activate();
  if (isRemembering() || isHashing()) {
    CqlSession session = getCassandra().getSession();
    List<DocKey> strandedDocs = new ArrayList<>();
    PreparedStatement preparedQuery = getCassandra().getPreparedQuery(RESET_PROCESSING_Q);
    BoundStatement statement =  preparedQuery.bind(getName());
    ResultSet procRs = session.execute(statement);
    strandedDocs.addAll(procRs.all().stream()
        .map((row) -> new DocKey(row.getString(0), row.getString(1))).collect(Collectors.toList()));
    preparedQuery = getCassandra().getPreparedQuery(RESET_ERROR_Q);
    statement = preparedQuery.bind(getName());
    ResultSet errorRs = session.execute(statement);
    strandedDocs.addAll(errorRs.all().stream()
        .map((row) -> new DocKey(row.getString(0), row.getString(1))).collect(Collectors.toList()));
    preparedQuery = getCassandra().getPreparedQuery(RESET_BATCHED_Q);
    statement = preparedQuery.bind(getName());
    ResultSet batchedRs = session.execute(statement);
    strandedDocs.addAll(batchedRs.all().stream()
        .map((row) -> new DocKey(row.getString(0), row.getString(1))).collect(Collectors.toList()));

    preparedQuery = getCassandra().getPreparedQuery(RESET_DOCS_U);
    // todo: batch
    for (DocKey docId : strandedDocs) {
      statement = preparedQuery.bind(docId.docid, docId.scanner);
      session.execute(statement);
    }

  }
}
 

public void start() {
    logger.info("Contact points {}", config.getContactPoints());

    CqlSessionBuilder builder = CqlSession.builder()
            .addContactPoints(Arrays.stream(config.getContactPoints().split(","))
                    .map(s -> new InetSocketAddress(s, config.getCqlPort()))
                    .collect(Collectors.toList()))
            .withLocalDatacenter(config.getLocalDC());

    if (config.getCqlUserName() != null)
        builder.withAuthCredentials(config.getCqlUserName(), config.getCqlPassword());

    if (config.isDockerCassandra()) {
        logger.info("Docker cassandra enabled in the yaml.");
        logger.info("Attempting to stand up container.");
        DockerHelper dh = new DockerHelper();
        dh.startDSE();

        //TODO
        try {
            Thread.sleep(60000);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

    }

    session = builder.build();


    logger.info("Preparing statements for " + CassandraManager.class.getSimpleName());
    stmts = new CassandraStatements.Prepared(session, config.getKeyspaceName(), config.getReplicationStrategy());

    refreshSchema();
}
 

protected List<Long> insertAndBackup(int records, CqlSession cqlSession, String[] backupArgs) {
    List<Long> executionTimes = insert(records, cqlSession);

    logger.info("Executing backup with arguments {}", asList(backupArgs));
    BackupRestoreCLI.main(backupArgs, false);

    return executionTimes;
}
 

protected void waitForCql() {
    await().until(() -> {
        try (CqlSession cqlSession = CqlSession.builder().build()) {
            return true;
        } catch (Exception ex) {
            return false;
        }
    });
}
 

protected CassandraLog4JManager(String name) {
  super(LoggerContext.getContext(), name);
  System.out.println(">>>> Creating CassandraLog4JManager");
  CassandraSupport cassandra = new CassandraSupport();

  Callable<Object> makeTables = new Callable<Object>() {


    @Override
    public Object call() throws Exception {
      System.out.println("Table and key space creation thread started");
      boolean tryAgain = true;
      int tryCount = 0;
      // ugly but effective fix for https://github.com/nsoft/jesterj/issues/1
      while(tryAgain) {
        try {
          CqlSession session = cassandra.getSession();
          session.execute(CREATE_LOG_KEYSPACE);
          session.execute(CREATE_LOG_TABLE);
          session.execute(CREATE_FT_TABLE);
          session.execute(FTI_STATUS_INDEX);
          session.execute(FTI_SCANNER_INDEX);
          tryAgain = false;
        } catch (Exception e) {
          tryCount++;
          // complain and die if we can't set up the tables in cassandra.
          e.printStackTrace();
          Thread.sleep(1000);
          if (tryCount > 10) {
            die(e);
            tryAgain = false;
          }
        }
      }
      return null;
    }
  };
  this.cassandraReady = cassandra.whenBooted(makeTables);
}
 

protected void dumpTable(final String keyspace, final String table, int expectedLength) {
    try (CqlSession session = CqlSession.builder().build()) {
        List<Row> rows = session.execute(selectFrom(keyspace, table).all().build()).all();

        for (Row row : rows) {
            logger.info(format("id: %s, date: %s, name: %s", row.getString(ID), uuidToDate(requireNonNull(row.getUuid(DATE))), row.getString(NAME)));
        }

        assertEquals(rows.size(), expectedLength);
    }
}
 

public CassandraActorSystemEventListenerRepository(String clusterName, CqlSession cassandraSession) {
    this.clusterName = clusterName;
    this.cassandraSession = cassandraSession;
    this.insertStatement = cassandraSession.prepare(INSERT_QUERY);
    this.deleteStatement = cassandraSession.prepare(DELETE_QUERY);
    this.selectStatement = cassandraSession.prepare(SELECT_QUERY);
}
 

static void cqlAlter() {
    FieldModel model = SampleModels.wrapper();
    try (CqlSession session = CqlSession.builder().build()) {
        model.getFieldInfos().stream()
                .filter(f -> !session.getMetadata()
                        .getKeyspace("sample").get()
                        .getTable("model").get()
                        .getColumn(f.id().code()).isPresent())
                .forEach(f -> session.execute(SchemaBuilder.alterTable("sample", "model")
                        .addColumn(f.id().code(), cqlType(f)).build()));
    }
}
 

@BeforeClass
public static synchronized void setUp() throws IOException, InterruptedException {
  System.setProperty("cassandra.skip_wait_for_gossip_to_settle", "0");
  System.setProperty("cassandra.load_ring_state", "false");
  System.setProperty("cassandra.initial_token", "0");
  System.setProperty("cassandra.num_tokens", "nil");
  System.setProperty("cassandra.allocate_tokens_for_local_replication_factor", "nil");
  EmbeddedCassandraServerHelper.startEmbeddedCassandra();
  CqlSession session = EmbeddedCassandraServerHelper.getSession();
  new CQLDataLoader(session).load(new ClassPathCQLDataSet("prepare_all.cql", "zeppelin"));

  Properties properties = new Properties();
  properties.setProperty(CASSANDRA_CLUSTER_NAME, EmbeddedCassandraServerHelper.getClusterName());
  properties.setProperty(CASSANDRA_COMPRESSION_PROTOCOL, "NONE");
  properties.setProperty(CASSANDRA_CREDENTIALS_USERNAME, "none");
  properties.setProperty(CASSANDRA_CREDENTIALS_PASSWORD, "none");

  properties.setProperty(CASSANDRA_LOAD_BALANCING_POLICY, "DEFAULT");
  properties.setProperty(CASSANDRA_RETRY_POLICY, "DEFAULT");
  properties.setProperty(CASSANDRA_RECONNECTION_POLICY, "DEFAULT");
  properties.setProperty(CASSANDRA_SPECULATIVE_EXECUTION_POLICY, "DEFAULT");

  properties.setProperty(CASSANDRA_POOLING_CONNECTION_PER_HOST_LOCAL, "2");
  properties.setProperty(CASSANDRA_POOLING_CONNECTION_PER_HOST_REMOTE, "1");
  properties.setProperty(CASSANDRA_POOLING_MAX_REQUESTS_PER_CONNECTION, "1024");

  properties.setProperty(CASSANDRA_POOLING_POOL_TIMEOUT_MILLIS, "5000");
  properties.setProperty(CASSANDRA_POOLING_HEARTBEAT_INTERVAL_SECONDS, "30");

  properties.setProperty(CASSANDRA_QUERY_DEFAULT_CONSISTENCY, "ONE");
  properties.setProperty(CASSANDRA_QUERY_DEFAULT_SERIAL_CONSISTENCY, "SERIAL");
  properties.setProperty(CASSANDRA_QUERY_DEFAULT_FETCH_SIZE, "5000");

  properties.setProperty(CASSANDRA_SOCKET_CONNECTION_TIMEOUT_MILLIS, "5000");
  properties.setProperty(CASSANDRA_SOCKET_READ_TIMEOUT_MILLIS, "12000");
  properties.setProperty(CASSANDRA_SOCKET_TCP_NO_DELAY, "true");

  properties.setProperty(CASSANDRA_HOSTS, EmbeddedCassandraServerHelper.getHost());
  properties.setProperty(CASSANDRA_PORT,
          Integer.toString(EmbeddedCassandraServerHelper.getNativeTransportPort()));
  properties.setProperty("datastax-java-driver.advanced.connection.pool.local.size", "1");
  interpreter = new CassandraInterpreter(properties);
  interpreter.open();
}
 

/**
 * Set up the Cassandra Cluster Connection with schema change listener to track the metadata changes
 *
 * @param schemaChangeListener the schema change listener
 */
public void startCassandraConnection(SchemaChangeListener schemaChangeListener) {

    String actionDescription = "start Cassandra Data Store Connection";
    try {
        CqlSessionBuilder builder = CqlSession.builder();
        builder.addContactPoint(new InetSocketAddress(serverAddresses, Integer.valueOf(port)));
        builder.withSchemaChangeListener(schemaChangeListener);
        builder.withAuthCredentials(username, password);
        this.cqlSession = builder.build();

    } catch (ExceptionInInitializerError e) {
        auditLog = CassandraDataStoreAuditCode.CONNECTOR_SERVER_CONNECTION_ERROR;
        omrsAuditLog.logRecord(actionDescription,
                auditLog.getLogMessageId(),
                auditLog.getSeverity(),
                auditLog.getFormattedLogMessage(),
                null,
                auditLog.getSystemAction(),
                auditLog.getUserAction());
    }
    auditLog = CassandraDataStoreAuditCode.CONNECTOR_INITIALIZED;
    omrsAuditLog.logRecord(actionDescription,
            auditLog.getLogMessageId(),
            auditLog.getSeverity(),
            auditLog.getFormattedLogMessage(),
            null,
            auditLog.getSystemAction(),
            auditLog.getUserAction());
}
 

@Test
public void shouldConnectUsingCassandraClient() {

	try (CqlSession session = CqlSession.builder().withLocalDatacenter("dc1")
			.addContactPoint(new InetSocketAddress(CASSANDRA_HOST, CASSANDRA_PORT))
			.withAuthCredentials(this.username, this.password).build()) {
		assertThat(session.getMetadata().getKeyspace("system")).isNotEmpty();
	}
}
 

SessionHolder connected(CqlSession session) {
  Objects.requireNonNull(session);
  if (this.session != null) {
    throw new IllegalStateException();
  }
  return new SessionHolder(connectionQueue, session, refCount);
}
 

synchronized Future<CqlSession> getSession(ContextInternal context) {
  if (closed) {
    return context.failedFuture("Client is closed");
  }
  SessionHolder holder = holders.get(clientName);
  if (holder.session != null) {
    return context.succeededFuture(holder.session);
  }
  return context.executeBlocking(promise -> {
    connect(promise);
  }, holder.connectionQueue);
}
 

/**
 * Use and create keyspace if required
 * @param keyspace name of keyspace
 * @return builder with keyspace
 */
public CasquatchTestDaoBuilder withTestKeyspace(String keyspace) {
    CqlSession session = this.session("system");
    try {
        session.execute(SimpleStatement.builder(String.format("CREATE KEYSPACE IF NOT EXISTS %s WITH replication = { 'class' : 'SimpleStrategy', 'replication_factor' : 1}  AND durable_writes = true;",keyspace)).setExecutionProfileName("ddl").build());
        session.close();
        log.info("Keyspace Created: {}",keyspace);
    }
    catch(Exception e) {
        log.error("Error running withTestKeyspace",e);
    }
    return (CasquatchTestDaoBuilder) this.withBasicSessionKeyspace(keyspace);
}
 

public void connect(final String node, final Integer port, final String dataCenter) {
    CqlSessionBuilder builder = CqlSession.builder();
    builder.addContactPoint(new InetSocketAddress(node, port));
    builder.withLocalDatacenter(dataCenter);

    session = builder.build();
}
 

/**
 * Produce a bound statement while applying object and query options
 * @param buildableQuery buildable query
 * @param obj populated object
 * @param queryOptions query options to apply
 * @param bindToSession override session
 * @return bound statement
 */
protected BoundStatement buildBoundStatement(BuildableQuery buildableQuery, Object obj, QueryOptions queryOptions, CqlSession bindToSession) {
    if(buildableQuery instanceof Select) {
        if(queryOptions!=null) {
            if(queryOptions.getLimit()!=null) {
                buildableQuery = ((Select) buildableQuery).limit(queryOptions.getLimit());
            }
        }
    }
    else if (buildableQuery instanceof RegularInsert) {
        if(queryOptions.getTtl()!=null) {
            buildableQuery = ((RegularInsert) buildableQuery).usingTtl(queryOptions.getTtl());
        }
    }
    SimpleStatement simpleStatement = buildableQuery.build();

    if(log.isTraceEnabled()) log.trace("Preparing Statement {}",simpleStatement.getQuery());
    BoundStatementBuilder boundStatementBuilder = bindToSession.prepare(simpleStatement).boundStatementBuilder();
    if(obj.getClass().equals(this.entityClass)) {
        //noinspection unchecked
        boundStatementBuilder=bindObject(boundStatementBuilder,(E) obj,queryOptions);
    }
    else if (obj.getClass().equals(SolrQueryEntity.class)) {
        //noinspection unchecked
        boundStatementBuilder=bindObject(boundStatementBuilder,(SolrQueryEntity) obj,queryOptions);
    }
    else {
        throw new DriverException(DriverException.CATEGORIES.CASQUATCH_MISSING_GENERATED_CLASS, "Unknown class");
    }
    if(queryOptions!=null) {
        if (queryOptions.getConsistencyLevel() != null) {
            boundStatementBuilder = boundStatementBuilder.setConsistencyLevel(queryOptions.getConsistencyLevel());
        }
        if (queryOptions.getProfile() != null) {
            boundStatementBuilder = boundStatementBuilder.setExecutionProfileName(queryOptions.getProfile());
        }
    }
    return boundStatementBuilder.build();
}
 

private static CqlSession getSession(String[] hosts, Configuration conf, int port) {
    return CqlSession.builder()
            .addContactPoints(
                    Arrays.stream(hosts)
                            .map(s -> new InetSocketAddress(s, port))
                            .collect(Collectors.toList())
            )
            .withLocalDatacenter("datacenter1")
            .build();
}
 

private Map<TokenRange, Long> describeSplits(String keyspace, String table, TokenRange tokenRange, int splitSize, int splitSizeMb, CqlSession session) {
    String query = String.format("SELECT mean_partition_size, partitions_count " +
                    "FROM %s.%s " +
                    "WHERE keyspace_name = ? AND table_name = ? AND range_start = ? AND range_end = ?",
            SchemaConstants.SYSTEM_KEYSPACE_NAME,
            SystemKeyspace.SIZE_ESTIMATES);

    ResultSet resultSet = session.execute(session.prepare(query).bind(keyspace, table, tokenRange.getStart().toString(), tokenRange.getEnd().toString()));

    Row row = resultSet.one();

    long meanPartitionSize = 0;
    long partitionCount = 0;
    int splitCount = 0;

    if (row != null) {
        meanPartitionSize = row.getLong("mean_partition_size");
        partitionCount = row.getLong("partitions_count");

        splitCount = splitSizeMb > 0
                ? (int) (meanPartitionSize * partitionCount / splitSizeMb / 1024 / 1024)
                : (int) (partitionCount / splitSize);
    }

    // If we have no data on this split or the size estimate is 0,
    // return the full split i.e., do not sub-split
    // Assume smallest granularity of partition count available from CASSANDRA-7688
    if (splitCount == 0) {
        Map<TokenRange, Long> wrappedTokenRange = new HashMap<>();
        wrappedTokenRange.put(tokenRange, (long) 128);
        return wrappedTokenRange;
    }

    List<TokenRange> splitRanges = tokenRange.splitEvenly(splitCount);
    Map<TokenRange, Long> rangesWithLength = new HashMap<>();
    for (TokenRange range : splitRanges) {
        rangesWithLength.put(range, partitionCount / splitCount);
    }
    return rangesWithLength;
}
 

private Map<TokenRange, Long> getSubSplits(String keyspace, String cfName, TokenRange range, Configuration conf, CqlSession session) {
    int splitSize = ConfigHelper.getInputSplitSize(conf);
    int splitSizeMb = ConfigHelper.getInputSplitSizeInMb(conf);
    try {
        return describeSplits(keyspace, cfName, range, splitSize, splitSizeMb, session);
    } catch (Exception e) {
        throw new RuntimeException(e);
    }
}