Skip to main content

mongodb/action/
aggregate.rs

1use std::{marker::PhantomData, time::Duration};
2
3use crate::bson::{Bson, Document};
4
5use crate::{
6    coll::options::{AggregateOptions, Hint},
7    collation::Collation,
8    error::Result,
9    operation::OperationTarget,
10    options::{ReadConcern, WriteConcern},
11    selection_criteria::SelectionCriteria,
12    Client,
13    ClientSession,
14    Collection,
15    Cursor,
16    Database,
17    SessionCursor,
18};
19
20use super::{
21    action_impl,
22    deeplink,
23    export_doc,
24    option_setters,
25    options_doc,
26    CollRef,
27    ExplicitSession,
28    ImplicitSession,
29};
30
31impl Database {
32    /// Runs an aggregation operation.
33    ///
34    /// See the documentation [here](https://www.mongodb.com/docs/manual/aggregation/) for more
35    /// information on aggregations.
36    ///
37    /// `await` will return d[`Result<Cursor<Document>>`]. If a [`ClientSession`] was provided, the
38    /// returned cursor will be a [`SessionCursor`]. If [`with_type`](Aggregate::with_type) was
39    /// called, the returned cursor will be generic over the `T` specified.
40    #[deeplink]
41    #[options_doc(aggregate)]
42    pub fn aggregate(&self, pipeline: impl IntoIterator<Item = Document>) -> Aggregate<'_> {
43        Aggregate::new(
44            AggregateTargetRef::Database(self),
45            pipeline.into_iter().collect(),
46        )
47    }
48}
49
50impl<T> Collection<T>
51where
52    T: Send + Sync,
53{
54    /// Runs an aggregation operation.
55    ///
56    /// See the documentation [here](https://www.mongodb.com/docs/manual/aggregation/) for more
57    /// information on aggregations.
58    ///
59    /// `await` will return d[`Result<Cursor<Document>>`]. If a [`ClientSession`] was provided, the
60    /// returned cursor will be a [`SessionCursor`]. If [`with_type`](Aggregate::with_type) was
61    /// called, the returned cursor will be generic over the `T` specified.
62    #[deeplink]
63    #[options_doc(aggregate)]
64    pub fn aggregate(&self, pipeline: impl IntoIterator<Item = Document>) -> Aggregate<'_> {
65        Aggregate::new(
66            AggregateTargetRef::Collection(CollRef::new(self)),
67            pipeline.into_iter().collect(),
68        )
69    }
70}
71
72#[cfg(feature = "sync")]
73impl crate::sync::Database {
74    /// Runs an aggregation operation.
75    ///
76    /// See the documentation [here](https://www.mongodb.com/docs/manual/aggregation/) for more
77    /// information on aggregations.
78    ///
79    /// [`run`](Aggregate::run) will return d[`Result<crate::sync::Cursor<Document>>`]. If a
80    /// [`crate::sync::ClientSession`] was provided, the returned cursor will be a
81    /// [`crate::sync::SessionCursor`]. If [`with_type`](Aggregate::with_type) was called, the
82    /// returned cursor will be generic over the `T` specified.
83    #[deeplink]
84    #[options_doc(aggregate, "run")]
85    pub fn aggregate(&self, pipeline: impl IntoIterator<Item = Document>) -> Aggregate<'_> {
86        self.async_database.aggregate(pipeline)
87    }
88}
89
90#[cfg(feature = "sync")]
91impl<T> crate::sync::Collection<T>
92where
93    T: Send + Sync,
94{
95    /// Runs an aggregation operation.
96    ///
97    /// See the documentation [here](https://www.mongodb.com/docs/manual/aggregation/) for more
98    /// information on aggregations.
99    ///
100    /// [`run`](Aggregate::run) will return d[`Result<crate::sync::Cursor<Document>>`]. If a
101    /// `crate::sync::ClientSession` was provided, the returned cursor will be a
102    /// `crate::sync::SessionCursor`. If [`with_type`](Aggregate::with_type) was called, the
103    /// returned cursor will be generic over the `T` specified.
104    #[deeplink]
105    #[options_doc(aggregate, "run")]
106    pub fn aggregate(&self, pipeline: impl IntoIterator<Item = Document>) -> Aggregate<'_> {
107        self.async_collection.aggregate(pipeline)
108    }
109}
110
111/// Run an aggregation operation.  Construct with [`Database::aggregate`] or
112/// [`Collection::aggregate`].
113#[must_use]
114pub struct Aggregate<'a, Session = ImplicitSession, T = Document> {
115    target: AggregateTargetRef<'a>,
116    pipeline: Vec<Document>,
117    options: Option<AggregateOptions>,
118    session: Session,
119    _phantom: PhantomData<fn() -> T>,
120}
121
122impl<'a> Aggregate<'a> {
123    fn new(target: AggregateTargetRef<'a>, pipeline: Vec<Document>) -> Self {
124        Self {
125            target,
126            pipeline,
127            options: None,
128            session: ImplicitSession,
129            _phantom: PhantomData,
130        }
131    }
132}
133
134#[option_setters(crate::coll::options::AggregateOptions)]
135#[export_doc(aggregate, extra = [session, batch])]
136impl<'a, Session, T> Aggregate<'a, Session, T> {
137    /// Use the provided type for the returned cursor.
138    ///
139    /// ```rust
140    /// # use futures_util::TryStreamExt;
141    /// # use mongodb::{bson::Document, error::Result, Cursor, Database};
142    /// # use serde::Deserialize;
143    /// # async fn run() -> Result<()> {
144    /// # let database: Database = todo!();
145    /// # let pipeline: Vec<Document> = todo!();
146    /// #[derive(Deserialize)]
147    /// struct PipelineOutput {
148    ///     len: usize,
149    /// }
150    ///
151    /// let aggregate_cursor = database
152    ///     .aggregate(pipeline)
153    ///     .with_type::<PipelineOutput>()
154    ///     .await?;
155    /// let aggregate_results: Vec<PipelineOutput> = aggregate_cursor.try_collect().await?;
156    /// # Ok(())
157    /// # }
158    /// ```
159    pub fn with_type<U>(self) -> Aggregate<'a, Session, U> {
160        Aggregate {
161            target: self.target,
162            pipeline: self.pipeline,
163            options: self.options,
164            session: self.session,
165            _phantom: PhantomData,
166        }
167    }
168}
169
170macro_rules! agg_exec_generic {
171    ($agg:expr) => {{
172        let mut aggregate = crate::operation::aggregate::Aggregate::new(
173            (&$agg.target).into(),
174            $agg.pipeline,
175            $agg.options,
176        );
177        let client = $agg.target.client();
178        client.execute_cursor_operation(&mut aggregate, None).await
179    }};
180}
181
182impl<'a, T> Aggregate<'a, ImplicitSession, T> {
183    /// Use the provided session when running the operation.
184    pub fn session(
185        self,
186        value: impl Into<&'a mut ClientSession>,
187    ) -> Aggregate<'a, ExplicitSession<'a>, T> {
188        Aggregate {
189            target: self.target,
190            pipeline: self.pipeline,
191            options: self.options,
192            session: ExplicitSession(value.into()),
193            _phantom: PhantomData,
194        }
195    }
196
197    /// Execute the aggregate command, returning a cursor that provides results in zero-copy raw
198    /// batches.
199    pub async fn batch(self) -> Result<crate::raw_batch_cursor::RawBatchCursor> {
200        agg_exec_generic!(self)
201    }
202}
203
204#[action_impl(sync = crate::sync::Cursor<T>)]
205impl<'a, T> Action for Aggregate<'a, ImplicitSession, T> {
206    type Future = AggregateFuture;
207
208    async fn execute(self) -> Result<Cursor<T>> {
209        agg_exec_generic!(self)
210    }
211}
212
213macro_rules! agg_exec_generic_session {
214    ($agg:expr) => {{
215        let mut aggregate = crate::operation::aggregate::Aggregate::new(
216            (&$agg.target).into(),
217            $agg.pipeline,
218            $agg.options,
219        );
220        let client = $agg.target.client();
221        let session = $agg.session;
222        client
223            .execute_cursor_operation(&mut aggregate, Some(session.0))
224            .await
225    }};
226}
227
228impl<'a, T> Aggregate<'a, ExplicitSession<'a>, T> {
229    /// Execute the aggregate command, returning a cursor that provides results in zero-copy raw
230    /// batches.
231    pub async fn batch(self) -> Result<crate::raw_batch_cursor::SessionRawBatchCursor> {
232        agg_exec_generic_session!(self)
233    }
234}
235
236#[action_impl(sync = crate::sync::SessionCursor<T>)]
237impl<'a, T> Action for Aggregate<'a, ExplicitSession<'a>, T> {
238    type Future = AggregateSessionFuture;
239
240    async fn execute(self) -> Result<SessionCursor<T>> {
241        agg_exec_generic_session!(self)
242    }
243}
244
245enum AggregateTargetRef<'a> {
246    Database(&'a Database),
247    Collection(CollRef<'a>),
248}
249
250impl AggregateTargetRef<'_> {
251    fn client(&self) -> &Client {
252        match self {
253            Self::Collection(cr) => cr.client(),
254            Self::Database(db) => db.client(),
255        }
256    }
257}
258
259impl From<&AggregateTargetRef<'_>> for OperationTarget {
260    fn from(value: &AggregateTargetRef<'_>) -> Self {
261        match value {
262            AggregateTargetRef::Collection(cr) => OperationTarget::Collection((*cr).clone()),
263            AggregateTargetRef::Database(db) => OperationTarget::Database((*db).clone()),
264        }
265    }
266}
267
268#[test]
269fn aggregate_session_type() {
270    // Assert that this code compiles but do not actually run it.
271    #[allow(
272        unreachable_code,
273        unused_variables,
274        dead_code,
275        clippy::diverging_sub_expression
276    )]
277    fn compile_ok() {
278        let agg: Aggregate = todo!();
279        let typed: Aggregate<'_, _, ()> = agg.with_type::<()>();
280        let mut session: ClientSession = todo!();
281        let typed_session: Aggregate<'_, _, ()> = typed.session(&mut session);
282    }
283}