final fix to org id setting in merge

src/merger/mod.rs

@@ -40,7 +40,7 @@ impl Merger {
       }
     };
 
-    let result = self.merge_internal(target_schema, data, &mut notifications_queue);
+    let result = self.merge_internal(target_schema, data, &mut notifications_queue, None, false);
 
     let val_resolved = match result {
       Ok(val) => val,
@@ -129,37 +129,16 @@ impl Merger {
     crate::drop::Drop::success_with_val(stripped_val)
   }
 
-  fn inject_organization_id(
-    relative: &mut serde_json::Map<String, Value>,
-    entity_fields: &serde_json::Map<String, Value>,
-    schema: &Arc<crate::database::schema::Schema>,
-  ) {
-    if !relative.contains_key("organization_id") {
-      let mut const_val: Option<Value> = None;
-      if let Some(compiled_props) = schema.obj.compiled_properties.get() {
-        if let Some(org_schema) = compiled_props.get("organization_id") {
-          if let Some(c) = &org_schema.obj.const_ {
-            const_val = Some(c.clone());
-          }
-        }
-      }
-
-      if let Some(c) = const_val {
-        relative.insert("organization_id".to_string(), c);
-      } else if let Some(org_id) = entity_fields.get("organization_id") {
-        relative.insert("organization_id".to_string(), org_id.clone());
-      }
-    }
-  }
-
   pub(crate) fn merge_internal(
     &self,
     mut schema: Arc<crate::database::schema::Schema>,
     data: Value,
     notifications: &mut Vec<String>,
+    parent_org_id: Option<String>,
+    is_child: bool,
   ) -> Result<Value, String> {
     match data {
-      Value::Array(items) => self.merge_array(schema, items, notifications),
+      Value::Array(items) => self.merge_array(schema, items, notifications, parent_org_id, is_child),
       Value::Object(map) => {
         if let Some(options) = schema.obj.compiled_options.get() {
           if let Some(disc) = schema.obj.compiled_discriminator.get() {
@@ -167,9 +146,7 @@ impl Merger {
             if let Some(v) = val {
               if let Some((idx_opt, target_id_opt)) = options.get(v) {
                 if let Some(target_id) = target_id_opt {
-                  if let Some(target_schema) =
+                  if let Some(target_schema) = self.db.schemas.get(target_id) {
-                    self.db.schemas.get(target_id)
-                  {
                     schema = target_schema.clone();
                   } else {
                     return Err(format!(
@@ -208,7 +185,7 @@ impl Merger {
             }
           }
         }
-        self.merge_object(schema, map, notifications)
+        self.merge_object(schema, map, notifications, parent_org_id, is_child)
       }
       _ => Err("Invalid merge payload: root must be an Object or Array".to_string()),
     }
@@ -219,6 +196,8 @@ impl Merger {
     schema: Arc<crate::database::schema::Schema>,
     items: Vec<Value>,
     notifications: &mut Vec<String>,
+    parent_org_id: Option<String>,
+    is_child: bool,
   ) -> Result<Value, String> {
     let mut item_schema = schema.clone();
     if let Some(crate::database::object::SchemaTypeOrArray::Single(t)) = &schema.obj.type_ {
@@ -231,7 +210,7 @@ impl Merger {
 
     let mut resolved_items = Vec::new();
     for item in items {
-      let resolved = self.merge_internal(item_schema.clone(), item, notifications)?;
+      let resolved = self.merge_internal(item_schema.clone(), item, notifications, parent_org_id.clone(), is_child)?;
       resolved_items.push(resolved);
     }
     Ok(Value::Array(resolved_items))
@@ -242,6 +221,8 @@ impl Merger {
     schema: Arc<crate::database::schema::Schema>,
     obj: serde_json::Map<String, Value>,
     notifications: &mut Vec<String>,
+    parent_org_id: Option<String>,
+    is_child: bool,
   ) -> Result<Value, String> {
     let queue_start = notifications.len();
 
@@ -301,6 +282,20 @@ impl Merger {
       }
     }
 
+    let mut current_org_id = None;
+    if let Some(compiled_props) = schema.obj.compiled_properties.get() {
+      if let Some(org_schema) = compiled_props.get("organization_id") {
+        if let Some(c) = &org_schema.obj.const_ {
+          if let Some(c_str) = c.as_str() {
+            current_org_id = Some(c_str.to_string());
+          }
+        }
+      }
+    }
+    if current_org_id.is_none() {
+      current_org_id = parent_org_id.clone();
+    }
+
     let user_id = self.db.auth_user_id()?;
     let timestamp = self.db.timestamp()?;
 
@@ -315,6 +310,16 @@ impl Merger {
       entity_change_kind = kind;
       entity_fetched = fetched;
       entity_replaces = replaces;
+
+      if entity_change_kind.as_deref() == Some("create") {
+        if is_child {
+          if !entity_fields.contains_key("organization_id") {
+            if let Some(ref org_id) = current_org_id {
+              entity_fields.insert("organization_id".to_string(), Value::String(org_id.clone()));
+            }
+          }
+        }
+      }
     }
 
     let mut entity_response = serde_json::Map::new();
@@ -335,13 +340,14 @@ impl Merger {
           if let Some(relation) = self.db.relations.get(&edge.constraint) {
             let parent_is_source = edge.forward;
 
+            let org_id_to_pass = entity_fields.get("organization_id").and_then(|v| v.as_str()).map(|s| s.to_string());
             if parent_is_source {
-              Self::inject_organization_id(&mut relative, &entity_fields, &rel_schema);
-
               let mut merged_relative = match self.merge_internal(
                 rel_schema.clone(),
                 Value::Object(relative),
                 notifications,
+                org_id_to_pass.clone(),
+                true,
               )? {
                 Value::Object(m) => m,
                 _ => continue,
@@ -357,8 +363,6 @@ impl Merger {
               );
               entity_response.insert(relation_name, Value::Object(merged_relative));
             } else {
-              Self::inject_organization_id(&mut relative, &entity_fields, &rel_schema);
-
               Self::apply_entity_relation(
                 &mut relative,
                 &relation.source_columns,
@@ -370,6 +374,8 @@ impl Merger {
                 rel_schema.clone(),
                 Value::Object(relative),
                 notifications,
+                org_id_to_pass.clone(),
+                true,
               )? {
                 Value::Object(m) => m,
                 _ => continue,
@@ -389,6 +395,16 @@ impl Merger {
       entity_change_kind = kind;
       entity_fetched = fetched;
       entity_replaces = replaces;
+
+      if entity_change_kind.as_deref() == Some("create") {
+        if is_child {
+          if !entity_fields.contains_key("organization_id") {
+            if let Some(ref org_id) = current_org_id {
+              entity_fields.insert("organization_id".to_string(), Value::String(org_id.clone()));
+            }
+          }
+        }
+      }
     }
 
     self.merge_entity_fields(
@@ -427,11 +443,10 @@ impl Merger {
               }
             }
 
+            let org_id_to_pass = entity_fields.get("organization_id").and_then(|v| v.as_str()).map(|s| s.to_string());
             let mut relative_responses = Vec::new();
             for relative_item_val in relative_arr {
               if let Value::Object(mut relative_item) = relative_item_val {
-                Self::inject_organization_id(&mut relative_item, &entity_fields, &item_schema);
-
                 Self::apply_entity_relation(
                   &mut relative_item,
                   &relation.source_columns,
@@ -443,6 +458,8 @@ impl Merger {
                   item_schema.clone(),
                   Value::Object(relative_item),
                   notifications,
+                  org_id_to_pass.clone(),
+                  true,
                 )? {
                   Value::Object(m) => m,
                   _ => continue,