icydb_core/db/query/intent/
mod.rs1use crate::{
2 db::query::{
3 ReadConsistency,
4 plan::{
5 AccessPath, AccessPlan, DeleteLimitSpec, ExecutablePlan, ExplainPlan, LogicalPlan,
6 OrderDirection, OrderSpec, PageSpec, PlanError, ProjectionSpec,
7 planner::{PlannerError, plan_access},
8 validate::validate_access_plan,
9 validate::validate_order,
10 },
11 predicate::{Predicate, SchemaInfo, ValidateError, normalize, validate},
12 },
13 error::InternalError,
14 key::Key,
15 traits::EntityKind,
16};
17use std::marker::PhantomData;
18use thiserror::Error as ThisError;
19
20#[derive(Clone, Copy, Debug, Eq, PartialEq)]
28pub enum QueryMode {
29 Load(LoadSpec),
30 Delete(DeleteSpec),
31}
32
33impl QueryMode {
34 #[must_use]
36 pub const fn is_load(&self) -> bool {
37 match self {
38 Self::Load(_) => true,
39 Self::Delete(_) => false,
40 }
41 }
42
43 #[must_use]
45 pub const fn is_delete(&self) -> bool {
46 match self {
47 Self::Delete(_) => true,
48 Self::Load(_) => false,
49 }
50 }
51}
52
53#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
60pub struct LoadSpec {
61 pub limit: Option<u32>,
62 pub offset: u64,
63}
64
65impl LoadSpec {
66 #[must_use]
68 pub const fn new() -> Self {
69 Self {
70 limit: None,
71 offset: 0,
72 }
73 }
74}
75
76#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
83pub struct DeleteSpec {
84 pub limit: Option<u32>,
85}
86
87impl DeleteSpec {
88 #[must_use]
90 pub const fn new() -> Self {
91 Self { limit: None }
92 }
93}
94
95#[derive(Debug)]
107pub struct Query<E: EntityKind> {
108 mode: QueryMode,
109 predicate: Option<Predicate>,
110 key_access: Option<KeyAccessState>,
111 key_access_conflict: bool,
112 order: Option<OrderSpec>,
113 projection: ProjectionSpec,
114 consistency: ReadConsistency,
115 _marker: PhantomData<E>,
116}
117
118impl<E: EntityKind> Query<E> {
119 #[must_use]
123 pub const fn new(consistency: ReadConsistency) -> Self {
124 Self {
125 mode: QueryMode::Load(LoadSpec::new()),
126 predicate: None,
127 key_access: None,
128 key_access_conflict: false,
129 order: None,
130 projection: ProjectionSpec::All,
131 consistency,
132 _marker: PhantomData,
133 }
134 }
135
136 #[must_use]
138 pub const fn mode(&self) -> QueryMode {
139 self.mode
140 }
141
142 #[must_use]
144 pub fn filter(mut self, predicate: Predicate) -> Self {
145 self.predicate = match self.predicate.take() {
146 Some(existing) => Some(Predicate::And(vec![existing, predicate])),
147 None => Some(predicate),
148 };
149 self
150 }
151
152 #[must_use]
154 pub fn order_by(mut self, field: impl AsRef<str>) -> Self {
155 self.order = Some(push_order(self.order, field.as_ref(), OrderDirection::Asc));
156 self
157 }
158
159 #[must_use]
161 pub fn order_by_desc(mut self, field: impl AsRef<str>) -> Self {
162 self.order = Some(push_order(self.order, field.as_ref(), OrderDirection::Desc));
163 self
164 }
165
166 fn set_key_access(mut self, kind: KeyAccessKind, access: KeyAccess) -> Self {
168 if let Some(existing) = &self.key_access
169 && existing.kind != kind
170 {
171 self.key_access_conflict = true;
172 }
173
174 self.key_access = Some(KeyAccessState { kind, access });
175
176 self
177 }
178
179 pub(crate) fn by_key(self, key: Key) -> Self {
181 self.set_key_access(KeyAccessKind::Single, KeyAccess::Single(key))
182 }
183
184 pub(crate) fn by_keys<I>(self, keys: I) -> Self
186 where
187 I: IntoIterator<Item = Key>,
188 {
189 self.set_key_access(
190 KeyAccessKind::Many,
191 KeyAccess::Many(keys.into_iter().collect()),
192 )
193 }
194
195 #[must_use]
197 pub const fn delete(mut self) -> Self {
198 if self.mode.is_load() {
199 self.mode = QueryMode::Delete(DeleteSpec::new());
200 }
201 self
202 }
203
204 #[must_use]
208 pub const fn limit(mut self, limit: u32) -> Self {
209 match self.mode {
210 QueryMode::Load(mut spec) => {
211 spec.limit = Some(limit);
212 self.mode = QueryMode::Load(spec);
213 }
214 QueryMode::Delete(mut spec) => {
215 spec.limit = Some(limit);
216 self.mode = QueryMode::Delete(spec);
217 }
218 }
219 self
220 }
221
222 #[must_use]
224 pub const fn offset(mut self, offset: u64) -> Self {
225 if let QueryMode::Load(mut spec) = self.mode {
226 spec.offset = offset;
227 self.mode = QueryMode::Load(spec);
228 }
229 self
230 }
231
232 pub fn explain(&self) -> Result<ExplainPlan, QueryError> {
234 let plan = self.build_plan::<E>()?;
235
236 Ok(plan.explain())
237 }
238
239 pub fn plan(&self) -> Result<ExecutablePlan<E>, QueryError> {
241 let plan = self.build_plan::<E>()?;
242
243 Ok(ExecutablePlan::new(plan))
244 }
245
246 fn build_plan<T: EntityKind>(&self) -> Result<LogicalPlan, QueryError> {
248 let model = T::MODEL;
250 let schema_info = SchemaInfo::from_entity_model(model)?;
251 self.validate_intent()?;
252
253 if let Some(order) = &self.order {
254 validate_order(&schema_info, order)?;
255 }
256
257 let normalized_predicate = self.predicate.as_ref().map(normalize);
259 let access_plan = match &self.key_access {
260 Some(state) => {
261 if let Some(predicate) = self.predicate.as_ref() {
262 validate(&schema_info, predicate)?;
263 }
264 access_plan_from_keys(&state.access)
265 }
266 None => plan_access::<T>(&schema_info, normalized_predicate.as_ref())?,
267 };
268
269 validate_access_plan(&schema_info, model, &access_plan)?;
270
271 let plan = LogicalPlan {
273 mode: self.mode,
274 access: access_plan,
275 predicate: normalized_predicate,
276 order: self.order.clone(),
277 delete_limit: match self.mode {
278 QueryMode::Delete(spec) => spec.limit.map(|max_rows| DeleteLimitSpec { max_rows }),
279 QueryMode::Load(_) => None,
280 },
281 page: match self.mode {
282 QueryMode::Load(spec) => {
283 if spec.limit.is_some() || spec.offset > 0 {
284 Some(PageSpec {
285 limit: spec.limit,
286 offset: spec.offset,
287 })
288 } else {
289 None
290 }
291 }
292 QueryMode::Delete(_) => None,
293 },
294 projection: self.projection.clone(),
295 consistency: self.consistency,
296 };
297
298 Ok(plan)
299 }
300
301 const fn validate_intent(&self) -> Result<(), IntentError> {
303 if self.key_access_conflict {
304 return Err(IntentError::KeyAccessConflict);
305 }
306
307 if let Some(state) = &self.key_access {
308 match state.kind {
309 KeyAccessKind::Many if self.predicate.is_some() => {
310 return Err(IntentError::ManyWithPredicate);
311 }
312 KeyAccessKind::Only if self.predicate.is_some() => {
313 return Err(IntentError::OnlyWithPredicate);
314 }
315 _ => {}
316 }
317 }
318
319 match self.mode {
320 QueryMode::Load(_) => {}
321 QueryMode::Delete(spec) => {
322 if spec.limit.is_some() && self.order.is_none() {
323 return Err(IntentError::DeleteLimitRequiresOrder);
324 }
325 }
326 }
327
328 Ok(())
329 }
330}
331
332impl<E: EntityKind<PrimaryKey = ()>> Query<E> {
333 pub(crate) fn only(self) -> Self {
335 self.set_key_access(KeyAccessKind::Only, KeyAccess::Single(Key::Unit))
336 }
337}
338
339#[derive(Debug, ThisError)]
344pub enum QueryError {
345 #[error("{0}")]
346 Validate(#[from] ValidateError),
347
348 #[error("{0}")]
349 Plan(#[from] PlanError),
350
351 #[error("{0}")]
352 Intent(#[from] IntentError),
353
354 #[error("{0}")]
355 Execute(#[from] InternalError),
356}
357
358impl From<PlannerError> for QueryError {
359 fn from(err: PlannerError) -> Self {
360 match err {
361 PlannerError::Plan(err) => Self::Plan(err),
362 PlannerError::Internal(err) => Self::Execute(err),
363 }
364 }
365}
366
367#[derive(Clone, Copy, Debug, ThisError)]
372pub enum IntentError {
373 #[error("delete limit requires an explicit ordering")]
374 DeleteLimitRequiresOrder,
375
376 #[error("many() cannot be combined with predicates")]
377 ManyWithPredicate,
378
379 #[error("only() cannot be combined with predicates")]
380 OnlyWithPredicate,
381
382 #[error("multiple key access methods were used on the same query")]
383 KeyAccessConflict,
384}
385
386#[derive(Clone, Debug, Eq, PartialEq)]
388enum KeyAccess {
389 Single(Key),
390 Many(Vec<Key>),
391}
392
393#[derive(Clone, Copy, Debug, Eq, PartialEq)]
395enum KeyAccessKind {
396 Single,
397 Many,
398 Only,
399}
400
401#[derive(Clone, Debug, Eq, PartialEq)]
403struct KeyAccessState {
404 kind: KeyAccessKind,
405 access: KeyAccess,
406}
407
408fn access_plan_from_keys(access: &KeyAccess) -> AccessPlan {
410 match access {
411 KeyAccess::Single(key) => AccessPlan::Path(AccessPath::ByKey(*key)),
412 KeyAccess::Many(keys) => {
413 if let Some((first, rest)) = keys.split_first()
414 && rest.is_empty()
415 {
416 return AccessPlan::Path(AccessPath::ByKey(*first));
417 }
418
419 AccessPlan::Path(AccessPath::ByKeys(keys.clone()))
420 }
421 }
422}
423
424fn push_order(order: Option<OrderSpec>, field: &str, direction: OrderDirection) -> OrderSpec {
426 match order {
427 Some(mut spec) => {
428 spec.fields.push((field.to_string(), direction));
429 spec
430 }
431 None => OrderSpec {
432 fields: vec![(field.to_string(), direction)],
433 },
434 }
435}
436
437#[cfg(test)]
438mod tests;