sqlx_cloudflare_do/
batch.rs

1use std::future::Future;
2
3use futures_core::Stream;
4use futures_util::{stream, TryFutureExt};
5use sqlx_cloudflare_core::BatchResult;
6use sqlx_core::error::Error;
7use sqlx_core::executor::Execute;
8use sqlx_core::query::Query;
9
10use crate::{js, Do, DoArguments, DoBatchResult, DoConnection, DoQueryResult};
11
12impl DoConnection {
13    /// Runs `queries` in order and atomically: if any statement fails, none
14    /// of them take effect.
15    ///
16    /// This is the Durable Object answer to a transaction, which `sql.exec`
17    /// does not allow (see [`DoTransactionManager`](crate::DoTransactionManager)).
18    /// The statements run inside [`transaction`](Self::transaction), and a
19    /// failing one rolls the others back. Returns what each statement did, in the order
20    /// given; for the rows each statement returns as well, see
21    /// [`fetch_batch`](Self::fetch_batch).
22    ///
23    /// ```no_run
24    /// # async fn transfer(
25    /// #     conn: &sqlx_cloudflare_do::DoConnection,
26    /// #     amount: i64,
27    /// #     from: i64,
28    /// #     to: i64,
29    /// # ) -> Result<(), sqlx::Error> {
30    /// let results = conn
31    ///     .execute_batch([
32    ///         sqlx::query("UPDATE accounts SET balance = balance - ?1 WHERE id = ?2")
33    ///             .bind(amount)
34    ///             .bind(from),
35    ///         sqlx::query("UPDATE accounts SET balance = balance + ?1 WHERE id = ?2")
36    ///             .bind(amount)
37    ///             .bind(to),
38    ///     ])
39    ///     .await?;
40    /// # Ok(())
41    /// # }
42    /// ```
43    ///
44    /// # Errors
45    ///
46    /// When binding an argument fails, in which case nothing ran, or when any
47    /// statement fails -- in which case the whole batch was rolled back.
48    pub fn execute_batch<'q, I>(
49        &self,
50        queries: I,
51    ) -> impl Future<Output = Result<Vec<DoQueryResult>, Error>> + Send + '_
52    where
53        I: IntoIterator<Item = Query<'q, Do, DoArguments>>,
54    {
55        self.run_batch(queries).map_ok(|results| {
56            results
57                .into_iter()
58                .map(|result| result.into_parts().1)
59                .collect()
60        })
61    }
62
63    /// Runs `queries` in order and atomically, as
64    /// [`execute_batch`](Self::execute_batch) does, and streams each
65    /// statement's rows along with what it did, in the order given.
66    ///
67    /// The shape sqlc's `:batchexec`, `:batchone` and `:batchmany` queries
68    /// take, each a stream with one item per statement -- but atomic. So
69    /// unlike a loop of queries it is not lazy: the whole batch runs when the
70    /// stream is first polled, and a failure is the one item it yields, with
71    /// nothing applied.
72    ///
73    /// ```no_run
74    /// # async fn users(conn: &sqlx_cloudflare_do::DoConnection) -> Result<(), sqlx::Error> {
75    /// use futures_util::TryStreamExt;
76    /// use sqlx::Row;
77    ///
78    /// let names: Vec<String> = conn
79    ///     .fetch_batch([1_i64, 2, 3].map(|id| {
80    ///         sqlx::query("SELECT name FROM users WHERE id = ?").bind(id)
81    ///     }))
82    ///     .map_ok(|result| result.rows()[0].get("name"))
83    ///     .try_collect()
84    ///     .await?;
85    /// # Ok(())
86    /// # }
87    /// ```
88    ///
89    /// # Errors
90    ///
91    /// The stream's one item is an error when binding an argument fails, in
92    /// which case nothing ran, or when any statement fails -- in which case
93    /// the whole batch was rolled back.
94    pub fn fetch_batch<'q, I>(
95        &self,
96        queries: I,
97    ) -> impl Stream<Item = Result<DoBatchResult, Error>> + Send + '_
98    where
99        I: IntoIterator<Item = Query<'q, Do, DoArguments>>,
100    {
101        self.run_batch(queries)
102            .map_ok(|results| stream::iter(results.into_iter().map(Ok)))
103            .try_flatten_stream()
104    }
105
106    fn run_batch<'q, I>(
107        &self,
108        queries: I,
109    ) -> impl Future<Output = Result<Vec<DoBatchResult>, Error>> + Send + '_
110    where
111        I: IntoIterator<Item = Query<'q, Do, DoArguments>>,
112    {
113        // Taken apart here, before the future, so it holds only owned SQL and
114        // arguments: the queries themselves need not outlive this call, and
115        // the transaction's callback has to own what it runs.
116        let queries = queries
117            .into_iter()
118            .map(|mut query| {
119                let arguments = query.take_arguments().map_err(Error::Encode)?;
120                Ok((query.sql(), arguments.unwrap_or_default()))
121            })
122            .collect::<Result<Vec<_>, Error>>();
123
124        async move {
125            let queries = queries?;
126
127            if queries.is_empty() {
128                return Ok(Vec::new());
129            }
130
131            // Inside a transaction, statements run as they would outside
132            // one; the transaction is what makes a failure undo the rest.
133            self.transaction(move |tx| async move {
134                queries
135                    .iter()
136                    .map(|(query, arguments)| {
137                        js::execute(&tx.sql, query.as_str(), arguments.values(), None)
138                            .map(|(rows, result)| BatchResult::new(rows, result))
139                    })
140                    .collect::<Result<Vec<_>, Error>>()
141            })
142            .await
143        }
144    }
145}
sqlx_cloudflare_do/batch.rs

sqlx_cloudflare_do/
batch.rs
