1use std::collections::HashMap;
13
14use omgbase_search::{EmbeddingProvider, f32_to_blob};
15use omgbase_store::Store;
16use oqx::ast::{Expr, Query, SelectItem};
17use oqx::walk::{Clause, Node, VisitContext, Visitor, transform, visit};
18use oqx::{Consumer, Engine, InMemoryEngine, Value, build, resolve_aliases};
19use serde_json::{Map, Value as Json};
20
21use crate::context::{SemanticVec, StoreContext, render_row_values, target_of};
22use crate::cursor::{decode_cursor, encode_cursor};
23use crate::error::{Result, SurfaceError};
24use crate::planner::{SqlitePlanner, root_target};
25
26pub const DEFAULT_LIMIT: usize = 50;
28
29const ROW_FNS: [&str; 12] = [
35 "text",
36 "semantic",
37 "under",
38 "under_heading",
39 "within",
40 "under_kind",
41 "yaml_path",
42 "json_pointer",
43 "has_edge",
44 "has_anchor",
45 "child_count",
46 "parent_type",
47];
48
49const ID_KEY: &str = "__oqx_id";
50const PATH_KEY: &str = "__oqx_path";
51const SELF_KEY: &str = "__oqx_self";
61const VALUE_KEY: &str = "__oqx_value";
65
66#[derive(Clone, Copy, Default)]
68pub struct QueryOptions<'a> {
69 pub limit: Option<usize>,
71 pub cursor: Option<&'a str>,
73 pub provider: Option<&'a dyn EmbeddingProvider>,
76 pub in_memory: bool,
79}
80
81#[derive(Clone, Debug, PartialEq)]
83pub struct OqxResult {
84 pub hits: Vec<Json>,
85 pub truncated: bool,
86 pub cursor: Option<String>,
87 pub consumer: Consumer,
88 pub count: Option<f64>,
89 pub exists: Option<bool>,
90 pub none: Option<bool>,
91 pub values: Option<Vec<Json>>,
93}
94
95impl OqxResult {
96 fn scalar(consumer: Consumer) -> Self {
97 Self {
98 hits: Vec::new(),
99 truncated: false,
100 cursor: None,
101 consumer,
102 count: None,
103 exists: None,
104 none: None,
105 values: None,
106 }
107 }
108
109 #[must_use]
112 pub fn to_json(&self) -> Json {
113 let mut m = Map::new();
114 m.insert("hits".to_owned(), Json::Array(self.hits.clone()));
115 m.insert("truncated".to_owned(), Json::Bool(self.truncated));
116 m.insert(
117 "cursor".to_owned(),
118 self.cursor.clone().map_or(Json::Null, Json::String),
119 );
120 m.insert(
121 "consumer".to_owned(),
122 Json::String(self.consumer.as_str().to_owned()),
123 );
124 if let Some(n) = self.count {
125 m.insert("count".to_owned(), Value::Number(n).to_canonical_json());
126 }
127 if let Some(b) = self.exists {
128 m.insert("exists".to_owned(), Json::Bool(b));
129 }
130 if let Some(b) = self.none {
131 m.insert("none".to_owned(), Json::Bool(b));
132 }
133 if let Some(v) = &self.values {
134 m.insert("values".to_owned(), Json::Array(v.clone()));
135 }
136 Json::Object(m)
137 }
138}
139
140#[must_use]
145pub fn rewrite_query(q: &Query) -> Query {
146 transform(q, &mut |e, _| match e {
147 Expr::Call {
148 recv: None,
149 name,
150 args,
151 span,
152 } if ROW_FNS.contains(&name.as_str()) => Expr::Call {
153 recv: Some(Box::new(build::ident("$self"))),
154 name,
155 args,
156 span,
157 },
158 other => other,
159 })
160}
161
162struct Phrases(Vec<String>);
166
167impl Visitor for Phrases {
168 fn enter(&mut self, node: Node<'_>, _ctx: &VisitContext<'_>) -> bool {
169 if let Node::Expr(Expr::Call {
170 recv: None,
171 name,
172 args,
173 ..
174 }) = node
175 && name == "semantic"
176 && let Some(Expr::Lit {
177 value: Value::Str(s),
178 ..
179 }) = args.first()
180 && !self.0.contains(s)
181 {
182 self.0.push(s.clone());
183 }
184 true
185 }
186}
187
188#[must_use]
191pub fn collect_semantic_phrases(source: &str) -> Vec<String> {
192 let Ok(q) = oqx::parse_string(source) else {
193 return Vec::new();
194 };
195 let mut phrases = Phrases(Vec::new());
196 visit(Node::Query(&q), &mut phrases);
197 phrases.0
198}
199
200fn mentions_outside_source(q: &Query, name: &str) -> bool {
206 struct Mentions<'n> {
207 name: &'n str,
208 found: bool,
209 }
210 impl Visitor for Mentions<'_> {
211 fn enter(&mut self, node: Node<'_>, ctx: &VisitContext<'_>) -> bool {
212 if self.found {
213 return false;
214 }
215 if ctx.clause == Some(Clause::Source) && ctx.path.len() == 1 {
217 return false;
218 }
219 if let Node::Expr(Expr::Ident { name, .. } | Expr::Outer { name, .. }) = node
220 && name == self.name
221 {
222 self.found = true;
223 }
224 !self.found
225 }
226 }
227 let mut m = Mentions { name, found: false };
228 visit(Node::Query(q), &mut m);
229 m.found
230}
231
232fn may_reach_non_rows(q: &Query) -> bool {
238 q.follow.is_some() || !q.from.is_empty() || root_target(&q.source).is_none()
239}
240
241fn describe_value(v: &Value) -> String {
243 match v {
244 Value::Undefined | Value::Null => "an absent value".to_owned(),
245 Value::Str(s) => format!(
246 "a string ({})",
247 serde_json::to_string(s).unwrap_or_default()
248 ),
249 Value::Number(_) => format!("a number ({v})"),
250 Value::Bool(b) => format!("a boolean ({b})"),
251 Value::Array(_) => "an array".to_owned(),
252 Value::Range(_) => "a range".to_owned(),
253 Value::Object(_) => "an object".to_owned(),
254 }
255}
256
257fn not_a_store_row(v: &Value) -> SurfaceError {
258 SurfaceError::filter_invalid(
259 format!(
260 "a hit must be a document, block, node or edge row — the query reached {}; to follow document references held in a property use refs(<field>)",
261 describe_value(v)
262 ),
263 "OQX",
264 )
265}
266
267fn to_hit(row: Value) -> Result<Value> {
275 let Value::Object(mut o) = row else {
276 return Ok(Value::Object(oqx::Object::new()));
277 };
278 if let Some(me) = o.remove(SELF_KEY) {
279 if target_of(&me).is_none() {
280 return Err(not_a_store_row(&me));
281 }
282 }
283 let Value::Object(o) = render_row_values(Value::Object(o)) else {
284 return Ok(Value::Object(oqx::Object::new()));
285 };
286 let mut id = Value::Undefined;
287 let mut path = Value::Undefined;
288 let mut rest = Vec::new();
289 for (k, v) in o {
290 match k.as_str() {
291 ID_KEY => id = v,
292 PATH_KEY => path = v,
293 _ => rest.push((k, v)),
294 }
295 }
296 let mut hit = oqx::Object::with_capacity(rest.len() + 2);
297 hit.insert("id", Value::Str(id.to_string()));
298 hit.insert(
299 "path",
300 Value::Str(if path.is_absent() {
301 String::new()
302 } else {
303 path.to_string()
304 }),
305 );
306 for (k, v) in rest {
307 hit.insert(k, v);
308 }
309 Ok(Value::Object(hit))
310}
311
312fn hit_str(hit: &Value, key: &str) -> String {
313 hit.as_object()
314 .and_then(|o| o.get(key))
315 .map(|v| v.to_string())
316 .unwrap_or_default()
317}
318
319fn dedup_hits_by_projection(hits: Vec<Value>) -> Vec<Value> {
323 let mut seen: Vec<String> = Vec::new();
324 let mut out = Vec::new();
325 for h in hits {
326 let mut pairs: Vec<(String, Value)> = h
327 .as_object()
328 .map(|o| {
329 o.iter()
330 .filter(|(k, _)| *k != "id" && *k != "path")
331 .map(|(k, v)| (k.to_owned(), v.clone()))
332 .collect()
333 })
334 .unwrap_or_default();
335 pairs.sort_by(|a, b| a.0.cmp(&b.0));
336 let key = Value::Array(
337 pairs
338 .into_iter()
339 .map(|(k, v)| Value::Array(vec![Value::Str(k), v]))
340 .collect(),
341 )
342 .to_canonical_json()
343 .to_string();
344 if seen.contains(&key) {
345 continue;
346 }
347 seen.push(key);
348 out.push(h);
349 }
350 out
351}
352
353fn const_bound(e: Option<&Expr>, word: &str) -> Result<Option<usize>> {
356 match e {
357 None => Ok(None),
358 Some(Expr::Lit {
359 value: Value::Number(n),
360 ..
361 }) if n.fract() == 0.0 && *n >= 0.0 && n.is_finite() => Ok(Some(*n as usize)),
362 Some(_) => Err(SurfaceError::filter_invalid(
363 format!("top-level {word} must be a non-negative integer literal"),
364 "OQX",
365 )),
366 }
367}
368
369fn value_of(hit: &Value) -> Value {
370 hit.as_object()
371 .and_then(|o| o.get(VALUE_KEY))
372 .cloned()
373 .unwrap_or(Value::Undefined)
374}
375
376fn without_value_key(hit: Value) -> Json {
377 hit.to_canonical_json()
378}
379
380pub fn query(
384 store: &Store,
385 repo_id: &str,
386 source: &str,
387 opts: QueryOptions<'_>,
388) -> Result<OqxResult> {
389 let phrases = collect_semantic_phrases(source);
394 let mut semantic: HashMap<String, SemanticVec> = HashMap::new();
395 if let Some(provider) = opts.provider.filter(|_| !phrases.is_empty()) {
396 for phrase in phrases {
397 let vec = provider
398 .embed_query(&phrase)
399 .map_err(|e| SurfaceError::new(e.code(), e.to_string()))?;
400 semantic.insert(
401 phrase,
402 SemanticVec {
403 model: provider.model().to_owned(),
404 vec: f32_to_blob(&vec),
405 },
406 );
407 }
408 }
409 let runner = Runner {
410 store,
411 repo_id,
412 semantic,
413 planned: !opts.in_memory,
414 };
415 run_inner(&runner, source, opts)
416}
417
418struct Runner<'a> {
424 store: &'a Store,
425 repo_id: &'a str,
426 semantic: HashMap<String, SemanticVec>,
427 planned: bool,
428}
429
430impl Runner<'_> {
431 fn run(&self, q: &Query) -> Result<oqx::OqxResult> {
438 let conn = self.store.conn();
439 let ctx = StoreContext::new(conn, self.repo_id, self.semantic.clone());
440 let plan = if self.planned {
441 SqlitePlanner::new(conn, self.repo_id)
442 .try_plan(q, &[])
443 .map_err(|e| SurfaceError::other(format!("sqlite: {e}")))?
444 } else {
445 None
446 };
447 let (ctx, residual) = match plan {
448 Some(plan) => {
453 let once = !mentions_outside_source(&plan.residual, oqx::ROWS_ROOT);
454 let ctx = if once {
455 ctx.with_rows_root_once(plan.rows)
456 } else {
457 ctx.with_rows_root(plan.rows)
458 };
459 (ctx, Some(plan.residual))
460 }
461 None => (ctx, None),
462 };
463 let engine = InMemoryEngine::new(ctx);
464 let out = engine.run(residual.as_ref().unwrap_or(q), &[]);
465 if let Some(failed) = engine.context().take_root_failure() {
468 return Err(failed.into());
469 }
470 Ok(out?)
473 }
474}
475
476fn run_inner(engine: &Runner<'_>, source: &str, opts: QueryOptions<'_>) -> Result<OqxResult> {
477 let parsed = rewrite_query(&resolve_aliases(&oqx::parse_string(source)?)?);
480 let consumer = parsed.consumer;
481
482 match consumer {
483 Consumer::Exists => {
484 let res = engine.run(&parsed)?;
485 let mut r = OqxResult::scalar(consumer);
486 r.exists = Some(matches!(res, oqx::OqxResult::Exists(true)));
487 return Ok(r);
488 }
489 Consumer::Count => {
490 let res = engine.run(&parsed)?;
491 let mut r = OqxResult::scalar(consumer);
492 r.count = Some(match res {
493 oqx::OqxResult::Count(n) => n,
494 _ => 0.0,
495 });
496 return Ok(r);
497 }
498 Consumer::None => {
499 let res = engine.run(&parsed)?;
500 let mut r = OqxResult::scalar(consumer);
501 r.none = Some(match res {
502 oqx::OqxResult::None(b) => b,
503 _ => true,
504 });
505 return Ok(r);
506 }
507 Consumer::Collect | Consumer::First | Consumer::Single => {}
508 }
509
510 let top_distinct = parsed.distinct;
516 let top_values = parsed.values;
517 let user_select: Vec<SelectItem> = if top_values {
518 parsed
519 .select
520 .first()
521 .map(|it| match it {
522 SelectItem::Field {
523 expr, lift, span, ..
524 } => SelectItem::Field {
525 name: VALUE_KEY.to_owned(),
526 expr: expr.clone(),
527 lift: *lift,
528 span: *span,
529 },
530 SelectItem::Collect { op, span, .. } => SelectItem::Collect {
531 name: VALUE_KEY.to_owned(),
532 op: op.clone(),
533 span: *span,
534 },
535 })
536 .into_iter()
537 .collect()
538 } else {
539 parsed.select.clone()
540 };
541 let id_item = build::field(ID_KEY, build::ident("$id"));
542 let path_item = build::field(PATH_KEY, build::ident("$path"));
543 let mut select = vec![id_item, path_item];
544 if !top_values && may_reach_non_rows(&parsed) {
545 select.push(build::field(SELF_KEY, build::ident("$it")));
546 }
547 select.extend(user_select);
548 let (top_limit, top_offset) = (parsed.limit.clone(), parsed.offset.clone());
552 let q = Query {
553 distinct: false,
554 values: false,
555 select,
556 limit: if consumer == Consumer::Collect {
557 None
558 } else {
559 parsed.limit.clone()
560 },
561 offset: if consumer == Consumer::Collect {
562 None
563 } else {
564 parsed.offset.clone()
565 },
566 ..parsed.clone()
567 };
568 let res = engine.run(&q)?;
569
570 if matches!(consumer, Consumer::First | Consumer::Single) {
571 let row = match res {
572 oqx::OqxResult::First(r) | oqx::OqxResult::Single(r) => r,
573 _ => None,
574 };
575 let mut out = OqxResult::scalar(consumer);
576 match row {
577 None => {
578 if top_values {
579 out.values = Some(Vec::new());
580 }
581 }
582 Some(r) => {
583 let hit = to_hit(r)?;
584 if top_values {
585 out.values = Some(vec![value_of(&hit).to_canonical_json()]);
586 } else {
587 out.hits = vec![without_value_key(hit)];
588 }
589 }
590 }
591 return Ok(out);
592 }
593
594 let mut rows: Vec<Value> = match res {
596 oqx::OqxResult::Collect(rows) => rows,
597 _ => Vec::new(),
598 };
599 let custom = parsed.order_by.as_ref().is_some_and(|o| !o.is_empty());
600 let cursor = opts.cursor.filter(|c| !c.is_empty() && !custom);
601 let eager = top_distinct || cursor.is_some();
606 if eager {
607 rows = rows.into_iter().map(to_hit).collect::<Result<_>>()?;
608 }
609 if top_distinct {
610 rows = dedup_hits_by_projection(rows);
611 }
612 let offset = const_bound(top_offset.as_ref(), "offset")?.unwrap_or(0);
613 let limit = const_bound(top_limit.as_ref(), "limit")?;
614 if offset > 0 || limit.is_some() {
615 let end = limit.map_or(rows.len(), |l| (offset + l).min(rows.len()));
616 if offset >= rows.len() {
617 rows.clear();
618 } else {
619 rows.truncate(end);
620 rows.drain(..offset);
621 }
622 }
623 let cap = opts.limit.unwrap_or(DEFAULT_LIMIT);
624 let mut page = rows;
625 if let Some(cursor) = cursor {
626 let parts = decode_cursor(cursor, "query", 2)?;
627 let (path, id) = (&parts[0], &parts[1]);
628 page.retain(|h| {
629 let hp = hit_str(h, "path");
630 let hi = hit_str(h, "id");
631 hp > *path || (hp == *path && hi > *id)
632 });
633 }
634 let truncated = page.len() > cap;
635 page.truncate(cap);
636 if !eager {
637 page = page.into_iter().map(to_hit).collect::<Result<_>>()?;
638 }
639 let cursor = if truncated && !custom {
640 page.last()
641 .map(|last| encode_cursor(&[&hit_str(last, "path"), &hit_str(last, "id")]))
642 } else {
643 None
644 };
645 let mut out = OqxResult::scalar(Consumer::Collect);
646 out.truncated = truncated;
647 out.cursor = cursor;
648 if top_values {
649 out.values = Some(
650 page.iter()
651 .map(|h| value_of(h).to_canonical_json())
652 .collect(),
653 );
654 } else {
655 out.hits = page.into_iter().map(without_value_key).collect();
656 }
657 Ok(out)
658}
659
660#[cfg(test)]
661mod tests {
662 use super::*;
663 use oqx::ast::Where;
664
665 #[test]
666 fn row_functions_become_self_methods() {
667 let q = oqx::parse_string(
668 "from blocks where text(\"x\") && under_heading(\"h\") && size(attrs) > 0 && doc.$path.startsWith(\"a\")",
669 )
670 .unwrap();
671 let r = rewrite_query(&q);
672 let Some(Where::And { parts, .. }) = &r.r#where else {
673 panic!("and")
674 };
675 let Where::Scalar { expr, .. } = &parts[0] else {
676 panic!("scalar")
677 };
678 assert!(
679 matches!(expr, Expr::Call { recv: Some(r), name, .. } if name == "text" && **r == build::ident("$self"))
680 );
681 let Where::Scalar { expr, .. } = &parts[2] else {
682 panic!("scalar")
683 };
684 assert!(
685 matches!(expr, Expr::Binary { left, .. } if matches!(&**left, Expr::Call { recv: None, name, .. } if name == "size"))
686 );
687 }
688
689 #[test]
690 fn semantic_phrases_are_collected_distinct() {
691 let phrases = collect_semantic_phrases(
692 "select s: semantic(\"alpha\") from docs where semantic(\"alpha\") > 0.5 || nodes exists { where semantic(\"beta\") > 0 } order by semantic(\"gamma\") desc",
693 );
694 assert_eq!(phrases, ["alpha", "beta", "gamma"]);
695 assert!(collect_semantic_phrases("not a query {{").is_empty());
696 assert!(collect_semantic_phrases("from docs").is_empty());
697 }
698
699 #[test]
700 fn hits_peel_the_injected_columns() {
701 let mut o = oqx::Object::new();
702 o.insert(ID_KEY, Value::Str("d_1".into()));
703 o.insert(PATH_KEY, Value::Null);
704 o.insert("layer", Value::Str("canon".into()));
705 let hit = to_hit(Value::Object(o)).unwrap();
706 let ho = hit.as_object().unwrap();
707 assert_eq!(ho.keys().collect::<Vec<_>>(), ["id", "path", "layer"]);
708 assert_eq!(ho.get("path"), Some(&Value::Str(String::new())));
709 let mut o = oqx::Object::new();
711 o.insert(ID_KEY, Value::Str("d_1".into()));
712 o.insert(PATH_KEY, Value::Str("a.md".into()));
713 o.insert("id", Value::Number(7.0));
714 let hit = to_hit(Value::Object(o)).unwrap();
715 let ho = hit.as_object().unwrap();
716 assert_eq!(ho.keys().collect::<Vec<_>>(), ["id", "path"]);
717 assert_eq!(ho.get("id"), Some(&Value::Number(7.0)));
718 }
719
720 #[test]
721 fn a_hit_that_is_not_a_store_row_fails_the_query() {
722 let mut o = oqx::Object::new();
725 o.insert(ID_KEY, Value::Undefined);
726 o.insert(PATH_KEY, Value::Undefined);
727 o.insert(SELF_KEY, Value::Str("/timeline/kickoff.md".into()));
728 let e = to_hit(Value::Object(o)).unwrap_err();
729 assert_eq!(e.code, "filter_invalid");
730 assert_eq!(
731 e.message,
732 "a hit must be a document, block, node or edge row — the query reached a string (\"/timeline/kickoff.md\"); to follow document references held in a property use refs(<field>)"
733 );
734 let mut row = oqx::Object::new();
736 row.insert("doc_id", Value::Str("d_1".into()));
737 row.insert("path", Value::Str("a.md".into()));
738 let row = crate::context::tag_row(row, crate::context::Target::Docs);
739 let mut o = oqx::Object::new();
740 o.insert(ID_KEY, Value::Str("d_1".into()));
741 o.insert(PATH_KEY, Value::Str("a.md".into()));
742 o.insert(SELF_KEY, row);
743 let hit = to_hit(Value::Object(o)).unwrap();
744 assert_eq!(
745 hit.as_object().unwrap().keys().collect::<Vec<_>>(),
746 ["id", "path"]
747 );
748 let parse = |s: &str| oqx::parse_string(s).unwrap();
750 assert!(!may_reach_non_rows(&parse(
751 "from docs where layer == \"canon\""
752 )));
753 assert!(may_reach_non_rows(&parse("from docs follow before")));
754 assert!(may_reach_non_rows(&parse("refs(\"/index.md\") first { }")));
755 }
756
757 #[test]
758 fn distinct_dedups_by_user_projection_first_wins() {
759 let mk = |id: &str, t: &str| {
760 let mut o = oqx::Object::new();
761 o.insert("id", Value::Str(id.into()));
762 o.insert("path", Value::Str("p".into()));
763 o.insert("type", Value::Str(t.into()));
764 Value::Object(o)
765 };
766 let out = dedup_hits_by_projection(vec![mk("1", "a"), mk("2", "b"), mk("3", "a")]);
767 assert_eq!(out.len(), 2);
768 assert_eq!(hit_str(&out[0], "id"), "1");
769 assert_eq!(hit_str(&out[1], "id"), "2");
770 }
771
772 #[test]
773 fn top_level_bounds_must_be_literals() {
774 assert_eq!(const_bound(None, "limit").unwrap(), None);
775 assert_eq!(
776 const_bound(
777 Some(&Expr::Lit {
778 value: Value::Number(3.0),
779 span: oqx::Span::EMPTY
780 }),
781 "limit"
782 )
783 .unwrap(),
784 Some(3)
785 );
786 let e = const_bound(
787 Some(&Expr::Lit {
788 value: Value::Number(-1.0),
789 span: oqx::Span::EMPTY,
790 }),
791 "offset",
792 )
793 .unwrap_err();
794 assert_eq!(e.code, "filter_invalid");
795 assert!(e.message.contains("top-level offset"));
796 }
797}