package com.datastax.oss.driver.api.core.cql;
import com.datastax.oss.driver.api.core.CqlIdentifier;
import com.datastax.oss.driver.internal.core.cql.DefaultBatchStatement;
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList;
import com.datastax.oss.driver.shaded.guava.common.collect.Iterables;
import edu.umd.cs.findbugs.annotations.NonNull;
import edu.umd.cs.findbugs.annotations.Nullable;
import java.util.Arrays;
import net.jcip.annotations.NotThreadSafe;
@NotThreadSafe
public class BatchStatementBuilder extends StatementBuilder<BatchStatementBuilder, BatchStatement> {
@NonNull private BatchType batchType;
@Nullable private CqlIdentifier keyspace;
@NonNull private ImmutableList.Builder<BatchableStatement<?>> statementsBuilder;
private int statementsCount;
public BatchStatementBuilder(@NonNull BatchType batchType) {
this.batchType = batchType;
this.statementsBuilder = ImmutableList.builder();
}
public BatchStatementBuilder(@NonNull BatchStatement template) {
super(template);
this.batchType = template.getBatchType();
this.statementsBuilder = ImmutableList.<BatchableStatement<?>>builder().addAll(template);
this.statementsCount = template.size();
}
@NonNull
public BatchStatementBuilder setKeyspace(@NonNull CqlIdentifier keyspace) {
this.keyspace = keyspace;
return this;
}
@NonNull
public BatchStatementBuilder setKeyspace(@NonNull String keyspaceName) {
return setKeyspace(CqlIdentifier.fromCql(keyspaceName));
}
@NonNull
public BatchStatementBuilder addStatement(@NonNull BatchableStatement<?> statement) {
if (statementsCount >= 0xFFFF) {
throw new IllegalStateException(
"Batch statement cannot contain more than " + 0xFFFF + " statements.");
}
statementsCount += 1;
statementsBuilder.add(statement);
return this;
}
@NonNull
public BatchStatementBuilder addStatements(@NonNull Iterable<BatchableStatement<?>> statements) {
int delta = Iterables.size(statements);
if (statementsCount + delta > 0xFFFF) {
throw new IllegalStateException(
"Batch statement cannot contain more than " + 0xFFFF + " statements.");
}
statementsCount += delta;
statementsBuilder.addAll(statements);
return this;
}
@NonNull
public BatchStatementBuilder addStatements(@NonNull BatchableStatement<?>... statements) {
return addStatements(Arrays.asList(statements));
}
@NonNull
public BatchStatementBuilder clearStatements() {
statementsBuilder = ImmutableList.builder();
statementsCount = 0;
return this;
}
@Override
@NonNull
public BatchStatement build() {
return new DefaultBatchStatement(
batchType,
statementsBuilder.build(),
executionProfileName,
executionProfile,
keyspace,
routingKeyspace,
routingKey,
routingToken,
buildCustomPayload(),
idempotent,
tracing,
timestamp,
pagingState,
pageSize,
consistencyLevel,
serialConsistencyLevel,
timeout,
node,
nowInSeconds);
}
public int getStatementsCount() {
return this.statementsCount;
}
}