bevy_erm_core/
database_query.rs

1use crate::*;
2use async_trait::async_trait;
3
4use bevy_ecs::query::{QueryItem, ROQueryItem, WorldQuery};
5use bevy_ecs::{
6    component::Component, prelude::*, system::SystemParam,
7    world::unsafe_world_cell::UnsafeWorldCell,
8};
9use bevy_mod_index::prelude::*;
10use bevy_utils::hashbrown::HashSet;
11use casey::lower;
12
13#[async_trait]
14pub trait CustomDatabaseQuery<DbResource: DatabaseResource, DerefItem: Send>: Sync {
15    async fn query(
16        &self,
17        tr: DatabaseTransaction<DbResource>,
18    ) -> Result<Vec<(DatabaseEntity, DerefItem)>, ()>;
19}
20
21pub trait ReturnSelector<'w> {
22    type ReturnItem;
23
24    fn load_components_from_entities(
25        world: UnsafeWorldCell<'w>,
26        entities: Vec<Entity>,
27    ) -> Vec<Self::ReturnItem>;
28}
29
30pub type DatabaseTransaction<'a, D> = <D as DatabaseResource>::Transaction;
31
32// The items to be returned as readonly or mutable
33// Use the World queries underlying specification
34pub type ReadOnlyItem<'a, Q, DbResource> =
35    ROQueryItem<'a, <Q as DBQueryInfo<DbResource>>::WorldQuery<'a>>;
36pub type Item<'a, Q, DbResource> = QueryItem<'a, <Q as DBQueryInfo<DbResource>>::WorldQuery<'a>>;
37#[async_trait]
38pub trait DBQueryInfo<DbResource: DatabaseResource> {
39    // type Mapper: ComponentMapper<Transaction = DbResource::Transaction>;
40    // Using the world query to specify the type of the item
41    // Allows for being able to use a bevy query from the DBQueryInfo
42    // Used when inserting or updating components for the Database Query
43    type WorldQuery<'a>: WorldQuery;
44    type DerefItem: Send;
45    type ReadOnlyItem<'a>: Send + From<ReadOnlyItem<'a, Self, DbResource>>;
46    type Item<'a>: Send;
47
48    async fn get<'w, D: DatabaseEntityWithRequest>(
49        db: &DbResource,
50        world: UnsafeWorldCell<'w>,
51        db_entity: &D,
52    ) -> Result<Self::ReadOnlyItem<'w>, ()>;
53    async fn get_mut<'w, D: DatabaseEntityWithRequest>(
54        db: &DbResource,
55        world: UnsafeWorldCell<'w>,
56        db_entity: &D,
57    ) -> Result<Self::Item<'w>, ()>;
58    async fn update_component<'w>(
59        db: &DbResource,
60        world: UnsafeWorldCell<'w>,
61        db_entity: &DatabaseEntity,
62        component: Self::ReadOnlyItem<'w>,
63    ) -> Result<(), ()>;
64    async fn insert_component<'w>(
65        db: &DbResource,
66        world: UnsafeWorldCell<'w>,
67        db_entity: &DatabaseEntity,
68        component: Self::ReadOnlyItem<'w>,
69    ) -> Result<(), ()>;
70    async fn load_components<'w, R: ReturnSelector<'w>, CDQ>(
71        db: &DbResource,
72        world: UnsafeWorldCell<'w>,
73        request: RequestId,
74        get_comp_from_db: CDQ,
75    ) -> Result<Vec<<R as ReturnSelector<'w>>::ReturnItem>, ()>
76    where
77        CDQ: CustomDatabaseQuery<DbResource, Self::DerefItem> + Send;
78    async fn create(
79        db: &DbResource,
80        world: UnsafeWorldCell<'_>,
81        component: Self::DerefItem,
82        request: RequestId,
83    ) -> Result<(), ()>;
84}
85
86pub struct QueryFetchState<'w, 's, DbResource: DatabaseResource> {
87    db_state: <Res<'w, DbResource> as SystemParam>::State,
88    phantom: std::marker::PhantomData<&'s ()>,
89}
90
91pub struct DatabaseQuery<'world, 'state, Q: DBQueryInfo<DbResource>, DbResource: DatabaseResource> {
92    // world and state will be needed later
93    db: Res<'world, DbResource>,
94    world: UnsafeWorldCell<'world>,
95    phantom2: std::marker::PhantomData<(&'state (), Q)>,
96}
97
98pub trait B {}
99
100pub trait A<B> {}
101
102// So query can be constructed by the system
103unsafe impl<'w, 's, I: DBQueryInfo<DbResource>, DbResource: DatabaseResource> SystemParam
104    for DatabaseQuery<'w, 's, I, DbResource>
105where
106    I: DBQueryInfo<DbResource> + 'static,
107{
108    type State = QueryFetchState<'static, 'static, DbResource>;
109
110    type Item<'world, 'state> = DatabaseQuery<'world, 'state, I, DbResource>;
111
112    fn init_state(
113        world: &mut World,
114        system_meta: &mut bevy_ecs::system::SystemMeta,
115    ) -> Self::State {
116        // https://github.com/chrisjuchem/bevy_mod_index/blob/15e9b4c9bbf26d4fc087ce056b07d1312464de2f/src/index.rs#L108
117        if !world.contains_resource::<DbResource>() {
118            world.init_resource::<DbResource>();
119        }
120
121        QueryFetchState {
122            db_state: <Res<'w, DbResource>>::init_state(world, system_meta),
123            phantom: std::marker::PhantomData,
124        }
125    }
126
127    unsafe fn get_param<'w2, 's2>(
128        state: &'s2 mut Self::State,
129        system_meta: &bevy_ecs::system::SystemMeta,
130        world: UnsafeWorldCell<'w2>,
131        change_tick: bevy_ecs::component::Tick,
132    ) -> Self::Item<'w2, 's2> {
133        let db_query = DatabaseQuery {
134            db: <Res<'w2, DbResource>>::get_param(
135                &mut state.db_state,
136                system_meta,
137                world,
138                change_tick,
139            ),
140            world,
141            phantom2: std::marker::PhantomData,
142        };
143
144        db_query
145    }
146}
147
148pub trait DatabaseEntityWithRequest: Sync + Send {
149    fn request(&self) -> &RequestId;
150    fn id(&self) -> &DatabaseEntityId;
151}
152
153impl<'w, 's, Q: DBQueryInfo<DbResource>, DbResource: DatabaseResource>
154    DatabaseQuery<'w, 's, Q, DbResource>
155{
156    pub async fn get<D: DatabaseEntityWithRequest>(
157        &self,
158        db_entity: &D,
159    ) -> Result<Q::ReadOnlyItem<'w>, ()> {
160        Q::get(self.db.as_ref(), self.world, db_entity).await
161    }
162
163    pub async fn get_mut<D: DatabaseEntityWithRequest>(
164        &self,
165        db_entity: &D,
166    ) -> Result<Q::Item<'w>, ()> {
167        Q::get_mut(self.db.as_ref(), self.world, db_entity).await
168    }
169
170    pub async fn update_component(
171        &self,
172        db_entity: &DatabaseEntity,
173        component: Q::ReadOnlyItem<'w>,
174    ) -> Result<(), ()> {
175        Q::update_component(self.db.as_ref(), self.world, db_entity, component).await
176    }
177
178    pub async fn insert_component(
179        &self,
180        db_entity: &DatabaseEntity,
181        component: Q::ReadOnlyItem<'w>,
182    ) -> Result<(), ()> {
183        Q::insert_component(self.db.as_ref(), self.world, db_entity, component).await
184    }
185
186    pub async fn load_components<R: ReturnSelector<'w>, CDQ>(
187        &self,
188        request: RequestId,
189        get_comp_from_db: CDQ,
190    ) -> Result<Vec<<R as ReturnSelector<'w>>::ReturnItem>, ()>
191    where
192        CDQ: CustomDatabaseQuery<DbResource, Q::DerefItem> + Send,
193    {
194        Q::load_components::<R, CDQ>(self.db.as_ref(), self.world, request, get_comp_from_db).await
195    }
196
197    pub async fn create(&self, component: Q::DerefItem, request: RequestId) -> Result<(), ()> {
198        Q::create(self.db.as_ref(), self.world, component, request).await
199    }
200
201    pub async fn update_or_insert_component(&self, entity: Entity) -> Result<(), ()> {
202        unsafe {
203            let mut q = self
204                .world
205                .world_mut()
206                .query::<(&DatabaseEntity, Q::WorldQuery<'w>)>();
207            let (db_entity, comp) = q.get(self.world.world(), entity).unwrap();
208
209            if db_entity.persisted.into() {
210                if db_entity.dirty {
211                    Q::update_component(self.db.as_ref(), self.world, db_entity, comp.into()).await
212                } else {
213                    Ok(())
214                }
215            } else {
216                Q::insert_component(self.db.as_ref(), self.world, db_entity, comp.into()).await
217            }
218        }
219    }
220
221    pub fn commit(&self, request: RequestId) -> Result<(), ()> {
222        self.db.commit_transaction(request);
223
224        Ok(())
225    }
226}
227
228#[async_trait]
229pub trait ComponentMapper {
230    type Component;
231    type Executor: Send;
232
233    async fn get<'c>(
234        e: &mut Self::Executor,
235        db_entity: &DatabaseEntityId,
236    ) -> Result<Self::Component, ()>;
237
238    async fn update_component<'c>(
239        tr: &mut Self::Executor,
240        db_entity: &DatabaseEntityId,
241        component: &Self::Component,
242    ) -> Result<(), ()>;
243
244    async fn insert_component<'c>(
245        tr: &mut Self::Executor,
246        db_entity: &DatabaseEntityId,
247        component: &Self::Component,
248    ) -> Result<(), ()>;
249}
250
251#[derive(Component)]
252pub struct NullComponent();
253
254impl RequestIdIndexInfo for NullComponent {}
255
256// To satisfy the type system when a DBQueryInfo is composed of other DBQueryInfos
257impl IndexInfo for NullComponent {
258    type Component = NullComponent;
259    type Value = RequestId;
260    type Storage = NoStorage<Self>;
261
262    fn value(_c: &Self::Component) -> Self::Value {
263        RequestId(generational_arena::Index::from_raw_parts(0, 0))
264    }
265}
266
267pub struct NullMapper;
268#[async_trait]
269impl ComponentMapper for NullMapper {
270    type Component = NullComponent;
271    type Executor = ();
272
273    async fn get<'c>(
274        _e: &mut Self::Executor,
275        _db_entity: &DatabaseEntityId,
276    ) -> Result<Self::Component, ()> {
277        unimplemented!()
278    }
279
280    async fn update_component<'c>(
281        _tr: &mut Self::Executor,
282        _db_entity: &DatabaseEntityId,
283        _component: &Self::Component,
284    ) -> Result<(), ()> {
285        unimplemented!()
286    }
287
288    async fn insert_component<'c>(
289        _tr: &mut Self::Executor,
290        _db_entity: &DatabaseEntityId,
291        _component: &Self::Component,
292    ) -> Result<(), ()> {
293        unimplemented!()
294    }
295}
296
297// Used to help speicfy whether the returned component is read only or mutable
298pub trait ReadMarker<DbResource: DatabaseResource>: DBQueryInfo<DbResource> {}
299
300// traits that are expected to be used in a tuple
301// &, &mut, Option
302pub trait TupleMarker<DbResource: DatabaseResource>: DBQueryInfo<DbResource> {}
303
304impl<T: DBQueryInfo<DbResource>, DbResource: DatabaseResource> TupleMarker<DbResource> for Option<T> where
305    for<'a> Option<<T as DBQueryInfo<DbResource>>::ReadOnlyItem<'a>>:
306        From<Option<ReadOnlyItem<'a, T, DbResource>>>
307{
308}
309
310#[async_trait]
311impl<T: DBQueryInfo<DbResource>, DbResource: DatabaseResource> DBQueryInfo<DbResource> for Option<T>
312where
313    for<'a> Option<<T as DBQueryInfo<DbResource>>::ReadOnlyItem<'a>>:
314        From<Option<ReadOnlyItem<'a, T, DbResource>>>,
315{
316    type DerefItem = Option<<T as DBQueryInfo<DbResource>>::DerefItem>;
317
318    // type Mapper = <T as DBQueryInfo<DbResource>>::Mapper;
319
320    type WorldQuery<'a> = Option<<T as DBQueryInfo<DbResource>>::WorldQuery<'a>>;
321
322    type ReadOnlyItem<'a> = Option<<T as DBQueryInfo<DbResource>>::ReadOnlyItem<'a>>;
323
324    type Item<'a> = Option<<T as DBQueryInfo<DbResource>>::Item<'a>>;
325
326    async fn get<'w, D: DatabaseEntityWithRequest>(
327        db: &DbResource,
328        world: UnsafeWorldCell<'w>,
329        db_entity: &D,
330    ) -> Result<Self::ReadOnlyItem<'w>, ()> {
331        T::get(db, world, db_entity).await.map(Some)
332    }
333
334    async fn get_mut<'w, D: DatabaseEntityWithRequest>(
335        db: &DbResource,
336        world: UnsafeWorldCell<'w>,
337        db_entity: &D,
338    ) -> Result<Self::Item<'w>, ()> {
339        T::get_mut(db, world, db_entity).await.map(Some)
340    }
341
342    async fn update_component<'w>(
343        db: &DbResource,
344        world: UnsafeWorldCell<'w>,
345        db_entity: &DatabaseEntity,
346        component: Self::ReadOnlyItem<'w>,
347    ) -> Result<(), ()> {
348        if let Some(component) = component {
349            T::update_component(db, world, db_entity, component).await
350        } else {
351            Ok(())
352        }
353    }
354
355    async fn insert_component<'w>(
356        db: &DbResource,
357        world: UnsafeWorldCell<'w>,
358        db_entity: &DatabaseEntity,
359        component: Self::ReadOnlyItem<'w>,
360    ) -> Result<(), ()> {
361        if let Some(component) = component {
362            T::insert_component(db, world, db_entity, component).await
363        } else {
364            Ok(())
365        }
366    }
367
368    async fn load_components<'w, R: ReturnSelector<'w>, CDQ>(
369        _db: &DbResource,
370        _world: UnsafeWorldCell<'w>,
371        _request: RequestId,
372        _get_comp_from_db: CDQ,
373    ) -> Result<Vec<<R as ReturnSelector<'w>>::ReturnItem>, ()>
374    where
375        CDQ: CustomDatabaseQuery<DbResource, Self::DerefItem> + Send,
376    {
377        Err(())
378    }
379
380    async fn create(
381        db: &DbResource,
382        world: UnsafeWorldCell<'_>,
383        component: Self::DerefItem,
384        request: RequestId,
385    ) -> Result<(), ()> {
386        if let Some(component) = component {
387            T::create(db, world, component, request).await
388        } else {
389            Ok(())
390        }
391    }
392}
393
394impl<'a, T: ComponentMapper, DbResource: DatabaseResource<Transaction = T::Executor>>
395    TupleMarker<DbResource> for &T
396where
397    <T as ComponentMapper>::Component: Component,
398{
399}
400
401impl<'a, T: ComponentMapper, DbResource: DatabaseResource<Transaction = T::Executor>>
402    ReadMarker<DbResource> for &T
403where
404    <T as ComponentMapper>::Component: Component,
405{
406}
407
408#[async_trait]
409impl<'b, T: ComponentMapper, DbResource: DatabaseResource<Transaction = T::Executor>>
410    DBQueryInfo<DbResource> for &T
411where
412    <T as ComponentMapper>::Component: Component,
413{
414    type DerefItem = <T as ComponentMapper>::Component;
415    // type Mapper = T;
416    type WorldQuery<'a> = &'a <T as ComponentMapper>::Component;
417
418    type ReadOnlyItem<'a> = &'a <T as ComponentMapper>::Component;
419
420    type Item<'a> = &'a <T as ComponentMapper>::Component;
421
422    async fn get<'w, D: DatabaseEntityWithRequest>(
423        db: &DbResource,
424        world: UnsafeWorldCell<'w>,
425        db_entity: &D,
426    ) -> Result<Self::ReadOnlyItem<'w>, ()> {
427        SingleComponentRetriever::<T, DbResource>::get(db, world, db_entity).await
428    }
429
430    async fn get_mut<'w, D: DatabaseEntityWithRequest>(
431        db: &DbResource,
432        world: UnsafeWorldCell<'w>,
433        db_entity: &D,
434    ) -> Result<Self::Item<'w>, ()> {
435        SingleComponentRetriever::<T, DbResource>::get(db, world, db_entity).await
436    }
437
438    async fn update_component<'w>(
439        db: &DbResource,
440        world: UnsafeWorldCell<'w>,
441        db_entity: &DatabaseEntity,
442        component: Self::ReadOnlyItem<'w>,
443    ) -> Result<(), ()> {
444        SingleComponentRetriever::<T, DbResource>::update_component(db, world, db_entity, component)
445            .await
446    }
447
448    async fn insert_component<'w>(
449        db: &DbResource,
450        world: UnsafeWorldCell<'w>,
451        db_entity: &DatabaseEntity,
452        component: Self::ReadOnlyItem<'w>,
453    ) -> Result<(), ()> {
454        SingleComponentRetriever::<T, DbResource>::insert_component(db, world, db_entity, component)
455            .await
456    }
457
458    async fn load_components<'w, R: ReturnSelector<'w>, CDQ>(
459        db: &DbResource,
460        world: UnsafeWorldCell<'w>,
461        request: RequestId,
462        get_comp_from_db: CDQ,
463    ) -> Result<Vec<<R as ReturnSelector<'w>>::ReturnItem>, ()>
464    where
465        CDQ: CustomDatabaseQuery<DbResource, Self::DerefItem> + Send,
466    {
467        SingleComponentRetriever::<T, DbResource>::load_components::<R, CDQ>(
468            db,
469            world,
470            request,
471            get_comp_from_db,
472        )
473        .await
474    }
475
476    async fn create(
477        db: &DbResource,
478        world: UnsafeWorldCell<'_>,
479        component: Self::DerefItem,
480        request: RequestId,
481    ) -> Result<(), ()> {
482        SingleComponentRetriever::<T, DbResource>::create(db, world, component, request).await
483    }
484}
485
486impl<'b, T: ComponentMapper, DbResource: DatabaseResource<Transaction = T::Executor>>
487    ReadMarker<DbResource> for &mut T
488where
489    <T as ComponentMapper>::Component: Component,
490{
491}
492
493#[async_trait]
494impl<'b, T: ComponentMapper, DbResource: DatabaseResource<Transaction = T::Executor>>
495    DBQueryInfo<DbResource> for &mut T
496where
497    <T as ComponentMapper>::Component: Component,
498{
499    type DerefItem = <T as ComponentMapper>::Component;
500    // type Mapper = T;
501    type WorldQuery<'a> = &'a mut <T as ComponentMapper>::Component;
502
503    type ReadOnlyItem<'a> = &'a <T as ComponentMapper>::Component;
504
505    type Item<'a> = Mut<'a, <T as ComponentMapper>::Component>;
506
507    async fn get<'w, D: DatabaseEntityWithRequest>(
508        db: &DbResource,
509        world: UnsafeWorldCell<'w>,
510        db_entity: &D,
511    ) -> Result<Self::ReadOnlyItem<'w>, ()> {
512        SingleComponentRetriever::<T, DbResource>::get(db, world, db_entity).await
513    }
514
515    async fn get_mut<'w, D: DatabaseEntityWithRequest>(
516        db: &DbResource,
517        world: UnsafeWorldCell<'w>,
518        db_entity: &D,
519    ) -> Result<Self::Item<'w>, ()> {
520        SingleComponentRetriever::<T, DbResource>::get_mut(db, world, db_entity).await
521    }
522
523    async fn update_component<'w>(
524        db: &DbResource,
525        world: UnsafeWorldCell<'w>,
526        db_entity: &DatabaseEntity,
527        component: Self::ReadOnlyItem<'w>,
528    ) -> Result<(), ()> {
529        SingleComponentRetriever::<T, DbResource>::update_component(db, world, db_entity, component)
530            .await
531    }
532
533    async fn insert_component<'w>(
534        db: &DbResource,
535        world: UnsafeWorldCell<'w>,
536        db_entity: &DatabaseEntity,
537        component: Self::ReadOnlyItem<'w>,
538    ) -> Result<(), ()> {
539        SingleComponentRetriever::<T, DbResource>::insert_component(db, world, db_entity, component)
540            .await
541    }
542
543    async fn load_components<'w, R: ReturnSelector<'w>, CDQ>(
544        db: &DbResource,
545        world: UnsafeWorldCell<'w>,
546        request: RequestId,
547        get_comp_from_db: CDQ,
548    ) -> Result<Vec<<R as ReturnSelector<'w>>::ReturnItem>, ()>
549    where
550        CDQ: CustomDatabaseQuery<DbResource, Self::DerefItem> + Send,
551    {
552        SingleComponentRetriever::<T, DbResource>::load_components::<R, CDQ>(
553            db,
554            world,
555            request,
556            get_comp_from_db,
557        )
558        .await
559    }
560
561    async fn create(
562        db: &DbResource,
563        world: UnsafeWorldCell<'_>,
564        component: Self::DerefItem,
565        request: RequestId,
566    ) -> Result<(), ()> {
567        SingleComponentRetriever::<T, DbResource>::create(db, world, component, request).await
568    }
569}
570
571/*
572 Like above but for tuples of DBQueryInfo
573 The macro takes a tuple of DBQueryInfo and creates
574 a new DBQueryInfo that returns a tuple of the items
575*/
576macro_rules! simple_composition_of_db_queries {
577    ( $( $name:ident )* ) => {
578        // include this Head Z so that can specify that the rest of them have the same database
579        // looks quite ugly but it works
580        #[async_trait]
581        impl<Z: TupleMarker<DbResource>, $($name: TupleMarker<DbResource>, )* DbResource: DatabaseResource> DBQueryInfo<DbResource> for (Z, $($name,)*)
582            where
583                for<'a>
584                    (
585                        <Z as DBQueryInfo<DbResource>>::ReadOnlyItem<'a>,
586                        $(<$name as DBQueryInfo<DbResource>>::ReadOnlyItem<'a>, )*
587                    )
588                        :
589                        From<(
590                            ReadOnlyItem<'a, Z, DbResource>,
591                            $(ReadOnlyItem<'a, $name, DbResource>,)*
592                        )>
593        {
594            type DerefItem = (Z::DerefItem, $($name::DerefItem,)*);
595            // type Mapper = NullMapper;
596            type WorldQuery<'a> = (Z::WorldQuery<'a>, $($name::WorldQuery<'a>,)*);
597
598            type ReadOnlyItem<'a> = (Z::ReadOnlyItem<'a>, $($name::ReadOnlyItem<'a>,)*);
599
600            type Item<'a> = (Z::Item<'a>, $($name::Item<'a>,)*);
601
602            async fn get<'w, DbEntity: DatabaseEntityWithRequest>(db: &DbResource, world: UnsafeWorldCell<'w>, db_entity: &DbEntity) -> Result<Self::ReadOnlyItem<'w>, ()> {
603            //returns a tuple of all the gets
604                Ok((
605                    Z::get(db, world, db_entity).await?,
606                    $(
607                    {
608                        $name::get(db, world, db_entity).await?
609                    },
610                )*))
611            }
612
613            async fn get_mut<'w, DbEntity: DatabaseEntityWithRequest>(db: &DbResource, world: UnsafeWorldCell<'w>, db_entity: &DbEntity) -> Result<Self::Item<'w>, ()> {
614                //returns a tuple of all the gets
615                Ok((
616                    Z::get_mut(db, world, db_entity).await?,
617
618                    $(
619                    {
620                        $name::get_mut(db, world, db_entity).await?
621                    },
622                )*))
623            }
624
625            async fn update_component<'w>(db: &DbResource, world: UnsafeWorldCell<'w>, db_entity: &DatabaseEntity, component: Self::ReadOnlyItem<'w>) -> Result<(), ()> {
626
627                let (z, $(lower!($name),)*) = component;
628
629                Z::update_component(db, world, db_entity, z).await?;
630                $($name::update_component(db, world, db_entity, lower!($name)).await?;)*
631
632                Ok(())
633            }
634
635            async fn insert_component<'w>(db: &DbResource, world: UnsafeWorldCell<'w>, db_entity: &DatabaseEntity, component: Self::ReadOnlyItem<'w>) -> Result<(), ()> {
636                let (z, $(lower!($name),)*) = component;
637
638                Z::insert_component(db, world, db_entity, z).await?;
639                $($name::insert_component(db, world, db_entity, lower!($name)).await?;)*
640
641                Ok(())
642            }
643
644            // CODE SMELL: probably should split up interface to avoid this method
645            async fn load_components<'w, R : ReturnSelector<'w>, CDQ>(
646                _db: &DbResource,
647                _world: UnsafeWorldCell<'w>,
648                _request: RequestId,
649                _get_comp_from_db: CDQ
650            ) -> Result<Vec<<R as ReturnSelector<'w>>::ReturnItem>, ()>
651            where
652                CDQ: CustomDatabaseQuery<DbResource, Self::DerefItem> + Send
653            {
654                unimplemented!()
655            }
656
657            async fn create(db: &DbResource, world: UnsafeWorldCell<'_>, component: Self::DerefItem, request: RequestId) -> Result<(), ()> {
658                let (z, $(lower!($name),)*) = component;
659
660                Z::create(db, world, z, request).await?;
661                $($name::create(db, world, lower!($name), request).await?;)*
662
663                Ok(())
664            }
665        }
666    }
667}
668
669// Create a simple composition of DBQueryInfo for tuples of length 1 to 10
670// Allows DBQueryInfo to be composed of other DBQueryInfo
671// eg. DBQuery<(User, Item)>
672// simple_composition_of_db_queries!{}
673simple_composition_of_db_queries! {A}
674simple_composition_of_db_queries! {A B}
675simple_composition_of_db_queries! {A B C}
676simple_composition_of_db_queries! {A B C D}
677// simple_composition_of_db_queries!{A B C D E}
678// simple_composition_of_db_queries!{A B C D E F}
679// simple_composition_of_db_queries!{A B C D E F G}
680// simple_composition_of_db_queries!{A B C D E F G H}
681// simple_composition_of_db_queries!{A B C D E F G H I}
682// simple_composition_of_db_queries!{A B C D E F G H I J}
683// simple_composition_of_db_queries!{A B C D E F G H I J K}
684
685#[derive(Default)]
686pub struct SingleComponentRetriever<Mapper, DbResource> {
687    phantom: std::marker::PhantomData<(Mapper, DbResource)>,
688}
689
690impl<
691        'a,
692        MyMapper: ComponentMapper,
693        DbResource: DatabaseResource<Transaction = <MyMapper as ComponentMapper>::Executor>,
694    > SingleComponentRetriever<MyMapper, DbResource>
695where
696    <MyMapper as ComponentMapper>::Component: Component,
697{
698    pub async fn get_internal<D: DatabaseEntityWithRequest>(
699        db: &DbResource,
700        world: UnsafeWorldCell<'_>,
701        db_entity: &D,
702        component_preloaded: Option<<MyMapper as ComponentMapper>::Component>,
703    ) -> Entity {
704        let mut tr = db.get_transaction(*db_entity.request());
705
706        let db_entity_id = *db_entity.id();
707
708        let mut reader = IntoSystem::into_system(
709            move |mut index: Index<DatabaseEntityIndex>| -> HashSet<Entity> {
710                index.lookup(&db_entity_id)
711            },
712        );
713
714        let entity_set: HashSet<Entity> = unsafe {
715            reader.initialize(world.world_mut());
716            reader.run((), world.world_mut())
717        };
718
719        match entity_set.iter().next() {
720            // Entity has been read into memory before
721            Some(entity) => {
722                match unsafe {
723                    world
724                        .world_mut()
725                        .get::<<MyMapper as ComponentMapper>::Component>(*entity)
726                } {
727                    // Entity also already has the desired component
728                    Some(_) => *entity,
729                    // Entity does not have the desired component (Load from database)
730                    None => {
731                        let db_component = match component_preloaded {
732                            Some(component) => component,
733                            None => MyMapper::get(&mut tr, db_entity.id()).await.unwrap(),
734                        };
735                        // write the component to the entity
736                        unsafe {
737                            let w = world.world_mut();
738                            w.entity_mut(*entity).insert(db_component);
739                            *entity
740                        }
741                    }
742                }
743            }
744            // Entity not found in world
745            None => {
746                let component = match component_preloaded {
747                    Some(component) => component,
748                    None => MyMapper::get(&mut tr, db_entity.id()).await.unwrap(),
749                };
750                unsafe {
751                    let w = world.world_mut();
752                    let entity = w
753                        .spawn((
754                            DatabaseEntity {
755                                id: *db_entity.id(),
756                                persisted: true.into(),
757                                dirty: false,
758                                request: *db_entity.request(),
759                            },
760                            component,
761                        ))
762                        .id();
763                    entity
764                }
765            }
766        }
767    }
768
769    pub async fn load_entities_for_components<CDQ>(
770        db: &DbResource,
771        world: UnsafeWorldCell<'_>,
772        request: RequestId,
773        get_comp_from_db: CDQ,
774    ) -> Result<Vec<Entity>, ()>
775    where
776        CDQ: CustomDatabaseQuery<DbResource, <MyMapper as ComponentMapper>::Component>,
777    {
778        let components = {
779            let tr = db.get_transaction(request);
780
781            get_comp_from_db.query(tr).await?
782        };
783
784        let mut entities = vec![];
785
786        for (db_entity, component) in components {
787            entities.push(Self::get_internal(db, world, &db_entity, Some(component)).await);
788        }
789
790        Ok(entities)
791    }
792
793    pub async fn update_or_insert_component(
794        db: &DbResource,
795        world: UnsafeWorldCell<'_>,
796        db_entity: &DatabaseEntity,
797        component: &MyMapper::Component,
798    ) -> Result<(), ()> {
799        if db_entity.persisted.into() {
800            if db_entity.dirty {
801                Self::update_component(db, world, db_entity, component).await
802            } else {
803                Ok(())
804            }
805        } else {
806            Self::insert_component(db, world, db_entity, component).await
807        }
808    }
809}
810
811impl<'w, C: WorldQuery> ReturnSelector<'w> for C {
812    type ReturnItem = C::Item<'w>;
813
814    fn load_components_from_entities(
815        world: UnsafeWorldCell<'w>,
816        entities: Vec<Entity>,
817    ) -> Vec<Self::ReturnItem> {
818        entities
819            .into_iter()
820            .map(|entity| unsafe {
821                let world = world.world_mut();
822                let mut q = world.query::<C>();
823                q.get_mut(world, entity).unwrap()
824            })
825            .collect()
826    }
827}
828
829#[async_trait]
830impl<
831        'b,
832        MyMapper: ComponentMapper,
833        DbResource: DatabaseResource<Transaction = <MyMapper as ComponentMapper>::Executor>,
834    > DBQueryInfo<DbResource> for SingleComponentRetriever<MyMapper, DbResource>
835where
836    <MyMapper as ComponentMapper>::Component: Component,
837{
838    type DerefItem = MyMapper::Component;
839    type WorldQuery<'a> = &'a mut MyMapper::Component;
840    type ReadOnlyItem<'a> = &'a MyMapper::Component;
841    type Item<'a> = Mut<'a, MyMapper::Component>;
842
843    async fn get<'w, D: DatabaseEntityWithRequest>(
844        db: &DbResource,
845        world: UnsafeWorldCell<'w>,
846        db_entity: &D,
847    ) -> Result<Self::ReadOnlyItem<'w>, ()> {
848        let entity = Self::get_internal(db, world, db_entity, None).await;
849
850        unsafe {
851            Ok(world
852                .world()
853                .get::<<MyMapper as ComponentMapper>::Component>(entity)
854                .unwrap())
855        }
856    }
857
858    async fn get_mut<'w, D: DatabaseEntityWithRequest>(
859        db: &DbResource,
860        world: UnsafeWorldCell<'w>,
861        db_entity: &D,
862    ) -> Result<Self::Item<'w>, ()> {
863        let entity = Self::get_internal(db, world, db_entity, None).await;
864
865        unsafe {
866            // set entity to dirty
867            world
868                .get_entity(entity)
869                .unwrap()
870                .get_mut::<DatabaseEntity>()
871                .unwrap()
872                .dirty = true;
873
874            Ok(world
875                .get_entity(entity)
876                .unwrap()
877                .get_mut::<<MyMapper as ComponentMapper>::Component>()
878                .unwrap())
879        }
880    }
881
882    async fn update_component<'w>(
883        db: &DbResource,
884        _world: UnsafeWorldCell<'w>,
885        db_entity: &DatabaseEntity,
886        component: ReadOnlyItem<'w, Self, DbResource>,
887    ) -> Result<(), ()> {
888        let mut tr = db.get_transaction(*db_entity.request());
889
890        MyMapper::update_component(&mut tr, db_entity.id(), component).await
891    }
892
893    async fn insert_component<'w>(
894        db: &DbResource,
895        _world: UnsafeWorldCell<'w>,
896        db_entity: &DatabaseEntity,
897        component: ReadOnlyItem<'w, Self, DbResource>,
898    ) -> Result<(), ()> {
899        let mut tr = db.get_transaction(*db_entity.request());
900        MyMapper::insert_component(&mut tr, db_entity.id(), component).await
901    }
902
903    async fn load_components<'w, R: ReturnSelector<'w>, CDQ>(
904        db: &DbResource,
905        world: UnsafeWorldCell<'w>,
906        request: RequestId,
907        get_comp_from_db: CDQ,
908    ) -> Result<Vec<<R as ReturnSelector<'w>>::ReturnItem>, ()>
909    where
910        CDQ: CustomDatabaseQuery<DbResource, Self::DerefItem> + Send,
911    {
912        let entities =
913            Self::load_entities_for_components::<CDQ>(db, world, request, get_comp_from_db).await?;
914
915        let components = R::load_components_from_entities(world, entities);
916        Ok(components)
917    }
918
919    async fn create(
920        db: &DbResource,
921        world: UnsafeWorldCell<'_>,
922        component: Self::DerefItem,
923        request: RequestId,
924    ) -> Result<(), ()> {
925        unsafe {
926            let w = world.world_mut();
927            w.spawn((
928                component,
929                DatabaseEntity {
930                    id: db.get_key(),
931                    persisted: false.into(),
932                    dirty: false,
933                    request,
934                },
935            ));
936        }
937
938        Ok(())
939    }
940}
bevy_erm_core/database_query.rs

bevy_erm_core/
database_query.rs
