client-libraries/lua/redis.lua

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