Skip to main content

icydb_core/db/query/session/
load.rs

1use crate::{
2    db::{
3        DbSession,
4        query::{
5            Query, QueryError,
6            expr::{FilterExpr, SortExpr},
7            plan::{ExecutablePlan, ExplainPlan},
8            predicate::Predicate,
9        },
10        response::{Response, Row},
11    },
12    key::Key,
13    traits::{CanisterKind, EntityKind, UnitKey},
14    view::View,
15};
16use std::collections::HashMap;
17
18///
19/// SessionLoadQuery
20///
21/// Fluent, session-bound load query wrapper that keeps intent pure
22/// while routing execution through the `DbSession` boundary.
23///
24
25pub struct SessionLoadQuery<'a, C: CanisterKind, E: EntityKind<Canister = C>> {
26    session: &'a DbSession<C>,
27    query: Query<E>,
28}
29
30impl<'a, C: CanisterKind, E: EntityKind<Canister = C>> SessionLoadQuery<'a, C, E> {
31    pub(crate) const fn new(session: &'a DbSession<C>, query: Query<E>) -> Self {
32        Self { session, query }
33    }
34
35    // ==================================================================
36    // Intent inspection
37    // ==================================================================
38
39    /// Return a reference to the underlying query intent.
40    #[must_use]
41    pub const fn query(&self) -> &Query<E> {
42        &self.query
43    }
44
45    // ==================================================================
46    // Intent builders (pure, no execution)
47    // ==================================================================
48
49    /// Filter by primary key.
50    #[must_use]
51    pub fn by_key(mut self, key: impl Into<Key>) -> Self {
52        self.query = self.query.by_key(key.into());
53        self
54    }
55
56    /// Load multiple entities by primary key.
57    ///
58    /// Semantics:
59    /// - Equivalent to `WHERE pk IN (…)`
60    /// - Uses key-based access (ByKey / ByKeys)
61    /// - Missing keys are ignored in MissingOk mode
62    /// - Strict mode treats missing rows as corruption
63    /// - Empty input yields an empty result set
64    #[must_use]
65    pub fn many<I>(mut self, keys: I) -> Self
66    where
67        I: IntoIterator<Item = E::PrimaryKey>,
68    {
69        self.query = self.query.by_keys(keys.into_iter().map(Into::into));
70        self
71    }
72
73    /// Add a predicate, implicitly AND-ing with any existing predicate.
74    #[must_use]
75    pub fn filter(mut self, predicate: Predicate) -> Self {
76        self.query = self.query.filter(predicate);
77        self
78    }
79
80    /// Apply a dynamic filter expression.
81    pub fn filter_expr(mut self, expr: FilterExpr) -> Result<Self, QueryError> {
82        self.query = self.query.filter_expr(expr)?;
83        Ok(self)
84    }
85
86    /// Apply a dynamic sort expression.
87    pub fn sort_expr(mut self, expr: SortExpr) -> Result<Self, QueryError> {
88        self.query = self.query.sort_expr(expr)?;
89        Ok(self)
90    }
91
92    /// Append an ascending sort key.
93    #[must_use]
94    pub fn order_by(mut self, field: impl AsRef<str>) -> Self {
95        self.query = self.query.order_by(field);
96        self
97    }
98
99    /// Append a descending sort key.
100    #[must_use]
101    pub fn order_by_desc(mut self, field: impl AsRef<str>) -> Self {
102        self.query = self.query.order_by_desc(field);
103        self
104    }
105
106    /// Apply a load limit.
107    #[must_use]
108    pub fn limit(mut self, limit: u32) -> Self {
109        self.query = self.query.limit(limit);
110        self
111    }
112
113    /// Apply a load offset.
114    #[must_use]
115    pub fn offset(mut self, offset: u32) -> Self {
116        self.query = self.query.offset(offset);
117        self
118    }
119
120    // ==================================================================
121    // Planning / diagnostics (no execution)
122    // ==================================================================
123
124    /// Explain this query without executing it.
125    pub fn explain(&self) -> Result<ExplainPlan, QueryError> {
126        self.query.explain()
127    }
128
129    /// Plan this query into an executor-ready plan.
130    pub fn plan(&self) -> Result<ExecutablePlan<E>, QueryError> {
131        self.query.plan()
132    }
133
134    // ==================================================================
135    // Execution boundary (single entry point)
136    // ==================================================================
137
138    /// Execute this query using the session's policy settings.
139    pub fn execute(&self) -> Result<Response<E>, QueryError> {
140        self.session.execute_query(self.query())
141    }
142
143    // ==================================================================
144    // Execution terminals — cardinality / existence
145    // ==================================================================
146
147    /// Return whether any rows match this query.
148    pub fn exists(&self) -> Result<bool, QueryError> {
149        Ok(self.count()? > 0)
150    }
151
152    /// Execute and return whether the response is empty.
153    pub fn is_empty(&self) -> Result<bool, QueryError> {
154        Ok(self.execute()?.is_empty())
155    }
156
157    /// Execute and return the number of matching rows.
158    pub fn count(&self) -> Result<u32, QueryError> {
159        Ok(self.execute()?.count())
160    }
161
162    /// Execute and require exactly one row.
163    pub fn require_one(&self) -> Result<(), QueryError> {
164        self.execute()?.require_one().map_err(QueryError::Response)
165    }
166
167    /// Execute and require at least one row.
168    pub fn require_some(&self) -> Result<(), QueryError> {
169        self.execute()?.require_some().map_err(QueryError::Response)
170    }
171
172    // ==================================================================
173    // Execution terminals — rows
174    // ==================================================================
175
176    pub fn row(&self) -> Result<Row<E>, QueryError> {
177        self.execute()?.row().map_err(QueryError::Response)
178    }
179
180    pub fn try_row(&self) -> Result<Option<Row<E>>, QueryError> {
181        self.execute()?.try_row().map_err(QueryError::Response)
182    }
183
184    pub fn rows(&self) -> Result<Vec<Row<E>>, QueryError> {
185        Ok(self.execute()?.rows())
186    }
187
188    // ==================================================================
189    // Execution terminals — entities
190    // ==================================================================
191
192    pub fn entity(&self) -> Result<E, QueryError> {
193        self.execute()?.entity().map_err(QueryError::Response)
194    }
195
196    pub fn try_entity(&self) -> Result<Option<E>, QueryError> {
197        self.execute()?.try_entity().map_err(QueryError::Response)
198    }
199
200    pub fn entities(&self) -> Result<Vec<E>, QueryError> {
201        Ok(self.execute()?.entities())
202    }
203
204    /// Execute and count entities grouped by the provided key selector.
205    pub fn group_count_by<K>(&self, key: impl Fn(&E) -> K) -> Result<HashMap<K, u32>, QueryError>
206    where
207        K: Eq + std::hash::Hash,
208    {
209        // Phase: materialize entities.
210        let entities = self.execute()?.entities();
211
212        // Phase: count by derived key.
213        let mut counts = HashMap::new();
214        for entity in entities {
215            *counts.entry(key(&entity)).or_insert(0) += 1;
216        }
217
218        Ok(counts)
219    }
220
221    /// Alias for `entity`.
222    pub fn one(&self) -> Result<E, QueryError> {
223        self.entity()
224    }
225
226    /// Alias for `try_entity`.
227    pub fn one_opt(&self) -> Result<Option<E>, QueryError> {
228        self.try_entity()
229    }
230
231    /// Alias for `entities`.
232    pub fn all(&self) -> Result<Vec<E>, QueryError> {
233        self.entities()
234    }
235
236    // ==================================================================
237    // Execution terminals — store keys
238    // ==================================================================
239
240    pub fn key(&self) -> Result<Option<Key>, QueryError> {
241        Ok(self.execute()?.key())
242    }
243
244    pub fn key_strict(&self) -> Result<Key, QueryError> {
245        self.execute()?.key_strict().map_err(QueryError::Response)
246    }
247
248    pub fn try_key(&self) -> Result<Option<Key>, QueryError> {
249        self.execute()?.try_key().map_err(QueryError::Response)
250    }
251
252    pub fn keys(&self) -> Result<Vec<Key>, QueryError> {
253        Ok(self.execute()?.keys())
254    }
255
256    pub fn contains_key(&self, key: &Key) -> Result<bool, QueryError> {
257        Ok(self.execute()?.contains_key(key))
258    }
259
260    // ==================================================================
261    // Execution terminals — primary keys
262    // ==================================================================
263
264    pub fn primary_key(&self) -> Result<E::PrimaryKey, QueryError> {
265        self.execute()?.primary_key().map_err(QueryError::Response)
266    }
267
268    pub fn try_primary_key(&self) -> Result<Option<E::PrimaryKey>, QueryError> {
269        self.execute()?
270            .try_primary_key()
271            .map_err(QueryError::Response)
272    }
273
274    pub fn primary_keys(&self) -> Result<Vec<E::PrimaryKey>, QueryError> {
275        Ok(self.execute()?.primary_keys())
276    }
277
278    // ==================================================================
279    // Execution terminals — views
280    // ==================================================================
281
282    pub fn views(&self) -> Result<Vec<View<E>>, QueryError> {
283        Ok(self.execute()?.views())
284    }
285
286    pub fn view(&self) -> Result<View<E>, QueryError> {
287        self.execute()?.view().map_err(QueryError::Response)
288    }
289
290    pub fn view_opt(&self) -> Result<Option<View<E>>, QueryError> {
291        self.execute()?.view_opt().map_err(QueryError::Response)
292    }
293}
294
295impl<C: CanisterKind, E: EntityKind<Canister = C>> SessionLoadQuery<'_, C, E>
296where
297    E::PrimaryKey: UnitKey,
298{
299    /// Load the singleton entity identified by the unit primary key `()`.
300    ///
301    /// Semantics:
302    /// - Equivalent to `WHERE pk = ()`
303    /// - Uses key-based access (ByKey)
304    /// - Does not allow predicates
305    /// - MissingOk mode returns empty
306    /// - Strict mode treats missing row as corruption
307    #[must_use]
308    pub fn only(mut self) -> Self {
309        self.query = self.query.only();
310        self
311    }
312}