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
26pub struct Db {
28 url: String,
29 url_pinned: bool,
31 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 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 pub fn sqlx_logging(mut self, on: bool) -> Self {
62 self.sqlx_logging = on;
63 self
64 }
65
66 pub fn migrate_on_startup(mut self) -> Self {
68 self.migrate_on_startup = true;
69 self
70 }
71
72 pub fn seed_on_startup(mut self) -> Self {
74 self.seed_on_startup = true;
75 self
76 }
77
78 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 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 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 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 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
246fn 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 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}