sqlx_pool_registry/
registry.rs

1use crate::{db_pools::DbPools, provider::PoolProvider};
2use std::{collections::HashMap, error::Error, fmt};
3
4/// Error returned when a named pool is not registered.
5///
6/// Available with the `with-named-pools` feature.
7#[derive(Clone, Debug, Eq, PartialEq)]
8pub struct UnknownPool {
9    name: String,
10}
11
12impl UnknownPool {
13    fn new(name: &str) -> Self {
14        Self {
15            name: name.to_owned(),
16        }
17    }
18
19    /// Return the name that was not found.
20    pub fn name(&self) -> &str {
21        &self.name
22    }
23}
24
25impl fmt::Display for UnknownPool {
26    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
27        write!(formatter, "unknown pool `{}`", self.name)
28    }
29}
30
31impl Error for UnknownPool {}
32
33/// A registry of database pool providers keyed by name.
34///
35/// `PoolRegistry` is available with the `with-named-pools` feature. It stores
36/// one homogeneous provider type and defaults to [`DbPools`]. Selecting a name
37/// returns the provider for that name; the registry itself does not keep a
38/// stateful "current" pool and does not implement [`PoolProvider`].
39///
40/// # Examples
41///
42/// ```rust
43/// use sqlx_pool_registry::sqlx::{self, postgres::PgPoolOptions};
44/// use sqlx_pool_registry::{DbPools, PoolProvider, PoolRegistry};
45///
46/// # fn example() -> Result<(), Box<dyn std::error::Error>> {
47/// let auth_primary = PgPoolOptions::new()
48///     .connect_lazy("postgresql://localhost/auth")?;
49/// let auth_replica = PgPoolOptions::new()
50///     .connect_lazy("postgresql://localhost/auth_replica")?;
51/// let analytics = PgPoolOptions::new()
52///     .connect_lazy("postgresql://localhost/analytics")?;
53///
54/// let mut pools = PoolRegistry::new();
55/// pools.insert("auth", DbPools::with_replica(auth_primary, auth_replica));
56/// pools.insert("analytics", DbPools::new(analytics));
57///
58/// let auth = pools.try_get("auth")?;
59/// let _primary = auth.primary();
60/// let _replica = auth.replica(); // Option<&PgPool>, with no fallback
61/// let _read = auth.read(); // Replica, or primary when no replica is configured
62/// # Ok(())
63/// # }
64/// ```
65#[derive(Clone, Debug)]
66pub struct PoolRegistry<P: PoolProvider = DbPools> {
67    pools: HashMap<String, P>,
68}
69
70impl<P: PoolProvider> PoolRegistry<P> {
71    /// Create an empty pool registry.
72    pub fn new() -> Self {
73        Self {
74            pools: HashMap::new(),
75        }
76    }
77
78    /// Insert a provider under `name`.
79    ///
80    /// If the name was already registered, the previous provider is replaced
81    /// and returned.
82    pub fn insert(&mut self, name: impl Into<String>, provider: P) -> Option<P> {
83        self.pools.insert(name.into(), provider)
84    }
85
86    /// Get a provider by name.
87    ///
88    /// Names are matched exactly and case-sensitively.
89    pub fn get(&self, name: &str) -> Option<&P> {
90        self.pools.get(name)
91    }
92
93    /// Get a provider by name, returning [`UnknownPool`] when it is absent.
94    pub fn try_get(&self, name: &str) -> Result<&P, UnknownPool> {
95        self.get(name).ok_or_else(|| UnknownPool::new(name))
96    }
97
98    /// Return whether a provider is registered under `name`.
99    pub fn contains_key(&self, name: &str) -> bool {
100        self.pools.contains_key(name)
101    }
102
103    /// Return the number of registered providers.
104    pub fn len(&self) -> usize {
105        self.pools.len()
106    }
107
108    /// Return whether the registry contains no providers.
109    pub fn is_empty(&self) -> bool {
110        self.pools.is_empty()
111    }
112
113    /// Iterate over registered names and providers in unspecified order.
114    pub fn iter(&self) -> impl Iterator<Item = (&str, &P)> {
115        self.pools
116            .iter()
117            .map(|(name, provider)| (name.as_str(), provider))
118    }
119}
120
121impl<P: PoolProvider> Default for PoolRegistry<P> {
122    fn default() -> Self {
123        Self::new()
124    }
125}
126
127impl<K, P> FromIterator<(K, P)> for PoolRegistry<P>
128where
129    K: Into<String>,
130    P: PoolProvider,
131{
132    fn from_iter<T: IntoIterator<Item = (K, P)>>(iter: T) -> Self {
133        let mut registry = Self::new();
134        registry.extend(iter);
135        registry
136    }
137}
138
139impl<K, P> Extend<(K, P)> for PoolRegistry<P>
140where
141    K: Into<String>,
142    P: PoolProvider,
143{
144    fn extend<T: IntoIterator<Item = (K, P)>>(&mut self, iter: T) {
145        self.pools.extend(
146            iter.into_iter()
147                .map(|(name, provider)| (name.into(), provider)),
148        );
149    }
150}
sqlx_pool_registry/registry.rs

sqlx_pool_registry/
registry.rs
