spring_sqlx/
lib.rs

1//! [![spring-rs](https://img.shields.io/github/stars/spring-rs/spring-rs)](https://spring-rs.github.io/docs/plugins/spring-sqlx)
2#![doc(html_favicon_url = "https://spring-rs.github.io/favicon.ico")]
3#![doc(html_logo_url = "https://spring-rs.github.io/logo.svg")]
4
5pub mod config;
6pub extern crate sqlx;
7use anyhow::Context;
8use config::SqlxConfig;
9use spring::app::AppBuilder;
10use spring::config::ConfigRegistry;
11use spring::error::Result;
12use spring::plugin::{ComponentRegistry, MutableComponentRegistry, Plugin};
13use spring::{async_trait, App};
14use sqlx::{Database, Pool};
15use std::sync::Arc;
16use std::time::Duration;
17
18#[cfg(not(feature = "postgres"))]
19pub type ConnectPool = sqlx::AnyPool;
20#[cfg(feature = "postgres")]
21pub type ConnectPool = sqlx::PgPool;
22
23pub struct SqlxPlugin;
24
25#[async_trait]
26impl Plugin for SqlxPlugin {
27    async fn build(&self, app: &mut AppBuilder) {
28        sqlx::any::install_default_drivers();
29        let config = app
30            .get_config::<SqlxConfig>()
31            .expect("sqlx plugin config load failed");
32
33        let connect_pool = Self::connect(&config)
34            .await
35            .expect("sqlx plugin load failed");
36
37        tracing::info!("sqlx connection success");
38
39        app.add_component(connect_pool)
40            .add_shutdown_hook(|app| Box::new(Self::close_db_connection(app)));
41    }
42}
43
44impl SqlxPlugin {
45    #[cfg(not(feature = "postgres"))]
46    pub async fn connect(config: &SqlxConfig) -> Result<ConnectPool> {
47        use sqlx::any::AnyPoolOptions;
48
49        let opt = Self::configure_pool(AnyPoolOptions::new(), config);
50        Self::establish_connection(opt, &config.uri).await
51    }
52
53    #[cfg(feature = "postgres")]
54    pub async fn connect(config: &SqlxConfig) -> Result<ConnectPool> {
55        use sqlx::postgres::PgPoolOptions;
56
57        let opt = Self::configure_pool(PgPoolOptions::new(), config);
58        Self::establish_connection(opt, &config.uri).await
59    }
60
61    fn configure_pool<T>(
62        mut opt: sqlx::pool::PoolOptions<T>,
63        config: &SqlxConfig,
64    ) -> sqlx::pool::PoolOptions<T>
65    where
66        T: Database,
67    {
68        opt = opt
69            .max_connections(config.max_connections)
70            .min_connections(config.min_connections);
71
72        if let Some(acquire_timeout) = config.acquire_timeout {
73            opt = opt.acquire_timeout(Duration::from_millis(acquire_timeout));
74        }
75        if let Some(idle_timeout) = config.idle_timeout {
76            opt = opt.idle_timeout(Duration::from_millis(idle_timeout));
77        }
78        if let Some(connect_timeout) = config.connect_timeout {
79            opt = opt.max_lifetime(Duration::from_millis(connect_timeout));
80        }
81
82        opt
83    }
84
85    async fn establish_connection<T>(opt: sqlx::pool::PoolOptions<T>, uri: &str) -> Result<Pool<T>>
86    where
87        T: Database,
88    {
89        Ok(opt
90            .connect(uri)
91            .await
92            .with_context(|| format!("Failed to connect to database: {}", uri))?)
93    }
94
95    async fn close_db_connection(app: Arc<App>) -> Result<String> {
96        app.get_component::<ConnectPool>()
97            .expect("sqlx connect pool not exists")
98            .close()
99            .await;
100        Ok("sqlx connection pool close successful".into())
101    }
102}
spring_sqlx/lib.rs

spring_sqlx/
lib.rs
