diff --git a/src/scripting/alternator/functions.rs b/src/scripting/alternator/functions.rs index 273c5d8..77a0344 100644 --- a/src/scripting/alternator/functions.rs +++ b/src/scripting/alternator/functions.rs @@ -1149,6 +1149,8 @@ pub async fn query( /// - `limit`: The maximum number of items to evaluate (optional). /// - `validation`: An optional item count validation. Look at [extract_validation_args] for details. /// - `with_result`: If true, the scan result is returned (default: false). +/// - `segment`: The segment to scan (optional). +/// - `total_segments`: The total number of segments (optional). #[rune::function(instance)] pub async fn scan( ctx: Ref, @@ -1171,9 +1173,33 @@ pub async fn scan( LIMIT_KEY, VALIDATION_KEY, WITH_RESULT_KEY, + SEGMENT_KEY, + TOTAL_SEGMENTS_KEY, ], )?; + if let Some(v) = params.get(SEGMENT_KEY) { + if let Ok(i) = v.as_signed() { + builder = builder.segment(match i32::try_from(i) { + Ok(val) => val, + Err(_) => return bad_input(format!("'{}' is out of range", SEGMENT_KEY)), + }); + } else { + return bad_input(format!("'{}' must be an integer", SEGMENT_KEY)); + } + } + + if let Some(v) = params.get(TOTAL_SEGMENTS_KEY) { + if let Ok(i) = v.as_signed() { + builder = builder.total_segments(match i32::try_from(i) { + Ok(val) => val, + Err(_) => return bad_input(format!("'{}' is out of range", TOTAL_SEGMENTS_KEY)), + }); + } else { + return bad_input(format!("'{}' must be an integer", TOTAL_SEGMENTS_KEY)); + } + } + if let Some(v) = params.get(FILTER_EXPRESSION_KEY) { if let Ok(s) = v.borrow_ref::() { builder = builder.filter_expression(s.as_str().to_string()); diff --git a/src/scripting/alternator/types.rs b/src/scripting/alternator/types.rs index cd16f3f..5e19012 100644 --- a/src/scripting/alternator/types.rs +++ b/src/scripting/alternator/types.rs @@ -22,6 +22,8 @@ pub const PROJECTION_EXPRESSION_KEY: &str = "projection_expression"; pub const UPDATE_EXPRESSION_KEY: &str = "update"; pub const QUERY_EXPRESSION_KEY: &str = "query"; pub const FILTER_EXPRESSION_KEY: &str = "filter"; +pub const SEGMENT_KEY: &str = "segment"; +pub const TOTAL_SEGMENTS_KEY: &str = "total_segments"; fn alternator_set_to_rune(key: &str, iter: I, wrapper: F) -> Result where diff --git a/workloads/alternator/api_demo.rn b/workloads/alternator/api_demo.rn index 724dbd6..d0cedd4 100644 --- a/workloads/alternator/api_demo.rn +++ b/workloads/alternator/api_demo.rn @@ -122,6 +122,16 @@ pub async fn run(db, i) { } }).await?; + // SCAN with segments + db.scan(TABLE, #{ + filter: "age > :min_age", + "attribute_values": #{ + ":min_age": 18 + }, + "segment": 1, + "total_segments": 5 + }).await?; + // DELETE Item db.delete(TABLE, key, ()).await?; } diff --git a/workloads/alternator/invalid_params.rn b/workloads/alternator/invalid_params.rn index bff4564..4e29df7 100644 --- a/workloads/alternator/invalid_params.rn +++ b/workloads/alternator/invalid_params.rn @@ -152,6 +152,24 @@ pub async fn run(db, i) { filter: 123 }).await, "scan", "filter"); + // Scan: segment must be integer and within range + assert_bad_param(db.scan(TABLE, #{ + segment: "1" + }).await, "scan", "segment"); + + assert_bad_param(db.scan(TABLE, #{ + segment: 1099511627776 + }).await, "scan", "segment is out of range"); + + // Scan: total_segments must be integer and within range + assert_bad_param(db.scan(TABLE, #{ + total_segments: "5" + }).await, "scan", "total_segments"); + + assert_bad_param(db.scan(TABLE, #{ + total_segments: 1099511627776 + }).await, "scan", "total_segments is out of range"); + // Create Table: primary_key must be object or string assert_bad_param(db.create_table("fail_table", #{ primary_key: 123