sqlx_cloudflare_d1/
batch.rs

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

sqlx_cloudflare_d1/
batch.rs
