Line data    Source code

       1              : use std::io::ErrorKind;
       2              : 
       3              : use anyhow::Ok;
       4              : 
       5              : use crate::pqproto::{CancelKeyData, id_to_cancel_key};
       6              : 
       7              : pub mod keyspace {
       8              :     pub const CANCEL_PREFIX: &str = "cancel";
       9              : }
      10              : 
      11              : #[derive(Clone, Debug, Eq, PartialEq)]
      12              : pub(crate) enum KeyPrefix {
      13              :     Cancel(CancelKeyData),
      14              : }
      15              : 
      16              : impl KeyPrefix {
      17            1 :     pub(crate) fn build_redis_key(&self) -> String {
      18            1 :         match self {
      19            1 :             KeyPrefix::Cancel(key) => {
      20            1 :                 let id = key.0.get();
      21            1 :                 let keyspace = keyspace::CANCEL_PREFIX;
      22            1 :                 format!("{keyspace}:{id:x}")
      23            1 :             }
      24            1 :         }
      25            1 :     }
      26              : 
      27              :     #[allow(dead_code)]
      28            2 :     pub(crate) fn as_str(&self) -> &'static str {
      29            2 :         match self {
      30            2 :             KeyPrefix::Cancel(_) => keyspace::CANCEL_PREFIX,
      31            2 :         }
      32            2 :     }
      33              : }
      34              : 
      35              : #[allow(dead_code)]
      36            1 : pub(crate) fn parse_redis_key(key: &str) -> anyhow::Result<KeyPrefix> {
      37            1 :     let (prefix, key_str) = key.split_once(':').ok_or_else(|| {
      38            0 :         anyhow::anyhow!(std::io::Error::new(
      39            0 :             ErrorKind::InvalidData,
      40            0 :             "missing prefix"
      41            0 :         ))
      42            1 :     })?;
      43              : 
      44            1 :     match prefix {
      45            1 :         keyspace::CANCEL_PREFIX => {
      46            1 :             let id = u64::from_str_radix(key_str, 16)?;
      47              : 
      48            1 :             Ok(KeyPrefix::Cancel(id_to_cancel_key(id)))
      49              :         }
      50            0 :         _ => Err(anyhow::anyhow!(std::io::Error::new(
      51            0 :             ErrorKind::InvalidData,
      52            0 :             "unknown prefix"
      53            0 :         ))),
      54              :     }
      55            1 : }
      56              : 
      57              : #[cfg(test)]
      58              : mod tests {
      59              :     use super::*;
      60              : 
      61              :     #[test]
      62            1 :     fn test_build_redis_key() {
      63            1 :         let cancel_key: KeyPrefix = KeyPrefix::Cancel(id_to_cancel_key(12345 << 32 | 54321));
      64            1 : 
      65            1 :         let redis_key = cancel_key.build_redis_key();
      66            1 :         assert_eq!(redis_key, "cancel:30390000d431");
      67            1 :     }
      68              : 
      69              :     #[test]
      70            1 :     fn test_parse_redis_key() {
      71            1 :         let redis_key = "cancel:30390000d431";
      72            1 :         let key: KeyPrefix = parse_redis_key(redis_key).expect("Failed to parse key");
      73            1 : 
      74            1 :         let ref_key = id_to_cancel_key(12345 << 32 | 54321);
      75            1 : 
      76            1 :         assert_eq!(key.as_str(), KeyPrefix::Cancel(ref_key).as_str());
      77            1 :         let KeyPrefix::Cancel(cancel_key) = key;
      78            1 :         assert_eq!(ref_key, cancel_key);
      79            1 :     }
      80              : }