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}