use std::collections::BTreeMap;
use std::future::Future;
use std::pin::Pin;
use autumn_web::pagination::{Page, PageRequest};
use autumn_web::search::{IndexDefinition, SearchDocument};
use serde::{Deserialize, Serialize};
use crate::error::{SearchError, SearchResult};
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub struct IndexedDocument {
pub document: SearchDocument,
pub embedding: Option<Vec<f32>>,
}
impl IndexedDocument {
#[must_use]
pub const fn new(document: SearchDocument) -> Self {
Self {
document,
embedding: None,
}
}
#[must_use]
pub fn with_embedding(mut self, embedding: Vec<f32>) -> Self {
self.embedding = Some(embedding);
self
}
#[must_use]
pub const fn id(&self) -> i64 {
self.document.id
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct SearchFilter {
pub tenant_id: Option<String>,
pub allowed_ids: Option<Vec<i64>>,
pub excluded_ids: Vec<i64>,
pub equals: BTreeMap<String, String>,
}
pub const TENANT_FILTER_KEY: &str = "tenant_id";
impl SearchFilter {
#[must_use]
pub fn tenant(mut self, tenant_id: impl Into<String>) -> Self {
self.tenant_id = Some(tenant_id.into());
self
}
#[must_use]
pub fn allow_ids(mut self, ids: impl IntoIterator<Item = i64>) -> Self {
self.allowed_ids = Some(ids.into_iter().collect());
self
}
#[must_use]
pub fn exclude_ids(mut self, ids: impl IntoIterator<Item = i64>) -> Self {
self.excluded_ids.extend(ids);
self
}
#[must_use]
pub fn equals(mut self, field: impl Into<String>, value: impl Into<String>) -> Self {
self.equals.insert(field.into(), value.into());
self
}
#[must_use]
pub fn matches_nothing(&self) -> bool {
self.allowed_ids.as_ref().is_some_and(Vec::is_empty)
}
#[must_use]
pub fn intersect(mut self, other: Self) -> Self {
let mut impossible = false;
match (&self.tenant_id, &other.tenant_id) {
(Some(a), Some(b)) if a != b => impossible = true,
(None, Some(b)) => self.tenant_id = Some(b.clone()),
_ => {}
}
self.allowed_ids = match (self.allowed_ids.take(), other.allowed_ids) {
(Some(a), Some(b)) => Some(a.into_iter().filter(|id| b.contains(id)).collect()),
(Some(a), None) => Some(a),
(None, b) => b,
};
self.excluded_ids.extend(other.excluded_ids);
self.excluded_ids.sort_unstable();
self.excluded_ids.dedup();
for (key, value) in other.equals {
match self.equals.get(&key) {
Some(existing) if *existing != value => impossible = true,
_ => {
self.equals.insert(key, value);
}
}
}
if impossible {
self.allowed_ids = Some(Vec::new());
}
self
}
#[must_use]
pub fn permits(&self, document: &SearchDocument) -> bool {
if let Some(tenant) = &self.tenant_id
&& document.tenant_id.as_deref() != Some(tenant.as_str())
{
return false;
}
if let Some(allowed) = &self.allowed_ids
&& !allowed.contains(&document.id)
{
return false;
}
if self.excluded_ids.contains(&document.id) {
return false;
}
for (field, expected) in &self.equals {
let actual = if field == TENANT_FILTER_KEY {
document.tenant_id.as_deref()
} else {
document.field(field)
};
if actual != Some(expected.as_str()) {
return false;
}
}
true
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct KeywordQuery {
pub text: String,
pub filter: SearchFilter,
pub page: PageRequest,
}
impl KeywordQuery {
#[must_use]
pub fn new(text: impl Into<String>, page: PageRequest) -> Self {
Self {
text: text.into(),
filter: SearchFilter::default(),
page,
}
}
#[must_use]
pub fn filter(mut self, filter: SearchFilter) -> Self {
self.filter = filter;
self
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct VectorQuery {
pub vector: Vec<f32>,
pub limit: usize,
pub filter: SearchFilter,
pub min_score: Option<f32>,
}
impl VectorQuery {
#[must_use]
pub fn new(vector: Vec<f32>, limit: usize) -> Self {
Self {
vector,
limit,
filter: SearchFilter::default(),
min_score: None,
}
}
#[must_use]
pub fn filter(mut self, filter: SearchFilter) -> Self {
self.filter = filter;
self
}
#[must_use]
pub const fn min_score(mut self, min_score: f32) -> Self {
self.min_score = Some(min_score);
self
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct SearchHit {
pub index: String,
pub id: i64,
pub score: f32,
}
impl SearchHit {
#[must_use]
pub fn new(index: impl Into<String>, id: i64, score: f32) -> Self {
Self {
index: index.into(),
id,
score,
}
}
}
#[allow(clippy::struct_excessive_bools)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct BackendCapabilities {
pub keyword: bool,
pub vector: bool,
pub weighted_fields: bool,
pub embedding_readback: bool,
pub conditional_index: bool,
}
impl Default for BackendCapabilities {
fn default() -> Self {
Self {
keyword: true,
vector: false,
weighted_fields: false,
embedding_readback: false,
conditional_index: false,
}
}
}
impl BackendCapabilities {
#[must_use]
pub const fn with_vector(mut self, supported: bool) -> Self {
self.vector = supported;
self
}
#[must_use]
pub const fn with_weighted_fields(mut self, supported: bool) -> Self {
self.weighted_fields = supported;
self
}
#[must_use]
pub const fn with_embedding_readback(mut self, supported: bool) -> Self {
self.embedding_readback = supported;
self
}
#[must_use]
pub const fn with_conditional_index(mut self, supported: bool) -> Self {
self.conditional_index = supported;
self
}
#[must_use]
pub const fn with_keyword(mut self, supported: bool) -> Self {
self.keyword = supported;
self
}
}
pub trait SearchBackend: Send + Sync + 'static {
fn name(&self) -> &'static str;
fn capabilities(&self) -> BackendCapabilities {
BackendCapabilities::default()
}
fn ensure_index<'a>(
&'a self,
definition: &'a IndexDefinition,
) -> BoxFuture<'a, SearchResult<()>>;
fn index<'a>(
&'a self,
definition: &'a IndexDefinition,
documents: &'a [IndexedDocument],
) -> BoxFuture<'a, SearchResult<()>>;
fn write_watermark(&self) -> BoxFuture<'_, SearchResult<Option<String>>> {
Box::pin(async { Ok(None) })
}
fn index_unless_newer<'a>(
&'a self,
definition: &'a IndexDefinition,
documents: &'a [IndexedDocument],
watermark: Option<&'a str>,
) -> BoxFuture<'a, SearchResult<()>> {
let _ = watermark;
self.index(definition, documents)
}
fn delete<'a>(
&'a self,
definition: &'a IndexDefinition,
ids: &'a [i64],
) -> BoxFuture<'a, SearchResult<()>>;
fn clear<'a>(&'a self, definition: &'a IndexDefinition) -> BoxFuture<'a, SearchResult<()>>;
fn keyword_search<'a>(
&'a self,
definition: &'a IndexDefinition,
query: &'a KeywordQuery,
) -> BoxFuture<'a, SearchResult<Page<SearchHit>>>;
fn vector_search<'a>(
&'a self,
_definition: &'a IndexDefinition,
_query: &'a VectorQuery,
) -> BoxFuture<'a, SearchResult<Vec<SearchHit>>> {
let name = self.name();
Box::pin(async move {
Err(SearchError::Unsupported {
backend: name,
operation: "vector search",
})
})
}
fn embedding<'a>(
&'a self,
_definition: &'a IndexDefinition,
_id: i64,
_filter: &'a SearchFilter,
) -> BoxFuture<'a, SearchResult<Option<Vec<f32>>>> {
let name = self.name();
Box::pin(async move {
Err(SearchError::Unsupported {
backend: name,
operation: "embedding read-back",
})
})
}
}
#[must_use]
pub fn empty_page(request: &PageRequest) -> Page<SearchHit> {
Page::new(Vec::new(), 0, request)
}
#[cfg(test)]
mod tests {
use super::*;
fn doc(id: i64) -> SearchDocument {
SearchDocument::new("articles", id).with_field("title", 'A', "Hello")
}
#[test]
fn an_empty_filter_permits_everything() {
assert!(SearchFilter::default().permits(&doc(1)));
assert!(!SearchFilter::default().matches_nothing());
}
#[test]
fn an_empty_allowlist_is_a_restriction_not_an_absence() {
let filter = SearchFilter::default().allow_ids(Vec::<i64>::new());
assert!(filter.matches_nothing());
assert!(!filter.permits(&doc(1)));
}
#[test]
fn tenant_filter_requires_an_exact_owner() {
let filter = SearchFilter::default().tenant("acme");
assert!(!filter.permits(&doc(1)), "an untenanted doc is not acme's");
assert!(filter.permits(&doc(1).with_tenant("acme")));
assert!(!filter.permits(&doc(1).with_tenant("globex")));
}
#[test]
fn excluded_ids_win_over_the_allowlist() {
let filter = SearchFilter::default()
.allow_ids([1_i64])
.exclude_ids([1_i64]);
assert!(!filter.permits(&doc(1)));
}
#[test]
fn equals_matches_indexed_fields_and_the_reserved_tenant_key() {
let filter = SearchFilter::default().equals("title", "Hello");
assert!(filter.permits(&doc(1)));
assert!(
!SearchFilter::default()
.equals("title", "Nope")
.permits(&doc(1))
);
assert!(
SearchFilter::default()
.equals(TENANT_FILTER_KEY, "acme")
.permits(&doc(1).with_tenant("acme"))
);
}
#[test]
fn equals_on_a_field_the_document_lacks_never_matches() {
assert!(
!SearchFilter::default()
.equals("missing", "x")
.permits(&doc(1))
);
}
#[test]
fn intersect_narrows_and_never_widens() {
let merged = SearchFilter::default()
.allow_ids([1_i64, 2, 3])
.intersect(SearchFilter::default().allow_ids([2_i64, 4]));
assert_eq!(merged.allowed_ids, Some(vec![2]));
let merged = SearchFilter::default()
.allow_ids([1_i64])
.intersect(SearchFilter::default());
assert_eq!(merged.allowed_ids, Some(vec![1]));
let merged = SearchFilter::default().intersect(SearchFilter::default().allow_ids([1_i64]));
assert_eq!(merged.allowed_ids, Some(vec![1]));
}
#[test]
fn intersect_collapses_conflicting_constraints_to_nothing() {
let merged = SearchFilter::default()
.tenant("acme")
.intersect(SearchFilter::default().tenant("globex"));
assert!(merged.matches_nothing());
let merged = SearchFilter::default()
.equals("state", "published")
.intersect(SearchFilter::default().equals("state", "draft"));
assert!(merged.matches_nothing());
}
#[test]
fn intersect_unions_and_dedups_exclusions() {
let merged = SearchFilter::default()
.exclude_ids([1_i64, 2])
.intersect(SearchFilter::default().exclude_ids([2_i64, 3]));
assert_eq!(merged.excluded_ids, vec![1, 2, 3]);
}
#[test]
fn default_capabilities_claim_only_keyword_search() {
let caps = BackendCapabilities::default();
assert!(caps.keyword);
assert!(!caps.vector);
assert!(!caps.weighted_fields);
assert!(!caps.embedding_readback);
}
#[test]
fn capabilities_are_declarable_from_outside_this_crate() {
let caps = BackendCapabilities::default()
.with_vector(true)
.with_weighted_fields(true)
.with_embedding_readback(true);
assert!(caps.keyword);
assert!(caps.vector);
assert!(caps.weighted_fields);
assert!(caps.embedding_readback);
assert!(!BackendCapabilities::default().with_keyword(false).keyword);
}
#[test]
fn empty_page_reports_a_zero_total() {
let page = empty_page(&PageRequest::default());
assert!(page.content.is_empty());
assert_eq!(page.total_elements, 0);
}
}