rorm_db/internal/
executor.rs

1use std::future;
2use std::future::{ready, Ready};
3use std::pin::Pin;
4use std::task::{ready, Context, Poll};
5
6use futures_core::stream;
7use rorm_sql::value::Value;
8use rorm_sql::DBImpl;
9use sqlx::{AssertSqlSafe, SqlSafeStr, SqlStr};
10use tracing::debug;
11
12use crate::executor::{
13    AffectedRows, All, DynamicExecutor, Executor, Nothing, One, Optional, QueryStrategy,
14    QueryStrategyResult, Stream,
15};
16use crate::futures_util::{BoxFuture, BoxStream};
17use crate::internal::any::{AnyExecutor, AnyPool, AnyQueryResult, AnyRow, AnyTransaction};
18use crate::internal::bind_params::bind_param;
19use crate::transaction::{Transaction, TransactionGuard};
20use crate::{Database, Error, Row};
21
22impl<'exe> Executor<'exe> for &'exe mut Transaction {
23    fn execute<Q>(self, query: String, values: Vec<Value<'_>>) -> Q::Result<'exe>
24    where
25        Q: QueryStrategy,
26    {
27        debug!(
28            target: "rorm_db::executor",
29            sql = query,
30            values.len = values.len(),
31            "Executing statement"
32        );
33        Q::execute(&mut self.sqlx, AssertSqlSafe(query).into_sql_str(), values)
34    }
35
36    fn into_dyn(self) -> DynamicExecutor<'exe> {
37        DynamicExecutor::Transaction(self)
38    }
39
40    fn dialect(&self) -> DBImpl {
41        match self.sqlx {
42            #[cfg(feature = "postgres")]
43            AnyTransaction::Postgres(_) => DBImpl::Postgres,
44            #[cfg(feature = "sqlite")]
45            AnyTransaction::Sqlite(_) => DBImpl::SQLite,
46        }
47    }
48
49    type EnsureTransactionFuture = Ready<Result<TransactionGuard<'exe>, Error>>;
50
51    fn ensure_transaction(self) -> BoxFuture<'exe, Result<TransactionGuard<'exe>, Error>> {
52        Box::pin(ready(Ok(TransactionGuard::Borrowed(self))))
53    }
54}
55
56impl<'exe> Executor<'exe> for &'exe Database {
57    fn execute<Q>(self, query: String, values: Vec<Value<'_>>) -> Q::Result<'exe>
58    where
59        Q: QueryStrategy,
60    {
61        debug!(
62            target: "rorm_db::executor",
63            sql = query,
64            values.len = values.len(),
65            "Executing statement"
66        );
67        Q::execute(&self.0, AssertSqlSafe(query).into_sql_str(), values)
68    }
69
70    fn into_dyn(self) -> DynamicExecutor<'exe> {
71        DynamicExecutor::Database(self)
72    }
73
74    fn dialect(&self) -> DBImpl {
75        match self.0 {
76            #[cfg(feature = "postgres")]
77            AnyPool::Postgres(_) => DBImpl::Postgres,
78            #[cfg(feature = "sqlite")]
79            AnyPool::Sqlite(_) => DBImpl::SQLite,
80        }
81    }
82
83    type EnsureTransactionFuture = BoxFuture<'exe, Result<TransactionGuard<'exe>, Error>>;
84
85    fn ensure_transaction(self) -> BoxFuture<'exe, Result<TransactionGuard<'exe>, Error>> {
86        Box::pin(async move { self.start_transaction().await.map(TransactionGuard::Owned) })
87    }
88}
89
90pub trait QueryStrategyImpl: QueryStrategyResult {
91    fn execute<'exe, E>(executor: E, query: SqlStr, values: Vec<Value<'_>>) -> Self::Result<'exe>
92    where
93        E: AnyExecutor<'exe>;
94}
95
96impl QueryStrategyResult for Nothing {
97    type Result<'query> = NothingFuture<'query>;
98}
99
100impl QueryStrategyImpl for Nothing {
101    fn execute<'exe, E>(executor: E, query: SqlStr, values: Vec<Value<'_>>) -> Self::Result<'exe>
102    where
103        E: AnyExecutor<'exe>,
104    {
105        let mut query = executor.query(query);
106        for x in values {
107            bind_param(&mut query, x);
108        }
109        NothingFuture {
110            stream: query.fetch_many(),
111        }
112    }
113}
114
115/// [`QueryStrategyResult::Result`] of [`Nothing`]
116pub struct NothingFuture<'stream> {
117    stream: BoxStream<'stream, sqlx::Result<sqlx::Either<AnyQueryResult, AnyRow>>>,
118}
119
120impl future::Future for NothingFuture<'_> {
121    type Output = Result<(), Error>;
122
123    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
124        loop {
125            return Poll::Ready(match ready!(self.stream.as_mut().poll_next(cx)) {
126                None => Ok(()),
127                Some(Err(error)) => Err(error.into()),
128                Some(_either) => continue,
129            });
130        }
131    }
132}
133
134impl QueryStrategyResult for AffectedRows {
135    type Result<'query> = BoxFuture<'query, Result<u64, Error>>;
136}
137
138impl QueryStrategyImpl for AffectedRows {
139    fn execute<'exe, E>(executor: E, query: SqlStr, values: Vec<Value<'_>>) -> Self::Result<'exe>
140    where
141        E: AnyExecutor<'exe>,
142    {
143        let mut query = executor.query(query);
144        for x in values {
145            bind_param(&mut query, x);
146        }
147        Box::pin(async move { Ok(query.fetch_affected_rows().await?) }) as BoxFuture<_>
148    }
149}
150
151impl QueryStrategyResult for One {
152    type Result<'query> = BoxFuture<'query, Result<Row, Error>>;
153}
154
155impl QueryStrategyImpl for One {
156    fn execute<'exe, E>(executor: E, query: SqlStr, values: Vec<Value<'_>>) -> Self::Result<'exe>
157    where
158        E: AnyExecutor<'exe>,
159    {
160        let mut query = executor.query(query);
161        for x in values {
162            bind_param(&mut query, x);
163        }
164        Box::pin(async move {
165            Ok(Row(query
166                .fetch_optional()
167                .await?
168                .ok_or(sqlx::Error::RowNotFound)?))
169        }) as BoxFuture<_>
170    }
171}
172
173impl QueryStrategyResult for Optional {
174    type Result<'query> = BoxFuture<'query, Result<Option<Row>, Error>>;
175}
176
177impl QueryStrategyImpl for Optional {
178    fn execute<'exe, E>(executor: E, query: SqlStr, values: Vec<Value<'_>>) -> Self::Result<'exe>
179    where
180        E: AnyExecutor<'exe>,
181    {
182        let mut query = executor.query(query);
183        for x in values {
184            bind_param(&mut query, x);
185        }
186        Box::pin(async move { Ok(query.fetch_optional().await?.map(Row)) }) as BoxFuture<_>
187    }
188}
189
190impl QueryStrategyResult for All {
191    type Result<'query> = BoxFuture<'query, Result<Vec<Row>, Error>>;
192}
193
194impl QueryStrategyImpl for All {
195    fn execute<'exe, E>(executor: E, query: SqlStr, values: Vec<Value<'_>>) -> Self::Result<'exe>
196    where
197        E: AnyExecutor<'exe>,
198    {
199        let mut query = executor.query(query);
200        for x in values {
201            bind_param(&mut query, x);
202        }
203        Box::pin(async move { Ok(query.fetch_all().await?.into_iter().map(Row).collect()) })
204            as BoxFuture<_>
205    }
206}
207
208impl QueryStrategyResult for Stream {
209    type Result<'query> = StreamResult<'query>;
210}
211
212impl QueryStrategyImpl for Stream {
213    fn execute<'exe, E>(executor: E, query: SqlStr, values: Vec<Value<'_>>) -> Self::Result<'exe>
214    where
215        E: AnyExecutor<'exe>,
216    {
217        let mut query = executor.query(query);
218        for x in values {
219            bind_param(&mut query, x);
220        }
221        StreamResult {
222            stream: query.fetch_many(),
223        }
224    }
225}
226
227/// [`QueryStrategyResult::Result`] of [`Stream`]
228pub struct StreamResult<'stream> {
229    stream: BoxStream<'stream, sqlx::Result<sqlx::Either<AnyQueryResult, AnyRow>>>,
230}
231
232impl Unpin for StreamResult<'_> {}
233impl stream::Stream for StreamResult<'_> {
234    type Item = Result<Row, Error>;
235
236    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
237        loop {
238            return Poll::Ready(match ready!(self.stream.as_mut().poll_next(cx)) {
239                None => None,
240                Some(Err(error)) => Some(Err(error.into())),
241                Some(Ok(sqlx::Either::Right(row))) => Some(Ok(Row(row))),
242                Some(Ok(sqlx::Either::Left(_result))) => continue,
243            });
244        }
245    }
246}
rorm_db/internal/executor.rs

rorm_db/internal/
executor.rs
