refactor: update fetchfromdatabase and associated APIs to new Entity trait family [WPB-22195]

crypto-macros/src/entity_derive/derive_impl.rs

@@ -38,7 +38,7 @@ impl KeyStoreEntityFlattened {
                 }
 
                 fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-                    crate::transaction::dynamic_dispatch::Entity::#struct_name(self)
+                    crate::transaction::dynamic_dispatch::Entity::#struct_name(self.into())
                 }
             }
         }

crypto-macros/src/entity_derive_new/column.rs

@@ -57,7 +57,11 @@ where
     pub(super) fn load_expression(&self) -> TokenStream {
         let column_name = self.sql_name();
 
-        let expr = quote!(row.get::<_, Vec<u8>>(#column_name)?);
+        let sql_data_type = match self.transformation {
+            None => self.column_type.get_as_type(),
+            Some(FieldTransformation::Hex) => quote!(String),
+        };
+        let expr = quote!(row.get::<_, #sql_data_type>(#column_name)?);
 
         let expr = match self.transformation {
             None => expr,

crypto-macros/src/entity_derive_new/column_type.rs

@@ -133,6 +133,9 @@ impl TryFrom<Type> for ColumnType {
 pub(super) trait EmitGetExpression {
     /// Emit an expression which wraps the input expression, appropriately parsing according to this column type.
     fn emit_get_expression(&self, input: TokenStream) -> TokenStream;
+
+    /// Emit an expression with the rust type which should be used in the rusqlite `get` expression
+    fn get_as_type(&self) -> TokenStream;
 }
 
 impl EmitGetExpression for IdColumnType {
@@ -142,17 +145,24 @@ impl EmitGetExpression for IdColumnType {
             Self::String => quote!(String::from_utf8(#input).map_err(|err| err.utf8_error())?),
         }
     }
+
+    fn get_as_type(&self) -> TokenStream {
+        quote!(Vec<u8>)
+    }
 }
 
 impl EmitGetExpression for ColumnType {
     fn emit_get_expression(&self, input: TokenStream) -> TokenStream {
         match self {
-            ColumnType::Bytes => input,
+            ColumnType::Bytes | ColumnType::OptionalBytes => input,
             ColumnType::String => quote!(String::from_utf8(#input).map_err(|err| err.utf8_error())?),
-            ColumnType::OptionalBytes => quote! {{
-                let data = #input;
-                (!data.is_empty()).then_some(data)
-            }},
+        }
+    }
+
+    fn get_as_type(&self) -> TokenStream {
+        match self {
+            ColumnType::Bytes | ColumnType::String => quote!(Vec<u8>),
+            ColumnType::OptionalBytes => quote!(Option<Vec<u8>>),
         }
     }
 }

crypto-macros/src/entity_derive_new/derive_impl.rs

@@ -3,14 +3,15 @@ use proc_macro2::{Span, TokenStream};
 use quote::quote;
 use syn::{Ident, Lifetime};
 
-use crate::entity_derive_new::{Entity, column_type::ColumnType};
+use crate::entity_derive_new::{Entity, FieldTransformation, column_type::ColumnType};
 
 impl quote::ToTokens for Entity {
     fn to_tokens(&self, tokens: &mut TokenStream) {
         tokens.extend(self.impl_entity_base());
+        tokens.extend(self.impl_primary_key());
         tokens.extend(self.impl_entity_generic());
         tokens.extend(self.impl_entity_wasm());
-        tokens.extend(self.impl_borrow_primary_key());
+        tokens.extend(self.impl_entity_get_borrowed());
         tokens.extend(self.impl_entity_database_mutation());
         tokens.extend(self.impl_entity_delete_borrowed());
         tokens.extend(self.impl_decrypting());
@@ -37,7 +38,36 @@ impl Entity {
                 const COLLECTION_NAME: &'static str = #collection_name;
 
                 fn to_transaction_entity(self) -> crate::transaction::dynamic_dispatch::Entity {
-                    crate::transaction::dynamic_dispatch::Entity::#struct_name(self)
+                    crate::transaction::dynamic_dispatch::Entity::#struct_name(self.into())
+                }
+            }
+        }
+    }
+
+    /// `impl PrimaryKey for MyEntity` and `impl BorrowPrimaryKey for MyEntity`
+    fn impl_primary_key(&self) -> TokenStream {
+        let Self {
+            struct_name, id_column, ..
+        } = self;
+
+        let primary_key = id_column.column_type.owned();
+        let borrowed_primary_key = id_column.column_type.borrowed();
+        let pk_field_name = &id_column.field_name;
+
+        quote! {
+            impl crate::traits::PrimaryKey for #struct_name {
+                type PrimaryKey = #primary_key;
+
+                fn primary_key(&self) -> Self::PrimaryKey {
+                    self.#pk_field_name.clone()
+                }
+            }
+
+            impl crate::traits::BorrowPrimaryKey for #struct_name {
+                type BorrowedPrimaryKey = #borrowed_primary_key;
+
+                fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
+                    &self.#pk_field_name
                 }
             }
         }
@@ -52,24 +82,15 @@ impl Entity {
             ..
         } = self;
 
-        let primary_key = id_column.column_type.owned();
-        let id_field_name = &id_column.field_name;
-
         let field_assignments = std::iter::once(id_column.field_assignment())
             .chain(other_columns.iter().map(|column| column.field_assignment()));
 
         quote! {
             #[cfg(not(target_family = "wasm"))]
             #[::async_trait::async_trait]
             impl crate::traits::Entity for #struct_name {
-                type PrimaryKey = #primary_key;
-
-                fn primary_key(&self) -> #primary_key {
-                    self.#id_field_name.clone()
-                }
-
                 async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> crate::CryptoKeystoreResult<Option<Self>> {
-                    <Self as crate::traits::BorrowPrimaryKey>::get_borrowed(conn, key).await
+                    <Self as crate::traits::EntityGetBorrowed>::get_borrowed(conn, key).await
                 }
 
                 async fn count(conn: &mut Self::ConnectionType) -> crate::CryptoKeystoreResult<u32> {
@@ -89,25 +110,14 @@ impl Entity {
 
     /// `#[cfg(target_family = "wasm")] impl Entity for MyEntity`
     fn impl_entity_wasm(&self) -> TokenStream {
-        let Self {
-            struct_name, id_column, ..
-        } = self;
-
-        let primary_key = id_column.column_type.owned();
-        let id_field_name = &id_column.field_name;
+        let Self { struct_name, .. } = self;
 
         quote! {
             #[cfg(target_family = "wasm")]
             #[::async_trait::async_trait(?Send)]
             impl crate::traits::Entity for #struct_name {
-                type PrimaryKey = #primary_key;
-
-                fn primary_key(&self) -> #primary_key {
-                    self.#id_field_name.clone()
-                }
-
                 async fn get(conn: &mut Self::ConnectionType, key: &Self::PrimaryKey) -> crate::CryptoKeystoreResult<Option<Self>> {
-                    <Self as crate::traits::BorrowPrimaryKey>::get_borrowed(conn, key).await
+                    <Self as crate::traits::EntityGetBorrowed>::get_borrowed(conn, key).await
                 }
 
                 async fn count(conn: &mut Self::ConnectionType) -> crate::CryptoKeystoreResult<u32> {
@@ -121,17 +131,15 @@ impl Entity {
         }
     }
 
-    /// `impl BorrowPrimaryKey for MyEntity`
-    fn impl_borrow_primary_key(&self) -> TokenStream {
+    /// `impl EntityGetBorrowed for MyEntity`
+    fn impl_entity_get_borrowed(&self) -> TokenStream {
         let Self {
             struct_name,
             id_column,
             other_columns,
             ..
         } = self;
 
-        let borrowed_primary_key = id_column.column_type.borrowed();
-        let pk_field_name = &id_column.field_name;
         let pk_column_name = id_column
             .column_name
             .clone()
@@ -140,29 +148,31 @@ impl Entity {
         let field_assignments = std::iter::once(id_column.field_assignment())
             .chain(other_columns.iter().map(|column| column.field_assignment()));
 
+        // if we ever add a second field transformation, we'll want this match pattern
+        #[allow(clippy::manual_map)]
+        let key_transform = match id_column.transformation {
+            None => None,
+            Some(FieldTransformation::Hex) => Some(quote! {let key = hex::encode(key);}),
+        };
+
         quote! {
             #[cfg_attr(target_family = "wasm", ::async_trait::async_trait(?Send))]
             #[cfg_attr(not(target_family = "wasm"), ::async_trait::async_trait)]
-            impl crate::traits::BorrowPrimaryKey for #struct_name {
-                type BorrowedPrimaryKey = #borrowed_primary_key;
-
-                fn borrow_primary_key(&self) -> &Self::BorrowedPrimaryKey {
-                    &self.#pk_field_name
-                }
-
-                async fn get_borrowed<Q>(conn: &mut Self::ConnectionType, key: &Q) -> crate::CryptoKeystoreResult<Option<Self>>
-                where
-                    Self::PrimaryKey: std::borrow::Borrow<Q>,
-                    Q: crate::traits::KeyType,
+            impl crate::traits::EntityGetBorrowed for #struct_name {
+                async fn get_borrowed(conn: &mut Self::ConnectionType, key: &Self::BorrowedPrimaryKey)
+                    -> crate::CryptoKeystoreResult<Option<Self>>
                 {
+                    let key = <&Self::BorrowedPrimaryKey as crate::traits::KeyType>::bytes(&key);
+                    let key = key.as_ref();
+                    #key_transform
                     #[cfg(target_family = "wasm")]
                     {
-                        conn.storage().new_get(key.bytes().as_ref()).await
+                        conn.storage().new_get(key).await
                     }
 
                     #[cfg(not(target_family = "wasm"))]
                     {
-                        crate::entities::platform::get_helper::<Self, _>(conn, #pk_column_name, key.bytes().as_ref(), |row| {
+                        crate::entities::platform::get_helper::<Self, _>(conn, #pk_column_name, key, |row| {
                             Ok(Self {
                                 #( #field_assignments, )*
                             })
@@ -196,6 +206,10 @@ impl Entity {
             .map(|tokens| quote!(#tokens,))
             .collect::<TokenStream>();
 
+        let sql_map_err = (!upsert).then_some(quote! {
+            .map_err(|_| CryptoKeystoreError::AlreadyExists(Self::COLLECTION_NAME))
+        });
+
         quote! {
             #[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
             #[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
@@ -211,7 +225,7 @@ impl Entity {
                         #[cfg(not(target_family = "wasm"))]
                         {
                             let mut stmt = tx.prepare_cached(#sql_statement)?;
-                            stmt.execute(rusqlite::params![#fields])?;
+                            stmt.execute(rusqlite::params![#fields])#sql_map_err?;
                             Ok(())
                         }
                 }
@@ -241,27 +255,35 @@ impl Entity {
         } = self;
 
         let id_column_name = id_column.sql_name();
+        // if we ever add a second field transformation, we'll want this match pattern
+        #[allow(clippy::manual_map)]
+        let key_transform = match id_column.transformation {
+            None => None,
+            Some(FieldTransformation::Hex) => Some(quote! {let key = hex::encode(key);}),
+        };
 
         quote! {
             #[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
             #[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
             impl<'a> crate::traits::EntityDeleteBorrowed<'a> for #struct_name {
-                async fn delete_borrowed<Q>(
+                async fn delete_borrowed(
                     tx: &<Self as crate::traits::EntityDatabaseMutation<'a>>::Transaction,
-                    id: &Q,
+                    id: &<Self as crate::traits::BorrowPrimaryKey>::BorrowedPrimaryKey,
                 ) -> crate::CryptoKeystoreResult<bool>
                 where
-                    Self::PrimaryKey: std::borrow::Borrow<Q>,
-                    Q: crate::traits::KeyType
+                    for<'pk> &'pk <Self as crate::traits::BorrowPrimaryKey>::BorrowedPrimaryKey: crate::traits::KeyType,
                 {
+                    let key = <&<Self as crate::traits::BorrowPrimaryKey>::BorrowedPrimaryKey as crate::traits::KeyType>::bytes(&id);
+                    let key = key.as_ref();
+                    #key_transform
                     #[cfg(target_family = "wasm")]
                     {
-                        tx.new_delete::<Self>(id.bytes().as_ref()).await
+                        tx.new_delete::<Self>(key).await
                     }
 
                     #[cfg(not(target_family = "wasm"))]
                     {
-                        crate::entities::platform::delete_helper::<Self>(tx, #id_column_name, id.bytes().as_ref()).await
+                        crate::entities::platform::delete_helper::<Self>(tx, #id_column_name, key).await
                     }
                 }
             }

crypto/src/e2e_identity/pki_env.rs

@@ -1,8 +1,8 @@
 use std::collections::HashSet;
 
 use core_crypto_keystore::{
-    connection::FetchFromDatabase,
     entities::{E2eiAcmeCA, E2eiCrl, E2eiIntermediateCert},
+    traits::FetchFromDatabase,
 };
 use wire_e2e_identity::prelude::x509::revocation::{PkiEnvironment, PkiEnvironmentParams};
 use x509_cert::der::Decode;
@@ -33,7 +33,7 @@ impl IntoIterator for NewCrlDistributionPoints {
 
 pub(crate) async fn restore_pki_env(data_provider: &impl FetchFromDatabase) -> Result<Option<PkiEnvironment>> {
     let mut trust_roots = vec![];
-    let Ok(ta_raw) = data_provider.find_unique::<E2eiAcmeCA>().await else {
+    let Ok(Some(ta_raw)) = data_provider.get_unique::<E2eiAcmeCA>().await else {
         return Ok(None);
     };
 
@@ -42,15 +42,15 @@ pub(crate) async fn restore_pki_env(data_provider: &impl FetchFromDatabase) -> R
     );
 
     let intermediates = data_provider
-        .find_all::<E2eiIntermediateCert>(Default::default())
+        .load_all::<E2eiIntermediateCert>()
         .await
         .map_err(KeystoreError::wrap("finding intermediate certificates"))?
         .into_iter()
         .map(|inter| x509_cert::Certificate::from_der(&inter.content))
         .collect::<Result<Vec<_>, _>>()?;
 
     let crls = data_provider
-        .find_all::<E2eiCrl>(Default::default())
+        .load_all::<E2eiCrl>()
         .await
         .map_err(KeystoreError::wrap("finding crls"))?
         .into_iter()

crypto/src/group_store.rs

@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use core_crypto_keystore::connection::FetchFromDatabase;
+use core_crypto_keystore::traits::FetchFromDatabase;
 
 use crate::{ConversationId, KeystoreError, MlsConversation, RecursiveError, Result};
 
@@ -31,7 +31,7 @@ impl GroupStoreEntity for MlsConversation {
         keystore: &impl FetchFromDatabase,
     ) -> crate::Result<Option<Self>> {
         let result = keystore
-            .find::<Self::RawStoreValue>(id)
+            .get_borrowed::<Self::RawStoreValue>(id.as_ref())
             .await
             .map_err(KeystoreError::wrap("finding mls conversation from keystore by id"))?;
         let Some(store_value) = result else {

crypto/src/mls/conversation/conversation_guard/decrypt/buffer_commit.rs

@@ -1,4 +1,4 @@
-use core_crypto_keystore::{connection::FetchFromDatabase as _, entities::StoredBufferedCommit};
+use core_crypto_keystore::{entities::StoredBufferedCommit, traits::FetchFromDatabase as _};
 use log::info;
 use openmls::framing::MlsMessageIn;
 use openmls_traits::OpenMlsCryptoProvider as _;
@@ -34,7 +34,7 @@ impl ConversationGuard {
         self.crypto_provider()
             .await?
             .keystore()
-            .find::<StoredBufferedCommit>(conversation.id())
+            .get_borrowed::<StoredBufferedCommit>(conversation.id().as_ref())
             .await
             .map(|option| option.map(StoredBufferedCommit::into_commit_data))
             .map_err(KeystoreError::wrap("attempting to retrieve buffered commit"))
@@ -69,7 +69,7 @@ impl ConversationGuard {
         self.crypto_provider()
             .await?
             .keystore()
-            .remove::<StoredBufferedCommit, _>(conversation.id())
+            .remove_borrowed::<StoredBufferedCommit>(conversation.id().as_ref())
             .await
             .map_err(KeystoreError::wrap("attempting to clear buffered commit"))
             .map_err(Into::into)

crypto/src/mls/conversation/merge.rs

@@ -33,7 +33,7 @@ impl MlsConversation {
         // ..so if there's any, we clear them after the commit is merged
         for oln in &previous_own_leaf_nodes {
             let ek = oln.encryption_key().as_slice();
-            let _ = backend.key_store().remove::<StoredEncryptionKeyPair, _>(ek).await;
+            let _ = backend.key_store().remove_borrowed::<StoredEncryptionKeyPair>(ek).await;
         }
 
         client