Skip to main content

sova_db/
plugin.rs

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