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

/*
 * Copyright DataStax, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.datastax.driver.core;

import static com.datastax.driver.core.ProtocolVersion.V4;

import com.datastax.driver.core.policies.RetryPolicy;
import com.google.common.collect.ImmutableMap;
import java.nio.ByteBuffer;
import java.util.List;
import java.util.Map;

public class DefaultPreparedStatement implements PreparedStatement {

  final PreparedId preparedId;

  final String query;
  final String queryKeyspace;
  final Map incomingPayload;
  final Cluster cluster;

  volatile ByteBuffer routingKey;

  volatile ConsistencyLevel consistency;
  volatile ConsistencyLevel serialConsistency;
  volatile boolean traceQuery;
  volatile RetryPolicy retryPolicy;
  volatile ImmutableMap outgoingPayload;
  volatile Boolean idempotent;

  private DefaultPreparedStatement(
      PreparedId id,
      String query,
      String queryKeyspace,
      Map incomingPayload,
      Cluster cluster) {
    this.preparedId = id;
    this.query = query;
    this.queryKeyspace = queryKeyspace;
    this.incomingPayload = incomingPayload;
    this.cluster = cluster;
  }

  static DefaultPreparedStatement fromMessage(
      Responses.Result.Prepared msg, Cluster cluster, String query, String queryKeyspace) {
    assert msg.metadata.columns != null;

    ColumnDefinitions defs = msg.metadata.columns;

    ProtocolVersion protocolVersion =
        cluster.getConfiguration().getProtocolOptions().getProtocolVersion();
    PreparedId.PreparedMetadata boundValuesMetadata =
        new PreparedId.PreparedMetadata(msg.statementId, defs);
    PreparedId.PreparedMetadata resultSetMetadata =
        new PreparedId.PreparedMetadata(msg.resultMetadataId, msg.resultMetadata.columns);

    int[] pkIndices = null;
    if (defs.size() > 0) {
      pkIndices =
          (protocolVersion.compareTo(V4) >= 0)
              ? msg.metadata.pkIndices
              : computePkIndices(cluster.getMetadata(), defs);
    }

    PreparedId preparedId =
        new PreparedId(boundValuesMetadata, resultSetMetadata, pkIndices, protocolVersion);
    return new DefaultPreparedStatement(
        preparedId, query, queryKeyspace, msg.getCustomPayload(), cluster);
  }

  private static int[] computePkIndices(Metadata clusterMetadata, ColumnDefinitions boundColumns) {
    List partitionKeyColumns = null;
    int[] pkIndexes = null;
    KeyspaceMetadata km = clusterMetadata.getKeyspace(Metadata.quote(boundColumns.getKeyspace(0)));
    if (km != null) {
      TableMetadata tm = km.getTable(Metadata.quote(boundColumns.getTable(0)));
      if (tm != null) {
        partitionKeyColumns = tm.getPartitionKey();
        pkIndexes = new int[partitionKeyColumns.size()];
        for (int i = 0; i < pkIndexes.length; ++i) pkIndexes[i] = -1;
      }
    }

    // Note: we rely on the fact CQL queries cannot span multiple tables. If that change, we'll have
    // to get smarter.
    for (int i = 0; i < boundColumns.size(); i++)
      maybeGetIndex(boundColumns.getName(i), i, partitionKeyColumns, pkIndexes);

    return allSet(pkIndexes) ? pkIndexes : null;
  }

  private static void maybeGetIndex(
      String name, int j, List pkColumns, int[] pkIndexes) {
    if (pkColumns == null) return;

    for (int i = 0; i < pkColumns.size(); ++i) {
      if (name.equals(pkColumns.get(i).getName())) {
        // We may have the same column prepared multiple times, but only pick the first value
        pkIndexes[i] = j;
        return;
      }
    }
  }

  private static boolean allSet(int[] pkColumns) {
    if (pkColumns == null) return false;

    for (int i = 0; i < pkColumns.length; ++i) if (pkColumns[i] < 0) return false;

    return true;
  }

  @Override
  public ColumnDefinitions getVariables() {
    return preparedId.boundValuesMetadata.variables;
  }

  @Override
  public BoundStatement bind(Object... values) {
    BoundStatement bs = new BoundStatement(this);
    return bs.bind(values);
  }

  @Override
  public BoundStatement bind() {
    return new BoundStatement(this);
  }

  @Override
  public PreparedStatement setRoutingKey(ByteBuffer routingKey) {
    this.routingKey = routingKey;
    return this;
  }

  @Override
  public PreparedStatement setRoutingKey(ByteBuffer... routingKeyComponents) {
    this.routingKey = SimpleStatement.compose(routingKeyComponents);
    return this;
  }

  @Override
  public ByteBuffer getRoutingKey() {
    return routingKey;
  }

  @Override
  public PreparedStatement setConsistencyLevel(ConsistencyLevel consistency) {
    this.consistency = consistency;
    return this;
  }

  @Override
  public ConsistencyLevel getConsistencyLevel() {
    return consistency;
  }

  @Override
  public PreparedStatement setSerialConsistencyLevel(ConsistencyLevel serialConsistency) {
    if (!serialConsistency.isSerial()) throw new IllegalArgumentException();
    this.serialConsistency = serialConsistency;
    return this;
  }

  @Override
  public ConsistencyLevel getSerialConsistencyLevel() {
    return serialConsistency;
  }

  @Override
  public String getQueryString() {
    return query;
  }

  @Override
  public String getQueryKeyspace() {
    return queryKeyspace;
  }

  @Override
  public PreparedStatement enableTracing() {
    this.traceQuery = true;
    return this;
  }

  @Override
  public PreparedStatement disableTracing() {
    this.traceQuery = false;
    return this;
  }

  @Override
  public boolean isTracing() {
    return traceQuery;
  }

  @Override
  public PreparedStatement setRetryPolicy(RetryPolicy policy) {
    this.retryPolicy = policy;
    return this;
  }

  @Override
  public RetryPolicy getRetryPolicy() {
    return retryPolicy;
  }

  @Override
  public PreparedId getPreparedId() {
    return preparedId;
  }

  @Override
  public Map getIncomingPayload() {
    return incomingPayload;
  }

  @Override
  public Map getOutgoingPayload() {
    return outgoingPayload;
  }

  @Override
  public PreparedStatement setOutgoingPayload(Map payload) {
    this.outgoingPayload = payload == null ? null : ImmutableMap.copyOf(payload);
    return this;
  }

  @Override
  public CodecRegistry getCodecRegistry() {
    return cluster.getConfiguration().getCodecRegistry();
  }

  /** {@inheritDoc} */
  @Override
  public PreparedStatement setIdempotent(Boolean idempotent) {
    this.idempotent = idempotent;
    return this;
  }

  /** {@inheritDoc} */
  @Override
  public Boolean isIdempotent() {
    return this.idempotent;
  }
}