Skip to main content

sqlserver_mcp_catalog/services/
sql_pool.rs

1// SQL Server 2025 - master/msdb/sandbox combined catalog MCP server.
2//
3// A process-wide cache of `bb8` connection pools, one per distinct
4// `tiberius::Config` (in practice: one, since `Config` is loaded once at
5// startup and reused for the process's lifetime) — mirrors
6// `data::store::cached_store_connection`'s process-wide-cache-by-key shape,
7// the existing convention in this codebase for "expensive resource, built
8// once, shared across tool calls" rather than introducing a different
9// pattern for this one case.
10
11use std::collections::HashMap;
12use std::sync::{Mutex, OnceLock};
13
14use bb8::Pool;
15use bb8_tiberius::ConnectionManager;
16
17pub type SqlPool = Pool<ConnectionManager>;
18
19/// Returns the cached pool for `cache_key`, building one via `config` on
20/// first use. `cache_key` is caller-supplied rather than derived from
21/// `config` here, since two `Config`s that are meaningfully different for
22/// pooling purposes (e.g. different resolved AAD tokens) don't necessarily
23/// differ in a way that's cheap to hash/compare.
24pub async fn cached_pool(
25    cache_key: &str,
26    config: tiberius::Config,
27    max_size: u32,
28) -> anyhow::Result<SqlPool> {
29    static POOLS: OnceLock<Mutex<HashMap<String, SqlPool>>> = OnceLock::new();
30    let pools = POOLS.get_or_init(|| Mutex::new(HashMap::new()));
31
32    if let Some(pool) = pools.lock().unwrap().get(cache_key) {
33        return Ok(pool.clone());
34    }
35
36    let manager = ConnectionManager::new(config);
37    let pool = Pool::builder().max_size(max_size).build(manager).await?;
38
39    // `bb8::Pool` is a cheap `Arc`-backed handle, so the pool built above
40    // and the one returned to the caller after this insert are the same
41    // underlying pool either way — a lost race here (two callers both
42    // missing the cache and both building) just means one extra pool gets
43    // built and then discarded once its `Arc` refcount drops, not a
44    // correctness issue.
45    let mut pools = pools.lock().unwrap();
46    Ok(pools.entry(cache_key.to_string()).or_insert(pool).clone())
47}