Skip to content

Commit

Permalink
feat: added hostname filtering support to find_bucket
Browse files Browse the repository at this point in the history
  • Loading branch information
ErikBjare committed Jun 8, 2021
1 parent 4a1943a commit 7ad9512
Show file tree
Hide file tree
Showing 3 changed files with 68 additions and 24 deletions.
28 changes: 18 additions & 10 deletions aw-query/src/functions.rs
Expand Up @@ -138,6 +138,7 @@ mod qfunctions {
) -> Result<DataType, QueryError> {
// Typecheck
validate::args_length(&args, 1)?;

let bucket_id: String = (&args[0]).try_into()?;
let interval = validate::get_timeinterval(env)?;

Expand Down Expand Up @@ -189,8 +190,14 @@ mod qfunctions {
_env: &VarEnv,
ds: &Datastore,
) -> Result<DataType, QueryError> {
validate::args_length(&args, 1)?;
validate::args_length(&args, 1).or_else(|_| validate::args_length(&args, 2))?;

let bucket_filter: String = (&args[0]).try_into()?;
let hostname_filter: Option<String> = match args.len() {
2 => Some((&args[1]).try_into()?),
_ => None,
};

let buckets = match ds.get_buckets() {
Ok(buckets) => buckets,
Err(e) => {
Expand All @@ -200,15 +207,16 @@ mod qfunctions {
)))
}
};
let bucketname = match aw_transform::find_bucket(&bucket_filter, buckets.keys()) {
Some(bucketname) => bucketname,
None => {
return Err(QueryError::BucketQueryError(format!(
"Couldn't find any bucket which starts with {}",
bucket_filter
)));
}
};
let bucketname =
match aw_transform::find_bucket(&bucket_filter, &hostname_filter, buckets.values()) {
Some(bucketname) => bucketname,
None => {
return Err(QueryError::BucketQueryError(format!(
"Couldn't find any bucket which starts with {}",
bucket_filter
)));
}
};
Ok(DataType::String(bucketname))
}

Expand Down
2 changes: 1 addition & 1 deletion aw-query/tests/query.rs
Expand Up @@ -351,7 +351,7 @@ mod query_tests {

let code = format!(
r#"
events = query_bucket(find_bucket("{}"));
events = query_bucket(find_bucket("{}", "testhost"));
events = flood(events);
events = sort_by_duration(events);
events = limit_events(events, 10000);
Expand Down
62 changes: 49 additions & 13 deletions aw-transform/src/find_bucket.rs
@@ -1,11 +1,21 @@
/// Finds the first bucket which starts with the specified string
use aw_models::Bucket;

/// Finds the first bucket which starts with the specified string, and optionally matches a
/// hostname.
pub fn find_bucket<'a>(
bucket_filter: &str,
bucketnames: impl IntoIterator<Item = &'a String>,
hostname_filter: &Option<String>,
buckets: impl IntoIterator<Item = &'a Bucket>,
) -> Option<String> {
for bucketname in bucketnames {
if bucketname.starts_with(bucket_filter) {
return Some(bucketname.to_string());
for bucket in buckets {
if bucket.id.starts_with(bucket_filter) {
if let Some(hostname) = hostname_filter {
if hostname == &bucket.hostname {
return Some(bucket.id.to_string());
}
} else {
return Some(bucket.id.to_string());
}
}
}
None
Expand All @@ -14,20 +24,46 @@ pub fn find_bucket<'a>(
#[cfg(test)]
mod tests {
use super::find_bucket;
use aw_models::Bucket;
use aw_models::BucketMetadata;

#[test]
fn test_find_bucket() {
let expected_bucketname = "aw-datastore-test_test-host".to_string();
let bucketnames = vec![
"no match".to_string(),
expected_bucketname.clone(),
"no match 2".to_string(),
];
let res = find_bucket("aw-datastore-test", &bucketnames);
let expected_hostname = "testhost".to_string();
let b1 = Bucket {
bid: None,
id: "no match".to_string(),
_type: "type".to_string(),
hostname: expected_hostname.clone(),
client: "testclient".to_string(),
created: None,
data: json_map! {},
metadata: BucketMetadata::default(),
events: None,
last_updated: None,
};
let mut b2 = b1.clone();
b2.id = expected_bucketname.clone();
let b3 = b1.clone();

let buckets = vec![b1.clone(), b2.clone(), b3.clone()];

// Check that it correctly finds bucket
let res = find_bucket("aw-datastore-test", &Some("testhost".to_string()), &buckets);
assert_eq!(res, Some(expected_bucketname));

let bucketnames = vec!["no match".to_string(), "no match 2".to_string()];
let res = find_bucket("aw-datastore-test", &bucketnames);
// Check that it doesn't find a bucket for an unavailable hostname
let res = find_bucket(
"aw-datastore-test",
&Some("unavailablehost".to_string()),
&buckets,
);
assert_eq!(res, None);

// Check that it doesn't find a bucket for any hostname
let buckets = vec![b1, b3];
let res = find_bucket("aw-datastore-test", &None, &buckets);
assert_eq!(res, None);
}
}

0 comments on commit 7ad9512

Please sign in to comment.