use async_trait::async_trait; use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite}; use super::types::{ GeminiFileMappingListQuery, GeminiFileMappingMimeTypeCount, GeminiFileMappingReadRepository, GeminiFileMappingStats, GeminiFileMappingWriteRepository, StoredGeminiFileMapping, StoredGeminiFileMappingListPage, UpsertGeminiFileMappingRecord, }; use crate::driver::sqlite::SqlitePool; use crate::error::SqlResultExt; use crate::DataLayerError; #[derive(Debug, Clone)] pub struct SqliteGeminiFileMappingRepository { pool: SqlitePool, } impl SqliteGeminiFileMappingRepository { pub fn new(pool: SqlitePool) -> Self { Self { pool } } async fn reload_by_file_name( &self, file_name: &str, ) -> Result { self.find_by_file_name(file_name).await?.ok_or_else(|| { DataLayerError::UnexpectedValue("gemini file mapping missing after write".to_string()) }) } } #[async_trait] impl GeminiFileMappingReadRepository for SqliteGeminiFileMappingRepository { async fn find_by_file_name( &self, file_name: &str, ) -> Result, DataLayerError> { let row = sqlx::query( r#" SELECT id, file_name, key_id, user_id, display_name, mime_type, source_hash, created_at AS created_at_unix_ms, expires_at AS expires_at_unix_secs FROM gemini_file_mappings WHERE file_name = ? LIMIT 1 "#, ) .bind(file_name) .fetch_optional(&self.pool) .await .map_sql_err()?; row.as_ref().map(map_row).transpose() } async fn list_mappings( &self, query: &GeminiFileMappingListQuery, ) -> Result { let total = build_list_count_query(query) .build_query_scalar::() .fetch_one(&self.pool) .await .map_sql_err()?; let rows = build_list_rows_query(query) .build() .fetch_all(&self.pool) .await .map_sql_err()?; let items = rows.iter().map(map_row).collect::, _>>()?; Ok(StoredGeminiFileMappingListPage { items, total: usize::try_from(total).unwrap_or_default(), }) } async fn summarize_mappings( &self, now_unix_secs: u64, ) -> Result { let totals = sqlx::query( r#" SELECT COUNT(*) AS total_mappings, SUM(CASE WHEN expires_at > ? THEN 1 ELSE 0 END) AS active_mappings FROM gemini_file_mappings "#, ) .bind(now_unix_secs as i64) .fetch_one(&self.pool) .await .map_sql_err()?; let total_mappings = usize::try_from(totals.try_get::("total_mappings").map_sql_err()?) .unwrap_or_default(); let active_mappings = usize::try_from( totals .try_get::, _>("active_mappings") .map_sql_err()? .unwrap_or(0), ) .unwrap_or_default(); let by_mime_type_rows = sqlx::query( r#" SELECT COALESCE(NULLIF(TRIM(mime_type), ''), 'unknown') AS mime_type, COUNT(*) AS count FROM gemini_file_mappings WHERE expires_at > ? GROUP BY COALESCE(NULLIF(TRIM(mime_type), ''), 'unknown') ORDER BY mime_type ASC "#, ) .bind(now_unix_secs as i64) .fetch_all(&self.pool) .await .map_sql_err()?; let by_mime_type = by_mime_type_rows .iter() .map(|row| { Ok(GeminiFileMappingMimeTypeCount { mime_type: row.try_get("mime_type").map_sql_err()?, count: usize::try_from(row.try_get::("count").map_sql_err()?) .unwrap_or_default(), }) }) .collect::, DataLayerError>>()?; Ok(GeminiFileMappingStats { total_mappings, active_mappings, expired_mappings: total_mappings.saturating_sub(active_mappings), by_mime_type, }) } } #[async_trait] impl GeminiFileMappingWriteRepository for SqliteGeminiFileMappingRepository { async fn upsert( &self, record: UpsertGeminiFileMappingRecord, ) -> Result { record.validate()?; sqlx::query( r#" INSERT INTO gemini_file_mappings ( id, file_name, key_id, user_id, display_name, mime_type, source_hash, created_at, expires_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(file_name) DO UPDATE SET key_id = excluded.key_id, user_id = excluded.user_id, display_name = excluded.display_name, mime_type = excluded.mime_type, source_hash = excluded.source_hash, expires_at = excluded.expires_at "#, ) .bind(&record.id) .bind(&record.file_name) .bind(&record.key_id) .bind(&record.user_id) .bind(&record.display_name) .bind(&record.mime_type) .bind(&record.source_hash) .bind(current_unix_secs() as i64) .bind(i64_from_u64( record.expires_at_unix_secs, "gemini_file_mappings.expires_at", )?) .execute(&self.pool) .await .map_sql_err()?; self.reload_by_file_name(&record.file_name).await } async fn delete_by_file_name(&self, file_name: &str) -> Result { let rows_affected = sqlx::query("DELETE FROM gemini_file_mappings WHERE file_name = ?") .bind(file_name) .execute(&self.pool) .await .map_sql_err()? .rows_affected(); Ok(rows_affected > 0) } async fn delete_by_id( &self, mapping_id: &str, ) -> Result, DataLayerError> { let existing = sqlx::query( r#" SELECT id, file_name, key_id, user_id, display_name, mime_type, source_hash, created_at AS created_at_unix_ms, expires_at AS expires_at_unix_secs FROM gemini_file_mappings WHERE id = ? LIMIT 1 "#, ) .bind(mapping_id) .fetch_optional(&self.pool) .await .map_sql_err()?; let Some(existing) = existing else { return Ok(None); }; sqlx::query("DELETE FROM gemini_file_mappings WHERE id = ?") .bind(mapping_id) .execute(&self.pool) .await .map_sql_err()?; Ok(Some(map_row(&existing)?)) } async fn delete_expired_before(&self, now_unix_secs: u64) -> Result { let rows_affected = sqlx::query("DELETE FROM gemini_file_mappings WHERE expires_at <= ?") .bind(now_unix_secs as i64) .execute(&self.pool) .await .map_sql_err()? .rows_affected(); Ok(usize::try_from(rows_affected).unwrap_or_default()) } } fn build_list_count_query(query: &GeminiFileMappingListQuery) -> QueryBuilder<'_, Sqlite> { let mut builder = QueryBuilder::::new("SELECT COUNT(*) AS total FROM gemini_file_mappings WHERE 1=1"); apply_list_filters(&mut builder, query); builder } fn build_list_rows_query(query: &GeminiFileMappingListQuery) -> QueryBuilder<'_, Sqlite> { let mut builder = QueryBuilder::::new( r#" SELECT id, file_name, key_id, user_id, display_name, mime_type, source_hash, created_at AS created_at_unix_ms, expires_at AS expires_at_unix_secs FROM gemini_file_mappings WHERE 1=1 "#, ); apply_list_filters(&mut builder, query); builder.push(" ORDER BY created_at DESC, file_name ASC LIMIT "); builder.push_bind(i64::try_from(query.limit).unwrap_or(i64::MAX)); builder.push(" OFFSET "); builder.push_bind(i64::try_from(query.offset).unwrap_or(i64::MAX)); builder } fn apply_list_filters(builder: &mut QueryBuilder<'_, Sqlite>, query: &GeminiFileMappingListQuery) { if !query.include_expired { builder.push(" AND expires_at > "); builder.push_bind(query.now_unix_secs as i64); } if let Some(search) = query .search .as_deref() .map(str::trim) .filter(|value| !value.is_empty()) { let pattern = format!("%{}%", search.to_ascii_lowercase()); builder.push(" AND (LOWER(file_name) LIKE "); builder.push_bind(pattern.clone()); builder.push(" OR LOWER(COALESCE(display_name, '')) LIKE "); builder.push_bind(pattern); builder.push(")"); } } fn current_unix_secs() -> u64 { chrono::Utc::now().timestamp().max(0) as u64 } fn i64_from_u64(value: u64, field_name: &str) -> Result { i64::try_from(value) .map_err(|_| DataLayerError::InvalidInput(format!("{field_name} exceeds i64: {value}"))) } fn map_row(row: &SqliteRow) -> Result { Ok(StoredGeminiFileMapping { id: row.try_get("id").map_sql_err()?, file_name: row.try_get("file_name").map_sql_err()?, key_id: row.try_get("key_id").map_sql_err()?, user_id: row.try_get("user_id").ok().flatten(), display_name: row.try_get("display_name").ok().flatten(), mime_type: row.try_get("mime_type").ok().flatten(), source_hash: row.try_get("source_hash").ok().flatten(), created_at_unix_ms: u64::try_from( row.try_get::("created_at_unix_ms").map_sql_err()?, ) .map_err(|_| { DataLayerError::UnexpectedValue( "gemini_file_mappings.created_at is invalid".to_string(), ) })?, expires_at_unix_secs: u64::try_from( row.try_get::("expires_at_unix_secs") .map_sql_err()?, ) .map_err(|_| { DataLayerError::UnexpectedValue( "gemini_file_mappings.expires_at is invalid".to_string(), ) })?, }) } #[cfg(test)] mod tests { use super::SqliteGeminiFileMappingRepository; use crate::lifecycle::migrate::run_sqlite_migrations; use crate::repository::gemini_file_mappings::{ GeminiFileMappingListQuery, GeminiFileMappingReadRepository, GeminiFileMappingWriteRepository, UpsertGeminiFileMappingRecord, }; #[tokio::test] async fn sqlite_repository_round_trips_gemini_file_mappings() { let pool = sqlx::sqlite::SqlitePoolOptions::new() .max_connections(1) .connect("sqlite::memory:") .await .expect("sqlite pool should connect"); run_sqlite_migrations(&pool) .await .expect("sqlite migrations should run"); let repository = SqliteGeminiFileMappingRepository::new(pool); let created = repository .upsert(UpsertGeminiFileMappingRecord { id: "mapping-1".to_string(), file_name: "files/example.png".to_string(), key_id: "key-1".to_string(), user_id: Some("user-1".to_string()), display_name: Some("Example".to_string()), mime_type: Some("image/png".to_string()), source_hash: Some("hash-1".to_string()), expires_at_unix_secs: 300, }) .await .expect("mapping should upsert"); assert_eq!(created.id, "mapping-1"); assert_eq!(created.mime_type, Some("image/png".to_string())); let updated = repository .upsert(UpsertGeminiFileMappingRecord { id: "mapping-replacement".to_string(), file_name: "files/example.png".to_string(), key_id: "key-2".to_string(), user_id: Some("user-2".to_string()), display_name: Some("Updated".to_string()), mime_type: Some("image/jpeg".to_string()), source_hash: Some("hash-2".to_string()), expires_at_unix_secs: 500, }) .await .expect("mapping should update"); assert_eq!(updated.id, "mapping-1"); assert_eq!(updated.key_id, "key-2"); let page = repository .list_mappings(&GeminiFileMappingListQuery { include_expired: false, search: Some("updated".to_string()), offset: 0, limit: 10, now_unix_secs: 400, }) .await .expect("mappings should list"); assert_eq!(page.total, 1); assert_eq!(page.items[0].file_name, "files/example.png"); let stats = repository .summarize_mappings(400) .await .expect("stats should load"); assert_eq!(stats.total_mappings, 1); assert_eq!(stats.active_mappings, 1); assert_eq!(stats.by_mime_type[0].mime_type, "image/jpeg"); assert_eq!( repository .delete_expired_before(600) .await .expect("expired mappings should delete"), 1 ); assert!(repository .find_by_file_name("files/example.png") .await .expect("find should run") .is_none()); } }