firestore 0.61.1

use crate::FirestoreInstant;
use crate::*;
use async_trait::async_trait;
use futures::stream::BoxStream;
use futures::StreamExt;
use futures::TryStreamExt;
use gcloud_sdk::google::firestore::v1::*;
use gcloud_sdk::RequestTimeouts;
use rsb_derive::*;
use serde::Deserialize;
use std::future;
use tracing::*;

#[derive(Debug, Eq, PartialEq, Clone, Builder)]
pub struct FirestoreListDocParams {
    pub collection_id: String,

    pub parent: Option<String>,

    #[default = "100"]
    pub page_size: usize,

    pub page_token: Option<String>,
    pub order_by: Option<Vec<FirestoreQueryOrder>>,
    pub return_only_fields: Option<Vec<String>>,

    /// Request options (e.g. request tags) for this listing operation.
    pub request_options: Option<FirestoreRequestOptions>,
}

#[derive(Debug, PartialEq, Clone, Builder)]
pub struct FirestoreListDocResult {
    pub documents: Vec<Document>,
    pub page_token: Option<String>,
}

#[derive(Debug, Eq, PartialEq, Clone, Builder)]
pub struct FirestoreListCollectionIdsParams {
    pub parent: Option<String>,

    #[default = "100"]
    pub page_size: usize,
    pub page_token: Option<String>,

    /// Request options (e.g. request tags) for this listing operation.
    pub request_options: Option<FirestoreRequestOptions>,
}

#[derive(Debug, PartialEq, Clone, Builder)]
pub struct FirestoreListCollectionIdsResult {
    pub collection_ids: Vec<String>,
    pub page_token: Option<String>,
}

#[async_trait]
impl FirestoreListingSupport for FirestoreDb {
    async fn list_doc(
        &self,
        params: FirestoreListDocParams,
    ) -> FirestoreResult<FirestoreListDocResult> {
        let span = span!(
            Level::DEBUG,
            "Firestore ListDocs",
            "/firestore/collection_name" = params.collection_id.as_str(),
            "/firestore/response_time" = field::Empty
        );

        let list_request = self.create_list_doc_request(params)?;
        self.list_doc_page(&list_request, &span).await
    }

    async fn stream_list_doc_with_errors<'b>(
        &self,
        params: FirestoreListDocParams,
    ) -> FirestoreResult<BoxStream<'b, FirestoreResult<Document>>> {
        self.stream_list_doc_with_retries(params).await
    }

    async fn stream_list_doc<'b>(
        &self,
        params: FirestoreListDocParams,
    ) -> FirestoreResult<BoxStream<'b, Document>> {
        let doc_stream = self.stream_list_doc_with_errors(params).await?;
        Ok(Box::pin(doc_stream.filter_map(|doc_res| {
            future::ready(match doc_res {
                Ok(doc) => Some(doc),
                Err(err) => {
                    error!(%err, "Error occurred while consuming documents.");
                    None
                }
            })
        })))
    }

    async fn stream_list_obj<'b, T>(
        &self,
        params: FirestoreListDocParams,
    ) -> FirestoreResult<BoxStream<'b, T>>
    where
        for<'de> T: Deserialize<'de> + 'b,
    {
        let doc_stream = self.stream_list_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 deserializing a document inside a stream. Document: {}",
                        doc.name
                    );
                    None
                }
            }
        })))
    }

    async fn stream_list_obj_with_errors<'b, T>(
        &self,
        params: FirestoreListDocParams,
    ) -> FirestoreResult<BoxStream<'b, FirestoreResult<T>>>
    where
        for<'de> T: Deserialize<'de> + 'b,
    {
        let doc_stream = self.stream_list_doc_with_errors(params).await?;

        Ok(Box::pin(doc_stream.and_then(|doc| async move {
            Self::deserialize_doc_to::<T>(&doc)
        })))
    }

    async fn list_collection_ids(
        &self,
        params: FirestoreListCollectionIdsParams,
    ) -> FirestoreResult<FirestoreListCollectionIdsResult> {
        let span = span!(
            Level::DEBUG,
            "Firestore ListCollectionIds",
            "/firestore/response_time" = field::Empty
        );

        self.list_collection_ids_page(&params, &span).await
    }

    async fn stream_list_collection_ids(
        &self,
        params: FirestoreListCollectionIdsParams,
    ) -> FirestoreResult<BoxStream<String>> {
        let stream = self.stream_list_collection_ids_with_errors(params).await?;
        Ok(Box::pin(stream.filter_map(|col_res| {
            future::ready(match col_res {
                Ok(col) => Some(col),
                Err(err) => {
                    error!(%err, "Error occurred while consuming collection IDs.");
                    None
                }
            })
        })))
    }

    async fn stream_list_collection_ids_with_errors(
        &self,
        params: FirestoreListCollectionIdsParams,
    ) -> FirestoreResult<BoxStream<FirestoreResult<String>>> {
        let stream: BoxStream<FirestoreResult<String>> = Box::pin(
            futures::stream::unfold(Some(params), move |maybe_params| async move {
                if let Some(params) = maybe_params {
                    let span = span!(
                        Level::DEBUG,
                        "Firestore Streaming ListCollections",
                        "/firestore/response_time" = field::Empty
                    );

                    match self.list_collection_ids_page(&params, &span).await {
                        Ok(results) => {
                            if let Some(next_page_token) = results.page_token.clone() {
                                Some((Ok(results), Some(params.with_page_token(next_page_token))))
                            } else {
                                Some((Ok(results), None))
                            }
                        }
                        Err(err) => {
                            error!(%err, "Error occurred while consuming documents.");
                            Some((Err(err), None))
                        }
                    }
                } else {
                    None
                }
            })
            .flat_map(|doc_res| {
                futures::stream::iter(match doc_res {
                    Ok(results) => results
                        .collection_ids
                        .into_iter()
                        .map(Ok::<String, FirestoreError>)
                        .collect(),
                    Err(err) => vec![Err(err)],
                })
            }),
        );

        Ok(stream)
    }
}

impl FirestoreDb {
    fn create_list_doc_request(
        &self,
        params: FirestoreListDocParams,
    ) -> FirestoreResult<ListDocumentsRequest> {
        Ok(ListDocumentsRequest {
            parent: params
                .parent
                .as_ref()
                .unwrap_or_else(|| self.get_documents_path())
                .clone(),
            collection_id: params.collection_id,
            page_size: params.page_size as i32,
            page_token: params.page_token.unwrap_or_default(),
            order_by: params
                .order_by
                .map(|fields| {
                    fields
                        .into_iter()
                        .map(|field| field.to_string_format())
                        .collect::<Vec<String>>()
                        .join(", ")
                })
                .unwrap_or_default(),
            mask: params
                .return_only_fields
                .map(|masks| DocumentMask { field_paths: masks }),
            consistency_selector: self
                .session_params
                .consistency_selector
                .as_ref()
                .map(|selector| selector.try_into())
                .transpose()?,
            request_options: self.resolve_request_options(params.request_options.as_ref()),
            show_missing: false,
        })
    }

    async fn list_doc_page(
        &self,
        list_request: &ListDocumentsRequest,
        span: &Span,
    ) -> FirestoreResult<FirestoreListDocResult> {
        let begin_utc: FirestoreInstant = FirestoreInstant::now();
        let list_inner = self
            .retry_read(
                span,
                "list documents",
                list_request,
                RequestTimeouts::with_total,
                |mut client, request| async move { client.list_documents(request).await },
            )
            .await?;
        let result = FirestoreListDocResult::new(list_inner.documents).opt_page_token(
            if !list_inner.next_page_token.is_empty() {
                Some(list_inner.next_page_token)
            } else {
                None
            },
        );
        let end_query_utc: FirestoreInstant = FirestoreInstant::now();
        let listing_duration = end_query_utc.duration_since(begin_utc);

        span.record("/firestore/response_time", listing_duration.as_millis());
        span.in_scope(|| {
            debug!(
                collection_id = list_request.collection_id.as_str(),
                duration_milliseconds = listing_duration.as_millis(),
                num_documents = result.documents.len(),
                "Listed documents.",
            );
        });

        Ok(result)
    }

    async fn stream_list_doc_with_retries<'b>(
        &self,
        params: FirestoreListDocParams,
    ) -> FirestoreResult<BoxStream<'b, FirestoreResult<Document>>> {
        #[cfg(feature = "caching")]
        {
            if let FirestoreCachedValue::UseCached(stream) =
                self.list_docs_from_cache(&params).await?
            {
                return Ok(stream);
            }
        }
        let list_request = self.create_list_doc_request(params.clone())?;
        Self::stream_list_doc_with_retries_inner(self.clone(), list_request)
    }

    fn stream_list_doc_with_retries_inner<'b>(
        db: FirestoreDb,
        list_request: ListDocumentsRequest,
    ) -> FirestoreResult<BoxStream<'b, FirestoreResult<Document>>> {
        let stream: BoxStream<FirestoreResult<Document>> = Box::pin(
            futures::stream::unfold(
                (db, Some(list_request)),
                move |(db, list_request)| async move {
                    if let Some(mut list_request) = list_request {
                        let span = span!(
                            Level::DEBUG,
                            "Firestore Streaming ListDocs",
                            "/firestore/collection_name" = list_request.collection_id.as_str(),
                            "/firestore/response_time" = field::Empty
                        );
                        match db.list_doc_page(&list_request, &span).await {
                            Ok(results) => {
                                if let Some(next_page_token) = results.page_token.clone() {
                                    list_request.page_token = next_page_token;
                                    Some((Ok(results), (db, Some(list_request))))
                                } else {
                                    Some((Ok(results), (db, None)))
                                }
                            }
                            Err(err) => {
                                error!(%err, "Error occurred while consuming documents.");
                                Some((Err(err), (db, None)))
                            }
                        }
                    } else {
                        None
                    }
                },
            )
            .flat_map(|doc_res| {
                futures::stream::iter(match doc_res {
                    Ok(results) => results
                        .documents
                        .into_iter()
                        .map(Ok::<Document, FirestoreError>)
                        .collect(),
                    Err(err) => vec![Err(err)],
                })
            }),
        );

        Ok(stream)
    }

    fn create_list_collection_ids_request(
        &self,
        params: &FirestoreListCollectionIdsParams,
    ) -> FirestoreResult<ListCollectionIdsRequest> {
        Ok(ListCollectionIdsRequest {
            parent: params
                .parent
                .as_ref()
                .unwrap_or_else(|| self.get_documents_path())
                .clone(),
            page_size: params.page_size as i32,
            page_token: params.page_token.clone().unwrap_or_default(),
            consistency_selector: self
                .session_params
                .consistency_selector
                .as_ref()
                .map(|selector| selector.try_into())
                .transpose()?,
            request_options: self.resolve_request_options(params.request_options.as_ref()),
        })
    }

    async fn list_collection_ids_page(
        &self,
        params: &FirestoreListCollectionIdsParams,
        span: &Span,
    ) -> FirestoreResult<FirestoreListCollectionIdsResult> {
        let list_request = self.create_list_collection_ids_request(params)?;
        let begin_utc: FirestoreInstant = FirestoreInstant::now();
        let list_inner = self
            .retry_read(
                span,
                "list collection IDs",
                &list_request,
                RequestTimeouts::with_total,
                |mut client, request| async move { client.list_collection_ids(request).await },
            )
            .await?;
        let result = FirestoreListCollectionIdsResult::new(list_inner.collection_ids)
            .opt_page_token(if !list_inner.next_page_token.is_empty() {
                Some(list_inner.next_page_token)
            } else {
                None
            });
        let end_query_utc: FirestoreInstant = FirestoreInstant::now();
        let listing_duration = end_query_utc.duration_since(begin_utc);

        span.record("/firestore/response_time", listing_duration.as_millis());
        span.in_scope(|| {
            debug!(
                duration_milliseconds = listing_duration.as_millis(),
                "Listed collections.",
            );
        });

        Ok(result)
    }
}

/// Returns `true` when a listing request can be served from a cache.
///
/// The cached listing path returns every document of a collection, unordered and unprojected, so
/// it can only answer requests that do not ask for ordering, a field projection or a page.
#[cfg(feature = "caching")]
#[inline]
fn list_params_supported(params: &FirestoreListDocParams) -> bool {
    params.order_by.is_none() && params.return_only_fields.is_none() && params.page_token.is_none()
}

impl FirestoreDb {
    /// Serves a `list` request from the cache when the session's cache mode allows it.
    ///
    /// Returns [`FirestoreCachedValue::SkipCache`] when there is no cache attached, the
    /// collection is not preloaded, or `params` asks for ordering, a field projection or a page -
    /// none of which the cached listing path supports. Returns an error only in
    /// [`FirestoreDbSessionCacheMode::ReadCachedOnly`], when the cache cannot serve the request
    /// completely and there is nowhere else to fall back to.
    #[cfg(feature = "caching")]
    #[inline]
    pub async fn list_docs_from_cache<'b>(
        &self,
        params: &FirestoreListDocParams,
    ) -> FirestoreResult<FirestoreCachedValue<BoxStream<'b, FirestoreResult<FirestoreDocument>>>>
    {
        if let FirestoreDbSessionCacheMode::ReadCachedOnly(ref cache)
        | FirestoreDbSessionCacheMode::ReadThroughCache(ref cache) =
            self.session_params.cache_mode
        {
            // The collection id is concatenated into the cache key below; an unvalidated `/`
            // would mis-key the cache into another collection's bucket even though the
            // server-bound listing itself sends the collection as a separate field.
            crate::db::validate_path_segment(params.collection_id.as_str(), "collection_id")?;

            let span = span!(
                Level::DEBUG,
                "Firestore List Cached",
                "/firestore/collection_name" = params.collection_id,
                "/firestore/cache_result" = field::Empty,
                "/firestore/response_time" = field::Empty
            );

            let begin_query_utc: FirestoreInstant = FirestoreInstant::now();

            let collection_path = if let Some(parent) = params.parent.as_ref() {
                format!("{}/{}", parent, params.collection_id.as_str())
            } else {
                format!(
                    "{}/{}",
                    self.get_documents_path(),
                    params.collection_id.as_str()
                )
            };

            // The cached listing path ignores ordering, field projection and paging, so it must
            // not be used when any of them is requested - it would silently return data that does
            // not match what was asked for.
            let cached_result = if list_params_supported(params) {
                cache.list_all_docs(&collection_path).await?
            } else {
                FirestoreCachedValue::SkipCache
            };

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

            match cached_result {
                FirestoreCachedValue::UseCached(stream) => {
                    span.record("/firestore/cache_result", "hit");
                    span.in_scope(|| {
                        debug!(
                            collection_id = params.collection_id,
                            "Reading all documents from cache."
                        );
                    });

                    Ok(FirestoreCachedValue::UseCached(stream))
                }
                FirestoreCachedValue::SkipCache => {
                    span.record("/firestore/cache_result", "miss");
                    if matches!(
                        self.session_params.cache_mode,
                        FirestoreDbSessionCacheMode::ReadCachedOnly(_)
                    ) {
                        let reason = if list_params_supported(params) {
                            "the collection is not configured to be preloaded, so the cache only \
                             holds an arbitrary subset of it"
                        } else {
                            "the request uses ordering, field projection or paging, which the \
                             cached listing path does not support"
                        };
                        span.in_scope(|| {
                            debug!(
                                collection_id = params.collection_id,
                                "Cache cannot serve this listing completely and cache mode is ReadCachedOnly, so returning an error.",
                            );
                        });
                        Err(crate::cache_incomplete_collection_error(
                            &params.collection_id,
                            reason,
                        ))
                    } else {
                        span.in_scope(|| {
                            debug!(
                                collection_id = params.collection_id,
                                "Cache doesn't have suitable documents for specified collection, so skipping cache and reading from Firestore.",
                            );
                        });
                        Ok(FirestoreCachedValue::SkipCache)
                    }
                }
            }
        } else {
            Ok(FirestoreCachedValue::SkipCache)
        }
    }
}