jspg/src/lib.rs

use pgrx::*;

pg_module_magic!();

use serde_json::{json, Value};
use std::{collections::HashMap, sync::RwLock};
use boon::{Compiler, Schemas, ValidationError, SchemaIndex, CompileError};
use lazy_static::lazy_static;

struct BoonCache {
  schemas: Schemas,
  id_to_index: HashMap<String, SchemaIndex>,
}

lazy_static! {
  static ref SCHEMA_CACHE: RwLock<BoonCache> = RwLock::new(BoonCache {
    schemas: Schemas::new(),
    id_to_index: HashMap::new(),
  });
}

#[pg_extern(strict)]
fn cache_json_schema(schema_id: &str, schema: JsonB) -> JsonB {
  let mut cache = SCHEMA_CACHE.write().unwrap();
  let schema_value: Value = schema.0;
  let schema_path = format!("urn:{}", schema_id);

  let mut compiler = Compiler::new();
  compiler.enable_format_assertions();

  // Use schema_path when adding the resource
  if let Err(e) = compiler.add_resource(&schema_path, schema_value.clone()) {
    return JsonB(json!({
      "success": false,
      "error": {
        "message": format!("Failed to add schema resource '{}': {}", schema_id, e),
        "schema_path": schema_path
      }
    }));
  }

  // Use schema_path when compiling
  match compiler.compile(&schema_path, &mut cache.schemas) {
    Ok(sch_index) => {
      // Store the index using the original schema_id as the key
      cache.id_to_index.insert(schema_id.to_string(), sch_index);
      JsonB(json!({ "success": true }))
    }
    Err(e) => {
      let error = match &e {
        CompileError::ValidationError { url: _url, src } => { 
          // Collect leaf errors from the meta-schema validation failure
          let mut error_list = Vec::new();
          collect_leaf_errors(src, &mut error_list);
          // Return the flat list directly
          json!(error_list)
        }
        _ => {
          // Keep existing handling for other compilation errors
          let _error_type = format!("{:?}", e).split('(').next().unwrap_or("Unknown").to_string(); 
          json!({
            "message": format!("Schema '{}' compilation failed: {}", schema_id, e),
            "schema_path": schema_path,
            "detail": format!("{:?}", e),
          })
        }
      };
      // Ensure the outer structure remains { success: false, error: ... }
      JsonB(json!({
        "success": false,
        "error": error
      }))
    }
  }
}

#[pg_extern(strict, parallel_safe)]
fn validate_json_schema(schema_id: &str, instance: JsonB) -> JsonB {
  let cache = SCHEMA_CACHE.read().unwrap();

  // Lookup uses the original schema_id
  match cache.id_to_index.get(schema_id) {
    None => JsonB(json!({
      "success": false,
      "error": {
        "message": format!("Schema with id '{}' not found in cache", schema_id)
      }
    })),
    Some(sch_index) => {
      let instance_value: Value = instance.0;
      match cache.schemas.validate(&instance_value, *sch_index) {
        Ok(_) => JsonB(json!({ "success": true })),
        Err(validation_error) => {
          // Collect all leaf errors first
          let mut raw_error_list = Vec::new();
          collect_leaf_errors(&validation_error, &mut raw_error_list);

          // Filter the errors (e.g., deduplicate by instance_path)
          let filtered_error_list = filter_boon_errors(raw_error_list);

          JsonB(json!({
            "success": false,
            "error": filtered_error_list // Return the filtered list
          }))
        }
      }
    }
  }
}

// Recursively collects leaf errors into a flat list
fn collect_leaf_errors(error: &ValidationError, errors_list: &mut Vec<Value>) {
  if error.causes.is_empty() {
    let default_message = format!("{}", error);
    let message = if let Some(start_index) = default_message.find("': ") {
        default_message[start_index + 3..].to_string()
    } else {
        default_message
    };

    errors_list.push(json!({
      "message": message,
      "schema_path": error.schema_url.to_string(),
      "instance_path": error.instance_location.to_string(),
    }));
  } else {
    for cause in &error.causes {
      collect_leaf_errors(cause, errors_list);
    }
  }
}

// Filters collected errors, e.g., deduplicating by instance_path
fn filter_boon_errors(raw_errors: Vec<Value>) -> Vec<Value> {
    use std::collections::HashMap;
    use std::collections::hash_map::Entry;

    // Use a HashMap to keep only the first error for each instance_path
    let mut unique_errors: HashMap<String, Value> = HashMap::new();

    for error_value in raw_errors {
        if let Some(instance_path_value) = error_value.get("instance_path") {
            if let Some(instance_path_str) = instance_path_value.as_str() {
                // Use Entry API to insert only if the key is not present
                if let Entry::Vacant(entry) = unique_errors.entry(instance_path_str.to_string()) {
                    entry.insert(error_value);
                }
            }
        }
        // If error doesn't have instance_path or it's not a string, we might ignore it or handle differently.
        // For now, we implicitly ignore errors without a valid string instance_path for deduplication.
    }

    // Collect the unique errors from the map values
    unique_errors.into_values().collect()
}

#[pg_extern(strict, parallel_safe)]
fn json_schema_cached(schema_id: &str) -> bool {
  let cache = SCHEMA_CACHE.read().unwrap();
  cache.id_to_index.contains_key(schema_id)
}

#[pg_extern(strict)]
fn clear_json_schemas() {
  let mut cache = SCHEMA_CACHE.write().unwrap();
  *cache = BoonCache {
    schemas: Schemas::new(),
    id_to_index: HashMap::new(),
  };
}

#[pg_extern(strict, parallel_safe)]
fn show_json_schemas() -> Vec<String> {
  let cache = SCHEMA_CACHE.read().unwrap();
  let ids: Vec<String> = cache.id_to_index.keys().cloned().collect();
  ids
}

/// This module is required by `cargo pgrx test` invocations.
/// It must be visible at the root of your extension crate.
#[cfg(test)]
pub mod pg_test {
  pub fn setup(_options: Vec<&str>) {
    // perform one-off initialization when the pg_test framework starts
  }

  #[must_use]
  pub fn postgresql_conf_options() -> Vec<&'static str> {
    // return any postgresql.conf settings that are required for your tests
    vec![]
  }
}


#[cfg(any(test, feature = "pg_test"))]
#[pg_schema]
mod tests {
  include!("tests.rs");
}