Expanded multi-dataset tests

Author: coderdan

src/encrypted_table/mod.rs

@@ -288,9 +288,7 @@ impl<D> EncryptedTable<D> {
         delete: PreparedDelete,
         dataset_id: Option<Uuid>,
     ) -> Result<DynamoRecordPatch, DeleteError> {
-        let scoped_cipher = ScopedZeroKmsCipher::init(self.cipher.clone(), dataset_id)
-            .await
-            .unwrap();
+        let scoped_cipher = ScopedZeroKmsCipher::init(self.cipher.clone(), dataset_id).await?;
 
         let PrimaryKeyParts { pk, sk } =
             encrypt_primary_key_parts(&scoped_cipher, delete.primary_key)?;

src/encrypted_table/query.rs

@@ -20,6 +20,7 @@ use super::{Dynamo, EncryptedTable, QueryError, ScopedZeroKmsCipher, SealError};
 pub struct QueryBuilder<S, B = ()> {
     parts: Vec<(String, SingleIndex, Plaintext)>,
     storage: B,
+    dataset_id: Option<Uuid>,
     __searchable: PhantomData<S>,
 }
 
@@ -93,6 +94,7 @@ impl<S> Default for QueryBuilder<S> {
         Self {
             parts: vec![],
             storage: Default::default(),
+            dataset_id: None,
             __searchable: Default::default(),
         }
     }
@@ -103,10 +105,17 @@ impl<S, B> QueryBuilder<S, B> {
         Self {
             parts: vec![],
             storage: backend,
+            dataset_id: None,
             __searchable: Default::default(),
         }
     }
 
+    /// Specify the dataset to query against.
+    pub fn via(mut self, dataset_id: Uuid) -> Self {
+        self.dataset_id = Some(dataset_id);
+        self
+    }
+
     pub fn eq(mut self, name: impl Into<String>, plaintext: impl Into<Plaintext>) -> Self {
         self.parts
             .push((name.into(), SingleIndex::Exact, plaintext.into()));
@@ -138,27 +147,12 @@ where
     ///
     /// While a client can decrypt records from any dataset it has access to,
     /// queries are always scoped to a single dataset.
-    pub async fn load<T>(self) -> Result<Vec<T>, QueryError>
-    where
-        T: Decryptable + Identifiable,
-    {
-        self.load_inner(None).await
-    }
-
-    /// Similar to `load`, but the query is scoped to a specific dataset.
-    pub async fn load_via<T>(self, dataset_id: Uuid) -> Result<Vec<T>, QueryError>
-    where
-        T: Decryptable + Identifiable,
-    {
-        self.load_inner(Some(dataset_id)).await
-    }
-
-    async fn load_inner<T>(self, dataset_id: Option<Uuid>) -> Result<Vec<T>, QueryError>
+    pub(crate) async fn load<T>(self) -> Result<Vec<T>, QueryError>
     where
         T: Decryptable + Identifiable,
     {
         let scoped_cipher =
-            ScopedZeroKmsCipher::init(self.storage.cipher.clone(), dataset_id).await?;
+            ScopedZeroKmsCipher::init(self.storage.cipher.clone(), self.dataset_id).await?;
 
         let storage = self.storage;
         let query = self.build()?;

src/errors/mod.rs

@@ -64,6 +64,9 @@ pub enum DeleteError {
     AwsBuildError(#[from] BuildError),
     #[error("AwsError: {0}")]
     Aws(String),
+
+    #[error("ZeroKMS Error: {0}")]
+    ZeroKMS(#[from] zerokms::Error),
 }
 
 /// Error returned by `EncryptedTable::query` when indexing, retrieving and decrypting records from DynamoDB

tests/common.rs

@@ -64,10 +64,10 @@ pub fn fail_not_found() -> CheckFailed {
 /// The name is used as a prefix in case its helpful to distinguish between tests.
 /// A random is appended to the name to ensure uniqueness for async tests.
 #[allow(dead_code)]
-pub async fn with_encrypted_table<F: Future<Output = miette::Result<()>>>(
+pub async fn with_encrypted_table<O, F: Future<Output = miette::Result<O>>>(
     table_name: &str,
     mut f: impl FnMut(EncryptedTable) -> F,
-) -> Result<(), Box<dyn std::error::Error>> {
+) -> Result<O, Box<dyn std::error::Error>> {
     let config = aws_config::from_env()
         .endpoint_url("http://localhost:8000")
         .load()