Files
jspg/src/lib.rs
2025-04-14 16:11:44 -04:00

238 lines
7.5 KiB
Rust

use pgrx::*;
use jsonschema::{Draft, Validator};
use serde_json::json;
use std::collections::HashMap;
use std::sync::RwLock;
use lazy_static::lazy_static;
pg_module_magic!();
// Global, thread-safe schema cache using the correct Validator type
lazy_static! {
static ref SCHEMA_CACHE: RwLock<HashMap<String, Validator>> = RwLock::new(HashMap::new());
}
// Cache a schema explicitly with a provided ID
#[pg_extern(immutable, strict, parallel_safe)]
fn cache_schema(schema_id: &str, schema: JsonB) -> JsonB {
let schema_value = schema.0;
// Compile the schema using the builder pattern
match jsonschema::options()
.with_draft(Draft::Draft7)
.should_validate_formats(true)
.build(&schema_value)
{
Ok(compiled_schema) => {
// If compilation succeeds, add it to the cache
let mut cache = SCHEMA_CACHE.write().unwrap();
cache.insert(schema_id.to_string(), compiled_schema);
JsonB(json!({ "success": true, "id": schema_id }))
}
Err(e) => {
// If compilation fails, return an error
JsonB(json!({
"success": false,
"error": format!("Failed to compile schema '{}': {}", schema_id, e)
}))
}
}
}
// Check if a schema is cached
#[pg_extern(immutable, strict, parallel_safe)]
fn schema_cached(schema_id: &str) -> bool {
SCHEMA_CACHE.read().unwrap().contains_key(schema_id)
}
// Validate JSONB instance against a cached schema by ID
#[pg_extern(immutable, strict, parallel_safe)]
fn validate_schema(schema_id: &str, instance: JsonB) -> JsonB {
let cache = SCHEMA_CACHE.read().unwrap();
let compiled_schema: &Validator = match cache.get(schema_id) {
Some(schema) => schema,
None => {
// Return the 'schema not cached' error in the standard object format
let error_msg = format!("Schema ID '{}' not cached", schema_id);
return JsonB(json!({
"valid": false,
"errors": [json!({
"kind": "SchemaNotFound", // Custom kind for this case
"error": error_msg
})]
}));
}
};
let instance_value = instance.0;
// Use iter_errors() to get all validation errors
let errors_iterator = compiled_schema.iter_errors(&instance_value);
// Collect errors into a vector first to check if any exist
let collected_errors_result: Vec<_> = errors_iterator.collect();
if collected_errors_result.is_empty() {
// No errors found, validation passed
JsonB(json!({ "valid": true }))
} else {
// Errors found, format them
let error_details = collect_all_errors(collected_errors_result.into_iter());
JsonB(json!({
"valid": false,
"errors": error_details
}))
}
}
fn format_validation_error(error: &jsonschema::ValidationError) -> serde_json::Value {
json!({
"instance_path": error.instance_path.to_string(),
"schema_path": error.schema_path.to_string(),
"kind": format!("{:?}", error.kind),
"error": error.to_string()
})
}
// Simplified: Collects all validation errors by formatting each one.
// Assumes the iterator provided by iter_errors() gives all necessary detail.
fn collect_all_errors<'a>(
errors: impl Iterator<Item = jsonschema::ValidationError<'a>>,
) -> Vec<serde_json::Value> {
errors.map(|e| format_validation_error(&e)).collect()
}
// Show the IDs of all schemas currently in the cache
#[pg_extern(immutable, parallel_safe)]
fn show_schema_cache() -> Vec<String> {
let cache = SCHEMA_CACHE.read().unwrap();
cache.keys().cloned().collect()
}
// Clear the entire schema cache explicitly
#[pg_extern(immutable, parallel_safe)]
fn clear_schema_cache() -> bool {
SCHEMA_CACHE.write().unwrap().clear();
true
}
#[pg_schema]
#[cfg(any(test, feature = "pg_test"))]
mod tests {
use pgrx::prelude::*;
use serde_json::json;
use pgrx::JsonB; // Import JsonB specifically for tests
// Helper to clear cache before tests that need it
fn setup_test() {
crate::clear_schema_cache();
}
#[pg_test]
fn test_cache_and_validate_schema() {
setup_test();
assert!(crate::cache_schema(
"test_schema",
JsonB(json!({ "type": "object" }))
).0["success"] == json!(true));
assert!(crate::schema_cached("test_schema"));
let result_valid = crate::validate_schema("test_schema", JsonB(json!({ "foo": "bar" })));
assert_eq!(result_valid.0["valid"], true);
let result_invalid = crate::validate_schema("test_schema", JsonB(json!(42)));
assert_eq!(result_invalid.0["valid"], false);
assert!(result_invalid.0["errors"][0]["error"].as_str().unwrap().contains("is not of type \"object\""));
}
#[pg_test]
fn test_schema_not_cached() {
setup_test();
let result = crate::validate_schema("unknown_schema", JsonB(json!({})));
assert_eq!(result.0["valid"], false);
assert!(result.0["errors"][0]["error"].as_str().unwrap().contains("not cached"));
}
#[pg_test]
fn test_clear_schema_cache() {
setup_test();
crate::cache_schema("clear_test", JsonB(json!({ "type": "object" })));
assert!(crate::schema_cached("clear_test"));
crate::clear_schema_cache();
assert!(!crate::schema_cached("clear_test"));
}
#[pg_test]
fn test_invalid_schema_cache() {
setup_test();
// Attempt to cache an invalid schema definition
let result = crate::cache_schema(
"bad_schema",
JsonB(json!({ "type": "unknown_type" }))
);
assert!(result.0["success"] == json!(false), "Caching an invalid schema should fail");
assert!(!crate::schema_cached("bad_schema"));
}
#[pg_test]
fn test_show_schema_cache() {
setup_test();
assert!(crate::cache_schema("schema1", JsonB(json!({ "type": "string" }))).0["success"] == json!(true));
assert!(crate::cache_schema("schema2", JsonB(json!({ "type": "number" }))).0["success"] == json!(true));
let mut cached_ids = crate::show_schema_cache();
cached_ids.sort(); // Sort for deterministic comparison
assert_eq!(cached_ids.len(), 2);
assert_eq!(cached_ids, vec!["schema1", "schema2"]);
crate::clear_schema_cache();
let empty_ids = crate::show_schema_cache();
assert!(empty_ids.is_empty());
}
#[pg_test]
fn test_detailed_validation_errors() {
setup_test();
let schema_id = "required_prop_schema";
let schema = JsonB(json!({
"title": "Test Required",
"type": "object",
"properties": {
"name": { "type": "string" },
"age": { "type": "integer" }
},
"required": ["name"]
}));
assert!(crate::cache_schema(schema_id, schema).0["success"] == json!(true));
// Instance missing the required 'name' property
let invalid_instance = JsonB(json!({ "age": 30 }));
let result = crate::validate_schema(schema_id, invalid_instance);
assert_eq!(result.0["valid"], false);
let errors = result.0["errors"].as_array().expect("Errors should be an array");
assert_eq!(errors.len(), 1, "Should have exactly one error");
let error = &errors[0];
eprintln!("Validation Error Details: {}", error);
assert_eq!(error["instance_path"].as_str().unwrap(), "", "Instance path should be root");
assert_eq!(error["schema_path"].as_str().unwrap(), "/required", "Schema path should point to required keyword");
assert!(error["kind"].as_str().unwrap().contains("Required"), "Error kind should be Required");
assert!(error["error"].as_str().unwrap().contains("is a required property"), "Error message mismatch");
}
}
#[cfg(test)]
pub mod pg_test {
pub fn setup(_options: Vec<&str>) {
// Initialization if needed
}
pub fn postgresql_conf_options() -> Vec<&'static str> {
vec![]
}
}