1use std::any::Any;
49use std::collections::HashMap;
50use std::future::Future;
51use std::marker::PhantomData;
52use std::pin::Pin;
53
54use super::pool::DbPool;
55use super::query::QueryBuilder;
56use super::repository::placeholder;
57use super::{DbError, Model, ModelRow, Value};
58use crate::pagination::Page;
59
60pub const IN_BATCH: usize = 500;
62
63const OWNER_KEY: &str = "rws_owner_key";
66
67#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
71#[non_exhaustive]
72pub enum RelationKind {
73 HasMany,
76 HasOne,
78 BelongsTo,
80 BelongsToMany,
82}
83
84pub struct Relation<O, R> {
99 kind: RelationKind,
100 foreign_key: &'static str,
104 join: Option<(&'static str, &'static str)>,
107 _models: PhantomData<fn() -> (O, R)>,
108}
109
110impl<O, R> Clone for Relation<O, R> {
111 fn clone(&self) -> Self {
112 *self
113 }
114}
115
116impl<O, R> Copy for Relation<O, R> {}
117
118impl<O, R> std::fmt::Debug for Relation<O, R> {
119 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120 f.debug_struct("Relation")
121 .field("kind", &self.kind)
122 .field("foreign_key", &self.foreign_key)
123 .field("join", &self.join)
124 .finish()
125 }
126}
127
128impl<O: Model, R: Model> Relation<O, R> {
129 pub fn has_many(foreign_key: &'static str) -> Self {
132 Relation { kind: RelationKind::HasMany, foreign_key, join: None, _models: PhantomData }
133 }
134
135 pub fn has_one(foreign_key: &'static str) -> Self {
138 Relation { kind: RelationKind::HasOne, foreign_key, join: None, _models: PhantomData }
139 }
140
141 pub fn belongs_to(foreign_key: &'static str) -> Self {
145 Relation { kind: RelationKind::BelongsTo, foreign_key, join: None, _models: PhantomData }
146 }
147
148 pub fn belongs_to_many(join_table: &'static str, owner_key: &'static str, related_key: &'static str) -> Self {
154 Relation { kind: RelationKind::BelongsToMany, foreign_key: owner_key, join: Some((join_table, related_key)), _models: PhantomData }
155 }
156
157 pub fn kind(&self) -> RelationKind {
159 self.kind
160 }
161
162 pub fn foreign_key(&self) -> &'static str {
166 self.foreign_key
167 }
168
169 pub fn join_table(&self) -> Option<(&'static str, &'static str)> {
171 self.join
172 }
173
174 fn key(&self) -> String {
176 format!("{:?}:{}:{}:{:?}", self.kind, R::table_name(), self.foreign_key, self.join)
177 }
178
179 pub async fn load(&self, pool: &DbPool, owner: &O) -> Result<Vec<R>, DbError>
181 where
182 R: Clone,
183 {
184 Ok(self.load_for(pool, std::slice::from_ref(owner)).await?.pop().unwrap_or_default())
185 }
186
187 pub async fn load_for(&self, pool: &DbPool, owners: &[O]) -> Result<Vec<Vec<R>>, DbError>
192 where
193 R: Clone,
194 {
195 let owner_keys: Vec<Value> = match self.kind {
197 RelationKind::BelongsTo => {
198 let mut keys = Vec::with_capacity(owners.len());
199 for owner in owners {
200 let value = owner.to_values().into_iter().find(|(c, _)| c.eq_ignore_ascii_case(self.foreign_key)).map(|(_, v)| v);
201 keys.push(value.ok_or_else(|| {
202 DbError::new(format!("belongs_to: {} has no column {}", O::table_name(), self.foreign_key))
203 })?);
204 }
205 keys
206 }
207 _ => owners.iter().map(|o| o.primary_key_value()).collect(),
208 };
209 let distinct = distinct_keys(&owner_keys);
210 let pk = R::primary_key_name();
211 let (rows, match_column) = match self.kind {
212 RelationKind::HasMany | RelationKind::HasOne => {
213 let live = live_and::<R>("");
214 let rows = fetch_in(pool, &distinct, |list| {
215 format!("SELECT * FROM {} WHERE {} IN ({list}){live} ORDER BY {pk}", R::table_name(), self.foreign_key)
216 })
217 .await?;
218 (rows, self.foreign_key)
219 }
220 RelationKind::BelongsTo => {
221 let live = live_and::<R>("");
222 let rows = fetch_in(pool, &distinct, |list| {
223 format!("SELECT * FROM {} WHERE {pk} IN ({list}){live} ORDER BY {pk}", R::table_name())
224 })
225 .await?;
226 (rows, pk)
227 }
228 RelationKind::BelongsToMany => {
229 let (join, related_key) = self.join.unwrap_or(("", ""));
230 let live = live_and::<R>("r.");
231 let rows = fetch_in(pool, &distinct, |list| {
232 format!(
233 "SELECT r.*, j.{owner} AS {OWNER_KEY} FROM {} r INNER JOIN {join} j ON j.{related_key} = r.{pk} \
234 WHERE j.{owner} IN ({list}){live} ORDER BY r.{pk}",
235 R::table_name(),
236 owner = self.foreign_key,
237 )
238 })
239 .await?;
240 (rows, OWNER_KEY)
241 }
242 };
243 let mut by_key: HashMap<String, Vec<R>> = HashMap::new();
244 for row in &rows {
245 let Some(key) = column(row, match_column).and_then(key_of) else { continue };
246 by_key.entry(key).or_default().push(R::from_row(row)?);
247 }
248 let one = matches!(self.kind, RelationKind::HasOne | RelationKind::BelongsTo);
249 Ok(owner_keys
250 .iter()
251 .map(|k| {
252 let mut related = key_of(k.clone()).and_then(|k| by_key.get(&k)).cloned().unwrap_or_default();
253 if one {
254 related.truncate(1);
255 }
256 related
257 })
258 .collect())
259 }
260}
261
262fn live_and<T: Model>(prefix: &str) -> String {
264 T::soft_delete_column().map(|c| format!(" AND {prefix}{c} IS NULL")).unwrap_or_default()
265}
266
267fn column(row: &ModelRow, name: &str) -> Option<Value> {
269 row.columns.iter().find(|(c, _)| c.eq_ignore_ascii_case(name)).map(|(_, v)| v.clone())
270}
271
272pub(crate) fn key_of(value: Value) -> Option<String> {
276 match value {
277 Value::Null => None,
278 Value::Bool(b) => Some(format!("i{}", b as i64)),
279 Value::Int(n) => Some(format!("i{n}")),
280 Value::Float(f) => Some(format!("f{f}")),
281 Value::Text(s) => Some(format!("t{s}")),
282 Value::Bytes(b) => Some(format!("b{}", b.iter().map(|x| format!("{x:02x}")).collect::<String>())),
283 }
284}
285
286pub(crate) fn distinct_keys(values: &[Value]) -> Vec<Value> {
288 let mut seen = std::collections::HashSet::new();
289 values.iter().filter(|v| key_of((*v).clone()).is_some_and(|k| seen.insert(k))).cloned().collect()
290}
291
292async fn fetch_in(pool: &DbPool, keys: &[Value], sql: impl Fn(String) -> String) -> Result<Vec<ModelRow>, DbError> {
295 let backend = pool.backend();
296 let mut rows = Vec::new();
297 for batch in keys.chunks(IN_BATCH) {
298 let list = (1..=batch.len()).map(|i| placeholder(backend, i)).collect::<Vec<_>>().join(", ");
299 rows.extend(pool.query_rows(&sql(list), batch).await?);
300 }
301 Ok(rows)
302}
303
304pub struct Loaded<T> {
309 model: T,
310 relations: Vec<(String, Box<dyn Any + Send + Sync>)>,
311}
312
313impl<T: Model> Loaded<T> {
314 pub fn model(&self) -> &T {
316 &self.model
317 }
318
319 pub fn into_inner(self) -> T {
321 self.model
322 }
323
324 pub fn related<R: Model + 'static>(&self, relation: &Relation<T, R>) -> &[R] {
327 self.get(relation).map(Vec::as_slice).unwrap_or(&[])
328 }
329
330 pub fn one<R: Model + 'static>(&self, relation: &Relation<T, R>) -> Option<&R> {
332 self.related(relation).first()
333 }
334
335 pub fn related_loaded<R: Model + 'static>(&self, relation: &Relation<T, R>) -> &[Loaded<R>] {
339 let key = nested_key(&relation.key());
340 self.relations.iter().find(|(k, _)| *k == key).and_then(|(_, v)| v.downcast_ref::<Vec<Loaded<R>>>()).map(Vec::as_slice).unwrap_or(&[])
341 }
342
343 pub fn is_loaded<R: Model + 'static>(&self, relation: &Relation<T, R>) -> bool {
345 self.get(relation).is_some()
346 }
347
348 pub fn take<R: Model + 'static>(&mut self, relation: &Relation<T, R>) -> Vec<R> {
350 let key = relation.key();
351 match self.relations.iter().position(|(k, _)| *k == key) {
352 Some(i) => self.relations.remove(i).1.downcast::<Vec<R>>().map(|b| *b).unwrap_or_default(),
353 None => Vec::new(),
354 }
355 }
356
357 fn get<R: Model + 'static>(&self, relation: &Relation<T, R>) -> Option<&Vec<R>> {
358 let key = relation.key();
359 self.relations.iter().find(|(k, _)| *k == key).and_then(|(_, v)| v.downcast_ref::<Vec<R>>())
360 }
361}
362
363impl<T> std::ops::Deref for Loaded<T> {
364 type Target = T;
365 fn deref(&self) -> &T {
366 &self.model
367 }
368}
369
370impl<T: std::fmt::Debug> std::fmt::Debug for Loaded<T> {
371 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
372 f.debug_struct("Loaded")
373 .field("model", &self.model)
374 .field("relations", &self.relations.iter().map(|(k, _)| k.as_str()).collect::<Vec<_>>())
375 .finish()
376 }
377}
378
379type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
382type Erased = Box<dyn Any + Send + Sync>;
383type PerOwner = Vec<Vec<(String, Erased)>>;
385
386trait Loader<O>: Send + Sync {
389 fn key(&self) -> String;
390 fn load<'a>(&'a self, pool: &'a DbPool, owners: &'a [O]) -> BoxFuture<'a, Result<PerOwner, DbError>>;
391}
392
393impl<O: Model + 'static, R: Model + Clone + 'static> Loader<O> for Relation<O, R> {
394 fn key(&self) -> String {
395 Relation::key(self)
396 }
397
398 fn load<'a>(&'a self, pool: &'a DbPool, owners: &'a [O]) -> BoxFuture<'a, Result<PerOwner, DbError>> {
399 Box::pin(async move {
400 let key = Relation::key(self);
401 let related = self.load_for(pool, owners).await?;
402 Ok(related.into_iter().map(|r| vec![(key.clone(), Box::new(r) as Erased)]).collect())
403 })
404 }
405}
406
407fn nested_key(key: &str) -> String {
409 format!("nested:{key}")
410}
411
412struct NestedLoader<O, R: Model> {
414 relation: Relation<O, R>,
415 loaders: Vec<Box<dyn Loader<R>>>,
416}
417
418impl<O: Model + 'static, R: Model + Clone + 'static> Loader<O> for NestedLoader<O, R> {
419 fn key(&self) -> String {
420 nested_key(&Relation::key(&self.relation))
421 }
422
423 fn load<'a>(&'a self, pool: &'a DbPool, owners: &'a [O]) -> BoxFuture<'a, Result<PerOwner, DbError>> {
424 Box::pin(async move {
425 let key = Relation::key(&self.relation);
426 let per_owner = self.relation.load_for(pool, owners).await?;
427 let counts: Vec<usize> = per_owner.iter().map(Vec::len).collect();
430 let flat: Vec<R> = per_owner.iter().flatten().cloned().collect();
431 let mut loaded = attach(pool, flat, &self.loaders).await?.into_iter();
432 let mut out = Vec::with_capacity(owners.len());
433 for (plain, n) in per_owner.into_iter().zip(counts) {
434 let nested: Vec<Loaded<R>> = loaded.by_ref().take(n).collect();
435 out.push(vec![(key.clone(), Box::new(plain) as Erased), (nested_key(&key), Box::new(nested) as Erased)]);
436 }
437 Ok(out)
438 })
439 }
440}
441
442pub struct Nested<R: Model> {
445 loaders: Vec<Box<dyn Loader<R>>>,
446}
447
448impl<R: Model + 'static> Nested<R> {
449 pub fn with<S: Model + Clone + 'static>(mut self, relation: Relation<R, S>) -> Self {
451 let key = relation.key();
452 if !self.loaders.iter().any(|l| l.key() == key) {
453 self.loaders.push(Box::new(relation));
454 }
455 self
456 }
457
458 pub fn with_nested<S: Model + Clone + 'static>(mut self, relation: Relation<R, S>, nested: impl FnOnce(Nested<S>) -> Nested<S>) -> Self {
460 let loaders = nested(Nested { loaders: Vec::new() }).loaders;
461 let loader = NestedLoader { relation, loaders };
462 let key = Loader::<R>::key(&loader);
463 if !self.loaders.iter().any(|l| l.key() == key) {
464 self.loaders.push(Box::new(loader));
465 }
466 self
467 }
468}
469
470pub struct EagerQuery<'a, T: Model> {
474 query: QueryBuilder<'a, T>,
475 loaders: Vec<Box<dyn Loader<T>>>,
476}
477
478impl<'a, T: Model + 'static> EagerQuery<'a, T> {
479 pub(crate) fn new(query: QueryBuilder<'a, T>) -> Self {
480 EagerQuery { query, loaders: Vec::new() }
481 }
482
483 pub fn with<R: Model + Clone + 'static>(mut self, relation: Relation<T, R>) -> Self {
485 let key = relation.key();
486 if !self.loaders.iter().any(|l| l.key() == key) {
487 self.loaders.push(Box::new(relation));
488 }
489 self
490 }
491
492 pub fn with_nested<R: Model + Clone + 'static>(mut self, relation: Relation<T, R>, nested: impl FnOnce(Nested<R>) -> Nested<R>) -> Self {
497 let loaders = nested(Nested { loaders: Vec::new() }).loaders;
498 let loader = NestedLoader { relation, loaders };
499 let key = Loader::<T>::key(&loader);
500 if !self.loaders.iter().any(|l| l.key() == key) {
501 self.loaders.push(Box::new(loader));
502 }
503 self
504 }
505
506 pub async fn fetch_all(self) -> Result<Vec<Loaded<T>>, DbError> {
509 let pool = self.query.pool();
510 let models = self.query.fetch_all().await?;
511 attach(pool, models, &self.loaders).await
512 }
513
514 pub async fn fetch_one(self) -> Result<Option<Loaded<T>>, DbError> {
516 let pool = self.query.pool();
517 let model = self.query.fetch_one().await?;
518 Ok(attach(pool, model.into_iter().collect(), &self.loaders).await?.pop())
519 }
520
521 pub async fn paginate(self, page: u64, per_page: u64) -> Result<Page<Loaded<T>>, DbError> {
523 let pool = self.query.pool();
524 let found = self.query.paginate(page, per_page).await?;
525 let (page, per_page, total) = (found.page, found.per_page, found.total_items);
526 let items = attach(pool, found.items, &self.loaders).await?;
527 Ok(Page::new(items, page, per_page, total))
528 }
529}
530
531async fn attach<T: Model>(pool: &DbPool, models: Vec<T>, loaders: &[Box<dyn Loader<T>>]) -> Result<Vec<Loaded<T>>, DbError> {
532 let mut relations: Vec<Vec<(String, Erased)>> = models.iter().map(|_| Vec::with_capacity(loaders.len())).collect();
533 if !models.is_empty() {
534 for loader in loaders {
535 for (slot, entries) in relations.iter_mut().zip(loader.load(pool, &models).await?) {
536 slot.extend(entries);
537 }
538 }
539 }
540 Ok(models.into_iter().zip(relations).map(|(model, relations)| Loaded { model, relations }).collect())
541}