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);
}
}
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);
}
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());
}
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();
}
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();
}
Aggregations