CNDB-15312: Fix multi-column partition key CQL generation in DataRange

Author: driftx

src/java/org/apache/cassandra/db/DataRange.java

@@ -288,9 +288,7 @@ public String toCQLString(TableMetadata metadata)
              * key are the same. If that is the case, we want to print the query as an equality on the partition key
              * rather than a token range, as if it was a partition query, for better readability.
              */
-            builder.append(ColumnMetadata.toCQLString(metadata.partitionKeyColumns()));
-            builder.append(" = ");
-            appendKeyString(builder, metadata.partitionKeyType, ((DecoratedKey) startKey()).getKey());
+            builder.append(((DecoratedKey) startKey()).toCQLString(metadata));
             needAnd = true;
         }
         else

src/java/org/apache/cassandra/db/DecoratedKey.java

@@ -19,11 +19,16 @@
 
 import java.nio.ByteBuffer;
 import java.util.Comparator;
+import java.util.List;
+import java.util.StringJoiner;
 import java.util.function.BiFunction;
 
+import org.apache.cassandra.db.marshal.CompositeType;
 import org.apache.cassandra.dht.IPartitioner;
 import org.apache.cassandra.dht.Token;
 import org.apache.cassandra.dht.Token.KeyBound;
+import org.apache.cassandra.schema.ColumnMetadata;
+import org.apache.cassandra.schema.TableMetadata;
 import org.apache.cassandra.utils.ByteBufferUtil;
 import org.apache.cassandra.utils.bytecomparable.ByteComparable;
 import org.apache.cassandra.utils.bytecomparable.ByteSource;
@@ -161,6 +166,32 @@ public String toString()
         return "DecoratedKey(" + getToken() + ", " + keystring + ")";
     }
 
+    /**
+     * Generate CQL representation of this partition key for the given table.
+     * For single-column keys: "k = 0"
+     * For multi-column keys: "k1 = 1 AND k2 = 2"
+     */
+    public String toCQLString(TableMetadata metadata)
+    {
+        List<ColumnMetadata> columns = metadata.partitionKeyColumns();
+
+        if (columns.size() == 1)
+            return toCQLString(columns.get(0), getKey());
+
+        ByteBuffer[] values = ((CompositeType) metadata.partitionKeyType).split(getKey());
+        StringJoiner joiner = new StringJoiner(" AND ");
+
+        for (int i = 0; i < columns.size(); i++)
+            joiner.add(toCQLString(columns.get(i), values[i]));
+
+        return joiner.toString();
+    }
+
+    private static String toCQLString(ColumnMetadata metadata, ByteBuffer key)
+    {
+        return String.format("%s = %s", metadata.name.toCQLString(), metadata.type.getString(key));
+    }
+
     public Token getToken()
     {
         return token;

src/java/org/apache/cassandra/db/SinglePartitionReadCommand.java

@@ -1151,8 +1151,7 @@ protected void appendCQLWhereClause(CqlBuilder builder)
     {
         builder.append(" WHERE ");
 
-        builder.append(ColumnMetadata.toCQLString(metadata().partitionKeyColumns())).append(" = ");
-        DataRange.appendKeyString(builder, metadata().partitionKeyType, partitionKey().getKey());
+        builder.append(partitionKey().toCQLString(metadata()));
 
         // We put the row filter first because the clustering index filter can end by "ORDER BY"
         if (!rowFilter().isEmpty())

test/unit/org/apache/cassandra/db/PartitionRangeReadCommandCQLTest.java

@@ -87,6 +87,12 @@ public void testToCQLString()
         assertToCQLString("SELECT * FROM %s WHERE v = 0 AND k = 0 AND c1 = 1 AND c2 < 2", "SELECT * FROM %s WHERE v = 0 AND k = 0 AND c1 = 1 AND c2 < 2 ALLOW FILTERING");
         assertToCQLString("SELECT * FROM %s WHERE v = 0 AND k = 0 AND c1 = 1 AND c2 >= 2", "SELECT * FROM %s WHERE v = 0 AND k = 0 AND c1 = 1 AND c2 >= 2 ALLOW FILTERING");
         assertToCQLString("SELECT * FROM %s WHERE v = 0 AND k = 0 AND c1 = 1 AND c2 <= 2", "SELECT * FROM %s WHERE v = 0 AND k = 0 AND c1 = 1 AND c2 <= 2 ALLOW FILTERING");
+
+        // test with multi-column partition key and index
+        createTable("CREATE TABLE %s (k1 int, k2 int, c int, v int, PRIMARY KEY ((k1, k2), c))");
+        createIndex("CREATE INDEX ON %s(v)");
+        assertToCQLString("SELECT * FROM %s WHERE v = 1 AND k1 = 1 AND k2 = 2", "SELECT * FROM %s WHERE v = 1 AND k1 = 1 AND k2 = 2 ALLOW FILTERING");
+        assertToCQLString("SELECT * FROM %s WHERE v = 1 AND k1 = 1 AND k2 = 2 AND c = 3", "SELECT * FROM %s WHERE v = 1 AND k1 = 1 AND k2 = 2 AND c = 3 ALLOW FILTERING");
     }
 
     @Override

test/unit/org/apache/cassandra/db/SinglePartitionReadCommandCQLTest.java

@@ -59,6 +59,11 @@ public void testToCQLString()
         assertToCQLString("SELECT * FROM %s WHERE k = 0 AND c < 0 AND v = 1 ALLOW FILTERING", "SELECT * FROM %s WHERE k = 0 AND v = 1 AND c < 0 ALLOW FILTERING");
         assertToCQLString("SELECT * FROM %s WHERE k = 0 AND c >= 0 AND v = 1 ALLOW FILTERING", "SELECT * FROM %s WHERE k = 0 AND v = 1 AND c >= 0 ALLOW FILTERING");
         assertToCQLString("SELECT * FROM %s WHERE k = 0 AND c <= 0 AND v = 1 ALLOW FILTERING", "SELECT * FROM %s WHERE k = 0 AND v = 1 AND c <= 0 ALLOW FILTERING");
+
+        // test with multi-column partition key (without index to test SinglePartitionReadCommand directly)
+        createTable("CREATE TABLE %s (k1 int, k2 int, c int, v int, PRIMARY KEY ((k1, k2), c))");
+        assertToCQLString("SELECT * FROM %s WHERE k1 = 1 AND k2 = 2", "SELECT * FROM %s WHERE k1 = 1 AND k2 = 2 ALLOW FILTERING");
+        assertToCQLString("SELECT * FROM %s WHERE k1 = 1 AND k2 = 2 AND c = 3", "SELECT * FROM %s WHERE k1 = 1 AND k2 = 2 AND c = 3 ALLOW FILTERING");
     }
 
     @Override