From 5b2feb5ea79e1b275b2be6bc310b763d31ca1e86 Mon Sep 17 00:00:00 2001
From: Alex Groleau <agroleau@teampalmtree.com>
Date: Thu, 14 May 2026 17:35:31 -0400
Subject: [PATCH] more ordering fixes

---
 src/lib.rs       | 41 ++++++++++++++++++++++++++++++-----------
 src/tests/mod.rs |  6 +++---
 2 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 78da49c..af8e47b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -4,6 +4,9 @@ use pgrx::*;
 #[cfg(not(test))]
 pg_module_magic!();
 
+#[cfg(test)]
+pub struct JsonB(pub serde_json::Value);
+
 #[cfg(test)]
 pub struct Json(pub serde_json::Value);
 
@@ -25,7 +28,7 @@ lazy_static::lazy_static! {
   static ref GLOBAL_JSPG: RwLock<Option<Arc<jspg::Jspg>>> = RwLock::new(None);
 }
 
-fn jspg_failure() -> Json {
+fn jspg_failure() -> JsonB {
   let error = crate::drop::Error {
     code: "ENGINE_NOT_INITIALIZED".to_string(),
     message: "JSPG extension has not been initialized via jspg_setup".to_string(),
@@ -37,7 +40,7 @@ fn jspg_failure() -> Json {
     },
   };
   let drop = crate::drop::Drop::with_errors(vec![error]);
-  Json(serde_json::to_value(drop).unwrap())
+  JsonB(serde_json::to_value(drop).unwrap())
 }
 
 #[cfg_attr(not(test), pg_extern(strict))]
@@ -55,24 +58,40 @@ pub fn jspg_setup(database: Json) -> Json {
 }
 
 #[cfg_attr(not(test), pg_extern)]
-pub fn jspg_merge(schema_id: &str, data: Json) -> Json {
+pub fn jspg_merge(schema_id: &str, data: JsonB) -> JsonB {
   // Try to acquire a read lock to get a clone of the Engine Arc
   let engine_opt = {
     let lock = GLOBAL_JSPG.read().unwrap();
     lock.clone()
   };
 
+  match engine_opt {
+    Some(engine) => {
+      let drop = engine.merger.merge(schema_id, data.0);
+      JsonB(serde_json::to_value(drop).unwrap())
+    }
+    None => jspg_failure(),
+  }
+}
+
+#[cfg_attr(not(test), pg_extern)]
+pub fn jspg_merge_ordered(schema_id: &str, data: Json) -> Json {
+  let engine_opt = {
+    let lock = GLOBAL_JSPG.read().unwrap();
+    lock.clone()
+  };
+
   match engine_opt {
     Some(engine) => {
       let drop = engine.merger.merge(schema_id, data.0);
       Json(serde_json::to_value(drop).unwrap())
     }
-    None => jspg_failure(),
+    None => Json(jspg_failure().0),
   }
 }
 
 #[cfg_attr(not(test), pg_extern)]
-pub fn jspg_query(schema_id: &str, filter: Option<Json>) -> Json {
+pub fn jspg_query(schema_id: &str, filter: Option<JsonB>) -> JsonB {
   let engine_opt = {
     let lock = GLOBAL_JSPG.read().unwrap();
     lock.clone()
@@ -83,7 +102,7 @@ pub fn jspg_query(schema_id: &str, filter: Option<Json>) -> Json {
       let drop = engine
         .queryer
         .query(schema_id, filter.as_ref().map(|f| &f.0));
-      Json(serde_json::to_value(drop).unwrap())
+      JsonB(serde_json::to_value(drop).unwrap())
     }
     None => jspg_failure(),
   }
@@ -92,7 +111,7 @@ pub fn jspg_query(schema_id: &str, filter: Option<Json>) -> Json {
 // `mask_json_schema` has been removed as the mask architecture is fully replaced by Spi string queries during DB interactions.
 
 #[cfg_attr(not(test), pg_extern(strict, parallel_safe))]
-pub fn jspg_validate(schema_id: &str, instance: Json) -> Json {
+pub fn jspg_validate(schema_id: &str, instance: JsonB) -> JsonB {
   // 1. Acquire Snapshot
   let jspg_arc = {
     let lock = GLOBAL_JSPG.read().unwrap();
@@ -102,7 +121,7 @@ pub fn jspg_validate(schema_id: &str, instance: Json) -> Json {
   // 2. Validate (Lock-Free)
   if let Some(engine) = jspg_arc {
     let drop = engine.validator.validate(schema_id, &instance.0);
-    Json(serde_json::to_value(drop).unwrap())
+    JsonB(serde_json::to_value(drop).unwrap())
   } else {
     jspg_failure()
   }
@@ -122,16 +141,16 @@ pub fn jspg_database() -> Json {
       let drop = crate::drop::Drop::success_with_val(database_json);
       Json(serde_json::to_value(drop).unwrap())
     }
-    None => jspg_failure(),
+    None => Json(jspg_failure().0),
   }
 }
 
 #[cfg_attr(not(test), pg_extern(strict))]
-pub fn jspg_teardown() -> Json {
+pub fn jspg_teardown() -> JsonB {
   let mut lock = GLOBAL_JSPG.write().unwrap();
   *lock = None;
   let drop = crate::drop::Drop::success();
-  Json(serde_json::to_value(drop).unwrap())
+  JsonB(serde_json::to_value(drop).unwrap())
 }
 
 #[cfg(test)]
diff --git a/src/tests/mod.rs b/src/tests/mod.rs
index 514437f..8c2501a 100644
--- a/src/tests/mod.rs
+++ b/src/tests/mod.rs
@@ -11,7 +11,7 @@ fn test_library_api() {
   // 1. Initially, schemas are not cached.
 
   // Expected uninitialized drop format: errors + null response
-  let uninitialized_drop = jspg_validate("source_schema", Json(json!({})));
+  let uninitialized_drop = jspg_validate("source_schema", JsonB(json!({})));
   assert_eq!(
     uninitialized_drop.0,
     json!({
@@ -226,7 +226,7 @@ fn test_library_api() {
   );
 
   // 4. Validate Happy Path
-  let happy_drop = jspg_validate("source_schema", Json(json!({"type": "source_schema", "name": "Neo"})));
+  let happy_drop = jspg_validate("source_schema", JsonB(json!({"type": "source_schema", "name": "Neo"})));
   assert_eq!(
     happy_drop.0,
     json!({
@@ -236,7 +236,7 @@ fn test_library_api() {
   );
 
   // 5. Validate Unhappy Path
-  let unhappy_drop = jspg_validate("source_schema", Json(json!({"type": "source_schema", "wrong": "data"})));
+  let unhappy_drop = jspg_validate("source_schema", JsonB(json!({"type": "source_schema", "wrong": "data"})));
   assert_eq!(
     unhappy_drop.0,
     json!({