firestore 0.57.1

#![allow(clippy::derive_partial_eq_without_eq)] // Since we may not be able to implement Eq for the changes coming from Firestore protos

use crate::db::support::FirestoreAggregatedQuerySupport;
use crate::FirestoreInstant;
use crate::{FirestoreDb, FirestoreQueryParams, FirestoreResult};
use async_trait::async_trait;
use futures::stream::BoxStream;
use futures::TryStreamExt;
use futures::{future, StreamExt};
use gcloud_sdk::google::firestore::v1::*;
use rsb_derive::*;
use serde::Deserialize;
use tracing::*;

#[derive(Debug, PartialEq, Clone, Builder)]
pub struct FirestoreAggregatedQueryParams {
    pub query_params: FirestoreQueryParams,
    pub aggregations: Vec<FirestoreAggregation>,
}

#[derive(Debug, PartialEq, Clone, Builder)]
pub struct FirestoreAggregation {
    pub alias: String,
    pub operator: Option<FirestoreAggregationOperator>,
}

impl From<&FirestoreAggregation> for structured_aggregation_query::Aggregation {
    fn from(aggregation: &FirestoreAggregation) -> Self {
        structured_aggregation_query::Aggregation {
            alias: aggregation.alias.clone(),
            operator: aggregation.operator.as_ref().map(|agg| agg.into()),
        }
    }
}

#[derive(Debug, PartialEq, Clone)]
pub enum FirestoreAggregationOperator {
    Count(FirestoreAggregationOperatorCount),
    Sum(FirestoreAggregationOperatorSum),
    Avg(FirestoreAggregationOperatorAvg),
}

impl From<&FirestoreAggregationOperator> for structured_aggregation_query::aggregation::Operator {
    fn from(op: &FirestoreAggregationOperator) -> Self {
        match op {
            FirestoreAggregationOperator::Count(opts) => {
                structured_aggregation_query::aggregation::Operator::Count(opts.into())
            }
            FirestoreAggregationOperator::Sum(opts) => {
                structured_aggregation_query::aggregation::Operator::Sum(opts.into())
            }
            FirestoreAggregationOperator::Avg(opts) => {
                structured_aggregation_query::aggregation::Operator::Avg(opts.into())
            }
        }
    }
}

#[derive(Debug, PartialEq, Clone, Builder)]
pub struct FirestoreAggregationOperatorCount {
    pub up_to: Option<usize>,
}

impl From<&FirestoreAggregationOperatorCount> for structured_aggregation_query::aggregation::Count {
    fn from(cnt: &FirestoreAggregationOperatorCount) -> Self {
        structured_aggregation_query::aggregation::Count {
            up_to: cnt.up_to.map(|v| v as i64),
        }
    }
}

#[derive(Debug, PartialEq, Clone, Builder)]
pub struct FirestoreAggregationOperatorSum {
    pub field_name: String,
}

impl From<&FirestoreAggregationOperatorSum> for structured_aggregation_query::aggregation::Sum {
    fn from(operator: &FirestoreAggregationOperatorSum) -> Self {
        structured_aggregation_query::aggregation::Sum {
            field: Some(structured_query::FieldReference {
                field_path: operator.field_name.clone(),
            }),
        }
    }
}

#[derive(Debug, PartialEq, Clone, Builder)]
pub struct FirestoreAggregationOperatorAvg {
    pub field_name: String,
}

impl From<&FirestoreAggregationOperatorAvg> for structured_aggregation_query::aggregation::Avg {
    fn from(operator: &FirestoreAggregationOperatorAvg) -> Self {
        structured_aggregation_query::aggregation::Avg {
            field: Some(structured_query::FieldReference {
                field_path: operator.field_name.clone(),
            }),
        }
    }
}

#[async_trait]
impl FirestoreAggregatedQuerySupport for FirestoreDb {
    async fn aggregated_query_doc(
        &self,
        params: FirestoreAggregatedQueryParams,
    ) -> FirestoreResult<Vec<Document>> {
        let collection_str = params.query_params.collection_id.to_string();

        let span = span!(
            Level::DEBUG,
            "Firestore Aggregated Query",
            "/firestore/collection_name" = collection_str.as_str(),
            "/firestore/response_time" = field::Empty
        );
        self.aggregated_query_doc_with_retries(params, &span).await
    }

    async fn stream_aggregated_query_doc<'b>(
        &self,
        params: FirestoreAggregatedQueryParams,
    ) -> FirestoreResult<BoxStream<'b, Document>> {
        let collection_str = params.query_params.collection_id.to_string();

        let span = span!(
            Level::DEBUG,
            "Firestore Streaming Aggregated Query",
            "/firestore/collection_name" = collection_str.as_str(),
            "/firestore/response_time" = field::Empty
        );

        let doc_stream = self
            .stream_aggregated_query_doc_with_retries(params, &span)
            .await?;

        Ok(Box::pin(doc_stream.filter_map(|doc_res| {
            future::ready(match doc_res {
                Ok(Some(doc)) => Some(doc),
                Ok(None) => None,
                Err(err) => {
                    error!(%err, "Error occurred while consuming query.");
                    None
                }
            })
        })))
    }

    async fn stream_aggregated_query_doc_with_errors<'b>(
        &self,
        params: FirestoreAggregatedQueryParams,
    ) -> FirestoreResult<BoxStream<'b, FirestoreResult<Document>>> {
        let collection_str = params.query_params.collection_id.to_string();

        let span = span!(
            Level::DEBUG,
            "Firestore Streaming Aggregated Query",
            "/firestore/collection_name" = collection_str.as_str(),
            "/firestore/response_time" = field::Empty
        );

        let doc_stream = self
            .stream_aggregated_query_doc_with_retries(params, &span)
            .await?;

        Ok(Box::pin(doc_stream.filter_map(|doc_res| {
            future::ready(match doc_res {
                Ok(Some(doc)) => Some(Ok(doc)),
                Ok(None) => None,
                Err(err) => {
                    error!(%err, "Error occurred while consuming query.");
                    Some(Err(err))
                }
            })
        })))
    }

    async fn aggregated_query_obj<T>(
        &self,
        params: FirestoreAggregatedQueryParams,
    ) -> FirestoreResult<Vec<T>>
    where
        for<'de> T: Deserialize<'de>,
    {
        let doc_vec = self.aggregated_query_doc(params).await?;
        doc_vec
            .iter()
            .map(|doc| Self::deserialize_doc_to(doc))
            .collect()
    }

    async fn stream_aggregated_query_obj<'b, T>(
        &self,
        params: FirestoreAggregatedQueryParams,
    ) -> FirestoreResult<BoxStream<'b, T>>
    where
        for<'de> T: Deserialize<'de>,
    {
        let doc_stream = self.stream_aggregated_query_doc(params).await?;
        Ok(Box::pin(doc_stream.filter_map(|doc| async move {
            match Self::deserialize_doc_to::<T>(&doc) {
                Ok(obj) => Some(obj),
                Err(err) => {
                    error!(
                        %err,
                        "Error occurred while consuming query document as a stream.",
                    );
                    None
                }
            }
        })))
    }

    async fn stream_aggregated_query_obj_with_errors<'b, T>(
        &self,
        params: FirestoreAggregatedQueryParams,
    ) -> FirestoreResult<BoxStream<'b, FirestoreResult<T>>>
    where
        for<'de> T: Deserialize<'de>,
        T: Send + 'b,
    {
        let doc_stream = self.stream_aggregated_query_doc_with_errors(params).await?;
        Ok(Box::pin(doc_stream.and_then(|doc| {
            future::ready(Self::deserialize_doc_to::<T>(&doc))
        })))
    }
}

impl FirestoreDb {
    fn create_aggregated_query_request(
        &self,
        params: FirestoreAggregatedQueryParams,
    ) -> FirestoreResult<RunAggregationQueryRequest> {
        Ok(RunAggregationQueryRequest {
            parent: params
                .query_params
                .parent
                .as_ref()
                .unwrap_or_else(|| self.get_documents_path())
                .clone(),
            consistency_selector: self
                .session_params
                .consistency_selector
                .as_ref()
                .map(|selector| selector.try_into())
                .transpose()?,
            request_options: self
                .resolve_request_options(params.query_params.request_options.as_ref()),
            query_type: Some(run_aggregation_query_request::QueryType::StructuredAggregationQuery(
                StructuredAggregationQuery {
                    aggregations: params.aggregations.iter().map(|agg| agg.into()).collect(),
                    query_type: Some(gcloud_sdk::google::firestore::v1::structured_aggregation_query::QueryType::StructuredQuery(params.query_params.try_into()?)),
                }
            )),
            explain_options: None,
        })
    }

    async fn stream_aggregated_query_doc_with_retries<'b>(
        &self,
        params: FirestoreAggregatedQueryParams,
        span: &Span,
    ) -> FirestoreResult<BoxStream<'b, FirestoreResult<Option<Document>>>> {
        let collection_id = params.query_params.collection_id.clone();
        let query_request = self.create_aggregated_query_request(params)?;
        let begin_query_utc: FirestoreInstant = FirestoreInstant::now();

        let query_stream = self
            .run_aggregation_query_with_retries(&query_request, span)
            .await?
            .map_ok(Self::aggregated_response_to_doc)
            .map_err(|e| e.into())
            .boxed();

        let end_query_utc: FirestoreInstant = FirestoreInstant::now();
        let query_duration = end_query_utc.duration_since(begin_query_utc);

        span.record("/firestore/response_time", query_duration.as_millis());
        span.in_scope(|| {
            debug!(
                ?collection_id,
                duration_milliseconds = query_duration.as_millis(),
                "Querying stream of documents in specified collection.",
            );
        });

        Ok(query_stream)
    }

    async fn aggregated_query_doc_with_retries(
        &self,
        params: FirestoreAggregatedQueryParams,
        span: &Span,
    ) -> FirestoreResult<Vec<Document>> {
        let collection_id = params.query_params.collection_id.clone();
        let query_request = self.create_aggregated_query_request(params)?;
        let begin_query_utc: FirestoreInstant = FirestoreInstant::now();

        let docs = self
            .run_aggregation_query_with_retries(&query_request, span)
            .await?
            .map_ok(Self::aggregated_response_to_doc)
            .try_collect::<Vec<Option<Document>>>()
            .await?
            .into_iter()
            .flatten()
            .collect();

        let end_query_utc: FirestoreInstant = FirestoreInstant::now();
        let query_duration = end_query_utc.duration_since(begin_query_utc);

        span.record("/firestore/response_time", query_duration.as_millis());
        span.in_scope(|| {
            debug!(
                ?collection_id,
                duration_milliseconds = query_duration.as_millis(),
                "Querying documents in specified collection.",
            );
        });

        Ok(docs)
    }

    async fn run_aggregation_query_with_retries(
        &self,
        query_request: &RunAggregationQueryRequest,
        span: &Span,
    ) -> FirestoreResult<gcloud_sdk::tonic::Streaming<RunAggregationQueryResponse>> {
        self.retry_read(
            span,
            "run aggregation query",
            query_request,
            |mut client, request| async move { client.run_aggregation_query(request).await },
        )
        .await
    }

    fn aggregated_response_to_doc(mut agg_res: RunAggregationQueryResponse) -> Option<Document> {
        agg_res.result.take().map(|agg_res_doc| Document {
            name: "".to_string(),
            fields: agg_res_doc.aggregate_fields,
            create_time: None,
            update_time: None,
        })
    }
}