Refactor comments and update beacon configurations

DynamoDbEncryption/runtimes/rust/examples/searchableencryption/basic_searchable_encryption.rs

@@ -72,10 +72,11 @@ pub async fn put_and_query_with_beacon(branch_key_id: &str) -> Result<(), crate:
         //     values are uniformly distributed across its range of possible values.
         //     In many use cases, the prefix of an identifier encodes some information
         //     about that identifier (e.g. zipcode and SSN prefixes encode geographic
-        //     information), while the suffix does not and is more uniformly distributed.
+        //     information), while the suffix does not encode such structure and tends 
+        //     to be closer to uniformly distributed, though not perfectly uniform.
         //     We will assume that the inspector ID field matches a similar use case.
-        //     So for this example, we only store and use the last
-        //     4 digits of the inspector ID, which we assume is uniformly distributed.
+        //     For this example, we use only the last four digits of the inspector ID and apply one partitions, 
+        //     assuming the suffix has sufficient uniformity for this purpose.
         // Since the full ID's range is divisible by the range of the last 4 digits,
         //     then the last 4 digits of the inspector ID are uniformly distributed
         //     over the range from 0 to 9,999.
@@ -114,12 +115,14 @@ pub async fn put_and_query_with_beacon(branch_key_id: &str) -> Result<(), crate:
         let last4_beacon = StandardBeacon::builder()
             .name("inspector_id_last4")
             .length(10)
+            .number_of_partitions(1)
             .build()?;
 
         // The configured DDB table has a GSI on the `aws_dbe_b_unit` AttributeName.
         // This field holds a unit serial number.
         // For this example, this is a 12-digit integer from 0 to 999,999,999,999 (10^12 possible values).
-        // We will assume values for this attribute are uniformly distributed across this range.
+        // We will assume values for this attribute are uniformly distributed across this range
+        // using one partitions.
         // A single unit serial number may be assigned to multiple `work_id`s.
         //
         // This link provides guidance for choosing a beacon length:
@@ -142,7 +145,7 @@ pub async fn put_and_query_with_beacon(branch_key_id: &str) -> Result<(), crate:
         // With a sufficiently large number of well-distributed inspector IDs,
         //    for a particular beacon we expect (10^12/2^30) ~= 931.3 unit serial numbers
         //    sharing that beacon value.
-        let unit_beacon = StandardBeacon::builder().name("unit").length(30).build()?;
+        let unit_beacon = StandardBeacon::builder().name("unit").length(30).number_of_partitions(1).build()?;
 
         let standard_beacon_list = vec![last4_beacon, unit_beacon];
 
@@ -169,6 +172,8 @@ pub async fn put_and_query_with_beacon(branch_key_id: &str) -> Result<(), crate:
         // 3. Create BeaconVersion.
         //    The BeaconVersion inside the list holds the list of beacons on the table.
         //    The BeaconVersion also stores information about the keystore.
+        //    The BeaconVersion parameter specifies the maximumNumberOfPartitions associated with a given beacon configuration. 
+        //    "maximumNumberOfPartitions" must be greater than "numberOfPartitions"; we set it to 8 to allow room for future expansion.
         //    BeaconVersion must be provided:
         //      - keyStore: The keystore configured in step 2.
         //      - keySource: A configuration for the key source.
@@ -182,6 +187,8 @@ pub async fn put_and_query_with_beacon(branch_key_id: &str) -> Result<(), crate:
         let beacon_version = BeaconVersion::builder()
             .standard_beacons(standard_beacon_list)
             .version(1) // MUST be 1
+            .maximum_number_of_partitions(8)
+            .default_number_of_partitions(1) //For beacons that do not require partitioning, only a single partition is used.
             .key_store(key_store.clone())
             .key_source(BeaconKeySource::Single(
                 SingleKeyStore::builder()
@@ -305,19 +312,29 @@ pub async fn put_and_query_with_beacon(branch_key_id: &str) -> Result<(), crate:
         //     This procedure is internal to the client and is abstracted away from the user;
         //     e.g. the user will only see "123456789012" and never
         //        "098765432109", though the actual query returned both.
+
+        let mut all_results: Vec<HashMap<String, AttributeValue>> = Vec::new();
+
         let expression_attributes_names = HashMap::from([
             ("#last4".to_string(), "inspector_id_last4".to_string()),
             ("#unit".to_string(), "unit".to_string()),
         ]);
 
+      // In this simple example, we know there are one buckets, 
+      // but in general the number can be obtained using transformClient.getNumberOfQueries(query)
+      let num_queries = 1; 
+
+     //We need to query for all possible parttions 
+
+       for partition in 0..num_queries {
+
         let expression_attribute_values = HashMap::from([
             (":last4".to_string(), AttributeValue::S("4321".to_string())),
-            (
-                ":unit".to_string(),
-                AttributeValue::S("123456789012".to_string()),
-            ),
+            (":unit".to_string(), AttributeValue::S("123456789012".to_string())),
+            ( ":aws_dbe_partition".to_string(), AttributeValue::N(partition.to_string())),
         ]);
 
+
         // GSIs do not update instantly
         // so if the results come back empty
         // we retry after a short sleep
@@ -332,27 +349,31 @@ pub async fn put_and_query_with_beacon(branch_key_id: &str) -> Result<(), crate:
                 .send()
                 .await?;
 
-            // if no results, sleep and try again
-            if query_response.items.is_none() || query_response.items.as_ref().unwrap().is_empty() {
+
+                 if let Some(items) = query_response.items {
+                   if !items.is_empty() {
+                      all_results.extend(items.clone());
+                      break;
+                    }
+                 }
+
                 std::thread::sleep(std::time::Duration::from_millis(20));
-                continue;
-            }
-
-            let attribute_values = query_response.items.unwrap();
-            // Validate only 1 item was returned: the item we just put
-            assert_eq!(attribute_values.len(), 1);
-            let returned_item = &attribute_values[0];
-            // Validate the item has the expected attributes
-            assert_eq!(
-                returned_item["inspector_id_last4"],
-                AttributeValue::S("4321".to_string())
-            );
-            assert_eq!(
-                returned_item["unit"],
-                AttributeValue::S("123456789012".to_string())
-            );
-            break;
+         }
         }
+        let attribute_values = all_results;
+        // Validate only 1 item was returned: the item we just put
+        assert_eq!(attribute_values.len(), 1);
+        let returned_item = &attribute_values[0];
+        // Validate the item has the expected attributes
+        assert_eq!(
+           returned_item["inspector_id_last4"],
+           AttributeValue::S("4321".to_string())
+        );
+        assert_eq!(
+            returned_item["unit"],
+            AttributeValue::S("123456789012".to_string())
+        );
+
         println!("basic_searchable_encryption successful.");
         Ok(())
     };