Add a direct API for building queries

Author: Bennett Hardwick

src/encrypted_table/mod.rs

@@ -14,7 +14,11 @@ use crate::{
 use aws_sdk_dynamodb::types::{AttributeValue, Delete, Put, TransactWriteItem};
 use cipherstash_client::{
     config::{console_config::ConsoleConfig, zero_kms_config::ZeroKMSConfig},
-    credentials::{auto_refresh::AutoRefresh, service_credentials::ServiceCredentials},
+    credentials::{
+        auto_refresh::AutoRefresh,
+        service_credentials::{ServiceCredentials, ServiceToken},
+        Credentials,
+    },
     encryption::Encryption,
     zero_kms::ZeroKMS,
 };
@@ -46,9 +50,9 @@ pub struct EncryptedTable<D = Dynamo> {
 }
 
 impl<D> EncryptedTable<D> {
-    // option here to generate query params
-
-    // option here to seal
+    pub fn cipher(&self) -> &Encryption<impl Credentials<Token = ServiceToken>> {
+        self.cipher.as_ref()
+    }
 }
 
 impl EncryptedTable<Headless> {
@@ -209,11 +213,11 @@ impl DynamoRecordPatch {
 }
 
 impl<D> EncryptedTable<D> {
-    pub fn query<S>(&self) -> QueryBuilder<S, D>
+    pub fn query<S>(&self) -> QueryBuilder<S, &Self>
     where
         S: Searchable,
     {
-        QueryBuilder::new(self)
+        QueryBuilder::with_backend(self)
     }
 
     pub async fn unseal_all(

src/encrypted_table/query.rs

@@ -23,10 +23,10 @@ use super::{Dynamo, EncryptedTable, QueryError};
  * D = Database
  *
  */
-pub struct QueryBuilder<'t, S, D = Dynamo> {
+pub struct QueryBuilder<S, B = ()> {
     parts: Vec<(String, SingleIndex, Plaintext)>,
-    table: &'t EncryptedTable<D>,
-    __table: PhantomData<S>,
+    backend: B,
+    __searchable: PhantomData<S>,
 }
 
 pub struct PreparedQuery {
@@ -94,15 +94,22 @@ impl PreparedQuery {
     }
 }
 
-impl<'t, S, D> QueryBuilder<'t, S, D>
-where
-    S: Searchable,
-{
-    pub fn new(table: &'t EncryptedTable<D>) -> Self {
+impl<S> QueryBuilder<S> {
+    pub fn new() -> Self {
         Self {
             parts: vec![],
-            table,
-            __table: Default::default(),
+            backend: Default::default(),
+            __searchable: Default::default(),
+        }
+    }
+}
+
+impl<S, B> QueryBuilder<S, B> {
+    pub fn with_backend(backend: B) -> Self {
+        Self {
+            parts: vec![],
+            backend,
+            __searchable: Default::default(),
         }
     }
 
@@ -117,8 +124,13 @@ where
             .push((name.into(), SingleIndex::Prefix, plaintext.into()));
         self
     }
+}
 
-    fn build(self) -> Result<PreparedQuery, QueryError> {
+impl<S, B> QueryBuilder<S, B>
+where
+    S: Searchable,
+{
+    pub fn build(self) -> Result<PreparedQuery, QueryError> {
         let items_len = self.parts.len();
 
         // this is the simplest way to brute force the index names but relies on some gross
@@ -184,12 +196,12 @@ where
     }
 }
 
-impl<'t, S> QueryBuilder<'t, S, Dynamo>
+impl<S> QueryBuilder<S, &EncryptedTable<Dynamo>>
 where
     S: Searchable + Identifiable,
 {
     pub async fn load<T: Decryptable>(self) -> Result<Vec<T>, QueryError> {
-        let table = self.table;
+        let table = self.backend;
         let query = self.build()?;
 
         let items = query.send(table).await?;
@@ -199,7 +211,7 @@ where
     }
 }
 
-impl<'t, S> QueryBuilder<'t, S, Dynamo>
+impl<S> QueryBuilder<S, &EncryptedTable<Dynamo>>
 where
     S: Searchable + Decryptable + Identifiable,
 {

tests/query_builder_direct.rs

@@ -0,0 +1,215 @@
+use cipherstash_dynamodb::{
+    Decryptable, Encryptable, EncryptedTable, Identifiable, QueryBuilder, Searchable,
+};
+use itertools::Itertools;
+use serial_test::serial;
+use std::future::Future;
+
+mod common;
+
+#[derive(
+    Identifiable, Encryptable, Decryptable, Searchable, Debug, PartialEq, Ord, PartialOrd, Eq,
+)]
+#[cipherstash(sort_key_prefix = "user")]
+pub struct User {
+    #[cipherstash(query = "exact", compound = "email#name")]
+    #[cipherstash(query = "exact")]
+    #[partition_key]
+    pub email: String,
+
+    #[cipherstash(query = "prefix", compound = "email#name")]
+    #[cipherstash(query = "prefix")]
+    pub name: String,
+
+    #[cipherstash(plaintext)]
+    pub tag: String,
+
+    #[cipherstash(skip)]
+    pub temp: bool,
+}
+
+impl User {
+    pub fn new(email: impl Into<String>, name: impl Into<String>, tag: impl Into<String>) -> Self {
+        Self {
+            name: name.into(),
+            email: email.into(),
+            tag: tag.into(),
+            temp: false,
+        }
+    }
+}
+
+async fn run_test<F: Future<Output = ()>>(
+    mut f: impl FnMut(aws_sdk_dynamodb::Client, String) -> F,
+) {
+    let config = aws_config::from_env()
+        .endpoint_url("http://localhost:8000")
+        .load()
+        .await;
+
+    let client = aws_sdk_dynamodb::Client::new(&config);
+
+    let table_name = "test-users-direct-query-builder";
+
+    common::create_table(&client, table_name).await;
+
+    let table = EncryptedTable::init(client.clone(), table_name)
+        .await
+        .expect("Failed to init table");
+
+    table
+        .put(User::new("[email protected]", "Dan Draper", "blue"))
+        .await
+        .expect("Failed to insert Dan");
+
+    table
+        .put(User::new("[email protected]", "Jane Smith", "red"))
+        .await
+        .expect("Failed to insert Jane");
+
+    table
+        .put(User::new("[email protected]", "Daniel Johnson", "green"))
+        .await
+        .expect("Failed to insert Daniel");
+
+    f(client, table_name.to_string()).await;
+}
+
+#[tokio::test]
+#[serial]
+async fn test_query_single_exact() {
+    run_test(|client, name| async move {
+        let table = EncryptedTable::init_headless()
+            .await
+            .expect("failed to init table");
+
+        let query = QueryBuilder::<User>::new()
+            .eq("email", "[email protected]")
+            .build()
+            .expect("failed to build query");
+
+        let term = query
+            .encrypt(table.cipher())
+            .await
+            .expect("failed to encrypt query");
+
+        let query = client
+            .query()
+            .table_name(name)
+            .index_name("TermIndex")
+            .key_condition_expression("term = :term")
+            .expression_attribute_values(":term", term);
+
+        let result = query.send().await.expect("failed to send");
+
+        let items = result.items.unwrap();
+
+        let res: Vec<User> = table
+            .decrypt_all(items)
+            .await
+            .expect("failed to decrypt")
+            .into_iter()
+            .sorted()
+            .collect_vec();
+
+        assert_eq!(
+            res,
+            vec![User::new("[email protected]", "Dan Draper", "blue")]
+        );
+    })
+    .await;
+}
+
+#[tokio::test]
+#[serial]
+async fn test_query_single_prefix() {
+    run_test(|client, name| async move {
+        let table = EncryptedTable::init_headless()
+            .await
+            .expect("failed to init table");
+
+        let query = QueryBuilder::<User>::new()
+            .starts_with("name", "Dan")
+            .build()
+            .expect("failed to build query");
+
+        let term = query
+            .encrypt(table.cipher())
+            .await
+            .expect("failed to encrypt query");
+
+        let query = client
+            .query()
+            .table_name(name)
+            .index_name("TermIndex")
+            .key_condition_expression("term = :term")
+            .expression_attribute_values(":term", term);
+
+        let result = query.send().await.expect("failed to send");
+
+        let items = result.items.unwrap();
+
+        let res: Vec<User> = table
+            .decrypt_all(items)
+            .await
+            .expect("failed to decrypt")
+            .into_iter()
+            .sorted()
+            .collect_vec();
+
+        assert_eq!(
+            res,
+            vec![
+                User::new("[email protected]", "Dan Draper", "blue"),
+                User::new("[email protected]", "Daniel Johnson", "green")
+            ]
+        );
+    })
+    .await;
+}
+
+#[tokio::test]
+#[serial]
+async fn test_query_compound() {
+    run_test(|client, name| async move {
+        let table = EncryptedTable::init_headless()
+            .await
+            .expect("failed to init table");
+
+        let query = QueryBuilder::<User>::new()
+            .starts_with("name", "Dan")
+            .eq("email", "[email protected]")
+            .build()
+            .expect("failed to build query");
+
+        let term = query
+            .encrypt(table.cipher())
+            .await
+            .expect("failed to encrypt query");
+
+        let query = client
+            .query()
+            .table_name(name)
+            .index_name("TermIndex")
+            .key_condition_expression("term = :term")
+            .expression_attribute_values(":term", term);
+
+        let result = query.send().await.expect("failed to send");
+
+        let items = result.items.unwrap();
+
+        let res: Vec<User> = table
+            .decrypt_all(items)
+            .await
+            .expect("failed to decrypt")
+            .into_iter()
+            .sorted()
+            .collect_vec();
+
+        assert_eq!(
+            res,
+            vec![User::new("[email protected]", "Dan Draper", "blue")]
+        );
+    })
+    .await;
+}