feat: added hostname filtering support to find_bucket

Author: ErikBjare

aw-query/src/functions.rs

@@ -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)?;
 
@@ -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) => {
@@ -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))
     }
 

aw-query/tests/query.rs

@@ -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);

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
@@ -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);
     }
 }