kettle/rds/script.go

package rds

import (
	"github.com/go-redis/redis"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

/**
 	提供对 redis 中 lua 脚本的简单封装

	lua 主要用于逻辑层面的批量 key 操作。 这和 multi 命令不同：
	1. lua 脚本是 atomic 的， 可以看作事务。(和DBMS事务不同，主要不支持回滚，这里仅仅是提供了“一次性”的逻辑)
	2. lua 脚本执行不会中途打断。
	3. lua 脚本相比 multi 有更小的数据传数量

	实现目标
	支持 lua 从 lua 文件加载
	支持从字符串加载
	支持从目录加载

	加载后，需要保存 redis 保存脚本后返回的对应的 hash

	調用方可根据文件名调用

	从字符串加载的，可自己设置名字
	名字不能重名
*/

// LuaBlock
type LuaBlock struct {
	*redis.Script
	name string
}

type luablocker interface {
	Eval(script string, keys []string, args ...interface{}) *redis.Cmd
	EvalSha(sha1 string, keys []string, args ...interface{}) *redis.Cmd
	ScriptExists(hashes ...string) *redis.BoolSliceCmd
	ScriptLoad(script string) *redis.StringCmd
}

var _ luablocker = (*redis.Client)(nil)
var _ luablocker = (*redis.Ring)(nil)
var _ luablocker = (*redis.ClusterClient)(nil)

func newLuaBlock(name, block string) *LuaBlock {
	return &LuaBlock{
		Script: redis.NewScript(block),
		name:   name,
	}
}

// LuaBlockManager
type LuaScripts struct {
	m map[string]*LuaBlock
	h luablocker
}

func NewLuaScripts(handle luablocker) *LuaScripts {
	return &LuaScripts{
		m: make(map[string]*LuaBlock),
		h: handle,
	}
}

func (self *LuaScripts) Exec(name string, keys []string, args ...interface{}) *redis.Cmd {
	luablock, ok := self.m[name]
	if !ok {
		panic("not load lua script:" + name)
	}
	return luablock.EvalSha(self.h, keys, args)
}

func (self *LuaScripts) LoadString(name string, block string) error {
	_, err := self.h.ScriptLoad(block).Result()
	if err != nil {
		return err
	}
	self.m[name] = newLuaBlock(name, block)
	return nil
}

func (self *LuaScripts) LoadFile(f string) error {
	data, err := ioutil.ReadFile(f)
	if err != nil {
		return err
	}
	filename := filepath.Base(f)            // xxx.lua
	name := strings.Split(filename, ".")[0] //xxx

	return self.LoadString(name, string(data))
}

func (self *LuaScripts) LoadPath(p string) error {
	err := filepath.Walk(p, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		fi, err := os.Stat(path)
		if err != nil {
			return err
		}
		if fi.IsDir() {
			return nil
		}
		return self.LoadFile(path)
	})
	return err
}