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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
use crate::errors::*;
use crate::{
FirestoreBatch, FirestoreBatchWriteResponse, FirestoreBatchWriter, FirestoreDb,
FirestoreRequestOptions, FirestoreResult, FirestoreWriteResult,
};
use async_trait::async_trait;
use futures::TryFutureExt;
use gcloud_sdk::google::firestore::v1::{BatchWriteRequest, Write};
use gcloud_sdk::RequestTimeouts;
use rsb_derive::*;
use std::collections::HashMap;
use std::time::Duration;
use tracing::*;
#[derive(Debug, Eq, PartialEq, Clone, Builder)]
pub struct FirestoreSimpleBatchWriteOptions {
/// How long one `write` keeps retrying a `BatchWrite` request that failed with a retryable
/// error, counted from its first attempt, before it returns that error.
///
/// `None`, the default, retries with no limit.
pub retry_max_elapsed_time: Option<Duration>,
/// Request options (e.g. request tags) for the batch write requests.
pub request_options: Option<FirestoreRequestOptions>,
}
pub struct FirestoreSimpleBatchWriter {
pub db: FirestoreDb,
pub options: FirestoreSimpleBatchWriteOptions,
pub batch_span: Span,
}
impl FirestoreSimpleBatchWriter {
/// Wraps `db` with `options` to write batches through Firestore's non-streaming `BatchWrite`
/// RPC.
///
/// Most callers get one from [`FirestoreDb::create_simple_batch_writer`] or
/// [`FirestoreDb::create_simple_batch_writer_with_options`] rather than calling this
/// directly.
///
/// # Errors
/// Never fails; returns `FirestoreResult<Self>` for symmetry with
/// [`FirestoreStreamingBatchWriter::new`](crate::FirestoreStreamingBatchWriter::new), which
/// can.
pub async fn new(
db: FirestoreDb,
options: FirestoreSimpleBatchWriteOptions,
) -> FirestoreResult<FirestoreSimpleBatchWriter> {
let batch_span = span!(Level::DEBUG, "Firestore Batch Write");
Ok(Self {
db,
options,
batch_span,
})
}
/// Starts a new batch of writes against this writer.
///
/// Queue writes on it with [`FirestoreBatch`] methods such as `update_object`, then send it
/// with [`FirestoreBatch::write`].
pub fn new_batch(&self) -> FirestoreBatch<'_, FirestoreSimpleBatchWriter> {
FirestoreBatch::new(&self.db, self)
}
}
#[async_trait]
impl FirestoreBatchWriter for FirestoreSimpleBatchWriter {
type WriteResult = FirestoreBatchWriteResponse;
async fn write(&self, writes: Vec<Write>) -> FirestoreResult<FirestoreBatchWriteResponse> {
let backoff = backoff::ExponentialBackoffBuilder::new()
.with_max_elapsed_time(self.options.retry_max_elapsed_time)
.build();
let request = BatchWriteRequest {
database: self.db.get_database_path().to_string(),
writes,
labels: HashMap::new(),
request_options: self
.db
.resolve_request_options(self.options.request_options.as_ref()),
};
backoff::future::retry(backoff, || {
async {
let response = self
.db
.client()
.get()
.batch_write(
self.db
.request_with_timeout(request.clone(), RequestTimeouts::with_total),
)
.await
.map_err(FirestoreError::from)?;
let batch_response = response.into_inner();
let write_results: FirestoreResult<Vec<FirestoreWriteResult>> = batch_response
.write_results
.into_iter()
.map(|s| s.try_into())
.collect();
Ok(FirestoreBatchWriteResponse::new(
0,
write_results?,
batch_response.status,
))
}
.map_err(firestore_err_to_backoff)
})
.await
}
}
impl FirestoreDb {
/// Opens a batch writer that sends each batch as one Firestore `BatchWrite` request, with
/// default options.
///
/// Prefer this over [`create_streaming_batch_writer`](Self::create_streaming_batch_writer)
/// for occasional or moderate-sized batches; reach for the streaming writer only once enough
/// batches are being sent back-to-back that one request per batch becomes the bottleneck.
/// Equivalent to
/// `create_simple_batch_writer_with_options(FirestoreSimpleBatchWriteOptions::new())`.
///
/// # Errors
/// Never fails; returns `FirestoreResult<FirestoreSimpleBatchWriter>` for symmetry with
/// [`create_streaming_batch_writer`](Self::create_streaming_batch_writer), which can.
pub async fn create_simple_batch_writer(&self) -> FirestoreResult<FirestoreSimpleBatchWriter> {
self.create_simple_batch_writer_with_options(FirestoreSimpleBatchWriteOptions::new())
.await
}
/// Same as [`create_simple_batch_writer`](Self::create_simple_batch_writer), with explicit
/// `options`.
///
/// Most usefully `options.retry_max_elapsed_time`, which bounds how long a `BatchWrite`
/// request keeps retrying on a transient failure before giving up.
///
/// # Errors
/// Never fails; returns `FirestoreResult<FirestoreSimpleBatchWriter>` for symmetry with
/// [the streaming writer's equivalent](Self::create_streaming_batch_writer_with_options),
/// which can.
pub async fn create_simple_batch_writer_with_options(
&self,
options: FirestoreSimpleBatchWriteOptions,
) -> FirestoreResult<FirestoreSimpleBatchWriter> {
FirestoreSimpleBatchWriter::new(self.clone(), options).await
}
}