PiCloud/crates/manager-core/src/kv_repo.rs

//! Low-level Postgres CRUD over `kv_entries`. Stays storage-only;
//! authorization, event emission, and empty-collection validation live
//! one layer up in `KvServiceImpl`.

use async_trait::async_trait;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use base64::Engine as _;
use picloud_shared::{AppId, KvListPage};
use sqlx::PgPool;

#[derive(Debug, thiserror::Error)]
pub enum KvRepoError {
    #[error("database error: {0}")]
    Db(#[from] sqlx::Error),

    #[error("invalid pagination cursor")]
    InvalidCursor,
}

/// Repo surface. The trait is exposed so tests can substitute an
/// in-memory backing without spinning up Postgres.
#[async_trait]
pub trait KvRepo: Send + Sync {
    async fn get(
        &self,
        app_id: AppId,
        collection: &str,
        key: &str,
    ) -> Result<Option<serde_json::Value>, KvRepoError>;

    /// Upserts the row. Returns the previous value (if any) so callers
    /// can determine whether this was an `insert` or an `update` for
    /// the emitted `ServiceEvent`.
    async fn set(
        &self,
        app_id: AppId,
        collection: &str,
        key: &str,
        value: serde_json::Value,
    ) -> Result<Option<serde_json::Value>, KvRepoError>;

    /// Returns the deleted value if present, `None` if the row didn't
    /// exist. The caller turns the `bool was-present` part into the
    /// SDK's return value; the `Option<value>` part feeds the
    /// `old_payload` field of the emitted delete event.
    async fn delete(
        &self,
        app_id: AppId,
        collection: &str,
        key: &str,
    ) -> Result<Option<serde_json::Value>, KvRepoError>;

    async fn has(&self, app_id: AppId, collection: &str, key: &str) -> Result<bool, KvRepoError>;

    async fn list(
        &self,
        app_id: AppId,
        collection: &str,
        cursor: Option<&str>,
        limit: u32,
    ) -> Result<KvListPage, KvRepoError>;
}

pub struct PostgresKvRepo {
    pool: PgPool,
}

impl PostgresKvRepo {
    #[must_use]
    pub fn new(pool: PgPool) -> Self {
        Self { pool }
    }
}

/// Hard ceiling on `list` page size — scripts that pass anything larger
/// silently get clamped to this. Cursor-style pagination keeps a single
/// request bounded; clients fetch the next page via the returned cursor.
const KV_LIST_MAX_LIMIT: u32 = 1_000;
const KV_LIST_DEFAULT_LIMIT: u32 = 100;

#[async_trait]
impl KvRepo for PostgresKvRepo {
    async fn get(
        &self,
        app_id: AppId,
        collection: &str,
        key: &str,
    ) -> Result<Option<serde_json::Value>, KvRepoError> {
        let row: Option<(serde_json::Value,)> = sqlx::query_as(
            "SELECT value FROM kv_entries \
             WHERE app_id = $1 AND collection = $2 AND key = $3",
        )
        .bind(app_id.into_inner())
        .bind(collection)
        .bind(key)
        .fetch_optional(&self.pool)
        .await?;
        Ok(row.map(|(v,)| v))
    }

    async fn set(
        &self,
        app_id: AppId,
        collection: &str,
        key: &str,
        value: serde_json::Value,
    ) -> Result<Option<serde_json::Value>, KvRepoError> {
        // `RETURNING` after `ON CONFLICT DO UPDATE` exposes the old
        // value via the `xmax`/old-row trick: capture the prior value
        // with a CTE so callers know whether this was insert vs update.
        let row: Option<(Option<serde_json::Value>,)> = sqlx::query_as(
            "WITH prev AS (\
                SELECT value FROM kv_entries \
                WHERE app_id = $1 AND collection = $2 AND key = $3\
             ), \
             upserted AS (\
                INSERT INTO kv_entries (app_id, collection, key, value) \
                VALUES ($1, $2, $3, $4) \
                ON CONFLICT (app_id, collection, key) DO UPDATE \
                  SET value = EXCLUDED.value, updated_at = NOW() \
                RETURNING 1\
             ) \
             SELECT (SELECT value FROM prev) FROM upserted",
        )
        .bind(app_id.into_inner())
        .bind(collection)
        .bind(key)
        .bind(value)
        .fetch_optional(&self.pool)
        .await?;
        Ok(row.and_then(|(v,)| v))
    }

    async fn delete(
        &self,
        app_id: AppId,
        collection: &str,
        key: &str,
    ) -> Result<Option<serde_json::Value>, KvRepoError> {
        let row: Option<(serde_json::Value,)> = sqlx::query_as(
            "DELETE FROM kv_entries \
             WHERE app_id = $1 AND collection = $2 AND key = $3 \
             RETURNING value",
        )
        .bind(app_id.into_inner())
        .bind(collection)
        .bind(key)
        .fetch_optional(&self.pool)
        .await?;
        Ok(row.map(|(v,)| v))
    }

    async fn has(&self, app_id: AppId, collection: &str, key: &str) -> Result<bool, KvRepoError> {
        let row: Option<(i64,)> = sqlx::query_as(
            "SELECT 1 FROM kv_entries \
             WHERE app_id = $1 AND collection = $2 AND key = $3",
        )
        .bind(app_id.into_inner())
        .bind(collection)
        .bind(key)
        .fetch_optional(&self.pool)
        .await?;
        Ok(row.is_some())
    }

    async fn list(
        &self,
        app_id: AppId,
        collection: &str,
        cursor: Option<&str>,
        limit: u32,
    ) -> Result<KvListPage, KvRepoError> {
        let limit = if limit == 0 {
            KV_LIST_DEFAULT_LIMIT
        } else {
            limit.min(KV_LIST_MAX_LIMIT)
        };

        let last_key = match cursor {
            Some(c) => Some(decode_cursor(c)?),
            None => None,
        };

        // Keyset pagination: rows beyond `last_key` ordered by key.
        // `+1` to detect a "more pages" condition without a separate
        // COUNT query.
        let take = i64::from(limit) + 1;
        let rows: Vec<(String,)> = sqlx::query_as(
            "SELECT key FROM kv_entries \
             WHERE app_id = $1 AND collection = $2 \
               AND ($3::text IS NULL OR key > $3) \
             ORDER BY key ASC \
             LIMIT $4",
        )
        .bind(app_id.into_inner())
        .bind(collection)
        .bind(last_key.as_deref())
        .bind(take)
        .fetch_all(&self.pool)
        .await?;

        let mut keys: Vec<String> = rows.into_iter().map(|(k,)| k).collect();
        let next_cursor = if keys.len() > limit as usize {
            keys.truncate(limit as usize);
            keys.last().map(|k| encode_cursor(k))
        } else {
            None
        };

        Ok(KvListPage { keys, next_cursor })
    }
}

fn encode_cursor(last_key: &str) -> String {
    URL_SAFE_NO_PAD.encode(last_key.as_bytes())
}

fn decode_cursor(cursor: &str) -> Result<String, KvRepoError> {
    let bytes = URL_SAFE_NO_PAD
        .decode(cursor)
        .map_err(|_| KvRepoError::InvalidCursor)?;
    String::from_utf8(bytes).map_err(|_| KvRepoError::InvalidCursor)
}