1use sqlx::any::{AnyArguments, AnyRow};
10use sqlx::{Any, Encode, Row, Type};
11
12use crate::migration::DbBackend;
13use crate::Db;
14
15pub trait AnyRowExt {
17 fn get_bool(&self, column: &str) -> Result<bool, sqlx::Error>;
20
21 fn get_text(&self, column: &str) -> Result<String, sqlx::Error>;
26
27 fn get_text_opt(&self, column: &str) -> Result<Option<String>, sqlx::Error>;
30}
31
32impl AnyRowExt for AnyRow {
33 fn get_bool(&self, column: &str) -> Result<bool, sqlx::Error> {
34 Ok(self.try_get::<i32, _>(column)? != 0)
35 }
36
37 fn get_text(&self, column: &str) -> Result<String, sqlx::Error> {
38 match self.try_get::<String, _>(column) {
39 Ok(s) => Ok(s),
40 Err(_) => decode_utf8(column, self.try_get::<Vec<u8>, _>(column)?),
41 }
42 }
43
44 fn get_text_opt(&self, column: &str) -> Result<Option<String>, sqlx::Error> {
45 match self.try_get::<Option<String>, _>(column) {
46 Ok(s) => Ok(s),
47 Err(_) => self
48 .try_get::<Option<Vec<u8>>, _>(column)?
49 .map(|b| decode_utf8(column, b))
50 .transpose(),
51 }
52 }
53}
54
55fn decode_utf8(column: &str, bytes: Vec<u8>) -> Result<String, sqlx::Error> {
58 String::from_utf8(bytes).map_err(|e| sqlx::Error::ColumnDecode {
59 index: column.to_string(),
60 source: Box::new(e),
61 })
62}
63
64pub fn on_conflict_ignore<C>(keys: impl IntoIterator<Item = C>) -> sea_query::OnConflict
71where
72 C: sea_query::IntoIden,
73{
74 use sea_query::IntoIden;
75 let keys: Vec<sea_query::DynIden> = keys.into_iter().map(IntoIden::into_iden).collect();
76 let first = keys[0].clone();
77 sea_query::OnConflict::columns(keys)
78 .update_column(first)
79 .to_owned()
80}
81
82pub fn text_cast(backend: DbBackend) -> &'static str {
88 match backend {
89 DbBackend::Mysql => "char",
90 DbBackend::Postgres | DbBackend::Sqlite => "text",
91 }
92}
93
94pub fn build<S>(backend: DbBackend, stmt: S) -> (String, sea_query::Values)
102where
103 S: sea_query::QueryStatementWriter,
104{
105 match backend {
106 DbBackend::Postgres => stmt.build(sea_query::PostgresQueryBuilder),
107 DbBackend::Mysql => stmt.build(sea_query::MysqlQueryBuilder),
108 DbBackend::Sqlite => stmt.build(sea_query::SqliteQueryBuilder),
109 }
110}
111
112pub fn insert_returning_id<I>(
127 db: &Db,
128 stmt: sea_query::InsertStatement,
129 id: I,
130) -> impl std::future::Future<Output = Result<i64, sqlx::Error>> + Send + '_
131where
132 I: sea_query::IntoIden + 'static,
133{
134 let (sql, values, returning) = render_insert(db.backend, stmt, id);
135 async move {
136 if returning {
137 bind_values(sqlx::query(&sql), values)
138 .fetch_one(&db.pool)
139 .await?
140 .try_get::<i64, _>(0)
141 } else {
142 bind_values(sqlx::query(&sql), values)
143 .execute(&db.pool)
144 .await?
145 .last_insert_id()
146 .ok_or(sqlx::Error::RowNotFound)
147 }
148 }
149}
150
151fn render_insert<I>(
155 backend: DbBackend,
156 mut stmt: sea_query::InsertStatement,
157 id: I,
158) -> (String, sea_query::Values, bool)
159where
160 I: sea_query::IntoIden + 'static,
161{
162 let returning = matches!(backend, DbBackend::Postgres | DbBackend::Sqlite);
163 if returning {
164 stmt.returning_col(id);
165 }
166 let (sql, values) = build(backend, stmt);
167 (sql, values, returning)
168}
169
170type AnyQuery<'q> = sqlx::query::Query<'q, Any, AnyArguments<'q>>;
171type AnyQueryAs<'q, O> = sqlx::query::QueryAs<'q, Any, O, AnyArguments<'q>>;
172
173fn bind_one<'q, T>(query: AnyQuery<'q>, value: T) -> AnyQuery<'q>
174where
175 T: 'q + Send + Encode<'q, Any> + Type<Any>,
176{
177 query.bind(value)
178}
179
180pub fn bind_values(mut query: AnyQuery<'_>, values: sea_query::Values) -> AnyQuery<'_> {
184 use sea_query::Value;
185 for value in values.0 {
186 query = match value {
187 Value::Bool(v) => bind_one(query, v.map(i32::from)),
191 Value::TinyInt(v) => bind_one(query, v.map(i32::from)),
192 Value::SmallInt(v) => bind_one(query, v),
193 Value::Int(v) => bind_one(query, v),
194 Value::BigInt(v) => bind_one(query, v),
195 Value::TinyUnsigned(v) => bind_one(query, v.map(i32::from)),
199 Value::SmallUnsigned(v) => bind_one(query, v.map(i32::from)),
200 Value::Unsigned(v) => bind_one(query, v.map(i64::from)),
201 Value::BigUnsigned(v) => bind_one(query, v.map(|n| n as i64)),
202 Value::Float(v) => bind_one(query, v),
203 Value::Double(v) => bind_one(query, v),
204 Value::String(v) => bind_one(query, v.map(|b| *b)),
205 Value::Char(v) => bind_one(query, v.map(|c| c.to_string())),
206 Value::Bytes(v) => bind_one(query, v.map(|b| *b)),
207 #[allow(unreachable_patterns)]
210 other => panic!("unsupported portable bind value: {other:?}"),
211 };
212 }
213 query
214}
215
216pub fn bind_values_as<O>(
218 mut query: AnyQueryAs<'_, O>,
219 values: sea_query::Values,
220) -> AnyQueryAs<'_, O> {
221 use sea_query::Value;
222 for value in values.0 {
223 query = match value {
224 Value::Bool(v) => query.bind(v.map(i32::from)),
225 Value::TinyInt(v) => query.bind(v.map(i32::from)),
226 Value::SmallInt(v) => query.bind(v),
227 Value::Int(v) => query.bind(v),
228 Value::BigInt(v) => query.bind(v),
229 Value::TinyUnsigned(v) => query.bind(v.map(i32::from)),
230 Value::SmallUnsigned(v) => query.bind(v.map(i32::from)),
231 Value::Unsigned(v) => query.bind(v.map(i64::from)),
232 Value::BigUnsigned(v) => query.bind(v.map(|n| n as i64)),
233 Value::Float(v) => query.bind(v),
234 Value::Double(v) => query.bind(v),
235 Value::String(v) => query.bind(v.map(|b| *b)),
236 Value::Char(v) => query.bind(v.map(|c| c.to_string())),
237 Value::Bytes(v) => query.bind(v.map(|b| *b)),
238 #[allow(unreachable_patterns)]
239 other => panic!("unsupported portable bind value: {other:?}"),
240 };
241 }
242 query
243}
244
245#[cfg(test)]
246mod tests {
247 use super::*;
248 use sea_query::{Alias, Expr, Iden, Query};
249 use sqlx::AnyPool;
250
251 #[derive(Iden)]
252 enum Widget {
253 Table,
254 Id,
255 Label,
256 Qty,
257 }
258
259 async fn pool() -> AnyPool {
260 sqlx::any::install_default_drivers();
261 let pool = sqlx::any::AnyPoolOptions::new()
262 .max_connections(1)
263 .connect("sqlite::memory:")
264 .await
265 .unwrap();
266 sqlx::raw_sql(
267 "create table widget (id text primary key, label text not null, qty integer not null)",
268 )
269 .execute(&pool)
270 .await
271 .unwrap();
272 pool
273 }
274
275 #[tokio::test]
276 async fn binds_parameters_on_insert_and_select() {
277 let pool = pool().await;
278 let backend = DbBackend::Sqlite;
279
280 let insert = Query::insert()
281 .into_table(Widget::Table)
282 .columns([Widget::Id, Widget::Label, Widget::Qty])
283 .values_panic(["w-1".into(), "Sprocket".into(), 7.into()])
284 .to_owned();
285 let (sql, values) = build(backend, insert);
286 bind_values(sqlx::query(&sql), values)
287 .execute(&pool)
288 .await
289 .unwrap();
290
291 let select = Query::select()
293 .column(Widget::Label)
294 .from(Widget::Table)
295 .and_where(Expr::col(Widget::Id).eq("w-1"))
296 .to_owned();
297 let (sql, values) = build(backend, select);
298 let label: String = bind_values_as(sqlx::query_as::<_, (String,)>(&sql), values)
299 .fetch_one(&pool)
300 .await
301 .unwrap()
302 .0;
303 assert_eq!(label, "Sprocket");
304
305 let count_stmt = Query::select()
307 .expr(Expr::col(Widget::Id).count())
308 .from(Widget::Table)
309 .and_where(Expr::col(Alias::new("qty")).eq(7))
310 .to_owned();
311 let (sql, values) = build(backend, count_stmt);
312 let count: i64 = bind_values_as(sqlx::query_as::<_, (i64,)>(&sql), values)
313 .fetch_one(&pool)
314 .await
315 .unwrap()
316 .0;
317 assert_eq!(count, 1);
318 }
319}