Restful key value store in ngx_lua

kv_store.lua

-- Simple KV store in Ngx+Lua

local SHARED_MEM = ngx.shared.mem
local MAX_ARGS = 10
local POST_CHUNK_SIZE = 4096
-- Max size for POST == 1MB
local MAX_POST_SIZE = 1024 * 1024

local DEBUG = true


-- Utils
local function print_table(t)
   table.foreach(t, function(k, v) ngx.say('\t', k, '\t', v) end)
end

local function read_post()
   local req_socket = ngx.req.socket

   local sock, err = req_socket()
   if not sock then
      ngx.say('ERR: Unable to get inbound socket: ', err)
      ngx.exit(500)
   end

   local res = { }
   local total_data = 0
   local data, err, partial

   -- read the incoming data in 4k chunks up to 1MB
   repeat
      data, err, partial = sock:receive(POST_CHUNK_SIZE)
      if not data then
         table.insert(res, partial)
         break
      end
      total_data = total_data + #data
      table.insert(res, data)
   until total_data >= MAX_POST_SIZE

   return res
end


-- Handlers
local function set(key, args)
   local ttl = tonumber(args['ttl']) or 0

   -- test content length
   local content_length = tonumber(ngx.var.http_content_length)
   if not content_length
      or content_length <= 0 or content_length >= MAX_POST_SIZE then

      ngx.say(false)
      return
   end

   local data_chunks = read_post()

   local data = table.concat(data_chunks)

   if DEBUG then
      ngx.say('data chunks: #', #data_chunks)
      ngx.say('data: #', #data)
   end

   local ok = SHARED_MEM:set(key, data, ttl)
   ngx.say(ok)
end

local function get(key, args)
   -- if 'key' isn't stored set value to _EMPTY_
   local value = SHARED_MEM:get(key) or '_EMPTY_'

   ngx.say(value)
end

local function delete(key, args)
   SHARED_MEM:delete(key)
   ngx.say(true)
end


local function default_handler(key, args)
   ngx.say('ERR: ', ngx.req.get_method())
   print_table(args)
end


-- Master table for handlers
local handlers = {
   POST = set,
   GET = get,
   DELETE = delete,
}


-- Entry point
function main()

   -- get request type: GET, POST, etc
   local req = ngx.req.get_method()

   -- get all the args
   local args = ngx.req.get_uri_args(MAX_ARGS)
   local key = args['key']

   if not key then
      ngx.say('"key=" arg is required.')
      return
   end

   local func = handlers[req] or default_handler

   func(key, args)
end

main()
ngx.exit(200)