Skip to main content

dynomite_search/
lib.rs

1//! RediSearch FT.* command surface for the
2//! [Dynomite](dynomite) cluster engine.
3//!
4//! `dynomite-search` is the layered search surface that sits
5//! on top of `dynomite-engine`. It owns:
6//!
7//! * the per-server [vector index registry](registry),
8//! * the [schema types](schema) that compile FT.CREATE
9//!   payloads into engine-level shapes,
10//! * the [FT.* dispatch layer](ft) plus the
11//!   [filter-expression grammar](ft_filter),
12//! * the cluster-coordinated k-NN [broadcast FSM](query_fsm),
13//! * the on-the-wire [codec](wire) the engine's DNODE plane
14//!   uses to fan a query out to every primary peer.
15//!
16//! The crate is designed to be wired into a Dynomite
17//! [`ServerBuilder`](dynomite::embed::ServerBuilder) via the
18//! [`CommandExtension`](dynomite::embed::CommandExtension)
19//! hook. The [`install`] helper does this in one call;
20//! [`SearchExtension`] is the underlying impl for embedders
21//! who want finer control.
22//!
23//! # Quickstart
24//!
25//! ```no_run
26//! use dynomite::embed::ServerBuilder;
27//! use dynomite::conf::DataStore;
28//! # tokio::runtime::Builder::new_multi_thread().enable_all().build().unwrap().block_on(async {
29//! let mut builder = ServerBuilder::new("dyn_o_mite")
30//!     .listen("127.0.0.1:0".parse().unwrap())
31//!     .dyn_listen("127.0.0.1:0".parse().unwrap())
32//!     .data_store(DataStore::Redis)
33//!     .servers(vec![dynomite::conf::ConfServer::parse("127.0.0.1:6379:1").unwrap()])
34//!     .tokens_str("0");
35//! let registry = dynomite_search::install(&mut builder);
36//! let handle = builder.build().unwrap().start().await.unwrap();
37//! let _ = registry; // hand off to admin tools, tests, ...
38//! handle.shutdown().await.unwrap();
39//! # });
40//! ```
41
42#![forbid(unsafe_code)]
43#![warn(missing_docs)]
44
45pub mod ft;
46pub mod ft_filter;
47pub mod query_fsm;
48pub mod registry;
49pub mod schema;
50pub mod sugest;
51pub mod sugest_registry;
52pub mod wire;
53
54use std::sync::Arc;
55
56use dynomite::embed::{CommandExtension, HsetOutcome, ServerBuilder};
57use dynomite::msg::MsgType;
58
59pub use crate::registry::{
60    RegistryError, TextFieldIndex, TextHit, TextRegexApproxResult, TextRegexResult, VectorRegistry,
61    VectorTable, VectorTableInfo,
62};
63pub use crate::schema::{
64    DistanceMetric, IndexAlgorithm, MetadataField, MetadataFieldType, VectorSchema, VectorType,
65};
66pub use crate::sugest::{SuggestionDict, SuggestionEntry, SuggestionHit};
67pub use crate::sugest_registry::SuggestionRegistry;
68
69/// [`CommandExtension`] implementation that routes FT.*
70/// commands and the HSET interception path through a shared
71/// [`VectorRegistry`] and [`SuggestionRegistry`].
72///
73/// Every cloneable handle to a `SearchExtension` references
74/// the same registries; embedders who want to inspect the
75/// live FT.* surface (admin paths, tests) can clone the
76/// registry handles out via [`SearchExtension::registry`]
77/// and [`SearchExtension::suggestions`].
78#[derive(Clone, Debug)]
79pub struct SearchExtension {
80    registry: Arc<VectorRegistry>,
81    suggestions: Arc<SuggestionRegistry>,
82}
83
84impl SearchExtension {
85    /// Wrap an existing registry in a [`SearchExtension`].
86    /// The suggestion-dictionary registry is allocated
87    /// fresh; callers that want to share it explicitly can
88    /// use [`Self::with_suggestions`].
89    #[must_use]
90    pub fn new(registry: Arc<VectorRegistry>) -> Self {
91        Self {
92            registry,
93            suggestions: Arc::new(SuggestionRegistry::new()),
94        }
95    }
96
97    /// Wrap both registries in a [`SearchExtension`].
98    #[must_use]
99    pub fn with_suggestions(
100        registry: Arc<VectorRegistry>,
101        suggestions: Arc<SuggestionRegistry>,
102    ) -> Self {
103        Self {
104            registry,
105            suggestions,
106        }
107    }
108
109    /// Borrow the wrapped vector-index registry.
110    #[must_use]
111    pub fn registry(&self) -> &Arc<VectorRegistry> {
112        &self.registry
113    }
114
115    /// Borrow the wrapped suggestion-dictionary registry.
116    #[must_use]
117    pub fn suggestions(&self) -> &Arc<SuggestionRegistry> {
118        &self.suggestions
119    }
120}
121
122impl Default for SearchExtension {
123    fn default() -> Self {
124        Self {
125            registry: Arc::new(VectorRegistry::new()),
126            suggestions: Arc::new(SuggestionRegistry::new()),
127        }
128    }
129}
130
131impl CommandExtension for SearchExtension {
132    fn handles_msg_type(&self, ty: MsgType) -> bool {
133        matches!(
134            ty,
135            MsgType::ReqRedisFtCreate
136                | MsgType::ReqRedisFtSearch
137                | MsgType::ReqRedisFtInfo
138                | MsgType::ReqRedisFtList
139                | MsgType::ReqRedisFtDropindex
140                | MsgType::ReqRedisFtRegex
141                | MsgType::ReqRedisFtSugadd
142                | MsgType::ReqRedisFtSugget
143                | MsgType::ReqRedisFtSugdel
144                | MsgType::ReqRedisFtSuglen
145                | MsgType::ReqRedisFtUnknown
146        )
147    }
148
149    fn try_dispatch(&self, args: &[&[u8]]) -> Option<Vec<u8>> {
150        // FT.SUG* commands route through the suggestion
151        // registry; everything else lands on the vector-
152        // index dispatcher. The keyword is `args[0]`.
153        if let Some(head) = args.first() {
154            let mut upper = [0u8; 16];
155            let n = head.len().min(upper.len());
156            for (i, &b) in head.iter().take(n).enumerate() {
157                upper[i] = b.to_ascii_uppercase();
158            }
159            if matches!(
160                &upper[..n],
161                b"FT.SUGADD" | b"FT.SUGGET" | b"FT.SUGDEL" | b"FT.SUGLEN"
162            ) {
163                return Some(crate::ft::dispatch_sugest(&self.suggestions, args));
164            }
165        }
166        Some(crate::ft::dispatch(&self.registry, args))
167    }
168
169    fn try_intercept_hset(&self, args: &[&[u8]]) -> HsetOutcome {
170        match crate::ft::maybe_index_hset(&self.registry, args) {
171            Ok(Some(_)) => HsetOutcome::Absorbed,
172            Ok(None) => HsetOutcome::NotIndexed,
173            Err(e) => HsetOutcome::Error(format!("{e}")),
174        }
175    }
176}
177
178/// Wire the FT.* command surface into `builder` via the
179/// [`CommandExtension`] hook. Returns an [`Arc`] handle to the
180/// shared [`VectorRegistry`] so the caller can hold a cloneable
181/// reference for admin paths / tests.
182///
183/// Equivalent to constructing a fresh [`SearchExtension`],
184/// installing it on the builder, and returning the registry
185/// handle:
186///
187/// ```no_run
188/// use std::sync::Arc;
189/// use dynomite::embed::ServerBuilder;
190/// use dynomite_search::{SearchExtension, VectorRegistry};
191/// let mut b = ServerBuilder::new("p");
192/// let registry = Arc::new(VectorRegistry::new());
193/// let ext = SearchExtension::new(registry.clone());
194/// b = b.with_command_extension(Arc::new(ext));
195/// ```
196pub fn install(builder: &mut ServerBuilder) -> Arc<VectorRegistry> {
197    let ext = SearchExtension::default();
198    let registry = Arc::clone(ext.registry());
199    builder.set_command_extension(Arc::new(ext));
200    registry
201}
202
203/// Take a [`ServerBuilder`] by value, install the FT.*
204/// extension, and return the wired builder plus the shared
205/// registry. Useful when the caller prefers to own the
206/// builder by value (the chained-call form):
207///
208/// ```no_run
209/// use dynomite::embed::ServerBuilder;
210/// let builder = ServerBuilder::new("p");
211/// let (builder, registry) = dynomite_search::install_owned(builder);
212/// let _ = (builder, registry);
213/// ```
214#[must_use]
215pub fn install_owned(builder: ServerBuilder) -> (ServerBuilder, Arc<VectorRegistry>) {
216    let ext = SearchExtension::default();
217    let registry = Arc::clone(ext.registry());
218    let builder = builder.with_command_extension(Arc::new(ext));
219    (builder, registry)
220}