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