prometheus_mcp/mcp/
prometheus_config.rs

1use std::env;
2use std::time::Duration;
3
4/// Configuration for Prometheus
5#[derive(Debug, Clone)]
6pub struct PrometheusConfig {
7    /// URL of the Prometheus server
8    pub url: String,
9    /// Timeout for Prometheus API requests in seconds
10    pub timeout: Duration,
11    /// Number of retries for Prometheus API requests
12    pub retries: u32,
13    /// Backoff between retries in milliseconds
14    pub retry_backoff_ms: u64,
15    /// Minimal interval between requests in milliseconds (simple rate limit)
16    pub min_request_interval_ms: Option<u64>,
17    /// Cache TTL for metadata/labels in seconds
18    pub cache_ttl_secs: Option<u64>,
19    /// Basic auth username
20    pub username: Option<String>,
21    /// Basic auth password
22    pub password: Option<String>,
23}
24
25impl Default for PrometheusConfig {
26    fn default() -> Self {
27        Self {
28            url: "http://localhost:9090".to_string(),
29            timeout: Duration::from_secs(10),
30            retries: 3,
31            retry_backoff_ms: 500,
32            min_request_interval_ms: None,
33            cache_ttl_secs: None,
34            username: None,
35            password: None,
36        }
37    }
38}
39
40impl PrometheusConfig {
41    /// Create a new PrometheusConfig from environment variables
42    pub fn from_env() -> Self {
43        let url =
44            env::var("PROMETHEUS_URL").unwrap_or_else(|_| "http://localhost:9090".to_string());
45
46        let timeout_secs = env::var("PROMETHEUS_TIMEOUT")
47            .ok()
48            .and_then(|s| s.parse::<u64>().ok())
49            .unwrap_or(10);
50
51        let retries = env::var("PROMETHEUS_RETRIES")
52            .ok()
53            .and_then(|s| s.parse::<u32>().ok())
54            .unwrap_or(3);
55
56        let retry_backoff_ms = env::var("PROMETHEUS_RETRY_BACKOFF_MS")
57            .ok()
58            .and_then(|s| s.parse::<u64>().ok())
59            .unwrap_or(500);
60
61        let min_request_interval_ms = env::var("PROMETHEUS_MIN_INTERVAL_MS")
62            .ok()
63            .and_then(|s| s.parse::<u64>().ok());
64
65        let cache_ttl_secs = env::var("PROMETHEUS_CACHE_TTL_SECS")
66            .ok()
67            .and_then(|s| s.parse::<u64>().ok());
68
69        let username = env::var("PROMETHEUS_USERNAME").ok();
70        let password = env::var("PROMETHEUS_PASSWORD").ok();
71
72        Self {
73            url,
74            timeout: Duration::from_secs(timeout_secs),
75            retries,
76            retry_backoff_ms,
77            min_request_interval_ms,
78            cache_ttl_secs,
79            username,
80            password,
81        }
82    }
83
84    /// Create a new PrometheusConfig from a map of values (useful for tests)
85    #[cfg(test)]
86    pub fn from_map(map: &std::collections::HashMap<&str, &str>) -> Self {
87        let url = map
88            .get("PROMETHEUS_URL")
89            .map(|s| s.to_string())
90            .unwrap_or_else(|| "http://localhost:9090".to_string());
91
92        let timeout_secs = map
93            .get("PROMETHEUS_TIMEOUT")
94            .and_then(|s| s.parse::<u64>().ok())
95            .unwrap_or(10);
96
97        let retries = map
98            .get("PROMETHEUS_RETRIES")
99            .and_then(|s| s.parse::<u32>().ok())
100            .unwrap_or(3);
101
102        let retry_backoff_ms = map
103            .get("PROMETHEUS_RETRY_BACKOFF_MS")
104            .and_then(|s| s.parse::<u64>().ok())
105            .unwrap_or(500);
106
107        let min_request_interval_ms = map
108            .get("PROMETHEUS_MIN_INTERVAL_MS")
109            .and_then(|s| s.parse::<u64>().ok());
110
111        let cache_ttl_secs = map
112            .get("PROMETHEUS_CACHE_TTL_SECS")
113            .and_then(|s| s.parse::<u64>().ok());
114
115        let username = map.get("PROMETHEUS_USERNAME").map(|s| s.to_string());
116        let password = map.get("PROMETHEUS_PASSWORD").map(|s| s.to_string());
117
118        Self {
119            url,
120            timeout: Duration::from_secs(timeout_secs),
121            retries,
122            retry_backoff_ms,
123            min_request_interval_ms,
124            cache_ttl_secs,
125            username,
126            password,
127        }
128    }
129
130    /// Create a new PrometheusConfig with the given URL
131    #[allow(dead_code)]
132    pub fn with_url(url: String) -> Self {
133        Self {
134            url,
135            ..Default::default()
136        }
137    }
138
139    /// Create a new PrometheusConfig with the given timeout
140    #[allow(dead_code)]
141    pub fn with_timeout(mut self, timeout_secs: u64) -> Self {
142        self.timeout = Duration::from_secs(timeout_secs);
143        self
144    }
145
146    /// Create a new PrometheusConfig with the given number of retries
147    #[allow(dead_code)]
148    pub fn with_retries(mut self, retries: u32) -> Self {
149        self.retries = retries;
150        self
151    }
152
153    /// Set backoff between retries
154    #[allow(dead_code)]
155    pub fn with_retry_backoff_ms(mut self, ms: u64) -> Self {
156        self.retry_backoff_ms = ms;
157        self
158    }
159
160    /// Set min request interval (rate limit)
161    #[allow(dead_code)]
162    pub fn with_min_interval_ms(mut self, ms: u64) -> Self {
163        self.min_request_interval_ms = Some(ms);
164        self
165    }
166
167    /// Set cache TTL (seconds)
168    #[allow(dead_code)]
169    pub fn with_cache_ttl_secs(mut self, secs: u64) -> Self {
170        self.cache_ttl_secs = Some(secs);
171        self
172    }
173
174    /// Set basic auth
175    #[allow(dead_code)]
176    pub fn with_basic_auth(
177        mut self,
178        username: impl Into<String>,
179        password: impl Into<String>,
180    ) -> Self {
181        self.username = Some(username.into());
182        self.password = Some(password.into());
183        self
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190    use std::collections::HashMap;
191
192    #[test]
193    fn test_from_env_defaults_and_overrides() {
194        // Use from_map to avoid mutating global process env in tests
195        let empty: HashMap<&str, &str> = HashMap::new();
196        let cfg = PrometheusConfig::from_map(&empty);
197        assert_eq!(cfg.url, "http://localhost:9090");
198
199        let mut vars: HashMap<&str, &str> = HashMap::new();
200        vars.insert("PROMETHEUS_URL", "http://example:9090");
201        vars.insert("PROMETHEUS_TIMEOUT", "5");
202        vars.insert("PROMETHEUS_RETRIES", "2");
203        vars.insert("PROMETHEUS_RETRY_BACKOFF_MS", "10");
204        vars.insert("PROMETHEUS_MIN_INTERVAL_MS", "20");
205        vars.insert("PROMETHEUS_CACHE_TTL_SECS", "30");
206        vars.insert("PROMETHEUS_USERNAME", "u");
207        vars.insert("PROMETHEUS_PASSWORD", "p");
208
209        let cfg = PrometheusConfig::from_map(&vars);
210        assert_eq!(cfg.url, "http://example:9090");
211        assert_eq!(cfg.timeout, std::time::Duration::from_secs(5));
212        assert_eq!(cfg.retries, 2);
213        assert_eq!(cfg.retry_backoff_ms, 10);
214        assert_eq!(cfg.min_request_interval_ms, Some(20));
215        assert_eq!(cfg.cache_ttl_secs, Some(30));
216        assert_eq!(cfg.username.as_deref(), Some("u"));
217        assert_eq!(cfg.password.as_deref(), Some("p"));
218    }
219}