Skip to main content

qm_pg/
lib.rs

1#![deny(missing_docs)]
2
3//! PostgreSQL connection and utilities.
4//!
5//! This crate provides PostgreSQL connection management with automatic database
6//! and user creation, connection pooling, and SeaORM integration.
7//!
8//! ## Features
9//!
10//! - **Connection Management**: Connect to PostgreSQL with connection pooling
11//! - **Auto-provisioning**: Create databases and users if they don't exist
12//! - **SeaORM Integration**: Convert pools to SeaORM database connections
13//! - **Configuration**: Environment-based configuration with prefix support
14//!
15//! ## Usage
16//!
17//! \```ignore
18//! use qm_pg::{DbConfig, DB};
19//!
20//! #[tokio::main]
21//! async fn main() -> anyhow::Result<()> {
22//!     let config = DbConfig::new()?;
23//!     let db = DB::new("my-app", &config).await?;
24//!     let conn = db.database_connection();
25//!     Ok(())
26//! }
27//! \```
28//!
29//! ## Environment Variables
30//!
31//! | Variable | Description | Default |
32//! |----------|-------------|---------|
33//! | `PG_HOST` | PostgreSQL host | `127.0.0.1` |
34//! | `PG_PORT` | PostgreSQL port | `5432` |
35//! | `PG_USERNAME` | Database username | (none) |
36//! | `PG_PASSWORD` | Database password | (none) |
37//! | `PG_DATABASE` | Database name | (none) |
38//! | `PG_ROOT_USERNAME` | Admin username | (none) |
39//! | `PG_ROOT_PASSWORD` | Admin password | (none) |
40//! | `PG_ROOT_DATABASE` | Admin database | (none) |
41//! | `PG_MAX_CONNECTIONS` | Max pool connections | `32` |
42//! | `PG_MIN_CONNECTIONS` | Min pool connections | `0` |
43
44mod config;
45mod db;
46
47use sqlx::Executor;
48
49/// PostgreSQL configuration.
50pub use crate::config::Config as DbConfig;
51pub use crate::db::DB;
52
53/// Ensure database and user exist, creating them if necessary.
54///
55/// This function connects as root and creates the database and user
56/// if they don't already exist.
57pub async fn ensure(app_name: &str, cfgs: &[&DbConfig]) -> anyhow::Result<()> {
58    for cfg in cfgs {
59        let db = DB::new_root(app_name, cfg).await?;
60        let query_args = ::core::result::Result::<_, ::sqlx::error::BoxDynError>::Ok(
61            <sqlx::postgres::Postgres as ::sqlx::database::Database>::Arguments::<'_>::default(),
62        );
63        let databases: Vec<String> =
64            sqlx::__query_with_result("SELECT datname FROM pg_database;", query_args)
65                .try_map(|v: sqlx::postgres::PgRow| {
66                    use ::sqlx::Row as _;
67                    v.try_get_unchecked::<String, _>(0usize)
68                })
69                .fetch_all(db.pool())
70                .await?;
71
72        if let Some(database) = cfg.database() {
73            if !databases.iter().any(|datname| datname == database) {
74                db.pool()
75                    .execute(format!("CREATE DATABASE {database}").as_str())
76                    .await?;
77            }
78        }
79
80        if let Some((username, password)) = cfg.username().zip(cfg.password()) {
81            let query_args = ::core::result::Result::<_, ::sqlx::error::BoxDynError>::Ok(
82                <sqlx::postgres::Postgres as ::sqlx::database::Database>::Arguments::<'_>::default(
83                ),
84            );
85            let users: Vec<Option<String>> =
86                sqlx::__query_with_result("SELECT usename FROM pg_catalog.pg_user;", query_args)
87                    .try_map(|v: sqlx::postgres::PgRow| {
88                        use ::sqlx::Row as _;
89                        v.try_get_unchecked::<Option<String>, _>(0usize)
90                    })
91                    .fetch_all(db.pool())
92                    .await?;
93            if !users
94                .iter()
95                .any(|usename| usename.as_deref() == Some(username))
96            {
97                db.pool()
98                    .execute(
99                        format!("CREATE USER {username} WITH SUPERUSER PASSWORD '{password}'")
100                            .as_str(),
101                    )
102                    .await?;
103            }
104        }
105    }
106    Ok(())
107}