kettle/rds/script.go

package rds

import (
	"github.com/go-redis/redis"
	"io/ioutil"
	"git.wanbits.io/joe/kettle/utl"
	"os"
	"path/filepath"
	"strings"
)

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

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

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

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

	調用方可根据文件名调用

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

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)

type LuaSrc struct {
	code, hash string
}

// LuaBlockManager
type LuaScriptManager struct {
	m map[string]*LuaSrc
	h luablocker
}

func NewLuaScriptManager(handle luablocker) *LuaScriptManager {
	return &LuaScriptManager{
		m: make(map[string]*LuaSrc),
		h: handle,
	}
}

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

func (self *LuaScriptManager) LoadString(name string, src string) error {
	hash, err := self.h.ScriptLoad(src).Result()
	if err != nil {
		return err
	}
	self.m[name] = &LuaSrc{
		code: src,
		hash: hash,
	}
	return nil
}

func (self *LuaScriptManager) LoadFile(f string) error {
	ext := filepath.Ext(f)
	if strings.ToLower(ext) != ".lua" {
		return utl.ErrParameters
	}
	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 *LuaScriptManager) 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
}