1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use crate::{
    bson::{doc, Document},
    change_stream::{event::ResumeToken, ChangeStreamData, WatchArgs},
    cmap::{Command, RawCommandResponse, StreamDescription},
    cursor::CursorSpecification,
    error::Result,
    operation::{append_options, OperationWithDefaults, Retryability},
    options::{ChangeStreamOptions, SelectionCriteria, WriteConcern},
};

use super::Aggregate;

pub(crate) struct ChangeStreamAggregate {
    inner: Aggregate,
    args: WatchArgs,
    resume_data: Option<ChangeStreamData>,
}

impl ChangeStreamAggregate {
    pub(crate) fn new(args: &WatchArgs, resume_data: Option<ChangeStreamData>) -> Result<Self> {
        Ok(Self {
            inner: Self::build_inner(args)?,
            args: args.clone(),
            resume_data,
        })
    }

    fn build_inner(args: &WatchArgs) -> Result<Aggregate> {
        let mut bson_options = Document::new();
        append_options(&mut bson_options, args.options.as_ref())?;

        let mut agg_pipeline = vec![doc! { "$changeStream": bson_options }];
        agg_pipeline.extend(args.pipeline.iter().cloned());
        Ok(Aggregate::new(
            args.target.clone(),
            agg_pipeline,
            args.options.as_ref().map(|o| o.aggregate_options()),
        ))
    }
}

impl OperationWithDefaults for ChangeStreamAggregate {
    type O = (CursorSpecification, ChangeStreamData);
    type Command = Document;

    const NAME: &'static str = "aggregate";

    fn build(&mut self, description: &StreamDescription) -> Result<Command> {
        if let Some(data) = &mut self.resume_data {
            let mut new_opts = self.args.options.clone().unwrap_or_default();
            if let Some(token) = data.resume_token.take() {
                if new_opts.start_after.is_some() && !data.document_returned {
                    new_opts.start_after = Some(token);
                    new_opts.start_at_operation_time = None;
                } else {
                    new_opts.resume_after = Some(token);
                    new_opts.start_after = None;
                    new_opts.start_at_operation_time = None;
                }
            } else {
                let saved_time = new_opts
                    .start_at_operation_time
                    .as_ref()
                    .or(data.initial_operation_time.as_ref());
                if saved_time.is_some() && description.max_wire_version.map_or(false, |v| v >= 7) {
                    new_opts.start_at_operation_time = saved_time.cloned();
                }
            }

            self.inner = Self::build_inner(&WatchArgs {
                options: Some(new_opts),
                ..self.args.clone()
            })?;
        }
        self.inner.build(description)
    }

    fn extract_at_cluster_time(
        &self,
        response: &bson::RawDocument,
    ) -> Result<Option<bson::Timestamp>> {
        self.inner.extract_at_cluster_time(response)
    }

    fn handle_response(
        &self,
        response: RawCommandResponse,
        description: &StreamDescription,
    ) -> Result<Self::O> {
        let op_time = response
            .raw_body()
            .get("operationTime")?
            .and_then(bson::RawBsonRef::as_timestamp);
        let spec = self.inner.handle_response(response, description)?;

        let mut data = ChangeStreamData {
            resume_token: ResumeToken::initial(self.args.options.as_ref(), &spec),
            ..ChangeStreamData::default()
        };
        let has_no_time = |o: &ChangeStreamOptions| {
            o.start_at_operation_time.is_none()
                && o.resume_after.is_none()
                && o.start_after.is_none()
        };
        if self.args.options.as_ref().map_or(true, has_no_time)
            && description.max_wire_version.map_or(false, |v| v >= 7)
            && spec.initial_buffer.is_empty()
            && spec.post_batch_resume_token.is_none()
        {
            data.initial_operation_time = op_time;
        }

        Ok((spec, data))
    }

    fn selection_criteria(&self) -> Option<&SelectionCriteria> {
        self.inner.selection_criteria()
    }

    fn supports_read_concern(&self, description: &StreamDescription) -> bool {
        self.inner.supports_read_concern(description)
    }

    fn write_concern(&self) -> Option<&WriteConcern> {
        self.inner.write_concern()
    }

    fn retryability(&self) -> Retryability {
        self.inner.retryability()
    }
}