Update to primary key abstraction

Bennett Hardwick · Bennett Hardwick · commit 65830a62365e · 2023-11-07T15:26:41.000+10:00
diff --git a/cryptonamo-derive/src/encryptable.rs b/cryptonamo-derive/src/encryptable.rs
@@ -40,9 +40,17 @@ pub(crate) fn derive_encryptable(input: DeriveInput) -> Result<TokenStream, syn:
         quote! { Self::type_name().into() }
     };
 
+    let primary_key_impl = if settings.sort_key_field.is_some() {
+        quote! { type PrimaryKey = cryptonamo::PkSk; }
+    } else {
+        quote! { type PrimaryKey = cryptonamo::Pk; }
+    };
+
     let expanded = quote! {
         #[automatically_derived]
         impl cryptonamo::traits::Encryptable for #ident {
+            #primary_key_impl
+
             fn type_name() -> &'static str {
                 #type_name
             }
diff --git a/examples/delete_user.rs b/examples/delete_user.rs
@@ -21,9 +21,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let table = EncryptedTable::init(client, "users").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?;
+    table.delete::<User>("jane@smith.org").await?;
+    table.delete::<User>("dan@coderdan.co").await?;
+    table.delete::<User>("daniel@example.com").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", None).await?;
+    let license: Option<License> = table.get("dan@coderdan.co").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", None).await?;
+    let user: Option<User> = table.get("dan@coderdan.co").await?;
 
     dbg!(user);
 
diff --git a/src/encrypted_table/mod.rs b/src/encrypted_table/mod.rs
@@ -6,7 +6,10 @@ pub use self::{
 };
 use crate::{
     crypto::*,
-    traits::{Decryptable, ReadConversionError, Searchable, WriteConversionError},
+    traits::{
+        Decryptable, PrimaryKey, PrimaryKeyParts, ReadConversionError, Searchable,
+        WriteConversionError,
+    },
 };
 use aws_sdk_dynamodb::{
     types::{AttributeValue, Delete, Put, TransactWriteItem},
@@ -112,14 +115,13 @@ impl EncryptedTable {
         QueryBuilder::new(self)
     }
 
-    pub async fn get<T>(&self, pk: &str, sk: Option<&str>) -> Result<Option<T>, GetError>
+    pub async fn get<T>(&self, k: impl Into<T::PrimaryKey>) -> Result<Option<T>, GetError>
     where
         T: Decryptable,
     {
-        let pk = encrypt_partition_key(pk, &self.cipher)?;
-        let sk = sk
-            .map(|sk| format!("{}#{}", T::type_name(), sk))
-            .unwrap_or_else(|| T::type_name().to_string());
+        let PrimaryKeyParts { pk, sk } = k.into().into_parts(T::type_name());
+
+        let pk = encrypt_partition_key(&pk, &self.cipher)?;
 
         let result = self
             .db
@@ -141,15 +143,10 @@ impl EncryptedTable {
     }
 
     // TODO: create PrimaryKey abstraction
-    pub async fn delete<E: Searchable>(
-        &self,
-        pk: &str,
-        sk: Option<&str>,
-    ) -> Result<(), DeleteError> {
-        let pk = AttributeValue::S(encrypt_partition_key(pk, &self.cipher)?);
-        let sk = sk
-            .map(|sk| format!("{}#{}", E::type_name(), sk))
-            .unwrap_or_else(|| E::type_name().to_string());
+    pub async fn delete<E: Searchable>(&self, k: impl Into<E::PrimaryKey>) -> Result<(), DeleteError> {
+        let PrimaryKeyParts { pk, sk } = k.into().into_parts(E::type_name());
+
+        let pk = AttributeValue::S(encrypt_partition_key(&pk, &self.cipher)?);
 
         let sk_to_delete = all_index_keys::<E>(&sk).into_iter().into_iter().chain([sk]);
 
diff --git a/src/lib.rs b/src/lib.rs
@@ -403,7 +403,7 @@ mod error;
 pub mod traits;
 pub use encrypted_table::{EncryptedTable, QueryBuilder};
 pub use error::Error;
-pub use traits::{Decryptable, Encryptable, Searchable};
+pub use traits::{Decryptable, Encryptable, Searchable, PrimaryKey, PkSk, Pk};
 
 #[doc(hidden)]
 pub use cryptonamo_derive::{Decryptable, Encryptable, Searchable};
diff --git a/src/traits/mod.rs b/src/traits/mod.rs
@@ -6,6 +6,10 @@ pub use cipherstash_client::encryption::{
     },
     Plaintext,
 };
+
+mod primary_key;
+pub use primary_key::*;
+
 use std::fmt::Debug;
 use thiserror::Error;
 
@@ -24,6 +28,8 @@ pub enum WriteConversionError {
 }
 
 pub trait Encryptable: Debug + Sized {
+    type PrimaryKey: PrimaryKey;
+
     // TODO: Add a function indicating that the root should be stored
     fn type_name() -> &'static str;
 
@@ -61,6 +67,7 @@ pub trait Searchable: Encryptable {
 
 pub trait Decryptable: Encryptable {
     /// Convert an `Unsealed` into a `Self`.
+
     fn from_unsealed(unsealed: Unsealed) -> Result<Self, SealError>;
 
     /// Defines which attributes are decryptable for this type.
diff --git a/src/traits/primary_key.rs b/src/traits/primary_key.rs
@@ -0,0 +1,63 @@
+pub struct PrimaryKeyParts {
+    pub(crate) pk: String,
+    pub(crate) sk: String,
+}
+
+pub trait PrimaryKey: private::Sealed {
+    fn into_parts(self, sk_prefix: &str) -> PrimaryKeyParts;
+}
+
+pub struct Pk(String);
+
+impl Pk {
+    pub fn new(pk: impl Into<String>) -> Self {
+        Self(pk.into())
+    }
+}
+
+impl From<String> for Pk {
+    fn from(value: String) -> Self {
+        Self(value)
+    }
+}
+
+impl From<&str> for Pk {
+    fn from(value: &str) -> Self {
+        Self::new(value)
+    }
+}
+
+pub struct PkSk(String, String);
+
+impl PkSk {
+    pub fn new(pk: impl Into<String>, sk: impl Into<String>) -> Self {
+        Self(pk.into(), sk.into())
+    }
+}
+
+mod private {
+    use super::*;
+
+    pub trait Sealed {}
+
+    impl Sealed for Pk {}
+    impl Sealed for PkSk {}
+}
+
+impl PrimaryKey for Pk {
+    fn into_parts(self, sk_prefix: &str) -> PrimaryKeyParts {
+        PrimaryKeyParts {
+            pk: self.0,
+            sk: sk_prefix.to_string(),
+        }
+    }
+}
+
+impl PrimaryKey for PkSk {
+    fn into_parts(self, sk_prefix: &str) -> PrimaryKeyParts {
+        PrimaryKeyParts {
+            pk: self.0,
+            sk: format!("{}#{}", sk_prefix, self.1),
+        }
+    }
+}
diff --git a/tests/compile_tests.rs b/tests/compile_tests.rs
@@ -7,6 +7,7 @@ fn ui_tests() {
     t.compile_fail("tests/ui/compound-index-too-many-fields.rs");
     t.compile_fail("tests/ui/index-unsupported.rs");
     t.compile_fail("tests/ui/compound-index-unsupported.rs");
+    t.compile_fail("tests/ui/using-pk-instead-of-pk-sk.rs");
 
     t.pass("tests/ui/pass.rs");
 }
diff --git a/tests/query_tests.rs b/tests/query_tests.rs
@@ -138,10 +138,7 @@ 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", None)
-            .await
-            .expect("Failed to send");
+        let res: Option<User> = table.get("dan@coderdan.co").await.expect("Failed to send");
         assert_eq!(
             res,
             Some(User::new("dan@coderdan.co", "Dan Draper", "blue"))
@@ -155,12 +152,12 @@ async fn test_get_by_partition_key() {
 async fn test_delete() {
     run_test(|table| async move {
         table
-            .delete::<User>("dan@coderdan.co", None)
+            .delete::<User>("dan@coderdan.co")
             .await
             .expect("Failed to send");
 
         let res = table
-            .get::<User>("dan@coderdan.co", None)
+            .get::<User>("dan@coderdan.co")
             .await
             .expect("Failed to send");
         assert_eq!(res, None);
diff --git a/tests/tenant_tests.rs b/tests/tenant_tests.rs
@@ -1,4 +1,4 @@
-use cryptonamo::{Decryptable, Encryptable, EncryptedTable, Searchable};
+use cryptonamo::{Decryptable, Encryptable, EncryptedTable, Searchable, PkSk};
 use itertools::Itertools;
 use serial_test::serial;
 use std::future::Future;
@@ -143,7 +143,7 @@ async fn test_query_compound() {
 async fn test_get_by_partition_key() {
     run_test(|table| async move {
         let res: Option<User> = table
-            .get("first-tenant", Some("dan@coderdan.co"))
+            .get(PkSk::new("first-tenant", "dan@coderdan.co"))
             .await
             .expect("Failed to send");
         assert_eq!(
@@ -159,12 +159,12 @@ async fn test_get_by_partition_key() {
 async fn test_delete() {
     run_test(|table| async move {
         table
-            .delete::<User>("first-tenant", Some("dan@coderdan.co"))
+            .delete::<User>(PkSk::new("first-tenant", "dan@coderdan.co"))
             .await
             .expect("Failed to send");
 
         let res = table
-            .get::<User>("first-tenant", Some("dan@coderdan.co"))
+            .get::<User>(PkSk::new("first-tenant", "dan@coderdan.co"))
             .await
             .expect("Failed to send");
         assert_eq!(res, None);
diff --git a/tests/ui/using-pk-instead-of-pk-sk.rs b/tests/ui/using-pk-instead-of-pk-sk.rs
@@ -0,0 +1,28 @@
+use cryptonamo::{Decryptable, Encryptable, Searchable};
+use cryptonamo::EncryptedTable;
+
+#[derive(Debug, Encryptable, Decryptable, Searchable)]
+struct User {
+    #[cryptonamo(query = "exact")]
+    #[partition_key]
+    email: String,
+
+    #[cryptonamo(query = "exact")]
+    #[sort_key]
+    name: String,
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let config = aws_config::from_env()
+        .endpoint_url("http://localhost:8000")
+        .load()
+        .await;
+
+    let client = aws_sdk_dynamodb::Client::new(&config);
+    let table = EncryptedTable::init(client, "users").await?;
+
+    let user: Option<User> = table.get("user@example.com").await?;
+
+    Ok(())
+}
diff --git a/tests/ui/using-pk-instead-of-pk-sk.stderr b/tests/ui/using-pk-instead-of-pk-sk.stderr
@@ -0,0 +1,14 @@
+error[E0277]: the trait bound `PkSk: From<&str>` is not satisfied
+  --> tests/ui/using-pk-instead-of-pk-sk.rs:25:40
+   |
+25 |     let user: Option<User> = table.get("user@example.com").await?;
+   |                                    --- ^^^^^^^^^^^^^^^^^^ the trait `From<&str>` is not implemented for `PkSk`
+   |                                    |
+   |                                    required by a bound introduced by this call
+   |
+   = note: required for `&str` to implement `Into<PkSk>`
+note: required by a bound in `EncryptedTable::get`
+  --> src/encrypted_table/mod.rs
+   |
+   |     pub async fn get<T>(&self, k: impl Into<T::PrimaryKey>) -> Result<Option<T>, GetError>
+   |                                        ^^^^^^^^^^^^^^^^^^^ required by this bound in `EncryptedTable::get`

Original file line number	Diff line number	Diff line change
`@@ -7,6 +7,7 @@ fn ui_tests() {`
`7`	`7`	`t.compile_fail("tests/ui/compound-index-too-many-fields.rs");`
`8`	`8`	`t.compile_fail("tests/ui/index-unsupported.rs");`
`9`	`9`	`t.compile_fail("tests/ui/compound-index-unsupported.rs");`
	`10`	`+ t.compile_fail("tests/ui/using-pk-instead-of-pk-sk.rs");`
`10`	`11`
`11`	`12`	`t.pass("tests/ui/pass.rs");`
`12`	`13`	`}`