From 926bdefe748b933a30196eddf0893302022d07d1 Mon Sep 17 00:00:00 2001 From: Ammar Khaku Date: Sun, 20 Nov 2022 19:14:07 -0800 Subject: [PATCH] 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. --- .../core/cql/DefaultBoundStatement.java | 3 ++- .../driver/core/cql/PreparedStatementIT.java | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultBoundStatement.java b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultBoundStatement.java index 4e0b8af3fac..6b404541224 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultBoundStatement.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/cql/DefaultBoundStatement.java @@ -358,7 +358,8 @@ public ByteBuffer getRoutingKey() { if (indices.isEmpty()) { return null; } else if (indices.size() == 1) { - return getBytesUnsafe(indices.get(0)); + int index = indices.get(0); + return isSet(index) ? getBytesUnsafe(index) : null; } else { ByteBuffer[] components = new ByteBuffer[indices.size()]; for (int i = 0; i < components.length; i++) { diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/cql/PreparedStatementIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/cql/PreparedStatementIT.java index 1b07edb53af..ddab089f871 100644 --- a/integration-tests/src/test/java/com/datastax/oss/driver/core/cql/PreparedStatementIT.java +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/cql/PreparedStatementIT.java @@ -475,6 +475,25 @@ private void should_infer_routing_information_when_partition_key_is_bound(String assertThat(tokenFactory.hash(boundStatement.getRoutingKey())).isEqualTo(expectedToken); } + @Test + public void should_return_null_routing_information_when_single_partition_key_is_unbound() { + should_return_null_routing_information_when_single_partition_key_is_unbound( + "SELECT a FROM prepared_statement_test WHERE a = ?"); + should_return_null_routing_information_when_single_partition_key_is_unbound( + "INSERT INTO prepared_statement_test (a) VALUES (?)"); + should_return_null_routing_information_when_single_partition_key_is_unbound( + "UPDATE prepared_statement_test SET b = 1 WHERE a = ?"); + should_return_null_routing_information_when_single_partition_key_is_unbound( + "DELETE FROM prepared_statement_test WHERE a = ?"); + } + + private void should_return_null_routing_information_when_single_partition_key_is_unbound( + String queryString) { + CqlSession session = sessionRule.session(); + BoundStatement boundStatement = session.prepare(queryString).bind(); + assertThat(boundStatement.getRoutingKey()).isNull(); + } + private static Iterable firstPageOf(CompletionStage stage) { return CompletableFutures.getUninterruptibly(stage).currentPage(); }