es_entity/

query.rs

//! Query execution infrastructure for event-sourced entities.
//!
//! This module provides the underlying query types used by the `es_query!` macro.
//! **These types are not intended to be used directly** - instead, use the `es_query!`
//! macro which provides a simpler interface for querying index tables and automatically
//! hydrating entities from their events.
//!
//! # Example
//!
//! Instead of using these types directly, use the `es_query!` macro:
//!
//! ```rust,ignore
//! es_query!(
//!     "SELECT id FROM users WHERE name = $1",
//!     name
//! ).fetch_optional(&pool).await
//! ```
//!
//! See the `es_query!` macro documentation for more details.

use crate::{
    db,
    error::EntityHydrationError,
    events::{EntityEvents, HydrationRow},
    one_time_executor::IntoOneTimeExecutor,
    snapshot::NO_SNAPSHOT_FINGERPRINT,
    traits::*,
    tree_query::{TreeQuerySource, build_tree_query, partition_by_tag, snapshot_fingerprints},
};

/// Query builder for event-sourced entities.
///
/// This type is generated by the `es_query!` macro and should not be constructed directly.
/// It wraps a SQLx query and provides methods to fetch and hydrate entities from their events.
///
/// `R` is the row type the underlying `sqlx::query_as!` decodes into —
/// `GenericEvent<Id>` for a plain repo, `SnapshotGenericEvent<Id>` for a
/// `#[es_repo(snapshot)]` repo — normalised into a `HydrationRow<Id>` before
/// an entity is built from it.
pub struct EsQuery<
    'q,
    Repo,
    Flavor,
    F,
    A,
    R = crate::events::GenericEvent<
        <<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId,
    >,
> where
    Repo: EsRepo,
{
    inner: sqlx::query::Map<'q, db::Db, F, A>,
    source: TreeQuerySource<<<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId>,
    _repo: std::marker::PhantomData<Repo>,
    _flavor: std::marker::PhantomData<Flavor>,
    _row: std::marker::PhantomData<fn() -> R>,
}

/// Query flavor for flat entities without nested relationships.
pub struct EsQueryFlavorFlat;

/// Query flavor for entities with nested relationships that need to be loaded recursively.
pub struct EsQueryFlavorNested;

impl<'q, Repo, Flavor, F, A, R> EsQuery<'q, Repo, Flavor, F, A, R>
where
    Repo: EsRepo,
    <<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId: Unpin,
    F: FnMut(db::Row) -> Result<R, sqlx::Error> + Send,
    R: Into<HydrationRow<<<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId>>
        + Send
        + Unpin,
    A: 'q + Send + sqlx::IntoArguments<'q, db::Db>,
{
    pub fn new(
        query: sqlx::query::Map<'q, db::Db, F, A>,
        source: TreeQuerySource<
            <<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId,
        >,
    ) -> Self {
        Self {
            inner: query,
            source,
            _repo: std::marker::PhantomData,
            _flavor: std::marker::PhantomData,
            _row: std::marker::PhantomData,
        }
    }

    async fn fetch_optional_inner<E: From<sqlx::Error> + From<EntityHydrationError>>(
        self,
        op: impl IntoOneTimeExecutor<'_>,
    ) -> Result<Option<<Repo as EsRepo>::Entity>, E> {
        let executor = op.into_executor();
        let rows = executor.fetch_all(self.inner).await?;
        if rows.is_empty() {
            return Ok(None);
        }

        Ok(EntityEvents::load_first(rows.into_iter())?)
    }

    async fn fetch_n_inner<E: From<sqlx::Error> + From<EntityHydrationError>>(
        self,
        op: impl IntoOneTimeExecutor<'_>,
        first: usize,
    ) -> Result<(Vec<<Repo as EsRepo>::Entity>, bool), E> {
        let executor = op.into_executor();
        let rows = executor.fetch_all(self.inner).await?;
        Ok(EntityEvents::load_n(rows.into_iter(), first)?)
    }

    async fn fetch_tree_rows<E: From<sqlx::Error>>(
        self,
        op: impl IntoOneTimeExecutor<'_>,
        include_deleted: bool,
    ) -> Result<
        (
            Vec<HydrationRow<<<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId>>,
            std::collections::HashMap<i32, Vec<db::Row>>,
        ),
        E,
    > {
        let executor = op.into_executor();
        let spec = <Repo as EsRepo>::nested_tree_spec();
        let sql = build_tree_query(
            self.source.user_sql,
            self.source.order_by_cols,
            &spec,
            include_deleted,
            self.source.n_user_args + 1,
        );

        let mut inner = self.inner;
        let mut args = sqlx::Execute::take_arguments(&mut inner)
            .map_err(sqlx::Error::Encode)?
            .unwrap_or_default();

        // A full-history load binds `NO_SNAPSHOT_FINGERPRINT` as the root's
        // own fingerprint, forcing every node in the tree to bind the
        // sentinel too. The `snapshot_table_name` check guards against a
        // false positive: `NoSnapshot::FINGERPRINT` *equals* that sentinel,
        // so an ordinary non-snapshot root would otherwise always match.
        let full_history = spec.snapshot_table_name.is_some()
            && self.source.snapshot_fingerprint == NO_SNAPSHOT_FINGERPRINT;
        // The root's own `es_query!` call site already bound its
        // fingerprint, so only descendants need binding here — walked in
        // the same DFS order `build_tree_query` assigned positions in.
        let mut fingerprints = snapshot_fingerprints(&spec).into_iter();
        if spec.snapshot_table_name.is_some() {
            fingerprints.next();
        }
        for fp in fingerprints {
            let bind = if full_history {
                NO_SNAPSHOT_FINGERPRINT
            } else {
                fp
            };
            sqlx::Arguments::add(&mut args, bind).map_err(sqlx::Error::Encode)?;
        }

        let rows: Vec<db::Row> = sqlx::query_with::<db::Db, _>(&sql, args)
            .fetch_all(executor)
            .await?;
        let mut by_tag = partition_by_tag(rows)?;
        let root_rows = by_tag.remove(&0).unwrap_or_default();
        let root = root_rows
            .iter()
            .map(self.source.decode)
            .collect::<Result<Vec<_>, sqlx::Error>>()?;
        Ok((root, by_tag))
    }
}

impl<'q, Repo, F, A, R> EsQuery<'q, Repo, EsQueryFlavorFlat, F, A, R>
where
    Repo: EsRepo,
    <<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId: Unpin,
    F: FnMut(db::Row) -> Result<R, sqlx::Error> + Send,
    R: Into<HydrationRow<<<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId>>
        + Send
        + Unpin,
    A: 'q + Send + sqlx::IntoArguments<'q, db::Db>,
{
    /// Fetches at most one entity from the query results.
    ///
    /// Returns `Ok(None)` if no entities match the query, or `Ok(Some(entity))` if found.
    pub async fn fetch_optional(
        self,
        op: impl IntoOneTimeExecutor<'_>,
    ) -> Result<Option<<Repo as EsRepo>::Entity>, <Repo as EsRepo>::QueryError> {
        self.fetch_optional_inner(op).await
    }

    /// Fetches up to `first` entities from the query results.
    ///
    /// Returns a tuple of (entities, has_more) where `has_more` indicates if there
    /// were more entities available beyond the requested limit.
    pub async fn fetch_n(
        self,
        op: impl IntoOneTimeExecutor<'_>,
        first: usize,
    ) -> Result<(Vec<<Repo as EsRepo>::Entity>, bool), <Repo as EsRepo>::QueryError> {
        self.fetch_n_inner(op, first).await
    }
}

impl<'q, Repo, F, A, R> EsQuery<'q, Repo, EsQueryFlavorNested, F, A, R>
where
    Repo: EsRepo,
    <<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId: Unpin,
    F: FnMut(db::Row) -> Result<R, sqlx::Error> + Send,
    R: Into<HydrationRow<<<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId>>
        + Send
        + Unpin,
    A: 'q + Send + sqlx::IntoArguments<'q, db::Db>,
{
    /// Fetches at most one entity and loads all nested relationships.
    ///
    /// Returns `Ok(None)` if no entities match, or `Ok(Some(entity))` with all
    /// nested entities loaded if found.
    pub async fn fetch_optional(
        self,
        op: impl IntoOneTimeExecutor<'_>,
    ) -> Result<Option<<Repo as EsRepo>::Entity>, <Repo as EsRepo>::QueryError> {
        self.fetch_optional_tree(op, false).await
    }

    /// Fetches up to `first` entities and loads all nested relationships.
    ///
    /// Returns a tuple of (entities, has_more) where all entities have their nested
    /// relationships loaded, and `has_more` indicates if more entities were available.
    pub async fn fetch_n(
        self,
        op: impl IntoOneTimeExecutor<'_>,
        first: usize,
    ) -> Result<(Vec<<Repo as EsRepo>::Entity>, bool), <Repo as EsRepo>::QueryError> {
        self.fetch_n_tree(op, first, false).await
    }

    /// Like [`fetch_optional`](EsQuery::fetch_optional) but transitively includes
    /// soft-deleted nested entities.
    pub async fn fetch_optional_include_deleted(
        self,
        op: impl IntoOneTimeExecutor<'_>,
    ) -> Result<Option<<Repo as EsRepo>::Entity>, <Repo as EsRepo>::QueryError> {
        self.fetch_optional_tree(op, true).await
    }

    /// Like [`fetch_n`](EsQuery::fetch_n) but transitively includes soft-deleted
    /// nested entities.
    pub async fn fetch_n_include_deleted(
        self,
        op: impl IntoOneTimeExecutor<'_>,
        first: usize,
    ) -> Result<(Vec<<Repo as EsRepo>::Entity>, bool), <Repo as EsRepo>::QueryError> {
        self.fetch_n_tree(op, first, true).await
    }

    async fn fetch_optional_tree(
        self,
        op: impl IntoOneTimeExecutor<'_>,
        include_deleted: bool,
    ) -> Result<Option<<Repo as EsRepo>::Entity>, <Repo as EsRepo>::QueryError> {
        let (root, mut by_tag) = self
            .fetch_tree_rows::<<Repo as EsRepo>::QueryError>(op, include_deleted)
            .await?;
        let Some(entity) = EntityEvents::load_first::<<Repo as EsRepo>::Entity>(root)? else {
            return Ok(None);
        };
        let mut entities = [entity];
        let mut cursor = 1i32;
        <Repo as EsRepo>::hydrate_nested_from_rows::<<Repo as EsRepo>::QueryError>(
            &mut by_tag,
            &mut cursor,
            &mut entities,
        )?;
        let [entity] = entities;
        Ok(Some(entity))
    }

    async fn fetch_n_tree(
        self,
        op: impl IntoOneTimeExecutor<'_>,
        first: usize,
        include_deleted: bool,
    ) -> Result<(Vec<<Repo as EsRepo>::Entity>, bool), <Repo as EsRepo>::QueryError> {
        let (root, mut by_tag) = self
            .fetch_tree_rows::<<Repo as EsRepo>::QueryError>(op, include_deleted)
            .await?;
        let (mut entities, more) = EntityEvents::load_n::<<Repo as EsRepo>::Entity>(root, first)?;
        let mut cursor = 1i32;
        <Repo as EsRepo>::hydrate_nested_from_rows::<<Repo as EsRepo>::QueryError>(
            &mut by_tag,
            &mut cursor,
            &mut entities,
        )?;
        Ok((entities, more))
    }
}