url shortener base on nginx lua and resty redis

urlshortener.lua

local _M = {
    _VERSION = '0.1'
}

local connect = function()
    local redis = require "resty.redis"
    local red = redis:new()
    red:set_timeout(1000) -- 1 sec
    local ok, err = red:connect("127.0.0.1", 6379)
    if not ok then
        ngx.say(err)
        return ngx.exit(ngx.ERROR)
    end

    return red
end

local basen = function(n)
    local digits = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
    local t = {}
    repeat
        local d = (n % 62) + 1
        n = math.floor(n / 62)
        table.insert(t, 1, digits:sub(d, d))
    until n == 0
    return table.concat(t,"");
end

-- ref https://github.com/AlexChittock/SimpleShortener
-- ref https://gist.github.com/MendelGusmao/2356310
-- "database scheme"
-- database 0: id ~> url
-- database 1: id ~> hits
-- database 2: id ~> [{referer|user_agent}]
-- database 3: id ~> hits (when id is not found)
-- database 4: id ~> [{referer|user_agent}] (when id is not found)
-- database 5: key "shorten.count" storing the number of shortened urls;
--             the id is generated by (this number + 1) converted to base 62
-- database 6: url md5 ~> id

function _M.pack(urlPrefix, checkAlready)
    local short_pre = urlPrefix or ("http://" .. ngx.var.host ..
                         ":" .. ngx.var.server_port .. "/")

    -- See if we have the url already
    local url = ngx.var.arg_url
    local hash = ngx.md5(url)
    local red = connect()

    if checkAlready or true then
        red:select(6)
        local id, err = red:get(hash)
        if id ~= ngx.null then
            ngx.say(short_pre .. id)
            return ngx.exit(ngx.OK)
        end
    end

    red:select(5)
    local count, err = red:incr("shorten.count")
    local id, err = basen(count)

    if checkAlready or true then
        red:select(6)
        red:set(hash, id)
    end

    red:select(0)
    red:set(id, url)

    ngx.say(short_pre .. id)
    return ngx.exit(ngx.OK)
end

function _M.unpack(notFoundRedirect, recordHits, recordReferAndUserAgent)
    local id = string.sub(ngx.var.request_uri, 2) -- remove prefix /

    local red = connect()
    red:select(0)
    local url, err = red:get(id)
    if not url then
        ngx.say(err)
        return ngx.exit(ngx.ERROR)
    end

    local referer = ngx.var.http_referer or ""
    if url == ngx.null then -- not found
        if recordHits then
            red:select(3)
            red:incr(id)
        end

        if recordReferAndUserAgent then
            red:select(4)
            red:rpush(id, referer .. "|" .. ngx.var.http_user_agent)
        end

        if notFoundRedirect then
            return ngx.redirect(notFoundRedirect)
        end

        return ngx.exit(ngx.HTTP_NOT_FOUND)
    end

    -- found
    if recordHits then
        red:select(1)
        red:incr(id)
    end

    if recordReferAndUserAgent then
        red:select(2)
        red:rpush(id, referer .. "|" .. ngx.var.http_user_agent)
    end

    return ngx.redirect(url)
end

return _M