Fix prepared routing inference override

Author: dkropachev

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

@@ -813,6 +813,9 @@ public RequestRoutingType getRequestRoutingType() {
     if (serialConsistencyLevel != null && serialConsistencyLevel.isSerial()) {
       return RequestRoutingType.LWT;
     }
+    if (preparedStatement instanceof RequestRoutingTypeAccessor) {
+      return ((RequestRoutingTypeAccessor) preparedStatement).getConfiguredRequestRoutingType();
+    }
     return preparedStatement.getRequestRoutingType();
   }
 

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

@@ -58,7 +58,7 @@
 import org.slf4j.LoggerFactory;
 
 @ThreadSafe
-public class DefaultPreparedStatement implements PreparedStatement {
+public class DefaultPreparedStatement implements PreparedStatement, RequestRoutingTypeAccessor {
   private static final Logger LOGGER = LoggerFactory.getLogger(DefaultPreparedStatement.class);
   private static final Splitter SPACE_SPLITTER = Splitter.onPattern("\\s+");
   private static final Splitter COMMA_SPLITTER = Splitter.onPattern(",");
@@ -211,6 +211,12 @@ public RequestRoutingType getRequestRoutingType() {
     return null;
   }
 
+  @Nullable
+  @Override
+  public RequestRoutingType getConfiguredRequestRoutingType() {
+    return requestRoutingType;
+  }
+
   @Override
   public void setResultMetadata(
       @NonNull ByteBuffer newResultMetadataId, @NonNull ColumnDefinitions newResultSetDefinitions) {

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

@@ -20,7 +20,7 @@
 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. */
+/** Internal hook to distinguish stored routing type from consistency-inferred routing type. */
 public interface RequestRoutingTypeAccessor {
   @Nullable
   RequestRoutingType getConfiguredRequestRoutingType();

core/src/test/java/com/datastax/oss/driver/internal/core/cql/DefaultPreparedStatementTest.java

@@ -0,0 +1,93 @@
+/*
+ * 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 static org.assertj.core.api.Assertions.assertThat;
+
+import com.datastax.oss.driver.api.core.ConsistencyLevel;
+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.cql.BoundStatement;
+import com.datastax.oss.driver.api.core.cql.ColumnDefinitions;
+import com.datastax.oss.driver.api.core.type.codec.registry.CodecRegistry;
+import com.datastax.oss.protocol.internal.util.Bytes;
+import java.util.Collections;
+import org.junit.Test;
+
+public class DefaultPreparedStatementTest {
+
+  @Test
+  public void should_not_keep_inferred_routing_type_after_bound_consistency_override() {
+    DefaultPreparedStatement preparedStatement =
+        newPreparedStatement(DefaultConsistencyLevel.LOCAL_SERIAL, null, null);
+
+    assertThat(preparedStatement.getRequestRoutingType()).isEqualTo(RequestRoutingType.LWT);
+
+    BoundStatement boundStatement =
+        preparedStatement.bind().setConsistencyLevel(DefaultConsistencyLevel.ONE);
+
+    assertThat(boundStatement.getConsistencyLevel()).isEqualTo(DefaultConsistencyLevel.ONE);
+    assertThat(boundStatement.getRequestRoutingType()).isNull();
+  }
+
+  @Test
+  public void should_keep_detected_lwt_routing_type_after_bound_consistency_override() {
+    DefaultPreparedStatement preparedStatement =
+        newPreparedStatement(DefaultConsistencyLevel.LOCAL_SERIAL, null, RequestRoutingType.LWT);
+
+    BoundStatement boundStatement =
+        preparedStatement.bind().setConsistencyLevel(DefaultConsistencyLevel.ONE);
+
+    assertThat(boundStatement.getRequestRoutingType()).isEqualTo(RequestRoutingType.LWT);
+  }
+
+  private DefaultPreparedStatement newPreparedStatement(
+      ConsistencyLevel consistencyLevel,
+      ConsistencyLevel serialConsistencyLevel,
+      RequestRoutingType requestRoutingType) {
+    ColumnDefinitions variableDefinitions =
+        DefaultColumnDefinitions.valueOf(Collections.emptyList());
+    return new DefaultPreparedStatement(
+        Bytes.fromHexString("0x"),
+        "SELECT * FROM test.foo WHERE pk = ?",
+        variableDefinitions,
+        Collections.emptyList(),
+        null,
+        null,
+        null,
+        null,
+        Collections.emptyMap(),
+        null,
+        null,
+        null,
+        null,
+        null,
+        Collections.emptyMap(),
+        null,
+        null,
+        null,
+        Integer.MIN_VALUE,
+        consistencyLevel,
+        serialConsistencyLevel,
+        false,
+        CodecRegistry.DEFAULT,
+        DefaultProtocolVersion.DEFAULT,
+        requestRoutingType);
+  }
+}