hehe_store/traits/
search.rs

1use crate::error::Result;
2use async_trait::async_trait;
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5use std::collections::HashMap;
6
7#[derive(Clone, Debug, Serialize, Deserialize)]
8pub struct Document {
9    pub id: String,
10    pub content: String,
11    #[serde(default)]
12    pub fields: HashMap<String, Value>,
13}
14
15impl Document {
16    pub fn new(id: impl Into<String>, content: impl Into<String>) -> Self {
17        Self {
18            id: id.into(),
19            content: content.into(),
20            fields: HashMap::new(),
21        }
22    }
23
24    pub fn with_field(mut self, key: impl Into<String>, value: impl Into<Value>) -> Self {
25        self.fields.insert(key.into(), value.into());
26        self
27    }
28}
29
30#[derive(Clone, Debug, Serialize, Deserialize)]
31pub struct SearchHit {
32    pub id: String,
33    pub score: f32,
34    pub content: String,
35    #[serde(default)]
36    pub highlights: Vec<String>,
37    #[serde(default)]
38    pub fields: HashMap<String, Value>,
39}
40
41#[derive(Clone, Debug, Default)]
42pub struct SearchFilter {
43    pub conditions: Vec<SearchCondition>,
44}
45
46#[derive(Clone, Debug)]
47pub enum SearchCondition {
48    Eq(String, Value),
49    Range(String, Option<Value>, Option<Value>),
50    In(String, Vec<Value>),
51}
52
53impl SearchFilter {
54    pub fn new() -> Self {
55        Self::default()
56    }
57
58    pub fn eq(mut self, field: impl Into<String>, value: impl Into<Value>) -> Self {
59        self.conditions
60            .push(SearchCondition::Eq(field.into(), value.into()));
61        self
62    }
63
64    pub fn range(
65        mut self,
66        field: impl Into<String>,
67        min: Option<Value>,
68        max: Option<Value>,
69    ) -> Self {
70        self.conditions
71            .push(SearchCondition::Range(field.into(), min, max));
72        self
73    }
74
75    pub fn is_empty(&self) -> bool {
76        self.conditions.is_empty()
77    }
78}
79
80#[derive(Clone, Debug)]
81pub struct IndexSchema {
82    pub fields: Vec<IndexField>,
83}
84
85#[derive(Clone, Debug)]
86pub struct IndexField {
87    pub name: String,
88    pub field_type: IndexFieldType,
89    pub stored: bool,
90    pub indexed: bool,
91}
92
93#[derive(Clone, Debug)]
94pub enum IndexFieldType {
95    Text,
96    Keyword,
97    Integer,
98    Float,
99    Boolean,
100    Date,
101}
102
103impl IndexSchema {
104    pub fn new() -> Self {
105        Self { fields: vec![] }
106    }
107
108    pub fn add_text(mut self, name: impl Into<String>) -> Self {
109        self.fields.push(IndexField {
110            name: name.into(),
111            field_type: IndexFieldType::Text,
112            stored: true,
113            indexed: true,
114        });
115        self
116    }
117
118    pub fn add_keyword(mut self, name: impl Into<String>) -> Self {
119        self.fields.push(IndexField {
120            name: name.into(),
121            field_type: IndexFieldType::Keyword,
122            stored: true,
123            indexed: true,
124        });
125        self
126    }
127
128    pub fn add_integer(mut self, name: impl Into<String>) -> Self {
129        self.fields.push(IndexField {
130            name: name.into(),
131            field_type: IndexFieldType::Integer,
132            stored: true,
133            indexed: true,
134        });
135        self
136    }
137}
138
139impl Default for IndexSchema {
140    fn default() -> Self {
141        Self::new()
142    }
143}
144
145#[async_trait]
146pub trait SearchStore: Send + Sync {
147    async fn create_index(&self, name: &str, schema: &IndexSchema) -> Result<()>;
148
149    async fn delete_index(&self, name: &str) -> Result<()>;
150
151    async fn index_exists(&self, name: &str) -> Result<bool>;
152
153    async fn list_indexes(&self) -> Result<Vec<String>>;
154
155    async fn index_documents(&self, index: &str, docs: &[Document]) -> Result<usize>;
156
157    async fn delete_documents(&self, index: &str, ids: &[String]) -> Result<usize>;
158
159    async fn search(&self, index: &str, query: &str, limit: usize) -> Result<Vec<SearchHit>>;
160
161    async fn search_with_filter(
162        &self,
163        index: &str,
164        query: &str,
165        filter: &SearchFilter,
166        limit: usize,
167    ) -> Result<Vec<SearchHit>>;
168
169    async fn count(&self, index: &str) -> Result<usize>;
170
171    fn backend_name(&self) -> &'static str;
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    #[test]
179    fn test_document() {
180        let doc = Document::new("doc1", "Hello world")
181            .with_field("category", "greeting")
182            .with_field("score", 0.9);
183
184        assert_eq!(doc.id, "doc1");
185        assert_eq!(doc.content, "Hello world");
186        assert_eq!(doc.fields.len(), 2);
187    }
188
189    #[test]
190    fn test_index_schema() {
191        let schema = IndexSchema::new()
192            .add_text("title")
193            .add_text("content")
194            .add_keyword("category");
195
196        assert_eq!(schema.fields.len(), 3);
197    }
198
199    #[test]
200    fn test_search_filter() {
201        let filter = SearchFilter::new()
202            .eq("category", "article")
203            .range("date", Some(Value::String("2024-01-01".into())), None);
204
205        assert_eq!(filter.conditions.len(), 2);
206    }
207}
hehe_store/traits/search.rs

hehe_store/traits/
search.rs
