Raw query

miso-connectors/src/lib.rs

@@ -317,5 +317,14 @@ pub trait Connector: Debug + Send + Sync {
         None
     }
 
+    fn raw_query(
+        &self,
+        _collection: &str,
+        _query: &str,
+        _handle: &dyn QueryHandle,
+    ) -> Option<Box<dyn QueryHandle>> {
+        None
+    }
+
     async fn close(&self);
 }

miso-connectors/src/quickwit.rs

@@ -93,6 +93,7 @@ struct QuickwitHandle {
     limit: Option<u64>,
     count: bool,
     collections: Vec<String>,
+    raw_query: Option<Value>,
 }
 
 #[typetag::serde]
@@ -163,12 +164,22 @@ impl QuickwitHandle {
         handle.collections.push(collection.to_string());
         handle
     }
+
+    fn with_raw_query(&self, query: Value) -> QuickwitHandle {
+        let mut handle = self.clone();
+        handle.raw_query = Some(query);
+        handle
+    }
 }
 
 impl fmt::Display for QuickwitHandle {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let mut items = Vec::new();
 
+        if let Some(raw) = &self.raw_query {
+            items.push(format!("raw_query={raw}"));
+        }
+
         if self.count {
             items.push("count".to_string());
         }
@@ -1159,42 +1170,44 @@ impl Connector for QuickwitConnector {
         collections.dedup();
         let collections = collections.join(",");
 
-        let mut query_map = Map::new();
-
-        if !handle.queries.is_empty() {
-            query_map.insert(
-                "query",
-                json!({
-                    "bool": {
-                        "must": handle.queries.clone(),
-                    }
-                }),
-            );
-        }
+        let query = if let Some(raw) = &handle.raw_query {
+            Some(raw.clone())
+        } else {
+            let mut query_map = Map::new();
 
-        if let Some(sorts) = &handle.sorts {
-            query_map.insert("sort", sorts.clone());
-        }
+            if !handle.queries.is_empty() {
+                query_map.insert(
+                    "query",
+                    json!({
+                        "bool": {
+                            "must": handle.queries.clone(),
+                        }
+                    }),
+                );
+            }
 
-        let is_aggregation_query = if let Some(aggs) = &handle.aggs {
-            query_map.insert("size", json!(0));
-            for (key, value) in aggs.as_object().unwrap() {
-                query_map.insert(key, value.clone());
+            if let Some(sorts) = &handle.sorts {
+                query_map.insert("sort", sorts.clone());
             }
-            true
-        } else {
-            if let Some(limit) = limit {
+
+            if let Some(aggs) = &handle.aggs {
+                query_map.insert("size", json!(0));
+                for (key, value) in aggs.as_object().unwrap() {
+                    query_map.insert(key, value.clone());
+                }
+            } else if let Some(limit) = limit {
                 query_map.insert("size", limit.into());
             }
-            false
-        };
 
-        let query = if !query_map.is_empty() {
-            Some(json!(query_map))
-        } else {
-            None
+            if !query_map.is_empty() {
+                Some(json!(query_map))
+            } else {
+                None
+            }
         };
 
+        let is_aggregation_query = handle.aggs.is_some();
+
         info!(
             count = handle.count,
             scroll_size,
@@ -1254,7 +1267,7 @@ impl Connector for QuickwitConnector {
 
     fn apply_filter(&self, ast: &Expr, handle: &dyn QueryHandle) -> Option<Box<dyn QueryHandle>> {
         let handle = downcast_unwrap!(handle, QuickwitHandle);
-        if handle.sorts.is_some() || !handle.group_by.is_empty() {
+        if handle.raw_query.is_some() || handle.sorts.is_some() || !handle.group_by.is_empty() {
             // Cannot filter over top-n / group by in Quickwit.
             return None;
         }
@@ -1267,7 +1280,7 @@ impl Connector for QuickwitConnector {
         handle: &dyn QueryHandle,
     ) -> Option<Box<dyn QueryHandle>> {
         let handle = downcast_unwrap!(handle, QuickwitHandle);
-        if handle.count || !handle.group_by.is_empty() {
+        if handle.raw_query.is_some() || handle.count || !handle.group_by.is_empty() {
             return None;
         }
 
@@ -1286,6 +1299,9 @@ impl Connector for QuickwitConnector {
 
     fn apply_limit(&self, mut max: u64, handle: &dyn QueryHandle) -> Option<Box<dyn QueryHandle>> {
         let handle = downcast_unwrap!(handle, QuickwitHandle);
+        if handle.raw_query.is_some() {
+            return None;
+        }
         if let Some(limit) = handle.limit
             && limit < max
         {
@@ -1301,7 +1317,7 @@ impl Connector for QuickwitConnector {
         handle: &dyn QueryHandle,
     ) -> Option<Box<dyn QueryHandle>> {
         let handle = downcast_unwrap!(handle, QuickwitHandle);
-        if handle.sorts.is_some() {
+        if handle.raw_query.is_some() || handle.sorts.is_some() {
             // Cannot top-n over top-n in Quickwit.
             return None;
         }
@@ -1337,7 +1353,7 @@ impl Connector for QuickwitConnector {
 
     fn apply_count(&self, handle: &dyn QueryHandle) -> Option<Box<dyn QueryHandle>> {
         let handle = downcast_unwrap!(handle, QuickwitHandle);
-        if !handle.group_by.is_empty() {
+        if handle.raw_query.is_some() || !handle.group_by.is_empty() {
             // Quickwit count query returns number of items instead of number of unique groups.
             // This is fine, as usually aggregation requests return few results, we can count
             // them ourselves.
@@ -1352,7 +1368,11 @@ impl Connector for QuickwitConnector {
         handle: &dyn QueryHandle,
     ) -> Option<Box<dyn QueryHandle>> {
         let handle = downcast_unwrap!(handle, QuickwitHandle);
-        if handle.limit.is_some() || handle.sorts.is_some() || !handle.group_by.is_empty() {
+        if handle.raw_query.is_some()
+            || handle.limit.is_some()
+            || handle.sorts.is_some()
+            || !handle.group_by.is_empty()
+        {
             // Quickwit's query (like Elasticsearch's) is not pipelined, most similar to SQL.
             // When you request it to both sort (or limit) and aggregate, it will always first
             // aggregate and then sort (or limit), no way to control the order of these 2 AFAIK.
@@ -1498,6 +1518,9 @@ impl Connector for QuickwitConnector {
         union_handle: &dyn QueryHandle,
     ) -> Option<Box<dyn QueryHandle>> {
         let handle = downcast_unwrap!(handle, QuickwitHandle);
+        if handle.raw_query.is_some() {
+            return None;
+        }
         let union_handle = union_handle.as_any().downcast_ref::<QuickwitHandle>()?;
 
         if handle != union_handle {
@@ -1525,6 +1548,17 @@ impl Connector for QuickwitConnector {
     async fn close(&self) {
         self.interval_task.shutdown().await;
     }
+
+    fn raw_query(
+        &self,
+        _collection: &str,
+        query: &str,
+        handle: &dyn QueryHandle,
+    ) -> Option<Box<dyn QueryHandle>> {
+        let handle = downcast_unwrap!(handle, QuickwitHandle);
+        let parsed: Value = serde_json::from_str(query).ok()?;
+        Some(Box::new(handle.with_raw_query(parsed)))
+    }
 }
 
 #[cfg(test)]

miso-kql/src/lexer.rs

@@ -180,6 +180,8 @@ pub enum Token {
     Iff,
     #[token("extract")]
     Extract,
+    #[token("raw")]
+    Raw,
 
     #[token(",")]
     Comma,

miso-kql/src/parser.rs

@@ -219,6 +219,7 @@ where
         Token::Case => "case".to_string(),
         Token::Iff => "iff".to_string(),
         Token::Extract => "extract".to_string(),
+        Token::Raw => "raw".to_string(),
     }
 }
 
@@ -1291,6 +1292,14 @@ impl KqlParser {
                 .labelled("let")
                 .boxed();
 
+            let string_literal = select! {
+                Token::String(StringValue::Text(s)) => s,
+            };
+
+            let raw_suffix = just(Token::Dot)
+                .ignore_then(just(Token::Raw))
+                .ignore_then(string_literal.delimited_by(just(Token::LParen), just(Token::RParen)));
+
             let scan_step = ident
                 .clone()
                 .then(
@@ -1302,13 +1311,19 @@ impl KqlParser {
                         )))
                         .or_not(),
                 )
-                .map(|(connector, collection)| {
-                    QueryStep::Scan(match collection {
-                        Some(collection) => ScanKind::Collection {
+                .then(raw_suffix.or_not())
+                .map(|((connector, collection), raw_query)| {
+                    QueryStep::Scan(match (collection, raw_query) {
+                        (Some(collection), Some(query)) => ScanKind::Raw {
+                            connector,
+                            collection,
+                            query,
+                        },
+                        (Some(collection), None) => ScanKind::Collection {
                             connector,
                             collection,
                         },
-                        None => ScanKind::Var(connector),
+                        (None, _) => ScanKind::Var(connector),
                     })
                 })
                 .labelled("scan")

miso-kql/src/parser_tests.rs

@@ -1291,6 +1291,39 @@ fn test_parse_extract_with_field_arguments() {
     }
 }
 
+#[test]
+fn test_raw_scan() {
+    let query = "connector.table.raw(\"{\\\"query\\\": {\\\"match_all\\\": {}}}\")";
+    let result = parse_unwrap!(query);
+
+    assert_eq!(result.len(), 1);
+    match &result[0] {
+        QueryStep::Scan(ScanKind::Raw {
+            connector,
+            collection,
+            query,
+        }) => {
+            assert_eq!(connector, "connector");
+            assert_eq!(collection, "table");
+            assert_eq!(query, r#"{"query": {"match_all": {}}}"#);
+        }
+        _ => panic!("Expected Raw scan step"),
+    }
+}
+
+#[test]
+fn test_raw_as_field_name() {
+    let query = r#"connector.table | where raw == "foo""#;
+    let result = parse_unwrap!(query);
+
+    assert_eq!(result.len(), 2);
+    assert!(matches!(
+        result[0],
+        QueryStep::Scan(ScanKind::Collection { .. })
+    ));
+    assert!(!matches!(result[0], QueryStep::Scan(ScanKind::Raw { .. })));
+}
+
 #[test]
 fn test_parse_extract_in_filter() {
     let query = r#"connector.table | where extract("(\\d+)", 1, message) == "123""#;