Skip to main content

laterite_core/
capabilities.rs

1//! Database capability declarations.
2//!
3//! A module declares the database capabilities it needs through
4//! [`Module::requires_db_capabilities`](crate::module::Module::requires_db_capabilities)
5//! and [`Module::optional_db_capabilities`](crate::module::Module::optional_db_capabilities).
6//! A required capability must be present or the application refuses to boot with a
7//! clear error, instead of failing deep inside a migration. An optional capability
8//! enables extra behaviour when present and is skipped when absent; a module
9//! queries the returned [`CapabilitySet`] to decide.
10//!
11//! Capabilities are currently PostgreSQL extensions, so availability means
12//! "installable on this server" (present in `pg_available_extensions`). On other
13//! backends none of them are available, so a module that requires one runs only
14//! on PostgreSQL. The naming is capability-oriented rather than extension-oriented
15//! so the check can grow other backends' equivalents later.
16
17use std::collections::BTreeSet;
18
19use sqlx::AnyPool;
20
21use crate::error::{CoreError, CoreResult};
22use crate::migration::DbBackend;
23use crate::module::{Capability, ModuleRegistry};
24
25/// The database capabilities found available at boot. Query it to gate optional
26/// features.
27#[derive(Debug, Default, Clone)]
28pub struct CapabilitySet {
29    available: BTreeSet<String>,
30}
31
32impl CapabilitySet {
33    /// Whether a capability is available on the connected database.
34    pub fn has(&self, name: &str) -> bool {
35        self.available.contains(name)
36    }
37
38    /// The available capability names, sorted.
39    pub fn names(&self) -> impl Iterator<Item = &str> {
40        self.available.iter().map(String::as_str)
41    }
42}
43
44/// Detects which declared capabilities are available, and errors if any module's
45/// required capability is missing. Returns the available set, for gating optional
46/// features. Run once at boot, before migrations.
47pub async fn check(
48    pool: &AnyPool,
49    backend: DbBackend,
50    registry: &ModuleRegistry,
51) -> CoreResult<CapabilitySet> {
52    let mut wanted: BTreeSet<&str> = BTreeSet::new();
53    for module in registry.iter() {
54        wanted.extend(
55            module
56                .requires_db_capabilities()
57                .iter()
58                .map(Capability::as_str),
59        );
60        wanted.extend(
61            module
62                .optional_db_capabilities()
63                .iter()
64                .map(Capability::as_str),
65        );
66    }
67
68    let available = detect(pool, backend, &wanted).await?;
69
70    for module in registry.iter() {
71        for cap in module.requires_db_capabilities() {
72            if !available.contains(cap.as_str()) {
73                return Err(CoreError::MissingDbCapability {
74                    module: module.id().to_string(),
75                    capability: cap.as_str().to_string(),
76                    backend: backend.name().to_string(),
77                });
78            }
79        }
80    }
81
82    Ok(CapabilitySet {
83        available: available.into_iter().map(str::to_string).collect(),
84    })
85}
86
87/// Which of `wanted` the backend can provide. Capabilities are PostgreSQL
88/// extensions; other backends provide none of them.
89async fn detect<'a>(
90    pool: &AnyPool,
91    backend: DbBackend,
92    wanted: &BTreeSet<&'a str>,
93) -> CoreResult<BTreeSet<&'a str>> {
94    if wanted.is_empty() || backend != DbBackend::Postgres {
95        return Ok(BTreeSet::new());
96    }
97    // Cast to text so the portable `Any` driver decodes the `name`-typed column.
98    let rows: Vec<String> = sqlx::query_scalar("SELECT name::text FROM pg_available_extensions")
99        .fetch_all(pool)
100        .await?;
101    let installable: BTreeSet<String> = rows.into_iter().collect();
102    Ok(wanted
103        .iter()
104        .copied()
105        .filter(|name| installable.contains(*name))
106        .collect())
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use crate::module::{Module, ModuleId};
113
114    async fn sqlite_pool() -> AnyPool {
115        sqlx::any::install_default_drivers();
116        sqlx::any::AnyPoolOptions::new()
117            .max_connections(1)
118            .connect("sqlite::memory:")
119            .await
120            .unwrap()
121    }
122
123    /// A test module: id, required capabilities, optional capabilities.
124    struct Needs(ModuleId, &'static [Capability], &'static [Capability]);
125
126    impl Module for Needs {
127        fn id(&self) -> ModuleId {
128            self.0
129        }
130        fn requires_db_capabilities(&self) -> &'static [Capability] {
131            self.1
132        }
133        fn optional_db_capabilities(&self) -> &'static [Capability] {
134            self.2
135        }
136    }
137
138    const POSTGIS: &[Capability] = &[Capability::new("postgis")];
139    const PG_TRGM: &[Capability] = &[Capability::new("pg_trgm")];
140    const NO_CAPS: &[Capability] = &[];
141
142    #[tokio::test]
143    async fn required_capability_unavailable_errors() {
144        let pool = sqlite_pool().await;
145        let mut reg = ModuleRegistry::new();
146        reg.register(Needs(ModuleId::new("acme.blog"), POSTGIS, NO_CAPS));
147        let err = check(&pool, DbBackend::Sqlite, &reg).await.unwrap_err();
148        assert!(matches!(err, CoreError::MissingDbCapability { .. }));
149    }
150
151    #[tokio::test]
152    async fn no_requirements_ok_and_absent_optional_not_in_set() {
153        let pool = sqlite_pool().await;
154        let mut reg = ModuleRegistry::new();
155        reg.register(Needs(ModuleId::new("acme.shop"), NO_CAPS, PG_TRGM));
156        let caps = check(&pool, DbBackend::Sqlite, &reg).await.unwrap();
157        assert!(!caps.has("pg_trgm"));
158    }
159}