Initial RediSearch support: JSON.INDEX ADD, JSON.QGET, JSON.SET ... INDEX

Author: gavrie

src/commands/index.rs

@@ -0,0 +1,161 @@
+use serde_json::Value;
+
+use redismodule::{Context, RedisError, RedisResult, RedisValue};
+use redismodule::{NextArg, REDIS_OK};
+
+use redisearch_api::{Document, FieldType};
+
+use crate::error::Error;
+use crate::redisjson::{Format, RedisJSON};
+use crate::schema::Schema;
+use crate::REDIS_JSON_TYPE;
+
+pub mod schema_map {
+    use crate::schema::Schema;
+    use std::collections::HashMap;
+
+    type SchemaMap = HashMap<String, Schema>;
+
+    /// We keep a static map, since it needs to be accessed from multiple Redis module commands
+    /// and there is no other obvious way to implement this (such as a "user data" pointer maintained
+    /// for us by Redis).
+    ///
+    /// The init function should be called only once.
+    /// Unwrapping the Option is thus safe afterwards.
+    ///
+    /// Since we have only one thread at the moment, getting a (mutable) reference to the map
+    /// is also safe. If we add threads, we can simply wrap the map: `Option<Mutex<SchemaMap>`.
+    ///
+    static mut SCHEMA_MAP: Option<SchemaMap> = None;
+
+    pub fn init() {
+        let map = HashMap::new();
+        unsafe {
+            SCHEMA_MAP = Some(map);
+        }
+    }
+
+    pub fn as_ref() -> &'static SchemaMap {
+        unsafe { SCHEMA_MAP.as_ref() }.unwrap()
+    }
+
+    pub fn as_mut() -> &'static mut SchemaMap {
+        unsafe { SCHEMA_MAP.as_mut() }.unwrap()
+    }
+}
+
+///////////////////////////
+
+fn add_field(index_name: &str, field_name: &str, path: String) -> RedisResult {
+    let map = schema_map::as_mut();
+
+    if !map.contains_key(index_name) {
+        let schema = Schema::new(&index_name);
+        map.insert(index_name.to_owned(), schema);
+    }
+
+    let schema = map.get_mut(index_name).unwrap();
+
+    if schema.fields.contains_key(field_name) {
+        Err("Field already exists".into())
+    } else {
+        schema.index.create_field(field_name);
+        schema.fields.insert(field_name.to_string(), path);
+        REDIS_OK
+    }
+}
+
+pub fn add_document(key: &str, index_name: &str, doc: &RedisJSON) -> RedisResult {
+    // TODO: Index the document with RediSearch:
+    // 1. Determine the index to use (how?)
+    // 2. Get the fields from the index along with their associated paths
+    // 3. Build a RS Document and populate the fields from our doc
+    // 4. Add the Document to the index
+
+    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
+        })
+}
+
+fn create_document(key: &str, schema: &Schema, doc: &RedisJSON) -> Result<Document, Error> {
+    let fields = &schema.fields;
+
+    let score = 1.0;
+    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),
+            _ => {}
+        }
+    }
+
+    Ok(rsdoc)
+}
+
+// JSON.INDEX ADD <index> <field> <path>
+// JSON.INDEX DEL <index> <field>
+// JSON.INDEX INFO <index> <field>
+pub fn index<I>(_ctx: &Context, args: I) -> RedisResult
+where
+    I: IntoIterator<Item = String>,
+{
+    let mut args = args.into_iter().skip(1);
+
+    let subcommand = args.next_string()?;
+    let index_name = args.next_string()?;
+    let field_name = args.next_string()?;
+
+    match subcommand.to_uppercase().as_str() {
+        "ADD" => {
+            let path = args.next_string()?;
+            add_field(&index_name, &field_name, path)
+        }
+        //"DEL" => {}
+        //"INFO" => {}
+        _ => Err("ERR unknown subcommand - try `JSON.INDEX HELP`".into()),
+    }
+}
+
+// JSON.QGET <index> <query> <path>
+pub fn qget<I>(ctx: &Context, args: I) -> RedisResult
+where
+    I: IntoIterator<Item = String>,
+{
+    let mut args = args.into_iter().skip(1);
+
+    let index_name = args.next_string()?;
+    let query = args.next_string()?;
+    let path = args.next().unwrap_or("$".to_string());
+
+    let map = schema_map::as_ref();
+
+    map.get(&index_name)
+        .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())
+        })
+}

src/commands/mod.rs

@@ -0,0 +1 @@
+pub mod index;

src/lib.rs

@@ -2,20 +2,23 @@
 extern crate redismodule;
 
 use redismodule::native_types::RedisType;
-use redismodule::raw as rawmod;
 use redismodule::raw::RedisModuleTypeMethods;
-use redismodule::{Context, NextArg, RedisError, RedisResult, RedisValue, REDIS_OK};
+use redismodule::{raw as rawmod, NextArg};
+use redismodule::{Context, RedisError, RedisResult, RedisValue, REDIS_OK};
 use serde_json::{Number, Value};
+
 use std::{i64, usize};
 
 mod array_index;
 mod backward;
+mod commands;
 mod error;
 mod nodevisitor;
 mod redisjson;
-mod schema;
+mod schema; // TODO: Remove
 
 use crate::array_index::ArrayIndex;
+use crate::commands::index;
 use crate::error::Error;
 use crate::redisjson::{Format, RedisJSON, SetOptions};
 
@@ -36,23 +39,6 @@ static REDIS_JSON_TYPE: RedisType = RedisType::new(
     },
 );
 
-static REDIS_JSON_SCHEMA_TYPE: RedisType = RedisType::new(
-    "ReJSON-SC",
-    1,
-    RedisModuleTypeMethods {
-        version: redismodule::TYPE_METHOD_VERSION,
-
-        rdb_load: Some(schema::type_methods::rdb_load),
-        rdb_save: Some(schema::type_methods::rdb_save),
-        aof_rewrite: None, // TODO add support
-        free: Some(schema::type_methods::free),
-
-        // Currently unused by Redis
-        mem_usage: None,
-        digest: None,
-    },
-);
-
 ///
 /// Backwards compatibility convertor for RedisJSON 1.x clients
 ///
@@ -94,7 +80,7 @@ fn json_del(ctx: &Context, args: Vec<String>) -> RedisResult {
 }
 
 ///
-/// JSON.SET <key> <path> <json> [NX | XX]
+/// JSON.SET <key> <path> <json> [NX | XX | FORMAT <format> | INDEX <index>]
 ///
 fn json_set(ctx: &Context, args: Vec<String>) -> RedisResult {
     let mut args = args.into_iter().skip(1);
@@ -105,6 +91,8 @@ fn json_set(ctx: &Context, args: Vec<String>) -> RedisResult {
 
     let mut format = Format::JSON;
     let mut set_option = SetOptions::None;
+    let mut index = None;
+
     loop {
         if let Some(s) = args.next() {
             match s.to_uppercase().as_str() {
@@ -113,19 +101,25 @@ fn json_set(ctx: &Context, args: Vec<String>) -> RedisResult {
                 "FORMAT" => {
                     format = Format::from_str(args.next_string()?.as_str())?;
                 }
+                "INDEX" => {
+                    index = Some(args.next_string()?);
+                }
                 _ => break,
             };
         } else {
             break;
         }
     }
 
-    let key = ctx.open_key_writable(&key);
-    let current = key.get_value::<RedisJSON>(&REDIS_JSON_TYPE)?;
+    let redis_key = ctx.open_key_writable(&key);
+    let current = redis_key.get_value::<RedisJSON>(&REDIS_JSON_TYPE)?;
 
     match (current, set_option) {
         (Some(ref mut doc), ref op) => {
             if doc.set_value(&value, &path, op, format)? {
+                if let Some(index) = index {
+                    index::add_document(&key, &index, &doc)?;
+                }
                 REDIS_OK
             } else {
                 Ok(RedisValue::None)
@@ -135,7 +129,16 @@ fn json_set(ctx: &Context, args: Vec<String>) -> RedisResult {
         (None, _) => {
             let doc = RedisJSON::from_str(&value, format)?;
             if path == "$" {
-                key.set_value(&REDIS_JSON_TYPE, doc)?;
+                redis_key.set_value(&REDIS_JSON_TYPE, doc)?;
+
+                if let Some(index) = index {
+                    // FIXME: We need to get the value even though we just set it,
+                    // since the original doc is consumed by set_value.
+                    // Can we do better than this?
+                    let doc = redis_key.get_value(&REDIS_JSON_TYPE)?.unwrap();
+                    index::add_document(&key, &index, doc)?;
+                }
+
                 REDIS_OK
             } else {
                 Err("ERR new objects must be created at the root".into())
@@ -721,10 +724,6 @@ fn json_len<F: Fn(&RedisJSON, &String) -> Result<usize, Error>>(
     Ok(length)
 }
 
-fn json_createindex(ctx: &Context, args: Vec<String>) -> RedisResult {
-    Err("Command was not implemented".into())
-}
-
 fn json_cache_info(_ctx: &Context, _args: Vec<String>) -> RedisResult {
     Err("Command was not implemented".into())
 }
@@ -735,6 +734,7 @@ fn json_cache_init(_ctx: &Context, _args: Vec<String>) -> RedisResult {
 //////////////////////////////////////////////////////
 
 pub extern "C" fn init(raw_ctx: *mut rawmod::RedisModuleCtx) -> c_int {
+    crate::commands::index::schema_map::init();
     redisearch_api::init(raw_ctx)
 }
 
@@ -743,7 +743,6 @@ redis_module! {
     version: 1,
     data_types: [
         REDIS_JSON_TYPE,
-        REDIS_JSON_SCHEMA_TYPE
     ],
     init: init,
     commands: [
@@ -768,7 +767,8 @@ redis_module! {
         ["json.debug", json_debug, ""],
         ["json.forget", json_del, "write"],
         ["json.resp", json_resp, ""],
-        ["json.createindex", json_createindex, "write"],
+        ["json.index", commands::index::index, "write"],
+        ["json.qget", commands::index::qget, ""],
         ["json._cacheinfo", json_cache_info, ""],
         ["json._cacheinit", json_cache_init, "write"],
     ],

src/redisjson.rs

@@ -166,6 +166,8 @@ impl RedisJSON {
         Ok(res)
     }
 
+    // FIXME: Implement this by manipulating serde_json::Value values,
+    // and then using serde to serialize to JSON instead of doing it ourselves with strings.
     pub fn to_json(&self, paths: &mut Vec<String>) -> Result<String, Error> {
         let mut selector = jsonpath_lib::selector(&self.data);
         let mut result = paths.drain(..).fold(String::from("{"), |mut acc, path| {
@@ -331,6 +333,7 @@ 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)?;
         match results.first() {