es_entity/
query.rs

1//! Query execution infrastructure for event-sourced entities.
2//!
3//! This module provides the underlying query types used by the `es_query!` macro.
4//! **These types are not intended to be used directly** - instead, use the `es_query!`
5//! macro which provides a simpler interface for querying index tables and automatically
6//! hydrating entities from their events.
7//!
8//! # Example
9//!
10//! Instead of using these types directly, use the `es_query!` macro:
11//!
12//! ```rust,ignore
13//! es_query!(
14//!     "SELECT id FROM users WHERE name = $1",
15//!     name
16//! ).fetch_optional(&pool).await
17//! ```
18//!
19//! See the `es_query!` macro documentation for more details.
20
21use crate::{
22    db,
23    error::EntityHydrationError,
24    events::{EntityEvents, GenericEvent},
25    one_time_executor::IntoOneTimeExecutor,
26    traits::*,
27    tree_query::{TreeQuerySource, build_tree_query, partition_by_tag},
28};
29
30/// Query builder for event-sourced entities.
31///
32/// This type is generated by the `es_query!` macro and should not be constructed directly.
33/// It wraps a SQLx query and provides methods to fetch and hydrate entities from their events.
34pub struct EsQuery<'q, Repo, Flavor, F, A>
35where
36    Repo: EsRepo,
37{
38    inner: sqlx::query::Map<'q, db::Db, F, A>,
39    source: TreeQuerySource<<<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId>,
40    _repo: std::marker::PhantomData<Repo>,
41    _flavor: std::marker::PhantomData<Flavor>,
42}
43
44/// Query flavor for flat entities without nested relationships.
45pub struct EsQueryFlavorFlat;
46
47/// Query flavor for entities with nested relationships that need to be loaded recursively.
48pub struct EsQueryFlavorNested;
49
50impl<'q, Repo, Flavor, F, A> EsQuery<'q, Repo, Flavor, F, A>
51where
52    Repo: EsRepo,
53    <<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId: Unpin,
54    F: FnMut(
55            db::Row,
56        ) -> Result<
57            GenericEvent<<<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId>,
58            sqlx::Error,
59        > + Send,
60    A: 'q + Send + sqlx::IntoArguments<'q, db::Db>,
61{
62    pub fn new(
63        query: sqlx::query::Map<'q, db::Db, F, A>,
64        source: TreeQuerySource<
65            <<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId,
66        >,
67    ) -> Self {
68        Self {
69            inner: query,
70            source,
71            _repo: std::marker::PhantomData,
72            _flavor: std::marker::PhantomData,
73        }
74    }
75
76    async fn fetch_optional_inner<E: From<sqlx::Error> + From<EntityHydrationError>>(
77        self,
78        op: impl IntoOneTimeExecutor<'_>,
79    ) -> Result<Option<<Repo as EsRepo>::Entity>, E> {
80        let executor = op.into_executor();
81        let rows = executor.fetch_all(self.inner).await?;
82        if rows.is_empty() {
83            return Ok(None);
84        }
85
86        Ok(EntityEvents::load_first(rows.into_iter())?)
87    }
88
89    async fn fetch_n_inner<E: From<sqlx::Error> + From<EntityHydrationError>>(
90        self,
91        op: impl IntoOneTimeExecutor<'_>,
92        first: usize,
93    ) -> Result<(Vec<<Repo as EsRepo>::Entity>, bool), E> {
94        let executor = op.into_executor();
95        let rows = executor.fetch_all(self.inner).await?;
96        Ok(EntityEvents::load_n(rows.into_iter(), first)?)
97    }
98
99    async fn fetch_tree_rows<E: From<sqlx::Error>>(
100        self,
101        op: impl IntoOneTimeExecutor<'_>,
102        include_deleted: bool,
103    ) -> Result<
104        (
105            Vec<GenericEvent<<<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId>>,
106            std::collections::HashMap<i32, Vec<db::Row>>,
107        ),
108        E,
109    > {
110        let executor = op.into_executor();
111        let spec = <Repo as EsRepo>::nested_tree_spec();
112        let sql = build_tree_query(
113            self.source.user_sql,
114            self.source.order_by_cols,
115            &spec,
116            include_deleted,
117            self.source.n_user_args + 1,
118        );
119
120        let mut inner = self.inner;
121        let args = sqlx::Execute::take_arguments(&mut inner)
122            .map_err(sqlx::Error::Encode)?
123            .unwrap_or_default();
124
125        let rows: Vec<db::Row> = sqlx::query_with::<db::Db, _>(&sql, args)
126            .fetch_all(executor)
127            .await?;
128        let mut by_tag = partition_by_tag(rows)?;
129        let root_rows = by_tag.remove(&0).unwrap_or_default();
130        let root = root_rows
131            .iter()
132            .map(self.source.decode)
133            .collect::<Result<Vec<_>, sqlx::Error>>()?;
134        Ok((root, by_tag))
135    }
136}
137
138impl<'q, Repo, F, A> EsQuery<'q, Repo, EsQueryFlavorFlat, F, A>
139where
140    Repo: EsRepo,
141    <<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId: Unpin,
142    F: FnMut(
143            db::Row,
144        ) -> Result<
145            GenericEvent<<<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId>,
146            sqlx::Error,
147        > + Send,
148    A: 'q + Send + sqlx::IntoArguments<'q, db::Db>,
149{
150    /// Fetches at most one entity from the query results.
151    ///
152    /// Returns `Ok(None)` if no entities match the query, or `Ok(Some(entity))` if found.
153    pub async fn fetch_optional(
154        self,
155        op: impl IntoOneTimeExecutor<'_>,
156    ) -> Result<Option<<Repo as EsRepo>::Entity>, <Repo as EsRepo>::QueryError> {
157        self.fetch_optional_inner(op).await
158    }
159
160    /// Fetches up to `first` entities from the query results.
161    ///
162    /// Returns a tuple of (entities, has_more) where `has_more` indicates if there
163    /// were more entities available beyond the requested limit.
164    pub async fn fetch_n(
165        self,
166        op: impl IntoOneTimeExecutor<'_>,
167        first: usize,
168    ) -> Result<(Vec<<Repo as EsRepo>::Entity>, bool), <Repo as EsRepo>::QueryError> {
169        self.fetch_n_inner(op, first).await
170    }
171}
172
173impl<'q, Repo, F, A> EsQuery<'q, Repo, EsQueryFlavorNested, F, A>
174where
175    Repo: EsRepo,
176    <<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId: Unpin,
177    F: FnMut(
178            db::Row,
179        ) -> Result<
180            GenericEvent<<<<Repo as EsRepo>::Entity as EsEntity>::Event as EsEvent>::EntityId>,
181            sqlx::Error,
182        > + Send,
183    A: 'q + Send + sqlx::IntoArguments<'q, db::Db>,
184{
185    /// Fetches at most one entity and loads all nested relationships.
186    ///
187    /// Returns `Ok(None)` if no entities match, or `Ok(Some(entity))` with all
188    /// nested entities loaded if found.
189    pub async fn fetch_optional(
190        self,
191        op: impl IntoOneTimeExecutor<'_>,
192    ) -> Result<Option<<Repo as EsRepo>::Entity>, <Repo as EsRepo>::QueryError> {
193        self.fetch_optional_tree(op, false).await
194    }
195
196    /// Fetches up to `first` entities and loads all nested relationships.
197    ///
198    /// Returns a tuple of (entities, has_more) where all entities have their nested
199    /// relationships loaded, and `has_more` indicates if more entities were available.
200    pub async fn fetch_n(
201        self,
202        op: impl IntoOneTimeExecutor<'_>,
203        first: usize,
204    ) -> Result<(Vec<<Repo as EsRepo>::Entity>, bool), <Repo as EsRepo>::QueryError> {
205        self.fetch_n_tree(op, first, false).await
206    }
207
208    /// Like [`fetch_optional`](EsQuery::fetch_optional) but transitively includes
209    /// soft-deleted nested entities.
210    pub async fn fetch_optional_include_deleted(
211        self,
212        op: impl IntoOneTimeExecutor<'_>,
213    ) -> Result<Option<<Repo as EsRepo>::Entity>, <Repo as EsRepo>::QueryError> {
214        self.fetch_optional_tree(op, true).await
215    }
216
217    /// Like [`fetch_n`](EsQuery::fetch_n) but transitively includes soft-deleted
218    /// nested entities.
219    pub async fn fetch_n_include_deleted(
220        self,
221        op: impl IntoOneTimeExecutor<'_>,
222        first: usize,
223    ) -> Result<(Vec<<Repo as EsRepo>::Entity>, bool), <Repo as EsRepo>::QueryError> {
224        self.fetch_n_tree(op, first, true).await
225    }
226
227    async fn fetch_optional_tree(
228        self,
229        op: impl IntoOneTimeExecutor<'_>,
230        include_deleted: bool,
231    ) -> Result<Option<<Repo as EsRepo>::Entity>, <Repo as EsRepo>::QueryError> {
232        let (root, mut by_tag) = self
233            .fetch_tree_rows::<<Repo as EsRepo>::QueryError>(op, include_deleted)
234            .await?;
235        let Some(entity) = EntityEvents::load_first::<<Repo as EsRepo>::Entity>(root)? else {
236            return Ok(None);
237        };
238        let mut entities = [entity];
239        let mut cursor = 1i32;
240        <Repo as EsRepo>::hydrate_nested_from_rows::<<Repo as EsRepo>::QueryError>(
241            &mut by_tag,
242            &mut cursor,
243            &mut entities,
244        )?;
245        let [entity] = entities;
246        Ok(Some(entity))
247    }
248
249    async fn fetch_n_tree(
250        self,
251        op: impl IntoOneTimeExecutor<'_>,
252        first: usize,
253        include_deleted: bool,
254    ) -> Result<(Vec<<Repo as EsRepo>::Entity>, bool), <Repo as EsRepo>::QueryError> {
255        let (root, mut by_tag) = self
256            .fetch_tree_rows::<<Repo as EsRepo>::QueryError>(op, include_deleted)
257            .await?;
258        let (mut entities, more) = EntityEvents::load_n::<<Repo as EsRepo>::Entity>(root, first)?;
259        let mut cursor = 1i32;
260        <Repo as EsRepo>::hydrate_nested_from_rows::<<Repo as EsRepo>::QueryError>(
261            &mut by_tag,
262            &mut cursor,
263            &mut entities,
264        )?;
265        Ok((entities, more))
266    }
267}
es_entity/query.rs

es_entity/
query.rs
