Redis中使用lua脚本

前言

公司中的爬虫项目作为一个接盘来的项目，项目中大致流程是，把要抓取的数据导入到redis中，然后不断pop再去抓取。为了避免redis中的数据丢失，大哥采取的是先pop再将拿到的数据通过set的形式在redis中记录下来。

	redis_client = redis.Redis("127.0.0.1", decode_responses=True)
	queue_name = "test_queue"
	data_key = "grab_data"
	data = redis_client.lpop(queue_name)
	if data is None:
		return
	redis_client.set(data_key, data)
	
	这样的操作看上去没有问题，但，有可能因为网络或者人为等原因，在pop之后还没有来的及去触发set命令。
	所以为了保证这两个操作能够缺一不可(原子性),  由此想到了redis中使用lua脚本。

正文

先上代码：

# 通过lua脚本将pop和set合并为一个原子操作
lua_script = """
local value = redis.call('lpop', ARGV[1])
if value then
    redis.call('set', ARGV[2], value)
end
return value
"""
lua_result = None


def pop_redis_queue(redis_client, queue_name, data_key):
    """
    从redis的队列中pop出一个值, 并且将pop到的值通过set的方式存起来
    :param redis_client: redis操作对象
    :param queue_name: 队列名称
    :data_key: 用于暂时存储获取值的key
    :return: 1. 当data_key对用的value不为空时, 直接返回它的value
             2. 当#{queue_name}队列的长度大于0时, 返回队首元素
             3. 当#{queue_name}队列的长度为0时, 返回None
    """
    # 先看是否有未处理的数据
    value = redis_client.get(data_key)
    if value:
        return str(value)
    global lua_result
    if not lua_result:
        lua_result = redis_client.register_script(lua_script)
    data = lua_result(args=[queue_name, data_key], client=redis_client)
    return data

No bb, just show me the code. 下次再见，再见来不及握手，主要还是我下班了