Support sk field

Author: Bennett Hardwick

cryptonamo-derive/src/settings/builder.rs

@@ -19,7 +19,7 @@ impl SortKeyPrefix {
     }
 }
 
-const RESERVED_FIELD_NAMES: &'static [&'static str] = &["term", "sk"];
+const RESERVED_FIELD_NAMES: &'static [&'static str] = &["term"];
 
 pub(crate) struct SettingsBuilder {
     ident: Ident,
@@ -117,7 +117,7 @@ impl SettingsBuilder {
                     .collect();
 
                 let explicit_pk = all_field_names.contains(&String::from("pk"));
-                // let explicit_sk = all_field_names.contains(&String::from("sk"));
+                let explicit_sk = all_field_names.contains(&String::from("sk"));
 
                 let mut compound_indexes: HashMap<String, Vec<(String, String, Span)>> =
                     Default::default();
@@ -172,6 +172,19 @@ impl SettingsBuilder {
                     // Parse the meta for the field
                     for attr in &field.attrs {
                         if attr.path().is_ident("sort_key") {
+                            if explicit_sk && field_name != "sk" {
+                                return Err(syn::Error::new_spanned(
+                                    field,
+                                    format!("field '{field_name}' cannot be used as sort key as struct contains field named 'sk' which must be used")
+                                ));
+                            }
+
+                            if explicit_sk {
+                                // if the 'sk' field is set then there should be no prefix
+                                // otherwise when deserialising the sk value would be incorrect
+                                self.sort_key_prefix = SortKeyPrefix::None;
+                            }
+
                             if let Some(f) = &self.sort_key_field {
                                 return Err(syn::Error::new_spanned(
                                     field,
@@ -182,7 +195,7 @@ impl SettingsBuilder {
                             self.sort_key_field = Some(field_name.clone());
                         }
 
-                        if attr.path().is_ident("partition_key"){
+                        if attr.path().is_ident("partition_key") {
                             if explicit_pk && field_name != "pk" {
                                 return Err(syn::Error::new_spanned(
                                     field,

src/crypto/sealed.rs

@@ -40,11 +40,11 @@ impl Sealed {
         let ciphertexts = T::decryptable_attributes()
             .into_iter()
             .map(|name| {
-                let attribute = if name == "pk" {
-                    self.inner().attributes.get("__pk")
-                } else {
-                    self.inner().attributes.get(name)
-                };
+                let attribute = self.inner().attributes.get(match name {
+                    "pk" => "__pk",
+                    "sk" => "__sk",
+                    _ => name,
+                });
 
                 attribute
                     .ok_or_else(|| SealError::MissingAttribute(name.to_string()))?

src/crypto/sealer.rs

@@ -90,11 +90,14 @@ impl<T> Sealer<T> {
             .zip(T::protected_attributes().into_iter())
             .for_each(|(enc, name)| {
                 if let Some(e) = enc {
-                    if name == "pk" {
-                        table_entry.add_attribute("__pk", e.into());
-                    } else {
-                        table_entry.add_attribute(name, e.into());
-                    }
+                    table_entry.add_attribute(
+                        match name {
+                            "pk" => "__pk",
+                            "sk" => "__sk",
+                            _ => name,
+                        },
+                        e.into(),
+                    );
                 }
             });
 

src/traits/primary_key.rs

@@ -2,7 +2,7 @@ use crate::Encryptable;
 
 pub struct PrimaryKeyParts {
     pub(crate) pk: String,
-    pub(crate) sk: String
+    pub(crate) sk: String,
 }
 
 pub trait PrimaryKey: private::Sealed {
@@ -17,20 +17,20 @@ impl Pk {
     }
 }
 
-impl From<String> for Pk {
-    fn from(value: String) -> Self {
-        Self(value)
-    }
-}
-
-impl From<&str> for Pk {
-    fn from(value: &str) -> Self {
+impl<P: Into<String>> From<P> for Pk {
+    fn from(value: P) -> Self {
         Self::new(value)
     }
 }
 
 pub struct PkSk(String, String);
 
+impl<Pk: Into<String>, Sk: Into<String>> From<(Pk, Sk)> for PkSk {
+    fn from(value: (Pk, Sk)) -> Self {
+        Self::new(value.0, value.1)
+    }
+}
+
 impl PkSk {
     pub fn new(pk: impl Into<String>, sk: impl Into<String>) -> Self {
         Self(pk.into(), sk.into())

tests/pk_tests.rs

@@ -7,31 +7,27 @@ use std::future::Future;
 mod common;
 
 #[derive(Encryptable, Decryptable, Searchable, Debug, PartialEq, Ord, PartialOrd, Eq)]
-#[cryptonamo(sort_key_prefix = "user")]
 pub struct User {
-    #[cryptonamo(query = "exact", compound = "pk#name")]
+    #[cryptonamo(query = "exact", compound = "pk#sk")]
     #[cryptonamo(query = "exact")]
     #[partition_key]
     pub pk: String,
 
-    #[cryptonamo(query = "prefix", compound = "pk#name")]
+    #[cryptonamo(query = "prefix", compound = "pk#sk")]
     #[cryptonamo(query = "prefix")]
-    pub name: String,
+    #[sort_key]
+    pub sk: String,
 
     #[cryptonamo(plaintext)]
     pub tag: String,
-
-    #[cryptonamo(skip)]
-    pub temp: bool,
 }
 
 impl User {
     pub fn new(email: impl Into<String>, name: impl Into<String>, tag: impl Into<String>) -> Self {
         Self {
-            name: name.into(),
             pk: email.into(),
+            sk: name.into(),
             tag: tag.into(),
-            temp: false,
         }
     }
 }
@@ -44,7 +40,7 @@ async fn run_test<F: Future<Output = ()>>(mut f: impl FnMut(EncryptedTable) -> F
 
     let client = aws_sdk_dynamodb::Client::new(&config);
 
-    let table_name = "test-users";
+    let table_name = "pk-sk-users";
 
     common::create_table(&client, table_name).await;
 
@@ -95,7 +91,7 @@ async fn test_query_single_prefix() {
     run_test(|table| async move {
         let res: Vec<User> = table
             .query()
-            .starts_with("name", "Dan")
+            .starts_with("sk", "Dan")
             .send()
             .await
             .expect("Failed to query")
@@ -120,7 +116,7 @@ async fn test_query_compound() {
     run_test(|table| async move {
         let res: Vec<User> = table
             .query()
-            .starts_with("name", "Dan")
+            .starts_with("sk", "Dan")
             .eq("pk", "[email protected]")
             .send()
             .await
@@ -138,7 +134,10 @@ async fn test_query_compound() {
 #[serial]
 async fn test_get_by_partition_key() {
     run_test(|table| async move {
-        let res: Option<User> = table.get("[email protected]").await.expect("Failed to send");
+        let res: Option<User> = table
+            .get(("[email protected]", "Dan Draper"))
+            .await
+            .expect("Failed to send");
         assert_eq!(
             res,
             Some(User::new("[email protected]", "Dan Draper", "blue"))
@@ -152,19 +151,19 @@ async fn test_get_by_partition_key() {
 async fn test_delete() {
     run_test(|table| async move {
         table
-            .delete::<User>("[email protected]")
+            .delete::<User>(("[email protected]", "Dan Draper"))
             .await
             .expect("Failed to send");
 
         let res = table
-            .get::<User>("[email protected]")
+            .get::<User>(("[email protected]", "Dan Draper"))
             .await
             .expect("Failed to send");
         assert_eq!(res, None);
 
         let res = table
             .query::<User>()
-            .starts_with("name", "Dan")
+            .starts_with("sk", "Dan")
             .send()
             .await
             .expect("Failed to send");
@@ -184,7 +183,7 @@ async fn test_delete() {
         let res = table
             .query::<User>()
             .eq("pk", "[email protected]")
-            .starts_with("name", "Dan")
+            .starts_with("sk", "Dan")
             .send()
             .await
             .expect("Failed to send");

tests/ui/using-pk-instead-of-pk-sk.stderr

@@ -6,6 +6,7 @@ error[E0277]: the trait bound `PkSk: From<&str>` is not satisfied
    |                                    |
    |                                    required by a bound introduced by this call
    |
+   = help: the trait `From<(Pk, Sk)>` is implemented for `PkSk`
    = note: required for `&str` to implement `Into<PkSk>`
 note: required by a bound in `EncryptedTable::get`
   --> src/encrypted_table/mod.rs