icydb_core/db/query/intent/

mod.rs

#![expect(clippy::used_underscore_binding)]
#[cfg(test)]
mod tests;

// Key-only access intent and helpers.
pub type DeleteSpec = crate::db::query::plan::DeleteSpec;
pub type LoadSpec = crate::db::query::plan::LoadSpec;
pub type QueryMode = crate::db::query::plan::QueryMode;

use crate::{
    db::{
        access::{AccessPath, AccessPlan, AccessPlanError},
        contracts::{Predicate, ReadConsistency, SchemaInfo, ValidateError},
        policy,
        query::{
            explain::ExplainPlan,
            expr::{FilterExpr, SortExpr, SortLowerError},
            plan::validate::validate_query_semantics,
            plan::{
                AccessPlannedQuery, DeleteLimitSpec, LogicalPlan, OrderDirection, OrderSpec,
                PageSpec, PlanError, PlannerError, canonical, plan_access,
            },
            predicate::{
                lower_to_execution_model, normalize, normalize_enum_literals,
                validate::reject_unsupported_query_features,
            },
        },
        response::ResponseError,
    },
    error::InternalError,
    model::entity::EntityModel,
    traits::{EntityKind, FieldValue, SingletonEntity},
    value::Value,
};
use std::marker::PhantomData;
use thiserror::Error as ThisError;

///
/// KeyAccess
/// Primary-key-only access hints for query planning.
///

#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) enum KeyAccess<K> {
    Single(K),
    Many(Vec<K>),
}

///
/// KeyAccessKind
/// Identifies which key-only builder set the access path.
///

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum KeyAccessKind {
    Single,
    Many,
    Only,
}

///
/// KeyAccessState
/// Tracks key-only access plus its origin for intent validation.
///

#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct KeyAccessState<K> {
    pub kind: KeyAccessKind,
    pub access: KeyAccess<K>,
}

// Build a model-level access plan for key-only intents.
pub(crate) fn access_plan_from_keys_value<K>(access: &KeyAccess<K>) -> AccessPlan<Value>
where
    K: FieldValue,
{
    match access {
        KeyAccess::Single(key) => AccessPlan::path(AccessPath::ByKey(key.to_value())),
        KeyAccess::Many(keys) => {
            let mut values: Vec<Value> = keys.iter().map(FieldValue::to_value).collect();
            canonical::canonicalize_key_values(&mut values);
            if let Some(first) = values.first()
                && values.len() == 1
            {
                return AccessPlan::path(AccessPath::ByKey(first.clone()));
            }

            AccessPlan::path(AccessPath::ByKeys(values))
        }
    }
}

// Convert model-level access plans into entity-keyed access plans.
pub(crate) fn access_plan_to_entity_keys<E: EntityKind>(
    model: &EntityModel,
    access: AccessPlan<Value>,
) -> Result<AccessPlan<E::Key>, PlanError> {
    access.into_executable::<E>(model)
}

// Convert model-level key values into typed entity keys.
pub(crate) fn coerce_entity_key<E: EntityKind>(
    model: &EntityModel,
    key: &Value,
) -> Result<E::Key, PlanError> {
    E::Key::from_value(key).ok_or_else(|| {
        PlanError::from(AccessPlanError::PrimaryKeyMismatch {
            field: model.primary_key.name.to_string(),
            key: key.clone(),
        })
    })
}

impl AccessPlan<Value> {
    /// Convert model-level access plans into typed executable access plans.
    pub(crate) fn into_executable<E: EntityKind>(
        self,
        model: &EntityModel,
    ) -> Result<AccessPlan<E::Key>, PlanError> {
        match self {
            Self::Path(path) => Ok(AccessPlan::path(path.into_executable::<E>(model)?)),
            Self::Union(children) => {
                let mut out = Vec::with_capacity(children.len());
                for child in children {
                    out.push(child.into_executable::<E>(model)?);
                }

                Ok(AccessPlan::Union(out))
            }
            Self::Intersection(children) => {
                let mut out = Vec::with_capacity(children.len());
                for child in children {
                    out.push(child.into_executable::<E>(model)?);
                }

                Ok(AccessPlan::Intersection(out))
            }
        }
    }
}

impl AccessPath<Value> {
    /// Convert one model-level access path into a typed executable access path.
    pub(crate) fn into_executable<E: EntityKind>(
        self,
        model: &EntityModel,
    ) -> Result<AccessPath<E::Key>, PlanError> {
        match self {
            Self::ByKey(key) => Ok(AccessPath::ByKey(coerce_entity_key::<E>(model, &key)?)),
            Self::ByKeys(keys) => {
                let mut out = Vec::with_capacity(keys.len());
                for key in keys {
                    out.push(coerce_entity_key::<E>(model, &key)?);
                }

                Ok(AccessPath::ByKeys(out))
            }
            Self::KeyRange { start, end } => Ok(AccessPath::KeyRange {
                start: coerce_entity_key::<E>(model, &start)?,
                end: coerce_entity_key::<E>(model, &end)?,
            }),
            Self::IndexPrefix { index, values } => Ok(AccessPath::IndexPrefix { index, values }),
            Self::IndexRange { spec } => Ok(AccessPath::IndexRange { spec }),
            Self::FullScan => Ok(AccessPath::FullScan),
        }
    }
}

///
/// QueryModel
///
/// Model-level query intent and planning context.
/// Consumes an `EntityModel` derived from typed entity definitions.
///

#[derive(Debug)]
pub(crate) struct QueryModel<'m, K> {
    model: &'m EntityModel,
    mode: QueryMode,
    predicate: Option<Predicate>,
    key_access: Option<KeyAccessState<K>>,
    key_access_conflict: bool,
    order: Option<OrderSpec>,
    distinct: bool,
    consistency: ReadConsistency,
}

impl<'m, K: FieldValue> QueryModel<'m, K> {
    #[must_use]
    pub(crate) const fn new(model: &'m EntityModel, consistency: ReadConsistency) -> Self {
        Self {
            model,
            mode: QueryMode::Load(LoadSpec::new()),
            predicate: None,
            key_access: None,
            key_access_conflict: false,
            order: None,
            distinct: false,
            consistency,
        }
    }

    /// Return the intent mode (load vs delete).
    #[must_use]
    pub(crate) const fn mode(&self) -> QueryMode {
        self.mode
    }

    #[must_use]
    fn has_explicit_order(&self) -> bool {
        policy::has_explicit_order(self.order.as_ref())
    }

    #[must_use]
    const fn load_spec(&self) -> Option<LoadSpec> {
        match self.mode {
            QueryMode::Load(spec) => Some(spec),
            QueryMode::Delete(_) => None,
        }
    }

    /// Add a predicate, implicitly AND-ing with any existing predicate.
    #[must_use]
    pub(crate) fn filter(mut self, predicate: Predicate) -> Self {
        self.predicate = match self.predicate.take() {
            Some(existing) => Some(Predicate::And(vec![existing, predicate])),
            None => Some(predicate),
        };
        self
    }

    /// Apply a dynamic filter expression using the model schema.
    pub(crate) fn filter_expr(self, expr: FilterExpr) -> Result<Self, QueryError> {
        let schema = SchemaInfo::from_entity_model(self.model)?;
        let predicate = expr.lower_with(&schema).map_err(QueryError::Validate)?;

        Ok(self.filter(predicate))
    }

    /// Apply a dynamic sort expression using the model schema.
    pub(crate) fn sort_expr(self, expr: SortExpr) -> Result<Self, QueryError> {
        let schema = SchemaInfo::from_entity_model(self.model)?;
        let order = match expr.lower_with(&schema) {
            Ok(order) => order,
            Err(SortLowerError::Validate(err)) => return Err(QueryError::Validate(err)),
            Err(SortLowerError::Plan(err)) => return Err(QueryError::from(*err)),
        };

        policy::validate_order_shape(Some(&order))
            .map_err(IntentError::from)
            .map_err(QueryError::from)?;

        Ok(self.order_spec(order))
    }

    /// Append an ascending sort key.
    #[must_use]
    pub(crate) fn order_by(mut self, field: impl AsRef<str>) -> Self {
        self.order = Some(push_order(self.order, field.as_ref(), OrderDirection::Asc));
        self
    }

    /// Append a descending sort key.
    #[must_use]
    pub(crate) fn order_by_desc(mut self, field: impl AsRef<str>) -> Self {
        self.order = Some(push_order(self.order, field.as_ref(), OrderDirection::Desc));
        self
    }

    /// Set a fully-specified order spec (validated before reaching this boundary).
    pub(crate) fn order_spec(mut self, order: OrderSpec) -> Self {
        self.order = Some(order);
        self
    }

    /// Enable DISTINCT semantics for this query intent.
    #[must_use]
    pub(crate) const fn distinct(mut self) -> Self {
        self.distinct = true;
        self
    }

    /// Track key-only access paths and detect conflicting key intents.
    fn set_key_access(mut self, kind: KeyAccessKind, access: KeyAccess<K>) -> Self {
        if let Some(existing) = &self.key_access
            && existing.kind != kind
        {
            self.key_access_conflict = true;
        }

        self.key_access = Some(KeyAccessState { kind, access });

        self
    }

    /// Set the access path to a single primary key lookup.
    pub(crate) fn by_id(self, id: K) -> Self {
        self.set_key_access(KeyAccessKind::Single, KeyAccess::Single(id))
    }

    /// Set the access path to a primary key batch lookup.
    pub(crate) fn by_ids<I>(self, ids: I) -> Self
    where
        I: IntoIterator<Item = K>,
    {
        self.set_key_access(
            KeyAccessKind::Many,
            KeyAccess::Many(ids.into_iter().collect()),
        )
    }

    /// Set the access path to the singleton primary key.
    pub(crate) fn only(self, id: K) -> Self {
        self.set_key_access(KeyAccessKind::Only, KeyAccess::Single(id))
    }

    /// Mark this intent as a delete query.
    #[must_use]
    pub(crate) const fn delete(mut self) -> Self {
        if self.mode.is_load() {
            self.mode = QueryMode::Delete(DeleteSpec::new());
        }
        self
    }

    /// Apply a limit to the current mode.
    ///
    /// Load limits bound result size; delete limits bound mutation size.
    #[must_use]
    pub(crate) const fn limit(mut self, limit: u32) -> Self {
        match self.mode {
            QueryMode::Load(mut spec) => {
                spec.limit = Some(limit);
                self.mode = QueryMode::Load(spec);
            }
            QueryMode::Delete(mut spec) => {
                spec.limit = Some(limit);
                self.mode = QueryMode::Delete(spec);
            }
        }
        self
    }

    /// Apply an offset to a load intent.
    #[must_use]
    pub(crate) const fn offset(mut self, offset: u32) -> Self {
        if let QueryMode::Load(mut spec) = self.mode {
            spec.offset = offset;
            self.mode = QueryMode::Load(spec);
        }
        self
    }

    /// Build a model-level logical plan using Value-based access keys.
    fn build_plan_model(&self) -> Result<AccessPlannedQuery<Value>, QueryError> {
        // Phase 1: schema surface and intent validation.
        let schema_info = SchemaInfo::from_entity_model(self.model)?;
        self.validate_intent()?;

        // Phase 2: predicate normalization and access planning.
        let normalized_predicate = self
            .predicate
            .as_ref()
            .map(|predicate| {
                reject_unsupported_query_features(predicate).map_err(ValidateError::from)?;
                let predicate = normalize_enum_literals(&schema_info, predicate)?;
                Ok::<Predicate, ValidateError>(normalize(&predicate))
            })
            .transpose()?;
        let access_plan_value = match &self.key_access {
            Some(state) => access_plan_from_keys_value(&state.access),
            None => plan_access(self.model, &schema_info, normalized_predicate.as_ref())?,
        };

        // Phase 3: assemble the executor-ready plan.
        let logical = LogicalPlan {
            mode: self.mode,
            predicate: normalized_predicate.map(lower_to_execution_model),
            // Canonicalize ORDER BY to include an explicit primary-key tie-break.
            // This ensures explain/fingerprint/execution share one deterministic order shape.
            order: canonicalize_order_spec(self.model, self.order.clone()),
            distinct: self.distinct,
            delete_limit: match self.mode {
                QueryMode::Delete(spec) => spec.limit.map(|max_rows| DeleteLimitSpec { max_rows }),
                QueryMode::Load(_) => None,
            },
            page: match self.mode {
                QueryMode::Load(spec) => {
                    if spec.limit.is_some() || spec.offset > 0 {
                        Some(PageSpec {
                            limit: spec.limit,
                            offset: spec.offset,
                        })
                    } else {
                        None
                    }
                }
                QueryMode::Delete(_) => None,
            },
            consistency: self.consistency,
        };
        let plan = AccessPlannedQuery::from_parts(logical, access_plan_value);

        validate_query_semantics(&schema_info, self.model, &plan)?;

        Ok(plan)
    }

    // Validate pre-plan policy invariants and key-access rules before planning.
    fn validate_intent(&self) -> Result<(), IntentError> {
        if self.key_access_conflict {
            return Err(IntentError::KeyAccessConflict);
        }

        policy::validate_intent_plan_shape(self.mode, self.order.as_ref())
            .map_err(IntentError::from)?;

        if let Some(state) = &self.key_access {
            match state.kind {
                KeyAccessKind::Many if self.predicate.is_some() => {
                    return Err(IntentError::ByIdsWithPredicate);
                }
                KeyAccessKind::Only if self.predicate.is_some() => {
                    return Err(IntentError::OnlyWithPredicate);
                }
                _ => {
                    // NOTE: Single/Many without predicates impose no additional constraints.
                }
            }
        }

        Ok(())
    }
}

///
/// Query
///
/// Typed, declarative query intent for a specific entity type.
///
/// This intent is:
/// - schema-agnostic at construction
/// - normalized and validated only during planning
/// - free of access-path decisions
///

///
/// PlannedQuery
///
/// Neutral query-owned planned contract produced by query planning.
/// Stores logical + access shape without executor compilation state.
///
#[derive(Debug)]
pub struct PlannedQuery<E: EntityKind> {
    plan: AccessPlannedQuery<E::Key>,
    _marker: PhantomData<E>,
}

impl<E: EntityKind> PlannedQuery<E> {
    #[must_use]
    pub(in crate::db) const fn new(plan: AccessPlannedQuery<E::Key>) -> Self {
        Self {
            plan,
            _marker: PhantomData,
        }
    }

    #[must_use]
    pub fn explain(&self) -> ExplainPlan {
        self.plan.explain_with_model(E::MODEL)
    }

    #[must_use]
    pub(in crate::db) fn into_inner(self) -> AccessPlannedQuery<E::Key> {
        self.plan
    }
}

#[derive(Debug)]
pub struct Query<E: EntityKind> {
    intent: QueryModel<'static, E::Key>,
    _marker: PhantomData<E>,
}

impl<E: EntityKind> Query<E> {
    /// Create a new intent with an explicit missing-row policy.
    /// MissingOk favors idempotency and may mask index/data divergence on deletes.
    /// Use Strict to surface missing rows during scan/delete execution.
    #[must_use]
    pub const fn new(consistency: ReadConsistency) -> Self {
        Self {
            intent: QueryModel::new(E::MODEL, consistency),
            _marker: PhantomData,
        }
    }

    /// Return the intent mode (load vs delete).
    #[must_use]
    pub const fn mode(&self) -> QueryMode {
        self.intent.mode()
    }

    #[must_use]
    pub(crate) fn has_explicit_order(&self) -> bool {
        self.intent.has_explicit_order()
    }

    #[must_use]
    pub(crate) const fn load_spec(&self) -> Option<LoadSpec> {
        self.intent.load_spec()
    }

    /// Add a predicate, implicitly AND-ing with any existing predicate.
    #[must_use]
    pub fn filter(mut self, predicate: Predicate) -> Self {
        self.intent = self.intent.filter(predicate);
        self
    }

    /// Apply a dynamic filter expression.
    pub fn filter_expr(self, expr: FilterExpr) -> Result<Self, QueryError> {
        let Self { intent, _marker } = self;
        let intent = intent.filter_expr(expr)?;

        Ok(Self { intent, _marker })
    }

    /// Apply a dynamic sort expression.
    pub fn sort_expr(self, expr: SortExpr) -> Result<Self, QueryError> {
        let Self { intent, _marker } = self;
        let intent = intent.sort_expr(expr)?;

        Ok(Self { intent, _marker })
    }

    /// Append an ascending sort key.
    #[must_use]
    pub fn order_by(mut self, field: impl AsRef<str>) -> Self {
        self.intent = self.intent.order_by(field);
        self
    }

    /// Append a descending sort key.
    #[must_use]
    pub fn order_by_desc(mut self, field: impl AsRef<str>) -> Self {
        self.intent = self.intent.order_by_desc(field);
        self
    }

    /// Enable DISTINCT semantics for this query.
    #[must_use]
    pub fn distinct(mut self) -> Self {
        self.intent = self.intent.distinct();
        self
    }

    /// Set the access path to a single primary key lookup.
    pub(crate) fn by_id(self, id: E::Key) -> Self {
        let Self { intent, _marker } = self;
        Self {
            intent: intent.by_id(id),
            _marker,
        }
    }

    /// Set the access path to a primary key batch lookup.
    pub(crate) fn by_ids<I>(self, ids: I) -> Self
    where
        I: IntoIterator<Item = E::Key>,
    {
        let Self { intent, _marker } = self;
        Self {
            intent: intent.by_ids(ids),
            _marker,
        }
    }

    /// Mark this intent as a delete query.
    #[must_use]
    pub fn delete(mut self) -> Self {
        self.intent = self.intent.delete();
        self
    }

    /// Apply a limit to the current mode.
    ///
    /// Load limits bound result size; delete limits bound mutation size.
    /// For load queries, any use of `limit` or `offset` requires an explicit
    /// `order_by(...)` so pagination is deterministic.
    #[must_use]
    pub fn limit(mut self, limit: u32) -> Self {
        self.intent = self.intent.limit(limit);
        self
    }

    /// Apply an offset to a load intent.
    ///
    /// Any use of `offset` or `limit` requires an explicit `order_by(...)`.
    #[must_use]
    pub fn offset(mut self, offset: u32) -> Self {
        self.intent = self.intent.offset(offset);
        self
    }

    /// Explain this intent without executing it.
    pub fn explain(&self) -> Result<ExplainPlan, QueryError> {
        let plan = self.planned()?;

        Ok(plan.explain())
    }

    /// Plan this intent into a neutral planned query contract.
    pub fn planned(&self) -> Result<PlannedQuery<E>, QueryError> {
        let plan = self.build_plan()?;

        Ok(PlannedQuery::new(plan))
    }

    // Build a logical plan for the current intent.
    fn build_plan(&self) -> Result<AccessPlannedQuery<E::Key>, QueryError> {
        let plan_value = self.intent.build_plan_model()?;
        let (logical, access) = plan_value.into_parts();
        let access = access_plan_to_entity_keys::<E>(E::MODEL, access)?;
        let plan = AccessPlannedQuery::from_parts(logical, access);

        Ok(plan)
    }
}

impl<E> Query<E>
where
    E: EntityKind + SingletonEntity,
    E::Key: Default,
{
    /// Set the access path to the singleton primary key.
    pub(crate) fn only(self) -> Self {
        let Self { intent, _marker } = self;

        Self {
            intent: intent.only(E::Key::default()),
            _marker,
        }
    }
}

///
/// QueryError
///

#[derive(Debug, ThisError)]
pub enum QueryError {
    #[error("{0}")]
    Validate(#[from] ValidateError),

    #[error("{0}")]
    Plan(Box<PlanError>),

    #[error("{0}")]
    Intent(#[from] IntentError),

    #[error("{0}")]
    Response(#[from] ResponseError),

    #[error("{0}")]
    Execute(#[from] InternalError),
}

impl From<PlannerError> for QueryError {
    fn from(err: PlannerError) -> Self {
        match err {
            PlannerError::Plan(err) => Self::from(*err),
            PlannerError::Internal(err) => Self::Execute(*err),
        }
    }
}

impl From<PlanError> for QueryError {
    fn from(err: PlanError) -> Self {
        Self::Plan(Box::new(err))
    }
}

///
/// IntentError
///

#[derive(Clone, Copy, Debug, ThisError)]
pub enum IntentError {
    #[error("{0}")]
    PlanShape(#[from] policy::PlanPolicyError),

    #[error("by_ids() cannot be combined with predicates")]
    ByIdsWithPredicate,

    #[error("only() cannot be combined with predicates")]
    OnlyWithPredicate,

    #[error("multiple key access methods were used on the same query")]
    KeyAccessConflict,

    #[error("cursor pagination requires an explicit ordering")]
    CursorRequiresOrder,

    #[error("cursor pagination requires an explicit limit")]
    CursorRequiresLimit,

    #[error("cursor tokens can only be used with .page().execute()")]
    CursorRequiresPagedExecution,
}

impl From<policy::CursorPagingPolicyError> for IntentError {
    fn from(err: policy::CursorPagingPolicyError) -> Self {
        match err {
            policy::CursorPagingPolicyError::CursorRequiresOrder => Self::CursorRequiresOrder,
            policy::CursorPagingPolicyError::CursorRequiresLimit => Self::CursorRequiresLimit,
        }
    }
}

/// Helper to append an ordering field while preserving existing order spec.
fn push_order(order: Option<OrderSpec>, field: &str, direction: OrderDirection) -> OrderSpec {
    match order {
        Some(mut spec) => {
            spec.fields.push((field.to_string(), direction));
            spec
        }
        None => OrderSpec {
            fields: vec![(field.to_string(), direction)],
        },
    }
}

// Normalize ORDER BY into a canonical, deterministic shape:
// - preserve user field order
// - remove explicit primary-key references from the user segment
// - append exactly one primary-key field as the terminal tie-break
fn canonicalize_order_spec(model: &EntityModel, order: Option<OrderSpec>) -> Option<OrderSpec> {
    let mut order = order?;
    if order.fields.is_empty() {
        return Some(order);
    }

    let pk_field = model.primary_key.name;
    let mut pk_direction = None;
    order.fields.retain(|(field, direction)| {
        if field == pk_field {
            if pk_direction.is_none() {
                pk_direction = Some(*direction);
            }
            false
        } else {
            true
        }
    });

    let pk_direction = pk_direction.unwrap_or(OrderDirection::Asc);
    order.fields.push((pk_field.to_string(), pk_direction));

    Some(order)
}