bevy_persistence_database/bevy/query/
persistence_query.rs

1//! A manual builder for creating and executing database queries that load results into a Bevy `World`.
2
3use crate::bevy::plugins::persistence_plugin::{PersistencePluginConfig, TokioRuntime};
4use crate::core::db::connection::{DatabaseConnectionResource, DocumentKind};
5use crate::core::db::{DatabaseConnection, read_version};
6use crate::core::persist::Persist;
7use crate::core::query::{FilterExpression, PersistenceQuerySpecification};
8use crate::core::session::PersistenceSession;
9use crate::core::versioning::version_manager::VersionKey;
10use bevy::prelude::{Component, Entity, World};
11use std::sync::Arc;
12
13/// Query builder: select which components and filters to apply.
14pub struct PersistenceQuery {
15    db: Option<Arc<dyn DatabaseConnection>>,
16    store: Option<String>,
17    pub component_names: Vec<&'static str>,
18    filter_expr: Option<FilterExpression>,
19
20    /// Track explicit absence filters for components.
21    pub(crate) without_component_names: Vec<&'static str>,
22
23    /// Components to fetch/deserialize without gating presence in backend.
24    pub(crate) fetch_only_component_names: Vec<&'static str>,
25
26    /// Whether to return full documents (internal use).
27    force_full_docs: bool,
28}
29
30impl PersistenceQuery {
31    /// Create a new query. The database connection and store are resolved from the
32    /// world's resources when [`run`](Self::run) is called, so no arguments are needed
33    /// for the common case.
34    ///
35    /// Use [`with_db`](Self::with_db) and [`store`](Self::store) to override either
36    /// value explicitly (required when calling [`fetch_into`](Self::fetch_into) or
37    /// [`fetch_ids`](Self::fetch_ids) directly without going through `run`).
38    pub fn new() -> Self {
39        Self {
40            db: None,
41            store: None,
42            component_names: Vec::new(),
43            filter_expr: None,
44            without_component_names: Vec::new(),
45            fetch_only_component_names: Vec::new(),
46            force_full_docs: false,
47        }
48    }
49
50    /// Explicitly supply the database connection instead of reading it from the world.
51    ///
52    /// Required when calling [`fetch_into`](Self::fetch_into) or
53    /// [`fetch_ids`](Self::fetch_ids) directly (i.e. outside of [`run`](Self::run)).
54    pub fn with_db(mut self, db: Arc<dyn DatabaseConnection>) -> Self {
55        self.db = Some(db);
56        self
57    }
58
59    /// Override the store to query against.
60    ///
61    /// Defaults to `PersistencePluginConfig::default_store` when not set.
62    pub fn store(mut self, store: impl Into<String>) -> Self {
63        self.store = Some(store.into());
64        self
65    }
66
67    /// Request loading component `T`.
68    pub fn with<T: Component + Persist>(mut self) -> Self {
69        self.component_names.push(T::name());
70        self
71    }
72
73    /// Request absence of component `T`.
74    pub fn without<T: Component + Persist>(mut self) -> Self {
75        self.without_component_names.push(T::name());
76        self
77    }
78
79    /// Internal: request fetching component by name without presence gating.
80    pub fn fetch_only_component(mut self, component_name: &'static str) -> Self {
81        self.fetch_only_component_names.push(component_name);
82        self
83    }
84
85    /// Combine current filter with OR.
86    pub fn or(mut self, expression: FilterExpression) -> Self {
87        self.filter_expr = Some(match self.filter_expr.take() {
88            Some(existing) => existing.or(expression),
89            None => expression,
90        });
91        self
92    }
93
94    /// Filter results to only documents whose primary key is in `guids`.
95    ///
96    /// Issues a single batched query — does not fire one query per key.
97    /// Combines with any existing filter using AND.
98    pub fn by_guids(mut self, guids: Vec<String>) -> Self {
99        let key_filter = FilterExpression::DocumentKey.in_(guids);
100        self.filter_expr = Some(match self.filter_expr.take() {
101            Some(existing) => existing.and(key_filter),
102            None => key_filter,
103        });
104        self
105    }
106
107    /// Execute this query synchronously against a live Bevy world.
108    ///
109    /// Reads the database connection and default store from the world's
110    /// `DatabaseConnectionResource` and `PersistencePluginConfig` resources if not
111    /// explicitly set via [`with_db`](Self::with_db) and [`store`](Self::store).
112    ///
113    /// Intended for use inside exclusive systems (`fn my_system(world: &mut World)`).
114    pub fn run(mut self, world: &mut World) -> Vec<Entity> {
115        if self.db.is_none() {
116            self.db = Some(
117                world
118                    .resource::<DatabaseConnectionResource>()
119                    .connection
120                    .clone(),
121            );
122        }
123        if self.store.is_none() {
124            self.store = Some(
125                world
126                    .resource::<PersistencePluginConfig>()
127                    .default_store
128                    .clone(),
129            );
130        }
131        let runtime = world.resource::<TokioRuntime>().runtime.clone();
132        runtime.block_on(self.fetch_into(world))
133    }
134
135    /// Sets the filter for the query using a `FilterExpression`.
136    pub fn filter(mut self, expression: FilterExpression) -> Self {
137        fn collect(expr: &FilterExpression, names: &mut Vec<&'static str>) {
138            match expr {
139                FilterExpression::Field { component_name, .. } => {
140                    if !names.contains(component_name) {
141                        names.push(component_name);
142                    }
143                }
144                FilterExpression::DocumentKey => {}
145                FilterExpression::BinaryOperator { lhs, rhs, .. } => {
146                    collect(lhs, names);
147                    collect(rhs, names);
148                }
149                FilterExpression::Literal(_) => {}
150            }
151        }
152
153        let mut names = self.component_names.clone();
154        collect(&expression, &mut names);
155        self.component_names = names;
156
157        self.filter_expr = Some(expression);
158        self
159    }
160
161    /// For component tests - internal use.
162    #[cfg(test)]
163    pub fn for_component<T: Component + Persist>(mut self) -> Self {
164        self.component_names.push(T::name());
165        self
166    }
167
168    /// Build a backend-agnostic spec.
169    pub fn build_spec(&self) -> PersistenceQuerySpecification {
170        let mut fetch_only = self.component_names.clone();
171        fetch_only.extend(self.fetch_only_component_names.iter().copied());
172        fetch_only.sort_unstable();
173        fetch_only.dedup();
174
175        let presence_with = self.component_names.clone();
176        let presence_without = self.without_component_names.clone();
177        let value_filters = self.filter_expr.clone();
178        let force_full_docs = self.force_full_docs;
179
180        let spec = PersistenceQuerySpecification {
181            store: self.store.clone().unwrap_or_default(),
182            kind: DocumentKind::Entity,
183            presence_with: presence_with.clone(),
184            presence_without: presence_without.clone(),
185            fetch_only: fetch_only.clone(),
186            value_filters: value_filters.clone(),
187            return_full_docs: force_full_docs
188                || (presence_with.is_empty() && presence_without.is_empty()),
189            pagination: None,
190        };
191
192        bevy::log::debug!(
193            "[builder] build_spec full_docs={} presence_with={:?} without={:?} fetch_only={:?} filter={:?}",
194            force_full_docs,
195            spec.presence_with,
196            spec.presence_without,
197            spec.fetch_only,
198            spec.value_filters
199        );
200
201        spec
202    }
203
204    /// Run the query for keys only.
205    pub async fn fetch_ids(&self) -> Vec<String> {
206        let db = self
207            .db
208            .as_ref()
209            .expect("PersistenceQuery: call with_db() before fetch_ids()");
210        let spec = self.build_spec();
211        db.execute_keys(&spec).await.expect("query failed")
212    }
213
214    /// Load matching entities into the World.
215    pub async fn fetch_into(&self, world: &mut World) -> Vec<bevy::prelude::Entity> {
216        let db = self
217            .db
218            .as_ref()
219            .expect("PersistenceQuery: call with_db() or use run() before fetch_into()");
220        let store = self
221            .store
222            .as_deref()
223            .expect("PersistenceQuery: call store() or use run() before fetch_into()");
224
225        let mut session = world
226            .remove_resource::<PersistenceSession>()
227            .expect("PersistenceSession missing");
228
229        let spec = self.build_spec();
230
231        bevy::log::debug!(
232            "[builder] fetch_into issuing execute_documents (partial_projection={})",
233            !spec.return_full_docs
234        );
235        let documents = db
236            .execute_documents(&spec)
237            .await
238            .expect("Batch document fetch failed");
239        bevy::log::debug!(
240            "[builder] fetch_into: backend returned {} documents",
241            documents.len()
242        );
243
244        let mut result = Vec::with_capacity(documents.len());
245        if !documents.is_empty() {
246            // When no components are explicitly requested the query is unconstrained;
247            // hydrate every registered component field found in each document.
248            let mut explicit_components: Vec<&'static str> = self.component_names.clone();
249            explicit_components.extend(self.fetch_only_component_names.iter().copied());
250            explicit_components.sort_unstable();
251            explicit_components.dedup();
252
253            for doc in documents {
254                let key_field = db.document_key_field();
255                let key = doc[key_field].as_str().unwrap_or_default().to_string();
256                if key.is_empty() {
257                    bevy::log::debug!(
258                        "[builder] fetch_into: skipping doc missing key '{}'",
259                        key_field
260                    );
261                    continue;
262                }
263
264                bevy::log::trace!(
265                    "[builder] deserializing {:?} for key={}",
266                    explicit_components,
267                    key
268                );
269                let entity = session
270                    .materialize_entity_document(
271                        world,
272                        &doc,
273                        key_field,
274                        &explicit_components,
275                        true,
276                    )
277                    .expect("component deserialization failed")
278                    .expect("document key should be present");
279
280                result.push(entity);
281            }
282        }
283
284        session
285            .fetch_and_insert_resources(&**db, store, world)
286            .await
287            .expect("resource deserialization failed");
288
289        world.insert_resource(session);
290        bevy::log::debug!(
291            "[builder] fetch_into: inserted {} entities into world",
292            result.len()
293        );
294        result
295    }
296
297    /// Re-read persisted versions from the database into the in-memory version cache,
298    /// without spawning entities or deserializing components.
299    ///
300    /// **Not for routine gameplay.** In a single-writer setup, an optimistic-concurrency
301    /// conflict during normal commits indicates a pipeline bug (version cache drift, load
302    /// dirt, etc.) and should be investigated — not silently papered over.
303    ///
304    /// Use this explicitly when the database was changed outside the running app (ops,
305    /// migration, manual repair, restore) and the in-memory version map must be realigned
306    /// to match. Live ECS state is left untouched; only the OCC version map is refreshed.
307    /// Returns the number of entity/resource versions updated.
308    ///
309    /// Intended for use inside exclusive systems (`fn my_system(world: &mut World)`).
310    pub fn reconcile_versions(mut self, world: &mut World) -> usize {
311        if self.db.is_none() {
312            self.db = Some(
313                world
314                    .resource::<DatabaseConnectionResource>()
315                    .connection
316                    .clone(),
317            );
318        }
319        if self.store.is_none() {
320            self.store = Some(
321                world
322                    .resource::<PersistencePluginConfig>()
323                    .default_store
324                    .clone(),
325            );
326        }
327        let runtime = world.resource::<TokioRuntime>().runtime.clone();
328        runtime.block_on(self.reconcile_versions_into(world))
329    }
330
331    async fn reconcile_versions_into(&self, world: &mut World) -> usize {
332        let db = self
333            .db
334            .as_ref()
335            .expect("PersistenceQuery: call with_db() or use run() before reconcile_versions()");
336
337        let mut query_with_full_docs = self.clone();
338        query_with_full_docs.force_full_docs = true;
339        let spec = query_with_full_docs.build_spec();
340
341        let documents = match db.execute_documents(&spec).await {
342            Ok(docs) => docs,
343            Err(e) => {
344                bevy::log::error!("version reconcile: document fetch failed: {e}");
345                return 0;
346            }
347        };
348
349        let mut session = world
350            .remove_resource::<PersistenceSession>()
351            .expect("PersistenceSession missing");
352
353        let key_field = db.document_key_field();
354        let store = self
355            .store
356            .as_deref()
357            .expect("PersistenceQuery: call store() or use run() before reconcile_versions()");
358        let mut realigned = 0usize;
359        for doc in &documents {
360            let key = doc[key_field].as_str().unwrap_or_default().to_string();
361            if key.is_empty() {
362                continue;
363            }
364            let version = read_version(doc).unwrap_or(1);
365            session
366                .version_manager_mut()
367                .set_version(VersionKey::Entity(key), version);
368            realigned += 1;
369        }
370
371        let resource_types: Vec<_> = session.persisted_resource_types().collect();
372        for type_id in resource_types {
373            let Some(res_name) = session.resource_name_for_type(type_id) else {
374                continue;
375            };
376            if let Ok(Some((_, version))) = db.fetch_resource(store, res_name).await {
377                session
378                    .version_manager_mut()
379                    .set_version(VersionKey::Resource(type_id), version);
380                realigned += 1;
381            }
382        }
383
384        world.insert_resource(session);
385        bevy::log::info!(
386            "version reconcile: realigned {realigned} entity/resource versions from the database"
387        );
388        realigned
389    }
390}
391
392impl Clone for PersistenceQuery {
393    fn clone(&self) -> Self {
394        Self {
395            db: self.db.clone(),
396            store: self.store.clone(),
397            component_names: self.component_names.clone(),
398            filter_expr: self.filter_expr.clone(),
399            without_component_names: self.without_component_names.clone(),
400            fetch_only_component_names: self.fetch_only_component_names.clone(),
401            force_full_docs: self.force_full_docs,
402        }
403    }
404}
405
406impl Default for PersistenceQuery {
407    fn default() -> Self {
408        Self::new()
409    }
410}
411
412/// Extension trait for `PersistenceQuery` to add component by name.
413pub trait WithComponentExt {
414    fn with_component(self, component_name: &'static str) -> Self;
415    fn without_component(self, component_name: &'static str) -> Self;
416    fn fetch_only_component(self, component_name: &'static str) -> Self;
417}
418
419impl WithComponentExt for PersistenceQuery {
420    fn with_component(mut self, component_name: &'static str) -> Self {
421        self.component_names.push(component_name);
422        self
423    }
424
425    fn without_component(mut self, component_name: &'static str) -> Self {
426        self.without_component_names.push(component_name);
427        self
428    }
429
430    fn fetch_only_component(mut self, component_name: &'static str) -> Self {
431        self.fetch_only_component_names.push(component_name);
432        self
433    }
434}
435
436#[cfg(test)]
437mod tests {
438    use super::*;
439    use crate::bevy::plugins::persistence_plugin::PersistencePluginCore;
440    use crate::core::db::{
441        MockDatabaseConnection,
442        connection::{
443            BEVY_PERSISTENCE_DATABASE_BEVY_TYPE_FIELD, BEVY_PERSISTENCE_DATABASE_METADATA_FIELD,
444            BEVY_PERSISTENCE_DATABASE_VERSION_FIELD, DocumentKind, PersistenceError,
445        },
446    };
447    use crate::core::session::PersistenceSession;
448    use bevy::MinimalPlugins;
449    use bevy::prelude::App;
450    use bevy_persistence_database_derive::persist;
451    use futures::executor::block_on;
452    use serde_json::json;
453    use std::sync::Arc;
454
455    const TEST_STORE: &str = "test_store";
456
457    #[persist(component)]
458    struct A {
459        value: i32,
460    }
461
462    #[persist(component)]
463    struct B {
464        name: String,
465    }
466
467    #[test]
468    fn build_spec_with_dsl() {
469        let db = Arc::new(MockDatabaseConnection::new());
470        let query = PersistenceQuery::new()
471            .with_db(db)
472            .store(TEST_STORE)
473            .with::<A>()
474            .filter(A::value().gt(10).and(B::name().eq("test")));
475
476        let spec = query.build_spec();
477        assert!(spec.presence_with.contains(&<A as Persist>::name()));
478        assert!(spec.presence_with.contains(&<B as Persist>::name()));
479        assert!(spec.value_filters.is_some());
480        assert!(!spec.return_full_docs);
481    }
482
483    #[test]
484    fn build_spec_with_or_combiner() {
485        let db = Arc::new(MockDatabaseConnection::new());
486        let query = PersistenceQuery::new()
487            .with_db(db)
488            .store(TEST_STORE)
489            .filter(A::value().gt(10))
490            .or(B::name().eq("foo"));
491        let spec = query.build_spec();
492
493        assert!(spec.value_filters.is_some());
494        assert!(!spec.return_full_docs);
495    }
496
497    #[persist(component)]
498    struct Health {
499        value: i32,
500    }
501
502    #[persist(component)]
503    struct Position {
504        x: f32,
505        y: f32,
506    }
507
508    #[tokio::test]
509    async fn fetch_into_loads_new_entities() {
510        let mut mock_db = MockDatabaseConnection::new();
511        mock_db.expect_document_key_field().return_const("_key");
512        mock_db.expect_execute_documents().returning(|spec| {
513            assert!(
514                !spec.return_full_docs,
515                "narrow query should use partial projection"
516            );
517            assert!(spec.fetch_only.contains(&"Health"));
518            assert!(spec.fetch_only.contains(&"Position"));
519            Box::pin(async {
520                Ok(vec![
521                    json!({
522                        "_key":"k1",
523                        BEVY_PERSISTENCE_DATABASE_METADATA_FIELD: {
524                            BEVY_PERSISTENCE_DATABASE_VERSION_FIELD: 1,
525                            BEVY_PERSISTENCE_DATABASE_BEVY_TYPE_FIELD: DocumentKind::Entity.as_str(),
526                        },
527                        "Health": 1,
528                        "Position": {"x": 1.0, "y": 2.0},
529                    }),
530                    json!({
531                        "_key":"k2",
532                        BEVY_PERSISTENCE_DATABASE_METADATA_FIELD: {
533                            BEVY_PERSISTENCE_DATABASE_VERSION_FIELD: 1,
534                            BEVY_PERSISTENCE_DATABASE_BEVY_TYPE_FIELD: DocumentKind::Entity.as_str(),
535                        },
536                        "Health": 3,
537                        "Position": {"x": 4.0, "y": 5.0},
538                    }),
539                ])
540            })
541        });
542        mock_db
543            .expect_fetch_resource()
544            .returning(|_, _| Box::pin(async { Ok(None) }));
545
546        let db = Arc::new(mock_db) as Arc<dyn DatabaseConnection>;
547
548        let mut app = App::new();
549        app.add_plugins(MinimalPlugins);
550        app.add_plugins(PersistencePluginCore::new(db.clone()));
551
552        {
553            let mut session = app
554                .world_mut()
555                .resource_mut::<PersistenceSession>();
556            session.register_component::<Health>();
557            session.register_component::<Position>();
558        }
559
560        let query = PersistenceQuery::new()
561            .with_db(db)
562            .store(TEST_STORE)
563            .with::<Health>()
564            .with::<Position>();
565        let loaded = query.fetch_into(app.world_mut()).await;
566
567        assert_eq!(loaded.len(), 2);
568    }
569
570    #[tokio::test]
571    async fn fetch_into_unconstrained_uses_full_docs() {
572        let mut mock_db = MockDatabaseConnection::new();
573        mock_db.expect_document_key_field().return_const("_key");
574        mock_db.expect_execute_documents().returning(|spec| {
575            assert!(
576                spec.return_full_docs,
577                "unconstrained query should fetch full documents"
578            );
579            Box::pin(async { Ok(vec![]) })
580        });
581        mock_db
582            .expect_fetch_resource()
583            .returning(|_, _| Box::pin(async { Ok(None) }));
584
585        let db = Arc::new(mock_db) as Arc<dyn DatabaseConnection>;
586        let mut app = App::new();
587        app.add_plugins(MinimalPlugins);
588        app.add_plugins(PersistencePluginCore::new(db.clone()));
589
590        let query = PersistenceQuery::new().with_db(db).store(TEST_STORE);
591        let _ = query.fetch_into(app.world_mut()).await;
592    }
593
594    #[test]
595    fn build_spec_empty_filters() {
596        #[persist(component)]
597        struct Comp1;
598
599        let db = Arc::new(MockDatabaseConnection::new());
600        let query = PersistenceQuery::new().with_db(db).store(TEST_STORE).for_component::<Comp1>();
601        let spec = query.build_spec();
602
603        assert!(!spec.presence_with.is_empty());
604        assert!(spec.presence_without.is_empty());
605        assert!(!spec.return_full_docs);
606        assert_eq!(spec.fetch_only, vec!["Comp1"]);
607    }
608
609    #[test]
610    #[should_panic(expected = "query failed")]
611    fn fetch_ids_panics_on_error() {
612        let mut mock_db = MockDatabaseConnection::new();
613        mock_db.expect_execute_keys().returning(|_spec| {
614            Box::pin(async { Err(PersistenceError::General("db error".into())) })
615        });
616        let db = Arc::new(mock_db);
617        let query = PersistenceQuery::new().with_db(db).store(TEST_STORE);
618        block_on(query.fetch_ids());
619    }
620}
bevy_persistence_database/bevy/query/persistence_query.rs

bevy_persistence_database/bevy/query/
persistence_query.rs
