feat: zscan (#8457)

* feat: zscan * fix: mongodb tests
NodeBB · Jul 2, 2020 · 723fe8e · 723fe8e
1 parent c279875
commit 723fe8e
Show file tree

Hide file tree

Showing 5 changed files with 167 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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) {