org.bdware.doip.RocksDBUtil Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of doip-audit-tool Show documentation
doip audit tool developed by bdware
There is a newer version: 1.5.4
package org.bdware.doip;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.rocksdb.Options;
import org.rocksdb.RocksDB;
import org.rocksdb.RocksDBException;
import org.rocksdb.RocksIterator;

import java.io.File;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

public class RocksDBUtil {
    private static final Logger LOGGER = LogManager.getLogger(RocksDBUtil.class);

    static Map cacheDB = new HashMap<>();
    static String dbPath = "./";

    static {
        RocksDB.loadLibrary();
    }

    RocksDB rocksDB;
    String path;

    public RocksDBUtil(String path, boolean readOnly) {
        try {
            this.path = path;
            Options options = new Options();
            options.setCreateIfMissing(true);
            File parent = new File(dbPath + "DBRepoStorage/");
            File dir = new File(parent, path);
            LOGGER.info("init RocksDB in " + dir.getAbsolutePath());
            if (!dir.exists()) {
                LOGGER.trace("create directory " + dir.getAbsolutePath() + ": " + dir.mkdirs());
            }
            File lockFile = new File(dir, "LOCK");
            LOGGER.trace("delete file" + lockFile.getAbsolutePath() + ": " + lockFile.delete());
            if (readOnly) {
                rocksDB = RocksDB.openReadOnly(options, dir.getAbsolutePath());
            } else {
                rocksDB = RocksDB.open(options, dir.getAbsolutePath());
            }
        } catch (RocksDBException e) {
            e.printStackTrace();
        }
    }

    public static RocksDBUtil loadDB(String path, boolean readOnly) {
        if (cacheDB.containsKey(path)) {
            return cacheDB.get(path);
        }
        RocksDBUtil ret = new RocksDBUtil(path, readOnly);
        cacheDB.put(path, ret);
        return ret;
    }

    public static RocksDBUtil loadDB(String path, String readOnly) {
        return loadDB(path, Boolean.parseBoolean(readOnly));
    }

    public void close() {
        rocksDB.close();
        cacheDB.remove(path);
    }

    public String get(String key) {
        try {
            return new String(rocksDB.get(key.getBytes()), StandardCharsets.UTF_8);
        } catch (Exception e) {
        }
        return null;
    }

    public byte[] getAsByte(String key) {
        try {
            return rocksDB.get(key.getBytes());
        } catch (Exception e) {
        }
        return null;
    }

    public long estimateKeySize() {
        try {
            return rocksDB.getLongProperty("rocksdb.estimate-num-keys");
        } catch (Exception e) {
            e.printStackTrace();
            return 0;
        }
    }

    public void put(String key, String value) {
        try {
            rocksDB.put(key.getBytes(), value.getBytes());
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public void put(String key, byte[] value) {
        try {
            rocksDB.put(key.getBytes(), value);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public String delete(String key) {
        try {
            rocksDB.delete(key.getBytes());

            return "success";
        } catch (Exception e) {
            e.printStackTrace();
        }
        return "failed";
    }

    public RocksIterator newIterator() {
        try {
            return rocksDB.newIterator();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }


}