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
28pub struct Db {
30 url: String,
31 url_pinned: bool,
33 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 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 pub fn sqlx_logging(mut self, on: bool) -> Self {
64 self.sqlx_logging = on;
65 self
66 }
67
68 pub fn migrate_on_startup(mut self) -> Self {
70 self.migrate_on_startup = true;
71 self
72 }
73
74 pub fn seed_on_startup(mut self) -> Self {
76 self.seed_on_startup = true;
77 self
78 }
79
80 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 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 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 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 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
247fn 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 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}