1use std::collections::HashMap;
26use std::marker::PhantomData;
27use std::pin::Pin;
28
29use futures_util::{Stream, TryStreamExt};
30use turso_orm_driver::{ConnectionTrait, Row, StreamTrait};
31use turso_sql::{
32 Build, Condition, Expr, Func, IntoCondition, IntoIden, JoinType, Order, Statement, TableRef,
33 Value,
34};
35
36use crate::entity::relation::column_of;
37use crate::entity::{
38 EntityTrait, FromQueryResult, IdenStatic, Iterable, Linked, ModelTrait, PartialModelTrait,
39 PrimaryKeyToColumn, Related, RelationDef,
40};
41use crate::{DbErr, Result};
42
43pub type ModelStream<'a, M> = Pin<Box<dyn Stream<Item = Result<M>> + Send + 'a>>;
49
50#[derive(Clone, Debug)]
52pub struct Select<E: EntityTrait> {
53 query: turso_sql::Select,
55 _e: PhantomData<E>,
57}
58
59fn qualified<E: EntityTrait>(column: E::Column) -> Expr {
61 Expr::col((E::TABLE_NAME, column.as_str()))
62}
63
64fn related_condition<S: EntityTrait>(
72 rel: &RelationDef,
73 model: &S::Model,
74 target_ref: &str,
75) -> Condition {
76 let mut cond = Condition::all();
77 for (f, t) in rel.from_col.iter().zip(&rel.to_col) {
78 let value = column_of::<S>(f).map_or(Value::Null, |c| model.get(c));
79 cond = cond.add(Expr::col((target_ref.to_owned(), *t)).eq(Expr::val(value)));
80 }
81 cond
82}
83
84impl<E: EntityTrait> Select<E> {
85 pub(crate) fn new() -> Self {
90 let mut query = turso_sql::Select::new().from(E::TABLE_NAME);
91 for c in E::Column::iter() {
92 query = query.expr(qualified::<E>(c));
93 }
94 Self {
95 query,
96 _e: PhantomData,
97 }
98 }
99
100 pub(crate) fn filter_by_pk(mut self, values: Vec<Value>) -> Self {
102 for (pk, value) in E::PrimaryKey::iter().zip(values) {
103 let column = pk.into_column();
104 self.query = self
105 .query
106 .and_where(qualified::<E>(column).eq(Expr::val(value)));
107 }
108 self
109 }
110
111 pub(crate) fn find_related_to<S>(model: &S::Model) -> Self
117 where
118 S: EntityTrait + Related<E>,
119 {
120 let to = <S as Related<E>>::to();
121 let mut select = Self::new();
122 match <S as Related<E>>::via() {
123 Some(via) => {
124 let junction = select.join_table(JoinType::Inner, to.from_tbl, |r| {
125 to.join_condition_refs(r, E::TABLE_NAME)
126 });
127 select.query = select
128 .query
129 .and_where(related_condition::<S>(&via, model, &junction));
130 }
131 None => {
132 select.query =
133 select
134 .query
135 .and_where(related_condition::<S>(&to, model, E::TABLE_NAME));
136 }
137 }
138 select
139 }
140
141 pub(crate) fn find_linked_to<L>(link: &L, model: &<L::FromEntity as EntityTrait>::Model) -> Self
148 where
149 L: Linked<ToEntity = E>,
150 {
151 let mut select = Self::new();
152 let mut previous = E::TABLE_NAME.to_owned();
153 for hop in link.link().into_iter().rev() {
154 previous = select.join_table(JoinType::Inner, hop.from_tbl, |r| {
155 hop.join_condition_refs(r, &previous)
156 });
157 }
158 for pk in <L::FromEntity as EntityTrait>::PrimaryKey::iter() {
159 let column = pk.into_column();
160 select.query = select.query.and_where(
161 Expr::col((previous.clone(), column.as_str())).eq(Expr::val(model.get(column))),
162 );
163 }
164 select
165 }
166
167 pub(crate) fn join_table(
174 &mut self,
175 kind: JoinType,
176 table: &'static str,
177 on: impl FnOnce(&str) -> Expr,
178 ) -> String {
179 let occurrences = self
180 .query
181 .from_tables()
182 .iter()
183 .chain(self.query.joins().iter().map(|j| &j.table))
184 .filter(|t| t.name.name() == table)
185 .count();
186 let (table_ref, reference) = if occurrences == 0 {
187 (TableRef::new(table), table.to_owned())
188 } else {
189 let alias = format!("{table}_{occurrences}");
190 (TableRef::new(table).alias(alias.clone()), alias)
191 };
192 let cond = on(&reference);
193 self.query = std::mem::take(&mut self.query).join(kind, table_ref, cond);
194 reference
195 }
196
197 fn join_related<R: EntityTrait>(&mut self, kind: JoinType) -> String
201 where
202 E: Related<R>,
203 {
204 let to = <E as Related<R>>::to();
205 match <E as Related<R>>::via() {
206 Some(via) => {
207 let junction = self.join_table(kind, via.to_tbl, |r| {
208 via.join_condition_refs(via.from_tbl, r)
209 });
210 self.join_table(kind, to.to_tbl, |r| to.join_condition_refs(&junction, r))
211 }
212 None => self.join_table(kind, to.to_tbl, |r| to.join_condition_refs(to.from_tbl, r)),
213 }
214 }
215
216 fn join_linked<L: Linked<FromEntity = E>>(&mut self, kind: JoinType, link: &L) -> String {
219 let mut previous = E::TABLE_NAME.to_owned();
220 for hop in link.link() {
221 previous = self.join_table(kind, hop.to_tbl, |r| hop.join_condition_refs(&previous, r));
222 }
223 previous
224 }
225
226 fn into_select_two<R: EntityTrait>(self, right_ref: &str) -> SelectTwo<E, R> {
229 let mut query = self.query.clear_items();
230 for c in E::Column::iter() {
231 query = query.expr_as(qualified::<E>(c), format!("A_{}", c.as_str()));
232 }
233 for c in R::Column::iter() {
234 query = query.expr_as(
235 Expr::col((right_ref.to_owned(), c.as_str())),
236 format!("B_{}", c.as_str()),
237 );
238 }
239 SelectTwo {
240 query,
241 _e: PhantomData,
242 }
243 }
244
245 #[must_use]
247 pub fn filter(mut self, cond: impl IntoCondition) -> Self {
248 self.query = self.query.and_where(cond);
249 self
250 }
251
252 #[must_use]
254 pub fn filter_option(self, cond: Option<impl IntoCondition>) -> Self {
255 match cond {
256 Some(c) => self.filter(c),
257 None => self,
258 }
259 }
260
261 #[must_use]
268 pub fn related_to<M: ModelTrait>(mut self, rel: &RelationDef, model: &M) -> Self {
269 self.query =
270 self.query
271 .and_where(related_condition::<M::Entity>(rel, model, E::TABLE_NAME));
272 self
273 }
274
275 #[must_use]
277 pub fn order_by(mut self, column: E::Column, order: Order) -> Self {
278 self.query = self.query.order_by_expr(qualified::<E>(column), order);
279 self
280 }
281
282 #[must_use]
284 pub fn order_by_asc(self, column: E::Column) -> Self {
285 self.order_by(column, Order::Asc)
286 }
287
288 #[must_use]
290 pub fn order_by_desc(self, column: E::Column) -> Self {
291 self.order_by(column, Order::Desc)
292 }
293
294 #[must_use]
296 pub fn order_by_expr(mut self, expr: Expr, order: Order) -> Self {
297 self.query = self.query.order_by_expr(expr, order);
298 self
299 }
300
301 #[must_use]
303 pub fn limit(mut self, limit: u64) -> Self {
304 self.query = self.query.limit(limit);
305 self
306 }
307
308 #[must_use]
310 pub fn offset(mut self, offset: u64) -> Self {
311 self.query = self.query.offset(offset);
312 self
313 }
314
315 #[must_use]
317 pub fn distinct(mut self) -> Self {
318 self.query = self.query.distinct();
319 self
320 }
321
322 #[must_use]
324 pub fn group_by(mut self, column: E::Column) -> Self {
325 self.query = self.query.group_by(qualified::<E>(column));
326 self
327 }
328
329 #[must_use]
331 pub fn group_by_expr(mut self, expr: Expr) -> Self {
332 self.query = self.query.group_by(expr);
333 self
334 }
335
336 #[must_use]
338 pub fn having(mut self, cond: impl IntoCondition) -> Self {
339 self.query = self.query.and_having(cond);
340 self
341 }
342
343 #[must_use]
348 pub fn select_only(mut self) -> Self {
349 self.query = self.query.clear_items();
350 self
351 }
352
353 #[must_use]
355 pub fn column(mut self, column: E::Column) -> Self {
356 self.query = self.query.expr(qualified::<E>(column));
357 self
358 }
359
360 #[must_use]
362 pub fn column_as(mut self, column: E::Column, alias: impl IntoIden) -> Self {
363 self.query = self.query.expr_as(qualified::<E>(column), alias);
364 self
365 }
366
367 #[must_use]
369 pub fn expr(mut self, expr: Expr) -> Self {
370 self.query = self.query.expr(expr);
371 self
372 }
373
374 #[must_use]
376 pub fn expr_as(mut self, expr: Expr, alias: impl IntoIden) -> Self {
377 self.query = self.query.expr_as(expr, alias);
378 self
379 }
380
381 #[must_use]
387 pub fn join(mut self, kind: JoinType, rel: &RelationDef) -> Self {
388 self.join_table(kind, rel.to_tbl, |r| {
389 rel.join_condition_refs(rel.from_tbl, r)
390 });
391 self
392 }
393
394 #[must_use]
397 pub fn join_as(mut self, kind: JoinType, rel: &RelationDef, alias: &'static str) -> Self {
398 let on = rel.join_condition_refs(rel.from_tbl, alias);
399 self.query = self
400 .query
401 .join(kind, TableRef::new(rel.to_tbl).alias(alias), on);
402 self
403 }
404
405 #[must_use]
407 pub fn join_rev(mut self, kind: JoinType, rel: &RelationDef) -> Self {
408 self.join_table(kind, rel.from_tbl, |r| {
409 rel.join_condition_refs(r, rel.to_tbl)
410 });
411 self
412 }
413
414 #[must_use]
417 pub fn inner_join<R: EntityTrait>(mut self, _: R) -> Self
418 where
419 E: Related<R>,
420 {
421 self.join_related::<R>(JoinType::Inner);
422 self
423 }
424
425 #[must_use]
428 pub fn left_join<R: EntityTrait>(mut self, _: R) -> Self
429 where
430 E: Related<R>,
431 {
432 self.join_related::<R>(JoinType::Left);
433 self
434 }
435
436 #[must_use]
443 pub fn find_also_related<R: EntityTrait>(mut self, _: R) -> SelectTwo<E, R>
444 where
445 E: Related<R>,
446 {
447 let right = self.join_related::<R>(JoinType::Left);
448 self.into_select_two::<R>(&right)
449 }
450
451 #[must_use]
457 pub fn find_with_related<R: EntityTrait>(self, r: R) -> SelectTwoMany<E, R>
458 where
459 E: Related<R>,
460 {
461 SelectTwoMany {
462 inner: self.find_also_related(r),
463 }
464 }
465
466 #[must_use]
471 pub fn find_also_linked<L>(mut self, link: &L) -> SelectTwo<E, L::ToEntity>
472 where
473 L: Linked<FromEntity = E>,
474 {
475 let right = self.join_linked(JoinType::Left, link);
476 self.into_select_two::<L::ToEntity>(&right)
477 }
478
479 #[must_use]
481 pub fn find_with_linked<L>(self, link: &L) -> SelectTwoMany<E, L::ToEntity>
482 where
483 L: Linked<FromEntity = E>,
484 {
485 SelectTwoMany {
486 inner: self.find_also_linked(link),
487 }
488 }
489
490 pub fn into_model<M: FromQueryResult>(self) -> Selector<M> {
492 Selector {
493 query: self.query,
494 _m: PhantomData,
495 }
496 }
497
498 pub fn into_partial_model<P: PartialModelTrait>(self) -> Selector<P> {
501 Selector {
502 query: P::select_cols(self.query.clear_items()),
503 _m: PhantomData,
504 }
505 }
506
507 pub fn into_tuple<T: FromQueryResult>(self) -> Selector<T> {
512 self.into_model::<T>()
513 }
514
515 #[cfg(feature = "with-json")]
517 #[cfg_attr(docsrs, doc(cfg(feature = "with-json")))]
518 pub fn into_json(self) -> Selector<serde_json::Value> {
519 self.into_model::<serde_json::Value>()
520 }
521
522 pub fn from_raw_sql(self, statement: Statement) -> RawSelector<E::Model> {
527 Selector::<E::Model>::from_statement(statement)
528 }
529
530 pub fn as_query(&self) -> &turso_sql::Select {
532 &self.query
533 }
534
535 pub fn query_mut(&mut self) -> &mut turso_sql::Select {
537 &mut self.query
538 }
539
540 pub fn into_query(self) -> turso_sql::Select {
542 self.query
543 }
544
545 pub fn build(&self) -> Statement {
547 self.query.to_statement()
548 }
549
550 pub async fn one<C: ConnectionTrait>(self, db: &C) -> Result<Option<E::Model>> {
556 self.into_model::<E::Model>().one(db).await
557 }
558
559 pub async fn all<C: ConnectionTrait>(self, db: &C) -> Result<Vec<E::Model>> {
565 self.into_model::<E::Model>().all(db).await
566 }
567
568 pub async fn stream<C: StreamTrait>(self, db: &C) -> Result<ModelStream<'_, E::Model>> {
575 self.into_model::<E::Model>().stream(db).await
576 }
577
578 pub async fn count<C: ConnectionTrait>(self, db: &C) -> Result<u64> {
585 self.into_model::<E::Model>().count(db).await
586 }
587
588 pub async fn exists<C: ConnectionTrait>(self, db: &C) -> Result<bool> {
596 self.into_model::<E::Model>().exists(db).await
597 }
598
599 pub fn paginate<C: ConnectionTrait>(
601 self,
602 db: &C,
603 page_size: u64,
604 ) -> Paginator<'_, C, E::Model> {
605 self.into_model::<E::Model>().paginate(db, page_size)
606 }
607
608 pub fn cursor_by<const N: usize>(
613 self,
614 columns: [E::Column; N],
615 ) -> crate::query::Cursor<E, E::Model> {
616 crate::query::Cursor::new(self.query, columns.to_vec())
617 }
618}
619
620#[derive(Clone, Debug)]
622pub struct Selector<M> {
623 query: turso_sql::Select,
625 _m: PhantomData<M>,
627}
628
629impl<M: FromQueryResult> Selector<M> {
630 pub fn from_query(query: turso_sql::Select) -> Self {
632 Self {
633 query,
634 _m: PhantomData,
635 }
636 }
637
638 pub fn from_statement(statement: Statement) -> RawSelector<M> {
640 RawSelector {
641 statement,
642 _m: PhantomData,
643 }
644 }
645
646 pub fn build(&self) -> Statement {
648 self.query.to_statement()
649 }
650
651 pub async fn one<C: ConnectionTrait>(mut self, db: &C) -> Result<Option<M>> {
657 self.query = self.query.limit(1);
658 let row = db.query_one(self.build()).await?;
659 row.map(|r| M::from_query_result(&r, "")).transpose()
660 }
661
662 pub async fn all<C: ConnectionTrait>(self, db: &C) -> Result<Vec<M>> {
668 let rows = db.query_all(self.build()).await?;
669 rows.iter().map(|r| M::from_query_result(r, "")).collect()
670 }
671
672 #[allow(
679 clippy::needless_lifetimes,
680 reason = "the stream borrows `db`, not `self`, and the explicit lifetime says so"
681 )]
682 pub async fn stream<'a, C: StreamTrait>(self, db: &'a C) -> Result<ModelStream<'a, M>>
683 where
684 M: 'a,
685 {
686 let stream = db.stream(self.build()).await?;
687 Ok(Box::pin(stream.map_err(DbErr::from).and_then(
688 |row| async move { M::from_query_result(&row, "") },
689 )))
690 }
691
692 fn unordered(&self) -> turso_sql::Select {
695 self.query
696 .clone()
697 .clear_order_by()
698 .reset_limit()
699 .reset_offset()
700 }
701
702 pub async fn count<C: ConnectionTrait>(self, db: &C) -> Result<u64> {
712 let outer = turso_sql::Select::new()
713 .expr_as(Func::count_star(), "num_items")
714 .from_subquery(self.unordered(), "sub");
715 let row = db
716 .query_one(outer.to_statement())
717 .await?
718 .ok_or(DbErr::RecordNotFound("count".into()))?;
719 Ok(row.get::<u64>("num_items")?)
720 }
721
722 pub async fn exists<C: ConnectionTrait>(self, db: &C) -> Result<bool> {
729 let probe = turso_sql::Select::new().expr_as(Expr::exists(self.unordered()), "found");
730 let row = db
731 .query_one(probe.to_statement())
732 .await?
733 .ok_or(DbErr::RecordNotFound("exists".into()))?;
734 Ok(row.get::<bool>("found")?)
735 }
736
737 pub fn paginate<C: ConnectionTrait>(self, db: &C, page_size: u64) -> Paginator<'_, C, M> {
739 Paginator {
740 query: self.query,
741 page_size: page_size.max(1),
742 db,
743 _m: PhantomData,
744 }
745 }
746}
747
748#[derive(Clone, Debug)]
750pub struct RawSelector<M> {
751 statement: Statement,
753 _m: PhantomData<M>,
755}
756
757impl<M: FromQueryResult> RawSelector<M> {
758 pub async fn one<C: ConnectionTrait>(self, db: &C) -> Result<Option<M>> {
764 let row = db.query_one(self.statement).await?;
765 row.map(|r| M::from_query_result(&r, "")).transpose()
766 }
767
768 pub async fn all<C: ConnectionTrait>(self, db: &C) -> Result<Vec<M>> {
774 let rows = db.query_all(self.statement).await?;
775 rows.iter().map(|r| M::from_query_result(r, "")).collect()
776 }
777}
778
779#[derive(Debug)]
781pub struct Paginator<'db, C, M> {
782 query: turso_sql::Select,
784 page_size: u64,
786 db: &'db C,
788 _m: PhantomData<M>,
790}
791
792impl<C: ConnectionTrait, M: FromQueryResult> Paginator<'_, C, M> {
793 pub async fn fetch_page(&self, page: u64) -> Result<Vec<M>> {
799 let query = self
800 .query
801 .clone()
802 .limit(self.page_size)
803 .offset(page.saturating_mul(self.page_size));
804 Selector::<M>::from_query(query).all(self.db).await
805 }
806
807 pub async fn num_items(&self) -> Result<u64> {
814 Selector::<M>::from_query(self.query.clone())
815 .count(self.db)
816 .await
817 }
818
819 pub async fn num_pages(&self) -> Result<u64> {
825 let items = self.num_items().await?;
826 Ok(items.div_ceil(self.page_size))
827 }
828
829 pub async fn num_items_and_pages(&self) -> Result<(u64, u64)> {
835 let items = self.num_items().await?;
836 Ok((items, items.div_ceil(self.page_size)))
837 }
838
839 pub async fn for_each_page(&self, mut f: impl FnMut(Vec<M>) -> bool) -> Result<()> {
846 let mut page = 0;
847 loop {
848 let rows = self.fetch_page(page).await?;
849 let full = u64::try_from(rows.len()).unwrap_or(u64::MAX) == self.page_size;
850 if !f(rows) || !full {
851 return Ok(());
852 }
853 page += 1;
854 }
855 }
856}
857
858#[derive(Clone, Debug)]
860pub struct SelectTwo<E: EntityTrait, R: EntityTrait> {
861 query: turso_sql::Select,
863 _e: PhantomData<(E, R)>,
865}
866
867impl<E: EntityTrait, R: EntityTrait> SelectTwo<E, R> {
868 #[must_use]
870 pub fn filter(mut self, cond: impl IntoCondition) -> Self {
871 self.query = self.query.and_where(cond);
872 self
873 }
874
875 #[must_use]
877 pub fn order_by(mut self, column: E::Column, order: Order) -> Self {
878 self.query = self.query.order_by_expr(qualified::<E>(column), order);
879 self
880 }
881
882 #[must_use]
888 pub fn order_by_related(mut self, column: R::Column, order: Order) -> Self {
889 self.query = self.query.order_by_expr(qualified::<R>(column), order);
890 self
891 }
892
893 #[must_use]
895 pub fn order_by_expr(mut self, expr: Expr, order: Order) -> Self {
896 self.query = self.query.order_by_expr(expr, order);
897 self
898 }
899
900 #[must_use]
902 pub fn limit(mut self, limit: u64) -> Self {
903 self.query = self.query.limit(limit);
904 self
905 }
906
907 #[must_use]
909 pub fn offset(mut self, offset: u64) -> Self {
910 self.query = self.query.offset(offset);
911 self
912 }
913
914 pub fn query_mut(&mut self) -> &mut turso_sql::Select {
916 &mut self.query
917 }
918
919 pub fn build(&self) -> Statement {
921 self.query.to_statement()
922 }
923
924 fn decode(row: &Row) -> Result<(E::Model, Option<R::Model>)> {
930 let a = E::Model::from_query_result(row, "A_")?;
931 let b = R::Model::from_query_result_optional(row, "B_")?;
932 Ok((a, b))
933 }
934
935 pub async fn one<C: ConnectionTrait>(
941 mut self,
942 db: &C,
943 ) -> Result<Option<(E::Model, Option<R::Model>)>> {
944 self.query = self.query.limit(1);
945 db.query_one(self.build())
946 .await?
947 .as_ref()
948 .map(Self::decode)
949 .transpose()
950 }
951
952 pub async fn all<C: ConnectionTrait>(
958 self,
959 db: &C,
960 ) -> Result<Vec<(E::Model, Option<R::Model>)>> {
961 db.query_all(self.build())
962 .await?
963 .iter()
964 .map(Self::decode)
965 .collect()
966 }
967}
968
969#[derive(Clone, Debug)]
978pub struct SelectTwoMany<E: EntityTrait, R: EntityTrait> {
979 inner: SelectTwo<E, R>,
981}
982
983impl<E: EntityTrait, R: EntityTrait> SelectTwoMany<E, R> {
984 #[must_use]
986 pub fn filter(mut self, cond: impl IntoCondition) -> Self {
987 self.inner = self.inner.filter(cond);
988 self
989 }
990
991 #[must_use]
993 pub fn order_by(mut self, column: E::Column, order: Order) -> Self {
994 self.inner = self.inner.order_by(column, order);
995 self
996 }
997
998 #[must_use]
1000 pub fn order_by_related(mut self, column: R::Column, order: Order) -> Self {
1001 self.inner = self.inner.order_by_related(column, order);
1002 self
1003 }
1004
1005 #[must_use]
1007 pub fn order_by_expr(mut self, expr: Expr, order: Order) -> Self {
1008 self.inner = self.inner.order_by_expr(expr, order);
1009 self
1010 }
1011
1012 pub fn build(&self) -> Statement {
1014 self.inner.build()
1015 }
1016
1017 pub async fn all<C: ConnectionTrait>(self, db: &C) -> Result<Vec<(E::Model, Vec<R::Model>)>> {
1026 let pairs = self.inner.all(db).await?;
1027 let mut groups: Vec<(E::Model, Vec<R::Model>)> = Vec::new();
1028 let mut index: HashMap<String, usize> = HashMap::new();
1029 for (left, right) in pairs {
1030 let key = E::PrimaryKey::iter()
1031 .map(|pk| left.get(pk.into_column()).to_literal())
1032 .collect::<Vec<_>>()
1033 .join("\u{1f}");
1034 let at = if let Some(&at) = index.get(&key) {
1035 at
1036 } else {
1037 groups.push((left, Vec::new()));
1038 index.insert(key, groups.len() - 1);
1039 groups.len() - 1
1040 };
1041 if let Some(r) = right {
1042 groups[at].1.push(r);
1043 }
1044 }
1045 Ok(groups)
1046 }
1047}
1048
1049pub(crate) fn pk_condition<E: EntityTrait>(values: Vec<Value>) -> Condition {
1054 let mut cond = Condition::all();
1055 for (pk, value) in E::PrimaryKey::iter().zip(values) {
1056 cond = cond.add(qualified::<E>(pk.into_column()).eq(Expr::val(value)));
1057 }
1058 cond
1059}