loading redis keys

plugins/dbgate-plugin-redis/src/backend/driver.js

@@ -1,4 +1,5 @@
 const _ = require('lodash');
+const async = require('async');
 const stream = require('stream');
 const driverBase = require('../frontend/driver');
 const Analyser = require('./Analyser');
@@ -9,11 +10,14 @@ const driver = {
   ...driverBase,
   analyserClass: Analyser,
   async connect({ server, port, password, database }) {
+    let db = 0;
+    if (_.isString(database) && database.startsWith('db')) db = parseInt(database.substring(2));
+    if (_.isNumber(database)) db = database;
     const pool = new Redis({
       host: server,
       port,
       password,
-      db: 0,
+      db,
     });
     return pool;
   },
@@ -62,6 +66,61 @@ const driver = {
 
     return _.range(16).map((index) => ({ name: `db${index}`, extInfo: info[`db${index}`], sortOrder: index }));
   },
+
+  async loadKeys(pool, root = '') {
+    const keys = await this.getKeys(pool, root);
+    const res = this.extractKeysFromLevel(root, keys);
+    await this.enrichKeyInfo(pool, res);
+    return res;
+  },
+
+  async getKeys(pool, root = '') {
+    const res = [];
+    let cursor = 0;
+    do {
+      const [strCursor, keys] = await pool.scan(cursor, 'MATCH', root ? `${root}:*` : '*', 'COUNT', 100);
+      res.push(...keys);
+      cursor = parseInt(strCursor);
+    } while (cursor > 0);
+    return res;
+  },
+
+  extractKeysFromLevel(root, keys) {
+    const prefix = root ? `${root}:` : '';
+    const rootSplit = _.compact(root.split(':'));
+    const res = {};
+    for (const key of keys) {
+      if (!key.startsWith(prefix)) continue;
+      const keySplit = key.split(':');
+      if (keySplit.length > rootSplit.length) {
+        if (keySplit.length == rootSplit.length + 1) {
+          res[keySplit[rootSplit.length]] = {
+            text: keySplit[rootSplit.length],
+            key,
+          };
+        } else {
+          res[keySplit[rootSplit.length]] = {
+            text: keySplit[rootSplit.length],
+            type: 'dir',
+          };
+        }
+      }
+    }
+    return Object.values(res);
+  },
+
+  async enrichOneKeyInfo(pool, item) {
+    const type = await pool.type(item.key);
+    item.type = type;
+  },
+
+  async enrichKeyInfo(pool, levelInfo) {
+    await async.eachLimit(
+      levelInfo.filter((x) => x.key),
+      10,
+      async (item) => await this.enrichOneKeyInfo(pool, item)
+    );
+  },
 };
 
 module.exports = driver;