Skip to main content

a2a_rs/adapter/storage/
database_config.rs

1//! Database configuration for SQLx storage
2
3#[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/// Supported database types, detected from the connection URL scheme.
11#[cfg(feature = "sqlx-storage")]
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum DatabaseType {
14    /// SQLite database (URLs starting with `sqlite:`)
15    Sqlite,
16    /// PostgreSQL database (URLs starting with `postgres:` or `postgresql:`)
17    Postgres,
18    /// MySQL database (URLs starting with `mysql:`)
19    Mysql,
20}
21
22#[cfg(feature = "sqlx-storage")]
23impl DatabaseType {
24    /// Detect the database type from a connection URL.
25    ///
26    /// Returns `None` if the URL scheme is not recognized.
27    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    /// Check whether this database type is supported by the currently compiled
40    /// features.
41    ///
42    /// MySQL is always `false`: the URL scheme is recognized so the error can
43    /// name it, and there is no schema or adapter behind it. There used to be a
44    /// `mysql` feature, which compiled a driver nothing used and made this
45    /// return `true` for a backend the storage could not talk to.
46    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    /// Returns the feature flag name needed to enable this database type.
55    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/// Database configuration with connection examples
77#[derive(Debug, Clone, Builder, Serialize, Deserialize)]
78pub struct DatabaseConfig {
79    /// Database connection URL
80    pub url: String,
81    /// Maximum number of connections in the pool
82    #[builder(default = 10)]
83    pub max_connections: u32,
84    /// Connection timeout in seconds
85    #[builder(default = 30)]
86    pub timeout_seconds: u64,
87    /// Whether to enable SQL query logging
88    #[builder(default = false)]
89    pub enable_logging: bool,
90}
91
92#[cfg(feature = "sqlx-storage")]
93impl DatabaseConfig {
94    /// Example configurations for different environments and databases
95    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    /// Create a new configuration from environment variables
143    ///
144    /// Expected environment variables:
145    /// - `DATABASE_URL`: Required - the database connection URL
146    /// - `DATABASE_MAX_CONNECTIONS`: Optional - defaults to 10
147    /// - `DATABASE_TIMEOUT_SECONDS`: Optional - defaults to 30  
148    /// - `DATABASE_ENABLE_LOGGING`: Optional - defaults to false
149    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    /// Validate the configuration
176    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        // Basic URL validation
190        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    /// Get the database type from the URL.
201    ///
202    /// Returns `None` if the URL scheme is not recognized.
203    pub fn database_type(&self) -> Option<DatabaseType> {
204        DatabaseType::from_url(&self.url)
205    }
206
207    /// Validate that the database URL scheme matches a compiled feature.
208    ///
209    /// Returns an error if the URL scheme is unrecognized or if the corresponding
210    /// feature flag is not enabled.
211    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        // Valid config
248        let config = DatabaseConfig::builder()
249            .url("sqlite:test.db".to_string())
250            .build();
251        assert!(config.validate().is_ok());
252
253        // Empty URL
254        let config = DatabaseConfig::builder().url("".to_string()).build();
255        assert!(config.validate().is_err());
256
257        // Invalid max connections
258        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        // Validate all examples
335        for (name, config) in examples {
336            assert!(
337                config.validate().is_ok(),
338                "Example '{}' failed validation",
339                name
340            );
341        }
342    }
343}