use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Query {
pub root: QueryNode,
}
impl Query {
pub fn new(root: QueryNode) -> Self {
Self { root }
}
pub fn field_eq(field: impl Into<String>, value: impl Into<String>) -> Self {
Self::new(QueryNode::Field(FieldQuery {
field: field.into(),
operator: FieldOperator::Equals,
value: QueryValue::Text(value.into()),
}))
}
pub fn tags(field: impl Into<String>, values: Vec<String>) -> Self {
Self::new(QueryNode::Field(FieldQuery {
field: field.into(),
operator: FieldOperator::In,
value: QueryValue::Tags(values),
}))
}
pub fn numeric_range(field: impl Into<String>, min: Option<f64>, max: Option<f64>) -> Self {
Self::new(QueryNode::Field(FieldQuery {
field: field.into(),
operator: FieldOperator::Range,
value: QueryValue::NumericRange { min, max },
}))
}
pub fn text_search(field: impl Into<String>, text: impl Into<String>) -> Self {
Self::new(QueryNode::Field(FieldQuery {
field: field.into(),
operator: FieldOperator::Contains,
value: QueryValue::Text(text.into()),
}))
}
pub fn prefix(field: impl Into<String>, prefix: impl Into<String>) -> Self {
Self::new(QueryNode::Field(FieldQuery {
field: field.into(),
operator: FieldOperator::Prefix,
value: QueryValue::Text(prefix.into()),
}))
}
pub fn fuzzy(field: impl Into<String>, text: impl Into<String>) -> Self {
Self::new(QueryNode::Field(FieldQuery {
field: field.into(),
operator: FieldOperator::Fuzzy,
value: QueryValue::Text(text.into()),
}))
}
pub fn vector(field: impl Into<String>, vector: Vec<f32>, k: usize) -> Self {
Self::new(QueryNode::Vector(VectorQuery {
field: field.into(),
vector,
k,
}))
}
pub fn vector_filtered(
filter: Query,
field: impl Into<String>,
vector: Vec<f32>,
k: usize,
) -> Self {
Self::new(QueryNode::And(vec![
filter.root,
QueryNode::Vector(VectorQuery {
field: field.into(),
vector,
k,
}),
]))
}
pub fn and(self, other: Query) -> Self {
Self::new(QueryNode::And(vec![self.root, other.root]))
}
pub fn or(self, other: Query) -> Self {
Self::new(QueryNode::Or(vec![self.root, other.root]))
}
pub fn negate(self) -> Self {
Self::new(QueryNode::Not(Box::new(self.root)))
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum QueryNode {
Field(FieldQuery),
And(Vec<QueryNode>),
Or(Vec<QueryNode>),
Not(Box<QueryNode>),
Vector(VectorQuery),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct VectorQuery {
pub field: String,
pub vector: Vec<f32>,
pub k: usize,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct FieldQuery {
pub field: String,
pub operator: FieldOperator,
pub value: QueryValue,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum FieldOperator {
Equals,
Contains,
Range,
In,
Prefix,
Fuzzy,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum QueryValue {
Text(String),
Numeric(f64),
NumericRange { min: Option<f64>, max: Option<f64> },
Tags(Vec<String>),
Boolean(bool),
}
#[derive(Default)]
pub struct QueryBuilder {
nodes: Vec<QueryNode>,
}
impl QueryBuilder {
pub fn new() -> Self {
Self { nodes: Vec::new() }
}
pub fn field_eq(mut self, field: impl Into<String>, value: impl Into<String>) -> Self {
self.nodes.push(QueryNode::Field(FieldQuery {
field: field.into(),
operator: FieldOperator::Equals,
value: QueryValue::Text(value.into()),
}));
self
}
pub fn numeric_eq(mut self, field: impl Into<String>, value: f64) -> Self {
self.nodes.push(QueryNode::Field(FieldQuery {
field: field.into(),
operator: FieldOperator::Equals,
value: QueryValue::Numeric(value),
}));
self
}
pub fn numeric_range(mut self, field: impl Into<String>, min: Option<f64>, max: Option<f64>) -> Self {
self.nodes.push(QueryNode::Field(FieldQuery {
field: field.into(),
operator: FieldOperator::Range,
value: QueryValue::NumericRange { min, max },
}));
self
}
pub fn tags(mut self, field: impl Into<String>, values: Vec<String>) -> Self {
self.nodes.push(QueryNode::Field(FieldQuery {
field: field.into(),
operator: FieldOperator::In,
value: QueryValue::Tags(values),
}));
self
}
pub fn contains(mut self, field: impl Into<String>, text: impl Into<String>) -> Self {
self.nodes.push(QueryNode::Field(FieldQuery {
field: field.into(),
operator: FieldOperator::Contains,
value: QueryValue::Text(text.into()),
}));
self
}
pub fn prefix(mut self, field: impl Into<String>, prefix: impl Into<String>) -> Self {
self.nodes.push(QueryNode::Field(FieldQuery {
field: field.into(),
operator: FieldOperator::Prefix,
value: QueryValue::Text(prefix.into()),
}));
self
}
pub fn fuzzy(mut self, field: impl Into<String>, text: impl Into<String>) -> Self {
self.nodes.push(QueryNode::Field(FieldQuery {
field: field.into(),
operator: FieldOperator::Fuzzy,
value: QueryValue::Text(text.into()),
}));
self
}
pub fn build_and(self) -> Query {
if self.nodes.is_empty() {
Query::new(QueryNode::Field(FieldQuery {
field: "*".to_string(),
operator: FieldOperator::Equals,
value: QueryValue::Text("*".to_string()),
}))
} else if self.nodes.len() == 1 {
Query::new(self.nodes.into_iter().next().unwrap())
} else {
Query::new(QueryNode::And(self.nodes))
}
}
pub fn build_or(self) -> Query {
if self.nodes.is_empty() {
Query::new(QueryNode::Field(FieldQuery {
field: "*".to_string(),
operator: FieldOperator::Equals,
value: QueryValue::Text("*".to_string()),
}))
} else if self.nodes.len() == 1 {
Query::new(self.nodes.into_iter().next().unwrap())
} else {
Query::new(QueryNode::Or(self.nodes))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_simple_field_query() {
let query = Query::field_eq("name", "Alice");
assert_eq!(
query.root,
QueryNode::Field(FieldQuery {
field: "name".to_string(),
operator: FieldOperator::Equals,
value: QueryValue::Text("Alice".to_string()),
})
);
}
#[test]
fn test_and_query() {
let query = Query::field_eq("name", "Alice")
.and(Query::numeric_range("age", Some(25.0), Some(40.0)));
match query.root {
QueryNode::And(nodes) => {
assert_eq!(nodes.len(), 2);
}
_ => panic!("Expected And node"),
}
}
#[test]
fn test_or_query() {
let query = Query::field_eq("status", "active")
.or(Query::field_eq("status", "pending"));
match query.root {
QueryNode::Or(nodes) => {
assert_eq!(nodes.len(), 2);
}
_ => panic!("Expected Or node"),
}
}
#[test]
fn test_not_query() {
let query = Query::field_eq("deleted", "true").negate();
match query.root {
QueryNode::Not(_) => {}
_ => panic!("Expected Not node"),
}
}
#[test]
fn test_tag_query() {
let query = Query::tags("tags", vec!["rust".to_string(), "database".to_string()]);
match query.root {
QueryNode::Field(FieldQuery { field, operator, value }) => {
assert_eq!(field, "tags");
assert_eq!(operator, FieldOperator::In);
assert_eq!(value, QueryValue::Tags(vec!["rust".to_string(), "database".to_string()]));
}
_ => panic!("Expected Field node"),
}
}
#[test]
fn test_query_builder_and() {
let query = QueryBuilder::new()
.field_eq("name", "Alice")
.numeric_range("age", Some(25.0), Some(40.0))
.tags("tags", vec!["rust".to_string()])
.build_and();
match query.root {
QueryNode::And(nodes) => {
assert_eq!(nodes.len(), 3);
}
_ => panic!("Expected And node"),
}
}
#[test]
fn test_query_builder_or() {
let query = QueryBuilder::new()
.field_eq("status", "active")
.field_eq("status", "pending")
.build_or();
match query.root {
QueryNode::Or(nodes) => {
assert_eq!(nodes.len(), 2);
}
_ => panic!("Expected Or node"),
}
}
#[test]
fn test_complex_query() {
let alice_query = Query::field_eq("name", "Alice")
.and(Query::numeric_range("age", Some(25.0), Some(40.0)));
let bob_query = Query::field_eq("name", "Bob")
.and(Query::numeric_range("age", Some(30.0), Some(50.0)));
let query = alice_query.or(bob_query);
match query.root {
QueryNode::Or(nodes) => {
assert_eq!(nodes.len(), 2);
for node in nodes {
match node {
QueryNode::And(inner_nodes) => {
assert_eq!(inner_nodes.len(), 2);
}
_ => panic!("Expected And node"),
}
}
}
_ => panic!("Expected Or node"),
}
}
#[test]
fn test_prefix_query() {
let query = Query::prefix("email", "admin@");
match query.root {
QueryNode::Field(FieldQuery { operator, .. }) => {
assert_eq!(operator, FieldOperator::Prefix);
}
_ => panic!("Expected Field node"),
}
}
#[test]
fn test_fuzzy_query() {
let query = Query::fuzzy("name", "alice");
match query.root {
QueryNode::Field(FieldQuery { operator, .. }) => {
assert_eq!(operator, FieldOperator::Fuzzy);
}
_ => panic!("Expected Field node"),
}
}
#[test]
fn test_empty_builder_and() {
let query = QueryBuilder::new().build_and();
match query.root {
QueryNode::Field(FieldQuery { field, .. }) => {
assert_eq!(field, "*");
}
_ => panic!("Expected Field node"),
}
}
}