demo-tabledata-redis/src/main/java/com/fr/plugin/db/redis/help/client/RedisNashornClient.java

package com.fr.plugin.db.redis.help.client;

import com.fanruan.api.log.LogKit;
import com.fr.base.TableData;
import com.fr.plugin.db.redis.core.emb.Redis;
import com.fr.plugin.db.redis.help.Console;
import com.fr.plugin.db.redis.help.ScriptBridge;
import redis.clients.jedis.Jedis;

import javax.script.ScriptEngine;
import javax.script.ScriptEngineFactory;
import javax.script.ScriptEngineManager;
import javax.script.ScriptException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;

/**
 * @author richie
 * @version 10.0
 * Created by richie on 2019-08-26
 */
public class RedisNashornClient extends BaseRedisClient<Object[], Map<Object, Object>> {

    private static final String PREPARE_SCRIPT =
            "function unique(array) {var res = [];for (var i = 0, arrayLen = array.length; i < arrayLen; i++) {for (var j = 0, resLen = res.length; j < resLen; j++ ) {if (array[i] === res[j]) {break;}} if (j === resLen) {res.push(array[i])}} return res;}\n" +
                    "function merge(table, column) {return {content:table, column:column}};";

    private ScriptEngine scriptEngine;

    public RedisNashornClient(Jedis jedis) {
        super(jedis);
        this.scriptEngine = findScriptEngine();
    }

    @Override
    @ScriptBridge
    public Object[] hkeys(String key) {
        long start = System.currentTimeMillis();
        Set<String> set = jedis.keys(key);
        LogKit.info("Fetch data from redis spend time {} ms.", System.currentTimeMillis() - start);
        return set.toArray(new String[0]);
    }

    @Override
    @ScriptBridge
    public Object[] hmget(String key, String... fields) {
        long start = System.currentTimeMillis();
        List<String> list = jedis.hmget(key, fields);
        LogKit.info("Fetch data from redis spend time {} ms.", System.currentTimeMillis() - start);
        return list.toArray(new String[0]);
    }

    @Override
    @ScriptBridge
    public Map<Object, Object> hgetAll(String key) {
        long start = System.currentTimeMillis();
        Map<String, String> map = jedis.hgetAll(key);
        LogKit.info("Fetch data from redis spend time {} ms.", System.currentTimeMillis() - start);
        return new TreeMap<Object, Object>(map);
    }

    @Override
    @ScriptBridge
    public Object[] keys(String pattern) {
        long start = System.currentTimeMillis();
        Set<String> set = jedis.keys(pattern);
        LogKit.info("Fetch data from redis spend time {} ms.", System.currentTimeMillis() - start);
        return set.toArray(new String[0]);
    }

    @Override
    @ScriptBridge
    public Object[] lrange(String key, long start, long end) {
        long start0 = System.currentTimeMillis();
        List<String> list = jedis.lrange(key, start, end);
        LogKit.info("Fetch data from redis spend time {} ms.", System.currentTimeMillis() - start0);
        return list.toArray(new String[0]);
    }

    @Override
    @ScriptBridge
    public Object[] mget(Object[] keys) {
        long start = System.currentTimeMillis();
        int len = keys.length;
        String[] parameters = new String[len];
        for (int i = 0; i < len; i++) {
            parameters[i] = String.valueOf(keys[i]);
        }
        List<String> list = jedis.mget(parameters);
        LogKit.info("Fetch data from redis spend time {} ms.", System.currentTimeMillis() - start);
        return list.toArray(new String[0]);
    }

    @Override
    @ScriptBridge
    public Object[] smembers(String key) {
        long start = System.currentTimeMillis();
        Set<String> set = jedis.smembers(key);
        LogKit.info("Fetch data from redis spend time {} ms.", System.currentTimeMillis() - start);
        return set.toArray(new String[0]);
    }

    @Override
    public SimpleDataModel build(Redis redis, String query, int limit) {
        scriptEngine.put("console", new Console());
        scriptEngine.put("Files", new NashornFiles(scriptEngine));
        scriptEngine.put("redis", new RedisNashornClient(redis.getClient()));
        try {
            scriptEngine.eval(PREPARE_SCRIPT);
            Map<String, Map<String, Object>> r = (Map<String, Map<String, Object>>) scriptEngine.eval("(function(){" + query + "})();");

            Map<String, Object> columns = r.get("column");
            int columnCount = columns.size();
            String[] columnNames = new String[columnCount];
            int k = 0;
            for (Map.Entry<String, Object> entry : columns.entrySet()) {
                columnNames[k] = String.valueOf(entry.getValue());
                k++;
            }

            Map<String, Object> content = r.get("content");
            int returnRowCount = content.size();
            int realCount = limit == TableData.RESULT_ALL ? returnRowCount : Math.min(returnRowCount, limit);

            List<List<Object>> data = new ArrayList<List<Object>>();
            int i = 0;
            for (Map.Entry<String, Object> entry : content.entrySet()) {
                if (i >= realCount) {
                    break;
                }
                List<Object> row = new ArrayList<Object>();

                Object el = entry.getValue();
                if (el instanceof Map) {
                    Map<String, Object> rowCollection = (Map<String, Object>) entry.getValue();
                    for (Map.Entry<String, Object> rowEntry : rowCollection.entrySet()) {
                        row.add(rowEntry.getValue());
                    }
                } else if (el instanceof Object[]) {
                    Object[] array = (Object[]) el;
                    row.addAll(Arrays.asList(array));
                }

                data.add(row);
                i++;
            }
            return new SimpleDataModel(columnNames, data, realCount);
        } catch (ScriptException e) {
            LogKit.error(e.getMessage(), e);
        }
        return null;
    }

    private static ScriptEngine findScriptEngine() {
        ScriptEngineManager scriptEngineManager = new ScriptEngineManager();
        ScriptEngine engine = scriptEngineManager.getEngineByName("nashorn");
        if (engine == null) {
            List<ScriptEngineFactory> factories = scriptEngineManager.getEngineFactories();
            for (ScriptEngineFactory factory : factories) {
                ScriptEngine current = factory.getScriptEngine();
                if (current != null) {
                    engine = current;
                }
            }
        }
        return engine;
    }
}