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