Aether/crates/aether-data/adapters/postgres/src/gemini_file_mappings.rs

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<StoredGeminiFileMapping, DataLayerError> {
        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::<i64, _>("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::<i64, _>("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<Option<StoredGeminiFileMapping>, 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<Option<StoredGeminiFileMapping>, 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<Option<StoredGeminiFileMapping>, 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<StoredGeminiFileMappingListPage, DataLayerError> {
        let total = build_list_count_query(query)
            .build_query_scalar::<i64>()
            .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<GeminiFileMappingStats, DataLayerError> {
        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::<i64, _>("total_mappings")
                .map_postgres_err()?,
        )
        .unwrap_or_default();
        let active_mappings = usize::try_from(
            totals
                .try_get::<i64, _>("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::<i64, _>("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<StoredGeminiFileMapping, 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
  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<Option<StoredGeminiFileMapping>, 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<bool, DataLayerError> {
        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<bool, DataLayerError> {
        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<bool, DataLayerError> {
        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<Option<StoredGeminiFileMapping>, 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<usize, DataLayerError> {
        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::<Postgres>::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::<Postgres>::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,
        );
    }
}