a2a_protocol_server/push/
sqlite_config_store.rs1use std::future::Future;
11use std::pin::Pin;
12
13use a2a_protocol_types::error::{A2aError, A2aResult};
14use a2a_protocol_types::push::TaskPushNotificationConfig;
15use sqlx::sqlite::SqlitePool;
16
17use super::config_store::PushConfigStore;
18
19#[derive(Debug, Clone)]
34pub struct SqlitePushConfigStore {
35 pool: SqlitePool,
36}
37
38async fn sqlite_pool(url: &str) -> Result<SqlitePool, sqlx::Error> {
40 use sqlx::sqlite::SqliteConnectOptions;
41 use std::str::FromStr;
42
43 let opts = SqliteConnectOptions::from_str(url)?
44 .pragma("journal_mode", "WAL")
45 .pragma("busy_timeout", "5000")
46 .pragma("synchronous", "NORMAL")
47 .pragma("foreign_keys", "ON")
48 .create_if_missing(true);
49
50 sqlx::sqlite::SqlitePoolOptions::new()
51 .max_connections(8)
52 .connect_with(opts)
53 .await
54}
55
56#[allow(clippy::needless_pass_by_value)]
58fn to_a2a_error(e: sqlx::Error) -> A2aError {
59 A2aError::internal(format!("sqlite error: {e}"))
60}
61
62impl SqlitePushConfigStore {
63 pub async fn new(url: &str) -> Result<Self, sqlx::Error> {
69 let pool = sqlite_pool(url).await?;
70 Self::from_pool(pool).await
71 }
72
73 pub async fn from_pool(pool: SqlitePool) -> Result<Self, sqlx::Error> {
79 sqlx::query(
80 "CREATE TABLE IF NOT EXISTS push_configs (
81 task_id TEXT NOT NULL,
82 id TEXT NOT NULL,
83 data TEXT NOT NULL,
84 PRIMARY KEY (task_id, id)
85 )",
86 )
87 .execute(&pool)
88 .await?;
89
90 Ok(Self { pool })
91 }
92}
93
94#[allow(clippy::manual_async_fn)]
95impl PushConfigStore for SqlitePushConfigStore {
96 fn set<'a>(
97 &'a self,
98 mut config: TaskPushNotificationConfig,
99 ) -> Pin<Box<dyn Future<Output = A2aResult<TaskPushNotificationConfig>> + Send + 'a>> {
100 Box::pin(async move {
101 let Some(task_id) = config.task_id.clone() else {
105 return Err(A2aError::invalid_params(
106 "taskId is required to store a push notification config",
107 ));
108 };
109 let id = config
110 .id
111 .clone()
112 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
113 config.id = Some(id.clone());
114
115 let data = serde_json::to_string(&config)
116 .map_err(|e| A2aError::internal(format!("serialize: {e}")))?;
117
118 sqlx::query(
119 "INSERT INTO push_configs (task_id, id, data)
120 VALUES (?1, ?2, ?3)
121 ON CONFLICT(task_id, id) DO UPDATE SET data = excluded.data",
122 )
123 .bind(&task_id)
124 .bind(&id)
125 .bind(&data)
126 .execute(&self.pool)
127 .await
128 .map_err(to_a2a_error)?;
129
130 Ok(config)
131 })
132 }
133
134 fn get<'a>(
135 &'a self,
136 task_id: &'a str,
137 id: &'a str,
138 ) -> Pin<Box<dyn Future<Output = A2aResult<Option<TaskPushNotificationConfig>>> + Send + 'a>>
139 {
140 Box::pin(async move {
141 let row: Option<(String,)> =
142 sqlx::query_as("SELECT data FROM push_configs WHERE task_id = ?1 AND id = ?2")
143 .bind(task_id)
144 .bind(id)
145 .fetch_optional(&self.pool)
146 .await
147 .map_err(to_a2a_error)?;
148
149 match row {
150 Some((data,)) => {
151 let config: TaskPushNotificationConfig = serde_json::from_str(&data)
152 .map_err(|e| A2aError::internal(format!("deserialize: {e}")))?;
153 Ok(Some(config))
154 }
155 None => Ok(None),
156 }
157 })
158 }
159
160 fn list<'a>(
161 &'a self,
162 task_id: &'a str,
163 ) -> Pin<Box<dyn Future<Output = A2aResult<Vec<TaskPushNotificationConfig>>> + Send + 'a>> {
164 Box::pin(async move {
165 let rows: Vec<(String,)> =
166 sqlx::query_as("SELECT data FROM push_configs WHERE task_id = ?1")
167 .bind(task_id)
168 .fetch_all(&self.pool)
169 .await
170 .map_err(to_a2a_error)?;
171
172 rows.into_iter()
173 .map(|(data,)| {
174 serde_json::from_str(&data)
175 .map_err(|e| A2aError::internal(format!("deserialize: {e}")))
176 })
177 .collect()
178 })
179 }
180
181 fn delete<'a>(
182 &'a self,
183 task_id: &'a str,
184 id: &'a str,
185 ) -> Pin<Box<dyn Future<Output = A2aResult<()>> + Send + 'a>> {
186 Box::pin(async move {
187 sqlx::query("DELETE FROM push_configs WHERE task_id = ?1 AND id = ?2")
188 .bind(task_id)
189 .bind(id)
190 .execute(&self.pool)
191 .await
192 .map_err(to_a2a_error)?;
193 Ok(())
194 })
195 }
196
197 fn count(&self) -> Pin<Box<dyn Future<Output = A2aResult<Option<usize>>> + Send + '_>> {
198 Box::pin(async move {
199 let (total,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM push_configs")
200 .fetch_one(&self.pool)
201 .await
202 .map_err(to_a2a_error)?;
203 Ok(Some(usize::try_from(total).unwrap_or(usize::MAX)))
204 })
205 }
206}
207
208#[cfg(test)]
209mod tests {
210 use super::*;
211 use a2a_protocol_types::push::TaskPushNotificationConfig;
212
213 async fn make_store() -> SqlitePushConfigStore {
214 SqlitePushConfigStore::new("sqlite::memory:")
215 .await
216 .expect("failed to create in-memory push config store")
217 }
218
219 fn make_config(task_id: &str, id: Option<&str>, url: &str) -> TaskPushNotificationConfig {
220 TaskPushNotificationConfig {
221 tenant: None,
222 id: id.map(String::from),
223 task_id: Some(task_id.to_string()),
224 url: url.to_string(),
225 token: None,
226 authentication: None,
227 }
228 }
229
230 #[tokio::test]
231 async fn set_assigns_id_when_none() {
232 let store = make_store().await;
233 let config = make_config("task-1", None, "https://example.com/hook");
234 let result = store.set(config).await.expect("set should succeed");
235 assert!(
236 result.id.is_some(),
237 "set should assign an id when none is provided"
238 );
239 }
240
241 #[tokio::test]
242 async fn count_reflects_total_configs_across_tasks() {
243 let store = make_store().await;
244 assert_eq!(store.count().await.unwrap(), Some(0));
245 store
246 .set(make_config("task-a", Some("c1"), "https://example.com/1"))
247 .await
248 .unwrap();
249 store
250 .set(make_config("task-b", Some("c1"), "https://example.com/2"))
251 .await
252 .unwrap();
253 assert_eq!(
254 store.count().await.unwrap(),
255 Some(2),
256 "count spans distinct task ids"
257 );
258 store.delete("task-a", "c1").await.unwrap();
259 assert_eq!(store.count().await.unwrap(), Some(1));
260 }
261
262 #[tokio::test]
263 async fn set_preserves_explicit_id() {
264 let store = make_store().await;
265 let config = make_config("task-1", Some("my-id"), "https://example.com/hook");
266 let result = store.set(config).await.expect("set should succeed");
267 assert_eq!(
268 result.id.as_deref(),
269 Some("my-id"),
270 "set should preserve the explicit id"
271 );
272 }
273
274 #[tokio::test]
275 async fn set_then_get_round_trip() {
276 let store = make_store().await;
277 let config = make_config("task-1", Some("cfg-1"), "https://example.com/hook");
278 store.set(config).await.unwrap();
279
280 let retrieved = store.get("task-1", "cfg-1").await.unwrap();
281 let retrieved = retrieved.expect("config should exist after set");
282 assert_eq!(retrieved.task_id.as_deref(), Some("task-1"));
283 assert_eq!(retrieved.url, "https://example.com/hook");
284 assert_eq!(retrieved.id.as_deref(), Some("cfg-1"));
285 }
286
287 #[tokio::test]
288 async fn get_returns_none_for_missing_config() {
289 let store = make_store().await;
290 let result = store
291 .get("no-task", "no-id")
292 .await
293 .expect("get should succeed");
294 assert!(
295 result.is_none(),
296 "get should return None for a missing config"
297 );
298 }
299
300 #[tokio::test]
301 async fn overwrite_existing_config() {
302 let store = make_store().await;
303 store
304 .set(make_config(
305 "task-1",
306 Some("cfg-1"),
307 "https://example.com/v1",
308 ))
309 .await
310 .unwrap();
311 store
312 .set(make_config(
313 "task-1",
314 Some("cfg-1"),
315 "https://example.com/v2",
316 ))
317 .await
318 .unwrap();
319
320 let retrieved = store.get("task-1", "cfg-1").await.unwrap().unwrap();
321 assert_eq!(
322 retrieved.url, "https://example.com/v2",
323 "overwrite should update the URL"
324 );
325 }
326
327 #[tokio::test]
328 async fn list_returns_empty_for_unknown_task() {
329 let store = make_store().await;
330 let configs = store.list("no-such-task").await.unwrap();
331 assert!(
332 configs.is_empty(),
333 "list should return empty vec for unknown task"
334 );
335 }
336
337 #[tokio::test]
338 async fn list_returns_only_configs_for_given_task() {
339 let store = make_store().await;
340 store
341 .set(make_config("task-a", Some("c1"), "https://a.com/1"))
342 .await
343 .unwrap();
344 store
345 .set(make_config("task-a", Some("c2"), "https://a.com/2"))
346 .await
347 .unwrap();
348 store
349 .set(make_config("task-b", Some("c3"), "https://b.com/1"))
350 .await
351 .unwrap();
352
353 let a_configs = store.list("task-a").await.unwrap();
354 assert_eq!(a_configs.len(), 2, "task-a should have exactly 2 configs");
355
356 let b_configs = store.list("task-b").await.unwrap();
357 assert_eq!(b_configs.len(), 1, "task-b should have exactly 1 config");
358 }
359
360 #[tokio::test]
361 async fn delete_removes_config() {
362 let store = make_store().await;
363 store
364 .set(make_config("task-1", Some("cfg-1"), "https://example.com"))
365 .await
366 .unwrap();
367
368 store
369 .delete("task-1", "cfg-1")
370 .await
371 .expect("delete should succeed");
372
373 let result = store.get("task-1", "cfg-1").await.unwrap();
374 assert!(result.is_none(), "config should be gone after delete");
375 }
376
377 #[tokio::test]
378 async fn delete_nonexistent_is_ok() {
379 let store = make_store().await;
380 let result = store.delete("no-task", "no-id").await;
381 assert!(
382 result.is_ok(),
383 "deleting a nonexistent config should not error"
384 );
385 }
386
387 #[tokio::test]
388 async fn delete_does_not_affect_other_configs() {
389 let store = make_store().await;
390 store
391 .set(make_config("task-1", Some("c1"), "https://a.com"))
392 .await
393 .unwrap();
394 store
395 .set(make_config("task-1", Some("c2"), "https://b.com"))
396 .await
397 .unwrap();
398
399 store.delete("task-1", "c1").await.unwrap();
400
401 let remaining = store.list("task-1").await.unwrap();
402 assert_eq!(
403 remaining.len(),
404 1,
405 "only the deleted config should be removed"
406 );
407 assert_eq!(remaining[0].id.as_deref(), Some("c2"));
408 }
409
410 #[test]
412 fn to_a2a_error_formats_message() {
413 let sqlite_err = sqlx::Error::RowNotFound;
414 let a2a_err = to_a2a_error(sqlite_err);
415 let msg = format!("{a2a_err}");
416 assert!(
417 msg.contains("sqlite error"),
418 "error message should contain 'sqlite error': {msg}"
419 );
420 }
421
422 #[tokio::test]
423 async fn multiple_tasks_independent_configs() {
424 let store = make_store().await;
425 store
427 .set(make_config("task-a", Some("cfg-1"), "https://a.com"))
428 .await
429 .unwrap();
430 store
431 .set(make_config("task-b", Some("cfg-1"), "https://b.com"))
432 .await
433 .unwrap();
434
435 let a = store.get("task-a", "cfg-1").await.unwrap().unwrap();
436 assert_eq!(a.url, "https://a.com");
437
438 let b = store.get("task-b", "cfg-1").await.unwrap().unwrap();
439 assert_eq!(b.url, "https://b.com");
440 }
441}