use oqx::Consumer;
use oqx::ast::{Expr, Follow, OpNode, Query, SelectItem, Subquery, Where};
use oqx::{Plan, QueryPlanner, Value, partition_pushable, residual_query};
use rusqlite::Connection;
use rusqlite::types::Value as SqlValue;
use crate::context::{Target, fetch_rows, tag_rows};
use crate::translate::{RESERVED_DOC_BASENAMES, TranslateCtx, translate_predicate};
fn aliases(t: Target) -> (&'static str, &'static str) {
match t {
Target::Docs => ("d", "d"),
Target::Blocks => ("b", "d"),
Target::Nodes => ("n", "d"),
Target::Edges => ("e", "d"),
}
}
fn from_clause(t: Target) -> &'static str {
match t {
Target::Docs => "docs d",
Target::Blocks => "blocks b JOIN docs d ON d.doc_id = b.doc_id",
Target::Nodes => "nodes n JOIN docs d ON d.doc_id = n.doc_id",
Target::Edges => "edges e JOIN docs d ON d.doc_id = e.src_doc",
}
}
fn columns(t: Target) -> &'static str {
match t {
Target::Docs => "d.*",
Target::Blocks => "b.*, d.path AS __path",
Target::Nodes => "n.*, d.path AS __path",
Target::Edges => "e.*, d.path AS __path",
}
}
fn order_clause(t: Target) -> &'static str {
match t {
Target::Docs => "d.path, d.doc_id",
Target::Blocks => "d.path, b.block_id",
Target::Nodes => "d.path, n.node_id",
Target::Edges => "d.path, e.edge_id",
}
}
fn guards(t: Target) -> &'static str {
match t {
Target::Docs => "d.repo_id = ? AND d.deleted_commit IS NULL",
Target::Blocks => "b.repo_id = ? AND b.deleted_commit IS NULL AND d.deleted_commit IS NULL",
Target::Nodes => "n.repo_id = ? AND d.deleted_commit IS NULL",
Target::Edges => "e.repo_id = ? AND e.to_commit IS NULL AND d.deleted_commit IS NULL",
}
}
fn root_target(source: &Expr) -> Option<Target> {
match source {
Expr::Ident { name } => Target::parse(name),
Expr::Member { recv, name } => match &**recv {
Expr::Ident { name: r } if r == "$repo" => Target::parse(name),
_ => None,
},
_ => None,
}
}
fn residual_may_raise(w: &Where, target: Target) -> bool {
where_may_raise(w, target, true)
}
fn where_may_raise(w: &Where, target: Target, root: bool) -> bool {
match w {
Where::And { parts } | Where::Or { parts } => {
parts.iter().any(|p| where_may_raise(p, target, root))
}
Where::Not { expr } => where_may_raise(expr, target, root),
Where::Scalar { expr } => expr_may_raise(expr, target, root),
Where::Op(op) => op_may_raise(op, target, root),
}
}
fn op_may_raise(op: &OpNode, target: Target, root: bool) -> bool {
op.op == Consumer::Single
|| expr_may_raise(&op.receiver, target, root)
|| subquery_may_raise(&op.sub, target)
}
fn subquery_may_raise(sub: &Subquery, target: Target) -> bool {
let inner = |e: &Expr| expr_may_raise(e, target, false);
sub.from.iter().any(inner)
|| sub
.r#where
.as_ref()
.is_some_and(|w| where_may_raise(w, target, false))
|| sub.select.iter().any(|item| match item {
SelectItem::Field { expr, lift, .. } => *lift > 0 || inner(expr),
SelectItem::Collect { op, .. } => op_may_raise(op, target, false),
})
|| sub.order_by.iter().flatten().any(|o| inner(&o.expr))
|| sub
.follow
.as_ref()
.is_some_and(|f| follow_may_raise(f, target))
|| sub.limit.as_ref().is_some_and(inner)
|| sub.offset.as_ref().is_some_and(inner)
}
fn follow_may_raise(f: &Follow, target: Target) -> bool {
let inner = |e: &Expr| expr_may_raise(e, target, false);
inner(&f.receiver)
|| f.r#where.as_ref().is_some_and(inner)
|| f.frontier.as_ref().is_some_and(inner)
|| f.by.as_ref().is_some_and(inner)
}
fn is_reserved(name: &str) -> bool {
RESERVED_DOC_BASENAMES.contains(&name)
}
fn expr_may_raise(e: &Expr, target: Target, root: bool) -> bool {
let again = |e: &Expr| expr_may_raise(e, target, root);
match e {
Expr::Lit(_) | Expr::Binding { .. } => false,
Expr::Ident { name } => root && target == Target::Docs && is_reserved(name),
Expr::Outer { .. } | Expr::Call { .. } => true,
Expr::Member { recv, name } => {
(is_reserved(name) && matches!(&**recv, Expr::Ident { name } if name == "doc"))
|| again(recv)
}
Expr::Index { recv, index } => again(recv) || again(index),
Expr::Unary { expr, .. } => again(expr),
Expr::Binary { left, right, .. }
| Expr::Logical { left, right, .. }
| Expr::In { left, right } => again(left) || again(right),
Expr::Range { lo, hi, .. } => {
lo.as_deref().is_some_and(again) || hi.as_deref().is_some_and(again)
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct Compiled {
pub target: Target,
pub sql: String,
pub params: Vec<SqlValue>,
pub residual: Query,
}
#[must_use]
pub fn compile(query: &Query, params: &[Value], repo_id: &str) -> Option<Compiled> {
if query.follow.is_some() || !query.from.is_empty() {
return None;
}
let target = root_target(&query.source)?;
let (self_alias, doc_alias) = aliases(target);
let ctx = TranslateCtx {
target,
self_alias,
doc_alias,
params,
};
let (pushed, residual) = partition_pushable(query.r#where.as_ref(), |e| {
translate_predicate(e, &ctx).is_some()
});
if pushed.is_empty() {
return None;
}
if residual
.as_ref()
.is_some_and(|w| residual_may_raise(w, target))
{
return None;
}
let mut where_sql = guards(target).to_owned();
let mut sql_params = vec![SqlValue::Text(repo_id.to_owned())];
for e in &pushed {
let frag = translate_predicate(e, &ctx).expect("accepted by partition_pushable");
where_sql.push_str(" AND (");
where_sql.push_str(&frag.sql);
where_sql.push(')');
sql_params.extend(frag.params);
}
let sql = format!(
"SELECT {} FROM {} WHERE {where_sql} ORDER BY {}",
columns(target),
from_clause(target),
order_clause(target)
);
Some(Compiled {
target,
sql,
params: sql_params,
residual: residual_query(query, residual),
})
}
pub struct SqlitePlanner<'a> {
conn: &'a Connection,
repo_id: String,
}
impl<'a> SqlitePlanner<'a> {
#[must_use]
pub fn new(conn: &'a Connection, repo_id: &str) -> Self {
Self {
conn,
repo_id: repo_id.to_owned(),
}
}
pub fn try_plan(&self, query: &Query, params: &[Value]) -> rusqlite::Result<Option<Plan>> {
let Some(compiled) = compile(query, params, &self.repo_id) else {
return Ok(None);
};
let rows = fetch_rows(self.conn, &compiled.sql, &compiled.params)?;
Ok(Some(Plan::new(
tag_rows(rows, compiled.target),
compiled.residual,
)))
}
}
impl QueryPlanner for SqlitePlanner<'_> {
fn plan(&self, query: &Query, params: &[Value]) -> Option<Plan> {
self.try_plan(query, params).ok().flatten()
}
}
#[cfg(test)]
mod tests {
use super::*;
use oqx::ROWS_ROOT;
use oqx::ast::Where;
fn parse(src: &str) -> Query {
oqx::parse_string(src).expect("parses")
}
fn text(s: &str) -> SqlValue {
SqlValue::Text(s.to_owned())
}
#[test]
fn a_pushable_scan_compiles_to_one_statement_in_root_order() {
let c =
compile(&parse("from docs where $path == \"index.md\""), &[], "r_1").expect("planned");
assert_eq!(c.target, Target::Docs);
assert_eq!(
c.sql,
"SELECT d.* FROM docs d WHERE d.repo_id = ? AND d.deleted_commit IS NULL AND ((d.path IS ?)) ORDER BY d.path, d.doc_id"
);
assert_eq!(c.params, vec![text("r_1"), text("index.md")]);
assert_eq!(
c.residual.source,
Expr::Ident {
name: ROWS_ROOT.to_owned()
}
);
assert_eq!(c.residual.r#where, None);
}
#[test]
fn every_target_has_its_join_columns_guards_and_order() {
let b = compile(
&parse("from blocks where $path.startsWith(\"lab/\")"),
&[],
"r",
)
.unwrap();
assert_eq!(
b.sql,
"SELECT b.*, d.path AS __path FROM blocks b JOIN docs d ON d.doc_id = b.doc_id \
WHERE b.repo_id = ? AND b.deleted_commit IS NULL AND d.deleted_commit IS NULL \
AND ((substr(d.path, 1, length(?)) = ?)) ORDER BY d.path, b.block_id"
);
assert_eq!(b.params, vec![text("r"), text("lab/"), text("lab/")]);
let n = compile(
&parse("$repo.nodes count { where kind == \"md:task\" }"),
&[],
"r",
)
.unwrap();
assert_eq!(n.target, Target::Nodes);
assert!(n.sql.starts_with(
"SELECT n.*, d.path AS __path FROM nodes n JOIN docs d ON d.doc_id = n.doc_id WHERE n.repo_id = ? AND d.deleted_commit IS NULL AND ((n.kind IS ?))"
));
assert!(n.sql.ends_with("ORDER BY d.path, n.node_id"));
let e = compile(
&parse("from edges where predicate == \"references\""),
&[],
"r",
)
.unwrap();
assert!(e.sql.starts_with(
"SELECT e.*, d.path AS __path FROM edges e JOIN docs d ON d.doc_id = e.src_doc WHERE e.repo_id = ? AND e.to_commit IS NULL AND d.deleted_commit IS NULL AND ((e.predicate IS ?))"
));
assert!(e.sql.ends_with("ORDER BY d.path, e.edge_id"));
}
#[test]
fn mixed_conjunctions_push_the_translatable_parts_and_keep_the_rest() {
let q = parse(
"from docs where $path.startsWith(\"processes/\") && nodes exists { where kind == \"md:task\" } && layer == \"canon\"",
);
let c = compile(&q, &[], "r").expect("planned");
assert!(
c.sql.contains("(substr(d.path, 1, length(?)) = ?)"),
"{}",
c.sql
);
assert!(c.sql.contains("p.key = 'layer'"), "{}", c.sql);
assert_eq!(
c.params,
vec![
text("r"),
text("processes/"),
text("processes/"),
text("canon")
]
);
assert!(
matches!(c.residual.r#where, Some(Where::Op(_))),
"{:?}",
c.residual.r#where
);
assert!(c.residual.from.is_empty());
assert_eq!(c.residual.select, q.select);
assert_eq!(c.residual.consumer, q.consumer);
}
#[test]
fn declined_shapes_return_none() {
assert!(compile(&parse("from docs"), &[], "r").is_none());
assert!(compile(&parse("from docs where era in 800..1680"), &[], "r").is_none());
assert!(compile(&parse("from docs where !verified"), &[], "r").is_none());
assert!(
compile(
&parse("from docs where $path == \"a\" || $path == \"b\""),
&[],
"r"
)
.is_none()
);
assert!(
compile(
&parse("from docs where nodes exists { where kind == \"md:task\" }"),
&[],
"r"
)
.is_none()
);
assert!(
compile(
&parse("from docs where $path == \"a.md\" follow distinct doc.out"),
&[],
"r"
)
.is_none()
);
assert!(compile(&parse("from things where $path == \"a.md\""), &[], "r").is_none());
assert!(compile(&parse("from $repo where $path == \"a.md\""), &[], "r").is_none());
assert!(compile(&parse("from docs.nodes where kind == \"x\""), &[], "r").is_none());
let mut q = parse("from docs where $path == \"a.md\"");
q.from.push(Expr::Ident {
name: "nodes".to_owned(),
});
assert!(compile(&q, &[], "r").is_none());
}
#[test]
fn a_residual_that_could_raise_declines_the_whole_query() {
let declined = |src: &str| {
assert!(
compile(&parse(src), &[], "r").is_none(),
"should decline: {src}"
);
};
let planned = |src: &str| {
assert!(
compile(&parse(src), &[], "r").is_some(),
"should plan: {src}"
);
};
declined("from docs where path == \"x\" && $path == \"nope.md\"");
declined("from docs where $path == \"nope.md\" && !body");
declined("from blocks where $path == \"x\" && doc.path == \"y\"");
declined("from blocks where $path == \"x\" && nodes exists { where doc.path == \"y\" }");
declined("from docs where $path.matches(\"[\") && $path == \"nope.md\"");
declined("from docs where nope(\"x\") && $path == \"nope.md\"");
declined("from docs where $path == \"x\" && size(tags) > 1");
declined("from docs where $path == \"x\" && nodes exists { where name.lower() == \"a\" }");
declined(
"from docs where $path == \"x\" && nodes count { where kind == \"a\" order by size(name) } > 1",
);
declined("from docs where $path == \"x\" && ^slug == \"y\"");
declined("from docs where $path == \"x\" && nodes exists { where name == ^title }");
declined("from docs where $path == \"x\" && nodes collect { ^first_task: name }");
declined(
"from docs where $path == \"x\" && blocks exists { select t: nodes single { where kind == \"md:task\" } }",
);
planned("from docs where $path.startsWith(\"lab/\") && $path == \"x\"");
planned("from docs where $path == \"x\" && era in 800..1680");
planned("from docs where $path == \"x\" && !verified");
planned("from docs where $path == \"x\" && (layer == \"a\" || layer == \"b\")");
planned("from docs where $path == \"x\" && verified == true");
planned("from docs where $path == \"x\" && nodes exists { where kind == \"md:task\" }");
planned(
"from docs where $path == \"x\" && nodes count { where kind == \"md:task\" limit 5 } > 1",
);
planned("from blocks where $path == \"x\" && !path");
planned("from docs where $path == \"x\" && nodes exists { where path == \"y\" }");
planned("from docs where $path == \"x\" && frontmatter.path == \"y\"");
}
#[test]
fn a_top_level_not_or_or_is_the_whole_residual_and_declines() {
assert!(compile(&parse("from docs where !($path == \"a\")"), &[], "r").is_none());
assert!(
compile(
&parse("from docs where ($path == \"a\" || $path == \"b\") && layer == \"canon\""),
&[],
"r"
)
.is_some()
);
}
}