jspg/src/lib.rs

use pgrx::*;

pg_module_magic!();

pub mod drop;
pub mod validator;

use serde_json::json;
use std::sync::{Arc, RwLock};

lazy_static::lazy_static! {
  // Global Atomic Swap Container:
  // - RwLock: To protect the SWAP of the Option.
  // - Option: Because it starts empty.
  // - Arc: Because multiple running threads might hold the OLD validator while we swap.
  // - Validator: It immutably owns the Registry.
  static ref GLOBAL_VALIDATOR: RwLock<Option<Arc<validator::Validator>>> = RwLock::new(None);
}

#[pg_extern(strict)]
pub fn cache_json_schemas(enums: JsonB, types: JsonB, puncs: JsonB) -> JsonB {
  // 1 & 2. Build Registry, Families, and Wrap in Validator all in one shot
  let new_validator = crate::validator::Validator::from_punc_definition(
    Some(&enums.0),
    Some(&types.0),
    Some(&puncs.0),
  );
  let new_arc = Arc::new(new_validator);

  // 3. ATOMIC SWAP
  {
    let mut lock = GLOBAL_VALIDATOR.write().unwrap();
    *lock = Some(new_arc);
  }

  let drop = crate::drop::Drop::success();
  JsonB(serde_json::to_value(drop).unwrap())
}

#[pg_extern(strict, parallel_safe)]
pub fn mask_json_schema(schema_id: &str, instance: JsonB) -> JsonB {
  // 1. Acquire Snapshot
  let validator_arc = {
    let lock = GLOBAL_VALIDATOR.read().unwrap();
    lock.clone()
  };

  // 2. Validate (Lock-Free)
  if let Some(validator) = validator_arc {
    // We need a mutable copy of the value to mask it
    let mut mutable_instance = instance.0.clone();

    match validator.mask(schema_id, &mut mutable_instance) {
      Ok(result) => {
        // If valid, return the MASKED instance
        if result.is_valid() {
          let drop = crate::drop::Drop::success_with_val(mutable_instance);
          JsonB(serde_json::to_value(drop).unwrap())
        } else {
          // If invalid, return errors (Schema Validation Errors)
          let errors: Vec<crate::drop::Error> = result
            .errors
            .into_iter()
            .map(|e| crate::drop::Error {
              code: e.code,
              message: e.message,
              details: crate::drop::ErrorDetails { path: e.path },
            })
            .collect();
          let drop = crate::drop::Drop::with_errors(errors);
          JsonB(serde_json::to_value(drop).unwrap())
        }
      }
      Err(e) => {
        // Schema Not Found or other fatal error
        let error = crate::drop::Error {
          code: e.code,
          message: e.message,
          details: crate::drop::ErrorDetails { path: e.path },
        };
        let drop = crate::drop::Drop::with_errors(vec![error]);
        JsonB(serde_json::to_value(drop).unwrap())
      }
    }
  } else {
    let error = crate::drop::Error {
      code: "VALIDATOR_NOT_INITIALIZED".to_string(),
      message: "JSON Schemas have not been cached yet. Run cache_json_schemas()".to_string(),
      details: crate::drop::ErrorDetails {
        path: "".to_string(),
      },
    };
    let drop = crate::drop::Drop::with_errors(vec![error]);
    JsonB(serde_json::to_value(drop).unwrap())
  }
}

#[pg_extern(strict, parallel_safe)]
pub fn validate_json_schema(schema_id: &str, instance: JsonB) -> JsonB {
  // 1. Acquire Snapshot
  let validator_arc = {
    let lock = GLOBAL_VALIDATOR.read().unwrap();
    lock.clone()
  };

  // 2. Validate (Lock-Free)
  if let Some(validator) = validator_arc {
    match validator.validate(schema_id, &instance.0) {
      Ok(result) => {
        if result.is_valid() {
          let drop = crate::drop::Drop::success();
          JsonB(serde_json::to_value(drop).unwrap())
        } else {
          let errors: Vec<crate::drop::Error> = result
            .errors
            .into_iter()
            .map(|e| crate::drop::Error {
              code: e.code,
              message: e.message,
              details: crate::drop::ErrorDetails { path: e.path },
            })
            .collect();
          let drop = crate::drop::Drop::with_errors(errors);
          JsonB(serde_json::to_value(drop).unwrap())
        }
      }
      Err(e) => {
        let error = crate::drop::Error {
          code: e.code,
          message: e.message,
          details: crate::drop::ErrorDetails { path: e.path },
        };
        let drop = crate::drop::Drop::with_errors(vec![error]);
        JsonB(serde_json::to_value(drop).unwrap())
      }
    }
  } else {
    let error = crate::drop::Error {
      code: "VALIDATOR_NOT_INITIALIZED".to_string(),
      message: "JSON Schemas have not been cached yet. Run cache_json_schemas()".to_string(),
      details: crate::drop::ErrorDetails {
        path: "".to_string(),
      },
    };
    let drop = crate::drop::Drop::with_errors(vec![error]);
    JsonB(serde_json::to_value(drop).unwrap())
  }
}

#[pg_extern(strict, parallel_safe)]
pub fn json_schema_cached(schema_id: &str) -> bool {
  if let Some(validator) = GLOBAL_VALIDATOR.read().unwrap().as_ref() {
    match validator.validate(schema_id, &serde_json::Value::Null) {
      Err(e) if e.code == "SCHEMA_NOT_FOUND" => false,
      _ => true,
    }
  } else {
    false
  }
}

#[pg_extern(strict)]
pub fn clear_json_schemas() -> JsonB {
  let mut lock = GLOBAL_VALIDATOR.write().unwrap();
  *lock = None;
  let drop = crate::drop::Drop::success();
  JsonB(serde_json::to_value(drop).unwrap())
}

#[pg_extern(strict, parallel_safe)]
pub fn show_json_schemas() -> JsonB {
  if let Some(validator) = GLOBAL_VALIDATOR.read().unwrap().as_ref() {
    let mut keys = validator.get_schema_ids();
    keys.sort();
    let drop = crate::drop::Drop::success_with_val(json!(keys));
    JsonB(serde_json::to_value(drop).unwrap())
  } else {
    let drop = crate::drop::Drop::success_with_val(json!([]));
    JsonB(serde_json::to_value(drop).unwrap())
  }
}

#[cfg(any(test, feature = "pg_test"))]
#[pg_schema]
mod tests {
  use pgrx::prelude::*;
  include!("tests/fixtures.rs");
}

#[cfg(test)]
pub mod pg_test {
  pub fn setup(_options: Vec<&str>) {
    // perform any initialization common to all tests
  }

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