use crate::collection::VectorCollection;
use crate::filter::{Condition, Filter};
use crate::observer::QueryOperationKind;
use crate::point::SearchResult;
use crate::{Error, Result};
use super::query_engine::RawGateOutcome;
use super::Database;
#[derive(Debug, Clone, Copy)]
pub enum GatedRead<'a> {
Dense {
query: &'a [f32],
k: usize,
ef: Option<usize>,
quality: Option<crate::SearchQuality>,
filter: Option<&'a Filter>,
},
Text {
query: &'a str,
k: usize,
filter: Option<&'a Filter>,
},
Hybrid {
vector: &'a [f32],
text: &'a str,
k: usize,
alpha: Option<f32>,
filter: Option<&'a Filter>,
},
}
impl GatedRead<'_> {
fn operation_kind(&self) -> QueryOperationKind {
match self {
GatedRead::Dense { .. } => QueryOperationKind::VectorSearch,
GatedRead::Text { .. } => QueryOperationKind::TextSearch,
GatedRead::Hybrid { .. } => QueryOperationKind::HybridSearch,
}
}
}
#[must_use]
pub(crate) fn scope_to_core_filter(condition: crate::velesql::Condition) -> Filter {
Filter::new(Condition::from(condition))
}
fn and_filters(caller: Option<&Filter>, scope: Option<Filter>) -> Option<Filter> {
match (caller, scope) {
(None, None) => None,
(Some(c), None) => Some(c.clone()),
(None, Some(s)) => Some(s),
(Some(c), Some(s)) => Some(Filter::new(Condition::And {
conditions: vec![c.condition.clone(), s.condition],
})),
}
}
fn run_dense(
coll: &VectorCollection,
query: &[f32],
k: usize,
ef: Option<usize>,
quality: Option<crate::SearchQuality>,
filter: Option<Filter>,
) -> Result<Vec<SearchResult>> {
match filter {
Some(f) => coll.search_with_filter(query, k, &f),
None => match (ef, quality) {
(Some(ef), _) => coll.search_with_ef(query, k, ef),
(None, Some(q)) => coll.search_with_quality(query, k, q),
(None, None) => coll.search(query, k),
},
}
}
fn dispatch_gated_read(
coll: &VectorCollection,
read: GatedRead<'_>,
scope_filter: Option<Filter>,
) -> Result<Vec<SearchResult>> {
match read {
GatedRead::Dense {
query,
k,
ef,
quality,
filter,
} => run_dense(
coll,
query,
k,
ef,
quality,
and_filters(filter, scope_filter),
),
GatedRead::Text { query, k, filter } => match and_filters(filter, scope_filter) {
Some(f) => coll.text_search_with_filter(query, k, &f),
None => coll.text_search(query, k),
},
GatedRead::Hybrid {
vector,
text,
k,
alpha,
filter,
} => match and_filters(filter, scope_filter) {
Some(f) => coll.hybrid_search_with_filter(vector, text, k, alpha, &f),
None => coll.hybrid_search(vector, text, k, alpha),
},
}
}
impl Database {
pub fn authorize_read(
&self,
collection: &str,
operation: QueryOperationKind,
principal: Option<&str>,
tenant_hint: Option<&str>,
) -> Result<Option<Filter>> {
match self.read_gate_raw(collection, operation, principal, tenant_hint)? {
RawGateOutcome::Allow => Ok(None),
RawGateOutcome::Deny(err) => Err(err),
RawGateOutcome::Scope(scope) => Ok(scope.filter.map(scope_to_core_filter)),
}
}
pub fn gated_search(
&self,
collection: &str,
principal: Option<&str>,
tenant_hint: Option<&str>,
read: GatedRead<'_>,
) -> Result<Vec<SearchResult>> {
let scope_filter =
match self.read_gate_raw(collection, read.operation_kind(), principal, tenant_hint)? {
RawGateOutcome::Allow => None,
RawGateOutcome::Deny(err) => return Err(err),
RawGateOutcome::Scope(scope) => scope.filter.map(scope_to_core_filter),
};
let coll = self
.get_vector_collection(collection)
.ok_or_else(|| Error::CollectionNotFound(collection.to_string()))?;
dispatch_gated_read(&coll, read, scope_filter)
}
}