a2a_rs/adapter/storage/
database_config.rs1#[cfg(feature = "sqlx-storage")]
4use bon::Builder;
5#[cfg(feature = "sqlx-storage")]
6use serde::{Deserialize, Serialize};
7#[cfg(feature = "sqlx-storage")]
8use std::collections::HashMap;
9
10#[cfg(feature = "sqlx-storage")]
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum DatabaseType {
14 Sqlite,
16 Postgres,
18 Mysql,
20}
21
22#[cfg(feature = "sqlx-storage")]
23impl DatabaseType {
24 pub fn from_url(url: &str) -> Option<Self> {
28 if url.starts_with("sqlite:") {
29 Some(Self::Sqlite)
30 } else if url.starts_with("postgres:") || url.starts_with("postgresql:") {
31 Some(Self::Postgres)
32 } else if url.starts_with("mysql:") || url.starts_with("mariadb:") {
33 Some(Self::Mysql)
34 } else {
35 None
36 }
37 }
38
39 pub fn is_feature_enabled(self) -> bool {
47 match self {
48 Self::Sqlite => cfg!(feature = "sqlite"),
49 Self::Postgres => cfg!(feature = "postgres"),
50 Self::Mysql => false,
51 }
52 }
53
54 pub fn feature_name(self) -> &'static str {
56 match self {
57 Self::Sqlite => "sqlite",
58 Self::Postgres => "postgres",
59 Self::Mysql => "mysql",
60 }
61 }
62}
63
64#[cfg(feature = "sqlx-storage")]
65impl std::fmt::Display for DatabaseType {
66 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67 match self {
68 Self::Sqlite => write!(f, "SQLite"),
69 Self::Postgres => write!(f, "PostgreSQL"),
70 Self::Mysql => write!(f, "MySQL"),
71 }
72 }
73}
74
75#[cfg(feature = "sqlx-storage")]
76#[derive(Debug, Clone, Builder, Serialize, Deserialize)]
78pub struct DatabaseConfig {
79 pub url: String,
81 #[builder(default = 10)]
83 pub max_connections: u32,
84 #[builder(default = 30)]
86 pub timeout_seconds: u64,
87 #[builder(default = false)]
89 pub enable_logging: bool,
90}
91
92#[cfg(feature = "sqlx-storage")]
93impl DatabaseConfig {
94 pub fn examples() -> HashMap<&'static str, Self> {
96 [
97 (
98 "sqlite_memory",
99 Self::builder()
100 .url("sqlite::memory:".to_string())
101 .max_connections(1)
102 .enable_logging(true)
103 .build(),
104 ),
105 (
106 "sqlite_file",
107 Self::builder()
108 .url("sqlite:a2a_tasks.db".to_string())
109 .max_connections(5)
110 .build(),
111 ),
112 (
113 "postgres_dev",
114 Self::builder()
115 .url("postgres://user:password@localhost/a2a_dev".to_string())
116 .max_connections(10)
117 .timeout_seconds(10)
118 .build(),
119 ),
120 (
121 "postgres_prod",
122 Self::builder()
123 .url("postgres://user:password@prod-db/a2a_prod".to_string())
124 .max_connections(50)
125 .timeout_seconds(5)
126 .enable_logging(false)
127 .build(),
128 ),
129 (
130 "mysql_dev",
131 Self::builder()
132 .url("mysql://user:password@localhost/a2a_dev".to_string())
133 .max_connections(10)
134 .timeout_seconds(10)
135 .build(),
136 ),
137 ]
138 .into_iter()
139 .collect()
140 }
141
142 pub fn from_env() -> Result<Self, std::env::VarError> {
150 let url = std::env::var("DATABASE_URL")?;
151
152 let max_connections = std::env::var("DATABASE_MAX_CONNECTIONS")
153 .ok()
154 .and_then(|s| s.parse().ok())
155 .unwrap_or(10);
156
157 let timeout_seconds = std::env::var("DATABASE_TIMEOUT_SECONDS")
158 .ok()
159 .and_then(|s| s.parse().ok())
160 .unwrap_or(30);
161
162 let enable_logging = std::env::var("DATABASE_ENABLE_LOGGING")
163 .ok()
164 .and_then(|s| s.parse().ok())
165 .unwrap_or(false);
166
167 Ok(Self::builder()
168 .url(url)
169 .max_connections(max_connections)
170 .timeout_seconds(timeout_seconds)
171 .enable_logging(enable_logging)
172 .build())
173 }
174
175 pub fn validate(&self) -> Result<(), String> {
177 if self.url.is_empty() {
178 return Err("Database URL cannot be empty".to_string());
179 }
180
181 if self.max_connections == 0 {
182 return Err("Max connections must be greater than 0".to_string());
183 }
184
185 if self.timeout_seconds == 0 {
186 return Err("Timeout must be greater than 0".to_string());
187 }
188
189 if !self.url.contains("://") && !self.url.starts_with("sqlite:") {
191 return Err(
192 "Database URL must contain a protocol (e.g., sqlite://, postgres://, mysql://)"
193 .to_string(),
194 );
195 }
196
197 Ok(())
198 }
199
200 pub fn database_type(&self) -> Option<DatabaseType> {
204 DatabaseType::from_url(&self.url)
205 }
206
207 pub fn validate_database_support(&self) -> Result<DatabaseType, String> {
212 let db_type = self.database_type().ok_or_else(|| {
213 format!(
214 "Unrecognized database URL scheme in '{}'. Expected sqlite:, postgres:, or mysql:",
215 self.url
216 )
217 })?;
218
219 if !db_type.is_feature_enabled() {
220 return Err(format!(
221 "{} database detected from URL but the '{}' feature is not enabled. \
222 Add `features = [\"{}\"]` to your a2a-rs dependency.",
223 db_type,
224 db_type.feature_name(),
225 db_type.feature_name(),
226 ));
227 }
228
229 Ok(db_type)
230 }
231}
232
233#[cfg(feature = "sqlx-storage")]
234impl Default for DatabaseConfig {
235 fn default() -> Self {
236 Self::builder().url("sqlite::memory:".to_string()).build()
237 }
238}
239
240#[cfg(test)]
241#[cfg(feature = "sqlx-storage")]
242mod tests {
243 use super::*;
244
245 #[test]
246 fn test_database_config_validation() {
247 let config = DatabaseConfig::builder()
249 .url("sqlite:test.db".to_string())
250 .build();
251 assert!(config.validate().is_ok());
252
253 let config = DatabaseConfig::builder().url("".to_string()).build();
255 assert!(config.validate().is_err());
256
257 let config = DatabaseConfig::builder()
259 .url("sqlite:test.db".to_string())
260 .max_connections(0)
261 .build();
262 assert!(config.validate().is_err());
263 }
264
265 #[test]
266 fn test_database_type_detection() {
267 let sqlite_config = DatabaseConfig::builder()
268 .url("sqlite:test.db".to_string())
269 .build();
270 assert_eq!(sqlite_config.database_type(), Some(DatabaseType::Sqlite));
271
272 let postgres_config = DatabaseConfig::builder()
273 .url("postgres://localhost/test".to_string())
274 .build();
275 assert_eq!(
276 postgres_config.database_type(),
277 Some(DatabaseType::Postgres)
278 );
279
280 let postgresql_config = DatabaseConfig::builder()
281 .url("postgresql://localhost/test".to_string())
282 .build();
283 assert_eq!(
284 postgresql_config.database_type(),
285 Some(DatabaseType::Postgres)
286 );
287
288 let mysql_config = DatabaseConfig::builder()
289 .url("mysql://localhost/test".to_string())
290 .build();
291 assert_eq!(mysql_config.database_type(), Some(DatabaseType::Mysql));
292
293 let unknown_config = DatabaseConfig::builder()
294 .url("http://localhost".to_string())
295 .build();
296 assert_eq!(unknown_config.database_type(), None);
297 }
298
299 #[test]
300 fn test_database_type_from_url() {
301 assert_eq!(
302 DatabaseType::from_url("sqlite::memory:"),
303 Some(DatabaseType::Sqlite)
304 );
305 assert_eq!(
306 DatabaseType::from_url("sqlite:data.db"),
307 Some(DatabaseType::Sqlite)
308 );
309 assert_eq!(
310 DatabaseType::from_url("postgres://user:pass@host/db"),
311 Some(DatabaseType::Postgres)
312 );
313 assert_eq!(
314 DatabaseType::from_url("postgresql://user:pass@host/db"),
315 Some(DatabaseType::Postgres)
316 );
317 assert_eq!(
318 DatabaseType::from_url("mysql://user:pass@host/db"),
319 Some(DatabaseType::Mysql)
320 );
321 assert_eq!(
322 DatabaseType::from_url("mariadb://user:pass@host/db"),
323 Some(DatabaseType::Mysql)
324 );
325 assert_eq!(DatabaseType::from_url("ftp://something"), None);
326 }
327
328 #[test]
329 fn test_examples() {
330 let examples = DatabaseConfig::examples();
331 assert!(examples.contains_key("sqlite_memory"));
332 assert!(examples.contains_key("postgres_dev"));
333
334 for (name, config) in examples {
336 assert!(
337 config.validate().is_ok(),
338 "Example '{}' failed validation",
339 name
340 );
341 }
342 }
343}