package com.datastax.oss.driver.internal.core.util;
import com.datastax.oss.driver.api.core.CqlIdentifier;
import com.datastax.oss.driver.api.core.ProtocolVersion;
import com.datastax.oss.driver.api.core.context.DriverContext;
import com.datastax.oss.driver.api.core.cql.BatchableStatement;
import com.datastax.oss.driver.api.core.cql.BoundStatement;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
import com.datastax.oss.driver.api.core.cql.Statement;
import com.datastax.oss.driver.api.core.session.Request;
import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry;
import com.datastax.oss.driver.internal.core.data.ValuesHelper;
import com.datastax.oss.protocol.internal.FrameCodec;
import com.datastax.oss.protocol.internal.PrimitiveSizes;
import com.datastax.oss.protocol.internal.request.query.QueryOptions;
import com.datastax.oss.protocol.internal.request.query.Values;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
public class Sizes {
public static int minimumRequestSize(Request request) {
int size = FrameCodec.V3_ENCODED_HEADER_SIZE;
if (!request.getCustomPayload().isEmpty()) {
size += PrimitiveSizes.sizeOfBytesMap(request.getCustomPayload());
}
return size;
}
public static int minimumStatementSize(Statement statement, DriverContext context) {
int size = minimumRequestSize(statement);
size += QueryOptions.queryFlagsSize(context.getProtocolVersion().getCode());
size += PrimitiveSizes.SHORT;
size += PrimitiveSizes.SHORT;
return size;
}
public static int sizeOfSimpleStatementValues(
SimpleStatement simpleStatement,
ProtocolVersion protocolVersion,
CodecRegistry codecRegistry) {
int size = 0;
if (!simpleStatement.getPositionalValues().isEmpty()) {
List<ByteBuffer> positionalValues =
new ArrayList<>(simpleStatement.getPositionalValues().size());
for (Object value : simpleStatement.getPositionalValues()) {
positionalValues.add(
ValuesHelper.encodeToDefaultCqlMapping(value, codecRegistry, protocolVersion));
}
size += Values.sizeOfPositionalValues(positionalValues);
} else if (!simpleStatement.getNamedValues().isEmpty()) {
Map<String, ByteBuffer> namedValues = new HashMap<>(simpleStatement.getNamedValues().size());
for (Map.Entry<CqlIdentifier, Object> value : simpleStatement.getNamedValues().entrySet()) {
namedValues.put(
value.getKey().asInternal(),
ValuesHelper.encodeToDefaultCqlMapping(
value.getValue(), codecRegistry, protocolVersion));
}
size += Values.sizeOfNamedValues(namedValues);
}
return size;
}
public static int sizeOfBoundStatementValues(BoundStatement boundStatement) {
return Values.sizeOfPositionalValues(boundStatement.getValues());
}
public static Integer sizeOfInnerBatchStatementInBytes(
BatchableStatement statement, ProtocolVersion protocolVersion, CodecRegistry codecRegistry) {
int size = 0;
size +=
PrimitiveSizes
.BYTE;
if (statement instanceof SimpleStatement) {
size += PrimitiveSizes.sizeOfLongString(((SimpleStatement) statement).getQuery());
size +=
sizeOfSimpleStatementValues(
((SimpleStatement) statement), protocolVersion, codecRegistry);
} else if (statement instanceof BoundStatement) {
size +=
PrimitiveSizes.sizeOfShortBytes(
((BoundStatement) statement).getPreparedStatement().getId());
size += sizeOfBoundStatementValues(((BoundStatement) statement));
}
return size;
}
}