@@ -72,10 +72,11 @@ pub async fn put_and_query_with_beacon(branch_key_id: &str) -> Result<(), crate:
7272 // values are uniformly distributed across its range of possible values.
7373 // In many use cases, the prefix of an identifier encodes some information
7474 // about that identifier (e.g. zipcode and SSN prefixes encode geographic
75- // information), while the suffix does not and is more uniformly distributed.
75+ // information), while the suffix does not encode such structure and tends
76+ // to be closer to uniformly distributed, though not perfectly uniform.
7677 // We will assume that the inspector ID field matches a similar use case.
77- // So for this example, we only store and use the last
78- // 4 digits of the inspector ID, which we assume is uniformly distributed .
78+ // For this example, we use only the last four digits of the inspector ID and apply one partitions,
79+ // assuming the suffix has sufficient uniformity for this purpose .
7980 // Since the full ID's range is divisible by the range of the last 4 digits,
8081 // then the last 4 digits of the inspector ID are uniformly distributed
8182 // 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:
114115 let last4_beacon = StandardBeacon :: builder ( )
115116 . name ( "inspector_id_last4" )
116117 . length ( 10 )
118+ . number_of_partitions ( 1 )
117119 . build ( ) ?;
118120
119121 // The configured DDB table has a GSI on the `aws_dbe_b_unit` AttributeName.
120122 // This field holds a unit serial number.
121123 // For this example, this is a 12-digit integer from 0 to 999,999,999,999 (10^12 possible values).
122- // We will assume values for this attribute are uniformly distributed across this range.
124+ // We will assume values for this attribute are uniformly distributed across this range
125+ // using one partitions.
123126 // A single unit serial number may be assigned to multiple `work_id`s.
124127 //
125128 // This link provides guidance for choosing a beacon length:
@@ -142,7 +145,11 @@ pub async fn put_and_query_with_beacon(branch_key_id: &str) -> Result<(), crate:
142145 // With a sufficiently large number of well-distributed inspector IDs,
143146 // for a particular beacon we expect (10^12/2^30) ~= 931.3 unit serial numbers
144147 // sharing that beacon value.
145- let unit_beacon = StandardBeacon :: builder ( ) . name ( "unit" ) . length ( 30 ) . build ( ) ?;
148+ let unit_beacon = StandardBeacon :: builder ( )
149+ . name ( "unit" )
150+ . length ( 30 )
151+ . number_of_partitions ( 1 )
152+ . build ( ) ?;
146153
147154 let standard_beacon_list = vec ! [ last4_beacon, unit_beacon] ;
148155
@@ -169,6 +176,8 @@ pub async fn put_and_query_with_beacon(branch_key_id: &str) -> Result<(), crate:
169176 // 3. Create BeaconVersion.
170177 // The BeaconVersion inside the list holds the list of beacons on the table.
171178 // The BeaconVersion also stores information about the keystore.
179+ // The BeaconVersion parameter specifies the maximumNumberOfPartitions associated with a given beacon configuration.
180+ // "maximumNumberOfPartitions" must be greater than "numberOfPartitions"; we set it to 8 to allow room for future expansion.
172181 // BeaconVersion must be provided:
173182 // - keyStore: The keystore configured in step 2.
174183 // - keySource: A configuration for the key source.
@@ -182,6 +191,8 @@ pub async fn put_and_query_with_beacon(branch_key_id: &str) -> Result<(), crate:
182191 let beacon_version = BeaconVersion :: builder ( )
183192 . standard_beacons ( standard_beacon_list)
184193 . version ( 1 ) // MUST be 1
194+ . maximum_number_of_partitions ( 8 )
195+ . default_number_of_partitions ( 1 ) //For beacons that do not require partitioning, only a single partition is used.
185196 . key_store ( key_store. clone ( ) )
186197 . key_source ( BeaconKeySource :: Single (
187198 SingleKeyStore :: builder ( )
@@ -305,54 +316,71 @@ pub async fn put_and_query_with_beacon(branch_key_id: &str) -> Result<(), crate:
305316 // This procedure is internal to the client and is abstracted away from the user;
306317 // e.g. the user will only see "123456789012" and never
307318 // "098765432109", though the actual query returned both.
319+
320+ let mut all_results: Vec < HashMap < String , AttributeValue > > = Vec :: new ( ) ;
321+
308322 let expression_attributes_names = HashMap :: from ( [
309323 ( "#last4" . to_string ( ) , "inspector_id_last4" . to_string ( ) ) ,
310324 ( "#unit" . to_string ( ) , "unit" . to_string ( ) ) ,
311325 ] ) ;
312326
313- let expression_attribute_values = HashMap :: from ( [
314- ( ":last4" . to_string ( ) , AttributeValue :: S ( "4321" . to_string ( ) ) ) ,
315- (
316- ":unit" . to_string ( ) ,
317- AttributeValue :: S ( "123456789012" . to_string ( ) ) ,
318- ) ,
319- ] ) ;
327+ // In this simple example, we know there are one buckets,
328+ // but in general the number can be obtained using transformClient.getNumberOfQueries(query)
329+ let num_queries = 1 ;
330+
331+ //We need to query for all possible parttions
332+
333+ for partition in 0 ..num_queries {
334+ let expression_attribute_values = HashMap :: from ( [
335+ ( ":last4" . to_string ( ) , AttributeValue :: S ( "4321" . to_string ( ) ) ) ,
336+ (
337+ ":unit" . to_string ( ) ,
338+ AttributeValue :: S ( "123456789012" . to_string ( ) ) ,
339+ ) ,
340+ (
341+ ":aws_dbe_partition" . to_string ( ) ,
342+ AttributeValue :: N ( partition. to_string ( ) ) ,
343+ ) ,
344+ ] ) ;
345+
346+ // GSIs do not update instantly
347+ // so if the results come back empty
348+ // we retry after a short sleep
349+ for _i in 0 ..10 {
350+ let query_response = ddb
351+ . query ( )
352+ . table_name ( ddb_table_name)
353+ . index_name ( GSI_NAME )
354+ . key_condition_expression ( "#last4 = :last4 and #unit = :unit" )
355+ . set_expression_attribute_names ( Some ( expression_attributes_names. clone ( ) ) )
356+ . set_expression_attribute_values ( Some ( expression_attribute_values. clone ( ) ) )
357+ . send ( )
358+ . await ?;
359+
360+ if let Some ( items) = query_response. items {
361+ if !items. is_empty ( ) {
362+ all_results. extend ( items. clone ( ) ) ;
363+ break ;
364+ }
365+ }
320366
321- // GSIs do not update instantly
322- // so if the results come back empty
323- // we retry after a short sleep
324- for _i in 0 ..10 {
325- let query_response = ddb
326- . query ( )
327- . table_name ( ddb_table_name)
328- . index_name ( GSI_NAME )
329- . key_condition_expression ( "#last4 = :last4 and #unit = :unit" )
330- . set_expression_attribute_names ( Some ( expression_attributes_names. clone ( ) ) )
331- . set_expression_attribute_values ( Some ( expression_attribute_values. clone ( ) ) )
332- . send ( )
333- . await ?;
334-
335- // if no results, sleep and try again
336- if query_response. items . is_none ( ) || query_response. items . as_ref ( ) . unwrap ( ) . is_empty ( ) {
337367 std:: thread:: sleep ( std:: time:: Duration :: from_millis ( 20 ) ) ;
338- continue ;
339368 }
340-
341- let attribute_values = query_response. items . unwrap ( ) ;
342- // Validate only 1 item was returned: the item we just put
343- assert_eq ! ( attribute_values. len( ) , 1 ) ;
344- let returned_item = & attribute_values[ 0 ] ;
345- // Validate the item has the expected attributes
346- assert_eq ! (
347- returned_item[ "inspector_id_last4" ] ,
348- AttributeValue :: S ( "4321" . to_string( ) )
349- ) ;
350- assert_eq ! (
351- returned_item[ "unit" ] ,
352- AttributeValue :: S ( "123456789012" . to_string( ) )
353- ) ;
354- break ;
355369 }
370+ let attribute_values = all_results;
371+ // Validate only 1 item was returned: the item we just put
372+ assert_eq ! ( attribute_values. len( ) , 1 ) ;
373+ let returned_item = & attribute_values[ 0 ] ;
374+ // Validate the item has the expected attributes
375+ assert_eq ! (
376+ returned_item[ "inspector_id_last4" ] ,
377+ AttributeValue :: S ( "4321" . to_string( ) )
378+ ) ;
379+ assert_eq ! (
380+ returned_item[ "unit" ] ,
381+ AttributeValue :: S ( "123456789012" . to_string( ) )
382+ ) ;
383+
356384 println ! ( "basic_searchable_encryption successful." ) ;
357385 Ok ( ( ) )
358386 } ;
0 commit comments