Merge pull request #795 from piodul/fix-793

prepared_statement: fix deserialization of compound partition keys
scylladb · Aug 22, 2023 · ffd374b · ffd374b
2 parents 48f393f + 61d65ca
commit ffd374b
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/scylla/src/statement/prepared_statement.rs b/scylla/src/statement/prepared_statement.rs
@@ -376,7 +376,16 @@ impl<'pk> Iterator for PartitionKeyDecoder<'pk> {
             return None;
         }
 
-        let col_idx = self.prepared_metadata.pk_indexes[self.value_index].index as usize;
+        // It should be safe to unwrap because PartitionKeyIndex::sequence
+        // fields of the pk_indexes form a permutation, but let's err
+        // on the safe side in case of bugs.
+        let pk_index = self
+            .prepared_metadata
+            .pk_indexes
+            .iter()
+            .find(|pki| pki.sequence == self.value_index as u16)?;
+
+        let col_idx = pk_index.index as usize;
         let spec = &self.prepared_metadata.col_specs[col_idx];
 
         let cell = if self.prepared_metadata.pk_indexes.len() == 1 {
@@ -432,14 +441,15 @@ mod tests {
                 typ,
             })
             .collect();
-        let pk_indexes = idx
+        let mut pk_indexes = idx
             .into_iter()
             .enumerate()
             .map(|(sequence, index)| PartitionKeyIndex {
                 index: index as u16,
                 sequence: sequence as u16,
             })
-            .collect();
+            .collect::<Vec<_>>();
+        pk_indexes.sort_unstable_by_key(|pki| pki.index);
         PreparedMetadata {
             flags: 0,
             col_count: col_specs.len(),