client-libraries/lua/redis.lua

   1 module('Redis', package.seeall)
   2
   3 require('socket')       -- requires LuaSocket as a dependency
   4
   5 -- ############################################################################
   6
   7 local protocol = {
   8     newline = '\r\n', ok = 'OK', err = 'ERR', null = 'nil',
   9 }
  10
  11 -- ############################################################################
  12
  13 local function toboolean(value)
  14     return value == 1
  15 end
  16
  17 local function _write(self, buffer)
  18     local _, err = self.socket:send(buffer)
  19     if err then error(err) end
  20 end
  21
  22 local function _read(self, len)
  23     if len == nil then len = '*l' end
  24     local line, err = self.socket:receive(len)
  25     if not err then return line else error('Connection error: ' .. err) end
  26 end
  27
  28 -- ############################################################################
  29
  30 local function _read_response(self)
  31     if options and options.close == true then return end
  32
  33     local res    = _read(self)
  34     local prefix = res:sub(1, -#res)
  35     local response_handler = protocol.prefixes[prefix]
  36
  37     if not response_handler then
  38         error("Unknown response prefix: " .. prefix)
  39     else
  40         return response_handler(self, res)
  41     end
  42 end
  43
  44
  45 local function _send_raw(self, buffer)
  46     -- TODO: optimize
  47     local bufferType = type(buffer)
  48
  49     if bufferType == 'string' then
  50         _write(self, buffer)
  51     elseif bufferType == 'table' then
  52         _write(self, table.concat(buffer))
  53     else
  54         error('Argument error: ' .. bufferType)
  55     end
  56
  57     return _read_response(self)
  58 end
  59
  60 local function _send_inline(self, command, ...)
  61     if arg.n == 0 then
  62         _write(self, command .. protocol.newline)
  63     else
  64         local arguments = arg
  65         arguments.n = nil
  66
  67         if #arguments > 0 then
  68             arguments = table.concat(arguments, ' ')
  69         else
  70             arguments = ''
  71         end
  72
  73         _write(self, command .. ' ' .. arguments .. protocol.newline)
  74     end
  75
  76     return _read_response(self)
  77 end
  78
  79 local function _send_bulk(self, command, ...)
  80     local arguments = arg
  81     local data      = tostring(table.remove(arguments))
  82     arguments.n = nil
  83
  84     -- TODO: optimize
  85     if #arguments > 0 then
  86         arguments = table.concat(arguments, ' ')
  87     else
  88         arguments = ''
  89     end
  90
  91     return _send_raw(self, {
  92         command, ' ', arguments, ' ', #data, protocol.newline, data, protocol.newline
  93     })
  94 end
  95
  96
  97 local function _read_line(self, response)
  98     return response:sub(2)
  99 end
 100
 101 local function _read_error(self, response)
 102     local err_line = response:sub(2)
 103
 104     if err_line:sub(1, 3) == protocol.err then
 105         error("Redis error: " .. err_line:sub(5))
 106     else
 107         error("Redis error: " .. err_line)
 108     end
 109 end
 110
 111 local function _read_bulk(self, response)
 112     local str = response:sub(2)
 113     local len = tonumber(str)
 114
 115     if not len then
 116         error('Cannot parse ' .. str .. ' as data length.')
 117     else
 118         if len == -1 then return nil end
 119         local data = _read(self, len + 2)
 120         return data:sub(1, -3);
 121     end
 122 end
 123
 124 local function _read_multibulk(self, response)
 125     local str = response:sub(2)
 126
 127     -- TODO: add a check if the returned value is indeed a number
 128     local list_count = tonumber(str)
 129
 130     if list_count == -1 then
 131         return nil
 132     else
 133         local list = {}
 134
 135         if list_count > 0 then
 136             for i = 1, list_count do
 137                 table.insert(list, i, _read_bulk(self, _read(self)))
 138             end
 139         end
 140
 141         return list
 142     end
 143 end
 144
 145 local function _read_integer(self, response)
 146     local res = response:sub(2)
 147     local number = tonumber(res)
 148
 149     if not number then
 150         if res == protocol.null then
 151             return nil
 152         else
 153             error('Cannot parse ' .. res .. ' as numeric response.')
 154         end
 155     end
 156
 157     return number
 158 end
 159
 160 -- ############################################################################
 161
 162 protocol.prefixes = {
 163     ['+'] = _read_line,
 164     ['-'] = _read_error,
 165     ['$'] = _read_bulk,
 166     ['*'] = _read_multibulk,
 167     [':'] = _read_integer,
 168 }
 169
 170 -- ############################################################################
 171
 172 local methods = {
 173     -- miscellaneous commands
 174     ping    = {
 175         'PING', _send_inline, function(response)
 176             if response == 'PONG' then return true else return false end
 177         end
 178     },
 179     echo    = { 'ECHO', _send_bulk },
 180     -- TODO: the server returns an empty -ERR on authentication failure
 181     auth    = { 'AUTH' },
 182
 183     -- connection handling
 184     quit    = { 'QUIT', function(self, command)
 185             _write(self, command .. protocol.newline)
 186         end
 187     },
 188
 189     -- commands operating on string values
 190     set             = { 'SET', _send_bulk },
 191     set_preserve    = { 'SETNX', _send_bulk, toboolean },
 192     get             = { 'GET' },
 193     get_multiple    = { 'MGET' },
 194     increment       = { 'INCR' },
 195     increment_by    = { 'INCRBY' },
 196     decrement       = { 'DECR' },
 197     decrement_by    = { 'DECRBY' },
 198     exists          = { 'EXISTS', _send_inline, toboolean },
 199     delete          = { 'DEL', _send_inline, toboolean },
 200     type            = { 'TYPE' },
 201
 202     -- commands operating on the key space
 203     keys            = {
 204         'KEYS',  _send_inline, function(response)
 205             local keys = {}
 206             response:gsub('%w+', function(key)
 207                 table.insert(keys, key)
 208             end)
 209             return keys
 210         end
 211     },
 212     random_key      = { 'RANDOMKEY' },
 213     rename          = { 'RENAME' },
 214     rename_preserve = { 'RENAMENX' },
 215     database_size   = { 'DBSIZE' },
 216
 217     -- commands operating on lists
 218     push_tail   = { 'RPUSH', _send_bulk },
 219     push_head   = { 'LPUSH', _send_bulk },
 220     list_length = { 'LLEN', _send_inline, function(response, key)
 221             --[[ TODO: redis seems to return a -ERR when the specified key does
 222                        not hold a list value, but this behaviour is not
 223                        consistent with the specs docs. This might be due to the
 224                        -ERR response paradigm being new, which supersedes the
 225                        check for negative numbers to identify errors. ]]
 226             if response == -2 then
 227                 error('Key ' .. key .. ' does not hold a list value')
 228             end
 229             return response
 230         end
 231     },
 232     list_range  = { 'LRANGE' },
 233     list_trim   = { 'LTRIM' },
 234     list_index  = { 'LINDEX' },
 235     list_set    = { 'LSET', _send_bulk },
 236     list_remove = { 'LREM', _send_bulk },
 237     pop_first   = { 'LPOP' },
 238     pop_last    = { 'RPOP' },
 239
 240     -- commands operating on sets
 241     set_add                = { 'SADD' },
 242     set_remove             = { 'SREM' },
 243     set_cardinality        = { 'SCARD' },
 244     set_is_member          = { 'SISMEMBER' },
 245     set_intersection       = { 'SINTER' },
 246     set_intersection_store = { 'SINTERSTORE' },
 247     set_members            = { 'SMEMBERS' },
 248
 249     -- multiple databases handling commands
 250     select_database = { 'SELECT' },
 251     move_key        = { 'MOVE' },
 252     flush_database  = { 'FLUSHDB' },
 253     flush_databases = { 'FLUSHALL' },
 254
 255     -- sorting
 256     --[[
 257         TODO: should we pass sort parameters as a table? e.g:
 258                 params = {
 259                     by    = 'weight_*',
 260                     get   = 'object_*',
 261                     limit = { 0, 10 },
 262                     sort  = { 'desc', 'alpha' }
 263                 }
 264     --]]
 265     sort    = { 'SORT' },
 266
 267     -- persistence control commands
 268     save            = { 'SAVE' },
 269     background_save = { 'BGSAVE' },
 270     last_save       = { 'LASTSAVE' },
 271     shutdown        = { 'SHUTDOWN', function(self, command)
 272             _write(self, command .. protocol.newline)
 273         end
 274     },
 275
 276     -- remote server control commands
 277     info    = {
 278         'INFO', _send_inline, function(response)
 279             local info = {}
 280             response:gsub('([^\r\n]*)\r\n', function(kv)
 281                 local k,v = kv:match(('([^:]*):([^:]*)'):rep(1))
 282                 info[k] = v
 283             end)
 284             return info
 285         end
 286     },
 287 }
 288
 289 function connect(host, port)
 290     local client_socket = socket.connect(host, port)
 291
 292     if not client_socket then
 293         error('Could not connect to ' .. host .. ':' .. port)
 294     end
 295
 296     local redis_client = {
 297         socket  = client_socket,
 298         raw_cmd = function(self, buffer)
 299             return _send_raw(self, buffer .. protocol.newline)
 300         end,
 301     }
 302
 303     return setmetatable(redis_client, {
 304         __index = function(self, method)
 305             local redis_meth = methods[method]
 306             if redis_meth then
 307                 return function(self, ...)
 308                     if not redis_meth[2] then
 309                         table.insert(redis_meth, 2, _send_inline)
 310                     end
 311
 312                     local response = redis_meth[2](self, redis_meth[1], ...)
 313                     if redis_meth[3] then
 314                         return redis_meth[3](response, ...)
 315                     else
 316                         return response
 317                     end
 318                 end
 319             end
 320         end
 321     })
 322 end