gsadmin/rustfs
1
0
Fork 0
mirror of https://github.com/rustfs/rustfs.git synced 2026-08-17 10:17:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(s3select): enforce query and resource limits (#5028)

Browse Source 
* fix(s3select): enforce query and resource limits

* fix(s3select): close query resource limit gaps

* fix(s3select): preserve timeout and stream invariants

* fix(s3select): enforce staged query limits

* fix(s3select): preserve policy error compatibility

* fix(s3select): bound error source traversal
This commit is contained in:
GatewayJ
2026-07-25 18:44:53 +08:00
committed by GitHub
parent 2dc4d0b651
commit 0364523dad
13 changed files with 3228 additions and 222 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/s3select-query/src/test/integration_test.rs
+95 -28
View File
@@ -15,7 +15,10 @@
#[cfg(test)]
mod integration_tests {
use crate::{create_fresh_db, get_global_db, instance::make_rustfsms};
use datafusion::arrow::array::{Array, StringArray};
use datafusion::arrow::{
array::{Array, Int64Array, StringArray},
record_batch::RecordBatch,
};
use rustfs_s3select_api::{
QueryError,
query::{Context, Query},
@@ -26,6 +29,52 @@ mod integration_tests {
};
use std::sync::Arc;
fn assert_ages_descending(output: &[RecordBatch]) {
let ages: Vec<i64> = output
.iter()
.flat_map(|batch| {
batch
.column(1)
.as_any()
.downcast_ref::<Int64Array>()
.expect("age column should be Int64")
.values()
.iter()
.copied()
.collect::<Vec<_>>()
})
.collect();
assert_eq!(ages, vec![40, 38, 35, 32, 30, 28, 26, 25, 24, 22]);
}
fn assert_department_counts(output: &[RecordBatch]) {
let mut counts: Vec<(&str, i64)> = output
.iter()
.flat_map(|batch| {
let departments = batch
.column(0)
.as_any()
.downcast_ref::<StringArray>()
.expect("department column should be Utf8");
let counts = batch
.column(1)
.as_any()
.downcast_ref::<Int64Array>()
.expect("count column should be Int64");
departments.iter().zip(counts.iter()).map(|(department, count)| {
(
department.expect("department should not be null"),
count.expect("count should not be null"),
)
})
})
.collect();
counts.sort_unstable();
assert_eq!(counts, vec![("Finance", 3), ("HR", 2), ("IT", 3), ("Marketing", 2)]);
}
fn create_test_input(sql: &str) -> SelectObjectContentInput {
SelectObjectContentInput {
bucket: "test-bucket".to_string(),
@@ -292,21 +341,21 @@ mod integration_tests {
#[tokio::test]
async fn test_select_with_aggregation() {
let sql = "SELECT department, COUNT(*) as count FROM S3Object GROUP BY department";
let sql = "SELECT department, COUNT(*) FROM S3Object GROUP BY department";
let input = create_test_input(sql);
let db = get_global_db(input.clone(), true).await.unwrap();
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
let result = db.execute(&query).await;
// Aggregation queries might fail due to lack of actual data, which is acceptable
match result {
Ok(_) => {
// If successful, that's great
}
Err(_) => {
// Expected to fail due to no actual data source
}
}
let output = db
.execute(&query)
.await
.expect("execute grouped CSV query")
.result()
.chunk_result()
.await
.expect("collect grouped CSV output");
assert_department_counts(&output);
}
#[tokio::test]
@@ -381,13 +430,22 @@ mod integration_tests {
#[tokio::test]
async fn test_query_with_order_by() {
let sql = "SELECT name, age FROM S3Object ORDER BY age DESC";
let sql = "SELECT name, CAST(age AS BIGINT) AS age FROM S3Object ORDER BY CAST(age AS BIGINT) DESC";
let input = create_test_input(sql);
let db = get_global_db(input.clone(), true).await.unwrap();
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
let result = db.execute(&query).await;
assert!(result.is_ok());
let output = db
.execute(&query)
.await
.expect("execute ordered CSV query")
.result()
.chunk_result()
.await
.expect("collect ordered CSV output");
assert_eq!(output.iter().map(|batch| batch.num_rows()).sum::<usize>(), 10);
assert_ages_descending(&output);
}
#[tokio::test]
@@ -582,21 +640,21 @@ mod integration_tests {
#[tokio::test]
async fn test_select_with_aggregation_json() {
let sql = "SELECT department, COUNT(*) as count FROM S3Object GROUP BY department";
let sql = "SELECT department, COUNT(*) FROM S3Object GROUP BY department";
let input = create_test_json_input(sql);
let db = get_global_db(input.clone(), true).await.unwrap();
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
let result = db.execute(&query).await;
// Aggregation queries may fail due to lack of actual data, which is acceptable
match result {
Ok(_) => {
// If successful, that's great
}
Err(_) => {
// Expected to fail due to no actual data source
}
}
let output = db
.execute(&query)
.await
.expect("execute grouped JSON query")
.result()
.chunk_result()
.await
.expect("collect grouped JSON output");
assert_department_counts(&output);
}
#[tokio::test]
@@ -672,8 +730,17 @@ mod integration_tests {
let db = get_global_db(input.clone(), true).await.unwrap();
let query = Query::new(Context { input: Arc::new(input) }, sql.to_string());
let result = db.execute(&query).await;
assert!(result.is_ok());
let output = db
.execute(&query)
.await
.expect("execute ordered JSON query")
.result()
.chunk_result()
.await
.expect("collect ordered JSON output");
assert_eq!(output.iter().map(|batch| batch.num_rows()).sum::<usize>(), 10);
assert_ages_descending(&output);
}
#[tokio::test]