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
95pub trait SearchStrategy: Clone + Send + Sync + Default {}
96
97/// Can retrieve documents given a SearchStrategy
98#[async_trait]
99pub trait Retrieve<S: SearchStrategy>: Send + Sync + DynClone {
100    async fn retrieve(
101        &self,
102        search_strategy: &S,
103        query: Query<states::Pending>,
104    ) -> Result<Query<states::Retrieved>>;
105
106    fn name(&self) -> &'static str {
107        let name = std::any::type_name::<Self>();
108        name.split("::").last().unwrap_or(name)
109    }
110}
111
112dyn_clone::clone_trait_object!(<S> Retrieve<S>);
113
114#[async_trait]
115impl<S: SearchStrategy> Retrieve<S> for Box<dyn Retrieve<S>> {
116    async fn retrieve(
117        &self,
118        search_strategy: &S,
119        query: Query<states::Pending>,
120    ) -> Result<Query<states::Retrieved>> {
121        self.as_ref().retrieve(search_strategy, query).await
122    }
123
124    fn name(&self) -> &'static str {
125        self.as_ref().name()
126    }
127}
128
129#[async_trait]
130impl<S: SearchStrategy> Retrieve<S> for Arc<dyn Retrieve<S>> {
131    async fn retrieve(
132        &self,
133        search_strategy: &S,
134        query: Query<states::Pending>,
135    ) -> Result<Query<states::Retrieved>> {
136        self.as_ref().retrieve(search_strategy, query).await
137    }
138
139    fn name(&self) -> &'static str {
140        self.as_ref().name()
141    }
142}
143
144#[async_trait]
145impl<S, F> Retrieve<S> for F
146where
147    S: SearchStrategy,
148    F: Fn(&S, Query<states::Pending>) -> Result<Query<states::Retrieved>> + Send + Sync + Clone,
149{
150    async fn retrieve(
151        &self,
152        search_strategy: &S,
153        query: Query<states::Pending>,
154    ) -> Result<Query<states::Retrieved>> {
155        (self)(search_strategy, query)
156    }
157}
158
159/// Can transform a response after retrieval
160#[async_trait]
161pub trait TransformResponse: Send + Sync + DynClone {
162    async fn transform_response(&self, query: Query<Retrieved>)
163    -> Result<Query<states::Retrieved>>;
164
165    fn name(&self) -> &'static str {
166        let name = std::any::type_name::<Self>();
167        name.split("::").last().unwrap_or(name)
168    }
169}
170
171dyn_clone::clone_trait_object!(TransformResponse);
172
173#[cfg(feature = "test-utils")]
174mock! {
175    #[derive(Debug)]
176    pub TransformResponse {}
177
178    #[async_trait]
179    impl TransformResponse for TransformResponse {
180        async fn transform_response(&self, query: Query<Retrieved>)
181            -> Result<Query<states::Retrieved>>;
182        fn name(&self) -> &'static str;
183    }
184
185    impl Clone for TransformResponse {
186        fn clone(&self) -> Self;
187    }
188}
189#[async_trait]
190impl<F> TransformResponse for F
191where
192    F: Fn(Query<Retrieved>) -> Result<Query<Retrieved>> + Send + Sync + Clone,
193{
194    async fn transform_response(&self, query: Query<Retrieved>) -> Result<Query<Retrieved>> {
195        (self)(query)
196    }
197}
198
199#[async_trait]
200impl TransformResponse for Box<dyn TransformResponse> {
201    async fn transform_response(&self, query: Query<Retrieved>) -> Result<Query<Retrieved>> {
202        self.as_ref().transform_response(query).await
203    }
204
205    fn name(&self) -> &'static str {
206        self.as_ref().name()
207    }
208}
209
210#[async_trait]
211impl TransformResponse for Arc<dyn TransformResponse> {
212    async fn transform_response(&self, query: Query<Retrieved>) -> Result<Query<Retrieved>> {
213        self.as_ref().transform_response(query).await
214    }
215
216    fn name(&self) -> &'static str {
217        self.as_ref().name()
218    }
219}
220
221/// Can answer the original query
222#[async_trait]
223pub trait Answer: Send + Sync + DynClone {
224    async fn answer(&self, query: Query<states::Retrieved>) -> Result<Query<states::Answered>>;
225
226    fn name(&self) -> &'static str {
227        let name = std::any::type_name::<Self>();
228        name.split("::").last().unwrap_or(name)
229    }
230}
231
232dyn_clone::clone_trait_object!(Answer);
233
234#[cfg(feature = "test-utils")]
235mock! {
236    #[derive(Debug)]
237    pub Answer {}
238
239    #[async_trait]
240    impl Answer for Answer {
241        async fn answer(&self, query: Query<states::Retrieved>) -> Result<Query<states::Answered>>;
242        fn name(&self) -> &'static str;
243    }
244
245    impl Clone for Answer {
246        fn clone(&self) -> Self;
247    }
248}
249#[async_trait]
250impl<F> Answer for F
251where
252    F: Fn(Query<Retrieved>) -> Result<Query<states::Answered>> + Send + Sync + Clone,
253{
254    async fn answer(&self, query: Query<Retrieved>) -> Result<Query<states::Answered>> {
255        (self)(query)
256    }
257}
258
259#[async_trait]
260impl Answer for Box<dyn Answer> {
261    async fn answer(&self, query: Query<Retrieved>) -> Result<Query<states::Answered>> {
262        self.as_ref().answer(query).await
263    }
264
265    fn name(&self) -> &'static str {
266        self.as_ref().name()
267    }
268}
269
270#[async_trait]
271impl Answer for Arc<dyn Answer> {
272    async fn answer(&self, query: Query<Retrieved>) -> Result<Query<states::Answered>> {
273        self.as_ref().answer(query).await
274    }
275
276    fn name(&self) -> &'static str {
277        self.as_ref().name()
278    }
279}
280
281/// Evaluates a query
282///
283/// An evaluator needs to be able to respond to each step in the query pipeline
284#[async_trait]
285pub trait EvaluateQuery: Send + Sync + DynClone {
286    async fn evaluate(&self, evaluation: QueryEvaluation) -> Result<()>;
287}
288
289dyn_clone::clone_trait_object!(EvaluateQuery);
290
291#[cfg(feature = "test-utils")]
292mock! {
293    #[derive(Debug)]
294    pub EvaluateQuery {}
295
296    #[async_trait]
297    impl EvaluateQuery for EvaluateQuery {
298        async fn evaluate(&self, evaluation: QueryEvaluation) -> Result<()>;
299    }
300
301    impl Clone for EvaluateQuery {
302        fn clone(&self) -> Self;
303    }
304}
305#[async_trait]
306impl EvaluateQuery for Box<dyn EvaluateQuery> {
307    async fn evaluate(&self, evaluation: QueryEvaluation) -> Result<()> {
308        self.as_ref().evaluate(evaluation).await
309    }
310}
311
312#[async_trait]
313impl EvaluateQuery for Arc<dyn EvaluateQuery> {
314    async fn evaluate(&self, evaluation: QueryEvaluation) -> Result<()> {
315        self.as_ref().evaluate(evaluation).await
316    }
317}