Skip to content

Commit

Permalink
prepared_statement: fix deserialization of compound partition keys
Browse files Browse the repository at this point in the history
When PartitionKeyDeserializer wants to deserialize i-th column of the
partition key, it needs to know the type of that column. It has access
to PreparedMetadata of the prepared query the partition key was
constructed from, which has a mapping between bind markers and types.
Hence, it needs to learn to which bind marker the i-th pk column
corresponds to.

The `self.pk_indexes` of type `Vec<PartitionKeyIndex>` is supposed to
help with that. The PartitionKeyIndex structure looks like this:

    pub struct PartitionKeyIndex {
        /// index in the serialized values
        pub index: u16,
        /// sequence number in partition key
        pub sequence: u16,
    }

So, we need to find an entry `pki` with `pki.sequence == i` and
`pki.index` will be the desired bind marker index.

Currently, we naively look at the `self.pk_indexes[i]` entry - which is
wrong. The `pk_indexes` vec is actually ordered by `index`, so we might
accidentally obtain an index of the wrong partition key column this way.
Although there was a test that was supposed to check this scenario, it
constructed `PreparedMetadata` manually and didn't sort the `pk_indexes`
by `index` field, which resulted in it accidentally passing.

In order to fix this, the `PartitionKeyIndex` is now doing a linear
lookup in order to find a `pki` with `pki.sequence == i`. This lookup
should be fine as partition keys usually have a little number of fields.

Fixes: #793
  • Loading branch information
piodul committed Aug 22, 2023
1 parent 9d76db1 commit 61d65ca
Showing 1 changed file with 13 additions and 3 deletions.
16 changes: 13 additions & 3 deletions scylla/src/statement/prepared_statement.rs
Original file line number Diff line number Diff line change
Expand Up @@ -382,7 +382,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 {
Expand Down Expand Up @@ -438,14 +447,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(),
Expand Down

0 comments on commit 61d65ca

Please sign in to comment.