cedros_data/store/

schema.rs

use std::collections::{BTreeMap, BTreeSet};

use sqlx::Row;
use sqlx::{Postgres, Transaction};

use crate::custom_schema::{
    apply_additive_schema, diff_custom_schema, parse_ref_table, validate_custom_schema,
};
use crate::defaults::SITE_SCHEMA_NAME;
use crate::error::{CedrosDataError, Result};
use crate::models::{CustomSchemaApplyReport, CustomSchemaDefinition, RegisterCustomSchemaRequest};

use super::CedrosData;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SchemaApplyDecision {
    Noop { version: i32 },
    Breaking { version: i32 },
    Apply { version: i32 },
}

impl CedrosData {
    pub async fn register_custom_schema(
        &self,
        request: RegisterCustomSchemaRequest,
    ) -> Result<CustomSchemaApplyReport> {
        self.ensure_site_exists().await?;
        let mut tx = self.pool().begin().await?;
        let report = self
            .register_custom_schema_in_transaction(&mut tx, &request.definition)
            .await?;
        if report.applied {
            tx.commit().await?;
        }
        Ok(report)
    }

    pub(super) async fn register_custom_schema_in_transaction(
        &self,
        tx: &mut Transaction<'_, Postgres>,
        definition: &CustomSchemaDefinition,
    ) -> Result<CustomSchemaApplyReport> {
        validate_custom_schema(definition)?;

        let (current_version, current_definition) = load_current_state(tx).await?;
        validate_foreign_key_targets(
            tx,
            SITE_SCHEMA_NAME,
            current_definition.as_ref(),
            definition,
        )
        .await?;
        let diff = diff_custom_schema(current_definition.as_ref(), definition);
        let next_version = match schema_apply_decision(
            current_version,
            !diff.additive_changes.is_empty(),
            !diff.breaking_changes.is_empty(),
        ) {
            SchemaApplyDecision::Noop { version } => {
                return Ok(CustomSchemaApplyReport {
                    applied: true,
                    version,
                    additive_changes: Vec::new(),
                    breaking_changes: Vec::new(),
                    generated_sql: Vec::new(),
                });
            }
            SchemaApplyDecision::Breaking { version } => {
                return Ok(CustomSchemaApplyReport {
                    applied: false,
                    version,
                    additive_changes: diff.additive_changes,
                    breaking_changes: diff.breaking_changes,
                    generated_sql: Vec::new(),
                });
            }
            SchemaApplyDecision::Apply { version } => version,
        };

        let generated_sql = apply_additive_schema(
            tx,
            SITE_SCHEMA_NAME,
            current_definition.as_ref(),
            definition,
        )
        .await?;
        self.invalidate_typed_table_columns_cache_for_schema(SITE_SCHEMA_NAME)
            .await;
        let report = CustomSchemaApplyReport {
            applied: true,
            version: next_version,
            additive_changes: diff.additive_changes,
            breaking_changes: Vec::new(),
            generated_sql: generated_sql.clone(),
        };

        upsert_schema_state(tx, next_version, definition).await?;
        record_schema_migration(tx, next_version, &report, &generated_sql).await?;
        Ok(report)
    }
}

fn schema_apply_decision(
    current_version: i32,
    has_additive_changes: bool,
    has_breaking_changes: bool,
) -> SchemaApplyDecision {
    if !has_additive_changes && !has_breaking_changes {
        return SchemaApplyDecision::Noop {
            version: current_version,
        };
    }

    let next_version = current_version + 1;
    if has_breaking_changes {
        return SchemaApplyDecision::Breaking {
            version: next_version,
        };
    }

    SchemaApplyDecision::Apply {
        version: next_version,
    }
}

async fn load_current_state(
    tx: &mut Transaction<'_, Postgres>,
) -> Result<(i32, Option<CustomSchemaDefinition>)> {
    let row = sqlx::query(
        "SELECT version, definition
         FROM cedros_data.custom_schema_state
         WHERE id = 1",
    )
    .fetch_optional(&mut **tx)
    .await?;

    parse_current_state(row)
}

fn parse_current_state(
    row: Option<sqlx::postgres::PgRow>,
) -> Result<(i32, Option<CustomSchemaDefinition>)> {
    let Some(row) = row else {
        return Ok((0, None));
    };

    let definition_value = row.get::<serde_json::Value, _>("definition");
    let definition = serde_json::from_value(definition_value)?;
    Ok((row.get::<i32, _>("version"), Some(definition)))
}

async fn validate_foreign_key_targets(
    tx: &mut Transaction<'_, Postgres>,
    schema_name: &str,
    current_definition: Option<&CustomSchemaDefinition>,
    incoming_definition: &CustomSchemaDefinition,
) -> Result<()> {
    let incoming_tables = table_columns_by_name(incoming_definition);
    let current_tables = current_definition
        .map(table_columns_by_name)
        .unwrap_or_default();
    let mut db_tables = BTreeMap::new();

    for table in &incoming_definition.tables {
        for foreign_key in &table.foreign_keys {
            let (ref_schema, ref_table) = parse_ref_table(schema_name, &foreign_key.ref_table);
            let ref_columns = match local_ref_table_columns(
                schema_name,
                &ref_schema,
                &ref_table,
                &incoming_tables,
                &current_tables,
            ) {
                Some(columns) => columns.clone(),
                None => load_table_columns(tx, &ref_schema, &ref_table, &mut db_tables).await?,
            };

            if ref_columns.is_empty() {
                return Err(CedrosDataError::InvalidRequest(format!(
                    "foreign key references unknown table {}.{}",
                    ref_schema, ref_table
                )));
            }

            for ref_column in &foreign_key.ref_columns {
                if ref_columns.contains(ref_column) {
                    continue;
                }
                return Err(CedrosDataError::InvalidRequest(format!(
                    "foreign key references unknown column {}.{}",
                    ref_table, ref_column
                )));
            }
        }
    }

    Ok(())
}

fn table_columns_by_name(
    definition: &CustomSchemaDefinition,
) -> BTreeMap<String, BTreeSet<String>> {
    definition
        .tables
        .iter()
        .map(|table| {
            (
                table.name.clone(),
                table
                    .columns
                    .iter()
                    .map(|column| column.name.clone())
                    .collect::<BTreeSet<String>>(),
            )
        })
        .collect()
}

fn local_ref_table_columns<'a>(
    schema_name: &str,
    ref_schema: &str,
    ref_table: &str,
    incoming_tables: &'a BTreeMap<String, BTreeSet<String>>,
    current_tables: &'a BTreeMap<String, BTreeSet<String>>,
) -> Option<&'a BTreeSet<String>> {
    if ref_schema != schema_name {
        return None;
    }

    incoming_tables
        .get(ref_table)
        .or_else(|| current_tables.get(ref_table))
}

async fn load_table_columns(
    tx: &mut Transaction<'_, Postgres>,
    table_schema: &str,
    table_name: &str,
    cache: &mut BTreeMap<(String, String), BTreeSet<String>>,
) -> Result<BTreeSet<String>> {
    let cache_key = (table_schema.to_string(), table_name.to_string());
    if let Some(columns) = cache.get(&cache_key) {
        return Ok(columns.clone());
    }

    let rows = sqlx::query(
        "SELECT column_name
         FROM information_schema.columns
         WHERE table_schema = $1 AND table_name = $2",
    )
    .bind(table_schema)
    .bind(table_name)
    .fetch_all(&mut **tx)
    .await?;

    let columns = rows
        .into_iter()
        .map(|row| row.get::<String, _>("column_name"))
        .collect::<BTreeSet<String>>();
    cache.insert(cache_key, columns.clone());
    Ok(columns)
}

async fn upsert_schema_state(
    tx: &mut Transaction<'_, Postgres>,
    version: i32,
    definition: &CustomSchemaDefinition,
) -> Result<()> {
    sqlx::query(
        "INSERT INTO cedros_data.custom_schema_state (id, version, definition, updated_at)
         VALUES (1, $1, $2, NOW())
         ON CONFLICT (id)
         DO UPDATE SET
             version = EXCLUDED.version,
             definition = EXCLUDED.definition,
             updated_at = NOW()",
    )
    .bind(version)
    .bind(serde_json::to_value(definition)?)
    .execute(&mut **tx)
    .await?;
    Ok(())
}

async fn record_schema_migration(
    tx: &mut Transaction<'_, Postgres>,
    version: i32,
    report: &CustomSchemaApplyReport,
    generated_sql: &[String],
) -> Result<()> {
    sqlx::query(
        "INSERT INTO cedros_data.custom_schema_migrations (version, report, generated_sql)
         VALUES ($1, $2, $3)
         ON CONFLICT (version)
         DO NOTHING",
    )
    .bind(version)
    .bind(serde_json::to_value(report)?)
    .bind(serde_json::to_value(generated_sql)?)
    .execute(&mut **tx)
    .await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::{schema_apply_decision, SchemaApplyDecision};

    #[test]
    fn schema_apply_decision_keeps_version_for_noop_diff() {
        assert_eq!(
            schema_apply_decision(3, false, false),
            SchemaApplyDecision::Noop { version: 3 }
        );
    }

    #[test]
    fn schema_apply_decision_advances_version_for_breaking_diff() {
        assert_eq!(
            schema_apply_decision(3, false, true),
            SchemaApplyDecision::Breaking { version: 4 }
        );
    }

    #[test]
    fn schema_apply_decision_advances_version_for_additive_diff() {
        assert_eq!(
            schema_apply_decision(3, true, false),
            SchemaApplyDecision::Apply { version: 4 }
        );
    }
}