Don't return empty routing key when partition key is unbound

DefaultBoundStatement#getRoutingKey has logic to infer the
routing key when no one has explicitly called setRoutingKey
or otherwise set the routing key on the statement.
It however doesn't check for cases where nothing has been
bound yet on the statement.
This causes more problems if the user decides to get a
BoundStatementBuilder from the PreparedStatement, set some
fields on it, and then copy it by constructing new
BoundStatementBuilder objects with the BoundStatement as a
parameter, since the empty ByteBuffer gets copied to all
bound statements, resulting in all requests being targeted
to the same Cassandra node in a token-aware load balancing
policy.

Author: akhaku

core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultBoundStatement.java

@@ -358,7 +358,7 @@ public ByteBuffer getRoutingKey() {
       if (indices.isEmpty()) {
         return null;
       } else if (indices.size() == 1) {
-        return getBytesUnsafe(indices.get(0));
+        return isSet(0) ? getBytesUnsafe(indices.get(0)) : null;
       } else {
         ByteBuffer[] components = new ByteBuffer[indices.size()];
         for (int i = 0; i < components.length; i++) {

core/src/test/java/com/datastax/oss/driver/api/core/cql/StatementBuilderTest.java

@@ -19,10 +19,30 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.datastax.oss.driver.api.core.DefaultProtocolVersion;
+import com.datastax.oss.driver.api.core.context.DriverContext;
+import com.datastax.oss.driver.api.core.type.DataType;
+import com.datastax.oss.driver.api.core.type.DataTypes;
+import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry;
+import com.datastax.oss.driver.internal.core.cql.DefaultColumnDefinition;
+import com.datastax.oss.driver.internal.core.cql.DefaultColumnDefinitions;
+import com.datastax.oss.driver.internal.core.cql.DefaultPreparedStatement;
 import com.datastax.oss.driver.shaded.guava.common.base.Charsets;
+import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList;
+import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
+import com.datastax.oss.protocol.internal.response.result.ColumnSpec;
+import com.datastax.oss.protocol.internal.response.result.RawType;
+import com.datastax.oss.protocol.internal.util.Bytes;
 import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Map;
+import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
 
+@RunWith(MockitoJUnitRunner.class)
 public class StatementBuilderTest {
 
   private static class MockSimpleStatementBuilder
@@ -46,6 +66,14 @@ public SimpleStatement build() {
     }
   }
 
+  @Mock private DriverContext context;
+
+  @Before
+  public void setup() {
+    when(context.getCodecRegistry()).thenReturn(CodecRegistry.DEFAULT);
+    when(context.getProtocolVersion()).thenReturn(DefaultProtocolVersion.V4);
+  }
+
   @Test
   public void should_handle_set_tracing_without_args() {
 
@@ -99,4 +127,72 @@ public void should_match_set_routing_key_vararg() {
     builderStmt = builder.setRoutingKey(buff2, buff1).build();
     assertThat(expectedStmt.getRoutingKey()).isNotEqualTo(builderStmt.getRoutingKey());
   }
+
+  @Test
+  public void should_correctly_set_routing_key_on_boundstatement() {
+    PreparedStatement preparedStatement = mockPreparedStatement("UPDATE foo SET v=? WHERE k=?",
+            ImmutableMap.of("v", DataTypes.INT, "k", DataTypes.INT));
+    BoundStatement boundStatement = preparedStatement.boundStatementBuilder().build();
+    // we should not have a routing key set
+    assertThat(boundStatement.getRoutingKey()).isNull();
+    boundStatement = boundStatement.set(0, 1, Integer.class).set(1, 2, Integer.class);
+    // we should have a non-empty routing key now
+    assertThat(boundStatement.getRoutingKey()).isNotNull();
+    assertThat(boundStatement.getRoutingKey().hasRemaining()).isTrue();
+  }
+
+  @Test
+  public void should_correctly_set_routing_key_on_boundstatementbuilder() {
+    PreparedStatement preparedStatement = mockPreparedStatement("UPDATE foo SET v=? WHERE k=?",
+            ImmutableMap.of("v", DataTypes.INT, "k", DataTypes.INT));
+    BoundStatement boundStatement = preparedStatement.boundStatementBuilder().build();
+    BoundStatement copy = new BoundStatementBuilder(boundStatement).build();
+    // we should not have a routing key set
+    assertThat(copy.getRoutingKey()).isNull();
+    copy = copy.set(0, 1, Integer.class).set(1, 2, Integer.class);
+    // we should have a non-empty routing key now
+    assertThat(copy.getRoutingKey()).isNotNull();
+    assertThat(copy.getRoutingKey().hasRemaining()).isTrue();
+  }
+
+  // copied from RequestLogFormatterTest, we should move somewhere to share b/w tests
+  private PreparedStatement mockPreparedStatement(String query, Map<String, DataType> variables) {
+    ImmutableList.Builder<ColumnDefinition> definitions = ImmutableList.builder();
+    int i = 0;
+    for (Map.Entry<String, DataType> entry : variables.entrySet()) {
+      definitions.add(
+              new DefaultColumnDefinition(
+                      new ColumnSpec(
+                              "test",
+                              "foo",
+                              entry.getKey(),
+                              i,
+                              RawType.PRIMITIVES.get(entry.getValue().getProtocolCode())),
+                      context));
+    }
+    return new DefaultPreparedStatement(
+            Bytes.fromHexString("0x"),
+            query,
+            DefaultColumnDefinitions.valueOf(definitions.build()),
+            Collections.singletonList(1), // note that this line is different from the one in RequestLogFormatterTest
+            null,
+            null,
+            null,
+            Collections.emptyMap(),
+            null,
+            null,
+            null,
+            null,
+            null,
+            Collections.emptyMap(),
+            null,
+            null,
+            null,
+            Integer.MIN_VALUE,
+            null,
+            null,
+            false,
+            context.getCodecRegistry(),
+            context.getProtocolVersion());
+  }
 }