Skip to main content

mll_axum_utils/database/
postgres.rs

1use crate::res::Res;
2use axum::async_trait;
3use axum::extract::FromRequestParts;
4use axum::http::request::Parts;
5use bb8::{Pool, PooledConnection};
6use diesel_async::pooled_connection::AsyncDieselConnectionManager;
7use diesel_async::AsyncPgConnection;
8
9pub type PgPool = AsyncDieselConnectionManager<AsyncPgConnection>;
10pub type Conn = PooledConnection<'static, PgPool>;
11
12/// 新建 pg 异步连接池
13pub async fn new_pg_pool(database_url: &str) -> Pool<PgPool> {
14    let config: PgPool = PgPool::new(database_url);
15    Pool::builder()
16        .build(config)
17        .await
18        .unwrap_or_else(|_| panic!("Error postgres connecting to {}", database_url))
19}
20
21/// 获取 pg 连接
22pub struct PgConn(pub Conn);
23
24#[async_trait]
25impl<S> FromRequestParts<S> for PgConn
26where
27    S: Send + Sync,
28{
29    type Rejection = Res<()>;
30
31    async fn from_request_parts(parts: &mut Parts, _: &S) -> Result<Self, Self::Rejection> {
32        let pool = parts
33            .extensions
34            .get::<Pool<PgPool>>()
35            .expect("未设置 PgPool")
36            .clone();
37        let conn = pool
38            .get_owned()
39            .await
40            .map_err(|_| Res::msg(500, "获取数据库连接失败"))?;
41
42        Ok(Self(conn))
43    }
44}