Search in sources :

Example 1 with ScriptGraphStatement

use of com.datastax.dse.driver.api.core.graph.ScriptGraphStatement in project java-driver by datastax.

the class GraphConversions method createContinuousMessageFromGraphStatement.

public static Message createContinuousMessageFromGraphStatement(GraphStatement<?> statement, GraphProtocol subProtocol, DriverExecutionProfile config, InternalDriverContext context, GraphBinaryModule graphBinaryModule) {
    final List<ByteBuffer> encodedQueryParams;
    if (!(statement instanceof ScriptGraphStatement) || ((ScriptGraphStatement) statement).getQueryParams().isEmpty()) {
        encodedQueryParams = Collections.emptyList();
    } else {
        try {
            Map<String, Object> queryParams = ((ScriptGraphStatement) statement).getQueryParams();
            if (subProtocol.isGraphBinary()) {
                Buffer graphBinaryParams = graphBinaryModule.serialize(queryParams);
                encodedQueryParams = Collections.singletonList(graphBinaryParams.nioBuffer());
                graphBinaryParams.release();
            } else {
                encodedQueryParams = Collections.singletonList(GraphSONUtils.serializeToByteBuffer(queryParams, subProtocol));
            }
        } catch (IOException e) {
            throw new UncheckedIOException("Couldn't serialize parameters for GraphStatement: " + statement, e);
        }
    }
    int consistencyLevel = DefaultConsistencyLevel.valueOf(config.getString(DefaultDriverOption.REQUEST_CONSISTENCY)).getProtocolCode();
    long timestamp = statement.getTimestamp();
    if (timestamp == Statement.NO_DEFAULT_TIMESTAMP) {
        timestamp = context.getTimestampGenerator().next();
    }
    int pageSize = config.getInt(DseDriverOption.GRAPH_CONTINUOUS_PAGING_PAGE_SIZE);
    int maxPages = config.getInt(DseDriverOption.GRAPH_CONTINUOUS_PAGING_MAX_PAGES);
    int maxPagesPerSecond = config.getInt(DseDriverOption.GRAPH_CONTINUOUS_PAGING_MAX_PAGES_PER_SECOND);
    int maxEnqueuedPages = config.getInt(DseDriverOption.GRAPH_CONTINUOUS_PAGING_MAX_ENQUEUED_PAGES);
    ContinuousPagingOptions options = new ContinuousPagingOptions(maxPages, maxPagesPerSecond, maxEnqueuedPages);
    DseQueryOptions queryOptions = new DseQueryOptions(consistencyLevel, encodedQueryParams, // ignored by the DSE Graph server
    Collections.emptyMap(), // also ignored
    true, pageSize, null, // also ignored
    ProtocolConstants.ConsistencyLevel.LOCAL_SERIAL, timestamp, // also ignored
    null, // graph CP does not support sizeInBytes
    false, options);
    if (statement instanceof ScriptGraphStatement) {
        return new Query(((ScriptGraphStatement) statement).getScript(), queryOptions);
    } else {
        return new RawBytesQuery(getQueryBytes(statement, subProtocol), queryOptions);
    }
}
Also used : ByteBuffer(java.nio.ByteBuffer) Buffer(org.apache.tinkerpop.gremlin.structure.io.Buffer) RawBytesQuery(com.datastax.dse.protocol.internal.request.RawBytesQuery) RawBytesQuery(com.datastax.dse.protocol.internal.request.RawBytesQuery) Query(com.datastax.oss.protocol.internal.request.Query) ContinuousPagingOptions(com.datastax.dse.protocol.internal.request.query.ContinuousPagingOptions) UncheckedIOException(java.io.UncheckedIOException) IOException(java.io.IOException) UncheckedIOException(java.io.UncheckedIOException) ByteBuffer(java.nio.ByteBuffer) ScriptGraphStatement(com.datastax.dse.driver.api.core.graph.ScriptGraphStatement) DseQueryOptions(com.datastax.dse.protocol.internal.request.query.DseQueryOptions)

Example 2 with ScriptGraphStatement

use of com.datastax.dse.driver.api.core.graph.ScriptGraphStatement in project java-driver by datastax.

the class DefaultReactiveGraphResultSetIT method should_retrieve_all_rows.

@Test
@DataProvider(value = { "1", "10", "100", "999", "1000", "1001", "2000" }, format = "%m [page size %p[0]]")
public void should_retrieve_all_rows(int pageSize) {
    DriverExecutionProfile profile = sessionRule.session().getContext().getConfig().getDefaultProfile().withInt(DseDriverOption.GRAPH_CONTINUOUS_PAGING_PAGE_SIZE, pageSize);
    ScriptGraphStatement statement = ScriptGraphStatement.builder("g.V()").setExecutionProfile(profile).build();
    ReactiveGraphResultSet rs = sessionRule.session().executeReactive(statement);
    List<ReactiveGraphNode> results = Flowable.fromPublisher(rs).toList().blockingGet();
    assertThat(results.size()).isEqualTo(1000);
    Set<ExecutionInfo> expectedExecInfos = new LinkedHashSet<>();
    for (ReactiveGraphNode row : results) {
        assertThat(row.getExecutionInfo()).isNotNull();
        assertThat(row.isVertex()).isTrue();
        expectedExecInfos.add(row.getExecutionInfo());
    }
    List<ExecutionInfo> execInfos = Flowable.<ExecutionInfo>fromPublisher(rs.getExecutionInfos()).toList().blockingGet();
    // DSE may send an empty page as it can't always know if it's done paging or not yet.
    // See: CASSANDRA-8871. In this case, this page's execution info appears in
    // rs.getExecutionInfos(), but is not present in expectedExecInfos since the page did not
    // contain any rows.
    assertThat(execInfos).containsAll(expectedExecInfos);
}
Also used : LinkedHashSet(java.util.LinkedHashSet) DriverExecutionProfile(com.datastax.oss.driver.api.core.config.DriverExecutionProfile) ScriptGraphStatement(com.datastax.dse.driver.api.core.graph.ScriptGraphStatement) ExecutionInfo(com.datastax.oss.driver.api.core.cql.ExecutionInfo) DataProvider(com.tngtech.java.junit.dataprovider.DataProvider) Test(org.junit.Test)

Example 3 with ScriptGraphStatement

use of com.datastax.dse.driver.api.core.graph.ScriptGraphStatement in project java-driver by datastax.

the class GraphRequestHandlerTest method should_honor_statement_consistency_level.

@Test
public void should_honor_statement_consistency_level() {
    // initialization
    GraphRequestHandlerTestHarness harness = GraphRequestHandlerTestHarness.builder().build();
    ScriptGraphStatement graphStatement = ScriptGraphStatement.builder("mockScript").setConsistencyLevel(DefaultConsistencyLevel.THREE).build();
    GraphBinaryModule module = createGraphBinaryModule(harness.getContext());
    // when
    DriverExecutionProfile executionProfile = Conversions.resolveExecutionProfile(graphStatement, harness.getContext());
    Message m = GraphConversions.createMessageFromGraphStatement(graphStatement, GRAPH_BINARY_1_0, executionProfile, harness.getContext(), module);
    // checks
    assertThat(m).isInstanceOf(Query.class);
    Query q = ((Query) m);
    assertThat(q.options.consistency).isEqualTo(DefaultConsistencyLevel.THREE.getProtocolCode());
}
Also used : Message(com.datastax.oss.protocol.internal.Message) Query(com.datastax.oss.protocol.internal.request.Query) RawBytesQuery(com.datastax.dse.protocol.internal.request.RawBytesQuery) DriverExecutionProfile(com.datastax.oss.driver.api.core.config.DriverExecutionProfile) ScriptGraphStatement(com.datastax.dse.driver.api.core.graph.ScriptGraphStatement) GraphTestUtils.createGraphBinaryModule(com.datastax.dse.driver.internal.core.graph.GraphTestUtils.createGraphBinaryModule) GraphBinaryModule(com.datastax.dse.driver.internal.core.graph.binary.GraphBinaryModule) Test(org.junit.Test)

Example 4 with ScriptGraphStatement

use of com.datastax.dse.driver.api.core.graph.ScriptGraphStatement in project java-driver by datastax.

the class GraphRequestHandlerTest method should_create_query_message_from_script_statement.

@Test
@UseDataProvider(location = DseTestDataProviders.class, value = "supportedGraphProtocols")
public void should_create_query_message_from_script_statement(GraphProtocol graphProtocol) throws IOException {
    // initialization
    GraphRequestHandlerTestHarness harness = GraphRequestHandlerTestHarness.builder().build();
    ScriptGraphStatement graphStatement = ScriptGraphStatement.newInstance("mockQuery").setQueryParam("p1", 1L).setQueryParam("p2", Uuids.random());
    GraphBinaryModule module = createGraphBinaryModule(harness.getContext());
    // when
    DriverExecutionProfile executionProfile = Conversions.resolveExecutionProfile(graphStatement, harness.getContext());
    Message m = GraphConversions.createMessageFromGraphStatement(graphStatement, graphProtocol, executionProfile, harness.getContext(), module);
    // checks
    assertThat(m).isInstanceOf(Query.class);
    Query q = ((Query) m);
    assertThat(q.query).isEqualTo("mockQuery");
    assertThat(q.options.positionalValues).containsExactly(serialize(graphStatement.getQueryParams(), graphProtocol, module));
    assertThat(q.options.namedValues).isEmpty();
}
Also used : Message(com.datastax.oss.protocol.internal.Message) Query(com.datastax.oss.protocol.internal.request.Query) RawBytesQuery(com.datastax.dse.protocol.internal.request.RawBytesQuery) DriverExecutionProfile(com.datastax.oss.driver.api.core.config.DriverExecutionProfile) ScriptGraphStatement(com.datastax.dse.driver.api.core.graph.ScriptGraphStatement) GraphTestUtils.createGraphBinaryModule(com.datastax.dse.driver.internal.core.graph.GraphTestUtils.createGraphBinaryModule) GraphBinaryModule(com.datastax.dse.driver.internal.core.graph.binary.GraphBinaryModule) Test(org.junit.Test) UseDataProvider(com.tngtech.java.junit.dataprovider.UseDataProvider)

Example 5 with ScriptGraphStatement

use of com.datastax.dse.driver.api.core.graph.ScriptGraphStatement in project java-driver by datastax.

the class GraphConversions method createCustomPayload.

public static Map<String, ByteBuffer> createCustomPayload(GraphStatement<?> statement, GraphProtocol subProtocol, DriverExecutionProfile config, InternalDriverContext context, GraphBinaryModule graphBinaryModule) {
    ProtocolVersion protocolVersion = context.getProtocolVersion();
    NullAllowingImmutableMap.Builder<String, ByteBuffer> payload = NullAllowingImmutableMap.builder();
    Map<String, ByteBuffer> statementOptions = statement.getCustomPayload();
    payload.putAll(statementOptions);
    final String graphLanguage;
    // Don't override anything that's already provided at the statement level
    if (!statementOptions.containsKey(GRAPH_LANG_OPTION_KEY)) {
        graphLanguage = statement instanceof ScriptGraphStatement ? LANGUAGE_GROOVY : LANGUAGE_BYTECODE;
        payload.put(GRAPH_LANG_OPTION_KEY, TypeCodecs.TEXT.encode(graphLanguage, protocolVersion));
    } else {
        graphLanguage = TypeCodecs.TEXT.decode(statementOptions.get(GRAPH_LANG_OPTION_KEY), protocolVersion);
        Preconditions.checkNotNull(graphLanguage, "A null value was set for the graph-language custom payload key.");
    }
    if (!isSystemQuery(statement, config)) {
        if (!statementOptions.containsKey(GRAPH_NAME_OPTION_KEY)) {
            String graphName = statement.getGraphName();
            if (graphName == null) {
                graphName = config.getString(DseDriverOption.GRAPH_NAME, null);
            }
            if (graphName != null) {
                payload.put(GRAPH_NAME_OPTION_KEY, TypeCodecs.TEXT.encode(graphName, protocolVersion));
            }
        }
        if (!statementOptions.containsKey(GRAPH_SOURCE_OPTION_KEY)) {
            String traversalSource = statement.getTraversalSource();
            if (traversalSource == null) {
                traversalSource = config.getString(DseDriverOption.GRAPH_TRAVERSAL_SOURCE, null);
            }
            if (traversalSource != null) {
                payload.put(GRAPH_SOURCE_OPTION_KEY, TypeCodecs.TEXT.encode(traversalSource, protocolVersion));
            }
        }
    }
    // the payload allows null entry values so doing a get directly here and checking for null
    final ByteBuffer payloadInitialProtocol = statementOptions.get(GRAPH_RESULTS_OPTION_KEY);
    if (payloadInitialProtocol == null) {
        Preconditions.checkNotNull(subProtocol);
        payload.put(GRAPH_RESULTS_OPTION_KEY, TypeCodecs.TEXT.encode(subProtocol.toInternalCode(), protocolVersion));
    } else {
        subProtocol = GraphProtocol.fromString(TypeCodecs.TEXT.decode(payloadInitialProtocol, protocolVersion));
    }
    if (subProtocol.isGraphBinary() && graphLanguage.equals(LANGUAGE_BYTECODE)) {
        Object bytecodeQuery = bytecodeToSerialize(statement);
        try {
            Buffer bytecodeByteBuf = graphBinaryModule.serialize(bytecodeQuery);
            payload.put(GRAPH_BINARY_QUERY_OPTION_KEY, bytecodeByteBuf.nioBuffer());
            bytecodeByteBuf.release();
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
    }
    if (!statementOptions.containsKey(GRAPH_READ_CONSISTENCY_LEVEL_OPTION_KEY)) {
        ConsistencyLevel readCl = statement.getReadConsistencyLevel();
        String readClString = readCl != null ? readCl.name() : config.getString(DseDriverOption.GRAPH_READ_CONSISTENCY_LEVEL, null);
        if (readClString != null) {
            payload.put(GRAPH_READ_CONSISTENCY_LEVEL_OPTION_KEY, TypeCodecs.TEXT.encode(readClString, protocolVersion));
        }
    }
    if (!statementOptions.containsKey(GRAPH_WRITE_CONSISTENCY_LEVEL_OPTION_KEY)) {
        ConsistencyLevel writeCl = statement.getWriteConsistencyLevel();
        String writeClString = writeCl != null ? writeCl.name() : config.getString(DseDriverOption.GRAPH_WRITE_CONSISTENCY_LEVEL, null);
        if (writeClString != null) {
            payload.put(GRAPH_WRITE_CONSISTENCY_LEVEL_OPTION_KEY, TypeCodecs.TEXT.encode(writeClString, protocolVersion));
        }
    }
    if (!statementOptions.containsKey(GRAPH_TIMEOUT_OPTION_KEY)) {
        Duration timeout = statement.getTimeout();
        if (timeout == null) {
            timeout = config.getDuration(DseDriverOption.GRAPH_TIMEOUT, Duration.ZERO);
        }
        if (timeout != null && !timeout.isZero()) {
            payload.put(GRAPH_TIMEOUT_OPTION_KEY, TypeCodecs.BIGINT.encode(timeout.toMillis(), protocolVersion));
        }
    }
    return payload.build();
}
Also used : ByteBuffer(java.nio.ByteBuffer) Buffer(org.apache.tinkerpop.gremlin.structure.io.Buffer) UncheckedIOException(java.io.UncheckedIOException) Duration(java.time.Duration) IOException(java.io.IOException) UncheckedIOException(java.io.UncheckedIOException) ProtocolVersion(com.datastax.oss.driver.api.core.ProtocolVersion) ByteBuffer(java.nio.ByteBuffer) DefaultConsistencyLevel(com.datastax.oss.driver.api.core.DefaultConsistencyLevel) ConsistencyLevel(com.datastax.oss.driver.api.core.ConsistencyLevel) NullAllowingImmutableMap(com.datastax.oss.protocol.internal.util.collection.NullAllowingImmutableMap) ScriptGraphStatement(com.datastax.dse.driver.api.core.graph.ScriptGraphStatement)

Aggregations

ScriptGraphStatement (com.datastax.dse.driver.api.core.graph.ScriptGraphStatement)6 RawBytesQuery (com.datastax.dse.protocol.internal.request.RawBytesQuery)4 Query (com.datastax.oss.protocol.internal.request.Query)4 DriverExecutionProfile (com.datastax.oss.driver.api.core.config.DriverExecutionProfile)3 IOException (java.io.IOException)3 UncheckedIOException (java.io.UncheckedIOException)3 ByteBuffer (java.nio.ByteBuffer)3 Buffer (org.apache.tinkerpop.gremlin.structure.io.Buffer)3 Test (org.junit.Test)3 GraphTestUtils.createGraphBinaryModule (com.datastax.dse.driver.internal.core.graph.GraphTestUtils.createGraphBinaryModule)2 GraphBinaryModule (com.datastax.dse.driver.internal.core.graph.binary.GraphBinaryModule)2 DseQueryOptions (com.datastax.dse.protocol.internal.request.query.DseQueryOptions)2 ConsistencyLevel (com.datastax.oss.driver.api.core.ConsistencyLevel)2 DefaultConsistencyLevel (com.datastax.oss.driver.api.core.DefaultConsistencyLevel)2 Message (com.datastax.oss.protocol.internal.Message)2 ContinuousPagingOptions (com.datastax.dse.protocol.internal.request.query.ContinuousPagingOptions)1 ProtocolVersion (com.datastax.oss.driver.api.core.ProtocolVersion)1 ExecutionInfo (com.datastax.oss.driver.api.core.cql.ExecutionInfo)1 NullAllowingImmutableMap (com.datastax.oss.protocol.internal.util.collection.NullAllowingImmutableMap)1 DataProvider (com.tngtech.java.junit.dataprovider.DataProvider)1