Change to dynamic sort key, new attrs

Author: Bennett Hardwick

README.md

@@ -253,7 +253,7 @@ To delete a record, use the [`EncryptedTable::delete`] method:
 
 ```rust
 #
-table.delete::<User>("jane@smith.org").await?;
+table.delete::<User>("jane@smith.org", User::type_name()).await?;
 ```
 
 #### Querying Records

cryptonamo-derive/src/encryptable.rs

@@ -33,13 +33,24 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
             }
         }));
 
+    let sort_key_impl = if let Some(sort_key_field) = &settings.sort_key_field {
+        let sort_key_attr = format_ident!("{sort_key_field}");
+        quote! { format!("{}#{}", Self::type_name(), self.#sort_key_attr) }
+    } else {
+        quote! { Self::type_name().into() }
+    };
+
     let expanded = quote! {
         #[automatically_derived]
         impl cryptonamo::traits::Encryptable for #ident {
             fn type_name() -> &'static str {
                 #type_name
             }
 
+            fn sort_key(&self) -> String {
+                #sort_key_impl
+            }
+
             fn partition_key(&self) -> String {
                 self.#partition_key.to_string()
             }

cryptonamo-derive/src/lib.rs

@@ -10,21 +10,21 @@ mod settings;
 use proc_macro::TokenStream;
 use syn::{parse_macro_input, DeriveInput};
 
-#[proc_macro_derive(Encryptable, attributes(cryptonamo))]
+#[proc_macro_derive(Encryptable, attributes(cryptonamo, sort_key, partition_key))]
 pub fn derive_encryptable(input: TokenStream) -> TokenStream {
     encryptable::derive_encryptable(parse_macro_input!(input as DeriveInput))
         .unwrap_or_else(syn::Error::into_compile_error)
         .into()
 }
 
-#[proc_macro_derive(Decryptable, attributes(cryptonamo))]
+#[proc_macro_derive(Decryptable, attributes(cryptonamo, sort_key, partition_key))]
 pub fn derive_decryptable(input: TokenStream) -> TokenStream {
     decryptable::derive_decryptable(parse_macro_input!(input as DeriveInput))
         .unwrap_or_else(syn::Error::into_compile_error)
         .into()
 }
 
-#[proc_macro_derive(Searchable, attributes(cryptonamo))]
+#[proc_macro_derive(Searchable, attributes(cryptonamo, sort_key, partition_key))]
 pub fn derive_searchable(input: TokenStream) -> TokenStream {
     searchable::derive_searchable(parse_macro_input!(input as DeriveInput))
         .unwrap_or_else(syn::Error::into_compile_error)

cryptonamo-derive/src/settings/builder.rs

@@ -7,7 +7,8 @@ pub(crate) struct SettingsBuilder {
     ident: Ident,
     type_name: String,
     sort_key_prefix: Option<String>,
-    partition_key: Option<String>,
+    sort_key_field: Option<String>,
+    partition_key_field: Option<String>,
     protected_attributes: Vec<String>,
     unprotected_attributes: Vec<String>,
     skipped_attributes: Vec<String>,
@@ -33,7 +34,8 @@ impl SettingsBuilder {
             ident: input.ident.clone(),
             type_name,
             sort_key_prefix: None,
-            partition_key: None,
+            sort_key_field: None,
+            partition_key_field: None,
             protected_attributes: Vec::new(),
             unprotected_attributes: Vec::new(),
             skipped_attributes: Vec::new(),
@@ -92,6 +94,48 @@ impl SettingsBuilder {
 
                     // Parse the meta for the field
                     for attr in &field.attrs {
+                        if attr.path().is_ident("sort_key") {
+                            let field_name = ident
+                                .as_ref()
+                                .ok_or_else(|| {
+                                    syn::Error::new_spanned(
+                                        field,
+                                        "internal error: identifier was not Some",
+                                    )
+                                })?
+                                .to_string();
+
+                            if let Some(f) = &self.sort_key_field {
+                                return Err(syn::Error::new_spanned(
+                                    field,
+                                    format!("sort key was already specified to be '{f}'"),
+                                ));
+                            }
+
+                            self.sort_key_field = Some(field_name);
+                        }
+
+                        if attr.path().is_ident("partition_key") {
+                            let field_name = ident
+                                .as_ref()
+                                .ok_or_else(|| {
+                                    syn::Error::new_spanned(
+                                        field,
+                                        "internal error: identifier was not Some",
+                                    )
+                                })?
+                                .to_string();
+
+                            if let Some(f) = &self.partition_key_field {
+                                return Err(syn::Error::new_spanned(
+                                    field,
+                                    format!("partition key was already specified to be '{f}'"),
+                                ));
+                            }
+
+                            self.partition_key_field = Some(field_name);
+                        }
+
                         if attr.path().is_ident("cryptonamo") {
                             let mut query: Option<(String, String, Span)> = None;
                             let mut compound_index_name: Option<(String, Span)> = None;
@@ -203,7 +247,8 @@ impl SettingsBuilder {
             ident,
             type_name,
             sort_key_prefix,
-            partition_key,
+            sort_key_field,
+            partition_key_field: partition_key,
             protected_attributes,
             unprotected_attributes,
             skipped_attributes,
@@ -222,7 +267,8 @@ impl SettingsBuilder {
         Ok(Settings {
             ident,
             sort_key_prefix,
-            partition_key: Some(partition_key), // TODO: Remove the Some
+            sort_key_field,
+            partition_key_field: Some(partition_key), // TODO: Remove the Some
             protected_attributes,
             unprotected_attributes,
             skipped_attributes,
@@ -236,7 +282,7 @@ impl SettingsBuilder {
     }
 
     pub(crate) fn set_partition_key(&mut self, value: String) -> Result<(), syn::Error> {
-        self.partition_key = Some(value);
+        self.partition_key_field = Some(value);
         Ok(())
     }
 

cryptonamo-derive/src/settings/mod.rs

@@ -16,7 +16,8 @@ pub(crate) enum AttributeMode {
 pub(crate) struct Settings {
     ident: Ident,
     pub(crate) sort_key_prefix: String,
-    pub(crate) partition_key: Option<String>,
+    pub(crate) sort_key_field: Option<String>,
+    pub(crate) partition_key_field: Option<String>,
     protected_attributes: Vec<String>,
     unprotected_attributes: Vec<String>,
 
@@ -70,7 +71,7 @@ impl Settings {
     }
 
     pub(crate) fn get_partition_key(&self) -> Result<String, syn::Error> {
-        self.partition_key.clone().ok_or_else(|| {
+        self.partition_key_field.clone().ok_or_else(|| {
             syn::Error::new_spanned(
                 &self.sort_key_prefix,
                 "No partition key defined for this struct",

examples/delete_user.rs

@@ -1,6 +1,6 @@
 mod common;
 use crate::common::User;
-use cryptonamo::EncryptedTable;
+use cryptonamo::{Encryptable, EncryptedTable};
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
@@ -21,9 +21,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let table = EncryptedTable::init(client, "users").await?;
 
-    table.delete::<User>("jane@smith.org").await?;
-    table.delete::<User>("dan@coderdan.co").await?;
-    table.delete::<User>("daniel@example.com").await?;
+    table
+        .delete::<User>("jane@smith.org", User::type_name())
+        .await?;
+    table
+        .delete::<User>("dan@coderdan.co", User::type_name())
+        .await?;
+    table
+        .delete::<User>("daniel@example.com", User::type_name())
+        .await?;
 
     Ok(())
 }

src/crypto/mod.rs

@@ -44,13 +44,13 @@ pub enum CryptoError {
     Other(String),
 }
 
-pub(crate) fn all_index_keys<E: Searchable + Encryptable>() -> Vec<String> {
+pub(crate) fn all_index_keys<E: Searchable + Encryptable>(sort_key: &str) -> Vec<String> {
     E::protected_indexes()
         .iter()
         .flat_map(|index_name| {
             (0..)
                 .take(MAX_TERMS_PER_INDEX)
-                .map(|i| format!("{}#{}#{}", E::type_name(), index_name, i))
+                .map(|i| format!("{}#{}#{}", sort_key, index_name, i))
                 .collect::<Vec<String>>()
         })
         .collect()

src/encrypted_table/mod.rs

@@ -138,12 +138,16 @@ impl EncryptedTable {
         }
     }
 
-    pub async fn delete<E: Searchable>(&self, pk: &str) -> Result<(), DeleteError> {
+    // TODO: create PrimaryKey abstraction
+    pub async fn delete<E: Searchable>(
+        &self,
+        pk: &str,
+        sk: impl Into<String>,
+    ) -> Result<(), DeleteError> {
         let pk = AttributeValue::S(encrypt_partition_key(pk, &self.cipher)?);
+        let sk: String = sk.into();
 
-        let sk_to_delete = [E::type_name().to_string()]
-            .into_iter()
-            .chain(all_index_keys::<E>().into_iter());
+        let sk_to_delete = all_index_keys::<E>(&sk).into_iter().into_iter().chain([sk]);
 
         let transact_items = sk_to_delete.map(|sk| {
             TransactWriteItem::builder()
@@ -177,6 +181,7 @@ impl EncryptedTable {
     {
         let mut seen_sk = HashSet::new();
 
+        let sk = record.sort_key();
         let sealer: Sealer<T> = record.into_sealer()?;
         let (pk, sealed) = sealer.seal(&self.cipher, 12).await?;
 
@@ -198,7 +203,7 @@ impl EncryptedTable {
             );
         }
 
-        for index_sk in all_index_keys::<T>() {
+        for index_sk in all_index_keys::<T>(&sk) {
             if seen_sk.contains(&index_sk) {
                 continue;
             }

src/traits/mod.rs

@@ -26,6 +26,11 @@ pub enum WriteConversionError {
 pub trait Encryptable: Debug + Sized {
     // TODO: Add a function indicating that the root should be stored
     fn type_name() -> &'static str;
+
+    fn sort_key(&self) -> String {
+        Self::type_name().into()
+    }
+
     fn partition_key(&self) -> String;
 
     fn protected_attributes() -> Vec<&'static str>;

tests/query_tests.rs

@@ -4,11 +4,11 @@ use serial_test::serial;
 use std::future::Future;
 
 #[derive(Encryptable, Decryptable, Searchable, Debug, PartialEq, Ord, PartialOrd, Eq)]
-#[cryptonamo(partition_key = "email")]
 #[cryptonamo(sort_key_prefix = "user")]
 pub struct User {
     #[cryptonamo(query = "exact", compound = "email#name")]
     #[cryptonamo(query = "exact")]
+    #[partition_key]
     pub email: String,
 
     #[cryptonamo(query = "prefix", compound = "email#name")]
@@ -33,7 +33,7 @@ impl User {
     }
 }
 
-async fn run_test<F: Future<Output = ()>>(f: impl FnOnce(EncryptedTable) -> F) {
+async fn run_test<F: Future<Output = ()>>(mut f: impl FnMut(EncryptedTable) -> F) {
     let config = aws_config::from_env()
         .endpoint_url("http://localhost:8000")
         .load()
@@ -145,7 +145,7 @@ async fn test_get_by_partition_key() {
 async fn test_delete() {
     run_test(|table| async move {
         table
-            .delete::<User>("dan@coderdan.co")
+            .delete::<User>("dan@coderdan.co", User::type_name())
             .await
             .expect("Failed to send");