Skip to main content

swiftide_core/
query_traits.rs

1use std::sync::Arc;
2
3use anyhow::Result;
4use async_trait::async_trait;
5use dyn_clone::DynClone;
6
7use crate::{
8    query::{
9        Query,
10        states::{self, Retrieved},
11    },
12    querying::QueryEvaluation,
13};
14
15#[cfg(feature = "test-utils")]
16use mockall::{mock, predicate::str};
17
18/// Can transform queries before retrieval
19#[async_trait]
20pub trait TransformQuery: Send + Sync + DynClone {
21    async fn transform_query(
22        &self,
23        query: Query<states::Pending>,
24    ) -> Result<Query<states::Pending>>;
25
26    fn name(&self) -> &'static str {
27        let name = std::any::type_name::<Self>();
28        name.split("::").last().unwrap_or(name)
29    }
30}
31
32dyn_clone::clone_trait_object!(TransformQuery);
33
34#[cfg(feature = "test-utils")]
35mock! {
36    #[derive(Debug)]
37    pub TransformQuery {}
38
39    #[async_trait]
40    impl TransformQuery for TransformQuery {
41        async fn transform_query(
42            &self,
43            query: Query<states::Pending>,
44        ) -> Result<Query<states::Pending>>;
45        fn name(&self) -> &'static str;
46    }
47
48    impl Clone for TransformQuery {
49        fn clone(&self) -> Self;
50    }
51}
52
53#[async_trait]
54impl<F> TransformQuery for F
55where
56    F: Fn(Query<states::Pending>) -> Result<Query<states::Pending>> + Send + Sync + Clone,
57{
58    async fn transform_query(
59        &self,
60        query: Query<states::Pending>,
61    ) -> Result<Query<states::Pending>> {
62        (self)(query)
63    }
64}
65
66#[async_trait]
67impl TransformQuery for Box<dyn TransformQuery> {
68    async fn transform_query(
69        &self,
70        query: Query<states::Pending>,
71    ) -> Result<Query<states::Pending>> {
72        self.as_ref().transform_query(query).await
73    }
74
75    fn name(&self) -> &'static str {
76        self.as_ref().name()
77    }
78}
79
80#[async_trait]
81impl TransformQuery for Arc<dyn TransformQuery> {
82    async fn transform_query(
83        &self,
84        query: Query<states::Pending>,
85    ) -> Result<Query<states::Pending>> {
86        self.as_ref().transform_query(query).await
87    }
88
89    fn name(&self) -> &'static str {
90        self.as_ref().name()
91    }
92}
93
94/// A search strategy for the query pipeline
95///
96/// Strategies do not need to implement `Default`; a strategy generic over a filter can then offer
97/// `default()` for its unfiltered form only, so the filter type is inferred.
98pub trait SearchStrategy: Clone + Send + Sync {}
99
100/// Can retrieve documents given a `SearchStrategy`
101#[async_trait]
102pub trait Retrieve<S: SearchStrategy>: Send + Sync + DynClone {
103    async fn retrieve(
104        &self,
105        search_strategy: &S,
106        query: Query<states::Pending>,
107    ) -> Result<Query<states::Retrieved>>;
108
109    fn name(&self) -> &'static str {
110        let name = std::any::type_name::<Self>();
111        name.split("::").last().unwrap_or(name)
112    }
113}
114
115dyn_clone::clone_trait_object!(<S> Retrieve<S>);
116
117#[async_trait]
118impl<S: SearchStrategy> Retrieve<S> for Box<dyn Retrieve<S>> {
119    async fn retrieve(
120        &self,
121        search_strategy: &S,
122        query: Query<states::Pending>,
123    ) -> Result<Query<states::Retrieved>> {
124        self.as_ref().retrieve(search_strategy, query).await
125    }
126
127    fn name(&self) -> &'static str {
128        self.as_ref().name()
129    }
130}
131
132#[async_trait]
133impl<S: SearchStrategy> Retrieve<S> for Arc<dyn Retrieve<S>> {
134    async fn retrieve(
135        &self,
136        search_strategy: &S,
137        query: Query<states::Pending>,
138    ) -> Result<Query<states::Retrieved>> {
139        self.as_ref().retrieve(search_strategy, query).await
140    }
141
142    fn name(&self) -> &'static str {
143        self.as_ref().name()
144    }
145}
146
147#[async_trait]
148impl<S, F> Retrieve<S> for F
149where
150    S: SearchStrategy,
151    F: Fn(&S, Query<states::Pending>) -> Result<Query<states::Retrieved>> + Send + Sync + Clone,
152{
153    async fn retrieve(
154        &self,
155        search_strategy: &S,
156        query: Query<states::Pending>,
157    ) -> Result<Query<states::Retrieved>> {
158        (self)(search_strategy, query)
159    }
160}
161
162/// Can transform a response after retrieval
163#[async_trait]
164pub trait TransformResponse: Send + Sync + DynClone {
165    async fn transform_response(&self, query: Query<Retrieved>)
166    -> Result<Query<states::Retrieved>>;
167
168    fn name(&self) -> &'static str {
169        let name = std::any::type_name::<Self>();
170        name.split("::").last().unwrap_or(name)
171    }
172}
173
174dyn_clone::clone_trait_object!(TransformResponse);
175
176#[cfg(feature = "test-utils")]
177mock! {
178    #[derive(Debug)]
179    pub TransformResponse {}
180
181    #[async_trait]
182    impl TransformResponse for TransformResponse {
183        async fn transform_response(&self, query: Query<Retrieved>)
184            -> Result<Query<states::Retrieved>>;
185        fn name(&self) -> &'static str;
186    }
187
188    impl Clone for TransformResponse {
189        fn clone(&self) -> Self;
190    }
191}
192#[async_trait]
193impl<F> TransformResponse for F
194where
195    F: Fn(Query<Retrieved>) -> Result<Query<Retrieved>> + Send + Sync + Clone,
196{
197    async fn transform_response(&self, query: Query<Retrieved>) -> Result<Query<Retrieved>> {
198        (self)(query)
199    }
200}
201
202#[async_trait]
203impl TransformResponse for Box<dyn TransformResponse> {
204    async fn transform_response(&self, query: Query<Retrieved>) -> Result<Query<Retrieved>> {
205        self.as_ref().transform_response(query).await
206    }
207
208    fn name(&self) -> &'static str {
209        self.as_ref().name()
210    }
211}
212
213#[async_trait]
214impl TransformResponse for Arc<dyn TransformResponse> {
215    async fn transform_response(&self, query: Query<Retrieved>) -> Result<Query<Retrieved>> {
216        self.as_ref().transform_response(query).await
217    }
218
219    fn name(&self) -> &'static str {
220        self.as_ref().name()
221    }
222}
223
224/// Can answer the original query
225#[async_trait]
226pub trait Answer: Send + Sync + DynClone {
227    async fn answer(&self, query: Query<states::Retrieved>) -> Result<Query<states::Answered>>;
228
229    fn name(&self) -> &'static str {
230        let name = std::any::type_name::<Self>();
231        name.split("::").last().unwrap_or(name)
232    }
233}
234
235dyn_clone::clone_trait_object!(Answer);
236
237#[cfg(feature = "test-utils")]
238mock! {
239    #[derive(Debug)]
240    pub Answer {}
241
242    #[async_trait]
243    impl Answer for Answer {
244        async fn answer(&self, query: Query<states::Retrieved>) -> Result<Query<states::Answered>>;
245        fn name(&self) -> &'static str;
246    }
247
248    impl Clone for Answer {
249        fn clone(&self) -> Self;
250    }
251}
252#[async_trait]
253impl<F> Answer for F
254where
255    F: Fn(Query<Retrieved>) -> Result<Query<states::Answered>> + Send + Sync + Clone,
256{
257    async fn answer(&self, query: Query<Retrieved>) -> Result<Query<states::Answered>> {
258        (self)(query)
259    }
260}
261
262#[async_trait]
263impl Answer for Box<dyn Answer> {
264    async fn answer(&self, query: Query<Retrieved>) -> Result<Query<states::Answered>> {
265        self.as_ref().answer(query).await
266    }
267
268    fn name(&self) -> &'static str {
269        self.as_ref().name()
270    }
271}
272
273#[async_trait]
274impl Answer for Arc<dyn Answer> {
275    async fn answer(&self, query: Query<Retrieved>) -> Result<Query<states::Answered>> {
276        self.as_ref().answer(query).await
277    }
278
279    fn name(&self) -> &'static str {
280        self.as_ref().name()
281    }
282}
283
284/// Evaluates a query
285///
286/// An evaluator needs to be able to respond to each step in the query pipeline
287#[async_trait]
288pub trait EvaluateQuery: Send + Sync + DynClone {
289    async fn evaluate(&self, evaluation: QueryEvaluation) -> Result<()>;
290}
291
292dyn_clone::clone_trait_object!(EvaluateQuery);
293
294#[cfg(feature = "test-utils")]
295mock! {
296    #[derive(Debug)]
297    pub EvaluateQuery {}
298
299    #[async_trait]
300    impl EvaluateQuery for EvaluateQuery {
301        async fn evaluate(&self, evaluation: QueryEvaluation) -> Result<()>;
302    }
303
304    impl Clone for EvaluateQuery {
305        fn clone(&self) -> Self;
306    }
307}
308#[async_trait]
309impl EvaluateQuery for Box<dyn EvaluateQuery> {
310    async fn evaluate(&self, evaluation: QueryEvaluation) -> Result<()> {
311        self.as_ref().evaluate(evaluation).await
312    }
313}
314
315#[async_trait]
316impl EvaluateQuery for Arc<dyn EvaluateQuery> {
317    async fn evaluate(&self, evaluation: QueryEvaluation) -> Result<()> {
318        self.as_ref().evaluate(evaluation).await
319    }
320}