Merge pull request #395 from unknowntpo/feat-align-rueidis-with-go-re…

…dis-gear feat: Add `GearsCmdable` to `go-redis` API adapter
redis · Oct 20, 2023 · 94323fd · 94323fd
2 parents 6ecaa43 + 543d05b
commit 94323fd
Show file tree

Hide file tree

Showing 6 changed files with 289 additions and 7 deletions.
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -30,7 +30,7 @@ services:
     ports:
       - "6377:6379"
   compat:
-    image: redis:7.2-rc-alpine
+    image: redis/redis-stack:7.2.0-v3
     ports:
       - "6378:6379"
   compat5:

diff --git a/hack/cmds/commands_gears2.json b/hack/cmds/commands_gears2.json
@@ -84,10 +84,13 @@
         "optional": true
       },
       {
-        "name": "config",
-        "type": "string",
-        "display_text": "config",
-        "token": "CONFIG",
+        "command": "CONFIG",
+        "name": [
+          "config"
+        ],
+        "type": [
+          "string"
+        ],
         "optional": true
       },
       {

diff --git a/internal/cmds/gen_triggers_and_functions.go b/internal/cmds/gen_triggers_and_functions.go
diff --git a/rueidiscompat/adapter.go b/rueidiscompat/adapter.go
@@ -387,13 +387,31 @@ type Cmdable interface {
 	ACLDryRun(ctx context.Context, username string, command ...any) *StringCmd
 
 	// TODO ModuleLoadex(ctx context.Context, conf *ModuleLoadexConfig) *StringCmd
+	GearsCmdable
 }
 
+// Align with go-redis
+// https://github.com/redis/go-redis/blob/f994ff1cd96299a5c8029ae3403af7b17ef06e8a/gears_commands.go#L9-L19
+type GearsCmdable interface {
+	TFunctionLoad(ctx context.Context, lib string) *StatusCmd
+	TFunctionLoadArgs(ctx context.Context, lib string, options *TFunctionLoadOptions) *StatusCmd
+	TFunctionDelete(ctx context.Context, libName string) *StatusCmd
+	TFunctionList(ctx context.Context) *MapStringInterfaceSliceCmd
+	TFunctionListArgs(ctx context.Context, options *TFunctionListOptions) *MapStringInterfaceSliceCmd
+	TFCall(ctx context.Context, libName string, funcName string, numKeys int) *Cmd
+	TFCallArgs(ctx context.Context, libName string, funcName string, numKeys int, options *TFCallOptions) *Cmd
+	TFCallASYNC(ctx context.Context, libName string, funcName string, numKeys int) *Cmd
+	TFCallASYNCArgs(ctx context.Context, libName string, funcName string, numKeys int, options *TFCallOptions) *Cmd
+}
+
+var _ Cmdable = (*Compat)(nil)
+
 type Compat struct {
 	client rueidis.Client
 	maxp   int
 }
 
+// CacheCompat implements commands that support client-side caching.
 type CacheCompat struct {
 	client rueidis.Client
 	ttl    time.Duration
@@ -2855,6 +2873,103 @@ func (c *Compat) doIntCmdPrimaries(ctx context.Context, fn func(c rueidis.Client
 	return ret
 }
 
+func (c *Compat) TFunctionLoad(ctx context.Context, lib string) *StatusCmd {
+	cmd := c.client.B().TfunctionLoad().LibraryCode(lib).Build()
+	resp := c.client.Do(ctx, cmd)
+	return newStatusCmd(resp)
+}
+
+func (c *Compat) TFunctionLoadArgs(ctx context.Context, lib string, options *TFunctionLoadOptions) *StatusCmd {
+	b := c.client.B()
+	var cmd cmds.Completed
+	if options.Replace {
+		cmd = b.TfunctionLoad().Replace().Config(options.Config).LibraryCode(lib).Build()
+	} else {
+		cmd = b.TfunctionLoad().Config(options.Config).LibraryCode(lib).Build()
+	}
+	resp := c.client.Do(ctx, cmd)
+	return newStatusCmd(resp)
+}
+
+func (c *Compat) TFunctionDelete(ctx context.Context, libName string) *StatusCmd {
+	cmd := c.client.B().TfunctionDelete().LibraryName(libName).Build()
+	resp := c.client.Do(ctx, cmd)
+	return newStatusCmd(resp)
+}
+
+func (c *Compat) TFunctionList(ctx context.Context) *MapStringInterfaceSliceCmd {
+	cmd := c.client.B().TfunctionList().Build()
+	resp := c.client.Do(ctx, cmd)
+	return newMapStringInterfaceSliceCmd(resp)
+}
+
+func (c *Compat) TFunctionListArgs(ctx context.Context, options *TFunctionListOptions) *MapStringInterfaceSliceCmd {
+	cmd := c.client.B().TfunctionList()
+	if options.Library != "" {
+		cmd.LibraryName(options.Library)
+	}
+	if options.Withcode {
+		cmd.Withcode()
+	}
+	if options.Verbose > 0 {
+		cmd.Verbose()
+		for i := 0; i < options.Verbose; i++ {
+			cmd.V()
+		}
+	}
+	cmdCompleted := cmd.Build()
+	resp := c.client.Do(ctx, cmdCompleted)
+	return newMapStringInterfaceSliceCmd(resp)
+}
+
+func (c *Compat) TFCall(ctx context.Context, libName string, funcName string, numKeys int) *Cmd {
+	cmd := c.client.
+		B().
+		Tfcall().
+		LibraryFunction(fmt.Sprintf("%s.%s", libName, funcName)).
+		Numkeys(int64(numKeys)).
+		Build()
+	resp := c.client.Do(ctx, cmd)
+	return newCmd(resp)
+}
+
+func (c *Compat) TFCallArgs(ctx context.Context, libName string, funcName string, numKeys int, options *TFCallOptions) *Cmd {
+	cmd := c.client.
+		B().
+		Tfcall().
+		LibraryFunction(fmt.Sprintf("%s.%s", libName, funcName)).
+		Numkeys(int64(numKeys)).
+		Key(options.Keys...).
+		Arg(options.Arguments...).
+		Build()
+	resp := c.client.Do(ctx, cmd)
+	return newCmd(resp)
+}
+
+func (c *Compat) TFCallASYNC(ctx context.Context, libName string, funcName string, numKeys int) *Cmd {
+	cmd := c.client.
+		B().
+		Tfcallasync().
+		LibraryFunction(fmt.Sprintf("%s.%s", libName, funcName)).
+		Numkeys(int64(numKeys)).
+		Build()
+	resp := c.client.Do(ctx, cmd)
+	return newCmd(resp)
+}
+
+func (c *Compat) TFCallASYNCArgs(ctx context.Context, libName string, funcName string, numKeys int, options *TFCallOptions) *Cmd {
+	cmd := c.client.
+		B().
+		Tfcallasync().
+		LibraryFunction(fmt.Sprintf("%s.%s", libName, funcName)).
+		Numkeys(int64(numKeys)).
+		Key(options.Keys...).
+		Arg(options.Arguments...).
+		Build()
+	resp := c.client.Do(ctx, cmd)
+	return newCmd(resp)
+}
+
 func (c CacheCompat) BitCount(ctx context.Context, key string, bitCount *BitCount) *IntCmd {
 	var resp rueidis.RedisResult
 	if bitCount == nil {

diff --git a/rueidiscompat/adapter_test.go b/rueidiscompat/adapter_test.go
@@ -8739,6 +8739,102 @@ func testAdapterCache(resp3 bool) {
 		//	Expect(value.Number).To(Equal(42))
 		//})
 	})
+
+	Describe("GearsCmdable", func() {
+		BeforeEach(func() {
+			Expect(adapter.FlushDB(ctx).Err()).NotTo(HaveOccurred())
+			adapter.TFunctionDelete(ctx, "lib1")
+		})
+		// Copied from go-redis
+		// https://github.com/redis/go-redis/blob/f994ff1cd96299a5c8029ae3403af7b17ef06e8a/gears_commands_test.go
+		It("should TFunctionLoad, TFunctionLoadArgs and TFunctionDelete ", Label("gears", "tfunctionload"), func() {
+			resultAdd, err := adapter.TFunctionLoad(ctx, libCode("lib1")).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resultAdd).To(BeEquivalentTo("OK"))
+			opt := &TFunctionLoadOptions{Replace: true, Config: `{"last_update_field_name":"last_update"}`}
+			resultAdd, err = adapter.TFunctionLoadArgs(ctx, libCodeWithConfig("lib1"), opt).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resultAdd).To(BeEquivalentTo("OK"))
+		})
+		It("should TFunctionList", Label("gears", "tfunctionlist"), func() {
+			resultAdd, err := adapter.TFunctionLoad(ctx, libCode("lib1")).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resultAdd).To(BeEquivalentTo("OK"))
+			resultList, err := adapter.TFunctionList(ctx).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resultList[0]["engine"]).To(BeEquivalentTo("js"))
+			opt := &TFunctionListOptions{Withcode: true, Verbose: 2}
+			resultListArgs, err := adapter.TFunctionListArgs(ctx, opt).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resultListArgs[0]["code"]).NotTo(BeEquivalentTo(""))
+		})
+
+		It("should TFCall", Label("gears", "tfcall"), func() {
+			var resultAdd interface{}
+			resultAdd, err := adapter.TFunctionLoad(ctx, libCode("lib1")).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resultAdd).To(BeEquivalentTo("OK"))
+			resultAdd, err = adapter.TFCall(ctx, "lib1", "foo", 0).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resultAdd).To(BeEquivalentTo("bar"))
+		})
+
+		It("should TFCallArgs", Label("gears", "tfcallargs"), func() {
+			var resultAdd interface{}
+			resultAdd, err := adapter.TFunctionLoad(ctx, libCode("lib1")).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resultAdd).To(BeEquivalentTo("OK"))
+			opt := &TFCallOptions{Arguments: []string{"foo", "bar"}}
+			resultAdd, err = adapter.TFCallArgs(ctx, "lib1", "foo", 0, opt).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resultAdd).To(BeEquivalentTo("bar"))
+		})
+
+		It("should TFCallASYNC", Label("gears", "TFCallASYNC"), func() {
+			var resultAdd interface{}
+			resultAdd, err := adapter.TFunctionLoad(ctx, libCode("lib1")).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resultAdd).To(BeEquivalentTo("OK"))
+			resultAdd, err = adapter.TFCallASYNC(ctx, "lib1", "foo", 0).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resultAdd).To(BeEquivalentTo("bar"))
+		})
+
+		It("should TFCallASYNCArgs", Label("gears", "TFCallASYNCargs"), func() {
+			var resultAdd interface{}
+			resultAdd, err := adapter.TFunctionLoad(ctx, libCode("lib1")).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resultAdd).To(BeEquivalentTo("OK"))
+			opt := &TFCallOptions{Arguments: []string{"foo", "bar"}}
+			resultAdd, err = adapter.TFCallASYNCArgs(ctx, "lib1", "foo", 0, opt).Result()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(resultAdd).To(BeEquivalentTo("bar"))
+		})
+	})
+}
+
+func libCode(libName string) string {
+	return fmt.Sprintf("#!js api_version=1.0 name=%s\n redis.registerFunction('foo', ()=>{{return 'bar'}})", libName)
+}
+
+func libCodeWithConfig(libName string) string {
+	lib := `#!js api_version=1.0 name=%s
+
+	var last_update_field_name = "__last_update__"
+	
+	if (redis.config.last_update_field_name !== undefined) {
+		if (typeof redis.config.last_update_field_name != 'string') {
+			throw "last_update_field_name must be a string";
+		}
+		last_update_field_name = redis.config.last_update_field_name
+	}
+	
+	redis.registerFunction("hset", function(client, key, field, val){
+		// get the current time in ms
+		var curr_time = client.call("time")[0];
+		return client.call('hset', key, field, val, last_update_field_name, curr_time);
+	});`
+	return fmt.Sprintf(lib, libName)
 }
 
 type numberStruct struct {

diff --git a/rueidiscompat/command.go b/rueidiscompat/command.go
@@ -2792,3 +2792,71 @@ func formatSec(dur time.Duration) int64 {
 	}
 	return int64(dur / time.Second)
 }
+
+// https://github.com/redis/go-redis/blob/f994ff1cd96299a5c8029ae3403af7b17ef06e8a/gears_commands.go#L21C1-L35C2
+type TFunctionLoadOptions struct {
+	Replace bool
+	Config  string
+}
+
+type TFunctionListOptions struct {
+	Withcode bool
+	Verbose  int
+	Library  string
+}
+
+type TFCallOptions struct {
+	Keys      []string
+	Arguments []string
+}
+
+type MapStringInterfaceSliceCmd struct {
+	err error
+	val []map[string]interface{}
+}
+
+func (cmd *MapStringInterfaceSliceCmd) SetVal(val []map[string]interface{}) {
+	cmd.val = val
+}
+
+func (cmd *MapStringInterfaceSliceCmd) Val() []map[string]interface{} {
+	return cmd.val
+}
+
+func (cmd *MapStringInterfaceSliceCmd) Err() error {
+	return cmd.err
+}
+
+func (cmd *MapStringInterfaceSliceCmd) Result() ([]map[string]interface{}, error) {
+	return cmd.Val(), cmd.Err()
+}
+
+func newMapStringInterfaceSliceCmd(res rueidis.RedisResult) *MapStringInterfaceSliceCmd {
+	arr, err := res.ToArray()
+	if err != nil {
+		return &MapStringInterfaceSliceCmd{err: err}
+	}
+	out := &MapStringInterfaceSliceCmd{val: make([]map[string]any, 0, len(arr))}
+	for _, ele := range arr {
+		m, err := ele.AsMap()
+		eleMap := make(map[string]any, len(m))
+		if err != nil {
+			out.err = err
+			return out
+		}
+		for k, v := range m {
+			var val any
+			if !v.IsNil() {
+				var err error
+				val, err = v.ToAny()
+				if err != nil {
+					out.err = err
+					return out
+				}
+			}
+			eleMap[k] = val
+		}
+		out.val = append(out.val, eleMap)
+	}
+	return out
+}