Fix inferred request routing copy in statement builders

Author: dkropachev

core/src/main/java/com/datastax/oss/driver/api/core/cql/StatementBuilder.java

@@ -23,6 +23,7 @@
 import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
 import com.datastax.oss.driver.api.core.metadata.Node;
 import com.datastax.oss.driver.api.core.metadata.token.Token;
+import com.datastax.oss.driver.internal.core.cql.RequestRoutingTypeAccessor;
 import com.datastax.oss.driver.internal.core.util.RoutingKey;
 import com.datastax.oss.protocol.internal.util.collection.NullAllowingImmutableMap;
 import edu.umd.cs.findbugs.annotations.NonNull;
@@ -89,7 +90,15 @@ protected StatementBuilder(StatementT template) {
     this.timeout = template.getTimeout();
     this.node = template.getNode();
     this.nowInSeconds = template.getNowInSeconds();
-    this.requestRoutingType = template.getRequestRoutingType();
+    this.requestRoutingType = getConfiguredRequestRoutingType(template);
+  }
+
+  @Nullable
+  private RequestRoutingType getConfiguredRequestRoutingType(StatementT template) {
+    if (template instanceof RequestRoutingTypeAccessor) {
+      return ((RequestRoutingTypeAccessor) template).getConfiguredRequestRoutingType();
+    }
+    return template.getRequestRoutingType();
   }
 
   /** @see Statement#setExecutionProfileName(String) */

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

@@ -49,7 +49,7 @@
 import org.slf4j.LoggerFactory;
 
 @Immutable
-public class DefaultBatchStatement implements BatchStatement {
+public class DefaultBatchStatement implements BatchStatement, RequestRoutingTypeAccessor {
   private static final Logger LOG = LoggerFactory.getLogger(DefaultBatchStatement.class);
 
   private final BatchType batchType;
@@ -873,6 +873,12 @@ public RequestRoutingType getRequestRoutingType() {
     return cachedStatementsRequestRoutingType;
   }
 
+  @Nullable
+  @Override
+  public RequestRoutingType getConfiguredRequestRoutingType() {
+    return requestRoutingType;
+  }
+
   @NonNull
   @Override
   public BatchStatement setRequestRoutingType(RequestRoutingType requestRoutingType) {

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

@@ -47,7 +47,7 @@
 import net.jcip.annotations.Immutable;
 
 @Immutable
-public class DefaultBoundStatement implements BoundStatement {
+public class DefaultBoundStatement implements BoundStatement, RequestRoutingTypeAccessor {
 
   private final PreparedStatement preparedStatement;
   private final ColumnDefinitions variableDefinitions;
@@ -816,6 +816,12 @@ public RequestRoutingType getRequestRoutingType() {
     return preparedStatement.getRequestRoutingType();
   }
 
+  @Nullable
+  @Override
+  public RequestRoutingType getConfiguredRequestRoutingType() {
+    return requestRoutingType;
+  }
+
   @NonNull
   @Override
   public BoundStatement setRequestRoutingType(@Nullable RequestRoutingType requestRoutingType) {

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

@@ -42,7 +42,7 @@
 import net.jcip.annotations.Immutable;
 
 @Immutable
-public class DefaultSimpleStatement implements SimpleStatement {
+public class DefaultSimpleStatement implements SimpleStatement, RequestRoutingTypeAccessor {
 
   private final String query;
   private final List<Object> positionalValues;
@@ -786,6 +786,12 @@ public RequestRoutingType getRequestRoutingType() {
     return null;
   }
 
+  @Nullable
+  @Override
+  public RequestRoutingType getConfiguredRequestRoutingType() {
+    return requestRoutingType;
+  }
+
   @NonNull
   @Override
   public SimpleStatement setRequestRoutingType(@Nullable RequestRoutingType requestRoutingType) {

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

@@ -0,0 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.oss.driver.internal.core.cql;
+
+import com.datastax.oss.driver.api.core.RequestRoutingType;
+import edu.umd.cs.findbugs.annotations.Nullable;
+
+/** Internal hook to distinguish explicitly configured routing type from inferred routing type. */
+public interface RequestRoutingTypeAccessor {
+  @Nullable
+  RequestRoutingType getConfiguredRequestRoutingType();
+}

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

@@ -21,9 +21,15 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.datastax.oss.driver.api.core.DefaultConsistencyLevel;
+import com.datastax.oss.driver.api.core.DefaultProtocolVersion;
+import com.datastax.oss.driver.api.core.RequestRoutingType;
+import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry;
+import com.datastax.oss.driver.internal.core.cql.DefaultBoundStatement;
 import com.datastax.oss.driver.shaded.guava.common.base.Charsets;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import java.nio.ByteBuffer;
+import java.util.Collections;
 import org.junit.Test;
 
 public class StatementBuilderTest {
@@ -103,4 +109,83 @@ public void should_match_set_routing_key_vararg() {
     builderStmt = builder.setRoutingKey(buff2, buff1).build();
     assertThat(expectedStmt.getRoutingKey()).isNotEqualTo(builderStmt.getRoutingKey());
   }
+
+  @Test
+  public void should_not_copy_inferred_simple_routing_type_as_explicit() {
+    SimpleStatement serialStatement =
+        SimpleStatement.builder("select * from test.foo")
+            .setConsistencyLevel(DefaultConsistencyLevel.LOCAL_SERIAL)
+            .build();
+
+    assertThat(serialStatement.getRequestRoutingType()).isEqualTo(RequestRoutingType.LWT);
+
+    SimpleStatement regularStatement =
+        SimpleStatement.builder(serialStatement)
+            .setConsistencyLevel(DefaultConsistencyLevel.ONE)
+            .build();
+
+    assertThat(regularStatement.getRequestRoutingType()).isNull();
+  }
+
+  @Test
+  public void should_not_copy_inferred_bound_routing_type_as_explicit() {
+    BoundStatement serialStatement = newRegularBoundStatement(DefaultConsistencyLevel.LOCAL_SERIAL);
+
+    assertThat(serialStatement.getRequestRoutingType()).isEqualTo(RequestRoutingType.LWT);
+
+    BoundStatement regularStatement =
+        new BoundStatementBuilder(serialStatement)
+            .setConsistencyLevel(DefaultConsistencyLevel.ONE)
+            .build();
+
+    assertThat(regularStatement.getRequestRoutingType()).isEqualTo(RequestRoutingType.REGULAR);
+  }
+
+  @Test
+  public void should_not_copy_inferred_batch_routing_type_as_explicit() {
+    BatchStatement serialStatement =
+        BatchStatement.builder(BatchType.LOGGED)
+            .setConsistencyLevel(DefaultConsistencyLevel.LOCAL_SERIAL)
+            .build();
+
+    assertThat(serialStatement.getRequestRoutingType()).isEqualTo(RequestRoutingType.LWT);
+
+    BatchStatement regularStatement =
+        BatchStatement.builder(serialStatement)
+            .setConsistencyLevel(DefaultConsistencyLevel.ONE)
+            .build();
+
+    assertThat(regularStatement.getRequestRoutingType()).isEqualTo(RequestRoutingType.REGULAR);
+  }
+
+  private BoundStatement newRegularBoundStatement(DefaultConsistencyLevel consistencyLevel) {
+    PreparedStatement preparedStatement = mock(PreparedStatement.class);
+    ColumnDefinitions variableDefinitions = mock(ColumnDefinitions.class);
+    when(preparedStatement.isLWT()).thenReturn(false);
+    when(preparedStatement.getRequestRoutingType()).thenReturn(RequestRoutingType.REGULAR);
+    when(preparedStatement.getVariableDefinitions()).thenReturn(variableDefinitions);
+    return new DefaultBoundStatement(
+        preparedStatement,
+        variableDefinitions,
+        new ByteBuffer[0],
+        null,
+        null,
+        null,
+        null,
+        null,
+        Collections.emptyMap(),
+        null,
+        false,
+        Statement.NO_DEFAULT_TIMESTAMP,
+        null,
+        Integer.MIN_VALUE,
+        consistencyLevel,
+        null,
+        null,
+        CodecRegistry.DEFAULT,
+        DefaultProtocolVersion.DEFAULT,
+        null,
+        Statement.NO_NOW_IN_SECONDS,
+        null);
+  }
 }