com.datastax.driver.core.Requests Maven / Gradle / Ivy

Go to download
/*
 * Copyright DataStax, Inc.
 *
 * This software can be used solely with DataStax Enterprise. Please consult the license at
 * http://www.datastax.com/terms/datastax-dse-driver-license-terms
 */
package com.datastax.driver.core;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
import io.netty.buffer.ByteBuf;
import java.nio.ByteBuffer;
import java.util.*;

class Requests {

  static final ByteBuffer[] EMPTY_BB_ARRAY = new ByteBuffer[0];

  private Requests() {}

  static class Startup extends Message.Request {
    private static final String CQL_VERSION_OPTION = "CQL_VERSION";
    private static final String CQL_VERSION = "3.0.0";
    private static final String DRIVER_VERSION_OPTION = "DRIVER_VERSION";
    private static final String DRIVER_NAME_OPTION = "DRIVER_NAME";
    private static final String DRIVER_NAME = "DataStax Enterprise Java Driver";

    static final String COMPRESSION_OPTION = "COMPRESSION";
    static final String NO_COMPACT_OPTION = "NO_COMPACT";

    static final Message.Coder coder =
        new Message.Coder() {
          @Override
          public void encode(Startup msg, ByteBuf dest, ProtocolVersion version) {
            CBUtil.writeStringMap(msg.options, dest);
          }

          @Override
          public int encodedSize(Startup msg, ProtocolVersion version) {
            return CBUtil.sizeOfStringMap(msg.options);
          }
        };

    private final Map options;
    private final Map additionalOptions;
    private final ProtocolOptions.Compression compression;
    private final boolean noCompact;

    Startup(
        ProtocolOptions.Compression compression,
        boolean noCompact,
        Map additionalOptions) {
      super(Message.Request.Type.STARTUP);
      this.compression = compression;
      this.noCompact = noCompact;
      this.additionalOptions = additionalOptions;

      ImmutableMap.Builder map = new ImmutableMap.Builder();
      map.put(CQL_VERSION_OPTION, CQL_VERSION);
      if (compression != ProtocolOptions.Compression.NONE)
        map.put(COMPRESSION_OPTION, compression.toString());
      if (noCompact) map.put(NO_COMPACT_OPTION, "true");

      map.put(DRIVER_VERSION_OPTION, Cluster.getDriverVersion());
      map.put(DRIVER_NAME_OPTION, DRIVER_NAME);

      map.putAll(additionalOptions);

      this.options = map.build();
    }

    @Override
    protected Request copyInternal() {
      return new Startup(compression, noCompact, additionalOptions);
    }

    @Override
    public String toString() {
      return "STARTUP " + options;
    }
  }

  // Only for protocol v1
  static class Credentials extends Message.Request {

    static final Message.Coder coder =
        new Message.Coder() {

          @Override
          public void encode(Credentials msg, ByteBuf dest, ProtocolVersion version) {
            assert version == ProtocolVersion.V1;
            CBUtil.writeStringMap(msg.credentials, dest);
          }

          @Override
          public int encodedSize(Credentials msg, ProtocolVersion version) {
            assert version == ProtocolVersion.V1;
            return CBUtil.sizeOfStringMap(msg.credentials);
          }
        };

    private final Map credentials;

    Credentials(Map credentials) {
      super(Message.Request.Type.CREDENTIALS);
      this.credentials = credentials;
    }

    @Override
    protected Request copyInternal() {
      return new Credentials(credentials);
    }
  }

  static class Options extends Message.Request {

    static final Message.Coder coder =
        new Message.Coder() {
          @Override
          public void encode(Options msg, ByteBuf dest, ProtocolVersion version) {}

          @Override
          public int encodedSize(Options msg, ProtocolVersion version) {
            return 0;
          }
        };

    Options() {
      super(Message.Request.Type.OPTIONS);
    }

    @Override
    protected Request copyInternal() {
      return new Options();
    }

    @Override
    public String toString() {
      return "OPTIONS";
    }
  }

  static class Query extends Message.Request {

    static final Message.Coder coder =
        new Message.Coder() {
          @Override
          public void encode(Query msg, ByteBuf dest, ProtocolVersion version) {
            CBUtil.writeLongString(msg.query, dest);
            msg.options.encode(dest, version);
          }

          @Override
          public int encodedSize(Query msg, ProtocolVersion version) {
            return CBUtil.sizeOfLongString(msg.query) + msg.options.encodedSize(version);
          }
        };

    final String query;
    final QueryProtocolOptions options;

    Query(String query) {
      this(query, QueryProtocolOptions.DEFAULT, false);
    }

    Query(String query, QueryProtocolOptions options, boolean tracingRequested) {
      super(Type.QUERY, tracingRequested);
      this.query = query;
      this.options = options;
    }

    @Override
    protected Request copyInternal() {
      return new Query(this.query, options, isTracingRequested());
    }

    @Override
    protected Request copyInternal(ConsistencyLevel newConsistencyLevel) {
      return new Query(this.query, options.copy(newConsistencyLevel), isTracingRequested());
    }

    @Override
    public String toString() {
      return "QUERY " + query + '(' + options + ')';
    }
  }

  static class Execute extends Message.Request {

    static final Message.Coder coder =
        new Message.Coder() {
          @Override
          public void encode(Execute msg, ByteBuf dest, ProtocolVersion version) {
            CBUtil.writeShortBytes(msg.statementId.bytes, dest);
            if (ProtocolFeature.PREPARED_METADATA_CHANGES.isSupportedBy(version))
              CBUtil.writeShortBytes(msg.resultMetadataId.bytes, dest);
            msg.options.encode(dest, version);
          }

          @Override
          public int encodedSize(Execute msg, ProtocolVersion version) {
            int size = CBUtil.sizeOfShortBytes(msg.statementId.bytes);
            if (ProtocolFeature.PREPARED_METADATA_CHANGES.isSupportedBy(version))
              size += CBUtil.sizeOfShortBytes(msg.resultMetadataId.bytes);
            size += msg.options.encodedSize(version);
            return size;
          }
        };

    final MD5Digest statementId;
    final MD5Digest resultMetadataId;
    final QueryProtocolOptions options;

    Execute(
        MD5Digest statementId,
        MD5Digest resultMetadataId,
        QueryProtocolOptions options,
        boolean tracingRequested) {
      super(Message.Request.Type.EXECUTE, tracingRequested);
      this.statementId = statementId;
      this.resultMetadataId = resultMetadataId;
      this.options = options;
    }

    @Override
    protected Request copyInternal() {
      return new Execute(statementId, resultMetadataId, options, isTracingRequested());
    }

    @Override
    protected Request copyInternal(ConsistencyLevel newConsistencyLevel) {
      return new Execute(
          statementId, resultMetadataId, options.copy(newConsistencyLevel), isTracingRequested());
    }

    @Override
    public String toString() {
      if (resultMetadataId != null)
        return "EXECUTE preparedId: "
            + statementId
            + " resultMetadataId: "
            + resultMetadataId
            + " ("
            + options
            + ')';
      else return "EXECUTE preparedId: " + statementId + " (" + options + ')';
    }
  }

  enum QueryFlag {
    // public flags
    VALUES(0),
    SKIP_METADATA(1),
    PAGE_SIZE(2),
    PAGING_STATE(3),
    SERIAL_CONSISTENCY(4),
    DEFAULT_TIMESTAMP(5),
    VALUE_NAMES(6),
    KEYSPACE(7),

    // private flags
    PAGE_SIZE_BYTES(30),
    CONTINUOUS_PAGING(31);

    /** The position of the bit that must be set to one to indicate a flag is set */
    private final int pos;

    QueryFlag(int pos) {
      this.pos = pos;
    }

    static EnumSet deserialize(int flags) {
      EnumSet set = EnumSet.noneOf(QueryFlag.class);
      QueryFlag[] values = QueryFlag.values();
      for (int n = 0; n < values.length; n++) {
        if ((flags & (1 << n)) != 0) set.add(values[n]);
      }
      return set;
    }

    static void serialize(EnumSet flags, ByteBuf dest, ProtocolVersion version) {
      int i = 0;
      boolean isV5 = version.compareTo(ProtocolVersion.V5) >= 0;
      for (QueryFlag flag : flags) {
        // don't include flag if would make value width greater than int and protocol < 5
        if (flag.pos < QueryFlag.KEYSPACE.pos || isV5) {
          // Handle case where keyspace is set but version doesn't support it (DSE_V1).
          if (flag == QueryFlag.KEYSPACE
              && !ProtocolFeature.KEYSPACE_ON_QUERY.isSupportedBy(version)) {
            continue;
          }
          i |= 1 << flag.pos;
        }
      }
      if (isV5) {
        dest.writeInt(i);
      } else {
        dest.writeByte((byte) i);
      }
    }

    static int serializedSize(ProtocolVersion version) {
      return version.compareTo(ProtocolVersion.V5) >= 0 ? 4 : 1;
    }
  }

  static class QueryProtocolOptions {

    static final QueryProtocolOptions DEFAULT =
        new QueryProtocolOptions(
            Message.Request.Type.QUERY,
            ConsistencyLevel.ONE,
            EMPTY_BB_ARRAY,
            Collections.emptyMap(),
            false,
            -1,
            null,
            ConsistencyLevel.SERIAL,
            Long.MIN_VALUE,
            null,
            null);

    private final EnumSet flags = EnumSet.noneOf(QueryFlag.class);
    private final Message.Request.Type requestType;
    final ConsistencyLevel consistency;
    final ByteBuffer[] positionalValues;
    final Map namedValues;
    final boolean skipMetadata;
    final int pageSize;
    final ByteBuffer pagingState;
    final ConsistencyLevel serialConsistency;
    final long defaultTimestamp;
    final String keyspace;
    final ContinuousPagingOptions continuousPagingOptions;

    QueryProtocolOptions(
        Message.Request.Type requestType,
        ConsistencyLevel consistency,
        ByteBuffer[] positionalValues,
        Map namedValues,
        boolean skipMetadata,
        int pageSize,
        ByteBuffer pagingState,
        ConsistencyLevel serialConsistency,
        long defaultTimestamp,
        String keyspace,
        ContinuousPagingOptions continuousPagingOptions) {

      Preconditions.checkArgument(positionalValues.length == 0 || namedValues.isEmpty());

      this.requestType = requestType;
      this.consistency = consistency;
      this.positionalValues = positionalValues;
      this.namedValues = namedValues;
      this.skipMetadata = skipMetadata;
      this.pageSize = pageSize;
      this.pagingState = pagingState;
      this.serialConsistency = serialConsistency;
      this.defaultTimestamp = defaultTimestamp;
      this.keyspace = keyspace;
      this.continuousPagingOptions = continuousPagingOptions;

      // Populate flags
      if (positionalValues.length > 0) {
        flags.add(QueryFlag.VALUES);
      }
      if (!namedValues.isEmpty()) {
        flags.add(QueryFlag.VALUES);
        flags.add(QueryFlag.VALUE_NAMES);
      }
      if (skipMetadata) flags.add(QueryFlag.SKIP_METADATA);
      if (pageSize >= 0) {
        flags.add(QueryFlag.PAGE_SIZE);
        if (continuousPagingOptions != null
            && continuousPagingOptions.getPageUnit() == ContinuousPagingOptions.PageUnit.BYTES)
          flags.add(QueryFlag.PAGE_SIZE_BYTES);
      }
      if (pagingState != null) flags.add(QueryFlag.PAGING_STATE);
      if (serialConsistency != ConsistencyLevel.SERIAL) flags.add(QueryFlag.SERIAL_CONSISTENCY);
      if (defaultTimestamp != Long.MIN_VALUE) flags.add(QueryFlag.DEFAULT_TIMESTAMP);
      if (keyspace != null) flags.add(QueryFlag.KEYSPACE);
      if (continuousPagingOptions != null) flags.add(QueryFlag.CONTINUOUS_PAGING);
    }

    QueryProtocolOptions copy(ConsistencyLevel newConsistencyLevel) {
      return new QueryProtocolOptions(
          requestType,
          newConsistencyLevel,
          positionalValues,
          namedValues,
          skipMetadata,
          pageSize,
          pagingState,
          serialConsistency,
          defaultTimestamp,
          keyspace,
          continuousPagingOptions);
    }

    void encode(ByteBuf dest, ProtocolVersion version) {
      switch (version) {
        case V1:
          // only EXECUTE messages have variables in V1, and their list must be written
          // even if it is empty; and they are never named
          if (requestType == Message.Request.Type.EXECUTE)
            CBUtil.writeValueList(positionalValues, dest);
          CBUtil.writeConsistencyLevel(consistency, dest);
          break;
        case V2:
        case V3:
        case V4:
        case V5:
        case DSE_V1:
        case DSE_V2:
          CBUtil.writeConsistencyLevel(consistency, dest);
          QueryFlag.serialize(flags, dest, version);
          if (flags.contains(QueryFlag.VALUES)) {
            if (flags.contains(QueryFlag.VALUE_NAMES)) {
              assert version.compareTo(ProtocolVersion.V3) >= 0;
              CBUtil.writeNamedValueList(namedValues, dest);
            } else {
              CBUtil.writeValueList(positionalValues, dest);
            }
          }
          if (flags.contains(QueryFlag.PAGE_SIZE)) dest.writeInt(pageSize);
          if (flags.contains(QueryFlag.PAGING_STATE)) CBUtil.writeValue(pagingState, dest);
          if (flags.contains(QueryFlag.SERIAL_CONSISTENCY))
            CBUtil.writeConsistencyLevel(serialConsistency, dest);
          if (version.compareTo(ProtocolVersion.V3) >= 0
              && flags.contains(QueryFlag.DEFAULT_TIMESTAMP)) dest.writeLong(defaultTimestamp);
          if (ProtocolFeature.KEYSPACE_ON_QUERY.isSupportedBy(version)
              && flags.contains(QueryFlag.KEYSPACE)) CBUtil.writeString(keyspace, dest);
          if (version.compareTo(ProtocolVersion.DSE_V1) >= 0
              && flags.contains(QueryFlag.CONTINUOUS_PAGING)) {
            dest.writeInt(continuousPagingOptions.getMaxPages());
            dest.writeInt(continuousPagingOptions.getMaxPagesPerSecond());
            if (version.compareTo(ProtocolVersion.DSE_V2) >= 0)
              dest.writeInt(continuousPagingOptions.getMaxEnqueuedPages());
          }
          break;
        default:
          throw version.unsupported();
      }
    }

    int encodedSize(ProtocolVersion version) {
      switch (version) {
        case V1:
          // only EXECUTE messages have variables in V1, and their list must be written
          // even if it is empty; and they are never named
          return (requestType == Message.Request.Type.EXECUTE
                  ? CBUtil.sizeOfValueList(positionalValues)
                  : 0)
              + CBUtil.sizeOfConsistencyLevel(consistency);
        case V2:
        case V3:
        case V4:
        case V5:
        case DSE_V1:
        case DSE_V2:
          int size = 0;
          size += CBUtil.sizeOfConsistencyLevel(consistency);
          size += QueryFlag.serializedSize(version); // flags
          if (flags.contains(QueryFlag.VALUES)) {
            if (flags.contains(QueryFlag.VALUE_NAMES)) {
              assert version.compareTo(ProtocolVersion.V3) >= 0;
              size += CBUtil.sizeOfNamedValueList(namedValues);
            } else {
              size += CBUtil.sizeOfValueList(positionalValues);
            }
          }
          if (flags.contains(QueryFlag.PAGE_SIZE)) size += 4;
          if (flags.contains(QueryFlag.PAGING_STATE)) size += CBUtil.sizeOfValue(pagingState);
          if (flags.contains(QueryFlag.SERIAL_CONSISTENCY))
            size += CBUtil.sizeOfConsistencyLevel(serialConsistency);
          if (version.compareTo(ProtocolVersion.V3) >= 0
              && flags.contains(QueryFlag.DEFAULT_TIMESTAMP)) size += 8;
          if (ProtocolFeature.KEYSPACE_ON_QUERY.isSupportedBy(version)
              && flags.contains(QueryFlag.KEYSPACE)) size += CBUtil.sizeOfString(keyspace);
          if (version.compareTo(ProtocolVersion.DSE_V1) >= 0
              && flags.contains(QueryFlag.CONTINUOUS_PAGING))
            size += version.compareTo(ProtocolVersion.DSE_V2) >= 0 ? 12 : 8;
          return size;
        default:
          throw version.unsupported();
      }
    }

    @Override
    public String toString() {
      return String.format(
          "[cl=%s, positionalVals=%s, namedVals=%s, skip=%b, psize=%d, state=%s, serialCl=%s]",
          consistency,
          Arrays.toString(positionalValues),
          namedValues,
          skipMetadata,
          pageSize,
          pagingState,
          serialConsistency);
    }
  }

  static class Batch extends Message.Request {

    static final Message.Coder coder =
        new Message.Coder() {
          @Override
          public void encode(Batch msg, ByteBuf dest, ProtocolVersion version) {
            int queries = msg.queryOrIdList.size();
            assert queries <= 0xFFFF;

            dest.writeByte(fromType(msg.type));
            dest.writeShort(queries);

            for (int i = 0; i < queries; i++) {
              Object q = msg.queryOrIdList.get(i);
              dest.writeByte((byte) (q instanceof String ? 0 : 1));
              if (q instanceof String) CBUtil.writeLongString((String) q, dest);
              else CBUtil.writeShortBytes(((MD5Digest) q).bytes, dest);

              CBUtil.writeValueList(msg.values[i], dest);
            }

            msg.options.encode(dest, version);
          }

          @Override
          public int encodedSize(Batch msg, ProtocolVersion version) {
            int size = 3; // type + nb queries
            for (int i = 0; i < msg.queryOrIdList.size(); i++) {
              Object q = msg.queryOrIdList.get(i);
              size +=
                  1
                      + (q instanceof String
                          ? CBUtil.sizeOfLongString((String) q)
                          : CBUtil.sizeOfShortBytes(((MD5Digest) q).bytes));

              size += CBUtil.sizeOfValueList(msg.values[i]);
            }
            size += msg.options.encodedSize(version);
            return size;
          }

          private byte fromType(BatchStatement.Type type) {
            switch (type) {
              case LOGGED:
                return 0;
              case UNLOGGED:
                return 1;
              case COUNTER:
                return 2;
              default:
                throw new AssertionError();
            }
          }
        };

    final BatchStatement.Type type;
    final List