/* * Copyright (C) 2012-2015 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 com.datastax.driver.core.exceptions.UnsupportedFeatureException; import com.datastax.driver.core.utils.CassandraVersion; import com.google.common.collect.ImmutableMap; import org.apache.log4j.Logger; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; import java.nio.ByteBuffer; import java.util.HashMap; import java.util.Map; import static com.datastax.driver.core.ProtocolVersion.V3; import static com.datastax.driver.core.querybuilder.QueryBuilder.eq; import static com.datastax.driver.core.querybuilder.QueryBuilder.select; import static org.apache.log4j.Level.TRACE; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Fail.fail; @CassandraVersion("2.2.0") @CCMConfig(jvmArgs = "-Dcassandra.custom_query_handler_class=org.apache.cassandra.cql3.CustomPayloadMirroringQueryHandler") public class CustomPayloadTest extends CCMTestsSupport { private Map<String, ByteBuffer> payload1; private Map<String, ByteBuffer> payload2; @BeforeMethod(groups = {"short", "unit"}) public void initPayloads() { payload1 = ImmutableMap.of( "k1", ByteBuffer.wrap(new byte[]{1, 2, 3}), "k2", ByteBuffer.wrap(new byte[]{4, 5, 6}) ); payload2 = ImmutableMap.of( "k2", ByteBuffer.wrap(new byte[]{1, 2}), "k3", ByteBuffer.wrap(new byte[]{3, 4}) ); } // execute @Test(groups = "short") public void should_echo_custom_payload_when_executing_statement() throws Exception { Statement statement = new SimpleStatement("SELECT c2 FROM t1 where c1 = ?", 1); statement.setOutgoingPayload(payload1); ResultSet rows = session().execute(statement); Map<String, ByteBuffer> actual = rows.getExecutionInfo().getIncomingPayload(); assertThat(actual).isEqualTo(payload1); } @Test(groups = "short") public void should_echo_custom_payload_when_executing_batch_statement() throws Exception { Statement statement = new BatchStatement().add(new SimpleStatement("INSERT INTO t1 (c1, c2) values (1, 'foo')")); statement.setOutgoingPayload(payload1); ResultSet rows = session().execute(statement); Map<String, ByteBuffer> actual = rows.getExecutionInfo().getIncomingPayload(); assertThat(actual).isEqualTo(payload1); } @Test(groups = "short") public void should_echo_custom_payload_when_building_statement() throws Exception { Statement statement = select("c2").from("t1").where(eq("c1", 1)).setOutgoingPayload(payload1); ResultSet rows = session().execute(statement); Map<String, ByteBuffer> actual = rows.getExecutionInfo().getIncomingPayload(); assertThat(actual).isEqualTo(payload1); } // prepare /** * Ensures that an incoming payload is propagated from prepared to bound statements. */ @Test(groups = "short") public void should_propagate_incoming_payload_to_bound_statement() throws Exception { RegularStatement statement = new SimpleStatement("SELECT c2 as col1 FROM t1 where c1 = ?"); statement.setOutgoingPayload(payload1); PreparedStatement ps = session().prepare(statement); // Prepared statement should inherit outgoing payload assertThat(ps.getOutgoingPayload()).isEqualTo(payload1); // Prepared statement should receive incoming payload assertThat(ps.getIncomingPayload()).isEqualTo(payload1); ps.setOutgoingPayload(null); // unset outgoing payload // bound statement should inherit from prepared statement's incoming payload BoundStatement bs = ps.bind(1); ResultSet rows = session().execute(bs); Map<String, ByteBuffer> actual = rows.getExecutionInfo().getIncomingPayload(); assertThat(actual).isEqualTo(payload1); bs = ps.bind(); bs.setInt(0, 1); rows = session().execute(bs); actual = rows.getExecutionInfo().getIncomingPayload(); assertThat(actual).isEqualTo(payload1); } /** * Ensures that an incoming payload is overridden by an explicitly set outgoing payload * when propagated to bound statements. */ @Test(groups = "short") public void should_override_incoming_payload_when_outgoing_payload_explicitly_set_on_preparing_statement() throws Exception { RegularStatement statement = new SimpleStatement("SELECT c2 as col2 FROM t1 where c1 = ?"); statement.setOutgoingPayload(payload1); PreparedStatement ps = session().prepare(statement); // Prepared statement should inherit outgoing payload assertThat(ps.getOutgoingPayload()).isEqualTo(payload1); // Prepared statement should receive incoming payload assertThat(ps.getIncomingPayload()).isEqualTo(payload1); ps.setOutgoingPayload(payload2); // override outgoing payload // bound statement should inherit from prepared statement's outgoing payload BoundStatement bs = ps.bind(1); ResultSet rows = session().execute(bs); Map<String, ByteBuffer> actual = rows.getExecutionInfo().getIncomingPayload(); assertThat(actual).isEqualTo(payload2); bs = ps.bind(); bs.setInt(0, 1); rows = session().execute(bs); actual = rows.getExecutionInfo().getIncomingPayload(); assertThat(actual).isEqualTo(payload2); } /** * Ensures that payloads can still be set individually on bound statements * if the prepared statement does not have a default payload. */ @Test(groups = "short") public void should_not_set_any_payload_on_bound_statement() throws Exception { RegularStatement statement = new SimpleStatement("SELECT c2 as col3 FROM t1 where c1 = ?"); PreparedStatement ps = session().prepare(statement); assertThat(ps.getOutgoingPayload()).isNull(); assertThat(ps.getIncomingPayload()).isNull(); // bound statement should not have outgoing payload BoundStatement bs = ps.bind(1); assertThat(bs.getOutgoingPayload()).isNull(); // explicitly set a payload for this boudn statement only bs.setOutgoingPayload(payload1); ResultSet rows = session().execute(bs); Map<String, ByteBuffer> actual = rows.getExecutionInfo().getIncomingPayload(); assertThat(actual).isEqualTo(payload1); // a second bound statement should not have any payload bs = ps.bind(); assertThat(bs.getOutgoingPayload()).isNull(); bs.setInt(0, 1); rows = session().execute(bs); actual = rows.getExecutionInfo().getIncomingPayload(); assertThat(actual).isNull(); } // pagination /** * Ensures that a custom payload is propagated throughout pages. */ @Test(groups = "short") public void should_echo_custom_payload_when_paginating() throws Exception { session().execute("INSERT INTO t1 (c1, c2) VALUES (1, 'a')"); session().execute("INSERT INTO t1 (c1, c2) VALUES (1, 'b')"); Statement statement = new SimpleStatement("SELECT c2 FROM t1 where c1 = 1"); statement.setFetchSize(1); statement.setOutgoingPayload(payload1); ResultSet rows = session().execute(statement); rows.all(); assertThat(rows.getAllExecutionInfo()).extracting("incomingPayload").containsOnly(payload1); } // TODO retries, spec execs // edge cases @Test(groups = "short") public void should_encode_null_values() throws Exception { Map<String, ByteBuffer> payload = new HashMap<String, ByteBuffer>(); payload.put("k1", Statement.NULL_PAYLOAD_VALUE); Statement statement = new SimpleStatement("SELECT c2 FROM t1 where c1 = ?", 1); statement.setOutgoingPayload(payload); ResultSet rows = session().execute(statement); Map<String, ByteBuffer> actual = rows.getExecutionInfo().getIncomingPayload(); assertThat(actual).isEqualTo(payload); } @Test(groups = "unit", expectedExceptions = NullPointerException.class) public void should_throw_npe_when_null_key_on_regular_statement() throws Exception { Map<String, ByteBuffer> payload = new HashMap<String, ByteBuffer>(); payload.put(null, ByteBuffer.wrap(new byte[]{1})); new SimpleStatement("SELECT c2 FROM t1 where c1 = ?", 1).setOutgoingPayload(payload); } @Test(groups = "unit", expectedExceptions = NullPointerException.class) public void should_throw_npe_when_null_value_on_regular_statement() throws Exception { Map<String, ByteBuffer> payload = new HashMap<String, ByteBuffer>(); payload.put("k1", null); new SimpleStatement("SELECT c2 FROM t1 where c1 = ?", 1).setOutgoingPayload(payload); } @Test(groups = "short", expectedExceptions = NullPointerException.class) public void should_throw_npe_when_null_key_on_prepared_statement() throws Exception { Map<String, ByteBuffer> payload = new HashMap<String, ByteBuffer>(); payload.put(null, ByteBuffer.wrap(new byte[]{1})); session().prepare(new SimpleStatement("SELECT c2 FROM t1 where c1 = 1")).setOutgoingPayload(payload); } @Test(groups = "short", expectedExceptions = NullPointerException.class) public void should_throw_npe_when_null_value_on_prepared_statement() throws Exception { Map<String, ByteBuffer> payload = new HashMap<String, ByteBuffer>(); payload.put("k1", null); session().prepare(new SimpleStatement("SELECT c2 FROM t1 where c1 = 2")).setOutgoingPayload(payload); } @Test(groups = "short") public void should_throw_ufe_when_protocol_version_lesser_than_4() throws Exception { try { Cluster v3cluster = register(Cluster.builder() .addContactPoints(getContactPoints()) .withPort(ccm().getBinaryPort()) .withProtocolVersion(V3) .build()) .init(); Session v3session = v3cluster.connect(); Statement statement = new SimpleStatement("SELECT c2 FROM t1 where c1 = ?", 1); statement.setOutgoingPayload(payload1); v3session.execute(statement); fail("Should not send custom payloads with protocol V3"); } catch (UnsupportedFeatureException e) { assertThat(e.getMessage()).isEqualTo( "Unsupported feature with the native protocol V3 (which is currently in use): Custom payloads are only supported since native protocol V4"); } } // log messages /** * Ensures that when debugging custom payloads, the driver will print appropriate log messages. */ @Test(groups = "short") public void should_print_log_message_when_level_trace() throws Exception { Logger logger = Logger.getLogger(Message.logger.getName()); MemoryAppender appender = new MemoryAppender(); try { logger.setLevel(TRACE); logger.addAppender(appender); Statement statement = new SimpleStatement("SELECT c2 FROM t1 where c1 = ?", 1); statement.setOutgoingPayload(payload1); session().execute(statement); String logs = appender.waitAndGet(10000); assertThat(logs) .contains("Sending payload: {k1:0x010203, k2:0x040506} (20 bytes total)") .contains("Received payload: {k1:0x010203, k2:0x040506} (20 bytes total)"); } finally { logger.setLevel(null); logger.removeAppender(appender); } } @Override public void onTestContextInitialized() { execute("CREATE TABLE t1 (c1 int, c2 text, PRIMARY KEY (c1, c2))"); } }