more subquery support

oximeter/db/src/oxql/query/ast.rs

@@ -106,9 +106,8 @@ pub struct Query2 {
 impl Query2 {
     // Check that this query (and any subqueries) start with a get table op.
     pub(crate) fn starts_with_get(&self) -> bool {
-        let Some(op) = self.ops.first() else {
-            return false;
-        };
+        let op =
+            self.ops.first().expect("should have parsed at least 1 operation");
         match op {
             TableOp::Basic(basic) => matches!(basic, BasicTableOp::Get(_)),
             TableOp::Grouped(GroupedTableOp { ops }) => {
@@ -122,10 +121,19 @@ impl Query2 {
     pub(crate) fn merge_ops_follow_subqueries(&self) -> bool {
         let (mut ops, mut next_ops) =
             (self.ops.iter(), self.ops.iter().skip(1));
-        while let (Some(op), Some(next_op)) = (ops.next(), next_ops.next()) {
+        while let (Some(op), maybe_next_op) = (ops.next(), next_ops.next()) {
             if let TableOp::Grouped(GroupedTableOp { ops: subq_ops }) = op {
-                if !next_op.is_merge() {
-                    return false;
+                // If the current subquery contains more than one operation,
+                // then the next op must both exist and be a merging operation.
+                if subq_ops.len() > 1 {
+                    match maybe_next_op {
+                        None => return false,
+                        Some(next_op) => {
+                            if !next_op.is_merge() {
+                                return false;
+                            }
+                        }
+                    }
                 }
                 if !subq_ops.iter().all(Query2::merge_ops_follow_subqueries) {
                     return false;
@@ -827,6 +835,10 @@ impl TableTransformation for Join {
 mod tests {
     use super::duration_consts::*;
     use super::duration_to_db_interval;
+    use crate::oxql::query::ast::BasicTableOp;
+    use crate::oxql::query::ast::Query2;
+    use crate::oxql::query::ast::TableOp;
+    use oximeter::TimeseriesName;
 
     #[test]
     fn test_duration_to_db_interval() {
@@ -878,4 +890,13 @@ mod tests {
             (31536000000000001, "NANOSECOND")
         );
     }
+
+    #[test]
+    fn test_table_op_starts_with_get() {
+        let name: TimeseriesName = "a:b".parse().unwrap();
+        assert!(Query2 {
+            ops: vec![TableOp::Basic(BasicTableOp::Get(name.clone()))]
+        }
+        .starts_with_get());
+    }
 }

oximeter/db/src/oxql/query/grammar.rs

@@ -351,7 +351,7 @@ peg::parser! {
 
         /// Parse a "filter" table operation.
         pub rule filter() -> Filter
-            = "filter" _ item:filter_item()
+            = "filter" _ item:filter_item() _?
         {
             item
         }
@@ -439,9 +439,14 @@ peg::parser! {
 
         /// TODO(ben) This is the way.
         pub rule query2() -> Query2
-            = ops:(basic_table_op() / grouped_table_op()) **<1,> query_delim()
-        {
-            Query2 { ops }
+            = ops:(basic_table_op() / grouped_table_op()) ++ query_delim()
+        {?
+            let query = Query2 { ops };
+            if query.starts_with_get() {
+                Ok(query)
+            } else {
+                Err("every subquery must start with a `get` operation")
+            }
         }
 
         rule grouped_table_op_delim() = quiet!{ _? ";" _? }
@@ -1043,4 +1048,28 @@ mod tests {
         );
         println!("{:#?}", query_parser::query2("{ get foo:bar | filter x == 0; get x:y } | join | group_by [a, b, c]"));
     }
+
+    #[test]
+    fn test_query2_starts_with_get() {
+        assert!(query_parser::query2("{ get a:b }").unwrap().starts_with_get());
+        assert!(query_parser::query2("{ get a:b; get a:b }")
+            .unwrap()
+            .starts_with_get());
+        assert!(query_parser::query2("{ get a:b; get a:b } | filter a == 0")
+            .unwrap()
+            .starts_with_get());
+
+        assert!(query_parser::query2("{ get a:b; filter foo == 0 }").is_err());
+        assert!(query_parser::query2("{ get a:b; filter foo == 0 }").is_err());
+    }
+
+    #[test]
+    fn test_query2_merge_ops_follow_subqueries() {
+        assert!(query_parser::query2("{ get a:b }")
+            .unwrap()
+            .merge_ops_follow_subqueries());
+        assert!(query_parser::query2("{ get a:b; get a:b }")
+            .unwrap()
+            .merge_ops_follow_subqueries());
+    }
 }