From 723fe8e8e09f85bff8756340f9ca915c26aef766 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 17:59:20 -0400 Subject: [PATCH] feat: zscan (#8457) * feat: zscan * fix: mongodb tests --- src/database/mongo/sorted.js | 42 +++++++++++++++++++++ src/database/postgres/sorted.js | 29 +++++++++++++++ src/database/redis/promisify.js | 2 +- src/database/redis/sorted.js | 30 +++++++++++++++ test/database/sorted.js | 65 +++++++++++++++++++++++++++++++++ 5 files changed, 167 insertions(+), 1 deletion(-) diff --git a/src/database/mongo/sorted.js b/src/database/mongo/sorted.js index 0f4c8a466a..c17dc06c53 100644 --- a/src/database/mongo/sorted.js +++ b/src/database/mongo/sorted.js @@ -466,6 +466,48 @@ module.exports = function (module) { } } + module.getSortedSetScan = async function (params) { + const project = { _id: 0, value: 1 }; + if (params.withScores) { + 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 += '$'; + } + let regex; + try { + regex = new RegExp(match); + } catch (err) { + return []; + } + + const cursor = module.client.collection('objects').find({ + _key: params.key, value: { $regex: regex }, + }, { projection: project }); + + if (params.limit) { + cursor.limit(params.limit); + } + + const data = await cursor.toArray(); + if (!params.withScores) { + return data.map(d => d.value); + } + return data; + }; + module.processSortedSet = async function (setKey, processFn, options) { var done = false; var ids = []; diff --git a/src/database/postgres/sorted.js b/src/database/postgres/sorted.js index 70e2641d45..f9a53706f8 100644 --- a/src/database/postgres/sorted.js +++ b/src/database/postgres/sorted.js @@ -610,6 +610,35 @@ DELETE FROM "legacy_zset" z return q; } + module.getSortedSetScan = 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 z."value", + z."score" + FROM "legacy_object_live" o + INNER JOIN "legacy_zset" z + ON o."_key" = z."_key" + AND o."type" = z."type" + WHERE o."_key" = $1::TEXT + AND z."value" LIKE '${match}' + LIMIT $2::INTEGER`, + values: [params.key, params.limit], + }); + if (!params.withScores) { + return res.rows.map(r => r.value); + } + return res.rows.map(r => ({ value: r.value, score: parseFloat(r.score) })); + }; + module.processSortedSet = async function (setKey, process, options) { const client = await module.pool.connect(); var batchSize = (options || {}).batch || 100; diff --git a/src/database/redis/promisify.js b/src/database/redis/promisify.js index b6a59e1cb8..320ba4dfe2 100644 --- a/src/database/redis/promisify.js +++ b/src/database/redis/promisify.js @@ -48,11 +48,11 @@ module.exports = function (redisClient) { zrank: util.promisify(redisClient.zrank).bind(redisClient), zrevrank: util.promisify(redisClient.zrevrank).bind(redisClient), zincrby: util.promisify(redisClient.zincrby).bind(redisClient), - zrangebylex: util.promisify(redisClient.zrangebylex).bind(redisClient), zrevrangebylex: util.promisify(redisClient.zrevrangebylex).bind(redisClient), zremrangebylex: util.promisify(redisClient.zremrangebylex).bind(redisClient), zlexcount: util.promisify(redisClient.zlexcount).bind(redisClient), + zscan: util.promisify(redisClient.zscan).bind(redisClient), lpush: util.promisify(redisClient.lpush).bind(redisClient), rpush: util.promisify(redisClient.rpush).bind(redisClient), diff --git a/src/database/redis/sorted.js b/src/database/redis/sorted.js index 443325f1a2..8aa5de5c38 100644 --- a/src/database/redis/sorted.js +++ b/src/database/redis/sorted.js @@ -276,4 +276,34 @@ module.exports = function (module) { } return await module.client.async[method](args); } + + module.getSortedSetScan = async function (params) { + let cursor = '0'; + + const returnData = []; + let done = false; + do { + /* eslint-disable no-await-in-loop */ + const res = await module.client.async.zscan(params.key, cursor, 'MATCH', params.match, 'COUNT', 100); + cursor = res[0]; + done = cursor === '0'; + const data = res[1]; + + for (let i = 0; i < data.length; i += 2) { + const value = data[i]; + const score = parseFloat(data[i + 1]); + if (params.withScores) { + returnData.push({ value: value, score: score }); + } else { + returnData.push(value); + } + if (params.limit && returnData.length >= params.limit) { + done = true; + break; + } + } + } while (!done); + + return returnData; + }; }; diff --git a/test/database/sorted.js b/test/database/sorted.js index eb7d578526..2e0a8d72fa 100644 --- a/test/database/sorted.js +++ b/test/database/sorted.js @@ -26,6 +26,71 @@ describe('Sorted Set methods', function () { ], done); }); + describe('sortedSetScan', function () { + it('should find matches in sorted set containing substring', async () => { + await db.sortedSetAdd('scanzset', [1, 2, 3, 4, 5, 6], ['aaaa', 'bbbb', 'bbcc', 'ddd', 'dddd', 'fghbc']); + const data = await db.getSortedSetScan({ + key: 'scanzset', + match: '*bc*', + }); + assert(data.includes('bbcc')); + assert(data.includes('fghbc')); + }); + + it('should find matches in sorted set with scores', async () => { + const data = await db.getSortedSetScan({ + key: 'scanzset', + match: '*bc*', + withScores: true, + }); + data.sort((a, b) => a.score - b.score); + assert.deepStrictEqual(data, [{ value: 'bbcc', score: 3 }, { value: 'fghbc', score: 6 }]); + }); + + it('should find matches in sorted set with a limit', async () => { + await db.sortedSetAdd('scanzset2', [1, 2, 3, 4, 5, 6], ['aaab', 'bbbb', 'bbcb', 'ddb', 'dddd', 'fghbc']); + const data = await db.getSortedSetScan({ + key: 'scanzset2', + match: '*b*', + limit: 2, + }); + assert.equal(data.length, 2); + }); + + it('should work for special characters', async () => { + await db.sortedSetAdd('scanzset3', [1, 2, 3, 4, 5], ['aaab{', 'bbbb', 'bbcb{', 'ddb', 'dddd']); + const data = await db.getSortedSetScan({ + key: 'scanzset3', + match: '*b{', + limit: 2, + }); + assert(data.includes('aaab{')); + assert(data.includes('bbcb{')); + }); + + it('should find everything starting with string', async () => { + await db.sortedSetAdd('scanzset4', [1, 2, 3, 4, 5], ['aaab{', 'bbbb', 'bbcb', 'ddb', 'dddd']); + const data = await db.getSortedSetScan({ + key: 'scanzset4', + match: 'b*', + limit: 2, + }); + assert(data.includes('bbbb')); + assert(data.includes('bbcb')); + }); + + it('should find everything ending with string', async () => { + await db.sortedSetAdd('scanzset5', [1, 2, 3, 4, 5, 6], ['aaab{', 'bbbb', 'bbcb', 'ddb', 'dddd', 'adb']); + const data = await db.getSortedSetScan({ + key: 'scanzset5', + match: '*db', + }); + assert.equal(data.length, 2); + assert(data.includes('ddb')); + assert(data.includes('adb')); + }); + }); + describe('sortedSetAdd()', function () { it('should add an element to a sorted set', function (done) { db.sortedSetAdd('sorted1', 1, 'value1', function (err) {