RedisJSON · gkorland · Oct 2, 2019 · Sep 16, 2019 · Sep 16, 2019 · Sep 18, 2019
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -8,16 +8,16 @@ edition = "2018"
 crate-type = ["cdylib"]
 
 [dependencies]
-bson = "0.13"
+bson = "0.14"
 serde_json = "1.0"
 libc = "0.2"
 jsonpath_lib = { git="https://github.com/gkorland/jsonpath.git", branch="reaplce_with_ownership_parser" }
 
 [dependencies.redismodule]
 git = "https://github.com/redislabsmodules/redismodule-rs.git"
-branch = "master"
+branch = "safe_context_rm_call"
 features = ["experimental-api"]
 
 [dependencies.redisearch_api]
 git = "https://github.com/RediSearch/redisearch-api-rs.git"
-branch = "master"
+branch = "safe_context_rm_call"
diff --git a/src/backward.rs b/src/backward.rs
@@ -34,7 +34,7 @@ impl From<u64> for NodeType {
     }
 }
 
-pub unsafe fn json_rdb_load(rdb: *mut raw::RedisModuleIO) -> Value {
+pub fn json_rdb_load(rdb: *mut raw::RedisModuleIO) -> Value {
     let node_type = raw::load_unsigned(rdb).into();
     match node_type {
         NodeType::Null => Value::Null,

diff --git a/src/commands/index.rs b/src/commands/index.rs
@@ -1,7 +1,8 @@
-use serde_json::Value;
+use std::thread;
 
-use redismodule::{Context, RedisError, RedisResult, RedisValue};
-use redismodule::{NextArg, REDIS_OK};
+use serde_json::{Map, Value};
+
+use redismodule::{Context, NextArg, RedisError, RedisResult, RedisValue, REDIS_OK};
 
 use redisearch_api::{Document, FieldType};
 
@@ -75,13 +76,11 @@ pub fn add_document(key: &str, index_name: &str, doc: &RedisJSON) -> RedisResult
 
     let map = schema_map::as_ref();
 
-    map.get(index_name)
-        .ok_or("ERR no such index".into())
-        .and_then(|schema| {
-            let rsdoc = create_document(key, schema, doc)?;
-            schema.index.add_document(&rsdoc)?;
-            REDIS_OK
-        })
+    if let Some(schema) = map.get(index_name) {
+        let rsdoc = create_document(key, schema, doc)?;
+        schema.index.add_document(&rsdoc)?;
+    }
+    REDIS_OK
 }
 
 fn create_document(key: &str, schema: &Schema, doc: &RedisJSON) -> Result<Document, Error> {
@@ -91,13 +90,14 @@ fn create_document(key: &str, schema: &Schema, doc: &RedisJSON) -> Result<Docume
     let rsdoc = Document::create(key, score);
 
     for (field_name, path) in fields {
-        let value = doc.get_doc(&path)?;
-
-        match value {
-            Value::String(v) => rsdoc.add_field(field_name, &v, FieldType::FULLTEXT),
-            Value::Number(v) => rsdoc.add_field(field_name, &v.to_string(), FieldType::NUMERIC),
-            Value::Bool(v) => rsdoc.add_field(field_name, &v.to_string(), FieldType::TAG),
-            _ => {}
+        let results = doc.get_values(path)?;
+        if let Some(value) = results.first() {
+            match value {
+                Value::String(v) => rsdoc.add_field(field_name, &v, FieldType::FULLTEXT),
+                Value::Number(v) => rsdoc.add_field(field_name, &v.to_string(), FieldType::NUMERIC),
+                Value::Bool(v) => rsdoc.add_field(field_name, &v.to_string(), FieldType::TAG),
+                _ => {}
+            }
         }
     }
 
@@ -120,14 +120,78 @@ where
     match subcommand.to_uppercase().as_str() {
         "ADD" => {
             let path = args.next_string()?;
-            add_field(&index_name, &field_name, &path)
+            add_field(&index_name, &field_name, &path)?;
+
+            // TODO handle another "ADD" calls in prallel a running call
+            thread::spawn(move || {
+                let schema = if let Some(stored_schema) = schema_map::as_ref().get(&index_name) {
+                    stored_schema
+                } else {
+                    return; // TODO handle this case
+                };
+
+                let ctx = Context::get_thread_safe_context();
+                let mut cursor: u64 = 0;
+                loop {
+                    ctx.lock();
+                    let res = scan_and_index(&ctx, &schema, cursor);
+                    ctx.unlock();
+
+                    match res {
+                        Ok(c) => cursor = c,
+                        Err(e) => {
+                            eprintln!("Err on index {:?}", e); // TODO hadnle this better
+                            return;
+                        }
+                    }
+                    if cursor == 0 {
+                        break;
+                    }
+                }
+            });
+
+            REDIS_OK
         }
         //"DEL" => {}
         //"INFO" => {}
         _ => Err("ERR unknown subcommand - try `JSON.INDEX HELP`".into()),
     }
 }
 
+fn scan_and_index(ctx: &Context, schema: &Schema, cursor: u64) -> Result<u64, RedisError> {
+    let values = ctx.call("scan", &[&cursor.to_string()]);
+    match values {
+        Ok(RedisValue::Array(arr)) => match (arr.get(0), arr.get(1)) {
+            (Some(RedisValue::SimpleString(next_cursor)), Some(RedisValue::Array(keys))) => {
+                let cursor = next_cursor.parse().unwrap();
+                let res = keys.iter().try_for_each(|k| {
+                    if let RedisValue::SimpleString(key) = k {
+                        ctx.open_key(&key)
+                            .get_value::<RedisJSON>(&REDIS_JSON_TYPE)
+                            .and_then(|doc| {
+                                if let Some(data) = doc {
+                                    if let Some(index) = &data.index {
+                                        if schema.name == *index {
+                                            add_document(key, index, data)?;
+                                        }
+                                    }
+                                    Ok(())
+                                } else {
+                                    Err("Error on get value from key".into())
+                                }
+                            })
+                    } else {
+                        Err("Error on parsing reply from scan".into()) 
+                    }
+                });
+                res.map(|_| cursor)
+            }
+            _ => Err("Error on parsing reply from scan".into()), 
+        },
+        _ => Err("Error on parsing reply from scan".into()), 
+    }
+}
+
 // JSON.QGET <index> <query> <path>
 pub fn qget<I>(ctx: &Context, args: I) -> RedisResult
 where
@@ -145,18 +209,29 @@ where
         .ok_or("ERR no such index".into())
         .map(|schema| &schema.index)
         .and_then(|index| {
-            let results: Result<Vec<_>, RedisError> = index
-                .search(&query)?
-                .map(|key| {
-                    let key = ctx.open_key_writable(&key);
-                    let value = match key.get_value::<RedisJSON>(&REDIS_JSON_TYPE)? {
-                        Some(doc) => doc.to_string(&path, Format::JSON)?.into(),
-                        None => RedisValue::None,
-                    };
-                    Ok(value)
-                })
-                .collect();
-
-            Ok(results?.into())
+            let result: Value =
+                index
+                    .search(&query)?
+                    .try_fold(Value::Object(Map::new()), |mut acc, key| {
+                        ctx.open_key(&key)
+                            .get_value::<RedisJSON>(&REDIS_JSON_TYPE)
+                            .and_then(|doc| {
+                                doc.map_or(Ok(Vec::new()), |data| {
+                                    data.get_values(&path)
+                                        .map_err(|e| e.into()) // Convert Error to RedisError
+                                        .map(|values| {
+                                            values.into_iter().map(|val| val.clone()).collect()
+                                        })
+                                })
+                            })
+                            .map(|r| {
+                                acc.as_object_mut()
+                                    .unwrap()
+                                    .insert(key.to_string(), Value::Array(r));
+                                acc
+                            })
+                    })?;
+
+            Ok(RedisJSON::serialize(&result, Format::JSON)?.into())
         })
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -132,7 +132,7 @@ fn json_set(ctx: &Context, args: Vec<String>) -> RedisResult {
         }
         (None, SetOptions::AlreadyExists) => Ok(RedisValue::None),
         (None, _) => {
-            let doc = RedisJSON::from_str(&value, format)?;
+            let doc = RedisJSON::from_str(&value, &index, format)?;
             if path == "$" {
                 redis_key.set_value(&REDIS_JSON_TYPE, doc)?;
 
@@ -692,7 +692,7 @@ fn json_resp(ctx: &Context, args: Vec<String>) -> RedisResult {
 
     let key = ctx.open_key(&key);
     match key.get_value::<RedisJSON>(&REDIS_JSON_TYPE)? {
-        Some(doc) => Ok(resp_serialize(doc.get_doc(&path)?)),
+        Some(doc) => Ok(resp_serialize(doc.get_first(&path)?)),
         None => Ok(RedisValue::None),
     }
 }

diff --git a/src/redisjson.rs b/src/redisjson.rs
@@ -41,6 +41,7 @@ impl Format {
 #[derive(Debug)]
 pub struct RedisJSON {
     data: Value,
+    pub index: Option<String>,
 }
 
 impl RedisJSON {
@@ -62,10 +63,14 @@ impl RedisJSON {
         }
     }
 
-    pub fn from_str(data: &str, format: Format) -> Result<Self, Error> {
+    pub fn from_str(data: &str, index: &Option<String>, format: Format) -> Result<Self, Error> {
         let value = RedisJSON::parse_str(data, format)?;
-        Ok(Self { data: value })
+        Ok(Self {
+            data: value,
+            index: index.clone(),
+        })
     }
+
     fn add_value(&mut self, path: &str, value: Value) -> Result<bool, Error> {
         if NodeVisitorImpl::check(path)? {
             let mut splits = path.rsplitn(2, '.');
@@ -154,7 +159,7 @@ impl RedisJSON {
     }
 
     pub fn to_string(&self, path: &str, format: Format) -> Result<String, Error> {
-        let results = self.get_doc(path)?;
+        let results = self.get_first(path)?;
         Self::serialize(results, format)
     }
 
@@ -193,35 +198,35 @@ impl RedisJSON {
     }
 
     pub fn str_len(&self, path: &str) -> Result<usize, Error> {
-        self.get_doc(path)?
+        self.get_first(path)?
             .as_str()
             .ok_or_else(|| "ERR wrong type of path value".into())
             .map(|s| s.len())
     }
 
     pub fn arr_len(&self, path: &str) -> Result<usize, Error> {
-        self.get_doc(path)?
+        self.get_first(path)?
             .as_array()
             .ok_or_else(|| "ERR wrong type of path value".into())
             .map(|arr| arr.len())
     }
 
     pub fn obj_len(&self, path: &str) -> Result<usize, Error> {
-        self.get_doc(path)?
+        self.get_first(path)?
             .as_object()
             .ok_or_else(|| "ERR wrong type of path value".into())
             .map(|obj| obj.len())
     }
 
     pub fn obj_keys<'a>(&'a self, path: &'a str) -> Result<Vec<&'a String>, Error> {
-        self.get_doc(path)?
+        self.get_first(path)?
             .as_object()
             .ok_or_else(|| "ERR wrong type of path value".into())
             .map(|obj| obj.keys().collect())
     }
 
     pub fn arr_index(&self, path: &str, scalar: &str, start: i64, end: i64) -> Result<i64, Error> {
-        if let Value::Array(arr) = self.get_doc(path)? {
+        if let Value::Array(arr) = self.get_first(path)? {
             // end=-1/0 means INFINITY to support backward with RedisJSON
             if arr.is_empty() || end < -1 {
                 return Ok(-1);
@@ -252,7 +257,7 @@ impl RedisJSON {
     }
 
     pub fn get_type(&self, path: &str) -> Result<String, Error> {
-        let s = RedisJSON::value_name(self.get_doc(path)?);
+        let s = RedisJSON::value_name(self.get_first(path)?);
         Ok(s.to_string())
     }
 
@@ -322,7 +327,7 @@ impl RedisJSON {
 
     pub fn get_memory<'a>(&'a self, path: &'a str) -> Result<usize, Error> {
         // TODO add better calculation, handle wrappers, internals and length
-        let res = match self.get_doc(path)? {
+        let res = match self.get_first(path)? {
             Value::Null => 0,
             Value::Bool(v) => mem::size_of_val(v),
             Value::Number(v) => mem::size_of_val(v),
@@ -333,26 +338,39 @@ impl RedisJSON {
         Ok(res.into())
     }
 
-    // TODO: Rename this to 'get_value', since 'doc' is overloaded.
-    pub fn get_doc<'a>(&'a self, path: &'a str) -> Result<&'a Value, Error> {
-        let results = jsonpath_lib::select(&self.data, path)?;
+    pub fn get_first<'a>(&'a self, path: &'a str) -> Result<&'a Value, Error> {
+        let results = self.get_values(path)?;
         match results.first() {
             Some(s) => Ok(s),
             None => Err("ERR path does not exist".into()),
         }
     }
+
+    pub fn get_values<'a>(&'a self, path: &'a str) -> Result<Vec<&'a Value>, Error> {
+        let results = jsonpath_lib::select(&self.data, path)?;
+        Ok(results)
+    }
 }
 
 pub mod type_methods {
     use super::*;
 
     #[allow(non_snake_case, unused)]
-    pub unsafe extern "C" fn rdb_load(rdb: *mut raw::RedisModuleIO, encver: c_int) -> *mut c_void {
+    pub extern "C" fn rdb_load(rdb: *mut raw::RedisModuleIO, encver: c_int) -> *mut c_void {
         let json = match encver {
             0 => RedisJSON {
                 data: backward::json_rdb_load(rdb),
+                index: None, // TODO handle load from rdb
             },
-            2 => RedisJSON::from_str(&raw::load_string(rdb), Format::JSON).unwrap(),
+            2 => {
+                let data = raw::load_string(rdb);
+                let schema = if raw::load_unsigned(rdb) > 0 {
+                    Some(raw::load_string(rdb))
+                } else {
+                    None
+                };
+                RedisJSON::from_str(&data, &schema, Format::JSON).unwrap()
+            }
             _ => panic!("Can't load old RedisJSON RDB"),
         };
         Box::into_raw(Box::new(json)) as *mut c_void
@@ -367,5 +385,11 @@ pub mod type_methods {
     pub unsafe extern "C" fn rdb_save(rdb: *mut raw::RedisModuleIO, value: *mut c_void) {
         let json = &*(value as *mut RedisJSON);
         raw::save_string(rdb, &json.data.to_string());
+        if let Some(index) = &json.index {
+            raw::save_unsigned(rdb, 1);
+            raw::save_string(rdb, &index);
+        } else {
+            raw::save_unsigned(rdb, 0);
+        }
     }
 }