com.datastax.dse.graph.internal.RemoteSourceSimpleGraphStatement 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.dse.graph.internal;

import com.datastax.driver.core.ConsistencyLevel;
import com.datastax.driver.core.SimpleStatement;
import com.datastax.driver.core.Statement;
import com.datastax.driver.dse.graph.GraphOptions;
import com.datastax.driver.dse.graph.GraphProtocol;
import com.datastax.driver.dse.graph.GraphStatement;
import com.datastax.dse.graph.internal.utils.GraphSONUtils;

/**
 * An implementation of {@link GraphStatement} specifically used for remote traversal sources
 * connected via a {@link com.datastax.driver.dse.DseSession}. This implementation will be used by
 * the {@link DseRemoteConnection}, the main reason for this class is that it only supports GraphSON
 * 2 sub-protocol and does not require configuration, which is not the case of the default {@link
 * com.datastax.driver.dse.graph.SimpleGraphStatement}.
 */
public class RemoteSourceSimpleGraphStatement extends GraphStatement {

  private final String query;

  private ConsistencyLevel nativeConsistencyLevel;

  private long defaultTimestamp = Long.MIN_VALUE;

  private volatile int readTimeoutMillis = Integer.MIN_VALUE;

  private String authorizationId;

  public RemoteSourceSimpleGraphStatement(String query) {
    this.query = query;
  }

  @Override
  public GraphStatement setConsistencyLevel(ConsistencyLevel consistencyLevel) {
    nativeConsistencyLevel = consistencyLevel;
    return this;
  }

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

  @Override
  public GraphStatement setDefaultTimestamp(long defaultTimestamp) {
    this.defaultTimestamp = defaultTimestamp;
    return this;
  }

  @Override
  public long getDefaultTimestamp() {
    return defaultTimestamp;
  }

  @Override
  public int getReadTimeoutMillis() {
    return readTimeoutMillis;
  }

  @Override
  public GraphStatement setReadTimeoutMillis(int readTimeoutMillis) {
    this.readTimeoutMillis = readTimeoutMillis;
    return this;
  }

  @Override
  public GraphStatement executingAs(String userOrRole) {
    this.authorizationId = userOrRole;
    return this;
  }

  @Override
  @Deprecated
  public Statement unwrap() {
    return unwrap(GraphProtocol.GRAPHSON_2_0);
  }

  @Override
  @Deprecated
  public Statement unwrap(GraphProtocol graphProtocol) {
    return unwrap(new GraphOptions().setGraphSubProtocol(graphProtocol));
  }

  @Override
  public Statement unwrap(GraphOptions graphOptions) {
    // serializing a `GraphTraversal` or a `Bytecode` will serialize as `Bytecode`.
    Statement statement = new SimpleStatement(query);

    if (getConsistencyLevel() != null) statement.setConsistencyLevel(nativeConsistencyLevel);
    if (getDefaultTimestamp() != Long.MIN_VALUE) statement.setDefaultTimestamp(defaultTimestamp);
    if (getReadTimeoutMillis() != Integer.MIN_VALUE)
      statement.setReadTimeoutMillis(readTimeoutMillis);
    if (this.authorizationId != null) statement.executingAs(authorizationId);

    this.setGraphLanguage(GraphSONUtils.BYTECODE_GRAPHSON_GRAPH_LANGUAGE);
    this.setTransformResultFunction(GraphSONUtils.ROW_TO_GRAPHSON2_TINKERPOP_OBJECTGRAPHNODE);

    return statement;
  }
}