siderite_orm/
db.rs

1//! [`Db`]: the database handle application code passes around.
2//!
3//! A `Db` is a cheap clone that points either at a connection pool or at an
4//! open transaction. Every ORM entry point takes `&Db`, so the same code runs
5//! inside or outside a transaction:
6//!
7//! ```ignore
8//! db.transaction(|tx| async move {
9//!     let user = User::objects(&tx).get(User::id.eq(1)).await?;
10//!     Post::objects(&tx).filter(Post::author.eq(user.id)).delete().await?;
11//!     Ok::<_, OrmError>(())
12//! }).await?;
13//! ```
14//!
15//! Inside the closure use `tx`, not the outer `db`: the outer handle would
16//! take a *second* connection, which deadlocks on a one-connection pool such
17//! as `sqlite::memory:`.
18
19use crate::backend::{Backend, ExecResult, QueryResult};
20use crate::capabilities::BackendCapabilities;
21use crate::error::{OrmError, QueryError};
22use crate::model::{Model, ModelMeta};
23use crate::plan::{PlanOrigin, QueryPlan};
24use crate::queryset::QuerySet;
25use crate::router::DatabaseRouter;
26use crate::signals::Signals;
27use crate::value::Value;
28use crate::write::WritePlan;
29use std::collections::BTreeMap;
30use std::future::Future;
31use std::sync::Arc;
32use std::time::Instant;
33use tracing::Instrument;
34
35mod transaction;
36use transaction::{ConnState, ExecutorRef, TxState};
37
38#[derive(Clone)]
39enum Target {
40    Pool(Arc<dyn Backend>),
41    Tx(Arc<TxState>),
42    Conn(Arc<ConnState>),
43}
44
45/// Handle to a database: a pool, or an open transaction.
46#[derive(Clone)]
47pub struct Db {
48    target: Target,
49    signals: Signals,
50}
51
52impl std::fmt::Debug for Db {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        let kind = self.capabilities().kind;
55        f.debug_struct("Db")
56            .field("backend", &kind)
57            .field("in_transaction", &self.in_transaction())
58            .finish()
59    }
60}
61
62impl Db {
63    /// Wrap a backend (connection pool).
64    pub fn new(backend: impl Backend) -> Self {
65        Self::from_arc(Arc::new(backend))
66    }
67
68    /// Wrap a shared backend.
69    pub fn from_arc(backend: Arc<dyn Backend>) -> Self {
70        Self {
71            target: Target::Pool(backend),
72            signals: Signals::default(),
73        }
74    }
75
76    /// Attach the [`Signals`] registry model operations dispatch to.
77    ///
78    /// The registry is shared by clones of this handle and by the transaction
79    /// handles opened from it.
80    #[must_use]
81    pub fn with_signals(mut self, signals: Signals) -> Self {
82        self.signals = signals;
83        self
84    }
85
86    /// The signal registry of this handle.
87    pub fn signals(&self) -> &Signals {
88        &self.signals
89    }
90
91    /// Whether `self` and `other` talk to the same database: two handles of
92    /// one pool, or a transaction and the pool it was opened on.
93    pub fn same_database(&self, other: &Db) -> bool {
94        std::ptr::addr_eq(Arc::as_ptr(self.pool()), Arc::as_ptr(other.pool()))
95    }
96
97    /// Identity tag for plans built against this database.
98    pub(crate) fn origin(&self) -> PlanOrigin {
99        PlanOrigin::new(Arc::as_ptr(self.pool()).cast::<()>() as usize)
100    }
101
102    /// Reject work carrying a subquery built against another database.
103    fn ensure_local(&self, foreign: bool) -> Result<(), OrmError> {
104        if foreign {
105            return Err(QueryError::InvalidPlan(
106                "a subquery built against another database cannot run here".into(),
107            )
108            .into());
109        }
110        Ok(())
111    }
112
113    fn pool(&self) -> &Arc<dyn Backend> {
114        match &self.target {
115            Target::Pool(pool) => pool,
116            Target::Tx(state) => &state.pool,
117            Target::Conn(state) => &state.pool,
118        }
119    }
120
121    fn executor(&self) -> Result<ExecutorRef<'_>, OrmError> {
122        match &self.target {
123            Target::Pool(pool) => Ok(ExecutorRef::Pool(pool.as_ref())),
124            Target::Tx(state) => {
125                let lease = state.control.lease(state.scope)?;
126                Ok(ExecutorRef::Leased(lease))
127            }
128            Target::Conn(state) => Ok(ExecutorRef::Leased(state.control.lease(0)?)),
129        }
130    }
131
132    /// Declared capabilities of the underlying backend.
133    pub fn capabilities(&self) -> BackendCapabilities {
134        self.pool().capabilities()
135    }
136
137    /// Count read-plan binds through the backend compiler without I/O.
138    ///
139    /// Args:
140    ///     plan: Complete target read plan, including existing filters.
141    ///
142    /// Returns:
143    ///     Exact count when supported; None for non-reporting backends.
144    ///
145    /// # Errors
146    /// Foreign database origin, invalid plan or unsupported capability.
147    pub fn read_parameter_count(&self, plan: &QueryPlan) -> Result<Option<usize>, OrmError> {
148        self.ensure_local(plan.has_foreign_origin(self.origin()))?;
149        self.pool().read_parameter_count(plan)
150    }
151
152    /// Whether this handle runs inside a transaction.
153    pub fn in_transaction(&self) -> bool {
154        matches!(self.target, Target::Tx(_))
155    }
156
157    /// Whether this handle owns a connection or transaction scope.
158    ///
159    /// Returns:
160    ///     True for scoped handles; false for a pool handle.
161    pub fn is_scoped(&self) -> bool {
162        !matches!(self.target, Target::Pool(_))
163    }
164
165    /// Execute a read plan.
166    ///
167    /// # Errors
168    /// Capability, backend or decode errors.
169    pub async fn fetch(&self, plan: &QueryPlan) -> Result<QueryResult, OrmError> {
170        self.ensure_local(plan.has_foreign_origin(self.origin()))?;
171        let mut executor = self.executor()?;
172        let query = executor.as_ref().fetch(plan);
173        let result = self.traced("select", &plan.source.name, query).await;
174        executor.finish();
175        result
176    }
177
178    /// Execute a write plan.
179    ///
180    /// # Errors
181    /// Capability or backend errors.
182    pub async fn execute(&self, plan: &WritePlan) -> Result<ExecResult, OrmError> {
183        let origin = self.origin();
184        let foreign = match plan {
185            WritePlan::Insert(_) => false,
186            WritePlan::Update(p) => p
187                .assignments
188                .iter()
189                .map(|(_, e)| e)
190                .chain(p.filter.iter())
191                .any(|e| e.has_foreign_origin(origin)),
192            WritePlan::Delete(p) => p.filter.iter().any(|e| e.has_foreign_origin(origin)),
193        };
194        self.ensure_local(foreign)?;
195        let (operation, table) = match plan {
196            WritePlan::Insert(p) => ("insert", &p.table),
197            WritePlan::Update(p) => ("update", &p.table),
198            WritePlan::Delete(p) => ("delete", &p.table),
199        };
200        let mut executor = self.executor()?;
201        let result = self
202            .traced(operation, table, executor.as_ref().execute(plan))
203            .await;
204        executor.finish();
205        result
206    }
207
208    /// Raw SQL returning rows (`db.raw_sql("SELECT .. WHERE id = ?", params![id])`).
209    ///
210    /// Placeholders follow the backend (`?` on SQLite, `$1` on PostgreSQL).
211    /// Parameters are always bound; never format untrusted input into `sql`.
212    ///
213    /// # Errors
214    /// Backend errors.
215    pub async fn raw_sql(&self, sql: &str, params: Vec<Value>) -> Result<QueryResult, OrmError> {
216        tracing::trace!(sql, "raw query");
217        let mut executor = self.executor()?;
218        let result = self
219            .traced("raw", "", executor.as_ref().fetch_raw(sql, params))
220            .await;
221        executor.finish();
222        result
223    }
224
225    /// Raw SQL returning the affected-row count. Parameters are bound.
226    ///
227    /// # Errors
228    /// Backend errors.
229    pub async fn raw_execute(&self, sql: &str, params: Vec<Value>) -> Result<u64, OrmError> {
230        tracing::trace!(sql, "raw statement");
231        let mut executor = self.executor()?;
232        let result = self
233            .traced("raw", "", executor.as_ref().execute_raw(sql, params))
234            .await;
235        executor.finish();
236        result
237    }
238
239    /// Run a parameterless multi-statement script (DDL).
240    ///
241    /// # Errors
242    /// Backend errors.
243    pub async fn execute_script(&self, sql: &str) -> Result<(), OrmError> {
244        tracing::trace!(sql, "script");
245        let mut executor = self.executor()?;
246        let result = self
247            .traced("script", "", executor.as_ref().execute_script(sql))
248            .await;
249        executor.finish();
250        result
251    }
252
253    /// Run `query` inside an `orm.query` span and record its duration.
254    ///
255    /// The span carries `db.system`, `db.operation`, `db.table` (empty for raw
256    /// SQL) and `elapsed_ms`. It is emitted at `debug` level. Bind parameters
257    /// are never recorded, and SQL text only by the raw entry points at
258    /// `trace` level.
259    async fn traced<T>(
260        &self,
261        operation: &'static str,
262        table: &str,
263        query: impl Future<Output = Result<T, OrmError>>,
264    ) -> Result<T, OrmError> {
265        let span = tracing::debug_span!(
266            "orm.query",
267            db.system = self.capabilities().kind.system_name(),
268            db.operation = operation,
269            db.table = table,
270            elapsed_ms = tracing::field::Empty,
271        );
272        let recorder = span.clone();
273        async move {
274            let started = Instant::now();
275            let result = query.await;
276            recorder.record("elapsed_ms", started.elapsed().as_secs_f64() * 1000.0);
277            result
278        }
279        .instrument(span)
280        .await
281    }
282}
283
284/// Named databases (`"default"`, `"analytics"`, ...) plus an optional
285/// [`DatabaseRouter`].
286///
287/// A queryset holds one [`Db`] and never spans aliases: pick the database
288/// with [`get`](Self::get), [`using`](Self::using) or the router-aware
289/// [`for_read`](Self::for_read) / [`for_write`](Self::for_write) /
290/// [`objects`](Self::objects). Combining querysets bound to different
291/// databases (`union` and friends) is an error.
292#[derive(Clone, Default)]
293pub struct Databases {
294    by_alias: BTreeMap<String, Db>,
295    router: Option<Arc<dyn DatabaseRouter>>,
296}
297
298impl std::fmt::Debug for Databases {
299    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
300        f.debug_struct("Databases")
301            .field("aliases", &self.by_alias.keys().collect::<Vec<_>>())
302            .field("router", &self.router.is_some())
303            .finish()
304    }
305}
306
307impl Databases {
308    /// Alias used by [`default_db`](Self::default_db).
309    pub const DEFAULT: &'static str = "default";
310
311    /// Empty registry.
312    pub fn new() -> Self {
313        Self::default()
314    }
315
316    /// Register `db` under `alias` (replacing an existing entry).
317    #[must_use]
318    pub fn with(mut self, alias: impl Into<String>, db: Db) -> Self {
319        self.by_alias.insert(alias.into(), db);
320        self
321    }
322
323    /// Route reads and writes through `router` (replacing a previous one).
324    #[must_use]
325    pub fn with_router(mut self, router: impl DatabaseRouter) -> Self {
326        self.router = Some(Arc::new(router));
327        self
328    }
329
330    /// Handle registered under `alias`.
331    pub fn get(&self, alias: &str) -> Option<&Db> {
332        self.by_alias.get(alias)
333    }
334
335    /// Handle registered as `"default"`.
336    pub fn default_db(&self) -> Option<&Db> {
337        self.get(Self::DEFAULT)
338    }
339
340    /// Registered aliases, sorted.
341    pub fn aliases(&self) -> impl Iterator<Item = &str> {
342        self.by_alias.keys().map(String::as_str)
343    }
344
345    /// The database `M` is read from: the router's choice, else `"default"`.
346    ///
347    /// # Errors
348    /// [`OrmError::UnknownDatabase`] if the chosen alias is not registered.
349    pub fn for_read<M: Model>(&self) -> Result<&Db, OrmError> {
350        self.resolve(self.router.as_deref().and_then(|r| r.db_for_read(M::META)))
351    }
352
353    /// The database `M` is written to: the router's choice, else `"default"`.
354    ///
355    /// # Errors
356    /// [`OrmError::UnknownDatabase`] if the chosen alias is not registered.
357    pub fn for_write<M: Model>(&self) -> Result<&Db, OrmError> {
358        self.resolve(self.router.as_deref().and_then(|r| r.db_for_write(M::META)))
359    }
360
361    /// Queryset over `M` on its read database (Django `Model.objects`
362    /// under a router).
363    ///
364    /// # Errors
365    /// As [`for_read`](Self::for_read).
366    pub fn objects<M: Model>(&self) -> Result<QuerySet<M>, OrmError> {
367        self.for_read::<M>().map(M::objects)
368    }
369
370    /// Queryset over `M` on the database registered as `alias` (Django
371    /// `Model.objects.using(alias)`), bypassing the router.
372    ///
373    /// # Errors
374    /// [`OrmError::UnknownDatabase`] if `alias` is not registered.
375    pub fn using<M: Model>(&self, alias: &str) -> Result<QuerySet<M>, OrmError> {
376        self.resolve(Some(alias)).map(M::objects)
377    }
378
379    /// Whether migrations may create `model` on `alias` (`true` without a
380    /// router).
381    pub fn allow_migrate(&self, alias: &str, model: &ModelMeta) -> bool {
382        self.router
383            .as_deref()
384            .is_none_or(|r| r.allow_migrate(alias, model))
385    }
386
387    fn resolve(&self, alias: Option<&str>) -> Result<&Db, OrmError> {
388        let alias = alias.unwrap_or(Self::DEFAULT);
389        self.get(alias)
390            .ok_or_else(|| OrmError::UnknownDatabase(alias.to_owned()))
391    }
392}
393
394/// Build a `Vec<Value>` of bind parameters: `params![id, "name"]`.
395#[macro_export]
396macro_rules! params {
397    ($($value:expr),* $(,)?) => {
398        ::std::vec![$($crate::Value::from($value)),*]
399    };
400}
401
402#[cfg(test)]
403mod tests {
404    use super::*;
405    use crate::backend::{ExecResult, Executor, QueryResult, Transaction};
406    use crate::capabilities::BackendCapabilities;
407    use crate::plan::QueryPlan;
408    use crate::write::DeletePlan;
409    use async_trait::async_trait;
410    use std::sync::Mutex;
411    use tracing::field::{Field, Visit};
412    use tracing::span::{Attributes, Id, Record};
413    use tracing::{Event, Subscriber};
414    use tracing_subscriber::Layer;
415    use tracing_subscriber::layer::{Context, SubscriberExt};
416    use tracing_subscriber::registry::LookupSpan;
417
418    /// Backend that answers every statement with an empty result.
419    struct Stub {
420        _id: u8,
421    }
422
423    #[async_trait]
424    impl Executor for Stub {
425        fn capabilities(&self) -> BackendCapabilities {
426            BackendCapabilities::sqlite()
427        }
428        async fn fetch(&self, _: &QueryPlan) -> Result<QueryResult, OrmError> {
429            Ok(QueryResult::default())
430        }
431        async fn execute(&self, _: &WritePlan) -> Result<ExecResult, OrmError> {
432            Ok(ExecResult::default())
433        }
434        async fn fetch_raw(&self, _: &str, _: Vec<Value>) -> Result<QueryResult, OrmError> {
435            Ok(QueryResult::default())
436        }
437        async fn execute_raw(&self, _: &str, _: Vec<Value>) -> Result<u64, OrmError> {
438            Ok(0)
439        }
440        async fn execute_script(&self, _: &str) -> Result<(), OrmError> {
441            Ok(())
442        }
443    }
444
445    #[async_trait]
446    impl Backend for Stub {
447        async fn begin(
448            &self,
449            _: Option<crate::capabilities::IsolationLevel>,
450        ) -> Result<Box<dyn Transaction>, OrmError> {
451            Ok(Box::new(Stub { _id: 0 }))
452        }
453    }
454
455    #[async_trait]
456    impl Transaction for Stub {
457        async fn commit(&self) -> Result<(), OrmError> {
458            Ok(())
459        }
460        async fn rollback(&self) -> Result<(), OrmError> {
461            Ok(())
462        }
463    }
464
465    fn stub_db() -> Db {
466        Db::new(Stub { _id: 1 })
467    }
468
469    /// Fields of one recorded span or event.
470    #[derive(Debug, Default, Clone)]
471    struct Captured {
472        name: String,
473        level: Option<tracing::Level>,
474        fields: Vec<(String, String)>,
475    }
476
477    impl Captured {
478        fn field(&self, name: &str) -> Option<&str> {
479            self.fields
480                .iter()
481                .rev()
482                .find(|(n, _)| n == name)
483                .map(|(_, v)| v.as_str())
484        }
485    }
486
487    struct Collect<'a>(&'a mut Vec<(String, String)>);
488
489    impl Visit for Collect<'_> {
490        fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
491            self.0.push((field.name().to_owned(), format!("{value:?}")));
492        }
493        fn record_str(&mut self, field: &Field, value: &str) {
494            self.0.push((field.name().to_owned(), value.to_owned()));
495        }
496    }
497
498    #[derive(Clone, Default)]
499    struct Recorder {
500        spans: Arc<Mutex<Vec<(Id, Captured)>>>,
501        events: Arc<Mutex<Vec<Captured>>>,
502    }
503
504    impl<S: Subscriber + for<'a> LookupSpan<'a>> Layer<S> for Recorder {
505        fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, _: Context<'_, S>) {
506            let mut captured = Captured {
507                name: attrs.metadata().name().to_owned(),
508                level: Some(*attrs.metadata().level()),
509                fields: Vec::new(),
510            };
511            attrs.record(&mut Collect(&mut captured.fields));
512            self.spans.lock().unwrap().push((id.clone(), captured));
513        }
514        fn on_record(&self, id: &Id, values: &Record<'_>, _: Context<'_, S>) {
515            let mut spans = self.spans.lock().unwrap();
516            if let Some((_, captured)) = spans.iter_mut().find(|(i, _)| i == id) {
517                values.record(&mut Collect(&mut captured.fields));
518            }
519        }
520        fn on_event(&self, event: &Event<'_>, _: Context<'_, S>) {
521            let mut captured = Captured {
522                name: event.metadata().name().to_owned(),
523                level: Some(*event.metadata().level()),
524                fields: Vec::new(),
525            };
526            event.record(&mut Collect(&mut captured.fields));
527            self.events.lock().unwrap().push(captured);
528        }
529    }
530
531    fn recording() -> (Recorder, tracing::subscriber::DefaultGuard) {
532        let recorder = Recorder::default();
533        let subscriber = tracing_subscriber::registry().with(recorder.clone());
534        (recorder, tracing::subscriber::set_default(subscriber))
535    }
536
537    fn query_spans(recorder: &Recorder) -> Vec<Captured> {
538        recorder
539            .spans
540            .lock()
541            .unwrap()
542            .iter()
543            .map(|(_, c)| c.clone())
544            .filter(|c| c.name == "orm.query")
545            .collect()
546    }
547
548    #[tokio::test]
549    async fn every_entry_point_runs_in_an_orm_query_span() {
550        let (recorder, _guard) = recording();
551        let db = stub_db();
552        db.fetch(&QueryPlan::from_table("books")).await.unwrap();
553        db.execute(&WritePlan::Delete(DeletePlan {
554            table: "authors".into(),
555            filter: None,
556            returning: Vec::new(),
557        }))
558        .await
559        .unwrap();
560        db.raw_sql("SELECT 1", Vec::new()).await.unwrap();
561        db.raw_execute("DELETE FROM t", Vec::new()).await.unwrap();
562        db.execute_script("CREATE TABLE t (id INTEGER)")
563            .await
564            .unwrap();
565
566        let spans = query_spans(&recorder);
567        let seen: Vec<(&str, &str)> = spans
568            .iter()
569            .map(|s| {
570                (
571                    s.field("db.operation").unwrap(),
572                    s.field("db.table").unwrap(),
573                )
574            })
575            .collect();
576        assert_eq!(
577            seen,
578            [
579                ("select", "books"),
580                ("delete", "authors"),
581                ("raw", ""),
582                ("raw", ""),
583                ("script", "")
584            ]
585        );
586        for span in &spans {
587            assert_eq!(span.field("db.system"), Some("sqlite"));
588            assert_eq!(span.level, Some(tracing::Level::DEBUG));
589            let elapsed: f64 = span.field("elapsed_ms").unwrap().parse().unwrap();
590            assert!(elapsed >= 0.0);
591        }
592    }
593
594    #[tokio::test]
595    async fn raw_sql_is_traced_but_bind_parameters_never_are() {
596        let (recorder, _guard) = recording();
597        let db = stub_db();
598        db.raw_sql(
599            "SELECT * FROM users WHERE token = ?",
600            vec![Value::from("s3cr3t-token")],
601        )
602        .await
603        .unwrap();
604
605        let events = recorder.events.lock().unwrap();
606        let sql_event = events
607            .iter()
608            .find(|e| e.field("sql").is_some())
609            .expect("sql event");
610        assert_eq!(sql_event.level, Some(tracing::Level::TRACE));
611        assert_eq!(
612            sql_event.field("sql"),
613            Some("SELECT * FROM users WHERE token = ?")
614        );
615        let everything = format!("{events:?} {:?}", query_spans(&recorder));
616        assert!(!everything.contains("s3cr3t-token"), "{everything}");
617    }
618
619    #[tokio::test]
620    async fn plan_calls_never_emit_sql_or_values() {
621        let (recorder, _guard) = recording();
622        let db = stub_db();
623        db.fetch(&QueryPlan::from_table("books")).await.unwrap();
624        assert!(recorder.events.lock().unwrap().is_empty());
625    }
626
627    #[tokio::test]
628    async fn same_database_follows_the_pool_across_clones_and_transactions() {
629        let (a, b) = (stub_db(), stub_db());
630        assert!(a.same_database(&a.clone()));
631        assert!(!a.same_database(&b));
632        let (outer, other) = (a.clone(), b.clone());
633        a.transaction(|tx| async move {
634            assert!(tx.same_database(&outer));
635            assert!(!tx.same_database(&other));
636            Ok::<_, OrmError>(())
637        })
638        .await
639        .unwrap();
640    }
641
642    #[tokio::test]
643    async fn signals_are_shared_by_clones_and_transaction_handles() {
644        let signals = Signals::new();
645        let db = stub_db().with_signals(signals.clone());
646        assert!(db.clone().signals().is_empty());
647        db.transaction(|tx| async move {
648            // The registry is shared, not copied.
649            let seen = tx.signals().clone();
650            tx.transaction(|nested| async move {
651                assert_eq!(nested.signals().len(), seen.len());
652                Ok::<_, OrmError>(())
653            })
654            .await
655        })
656        .await
657        .unwrap();
658    }
659
660    #[test]
661    fn databases_lists_sorted_aliases_and_hides_the_router_in_debug() {
662        let dbs = Databases::new()
663            .with("replica", stub_db())
664            .with("default", stub_db())
665            .with("analytics", stub_db());
666        assert_eq!(
667            dbs.aliases().collect::<Vec<_>>(),
668            ["analytics", "default", "replica"]
669        );
670        assert!(dbs.default_db().is_some());
671        assert!(dbs.get("missing").is_none());
672        assert!(format!("{dbs:?}").contains("analytics"));
673    }
674}
siderite_orm/db.rs

siderite_orm/
db.rs
