spring_sqlx/
lib.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
//! [spring-sqlx](https://spring-rs.github.io/docs/plugins/spring-sqlx/)
#![doc(html_favicon_url = "https://spring-rs.github.io/favicon.ico")]
#![doc(html_logo_url = "https://spring-rs.github.io/logo.svg")]

pub mod config;
pub extern crate sqlx;
use anyhow::Context;
use config::SqlxConfig;
use spring::app::AppBuilder;
use spring::config::ConfigRegistry;
use spring::error::Result;
use spring::plugin::Plugin;
use spring::{async_trait, App};
use sqlx::{Database, Pool};
use std::sync::Arc;
use std::time::Duration;

#[cfg(not(feature = "postgres"))]
pub type ConnectPool = sqlx::AnyPool;
#[cfg(feature = "postgres")]
pub type ConnectPool = sqlx::PgPool;

pub struct SqlxPlugin;

#[async_trait]
impl Plugin for SqlxPlugin {
    async fn build(&self, app: &mut AppBuilder) {
        sqlx::any::install_default_drivers();
        let config = app
            .get_config::<SqlxConfig>()
            .expect("sqlx plugin config load failed");

        let connect_pool = Self::connect(&config)
            .await
            .expect("sqlx plugin load failed");

        tracing::info!("sqlx connection success");

        app.add_component(connect_pool)
            .add_shutdown_hook(|app| Box::new(Self::close_db_connection(app)));
    }
}

impl SqlxPlugin {
    #[cfg(not(feature = "postgres"))]
    pub async fn connect(config: &SqlxConfig) -> Result<ConnectPool> {
        use sqlx::any::AnyPoolOptions;

        let opt = Self::configure_pool(AnyPoolOptions::new(), config);
        Self::establish_connection(opt, &config.uri).await
    }

    #[cfg(feature = "postgres")]
    pub async fn connect(config: &SqlxConfig) -> Result<ConnectPool> {
        use sqlx::postgres::PgPoolOptions;

        let opt = Self::configure_pool(PgPoolOptions::new(), config);
        Self::establish_connection(opt, &config.uri).await
    }

    fn configure_pool<T>(mut opt: sqlx::pool::PoolOptions<T>, config: &SqlxConfig) -> sqlx::pool::PoolOptions<T>
    where
        T: Database
    {
        opt = opt
            .max_connections(config.max_connections)
            .min_connections(config.min_connections);

        if let Some(acquire_timeout) = config.acquire_timeout {
            opt = opt.acquire_timeout(Duration::from_millis(acquire_timeout));
        }
        if let Some(idle_timeout) = config.idle_timeout {
            opt = opt.idle_timeout(Duration::from_millis(idle_timeout));
        }
        if let Some(connect_timeout) = config.connect_timeout {
            opt = opt.max_lifetime(Duration::from_millis(connect_timeout));
        }

        opt
    }

    async fn establish_connection<T>(opt: sqlx::pool::PoolOptions<T>, uri: &str) -> Result<Pool<T>>
        where
        T: Database
    {
        Ok(opt.connect(uri)
            .await
            .with_context(|| format!("Failed to connect to database: {}", uri))?)
    }

    async fn close_db_connection(app: Arc<App>) -> Result<String> {
        app.get_component::<ConnectPool>()
            .expect("sqlx connect pool not exists")
            .close()
            .await;
        Ok("sqlx connection pool close successful".into())
    }
}