gearbox_rs_postgres/
client.rs

1use crate::config::{PgConfig, create_pool_for_schema};
2use gearbox_rs_core::{Error, Hub};
3use gearbox_rs_macros::cog;
4use sqlx::PgPool;
5use std::collections::HashMap;
6use std::sync::Arc;
7
8/// PostgreSQL client Cog that manages a map of connection pools keyed by schema name.
9///
10/// Created automatically from [`PgConfig`](crate::config::PgConfig) during
11/// `Gearbox::crank()`. Inject it into other Cogs or route handlers to access
12/// the database:
13///
14/// ```ignore
15/// #[cog]
16/// pub struct UserRepo {
17///     #[inject]
18///     pg: Arc<PgClient>,
19/// }
20/// ```
21#[cog]
22pub struct PgClient {
23    #[default_async(load_pools)]
24    pub pools: Arc<HashMap<String, PgPool>>,
25}
26
27impl PgClient {
28    /// Returns the connection pool for the given schema name.
29    ///
30    /// Returns [`PgError::SchemaNotFound`](crate::error::PgError::SchemaNotFound)
31    /// if no pool is configured for that schema.
32    pub fn get_pool(&self, schema: &str) -> Result<&PgPool, crate::error::PgError> {
33        self.pools
34            .get(schema)
35            .ok_or_else(|| crate::error::PgError::SchemaNotFound(schema.to_string()))
36    }
37
38    /// Convenience method for single-schema setups.
39    ///
40    /// If exactly one pool exists, returns it. Otherwise looks up the `"default"` pool.
41    /// Returns [`PgError::SchemaNotFound`](crate::error::PgError::SchemaNotFound)
42    /// if no suitable pool is found.
43    pub fn pool(&self) -> Result<&PgPool, crate::error::PgError> {
44        if self.pools.len() == 1 {
45            self.pools
46                .values()
47                .next()
48                .ok_or_else(|| crate::error::PgError::SchemaNotFound("default".to_string()))
49        } else {
50            self.get_pool("default")
51        }
52    }
53}
54
55async fn load_pools(hub: Arc<Hub>) -> Result<Arc<HashMap<String, PgPool>>, Error> {
56    let config = hub.get_config::<PgConfig>()?;
57    let schemas = config.resolved_schemas();
58    let mut pools = HashMap::new();
59
60    for (key, schema) in &schemas {
61        let pool = create_pool_for_schema(&config, schema)
62            .await
63            .map_err(|e| Error::External {
64                context: format!("Failed to create pool for schema '{}'", key),
65                source: Box::new(e),
66            })?;
67        pools.insert(key.clone(), pool);
68    }
69
70    Ok(Arc::new(pools))
71}
gearbox_rs_postgres/client.rs

gearbox_rs_postgres/
client.rs
