From e95cd28f6f9199b1973e54e0668805e28cb477d2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bar=C4=B1=C5=9F=20Soner=20U=C5=9Fakl=C4=B1?= Date: Thu, 2 Jul 2020 20:11:53 -0400 Subject: [PATCH] Zscan (#8458) * feat: zscan * fix: mongodb tests * feat: scan, ip search starts with --- src/database/mongo/helpers.js | 21 ++++++++++++++++++++- src/database/mongo/main.js | 8 ++++++++ src/database/mongo/sorted.js | 15 +-------------- src/database/postgres/main.js | 19 +++++++++++++++++++ src/database/redis/main.js | 12 ++++++++++++ src/database/redis/promisify.js | 1 + src/user/search.js | 6 +++++- test/database/keys.js | 15 +++++++++++++++ 8 files changed, 81 insertions(+), 16 deletions(-) diff --git a/src/database/mongo/helpers.js b/src/database/mongo/helpers.js index 928f7516f2..69b750572c 100644 --- a/src/database/mongo/helpers.js +++ b/src/database/mongo/helpers.js @@ -1,6 +1,7 @@ 'use strict'; -var helpers = module.exports; +const helpers = module.exports; +const utils = require('../../utils'); helpers.noop = function () {}; @@ -48,3 +49,21 @@ helpers.deserializeData = function (data) { helpers.valueToString = function (value) { return String(value); }; + +helpers.buildMatchQuery = function (match) { + let _match = match; + if (match.startsWith('*')) { + _match = _match.substring(1); + } + if (match.endsWith('*')) { + _match = _match.substring(0, _match.length - 1); + } + _match = utils.escapeRegexChars(_match); + if (!match.startsWith('*')) { + _match = '^' + _match; + } + if (!match.endsWith('*')) { + _match += '$'; + } + return _match; +}; diff --git a/src/database/mongo/main.js b/src/database/mongo/main.js index fbd53d088a..cff34cc429 100644 --- a/src/database/mongo/main.js +++ b/src/database/mongo/main.js @@ -1,6 +1,7 @@ 'use strict'; module.exports = function (module) { + const helpers = require('./helpers'); module.flushdb = async function () { await module.client.dropDatabase(); }; @@ -27,6 +28,13 @@ module.exports = function (module) { return item !== undefined && item !== null; }; + module.scan = async function (params) { + const match = helpers.buildMatchQuery(params.match); + return await module.client.collection('objects').distinct( + '_key', { _key: { $regex: new RegExp(match) } } + ); + }; + module.delete = async function (key) { if (!key) { return; diff --git a/src/database/mongo/sorted.js b/src/database/mongo/sorted.js index c17dc06c53..4721c0bc14 100644 --- a/src/database/mongo/sorted.js +++ b/src/database/mongo/sorted.js @@ -472,20 +472,7 @@ module.exports = function (module) { project.score = 1; } - let match = params.match; - if (params.match.startsWith('*')) { - match = match.substring(1); - } - if (params.match.endsWith('*')) { - match = match.substring(0, match.length - 1); - } - match = utils.escapeRegexChars(match); - if (!params.match.startsWith('*')) { - match = '^' + match; - } - if (!params.match.endsWith('*')) { - match += '$'; - } + const match = helpers.buildMatchQuery(params.match); let regex; try { regex = new RegExp(match); diff --git a/src/database/postgres/main.js b/src/database/postgres/main.js index 85c774efe2..37e8f695a1 100644 --- a/src/database/postgres/main.js +++ b/src/database/postgres/main.js @@ -42,6 +42,25 @@ module.exports = function (module) { return res.rows[0].e; }; + module.scan = async function (params) { + let match = params.match; + if (match.startsWith('*')) { + match = '%' + match.substring(1); + } + if (match.endsWith('*')) { + match = match.substring(0, match.length - 1) + '%'; + } + + const res = await module.pool.query({ + text: ` + SELECT o."_key" + FROM "legacy_object_live" o + WHERE o."_key" LIKE '${match}'`, + }); + + return res.rows.map(r => r._key); + }; + module.delete = async function (key) { if (!key) { return; diff --git a/src/database/redis/main.js b/src/database/redis/main.js index bb1f074bca..7313f1c10e 100644 --- a/src/database/redis/main.js +++ b/src/database/redis/main.js @@ -23,6 +23,18 @@ module.exports = function (module) { return exists === 1; }; + module.scan = async function (params) { + let cursor = '0'; + let returnData = []; + do { + /* eslint-disable no-await-in-loop */ + const res = await module.client.async.scan(cursor, 'MATCH', params.match, 'COUNT', 10000); + cursor = res[0]; + returnData = returnData.concat(res[1]); + } while (cursor !== '0'); + return returnData; + }; + module.delete = async function (key) { await module.client.async.del(key); module.objectCache.delObjectCache(key); diff --git a/src/database/redis/promisify.js b/src/database/redis/promisify.js index 320ba4dfe2..d768baad59 100644 --- a/src/database/redis/promisify.js +++ b/src/database/redis/promisify.js @@ -7,6 +7,7 @@ module.exports = function (redisClient) { send_command: util.promisify(redisClient.send_command).bind(redisClient), exists: util.promisify(redisClient.exists).bind(redisClient), + scan: util.promisify(redisClient.scan).bind(redisClient), del: util.promisify(redisClient.del).bind(redisClient), get: util.promisify(redisClient.get).bind(redisClient), diff --git a/src/user/search.js b/src/user/search.js index e2d4c50fb3..6c7d4724eb 100644 --- a/src/user/search.js +++ b/src/user/search.js @@ -1,6 +1,8 @@ 'use strict'; +const _ = require('lodash'); + const meta = require('../meta'); const plugins = require('../plugins'); const db = require('../database'); @@ -128,6 +130,8 @@ module.exports = function (User) { } async function searchByIP(ip) { - return await db.getSortedSetRevRange('ip:' + ip + ':uid', 0, -1); + const ipKeys = await db.scan({ match: 'ip:' + ip + '*' }); + const uids = await db.getSortedSetRevRange(ipKeys, 0, -1); + return _.uniq(uids); } }; diff --git a/test/database/keys.js b/test/database/keys.js index d603d6a4e6..7c5fe7cc25 100644 --- a/test/database/keys.js +++ b/test/database/keys.js @@ -61,6 +61,21 @@ describe('Key methods', function () { }); }); + describe('scan', function () { + it('should scan keys for pattern', async function () { + await db.sortedSetAdd('ip:123:uid', 1, 'a'); + await db.sortedSetAdd('ip:123:uid', 2, 'b'); + await db.sortedSetAdd('ip:124:uid', 2, 'b'); + await db.sortedSetAdd('ip:1:uid', 1, 'a'); + await db.sortedSetAdd('ip:23:uid', 1, 'a'); + const data = await db.scan({ match: 'ip:1*' }); + assert.equal(data.length, 3); + assert(data.includes('ip:123:uid')); + assert(data.includes('ip:124:uid')); + assert(data.includes('ip:1:uid')); + }); + }); + it('should delete a key without error', function (done) { db.delete('testKey', function (err) { assert.ifError(err);