siderite_orm/queryset/
fetch.rs

1//! Terminal methods that read models.
2
3use super::QuerySet;
4use super::related::attach_all;
5use crate::backend::{QueryResult, Row};
6use crate::error::{OrmError, QueryError};
7use crate::expr::Expr;
8use crate::model::Model;
9use crate::persist::{pk_field, pk_filter};
10use crate::plan::{DistinctMode, QueryPlan, SelectExpr};
11use crate::types::DbType;
12use std::collections::HashMap;
13use std::hash::Hash;
14
15/// Models decoded from fetched rows, each with its row.
16pub(super) type Decoded<M> = Result<Vec<(M, Row)>, OrmError>;
17
18/// One page of a [`QuerySet::paginate`] result.
19#[derive(Debug, Clone, PartialEq)]
20pub struct Page<M> {
21    /// Rows of this page.
22    pub items: Vec<M>,
23    /// Rows matching the queryset over all pages.
24    pub total: u64,
25    /// 1-based page number.
26    pub page: u64,
27    /// Requested page size.
28    pub per_page: u64,
29}
30
31impl<M> Page<M> {
32    /// Number of pages needed for [`total`](Self::total) rows.
33    pub fn total_pages(&self) -> u64 {
34        self.total.div_ceil(self.per_page.max(1))
35    }
36}
37
38/// Whether `plan` can be counted / tested for existence with a plain
39/// aggregate instead of a derived table.
40pub(super) fn is_flat(plan: &QueryPlan) -> bool {
41    plan.grouping.is_empty()
42        && plan.having.is_none()
43        && plan.distinct == DistinctMode::None
44        && plan.limit.is_none()
45        && plan.offset.is_none()
46        && plan.compound.is_empty()
47}
48
49/// `plan` without row locks and, unless a limit makes it meaningful, without
50/// ordering (PostgreSQL rejects `ORDER BY` next to aggregates).
51pub(super) fn unordered(plan: &QueryPlan) -> QueryPlan {
52    let mut plan = plan.clone();
53    plan.lock = None;
54    if plan.limit.is_none() && plan.offset.is_none() {
55        plan.ordering.clear();
56    }
57    plan
58}
59
60/// The first column of the first row as an unsigned count.
61pub(super) fn scalar_count(result: &QueryResult) -> Result<u64, QueryError> {
62    let row = result.rows.first().ok_or_else(|| QueryError::Decode {
63        column: "count".into(),
64        reason: "no row returned".into(),
65    })?;
66    row.decode_at::<i64>(0)
67        .map(|n| u64::try_from(n).unwrap_or(0))
68}
69
70impl<M: Model> QuerySet<M> {
71    /// Run the query and decode every row.
72    ///
73    /// # Errors
74    /// Capability, backend or decode errors.
75    pub async fn all(self) -> Result<Vec<M>, OrmError> {
76        Ok(self
77            .all_annotated()
78            .await?
79            .into_iter()
80            .map(|(model, _)| model)
81            .collect())
82    }
83
84    /// Like [`all`](Self::all), pairing each model with its full result row,
85    /// so `annotate`d columns can be read with [`Row::get_as`].
86    ///
87    /// # Errors
88    /// As [`all`](Self::all).
89    pub async fn all_annotated(self) -> Result<Vec<(M, Row)>, OrmError> {
90        self.ready()?;
91        if self.empty {
92            return Ok(Vec::new());
93        }
94        let result = self.db.fetch(&self.plan).await?;
95        self.decode_rows(result.rows).await
96    }
97
98    /// Decode fetched `rows` into models, attaching `select_related`
99    /// objects and running the prefetches.
100    pub(super) async fn decode_rows(&self, rows: Vec<Row>) -> Decoded<M> {
101        let mut models = Vec::with_capacity(rows.len());
102        for row in &rows {
103            let mut model = M::from_row(row, "")?;
104            attach_all(&mut model, &self.related, row)?;
105            models.push(model);
106        }
107        self.run_prefetches(&mut models).await?;
108        Ok(models.into_iter().zip(rows).collect())
109    }
110
111    /// The first row, if any. An unordered queryset is ordered by primary key.
112    ///
113    /// # Errors
114    /// As [`all`](Self::all).
115    pub async fn first(self) -> Result<Option<M>, OrmError> {
116        let ordered = if self.plan.ordering.is_empty() {
117            let pk = pk_field::<M>()?.column;
118            self.with_plan(|p| p.order_by(Expr::col(pk), crate::OrderDirection::Asc))
119        } else {
120            self
121        };
122        Ok(ordered.limit(1).all().await?.into_iter().next())
123    }
124
125    /// The last row, if any: [`first`](Self::first) on the reversed ordering.
126    ///
127    /// # Errors
128    /// As [`all`](Self::all).
129    pub async fn last(self) -> Result<Option<M>, OrmError> {
130        self.reverse().first().await
131    }
132
133    /// The row with the smallest `field`.
134    ///
135    /// # Errors
136    /// [`QueryError::DoesNotExist`] on an empty queryset, or as [`all`](Self::all).
137    pub async fn earliest(self, field: impl Into<Expr>) -> Result<M, OrmError> {
138        self.order_by([field.into().asc()]).first_or_missing().await
139    }
140
141    /// The row with the largest `field`.
142    ///
143    /// # Errors
144    /// [`QueryError::DoesNotExist`] on an empty queryset, or as [`all`](Self::all).
145    pub async fn latest(self, field: impl Into<Expr>) -> Result<M, OrmError> {
146        self.order_by([field.into().desc()])
147            .first_or_missing()
148            .await
149    }
150
151    async fn first_or_missing(self) -> Result<M, OrmError> {
152        self.first()
153            .await?
154            .ok_or_else(|| QueryError::DoesNotExist.into())
155    }
156
157    /// Exactly one row matching `predicate`.
158    ///
159    /// # Errors
160    /// [`QueryError::DoesNotExist`], [`QueryError::MultipleObjectsReturned`],
161    /// or as [`all`](Self::all).
162    pub async fn get(self, predicate: Expr) -> Result<M, OrmError> {
163        self.filter(predicate).one().await
164    }
165
166    /// The only row of this queryset.
167    pub(super) async fn one(self) -> Result<M, OrmError> {
168        let mut rows = self.limit(2).all().await?;
169        match rows.len() {
170            0 => Err(QueryError::DoesNotExist.into()),
171            1 => rows.pop().ok_or_else(|| QueryError::DoesNotExist.into()),
172            n => Err(QueryError::MultipleObjectsReturned(n as u64).into()),
173        }
174    }
175
176    /// Number of matching rows.
177    ///
178    /// # Errors
179    /// Capability or backend errors.
180    pub async fn count(self) -> Result<u64, OrmError> {
181        self.ready()?;
182        if self.empty {
183            return Ok(0);
184        }
185        let count = Expr::from(crate::Count::all());
186        let plan = if is_flat(&self.plan) {
187            let mut plan = unordered(&self.plan);
188            plan.projection = vec![SelectExpr::new(count, Some("count".into()))];
189            plan
190        } else {
191            QueryPlan::from_subquery(unordered(&self.plan), "counted").select(count, Some("count"))
192        };
193        Ok(scalar_count(&self.db.fetch(&plan).await?)?)
194    }
195
196    /// Whether at least one row matches.
197    ///
198    /// # Errors
199    /// Capability or backend errors.
200    pub async fn exists(self) -> Result<bool, OrmError> {
201        self.ready()?;
202        if self.empty {
203            return Ok(false);
204        }
205        if !is_flat(&self.plan) {
206            return Ok(self.count().await? > 0);
207        }
208        let mut plan = unordered(&self.plan);
209        plan.projection = vec![SelectExpr::new(Expr::val(1), Some("one".into()))];
210        plan.limit = Some(1);
211        Ok(!self.db.fetch(&plan).await?.rows.is_empty())
212    }
213
214    /// Whether `object` is among the matching rows (by primary key).
215    ///
216    /// # Errors
217    /// Capability or backend errors.
218    pub async fn contains(self, object: &M) -> Result<bool, OrmError> {
219        let predicate = pk_filter::<M>(&object.pk())?;
220        self.filter(predicate).exists().await
221    }
222
223    /// One page of the queryset (`page` starts at 1) plus the total count.
224    ///
225    /// The plan is checked against the backend before the count runs.
226    ///
227    /// # Errors
228    /// [`QueryError::InvalidPlan`] for `page == 0` or `per_page == 0`;
229    /// capability or backend errors.
230    pub async fn paginate(self, page: u64, per_page: u64) -> Result<Page<M>, OrmError> {
231        if page == 0 || per_page == 0 {
232            return Err(QueryError::InvalidPlan("page and per_page start at 1".into()).into());
233        }
234        self.ready()?;
235        self.plan.check(&self.db.capabilities())?;
236        let total = self.clone().count().await?;
237        let items = self
238            .offset((page - 1).saturating_mul(per_page))
239            .limit(per_page)
240            .all()
241            .await?;
242        Ok(Page {
243            items,
244            total,
245            page,
246            per_page,
247        })
248    }
249
250    /// Rows whose primary key is in `ids`, keyed by primary key
251    /// (Django `in_bulk`). Missing keys are absent from the map.
252    ///
253    /// # Errors
254    /// Capability or backend errors.
255    pub async fn in_bulk(
256        self,
257        ids: impl IntoIterator<Item = M::Pk>,
258    ) -> Result<HashMap<M::Pk, M>, OrmError>
259    where
260        M::Pk: Eq + Hash,
261    {
262        let ids: Vec<M::Pk> = ids.into_iter().collect();
263        let column = pk_field::<M>()?.column;
264        let chunk = self.db.capabilities().max_params.max(1);
265        let mut found = HashMap::with_capacity(ids.len());
266        for ids in ids.chunks(chunk) {
267            let predicate = Expr::col(column).is_in(ids.iter().map(DbType::to_value));
268            for model in self.clone().filter(predicate).all().await? {
269                found.insert(model.pk(), model);
270            }
271        }
272        Ok(found)
273    }
274}
siderite_orm/queryset/fetch.rs

siderite_orm/queryset/
fetch.rs
