mongodb 3.6.0

The official MongoDB driver for Rust
Documentation
use crate::{
    bson::{doc, Document},
    change_stream::{
        common::{ChangeStreamData, WatchArgs},
        event::ResumeToken,
    },
    client::Retry,
    cmap::{Command, RawCommandResponse, StreamDescription},
    cursor::common::CursorSpecification,
    error::Result,
    operation::{append_options, ExecutionContext, Operation, Retryability},
    options::{ChangeStreamOptions, ClientOptions, 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 Operation for ChangeStreamAggregate {
    type O = (CursorSpecification, ChangeStreamData);

    const NAME: &'static crate::bson_compat::CStr = Aggregate::NAME;

    const ZERO_COPY: bool = true;

    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.is_some_and(|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: &crate::bson::RawDocument,
    ) -> Result<Option<crate::bson::Timestamp>> {
        self.inner.extract_at_cluster_time(response)
    }

    fn handle_response<'a>(
        &'a self,
        response: std::borrow::Cow<'a, RawCommandResponse>,
        mut context: ExecutionContext<'a>,
    ) -> crate::BoxFuture<'a, Result<Self::O>> {
        use futures_util::FutureExt;
        async move {
            let op_time = response
                .raw_body()
                .get("operationTime")?
                .and_then(crate::bson::RawBsonRef::as_timestamp);

            let inner_context = ExecutionContext {
                connection: context.connection,
                session: context.session.as_deref_mut(),
                effective_criteria: context.effective_criteria,
            };
            let spec = {
                use crate::operation::OperationWithDefaults;
                self.inner.handle_response_cow(response, inner_context)?
            };

            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()
            };

            let description = context.connection.stream_description()?;
            if self.args.options.as_ref().is_none_or(has_no_time)
                && description.max_wire_version.is_some_and(|v| v >= 7)
                && spec.is_empty
                && spec.post_batch_resume_token.is_none()
            {
                data.initial_operation_time = op_time;
            }

            Ok((spec, data))
        }
        .boxed()
    }

    fn selection_criteria(&self) -> crate::operation::Feature<&SelectionCriteria> {
        self.inner.selection_criteria()
    }

    fn write_concern(&self) -> crate::operation::Feature<&WriteConcern> {
        self.inner.write_concern()
    }

    fn retryability(&self, options: &ClientOptions) -> Retryability {
        self.inner.retryability(options)
    }

    fn is_backpressure_retryable(&self, options: &ClientOptions) -> bool {
        self.inner.is_backpressure_retryable(options)
    }

    fn target(&self) -> crate::operation::OperationTarget {
        self.inner.target()
    }

    fn handle_error(&self, error: crate::error::Error) -> Result<Self::O> {
        Err(error)
    }

    fn read_concern(&self) -> crate::operation::Feature<&crate::options::ReadConcern> {
        self.inner.read_concern()
    }

    fn supports_sessions(&self) -> bool {
        self.inner.supports_sessions()
    }

    fn update_for_retry(&mut self, retry: Option<&Retry>) {
        self.inner.update_for_retry(retry);
    }

    fn override_criteria(&self) -> crate::operation::OverrideCriteriaFn {
        self.inner.override_criteria()
    }

    fn pinned_connection(&self) -> Option<&crate::cmap::conn::PinnedConnectionHandle> {
        self.inner.pinned_connection()
    }

    fn name(&self) -> &crate::bson_compat::CStr {
        self.inner.name()
    }

    #[cfg(feature = "opentelemetry")]
    type Otel = crate::otel::Witness<Self>;
}

#[cfg(feature = "opentelemetry")]
impl crate::otel::OtelInfoDefaults for ChangeStreamAggregate {}