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