Skip to main content

sova_db/
plugin.rs

1use crate::handle::DbPool;
2use crate::migrate_cli::run_migrate;
3use crate::tx::inject_conn;
4use sova_core::extend::StateMap;
5use sova_core::{App, Error, Plugin};
6use sea_orm::{ConnectOptions, Database, DatabaseConnection};
7use sea_orm_migration::MigratorTrait;
8use std::future::Future;
9use std::path::Path;
10use std::pin::Pin;
11use std::sync::Arc;
12
13type MigrateFn = Arc<
14    dyn Fn(
15            DatabaseConnection,
16            Vec<String>,
17        ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send>>
18        + Send
19        + Sync,
20>;
21
22type SeedFn = Arc<
23    dyn Fn(Arc<StateMap>) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send>>
24        + Send
25        + Sync,
26>;
27
28/// SeaORM pool plugin (backend selected by URL + Cargo features).
29pub struct Db {
30    url: String,
31    /// When true, [`Self::url`] wins over `DATABASE_URL` / toml at install time.
32    url_pinned: bool,
33    /// Emit sqlx query tracing events (for DevTools / RUST_LOG=sqlx=debug).
34    sqlx_logging: bool,
35    migrate: Option<MigrateFn>,
36    seed: Option<SeedFn>,
37    migrate_on_startup: bool,
38    seed_on_startup: bool,
39}
40
41impl Db {
42    pub fn from_env() -> Self {
43        let url = std::env::var("DATABASE_URL").unwrap_or_default();
44        Self {
45            url,
46            url_pinned: false,
47            sqlx_logging: false,
48            migrate: None,
49            seed: None,
50            migrate_on_startup: false,
51            seed_on_startup: false,
52        }
53    }
54
55    /// Pin the connection URL (takes precedence over `DATABASE_URL` and `[db] url`).
56    pub fn url(mut self, url: impl Into<String>) -> Self {
57        self.url = url.into();
58        self.url_pinned = true;
59        self
60    }
61
62    /// Enable sqlx statement logging via tracing (DevTools DB tab / `RUST_LOG=sqlx=debug`).
63    pub fn sqlx_logging(mut self, on: bool) -> Self {
64        self.sqlx_logging = on;
65        self
66    }
67
68    /// Apply pending migrations after connect (also still available as `migrate` CLI).
69    pub fn migrate_on_startup(mut self) -> Self {
70        self.migrate_on_startup = true;
71        self
72    }
73
74    /// Run [`Self::seed`] after connect / migrate (also still available as `seed` CLI).
75    pub fn seed_on_startup(mut self) -> Self {
76        self.seed_on_startup = true;
77        self
78    }
79
80    /// Register `myapp migrate [up|down|status] [N]` CLI hooks.
81    pub fn migrations<M: MigratorTrait + 'static>(mut self) -> Self {
82        self.migrate = Some(Arc::new(move |conn, args| {
83            Box::pin(async move { run_migrate::<M>(conn, &args).await })
84        }));
85        self
86    }
87
88    /// Register `myapp seed` CLI (and optionally [`Self::seed_on_startup`]).
89    ///
90    /// Accepts `Result<(), E>` where `E: Into<Error>` so facade `AppError` works with `?`.
91    pub fn seed<F, Fut, E>(mut self, f: F) -> Self
92    where
93        F: Fn(Arc<StateMap>) -> Fut + Send + Sync + 'static,
94        Fut: Future<Output = Result<(), E>> + Send + 'static,
95        E: Into<Error> + Send + 'static,
96    {
97        self.seed = Some(Arc::new(move |state| {
98            let fut = f(state);
99            Box::pin(async move { fut.await.map_err(Into::into) })
100        }));
101        self
102    }
103}
104
105impl Plugin for Db {
106    fn id(&self) -> &'static str {
107        "db"
108    }
109
110    fn meta(&self) -> sova_core::PluginMeta {
111        sova_core::PluginMeta::new("Database")
112            .description("SeaORM pool, migrate CLI, optional seed CLI")
113            .version(env!("CARGO_PKG_VERSION"))
114    }
115
116    fn install(mut self, app: &mut App) {
117        // Pinned `.url()` wins; else `DATABASE_URL`, then `[db] url` in toml.
118        if !self.url_pinned {
119            if let Ok(u) = std::env::var("DATABASE_URL") {
120                if !u.is_empty() {
121                    self.url = u;
122                }
123            }
124            if self.url.is_empty() {
125                if let Some(doc) = app.config_doc() {
126                    if let Some(u) = doc
127                        .section("db")
128                        .and_then(|s| s.get("url").and_then(|v| v.as_str()).map(str::to_string))
129                    {
130                        self.url = resolve_sqlite_url(&u, doc.source_dir.as_deref());
131                    }
132                }
133            }
134        }
135
136        // Toml can enable auto migrate/seed without code changes.
137        if let Some(doc) = app.config_doc() {
138            if let Some(section) = doc.section("db") {
139                if section
140                    .get("migrate_on_startup")
141                    .and_then(|v| v.as_bool())
142                    .unwrap_or(false)
143                {
144                    self.migrate_on_startup = true;
145                }
146                if section
147                    .get("seed_on_startup")
148                    .and_then(|v| v.as_bool())
149                    .unwrap_or(false)
150                {
151                    self.seed_on_startup = true;
152                }
153            }
154        }
155
156        if self.url.is_empty() {
157            app.on_startup(|_state| async {
158                Err(Error::Internal(
159                    "database url is empty; set DATABASE_URL or [db] url in sova.toml".into(),
160                ))
161            });
162            return;
163        }
164
165        let pool = DbPool::new();
166        app.state(pool.clone());
167
168        let url = self.url.clone();
169        let pool_start = pool.clone();
170        let sqlx_logging = self.sqlx_logging;
171        let migrate_boot = self
172            .migrate_on_startup
173            .then(|| self.migrate.clone())
174            .flatten();
175        app.on_startup(move |_state| {
176            let url = url.clone();
177            let pool = pool_start.clone();
178            let migrate_boot = migrate_boot.clone();
179            async move {
180                let mut opt = ConnectOptions::new(url);
181                opt.sqlx_logging(sqlx_logging);
182                let conn = Database::connect(opt)
183                    .await
184                    .map_err(|e| Error::Internal(format!("db connect: {e}")))?;
185                conn.ping()
186                    .await
187                    .map_err(|e| Error::Internal(format!("db ping: {e}")))?;
188                if let Some(migrate) = migrate_boot {
189                    // empty args → migrate up (all pending)
190                    migrate(conn.clone(), Vec::new()).await?;
191                }
192                pool.set(conn);
193                Ok(())
194            }
195        });
196
197        let pool_stop = pool.clone();
198        app.on_shutdown(move || {
199            let pool = pool_stop.clone();
200            async move {
201                pool.clear();
202            }
203        });
204
205        app.use_middleware(inject_conn(pool.clone()));
206
207        let pool_check = pool.clone();
208        app.register_check("db", move |_state| {
209            let pool = pool_check.clone();
210            async move {
211                let conn = pool.get().map_err(Error::from)?;
212                conn.ping()
213                    .await
214                    .map_err(|e| Error::Internal(format!("db ping: {e}")))?;
215                Ok(())
216            }
217        });
218
219        if let Some(migrate) = self.migrate.clone() {
220            let pool_cli = pool.clone();
221            app.register_cli("migrate", move |_state, args| {
222                let pool = pool_cli.clone();
223                let migrate = Arc::clone(&migrate);
224                async move {
225                    let conn = pool.get().map_err(Error::from)?;
226                    migrate(conn, args).await
227                }
228            });
229        }
230
231        if let Some(seed) = self.seed.clone() {
232            let seed_cli = Arc::clone(&seed);
233            app.register_cli("seed", move |state, _args| {
234                let seed = Arc::clone(&seed_cli);
235                async move { seed(state).await }
236            });
237            if self.seed_on_startup {
238                app.on_startup(move |state| {
239                    let seed = Arc::clone(&seed);
240                    async move { seed(state).await }
241                });
242            }
243        }
244    }
245}
246
247/// If `sqlite:` path is relative, resolve it against the directory of `sova.toml`.
248fn resolve_sqlite_url(url: &str, source_dir: Option<&Path>) -> String {
249    let Some(dir) = source_dir else {
250        return url.to_string();
251    };
252    let rest = match url.strip_prefix("sqlite:") {
253        Some(r) if !r.starts_with('/') && !r.starts_with("//") && r != ":memory:" && !r.starts_with(":memory:") => r,
254        _ => return url.to_string(),
255    };
256    // `sqlite:hn.db?mode=rwc` or `sqlite:./data/x.db?mode=rwc`
257    let (path_part, query) = match rest.split_once('?') {
258        Some((p, q)) => (p, Some(q)),
259        None => (rest, None),
260    };
261    let path = Path::new(path_part);
262    if path.is_absolute() {
263        return url.to_string();
264    }
265    let abs = dir.join(path);
266    match query {
267        Some(q) => format!("sqlite://{}?{q}", abs.display()),
268        None => format!("sqlite://{}", abs.display()),
269    }
270}