firestore 0.57.1

// Deliberately public: `FirestoreTransactionOps` exists so that transaction operations are
// available on both `FirestoreTransaction` and `FirestoreTransactionData`, which lets callers
// write transaction-agnostic abstractions over the trait. See
// https://github.com/abdolence/firestore-rs/issues/206.
pub use crate::db::transaction_ops::FirestoreTransactionOps;

use crate::errors::*;
use crate::timestamp_utils::from_timestamp;
use crate::{
    FirestoreConsistencySelector, FirestoreDb, FirestoreError, FirestoreRequestOptions,
    FirestoreResult, FirestoreTransactionId, FirestoreTransactionMode, FirestoreTransactionOptions,
    FirestoreTransactionResponse, FirestoreWriteResult,
};
use backoff::backoff::Backoff;
use backoff::ExponentialBackoffBuilder;
use futures::future::BoxFuture;
use gcloud_sdk::google::firestore::v1::{BeginTransactionRequest, CommitRequest, RollbackRequest};
use gcloud_sdk::tonic::Code;
use std::time::Duration;
use tracing::*;

#[cfg(test)]
mod commit_tests;

#[derive(Debug, Clone)]
pub struct FirestoreTransactionData {
    transaction_id: FirestoreTransactionId,
    document_path: String,
    transaction_span: Span,
    writes: Vec<gcloud_sdk::google::firestore::v1::Write>,
}

impl FirestoreTransactionData {
    /// Assembles a transaction snapshot from its parts, bypassing `begin_transaction`.
    ///
    /// Most callers get a value of this type from [`FirestoreTransaction::into_data`] instead of
    /// calling this directly.
    pub fn new(
        transaction_id: FirestoreTransactionId,
        document_path: String,
        transaction_span: Span,
        writes: Vec<gcloud_sdk::google::firestore::v1::Write>,
    ) -> Self {
        Self {
            transaction_id,
            document_path,
            transaction_span,
            writes,
        }
    }

    /// Returns the transaction ID Firestore assigned when the transaction began.
    #[inline]
    pub fn transaction_id(&self) -> &FirestoreTransactionId {
        &self.transaction_id
    }

    /// Returns the documents path that this transaction's writes are resolved against.
    #[inline]
    pub fn documents_path(&self) -> &String {
        &self.document_path
    }

    /// Returns the tracing span opened for this transaction.
    #[inline]
    pub fn transaction_span(&self) -> &Span {
        &self.transaction_span
    }

    /// Returns the writes queued on this transaction so far, in the order they were added.
    #[inline]
    pub fn writes(&self) -> &Vec<gcloud_sdk::google::firestore::v1::Write> {
        &self.writes
    }

    /// Returns whether the transaction has no queued writes yet.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.writes.is_empty()
    }
}

impl From<FirestoreTransaction<'_>> for FirestoreTransactionData {
    fn from(transaction: FirestoreTransaction) -> Self {
        transaction.into_data()
    }
}

impl FirestoreTransactionOps for FirestoreTransactionData {
    fn add<I>(&mut self, write: I) -> FirestoreResult<&mut Self>
    where
        I: TryInto<gcloud_sdk::google::firestore::v1::Write, Error = FirestoreError>,
    {
        let write = write.try_into()?;
        self.writes.push(write);
        Ok(self)
    }

    fn get_documents_path(&self) -> &String {
        &self.document_path
    }
}

#[derive(Debug)]
pub struct FirestoreTransaction<'a> {
    db: &'a FirestoreDb,
    data: FirestoreTransactionData,
    finished: bool,
    request_options: Option<FirestoreRequestOptions>,
}

impl<'a> FirestoreTransaction<'a> {
    /// Opens a new transaction against `db` via Firestore's `BeginTransaction` RPC.
    ///
    /// Most callers reach this through [`FirestoreDb::begin_transaction`] or
    /// [`FirestoreDb::begin_transaction_with_options`] rather than calling it directly.
    ///
    /// Returns an error if converting a read-only `options.mode`'s consistency selector fails,
    /// or if the `BeginTransaction` request itself fails.
    pub async fn new(
        db: &'a FirestoreDb,
        options: FirestoreTransactionOptions,
    ) -> FirestoreResult<FirestoreTransaction<'a>> {
        let transaction_span = span!(
            Level::DEBUG,
            "Firestore Transaction",
            "/firestore/transaction_id" = field::Empty,
            "/firestore/commit_time" = field::Empty
        );

        let request = gcloud_sdk::tonic::Request::new(BeginTransactionRequest {
            database: db.get_database_path().clone(),
            options: Some(options.clone().try_into()?),
            request_options: db.resolve_request_options(options.request_options.as_ref()),
        });

        let response = db
            .client()
            .get()
            .begin_transaction(request)
            .await?
            .into_inner();

        let mut hex_trans_id = hex::encode(&response.transaction);
        hex_trans_id.truncate(16);

        transaction_span.record("/firestore/transaction_id", hex_trans_id);

        transaction_span.in_scope(|| {
            debug!(mode = ?options.mode, "Created a new transaction.");
        });

        let data = FirestoreTransactionData {
            transaction_id: response.transaction,
            document_path: db.get_documents_path().clone(),
            transaction_span,
            writes: Vec::new(),
        };

        Ok(Self {
            db,
            data,
            finished: false,
            request_options: options.request_options,
        })
    }

    /// Returns the transaction ID Firestore assigned when the transaction began.
    #[inline]
    pub fn transaction_id(&self) -> &FirestoreTransactionId {
        &self.data.transaction_id
    }

    /// Returns the client the transaction was opened on.
    #[inline]
    pub fn db(&self) -> &'a FirestoreDb {
        self.db
    }

    /// Commits every write queued on the transaction and consumes it.
    ///
    /// This is the only one of [`commit`](Self::commit), [`rollback`](Self::rollback) and
    /// [`finish`](Self::finish) that sends the queued writes to Firestore; the other two discard
    /// them. A transaction with no queued writes still commits successfully - useful for a
    /// read-only transaction that only needed a consistent snapshot.
    ///
    /// Returns an error if the `Commit` request fails. Only `ABORTED` sets `retry_possible`,
    /// meaning the whole transaction may run again. Any other failure, a dropped connection
    /// included, leaves the outcome unknown: the writes may already be applied.
    pub async fn commit(mut self) -> FirestoreResult<FirestoreTransactionResponse> {
        self.finished = true;

        if self.data.writes.is_empty() {
            self.data.transaction_span.in_scope(|| {
                debug!("Transaction has been committed without any writes.");
            });
        }

        let request = gcloud_sdk::tonic::Request::new(CommitRequest {
            database: self.db.get_database_path().clone(),
            writes: std::mem::take(&mut self.data.writes),
            transaction: self.data.transaction_id.clone(),
            request_options: self
                .db
                .resolve_request_options(self.request_options.as_ref()),
        });

        let response = self
            .db
            .client()
            .get()
            .commit(request)
            .await
            .map_err(|status| {
                let retry_possible = status.code() == Code::Aborted;
                match FirestoreError::from(status) {
                    FirestoreError::DatabaseError(error) => {
                        FirestoreError::DatabaseError(FirestoreDatabaseError {
                            retry_possible,
                            ..error
                        })
                    }
                    other => other,
                }
            })?
            .into_inner();

        let result = FirestoreTransactionResponse::new(
            response
                .write_results
                .into_iter()
                .map(|s| s.try_into())
                .collect::<FirestoreResult<Vec<FirestoreWriteResult>>>()?,
        )
        .opt_commit_time(response.commit_time.map(from_timestamp).transpose()?);

        if let Some(ref commit_time) = result.commit_time {
            self.data
                .transaction_span
                .record("/firestore/commit_time", commit_time.to_string());
        }

        self.data.transaction_span.in_scope(|| {
            debug!("Transaction has been committed.");
        });

        Ok(result)
    }

    /// Discards every write queued on the transaction and consumes it.
    ///
    /// Unlike [`finish`](Self::finish), this tells Firestore to release the transaction's locks
    /// immediately, via the `Rollback` request, rather than waiting for the transaction to
    /// expire on its own.
    ///
    /// Returns an error if the `Rollback` request fails; the queued writes are discarded locally
    /// either way.
    pub async fn rollback(mut self) -> FirestoreResult<()> {
        self.finished = true;
        let request = gcloud_sdk::tonic::Request::new(RollbackRequest {
            database: self.db.get_database_path().clone(),
            transaction: self.data.transaction_id.clone(),
            request_options: self
                .db
                .resolve_request_options(self.request_options.as_ref()),
        });

        self.db.client().get().rollback(request).await?;

        self.data.transaction_span.in_scope(|| {
            debug!("Transaction has been rolled back.");
        });

        Ok(())
    }

    /// Marks the transaction finished locally, without telling Firestore.
    ///
    /// Unlike [`rollback`](Self::rollback), this sends no request and leaves the transaction
    /// open on the server until Firestore expires it. Call [`rollback`](Self::rollback) to
    /// release its locks.
    ///
    /// # Errors
    /// Never fails; returns `FirestoreResult<()>` for symmetry with [`commit`](Self::commit) and
    /// [`rollback`](Self::rollback).
    pub fn finish(&mut self) -> FirestoreResult<()> {
        self.finished = true;
        self.data.transaction_span.in_scope(|| {
            debug!("Transaction has been finished locally without rolling back to be able to retry it again.");
        });
        Ok(())
    }

    /// Returns whether the transaction has no queued writes yet.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Returns the transaction's underlying data without consuming it.
    #[inline]
    pub fn transaction_data(&self) -> &FirestoreTransactionData {
        &self.data
    }

    /// Reattaches a [`FirestoreTransactionData`] captured by [`into_data`](Self::into_data) to
    /// `db`, resuming the same transaction.
    ///
    /// The rebuilt value's drop warning is armed again, so commit, rollback or finish it before
    /// it goes out of scope.
    pub fn from_data(
        db: &'a FirestoreDb,
        data: FirestoreTransactionData,
    ) -> FirestoreTransaction<'a> {
        Self {
            db,
            data,
            finished: false,
            request_options: None,
        }
    }

    /// Detaches the transaction's data from its `db` borrow, consuming self.
    ///
    /// Neither commits nor rolls back. Use this to carry a transaction across an `.await`
    /// boundary or a function call that cannot hold the `&FirestoreDb` borrow, and reattach it
    /// later with [`from_data`](Self::from_data).
    #[inline]
    pub fn into_data(mut self) -> FirestoreTransactionData {
        self.finished = true;
        FirestoreTransactionData::new(
            self.data.transaction_id.clone(),
            self.data.document_path.clone(),
            self.data.transaction_span.clone(),
            std::mem::take(&mut self.data.writes),
        )
    }
}

impl FirestoreTransactionOps for FirestoreTransaction<'_> {
    fn add<I>(&mut self, write: I) -> FirestoreResult<&mut Self>
    where
        I: TryInto<gcloud_sdk::google::firestore::v1::Write, Error = FirestoreError>,
    {
        self.data.add(write)?;
        Ok(self)
    }

    fn get_documents_path(&self) -> &String {
        self.data.get_documents_path()
    }
}

impl<'a> Drop for FirestoreTransaction<'a> {
    fn drop(&mut self) {
        if !self.finished {
            self.data
                .transaction_span
                .in_scope(|| warn!("Transaction was neither committed nor rolled back."));
        }
    }
}

impl FirestoreDb {
    /// Opens a new read-write transaction with default options.
    ///
    /// Equivalent to `begin_transaction_with_options(FirestoreTransactionOptions::new())`; see
    /// [`begin_transaction_with_options`](Self::begin_transaction_with_options) for the caller's
    /// obligations and the failure mode, and [`run_transaction`](Self::run_transaction) for a
    /// version that commits and retries for you.
    pub async fn begin_transaction(&self) -> FirestoreResult<FirestoreTransaction<'_>> {
        Self::begin_transaction_with_options(self, FirestoreTransactionOptions::new()).await
    }

    /// Opens a new transaction with `options`, borrowing `self` for the transaction's lifetime.
    ///
    /// The returned [`FirestoreTransaction`] must be committed, rolled back or finished
    /// explicitly - dropping it without doing so only logs a warning and leaves the transaction
    /// open on the server until Firestore's own expiry closes it. Prefer
    /// [`run_transaction_with_options`](Self::run_transaction_with_options) unless you need
    /// direct control over commit and retry.
    ///
    /// Returns an error if converting a read-only `options.mode`'s consistency selector fails,
    /// or if the `BeginTransaction` request itself fails.
    pub async fn begin_transaction_with_options(
        &self,
        options: FirestoreTransactionOptions,
    ) -> FirestoreResult<FirestoreTransaction<'_>> {
        FirestoreTransaction::new(self, options).await
    }

    /// Runs `func` inside a transaction, retrying and committing automatically.
    ///
    /// Equivalent to [`run_transaction_with_options`](Self::run_transaction_with_options) with
    /// [`FirestoreTransactionOptions::new()`]. `func` receives a [`FirestoreDb`] bound to the
    /// transaction's read consistency and the open [`FirestoreTransaction`]; queue writes on the
    /// latter through [`FirestoreTransactionOps`] (`update_object`, `delete_by_id`, `transform`,
    /// ...) or a fluent chain's `add_to_transaction`. A Firestore transaction cannot create a
    /// document with a server-generated ID, so `update_object` with an explicit document ID -
    /// not an insert - is how a transaction creates one.
    ///
    /// On a transient failure - `func` returning [`BackoffError::Transient`], or an `ABORTED`
    /// answer to `Commit` (see [`FirestoreTransaction::commit`]) - the whole transaction is
    /// retried from the start: after the error's `retry_after` when it names one, and after an
    /// exponential backoff otherwise. It is retried at most 4 times, 5 attempts in all, with no
    /// time limit; [`run_transaction_with_options`](Self::run_transaction_with_options) changes
    /// both bounds. Once they are reached, the last attempt's error is returned. Any other commit
    /// error is returned without running `func` again.
    /// `func` must therefore be safe to run more than once for the same call: read the state it
    /// needs from the transaction-scoped `db` argument on every invocation rather than closing over
    /// state read before the transaction started, and avoid side effects inside the closure that
    /// are not themselves safe to repeat, such as an external HTTP call. A convenience of this: `?`
    /// on a plain [`FirestoreResult`] inside the closure already converts to
    /// `BackoffError::Transient` for free, through `backoff::Error`'s blanket `From` impl, so the
    /// ordinary error paths already retry; return `Err(BackoffError::Permanent(err))` explicitly
    /// for an error that must not be retried. A permanent error rolls the transaction back and is
    /// returned wrapped in [`FirestoreError::ErrorInTransaction`].
    ///
    /// # Examples
    /// ```rust,no_run
    /// use firestore::*;
    /// use futures::FutureExt;
    /// use serde::{Deserialize, Serialize};
    ///
    /// #[derive(Debug, Clone, Deserialize, Serialize)]
    /// struct Counter {
    ///     value: i64,
    /// }
    ///
    /// # async fn example(db: FirestoreDb) -> FirestoreResult<()> {
    /// db.run_transaction(|db, transaction| {
    ///     async move {
    ///         let mut counter: Counter = db
    ///             .fluent()
    ///             .select()
    ///             .by_id_in("counters")
    ///             .obj()
    ///             .one("visits")
    ///             .await?
    ///             .unwrap_or(Counter { value: 0 });
    ///
    ///         counter.value += 1;
    ///
    ///         db.fluent()
    ///             .update()
    ///             .in_col("counters")
    ///             .document_id("visits")
    ///             .object(&counter)
    ///             .add_to_transaction(transaction)?;
    ///
    ///         Ok(())
    ///     }
    ///     .boxed()
    /// })
    /// .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn run_transaction<T, FN, E>(&self, func: FN) -> FirestoreResult<T>
    where
        for<'b> FN: Fn(
            FirestoreDb,
            &'b mut FirestoreTransaction,
        ) -> BoxFuture<'b, std::result::Result<T, BackoffError<E>>>,
        E: std::error::Error + Send + Sync + 'static,
    {
        self.run_transaction_with_options(func, FirestoreTransactionOptions::new())
            .await
    }

    /// Same as [`run_transaction`](Self::run_transaction), with explicit `options`.
    ///
    /// Two options bound the retries described on [`run_transaction`](Self::run_transaction):
    /// `options.max_retries` caps how many retries follow the first attempt (4 by default, 5
    /// attempts in all), and `options.max_elapsed_time`, when set, also caps how long retrying
    /// may go on, counted from the first failure. Retrying stops at whichever is reached first.
    ///
    /// Returns an error if `BeginTransaction` fails, `func` returns [`BackoffError::Permanent`],
    /// the commit fails without `retry_possible`, or retries are exhausted. Callback errors are
    /// wrapped in [`FirestoreError::ErrorInTransaction`]; rollback is attempted on the failed
    /// transaction before retrying or returning the error. A rollback error does not replace the
    /// callback error.
    pub async fn run_transaction_with_options<T, FN, E>(
        &self,
        func: FN,
        options: FirestoreTransactionOptions,
    ) -> FirestoreResult<T>
    where
        for<'b> FN: Fn(
            FirestoreDb,
            &'b mut FirestoreTransaction,
        ) -> BoxFuture<'b, std::result::Result<T, BackoffError<E>>>,
        E: std::error::Error + Send + Sync + 'static,
    {
        let max_elapsed_time = options
            .max_elapsed_time
            .map(Duration::try_from)
            .transpose()?;
        // The first BeginTransaction error is returned as it is, and every retry names the first
        // transaction's ID in `ReadWriteRetry`.
        let transaction = self.begin_transaction_with_options(options.clone()).await?;
        let retry_options = FirestoreTransactionOptions {
            mode: FirestoreTransactionMode::ReadWriteRetry(transaction.transaction_id().clone()),
            ..options
        };
        let mut result = self.run_transaction_attempt(transaction, &func).await;
        // Built only now, so that `max_elapsed_time` counts from the first failure.
        let mut backoff = ExponentialBackoffBuilder::new()
            .with_max_elapsed_time(max_elapsed_time)
            .build();
        let mut retries = 0;

        loop {
            let (err, retry_after) = match result {
                Ok(value) => return Ok(value),
                Err(BackoffError::Permanent(err)) => return Err(err),
                Err(BackoffError::Transient { err, retry_after }) => (err, retry_after),
            };
            if retries == retry_options.max_retries {
                return Err(err);
            }
            let Some(delay) = retry_after.or_else(|| backoff.next_backoff()) else {
                return Err(err);
            };
            tokio::time::sleep(delay).await;
            retries += 1;
            result = match self
                .begin_transaction_with_options(retry_options.clone())
                .await
            {
                Ok(transaction) => self.run_transaction_attempt(transaction, &func).await,
                Err(err) => Err(firestore_err_to_backoff(err)),
            };
        }
    }

    /// Runs `func` in `transaction` and commits it, or rolls it back if `func` fails.
    ///
    /// A `func` error comes back wrapped in [`FirestoreError::ErrorInTransaction`], still
    /// transient or permanent as `func` returned it; a commit error is transient only when it is
    /// `retry_possible`.
    async fn run_transaction_attempt<T, FN, E>(
        &self,
        mut transaction: FirestoreTransaction<'_>,
        func: &FN,
    ) -> std::result::Result<T, BackoffError<FirestoreError>>
    where
        for<'b> FN: Fn(
            FirestoreDb,
            &'b mut FirestoreTransaction,
        ) -> BoxFuture<'b, std::result::Result<T, BackoffError<E>>>,
        E: std::error::Error + Send + Sync + 'static,
    {
        let transaction_id = transaction.transaction_id().clone();
        let transaction_span = transaction.data.transaction_span.clone();
        let db = self.clone_with_consistency_selector(FirestoreConsistencySelector::Transaction(
            transaction_id.clone(),
        ));

        let result = match func(db, &mut transaction).await {
            Ok(value) => transaction
                .commit()
                .await
                .map(|_| value)
                .map_err(firestore_err_to_backoff),
            Err(err) => {
                // Reads can hold locks even without queued writes, so release them now rather
                // than leaving them until the transaction expires.
                transaction.rollback().await.ok();
                let in_transaction = |err: E| {
                    FirestoreError::ErrorInTransaction(FirestoreErrorInTransaction::new(
                        transaction_id,
                        Box::new(err),
                    ))
                };
                Err(match err {
                    BackoffError::Transient { err, retry_after } => BackoffError::Transient {
                        err: in_transaction(err),
                        retry_after,
                    },
                    BackoffError::Permanent(err) => BackoffError::Permanent(in_transaction(err)),
                })
            }
        };

        if let Err(BackoffError::Transient { err, retry_after }) = &result {
            transaction_span.in_scope(|| {
                warn!(%err, delay = ?retry_after, "Transient error occurred in transaction.");
            });
        }
        result
    }
}

#[cfg(test)]
mod cleanup_tests;