use async_trait::async_trait; use futures_util::TryStreamExt; use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row}; use aether_data_contracts::repository::gemini_file_mappings::{ GeminiFileMappingListQuery, GeminiFileMappingMimeTypeCount, GeminiFileMappingReadRepository, GeminiFileMappingStats, GeminiFileMappingWriteRepository, StoredGeminiFileMapping, StoredGeminiFileMappingListPage, UpsertGeminiFileMappingRecord, }; use aether_data_contracts::DataLayerError; use aether_data_query::{push_ci_contains_any, push_limit_offset, SqlDialect, WhereClause}; use crate::error::SqlxResultExt; #[derive(Debug, Clone)] pub struct SqlxGeminiFileMappingRepository { pool: PgPool, } impl SqlxGeminiFileMappingRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } fn map_row(row: &PgRow) -> Result { Ok(StoredGeminiFileMapping { id: row.try_get("id").map_postgres_err()?, file_name: row.try_get("file_name").map_postgres_err()?, key_id: row.try_get("key_id").map_postgres_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_postgres_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_postgres_err()?, ) .map_err(|_| { DataLayerError::UnexpectedValue( "gemini_file_mappings.expires_at is invalid".to_string(), ) })?, }) } } #[async_trait] impl GeminiFileMappingReadRepository for SqlxGeminiFileMappingRepository { 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, EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms, EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs FROM gemini_file_mappings WHERE file_name = $1 "#, ) .bind(file_name) .fetch_optional(&self.pool) .await .map_postgres_err()?; match row { Some(row) => Ok(Some(Self::map_row(&row)?)), None => Ok(None), } } async fn find_active_by_file_name_for_user( &self, file_name: &str, user_id: &str, now_unix_secs: u64, ) -> Result, DataLayerError> { let row = sqlx::query( r#" SELECT id, file_name, key_id, user_id, display_name, mime_type, source_hash, EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms, EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs FROM gemini_file_mappings WHERE file_name = $1 AND user_id = $2 AND expires_at > TO_TIMESTAMP($3::double precision) "#, ) .bind(file_name) .bind(user_id) .bind(now_unix_secs as f64) .fetch_optional(&self.pool) .await .map_postgres_err()?; row.as_ref().map(Self::map_row).transpose() } async fn find_active_by_file_name_for_owner( &self, file_name: &str, key_id: &str, user_id: &str, now_unix_secs: u64, ) -> Result, DataLayerError> { let row = sqlx::query( r#" SELECT id, file_name, key_id, user_id, display_name, mime_type, source_hash, EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms, EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs FROM gemini_file_mappings WHERE file_name = $1 AND key_id = $2 AND user_id = $3 AND expires_at > TO_TIMESTAMP($4::double precision) "#, ) .bind(file_name) .bind(key_id) .bind(user_id) .bind(now_unix_secs as f64) .fetch_optional(&self.pool) .await .map_postgres_err()?; row.as_ref().map(Self::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_postgres_err()?; let mut builder = build_list_rows_query(query); let built_query = builder.build(); let mut rows = built_query.fetch(&self.pool); let mut items = Vec::new(); while let Some(row) = rows.try_next().await.map_postgres_err()? { items.push(Self::map_row(&row)?); } 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(*)::bigint AS total_mappings, COUNT(*) FILTER (WHERE expires_at > TO_TIMESTAMP($1::double precision))::bigint AS active_mappings FROM gemini_file_mappings "#, ) .bind(now_unix_secs as f64) .fetch_one(&self.pool) .await .map_postgres_err()?; let total_mappings = usize::try_from( totals .try_get::("total_mappings") .map_postgres_err()?, ) .unwrap_or_default(); let active_mappings = usize::try_from( totals .try_get::("active_mappings") .map_postgres_err()?, ) .unwrap_or_default(); let mut by_mime_type_rows = sqlx::query( r#" SELECT COALESCE(NULLIF(TRIM(mime_type), ''), 'unknown') AS mime_type, COUNT(*)::bigint AS count FROM gemini_file_mappings WHERE expires_at > TO_TIMESTAMP($1::double precision) GROUP BY COALESCE(NULLIF(TRIM(mime_type), ''), 'unknown') ORDER BY mime_type ASC "#, ) .bind(now_unix_secs as f64) .fetch(&self.pool); let mut by_mime_type = Vec::new(); while let Some(row) = by_mime_type_rows.try_next().await.map_postgres_err()? { by_mime_type.push(GeminiFileMappingMimeTypeCount { mime_type: row.try_get("mime_type").map_postgres_err()?, count: usize::try_from(row.try_get::("count").map_postgres_err()?) .unwrap_or_default(), }); } Ok(GeminiFileMappingStats { total_mappings, active_mappings, expired_mappings: total_mappings.saturating_sub(active_mappings), by_mime_type, }) } } #[async_trait] impl GeminiFileMappingWriteRepository for SqlxGeminiFileMappingRepository { async fn upsert( &self, record: UpsertGeminiFileMappingRecord, ) -> Result { record.validate()?; let row = 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 ($1,$2,$3,$4,$5,$6,$7,NOW(),TO_TIMESTAMP($8::double precision)) 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 RETURNING id, file_name, key_id, user_id, display_name, mime_type, source_hash, EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms, EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs "#, ) .bind(record.id.clone()) .bind(record.file_name.clone()) .bind(record.key_id.clone()) .bind(record.user_id.clone()) .bind(record.display_name.clone()) .bind(record.mime_type.clone()) .bind(record.source_hash.clone()) .bind(record.expires_at_unix_secs as f64) .fetch_one(&self.pool) .await .map_postgres_err()?; Self::map_row(&row) } async fn upsert_if_owner_matches( &self, record: UpsertGeminiFileMappingRecord, ) -> Result, DataLayerError> { record.validate()?; let row = 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 ($1,$2,$3,$4,$5,$6,$7,NOW(),TO_TIMESTAMP($8::double precision)) ON CONFLICT (file_name) DO UPDATE SET display_name = EXCLUDED.display_name, mime_type = EXCLUDED.mime_type, source_hash = EXCLUDED.source_hash, expires_at = EXCLUDED.expires_at WHERE gemini_file_mappings.key_id = EXCLUDED.key_id AND gemini_file_mappings.user_id IS NOT DISTINCT FROM EXCLUDED.user_id RETURNING id, file_name, key_id, user_id, display_name, mime_type, source_hash, EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms, EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs "#, ) .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(record.expires_at_unix_secs as f64) .fetch_optional(&self.pool) .await .map_postgres_err()?; row.as_ref().map(Self::map_row).transpose() } async fn delete_by_file_name(&self, file_name: &str) -> Result { let result = sqlx::query( r#" DELETE FROM gemini_file_mappings WHERE file_name = $1 #"#, ) .bind(file_name) .execute(&self.pool) .await .map_postgres_err()?; Ok(result.rows_affected() > 0) } async fn delete_by_file_name_for_user( &self, file_name: &str, user_id: &str, ) -> Result { let result = sqlx::query( r#" DELETE FROM gemini_file_mappings WHERE file_name = $1 AND user_id = $2 "#, ) .bind(file_name) .bind(user_id) .execute(&self.pool) .await .map_postgres_err()?; Ok(result.rows_affected() > 0) } async fn delete_by_file_name_for_owner( &self, file_name: &str, key_id: &str, user_id: &str, ) -> Result { let result = sqlx::query( r#" DELETE FROM gemini_file_mappings WHERE file_name = $1 AND key_id = $2 AND user_id = $3 "#, ) .bind(file_name) .bind(key_id) .bind(user_id) .execute(&self.pool) .await .map_postgres_err()?; Ok(result.rows_affected() > 0) } async fn delete_by_id( &self, mapping_id: &str, ) -> Result, DataLayerError> { let row = sqlx::query( r#" DELETE FROM gemini_file_mappings WHERE id = $1 RETURNING id, file_name, key_id, user_id, display_name, mime_type, source_hash, EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms, EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs "#, ) .bind(mapping_id) .fetch_optional(&self.pool) .await .map_postgres_err()?; match row { Some(row) => Ok(Some(Self::map_row(&row)?)), None => Ok(None), } } async fn delete_expired_before(&self, now_unix_secs: u64) -> Result { let result = sqlx::query( r#" DELETE FROM gemini_file_mappings WHERE expires_at <= TO_TIMESTAMP($1::double precision) "#, ) .bind(now_unix_secs as f64) .execute(&self.pool) .await .map_postgres_err()?; Ok(usize::try_from(result.rows_affected()).unwrap_or_default()) } } fn build_list_count_query(query: &GeminiFileMappingListQuery) -> QueryBuilder<'_, Postgres> { let mut builder = QueryBuilder::::new("SELECT COUNT(*)::bigint AS total FROM gemini_file_mappings"); let mut where_clause = WhereClause::new(); apply_list_filters(&mut builder, &mut where_clause, query); builder } fn build_list_rows_query(query: &GeminiFileMappingListQuery) -> QueryBuilder<'_, Postgres> { let mut builder = QueryBuilder::::new( r#" SELECT id, file_name, key_id, user_id, display_name, mime_type, source_hash, EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms, EXTRACT(EPOCH FROM expires_at)::bigint AS expires_at_unix_secs FROM gemini_file_mappings "#, ); let mut where_clause = WhereClause::new(); apply_list_filters(&mut builder, &mut where_clause, query); builder.push(" ORDER BY created_at DESC, file_name ASC"); push_limit_offset( &mut builder, i64::try_from(query.limit).unwrap_or(i64::MAX), i64::try_from(query.offset).unwrap_or(i64::MAX), ); builder } fn apply_list_filters( builder: &mut QueryBuilder<'_, Postgres>, where_clause: &mut WhereClause, query: &GeminiFileMappingListQuery, ) { if let Some(user_id) = query.user_id.as_deref() { where_clause.push_next(builder); builder.push("user_id = "); builder.push_bind(user_id.to_string()); } if !query.include_expired { where_clause.push_next(builder); builder.push("expires_at > TO_TIMESTAMP("); builder.push_bind(query.now_unix_secs as f64); builder.push("::double precision)"); } if let Some(search) = query .search .as_deref() .map(str::trim) .filter(|value| !value.is_empty()) { push_ci_contains_any( builder, where_clause, SqlDialect::Postgres, &["file_name", "COALESCE(display_name, '')"], search, ); } }