Merge pull request #246 from richardsun0713/rsun-multi-compute-groupby

feat(logs): support multiple --compute and --group-by values

Author: platinummonkey

docs/EXAMPLES.md

@@ -128,6 +128,13 @@ pup logs aggregate \
   --from="30m" \
   --compute="percentile(@duration, 99)" \
   --group-by="service"
+
+# Multiple metrics in one query (comma-separated)
+pup logs aggregate \
+  --query="service:web-app" \
+  --from="1h" \
+  --compute="count,avg(@duration),percentile(@duration, 95)" \
+  --group-by="service,status"
 ```
 
 ### Search Logs in Specific Storage Tier

src/commands/logs.rs

@@ -22,8 +22,8 @@ pub struct AggregateArgs {
     pub query: String,
     pub from: String,
     pub to: String,
-    pub compute: String,
-    pub group_by: Option<String>,
+    pub compute: Vec<String>,
+    pub group_by: Vec<String>,
     pub limit: i32,
     pub storage: Option<String>,
 }
@@ -56,6 +56,39 @@ fn parse_storage_tier(storage: Option<String>) -> Result<Option<LogsStorageTier>
     }
 }
 
+/// Split a comma-separated compute string into individual compute expressions,
+/// respecting parentheses so that `percentile(@duration, 95)` is not split.
+pub fn split_compute_args(input: &str) -> Vec<String> {
+    let mut result = Vec::new();
+    let mut current = String::new();
+    let mut depth = 0u32;
+    for ch in input.chars() {
+        match ch {
+            '(' => {
+                depth += 1;
+                current.push(ch);
+            }
+            ')' => {
+                depth = depth.saturating_sub(1);
+                current.push(ch);
+            }
+            ',' if depth == 0 => {
+                let trimmed = current.trim().to_string();
+                if !trimmed.is_empty() {
+                    result.push(trimmed);
+                }
+                current.clear();
+            }
+            _ => current.push(ch),
+        }
+    }
+    let trimmed = current.trim().to_string();
+    if !trimmed.is_empty() {
+        result.push(trimmed);
+    }
+    result
+}
+
 fn parse_compute_raw(input: &str) -> Result<(String, Option<String>)> {
     let input = input.trim();
     if input.is_empty() {
@@ -110,12 +143,11 @@ fn build_aggregate_body(
     query: String,
     from_ms: i64,
     to_ms: i64,
-    compute: String,
-    group_by: Option<String>,
+    computes: Vec<String>,
+    group_bys: Vec<String>,
     limit: i32,
     storage: Option<String>,
 ) -> Result<serde_json::Value> {
-    let (aggregation, metric) = parse_compute_raw(&compute)?;
     let storage_tier = normalize_storage_tier(storage)?;
 
     let mut filter = serde_json::json!({
@@ -127,22 +159,35 @@ fn build_aggregate_body(
         filter["storage_tier"] = serde_json::Value::String(tier);
     }
 
-    let mut compute_obj = serde_json::json!({ "aggregation": aggregation });
-    if let Some(metric) = metric {
-        compute_obj["metric"] = serde_json::Value::String(metric);
-    }
+    let compute_arr: Vec<serde_json::Value> = computes
+        .iter()
+        .map(|c| {
+            let (aggregation, metric) = parse_compute_raw(c)?;
+            let mut obj = serde_json::json!({ "aggregation": aggregation });
+            if let Some(m) = metric {
+                obj["metric"] = serde_json::Value::String(m);
+            }
+            Ok(obj)
+        })
+        .collect::<Result<Vec<_>>>()?;
 
     let mut body = serde_json::json!({
         "filter": filter,
-        "compute": [compute_obj]
+        "compute": compute_arr
     });
 
-    if let Some(facet) = group_by {
-        let mut group_by_obj = serde_json::json!({ "facet": facet });
-        if limit > 0 {
-            group_by_obj["limit"] = serde_json::json!(limit);
-        }
-        body["group_by"] = serde_json::json!([group_by_obj]);
+    if !group_bys.is_empty() {
+        let group_by_arr: Vec<serde_json::Value> = group_bys
+            .iter()
+            .map(|facet| {
+                let mut obj = serde_json::json!({ "facet": facet });
+                if limit > 0 {
+                    obj["limit"] = serde_json::json!(limit);
+                }
+                obj
+            })
+            .collect();
+        body["group_by"] = serde_json::json!(group_by_arr);
     }
 
     Ok(body)
@@ -285,11 +330,14 @@ pub async fn aggregate(cfg: &Config, args: AggregateArgs) -> Result<()> {
         query,
         from,
         to,
-        compute,
+        mut compute,
         group_by,
         limit,
         storage,
     } = args;
+    if compute.is_empty() {
+        compute.push("count".into());
+    }
     let from_ms = util::parse_time_to_unix_millis(&from)?;
     let to_ms = util::parse_time_to_unix_millis(&to)?;
     let body = build_aggregate_body(query, from_ms, to_ms, compute, group_by, limit, storage)?;
@@ -304,11 +352,14 @@ pub async fn aggregate(cfg: &Config, args: AggregateArgs) -> Result<()> {
         query,
         from,
         to,
-        compute,
+        mut compute,
         group_by,
         limit,
         storage,
     } = args;
+    if compute.is_empty() {
+        compute.push("count".into());
+    }
     let from_ms = util::parse_time_to_unix_millis(&from)?;
     let to_ms = util::parse_time_to_unix_millis(&to)?;
     let body = build_aggregate_body(query, from_ms, to_ms, compute, group_by, limit, storage)?;
@@ -626,7 +677,7 @@ mod tests {
 
     #[test]
     fn test_parse_compute_unsupported_percentile() {
-        assert!(parse_compute_raw("percentile(@duration, 50)").is_err());
+        assert!(parse_compute_raw("percentile(@duration, 42)").is_err());
     }
 
     #[test]
@@ -652,8 +703,8 @@ mod tests {
             "service:web".into(),
             1,
             2,
-            "avg(@duration)".into(),
-            Some("service".into()),
+            vec!["avg(@duration)".into()],
+            vec!["service".into()],
             3,
             Some("flex".into()),
         )
@@ -682,7 +733,8 @@ mod tests {
 
     #[test]
     fn test_build_aggregate_body_omits_group_by_for_plain_count() {
-        let body = build_aggregate_body("*".into(), 1, 2, "count".into(), None, 10, None).unwrap();
+        let body =
+            build_aggregate_body("*".into(), 1, 2, vec!["count".into()], vec![], 10, None).unwrap();
 
         assert_eq!(
             body,
@@ -698,4 +750,102 @@ mod tests {
             })
         );
     }
+
+    #[test]
+    fn test_build_aggregate_body_multiple_computes() {
+        let body = build_aggregate_body(
+            "*".into(),
+            1,
+            2,
+            vec![
+                "count".into(),
+                "avg(@duration)".into(),
+                "percentile(@duration, 95)".into(),
+            ],
+            vec![],
+            10,
+            None,
+        )
+        .unwrap();
+
+        assert_eq!(
+            body,
+            serde_json::json!({
+                "filter": {
+                    "query": "*",
+                    "from": "1",
+                    "to": "2"
+                },
+                "compute": [
+                    { "aggregation": "count" },
+                    { "aggregation": "avg", "metric": "@duration" },
+                    { "aggregation": "pc95", "metric": "@duration" }
+                ]
+            })
+        );
+    }
+
+    #[test]
+    fn test_build_aggregate_body_multiple_group_bys() {
+        let body = build_aggregate_body(
+            "*".into(),
+            1,
+            2,
+            vec!["count".into()],
+            vec!["service".into(), "status".into()],
+            5,
+            None,
+        )
+        .unwrap();
+
+        assert_eq!(
+            body,
+            serde_json::json!({
+                "filter": {
+                    "query": "*",
+                    "from": "1",
+                    "to": "2"
+                },
+                "compute": [{ "aggregation": "count" }],
+                "group_by": [
+                    { "facet": "service", "limit": 5 },
+                    { "facet": "status", "limit": 5 }
+                ]
+            })
+        );
+    }
+
+    #[test]
+    fn test_split_compute_args_single() {
+        assert_eq!(split_compute_args("count"), vec!["count"]);
+    }
+
+    #[test]
+    fn test_split_compute_args_multiple() {
+        assert_eq!(
+            split_compute_args("count,avg(@duration),max(@duration)"),
+            vec!["count", "avg(@duration)", "max(@duration)"]
+        );
+    }
+
+    #[test]
+    fn test_split_compute_args_preserves_parens_with_comma() {
+        assert_eq!(
+            split_compute_args("count,percentile(@duration, 95)"),
+            vec!["count", "percentile(@duration, 95)"]
+        );
+    }
+
+    #[test]
+    fn test_split_compute_args_trims_whitespace() {
+        assert_eq!(
+            split_compute_args(" count , avg(@duration) "),
+            vec!["count", "avg(@duration)"]
+        );
+    }
+
+    #[test]
+    fn test_split_compute_args_empty() {
+        assert!(split_compute_args("").is_empty());
+    }
 }

src/main.rs

@@ -1182,6 +1182,12 @@ enum Commands {
     ///   # Aggregate logs by status
     ///   pup logs aggregate --query="*" --compute="count" --group-by="status"
     ///
+    ///   # Multiple computes in one query (comma-separated)
+    ///   pup logs aggregate --query="*" --compute="count,avg(@duration),percentile(@duration, 95)"
+    ///
+    ///   # Multiple group-by dimensions (comma-separated)
+    ///   pup logs aggregate --query="*" --compute="count" --group-by="service,status"
+    ///
     ///   # List log archives
     ///   pup logs archives list
     ///
@@ -2137,11 +2143,18 @@ enum LogActions {
         from: String,
         #[arg(long, default_value = "now", help = "End time")]
         to: String,
-        #[arg(long, default_value = "count", help = "Metric to compute")]
+        #[arg(
+            long,
+            default_value = "count",
+            help = "Metrics to compute (comma-separated, e.g. count,avg(@duration),percentile(@duration, 95))"
+        )]
         compute: String,
-        #[arg(long, help = "Field to group by")]
+        #[arg(
+            long,
+            help = "Fields to group by (comma-separated, e.g. service,status)"
+        )]
         group_by: Option<String>,
-        #[arg(long, default_value_t = 10, help = "Maximum groups")]
+        #[arg(long, default_value_t = 10, help = "Maximum groups per facet")]
         limit: i32,
         #[arg(long, help = "Storage tier: indexes, online-archives, or flex")]
         storage: Option<String>,
@@ -6337,8 +6350,15 @@ async fn main_inner() -> anyhow::Result<()> {
                             query: query.unwrap_or_default(),
                             from,
                             to,
-                            compute,
-                            group_by,
+                            compute: commands::logs::split_compute_args(&compute),
+                            group_by: group_by
+                                .map(|g| {
+                                    g.split(',')
+                                        .map(|s| s.trim().to_string())
+                                        .filter(|s| !s.is_empty())
+                                        .collect()
+                                })
+                                .unwrap_or_default(),
                             limit,
                             storage,
                         },