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
27pub struct Db {
29 url: String,
30 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 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 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 pub fn seed<F, Fut, E>(mut self, f: F) -> Self
66 where
67 F: Fn(Arc<StateMap>) -> Fut + Send + Sync + 'static,
68 Fut: Future<Output = Result<(), E>> + Send + 'static,
69 E: Into<Error> + Send + 'static,
70 {
71 self.seed = Some(Arc::new(move |state| {
72 let fut = f(state);
73 Box::pin(async move { fut.await.map_err(Into::into) })
74 }));
75 self
76 }
77}
78
79impl Plugin for Db {
80 fn id(&self) -> &'static str {
81 "db"
82 }
83
84 fn meta(&self) -> sova_core::PluginMeta {
85 sova_core::PluginMeta::new("Database")
86 .description("SeaORM pool, migrate CLI, optional seed CLI")
87 .version(env!("CARGO_PKG_VERSION"))
88 }
89
90 fn install(mut self, app: &mut App) {
91 if !self.url_pinned {
93 if let Ok(u) = std::env::var("DATABASE_URL") {
94 if !u.is_empty() {
95 self.url = u;
96 }
97 }
98 if self.url.is_empty() {
99 if let Some(u) = app
100 .config_doc()
101 .and_then(|d| d.section("db"))
102 .and_then(|s| s.get("url").and_then(|v| v.as_str()).map(str::to_string))
103 {
104 self.url = u;
105 }
106 }
107 }
108
109 if self.url.is_empty() {
110 app.on_startup(|_state| async {
111 Err(Error::Internal(
112 "database url is empty; set DATABASE_URL or [db] url in sova.toml".into(),
113 ))
114 });
115 return;
116 }
117
118 let pool = DbPool::new();
119 app.state(pool.clone());
120
121 let url = self.url.clone();
122 let pool_start = pool.clone();
123 app.on_startup(move |_state| {
124 let url = url.clone();
125 let pool = pool_start.clone();
126 async move {
127 let conn = Database::connect(&url)
128 .await
129 .map_err(|e| Error::Internal(format!("db connect: {e}")))?;
130 conn.ping()
131 .await
132 .map_err(|e| Error::Internal(format!("db ping: {e}")))?;
133 pool.set(conn).await;
134 Ok(())
135 }
136 });
137
138 let pool_stop = pool.clone();
139 app.on_shutdown(move || {
140 let pool = pool_stop.clone();
141 async move {
142 pool.clear().await;
143 }
144 });
145
146 app.use_middleware(inject_conn(pool.clone()));
147
148 let pool_check = pool.clone();
149 app.register_check("db", move |_state| {
150 let pool = pool_check.clone();
151 async move {
152 let conn = pool.get().await.map_err(Error::from)?;
153 conn.ping()
154 .await
155 .map_err(|e| Error::Internal(format!("db ping: {e}")))?;
156 Ok(())
157 }
158 });
159
160 if let Some(migrate) = self.migrate {
161 let pool_cli = pool.clone();
162 app.register_cli("migrate", move |_state, args| {
163 let pool = pool_cli.clone();
164 let migrate = Arc::clone(&migrate);
165 async move {
166 let conn = pool.get().await.map_err(Error::from)?;
167 migrate(conn, args).await
168 }
169 });
170 }
171
172 if let Some(seed) = self.seed {
173 app.register_cli("seed", move |state, _args| {
174 let seed = Arc::clone(&seed);
175 async move { seed(state).await }
176 });
177 }
178 }
179}