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 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 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 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 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}