cedros_data/store/

mod.rs

#[cfg(feature = "storage")]
mod assets;
mod contracts;
mod defaults;
mod entries;
mod imex;
pub(crate) mod metered_reads;
mod migrate;
mod schema;
mod site;

use std::{collections::BTreeMap, sync::Arc, time::Instant};

use sqlx::{postgres::PgPoolOptions, Executor, PgPool, Postgres, Row, Transaction};
use tokio::sync::RwLock;

use crate::config::CedrosDataConfig;
use crate::error::{CedrosDataError, Result};
use crate::models::{Collection, CollectionMode, ContractSchema, Site};

use self::entries::typed_table_columns_with;

type TypedTableColumns = BTreeMap<String, bool>;
type TypedTableCacheKey = (String, String);
type TypedTableColumnsCache = BTreeMap<TypedTableCacheKey, TypedTableColumns>;

#[derive(Clone)]
pub struct CedrosData {
    pool: PgPool,
    typed_table_columns_cache: Arc<RwLock<TypedTableColumnsCache>>,
    /// Cached (timestamp, value) for freeReadsPerMonth with 5-min TTL.
    free_reads_cache: Arc<RwLock<Option<(Instant, i64)>>>,
}

impl CedrosData {
    pub async fn from_env() -> Result<Self> {
        let config = CedrosDataConfig::from_env()?;
        Self::connect(&config.postgres_uri).await
    }

    pub async fn connect(postgres_uri: &str) -> Result<Self> {
        let pool = PgPoolOptions::new()
            .max_connections(10)
            .connect(postgres_uri)
            .await?;
        Ok(Self {
            pool,
            typed_table_columns_cache: Arc::new(RwLock::new(BTreeMap::new())),
            free_reads_cache: Arc::new(RwLock::new(None)),
        })
    }

    pub fn pool(&self) -> &PgPool {
        &self.pool
    }

    async fn current_site(&self) -> Result<Site> {
        self.ensure_site_exists().await?;
        load_current_site_with(&self.pool).await
    }

    async fn find_collection(&self, collection_name: &str) -> Result<Collection> {
        find_collection_with(&self.pool, collection_name).await
    }

    async fn latest_contract(&self, collection_name: &str) -> Result<Option<ContractSchema>> {
        latest_contract_with(&self.pool, collection_name).await
    }

    pub(super) async fn typed_table_columns_in_transaction(
        &self,
        tx: &mut Transaction<'_, Postgres>,
        schema_name: &str,
        table_name: &str,
    ) -> Result<TypedTableColumns> {
        let cache_key = typed_table_cache_key(schema_name, table_name);
        if let Some(columns) = self
            .typed_table_columns_cache
            .read()
            .await
            .get(&cache_key)
            .cloned()
        {
            return Ok(columns);
        }

        let columns = typed_table_columns_with(&mut **tx, schema_name, table_name).await?;
        self.typed_table_columns_cache
            .write()
            .await
            .insert(cache_key, columns.clone());
        Ok(columns)
    }

    pub(super) async fn validate_typed_table_compatibility_in_transaction(
        &self,
        tx: &mut Transaction<'_, Postgres>,
        schema_name: &str,
        table_name: &str,
    ) -> Result<()> {
        self.typed_table_columns_in_transaction(tx, schema_name, table_name)
            .await
            .map(|_| ())
    }

    pub(super) async fn invalidate_typed_table_columns_cache(
        &self,
        schema_name: &str,
        table_name: &str,
    ) {
        let mut cache = self.typed_table_columns_cache.write().await;
        invalidate_typed_table_columns_cache_entry(&mut cache, schema_name, table_name);
    }

    pub(super) async fn invalidate_typed_table_columns_cache_for_schema(&self, schema_name: &str) {
        let mut cache = self.typed_table_columns_cache.write().await;
        invalidate_typed_table_columns_cache_entries_for_schema(&mut cache, schema_name);
    }
}

fn typed_table_cache_key(schema_name: &str, table_name: &str) -> TypedTableCacheKey {
    (schema_name.to_string(), table_name.to_string())
}

fn invalidate_typed_table_columns_cache_entry(
    cache: &mut TypedTableColumnsCache,
    schema_name: &str,
    table_name: &str,
) {
    cache.remove(&typed_table_cache_key(schema_name, table_name));
}

fn invalidate_typed_table_columns_cache_entries_for_schema(
    cache: &mut TypedTableColumnsCache,
    schema_name: &str,
) {
    cache.retain(|(cached_schema, _), _| cached_schema != schema_name);
}

pub(super) async fn load_site_if_configured_with<'e, E>(executor: E) -> Result<Option<Site>>
where
    E: Executor<'e, Database = Postgres>,
{
    let row = sqlx::query(
        "SELECT display_name, metadata
         FROM site
         WHERE id = 1",
    )
    .fetch_optional(executor)
    .await?;

    row.map(map_site_row).transpose()
}

pub(super) async fn load_current_site_with<'e, E>(executor: E) -> Result<Site>
where
    E: Executor<'e, Database = Postgres>,
{
    load_site_if_configured_with(executor)
        .await?
        .ok_or(CedrosDataError::SiteNotConfigured)
}

pub(super) async fn find_collection_with<'e, E>(
    executor: E,
    collection_name: &str,
) -> Result<Collection>
where
    E: Executor<'e, Database = Postgres>,
{
    let row = sqlx::query(
        "SELECT collection_name, mode, table_name, strict_contract
         FROM collections
         WHERE collection_name = $1",
    )
    .bind(collection_name)
    .fetch_optional(executor)
    .await?;

    row.map(map_collection_row)
        .transpose()?
        .ok_or_else(|| CedrosDataError::CollectionNotFound(collection_name.to_string()))
}

pub(super) async fn latest_contract_with<'e, E>(
    executor: E,
    collection_name: &str,
) -> Result<Option<ContractSchema>>
where
    E: Executor<'e, Database = Postgres>,
{
    let row = sqlx::query(
        "SELECT contract
         FROM collection_contracts
         WHERE collection_name = $1
         ORDER BY version DESC
         LIMIT 1",
    )
    .bind(collection_name)
    .fetch_optional(executor)
    .await?;

    let Some(row) = row else {
        return Ok(None);
    };

    let contract_value = row.get::<serde_json::Value, _>("contract");
    Ok(Some(serde_json::from_value(contract_value)?))
}

pub(super) async fn next_contract_version_with<'e, E>(
    executor: E,
    collection_name: &str,
) -> Result<i32>
where
    E: Executor<'e, Database = Postgres>,
{
    let row = sqlx::query(
        "SELECT COALESCE(MAX(version), 0) AS version
         FROM collection_contracts
         WHERE collection_name = $1",
    )
    .bind(collection_name)
    .fetch_one(executor)
    .await?;

    Ok(row.get::<i32, _>("version") + 1)
}

pub(super) async fn insert_contract_in_transaction(
    tx: &mut Transaction<'_, Postgres>,
    collection_name: &str,
    contract: &ContractSchema,
) -> Result<i32> {
    lock_collection_for_contracts(tx, collection_name).await?;
    let version = next_contract_version_with(&mut **tx, collection_name).await?;
    insert_contract_version_in_transaction(tx, collection_name, version, contract).await?;
    Ok(version)
}

pub(super) async fn insert_contract_version_in_transaction(
    tx: &mut Transaction<'_, Postgres>,
    collection_name: &str,
    version: i32,
    contract: &ContractSchema,
) -> Result<()> {
    let existing = sqlx::query(
        "SELECT contract
         FROM collection_contracts
         WHERE collection_name = $1 AND version = $2",
    )
    .bind(collection_name)
    .bind(version)
    .fetch_optional(&mut **tx)
    .await?;

    if let Some(existing_row) = existing {
        let existing_contract =
            serde_json::from_value(existing_row.get::<serde_json::Value, _>("contract"))?;
        ensure_matching_contract_version(collection_name, version, &existing_contract, contract)?;
        return Ok(());
    }

    sqlx::query(
        "INSERT INTO collection_contracts (collection_name, version, contract)
         VALUES ($1, $2, $3)",
    )
    .bind(collection_name)
    .bind(version)
    .bind(serde_json::to_value(contract)?)
    .execute(&mut **tx)
    .await?;
    Ok(())
}

async fn lock_collection_for_contracts(
    tx: &mut Transaction<'_, Postgres>,
    collection_name: &str,
) -> Result<()> {
    let row = sqlx::query(
        "SELECT 1
         FROM collections
         WHERE collection_name = $1
         FOR UPDATE",
    )
    .bind(collection_name)
    .fetch_optional(&mut **tx)
    .await?;

    if row.is_some() {
        return Ok(());
    }

    Err(CedrosDataError::InvalidRequest(format!(
        "collection {} is missing during contract version allocation",
        collection_name
    )))
}

fn ensure_matching_contract_version(
    collection_name: &str,
    version: i32,
    existing: &ContractSchema,
    incoming: &ContractSchema,
) -> Result<()> {
    if serde_json::to_value(existing)? == serde_json::to_value(incoming)? {
        return Ok(());
    }

    Err(CedrosDataError::InvalidRequest(format!(
        "import contract version conflict for {collection_name} v{version}"
    )))
}

pub(super) fn map_site_row(row: sqlx::postgres::PgRow) -> Result<Site> {
    Ok(Site {
        display_name: row.get("display_name"),
        metadata: row.get("metadata"),
    })
}

pub(super) fn map_collection_row(row: sqlx::postgres::PgRow) -> Result<Collection> {
    let mode = match row.get::<String, _>("mode").as_str() {
        "jsonb" => CollectionMode::Jsonb,
        "typed" => CollectionMode::Typed,
        other => {
            return Err(CedrosDataError::InvalidRequest(format!(
                "unknown collection mode: {other}"
            )));
        }
    };

    let strict_contract = row
        .get::<Option<serde_json::Value>, _>("strict_contract")
        .map(serde_json::from_value)
        .transpose()?;

    Ok(Collection {
        collection_name: row.get("collection_name"),
        mode,
        table_name: row.get("table_name"),
        strict_contract,
    })
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use crate::models::{ContractField, ContractSchema, ValueType};

    use super::{
        ensure_matching_contract_version, invalidate_typed_table_columns_cache_entries_for_schema,
        invalidate_typed_table_columns_cache_entry, typed_table_cache_key,
    };

    #[test]
    fn matching_contract_versions_are_idempotent() {
        let contract = sample_contract("title");
        assert!(ensure_matching_contract_version("articles", 2, &contract, &contract).is_ok());
    }

    #[test]
    fn conflicting_contract_versions_are_rejected() {
        let existing = sample_contract("title");
        let incoming = sample_contract("summary");

        assert!(ensure_matching_contract_version("articles", 2, &existing, &incoming).is_err());
    }

    #[test]
    fn invalidating_one_typed_table_cache_entry_preserves_other_entries() {
        let mut cache = sample_typed_table_cache();

        invalidate_typed_table_columns_cache_entry(&mut cache, "site_data", "pages");

        assert!(!cache.contains_key(&typed_table_cache_key("site_data", "pages")));
        assert!(cache.contains_key(&typed_table_cache_key("site_data", "articles")));
        assert!(cache.contains_key(&typed_table_cache_key("site_blog", "pages")));
    }

    #[test]
    fn invalidating_typed_table_cache_by_schema_removes_matching_entries() {
        let mut cache = sample_typed_table_cache();

        invalidate_typed_table_columns_cache_entries_for_schema(&mut cache, "site_data");

        assert!(!cache.contains_key(&typed_table_cache_key("site_data", "pages")));
        assert!(!cache.contains_key(&typed_table_cache_key("site_data", "articles")));
        assert!(cache.contains_key(&typed_table_cache_key("site_blog", "pages")));
    }

    fn sample_contract(path: &str) -> ContractSchema {
        ContractSchema {
            fields: vec![ContractField {
                path: path.to_string(),
                required: true,
                types: vec![ValueType::String],
            }],
        }
    }

    fn sample_typed_table_cache() -> BTreeMap<(String, String), BTreeMap<String, bool>> {
        BTreeMap::from([
            (
                typed_table_cache_key("site_data", "pages"),
                BTreeMap::from([(String::from("entry_key"), false)]),
            ),
            (
                typed_table_cache_key("site_data", "articles"),
                BTreeMap::from([(String::from("entry_key"), false)]),
            ),
            (
                typed_table_cache_key("site_blog", "pages"),
                BTreeMap::from([(String::from("entry_key"), false)]),
            ),
        ])
    }
}