Actually get things working

Bennett Hardwick · Bennett Hardwick · commit 7216d141eea3 · 2023-11-07T14:20:43.000+10:00
diff --git a/README.md b/README.md
@@ -253,7 +253,7 @@ To delete a record, use the [`EncryptedTable::delete`] method:
 
 ```rust
 #
-table.delete::<User>("jane@smith.org", User::type_name()).await?;
+table.delete::<User>("jane@smith.org", None).await?;
 ```
 
 #### Querying Records
diff --git a/examples/delete_user.rs b/examples/delete_user.rs
@@ -1,6 +1,6 @@
 mod common;
 use crate::common::User;
-use cryptonamo::{Encryptable, EncryptedTable};
+use cryptonamo::EncryptedTable;
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
@@ -21,15 +21,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let table = EncryptedTable::init(client, "users").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?;
+    table.delete::<User>("jane@smith.org", None).await?;
+    table.delete::<User>("dan@coderdan.co", None).await?;
+    table.delete::<User>("daniel@example.com", None).await?;
 
     Ok(())
 }
diff --git a/examples/get_license.rs b/examples/get_license.rs
@@ -20,7 +20,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let client = aws_sdk_dynamodb::Client::new(&config);
 
     let table = EncryptedTable::init(client, "users").await?;
-    let license: Option<License> = table.get("dan@coderdan.co").await?;
+    let license: Option<License> = table.get("dan@coderdan.co", None).await?;
 
     dbg!(license);
 
diff --git a/examples/get_user.rs b/examples/get_user.rs
@@ -20,7 +20,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let client = aws_sdk_dynamodb::Client::new(&config);
 
     let table = EncryptedTable::init(client, "users").await?;
-    let user: Option<User> = table.get("dan@coderdan.co").await?;
+    let user: Option<User> = table.get("dan@coderdan.co", None).await?;
 
     dbg!(user);
 
diff --git a/src/crypto/sealer.rs b/src/crypto/sealer.rs
@@ -60,13 +60,10 @@ impl<T> Sealer<T> {
         T: Searchable,
     {
         let pk = encrypt_partition_key(&self.inner.partition_key(), cipher).unwrap(); // FIXME
+        let sk = self.inner.sort_key();
 
-        let mut table_entry = TableEntry::new_with_attributes(
-            pk.clone(),
-            T::type_name().to_string(),
-            None,
-            self.unsealed.unprotected(),
-        );
+        let mut table_entry =
+            TableEntry::new_with_attributes(pk.clone(), sk, None, self.unsealed.unprotected());
 
         let protected = T::protected_attributes()
             .iter()
@@ -84,6 +81,8 @@ impl<T> Sealer<T> {
                 }
             });
 
+        let sort_key = self.inner.sort_key();
+
         let protected_indexes = T::protected_indexes();
         let terms: Vec<(&&str, Vec<u8>)> = protected_indexes
             .iter()
@@ -126,7 +125,7 @@ impl<T> Sealer<T> {
                         .clone()
                         .set_term(hex::encode(term))
                         // TODO: HMAC the sort key, too (users#index_name#pk)
-                        .set_sk(format!("{}#{}#{}", T::type_name(), index_name, i)),
+                        .set_sk(format!("{}#{}#{}", &sort_key, index_name, i)),
                 )
             })
             .chain(once(Sealed(table_entry.clone())))
diff --git a/src/encrypted_table/mod.rs b/src/encrypted_table/mod.rs
@@ -112,12 +112,14 @@ impl EncryptedTable {
         QueryBuilder::new(self)
     }
 
-    pub async fn get<T>(&self, pk: &str) -> Result<Option<T>, GetError>
+    pub async fn get<T>(&self, pk: &str, sk: Option<&str>) -> Result<Option<T>, GetError>
     where
         T: Decryptable,
     {
         let pk = encrypt_partition_key(pk, &self.cipher)?;
-        let sk = T::type_name().to_string();
+        let sk = sk
+            .map(|sk| format!("{}#{}", T::type_name(), sk))
+            .unwrap_or_else(|| T::type_name().to_string());
 
         let result = self
             .db
@@ -142,10 +144,12 @@ impl EncryptedTable {
     pub async fn delete<E: Searchable>(
         &self,
         pk: &str,
-        sk: impl Into<String>,
+        sk: Option<&str>,
     ) -> Result<(), DeleteError> {
         let pk = AttributeValue::S(encrypt_partition_key(pk, &self.cipher)?);
-        let sk: String = sk.into();
+        let sk = sk
+            .map(|sk| format!("{}#{}", E::type_name(), sk))
+            .unwrap_or_else(|| E::type_name().to_string());
 
         let sk_to_delete = all_index_keys::<E>(&sk).into_iter().into_iter().chain([sk]);
 
diff --git a/tests/common.rs b/tests/common.rs
@@ -0,0 +1,76 @@
+use aws_sdk_dynamodb::{
+    types::{
+        AttributeDefinition, GlobalSecondaryIndex, KeySchemaElement, KeyType, Projection,
+        ProjectionType, ProvisionedThroughput, ScalarAttributeType,
+    },
+    Client,
+};
+
+pub async fn create_table(client: &Client, table_name: &str) {
+    let _ = client.delete_table().table_name(table_name).send().await;
+
+    client
+        .create_table()
+        .table_name(table_name)
+        .attribute_definitions(
+            AttributeDefinition::builder()
+                .attribute_name("pk")
+                .attribute_type(ScalarAttributeType::S)
+                .build(),
+        )
+        .attribute_definitions(
+            AttributeDefinition::builder()
+                .attribute_name("sk")
+                .attribute_type(ScalarAttributeType::S)
+                .build(),
+        )
+        .attribute_definitions(
+            AttributeDefinition::builder()
+                .attribute_name("term")
+                .attribute_type(ScalarAttributeType::S)
+                .build(),
+        )
+        .key_schema(
+            KeySchemaElement::builder()
+                .attribute_name("pk")
+                .key_type(KeyType::Hash)
+                .build(),
+        )
+        .key_schema(
+            KeySchemaElement::builder()
+                .attribute_name("sk")
+                .key_type(KeyType::Range)
+                .build(),
+        )
+        .provisioned_throughput(
+            ProvisionedThroughput::builder()
+                .read_capacity_units(5)
+                .write_capacity_units(5)
+                .build(),
+        )
+        .global_secondary_indexes(
+            GlobalSecondaryIndex::builder()
+                .index_name("TermIndex")
+                .key_schema(
+                    KeySchemaElement::builder()
+                        .attribute_name("term")
+                        .key_type(KeyType::Hash)
+                        .build(),
+                )
+                .projection(
+                    Projection::builder()
+                        .projection_type(ProjectionType::All)
+                        .build(),
+                )
+                .provisioned_throughput(
+                    ProvisionedThroughput::builder()
+                        .read_capacity_units(5)
+                        .write_capacity_units(5)
+                        .build(),
+                )
+                .build(),
+        )
+        .send()
+        .await
+        .expect("Failed to create table");
+}
diff --git a/tests/query_tests.rs b/tests/query_tests.rs
@@ -1,8 +1,11 @@
-use cryptonamo::{Decryptable, Encryptable, EncryptedTable, Searchable};
+use cryptonamo::{Encryptable, EncryptedTable};
+use cryptonamo_derive::{Decryptable, Searchable};
 use itertools::Itertools;
 use serial_test::serial;
 use std::future::Future;
 
+mod common;
+
 #[derive(Encryptable, Decryptable, Searchable, Debug, PartialEq, Ord, PartialOrd, Eq)]
 #[cryptonamo(sort_key_prefix = "user")]
 pub struct User {
@@ -23,7 +26,7 @@ pub struct User {
 }
 
 impl User {
-    fn new(email: impl Into<String>, name: impl Into<String>, tag: impl Into<String>) -> Self {
+    pub fn new(email: impl Into<String>, name: impl Into<String>, tag: impl Into<String>) -> Self {
         Self {
             name: name.into(),
             email: email.into(),
@@ -41,7 +44,11 @@ async fn run_test<F: Future<Output = ()>>(mut f: impl FnMut(EncryptedTable) -> F
 
     let client = aws_sdk_dynamodb::Client::new(&config);
 
-    let table = EncryptedTable::init(client, "users")
+    let table_name = "test-users";
+
+    common::create_table(&client, table_name).await;
+
+    let table = EncryptedTable::init(client, table_name)
         .await
         .expect("Failed to init table");
 
@@ -131,7 +138,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("dan@coderdan.co").await.expect("Failed to send");
+        let res: Option<User> = table
+            .get("dan@coderdan.co", None)
+            .await
+            .expect("Failed to send");
         assert_eq!(
             res,
             Some(User::new("dan@coderdan.co", "Dan Draper", "blue"))
@@ -145,12 +155,12 @@ async fn test_get_by_partition_key() {
 async fn test_delete() {
     run_test(|table| async move {
         table
-            .delete::<User>("dan@coderdan.co", User::type_name())
+            .delete::<User>("dan@coderdan.co", None)
             .await
             .expect("Failed to send");
 
         let res = table
-            .get::<User>("dan@coderdan.co")
+            .get::<User>("dan@coderdan.co", None)
             .await
             .expect("Failed to send");
         assert_eq!(res, None);
diff --git a/tests/tenant_query_tests.rs b/tests/tenant_query_tests.rs
diff --git a/tests/tenant_tests.rs b/tests/tenant_tests.rs

Original file line number	Diff line number	Diff line change
@@ -253,7 +253,7 @@ To delete a record, use the [`EncryptedTable::delete`] method:
`253`	`253`
`254`	`254`	```rust
`255`	`255`	`#`
`256`		`-table.delete::<User>("jane@smith.org", User::type_name()).await?;`
	`256`	`+table.delete::<User>("jane@smith.org", None).await?;`
`257`	`257`	```
`258`	`258`
`259`	`259`	`#### Querying Records`