AuthPlatform/src/util/database.lua

local pgmoon = require('pgmoon')
local ngx = ngx

local _M = {}
local WRITE = 'WRITE'
local READ = 'READ'

local mt = { __index = _M }

--[[
    Get db connection from connection pool
    @return bool, db_context, err
--]]
function _M:get_connection()
    local code = 0
    if ngx.ctx[self.db_type] then
        -- if write before read, make sure write read connection the same
        if ngx.ctx[WRITE] then
            return code, ngx.ctx[WRITE]
        end
        return code, ngx.ctx[self.db_type]
    end
    -- 创建一个新的连接
    local conn = pgmoon.new({
        host = self.host, -- postgres host
        port = self.port,        -- postgres port
        user = self.user,
        password = self.password,      -- postgres password
        database = self.database
    });
    --conn.set_timeout(self.timeout or 1000)
    ---- 连接到数据库
    local ok, err = conn:connect()
    if not ok then
        print("Connection failed: " .. err)
        code = 0x000002
    end
    ngx.log(ngx.INFO, 'Connection success')
    --ngx.say("Connection success")
    ngx.ctx[self.db_type] = conn
    return code,conn
end

--[[
    把连接返回到连接池
    用set_keepalive代替close() 将开启连接池特性,可以为每个nginx工作进程，指定连接最大空闲时间，和连接池最大连接数
    @return void
--]]
function _M.close(self)
    if ngx.ctx[READ] then
        ngx.ctx[READ]:set_keepalive(self.db_pool_timeout, self.db_pool_size)
        ngx.ctx[READ] = nil
    end
    if ngx.ctx[WRITE] then
        ngx.ctx[WRITE]:set_keepalive(self.db_pool_timeout, self.db_pool_size)
        ngx.ctx[WRITE] = nil
    end
end

--[[
    执行数据库语句
    @param sql
    @return bool, data, err
--]]
function _M.db_query(self, sql)
    local code, conn = self:get_connection()
    if code ~= 0 then
        return code,nil
    end
    ngx.log(ngx.INFO, 'begin db query :' .. sql)
    -- 执行查询
    local res, err = conn:query(sql)
    if not res then
        ngx.log(ngx.ERR, 'Query failed:' .. sql)
        return 2,nil
    end
    return 0, res
end

function _M.new(self, opts)
    return setmetatable({
        host = opts.host or '127.0.0.1',
        port = opts.port or 5432,
        user = opts.user or 'postgres',
        password = opts.password or '',
        database = opts.database or 'postgres',
        --charset = opts.charset or 'utf8mb4',
        --timeout = opts.timeout or 1000,
        --max_packet_size = 1024 * 1024,
        db_pool_timeout = opts.pool_timeout or 1000,
        db_pool_size = opts.pool_size or 1000,
        db_type = opts.db_type,
        }, mt)
end

return _M