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::{Database, DatabaseConnection};
7use sea_orm_migration::MigratorTrait;
8use std::future::Future;
9use std::pin::Pin;
10use std::sync::Arc;
11
12type MigrateFn = Arc<
13    dyn Fn(
14            DatabaseConnection,
15            Vec<String>,
16        ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send>>
17        + Send
18        + Sync,
19>;
20
21type SeedFn = Arc<
22    dyn Fn(Arc<StateMap>) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send>>
23        + Send
24        + Sync,
25>;
26
27/// SeaORM pool plugin (backend selected by URL + Cargo features).
28pub struct Db {
29    url: String,
30    migrate: Option<MigrateFn>,
31    seed: Option<SeedFn>,
32}
33
34impl Db {
35    pub fn from_env() -> Self {
36        let url = std::env::var("DATABASE_URL").unwrap_or_default();
37        Self {
38            url,
39            migrate: None,
40            seed: None,
41        }
42    }
43
44    pub fn url(mut self, url: impl Into<String>) -> Self {
45        self.url = url.into();
46        self
47    }
48
49    /// Register `myapp migrate [up|down|status] [N]` CLI hooks.
50    pub fn migrations<M: MigratorTrait + 'static>(mut self) -> Self {
51        self.migrate = Some(Arc::new(move |conn, args| {
52            Box::pin(async move { run_migrate::<M>(conn, &args).await })
53        }));
54        self
55    }
56
57    /// Register `myapp seed` CLI (runs after DB startup; not on every server start).
58    pub fn seed<F, Fut>(mut self, f: F) -> Self
59    where
60        F: Fn(Arc<StateMap>) -> Fut + Send + Sync + 'static,
61        Fut: Future<Output = Result<(), Error>> + Send + 'static,
62    {
63        self.seed = Some(Arc::new(move |state| Box::pin(f(state))));
64        self
65    }
66}
67
68impl Plugin for Db {
69    fn id(&self) -> &'static str {
70        "db"
71    }
72
73    fn meta(&self) -> sova_core::PluginMeta {
74        sova_core::PluginMeta::new("Database")
75            .description("SeaORM pool, migrate CLI, optional seed CLI")
76            .version(env!("CARGO_PKG_VERSION"))
77    }
78
79    fn install(mut self, app: &mut App) {
80        // Env wins, then builder `.url()`, then `[db] url` in toml.
81        if let Ok(u) = std::env::var("DATABASE_URL") {
82            if !u.is_empty() {
83                self.url = u;
84            }
85        }
86        if self.url.is_empty() {
87            if let Some(u) = app
88                .config_doc()
89                .and_then(|d| d.section("db"))
90                .and_then(|s| s.get("url").and_then(|v| v.as_str()).map(str::to_string))
91            {
92                self.url = u;
93            }
94        }
95
96        if self.url.is_empty() {
97            app.on_startup(|_state| async {
98                Err(Error::Internal(
99                    "database url is empty; set DATABASE_URL or [db] url in sova.toml".into(),
100                ))
101            });
102            return;
103        }
104
105        let pool = DbPool::new();
106        app.state(pool.clone());
107
108        let url = self.url.clone();
109        let pool_start = pool.clone();
110        app.on_startup(move |_state| {
111            let url = url.clone();
112            let pool = pool_start.clone();
113            async move {
114                let conn = Database::connect(&url)
115                    .await
116                    .map_err(|e| Error::Internal(format!("db connect: {e}")))?;
117                conn.ping()
118                    .await
119                    .map_err(|e| Error::Internal(format!("db ping: {e}")))?;
120                pool.set(conn).await;
121                Ok(())
122            }
123        });
124
125        let pool_stop = pool.clone();
126        app.on_shutdown(move || {
127            let pool = pool_stop.clone();
128            async move {
129                pool.clear().await;
130            }
131        });
132
133        app.use_middleware(inject_conn(pool.clone()));
134
135        let pool_check = pool.clone();
136        app.register_check("db", move |_state| {
137            let pool = pool_check.clone();
138            async move {
139                let conn = pool.get().await.map_err(Error::from)?;
140                conn.ping()
141                    .await
142                    .map_err(|e| Error::Internal(format!("db ping: {e}")))?;
143                Ok(())
144            }
145        });
146
147        if let Some(migrate) = self.migrate {
148            let pool_cli = pool.clone();
149            app.register_cli("migrate", move |_state, args| {
150                let pool = pool_cli.clone();
151                let migrate = Arc::clone(&migrate);
152                async move {
153                    let conn = pool.get().await.map_err(Error::from)?;
154                    migrate(conn, args).await
155                }
156            });
157        }
158
159        if let Some(seed) = self.seed {
160            app.register_cli("seed", move |state, _args| {
161                let seed = Arc::clone(&seed);
162                async move { seed(state).await }
163            });
164        }
165    }
166}