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     increment     = inline('INCR'),
 243     increment_by  = inline('INCRBY'),
 244     decrement     = inline('DECR'),
 245     decrement_by  = inline('DECRBY'),
 246     exists        = inline('EXISTS', toboolean),
 247     delete        = inline('DEL', toboolean),
 248     type          = inline('TYPE'),
 249
 250     -- commands operating on the key space
 251     keys          = inline('KEYS',
 252         function(response)
 253             local keys = {}
 254             response:gsub('%w+', function(key)
 255                 table.insert(keys, key)
 256             end)
 257             return keys
 258         end
 259     ),
 260     random_key       = inline('RANDOMKEY'),
 261     rename           = inline('RENAME'),
 262     rename_preserve  = inline('RENAMENX'),
 263     expire           = inline('EXPIRE', toboolean),
 264     database_size    = inline('DBSIZE'),
 265
 266     -- commands operating on lists
 267     push_tail    = bulk('RPUSH'),
 268     push_head    = bulk('LPUSH'),
 269     list_length  = inline('LLEN'),
 270     list_range   = inline('LRANGE'),
 271     list_trim    = inline('LTRIM'),
 272     list_index   = inline('LINDEX'),
 273     list_set     = bulk('LSET'),
 274     list_remove  = bulk('LREM'),
 275     pop_first    = inline('LPOP'),
 276     pop_last     = inline('RPOP'),
 277
 278     -- commands operating on sets
 279     set_add                 = inline('SADD'),
 280     set_remove              = inline('SREM'),
 281     set_cardinality         = inline('SCARD'),
 282     set_is_member           = inline('SISMEMBER'),
 283     set_intersection        = inline('SINTER'),
 284     set_intersection_store  = inline('SINTERSTORE'),
 285     set_members             = inline('SMEMBERS'),
 286
 287     -- multiple databases handling commands
 288     select_database  = inline('SELECT'),
 289     move_key         = inline('MOVE'),
 290     flush_database   = inline('FLUSHDB'),
 291     flush_databases  = inline('FLUSHALL'),
 292
 293     -- sorting
 294     --[[
 295         TODO: should we pass sort parameters as a table? e.g:
 296                 params = {
 297                     by    = 'weight_*',
 298                     get   = 'object_*',
 299                     limit = { 0, 10 },
 300                     sort  = { 'desc', 'alpha' }
 301                 }
 302     --]]
 303     sort  = custom('SORT',
 304         function(client, command, params)
 305             -- TODO: here we will put the logic needed to serialize the params
 306             --       table to be sent as the argument of the SORT command.
 307             return request.inline(client, command, params)
 308         end
 309     ),
 310
 311     -- persistence control commands
 312     save             = inline('SAVE'),
 313     background_save  = inline('BGSAVE'),
 314     last_save        = inline('LASTSAVE'),
 315     shutdown         = custom('SHUTDOWN',
 316         function(client, command)
 317             -- let's fire and forget! the connection is closed as soon
 318             -- as the SHUTDOWN command is received by the server.
 319             network.write(command .. protocol.newline)
 320         end
 321     ),
 322
 323     -- remote server control commands
 324     info  = inline('INFO',
 325         function(response)
 326             local info = {}
 327             response:gsub('([^\r\n]*)\r\n', function(kv)
 328                 local k,v = kv:match(('([^:]*):([^:]*)'):rep(1))
 329                 info[k] = v
 330             end)
 331             return info
 332         end
 333     ),
 334 }