From 2a9e179ea8d7456a1cf22f9171bc44de0ae998c4 Mon Sep 17 00:00:00 2001
From: Jens Elstner <jens.elstner@technicallyfeasible.com>
Date: Thu, 13 Feb 2020 13:40:34 +0000
Subject: [PATCH] feat: ability force custom scripts to be readOnly and execute
 on slaves

---
 lib/cluster/index.ts                    |  5 +++--
 lib/command.ts                          |  9 +++++++++
 lib/commander.ts                        |  4 +++-
 lib/script.ts                           |  6 +++++-
 test/functional/cluster/index.ts        | 26 ++++++++++++++++++++++++-
 test/functional/maxRetriesPerRequest.ts |  1 +
 6 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/lib/cluster/index.ts b/lib/cluster/index.ts
index c26fd28d..7636eb66 100644
--- a/lib/cluster/index.ts
+++ b/lib/cluster/index.ts
@@ -554,8 +554,9 @@ class Cluster extends EventEmitter {
     let to = this.options.scaleReads;
     if (to !== "master") {
       const isCommandReadOnly =
-        commands.exists(command.name) &&
-        commands.hasFlag(command.name, "readonly");
+        command.isReadOnly ||
+        (commands.exists(command.name) &&
+          commands.hasFlag(command.name, "readonly"));
       if (!isCommandReadOnly) {
         to = "master";
       }
diff --git a/lib/command.ts b/lib/command.ts
index e26b1b2d..90e36252 100644
--- a/lib/command.ts
+++ b/lib/command.ts
@@ -22,6 +22,10 @@ interface ICommandOptions {
   replyEncoding?: string | null;
   errorStack?: string;
   keyPrefix?: string;
+  /**
+   * Force the command to be readOnly so it will also execute on slaves
+   */
+  readOnly?: boolean;
 }
 
 type ArgumentTransformer = (args: any[]) => any[];
@@ -142,6 +146,7 @@ export default class Command implements ICommand {
   }
 
   public ignore?: boolean;
+  public isReadOnly?: boolean;
 
   private replyEncoding: string | null;
   private errorStack: string;
@@ -185,6 +190,10 @@ export default class Command implements ICommand {
     if (options.keyPrefix) {
       this._iterateKeys(key => options.keyPrefix + key);
     }
+
+    if (options.readOnly) {
+      this.isReadOnly = true;
+    }
   }
 
   private initPromise() {
diff --git a/lib/commander.ts b/lib/commander.ts
index 17c73e35..f499e1e0 100644
--- a/lib/commander.ts
+++ b/lib/commander.ts
@@ -78,13 +78,15 @@ Commander.prototype.send_command = Commander.prototype.call;
  * @param {object} definition
  * @param {string} definition.lua - the lua code
  * @param {number} [definition.numberOfKeys=null] - the number of keys.
+ * @param {boolean} [definition.readOnly=false] - force this script to be readonly so it executes on slaves as well.
  * If omit, you have to pass the number of keys as the first argument every time you invoke the command
  */
 Commander.prototype.defineCommand = function(name, definition) {
   var script = new Script(
     definition.lua,
     definition.numberOfKeys,
-    this.options.keyPrefix
+    this.options.keyPrefix,
+    definition.readOnly
   );
   this.scriptsSet[name] = script;
   this[name] = generateScriptingFunction(script, "utf8");
diff --git a/lib/script.ts b/lib/script.ts
index f8e3d664..fa32092a 100644
--- a/lib/script.ts
+++ b/lib/script.ts
@@ -10,7 +10,8 @@ export default class Script {
   constructor(
     private lua: string,
     private numberOfKeys: number = null,
-    private keyPrefix: string = ""
+    private keyPrefix: string = "",
+    private readOnly: boolean = false
   ) {
     this.sha = createHash("sha1")
       .update(lua)
@@ -29,6 +30,9 @@ export default class Script {
     if (this.keyPrefix) {
       options.keyPrefix = this.keyPrefix;
     }
+    if (this.readOnly) {
+      options.readOnly = true;
+    }
 
     const evalsha = new Command("evalsha", [this.sha].concat(args), options);
     evalsha.isCustomCommand = true;
diff --git a/test/functional/cluster/index.ts b/test/functional/cluster/index.ts
index 13f4c6ed..0b2e4f91 100644
--- a/test/functional/cluster/index.ts
+++ b/test/functional/cluster/index.ts
@@ -1,6 +1,6 @@
 import MockServer from "../../helpers/mock_server";
 import { expect } from "chai";
-import { Cluster } from "../../../lib";
+import { Cluster, default as Redis } from "../../../lib";
 import * as utils from "../../../lib/utils";
 import * as sinon from "sinon";
 
@@ -193,6 +193,30 @@ describe("cluster", function() {
           });
         });
       });
+
+      it("should send custom readOnly scripts to a slave", function(done) {
+        var cluster = new Cluster([{ host: "127.0.0.1", port: "30001" }], {
+          scaleReads: "slave"
+        });
+        cluster.on("ready", function() {
+          const redis: Redis = cluster;
+          redis.defineCommand("test", {
+            numberOfKeys: 1,
+            lua: "return {KEYS[1],ARGV[1],ARGV[2]}",
+            readOnly: true
+          });
+
+          const stub = sinon.stub(utils, "sample").returns("127.0.0.1:30003");
+          redis.test("k1", "a1", "a2", function(err, result) {
+            stub.restore();
+            expect(stub.callCount).to.eql(1);
+            // because of the beforeEach handler this will be the port of the slave called
+            expect(result).to.eql(30003);
+            cluster.disconnect();
+            done();
+          });
+        });
+      });
     });
 
     context("custom", function() {
diff --git a/test/functional/maxRetriesPerRequest.ts b/test/functional/maxRetriesPerRequest.ts
index aca176ce..3c73a71d 100644
--- a/test/functional/maxRetriesPerRequest.ts
+++ b/test/functional/maxRetriesPerRequest.ts
@@ -5,6 +5,7 @@ import { MaxRetriesPerRequestError } from "../../lib/errors";
 describe("maxRetriesPerRequest", function() {
   it("throw the correct error when reached the limit", function(done) {
     var redis = new Redis(9999, {
+      connectTimeout: 1,
       retryStrategy() {
         return 1;
       }