resty-mongol3的简单封装

resty-mongol3的简单封装

近期项目中需要用到 openresty 来操作 mongodb 数据库，在 github 上面找了给开源的 resty-mongol 发现支持不了 mongo3.0 的认证，后来发现 resty-mongol3 对其改进了一下，能够支持 mongo3.0 认证，因为业务中经常用到 mongo 来进行操作，所以参照 mongo shell 的方式对其进行了一下简单的封装，方便以后使用。（因为业务需要，有的方法不要被支持，固没有封装）如有 bug ，望大家批评指正。

local mongo = require("resty.mongol")
local object_id = require("resty.mongol.object_id")
local _M = {
    _VERSION = "0.0.1"
}

local metatable = { __index = _M }

-- 创建objectId
_M.objectId = function(str)
    local buf = (str:gsub('..', function (cc)
        return string.char(tonumber(cc, 16))
        end))
    return object_id.new(buf)
end

--[[
    @desc
        Creates a MongoClient instance. 
    @params
        opts            @type     table
    @return
        table             @type     table     A MongoClient instance
 ]]
function _M.new(self, opts)
    opts = opts or {}
    local timeout     = opts.timeout or 3000
    local host         = opts.host or "localhost"
    local port         = opts.port or 27017
    local passwd    = opts.passwd or ""
    local user    =  opts.user or ""
    local database  = opts.database or "admin"
    local keepalive = (opts.keepalive and opts.keepalive * 1000) or 60000
    local poolSize  = opts.poolSize or 1000
  
    return setmetatable({
            timeout      = timeout,
            host          = host,
            port         = tostring(port),
            user          = user,
            passwd         = passwd,
            database     = database,
            keepalive     = keepalive,
            poolSize     = poolSize,
            _db            = database,
            _user        = user,
            _passwd     = passwd,
            _sort       = {},
            _limit      = 100,
            _skip       = 0,
            }, metatable)
end

--[[
    @desc
        get mongodb's connection objects. 
 ]]
local function getMgoConn(mgoConfig)
    -- 获取连接对象
    local mgoConn = mongo:new()
    if not mgoConn then
        return nil, "get mongo connection occur error"
    end
    
    -- 设置链接超时
    mgoConn:set_timeout(mgoConfig.timeout)
    
    --获取连接客户端
    local ok, err =mgoConn:connect(mgoConfig.host, mgoConfig.port)
    if not ok then 
        return nil, err
    end 
    return mgoConn, nil
end

--[[
    @desc
        pack connection commands. 
 ]]
local function packConnCmd(self, mgoConn, cmd, ... )    
    local result, err = mgoConn[cmd](mgoConn, ... )

    mgoConn:set_keepalive(self.keepalive, self.poolSize)
    
    return result, err
end

--[[
    @desc
        this is a map of mongol.conn's command. 
 ]]
local connCmd = {
    isMaster         = "ismaster",
    getPrimary         = "getprimary",
    getReusedTime     = "get_reused_times",
    dbs             = "databases",           
}
for k,v in pairs(connCmd) do
    
    _M[k] =
            function (self, ...)
                   --获取连接客户端
                local mgoConn, err = getMgoConn(self)
                if not mgoConn then 
                    return nil, err
                end 
                return packConnCmd(self, mgoConn, v, ...)
            end
end


--[[
    @desc
        switch db by dbName and auth your id 
    @params

    @return
        Returns a database object, or nil.
 ]]
function _M.useDatabase(self, dbName, user, passwd)
       --获取连接客户端
       self._db = dbName
       self._user = user or self._user
       self._passwd = passwd or self._psswd
       return string.format("%s %s","current database is", dbName)
end

function _M.ping( self )
       --获取连接客户端
    local mgoConn, err = getMgoConn(self)
    if not mgoConn then 
        return nil, err
    end    
        
    local db = mgoConn:new_db_handle(self._db)
    if not db then
        return nil, "get database occur error"
    end

    --用户授权
    local count, err = mgoConn:get_reused_times()

    if (count == 0) or err then
        if self._user and self._passwd then
            local ok, err = db:auth_scram_sha1(self._user, self._passwd)
            if not ok then 
                return nil, err
            end 
        end
    end
    return "ok", nil
end

--[[
    @desc
        switch db by self.dbName and auth your id 
    @params
        dbName            @type     table     @default    self.database
        user             @type     string     @default    self.user
        passwd             @type     string  @default     self.passwd
    @return
        Returns a database object, or nil.
 ]]
local function getDB(self)
       --获取连接客户端
    local mgoConn, err = getMgoConn(self)
    if not mgoConn then 
        return nil, nil, err
    end    

    local db = mgoConn:new_db_handle(self._db)
    if not db then
        return nil, nil, "get database occur error"
    end

    --用户授权
    local count, err = mgoConn:get_reused_times()

    if (count == 0) or err then
        if self._user and self._passwd then
            local ok, err = db:auth_scram_sha1(self._user, self._passwd)
            if not ok then 
                return nil, nil, err
            end 
        end
    end
    return db, mgoConn, nil
end

local dbCmd = {
    addUser         = "add_user",
    -- getColl         = "get_col",
    -- getGrid         = "get_grid", 
    dropDatabase     = "dropDatabase", 
}

--[[
    @desc
        pack database commands. 
 ]]
local function packDBCmd(self, db, mgoConn, cmd, ... )    
    local result, err = db[cmd](db, ... )

    mgoConn:set_keepalive(self.keepalive, self.poolSize)
    
    return result, err
end

for k,v in pairs(dbCmd) do
    
    _M[k] =
            function (self, ...)
                   --获取连接客户端
                local db, mgoConn, err = getDB(self)
                if not db then
                    return nil, "get current database occur error " .. err
                end
                return packDBCmd(self, db, mgoConn, v, ...)
            end
end

function _M.list( self )
       --获取连接客户端
    local db, mgoConn, err = getDB(self)
    if not db then
        return nil, "get current database occur error " .. err
    end
    -- return packDBCmd(self, db, mgoConn, v, ...)
    local cursor, err = db:listcollections()
    if not cursor then
        return nil, string.format("%s %s %s", "system.namespaces", "find occur error", err)
    end
    local results = {}
    for index, item in cursor:pairs() do
        table.insert(results, item)
    end
    
    mgoConn:set_keepalive(self.keepalive, self.poolSize)
    
    return result, err
end

function _M.getCollection( self, collName )
    self.collName = collName
    return self
end

local collCmd = {
    count     = "count",
    drop     = "drop",
    update     = "update",    
    insert  = "insert",  
    delete     = "delete", 
}

--[[
    @desc
        pack collection's commands. 
 ]]
local function packCollCmd(self, db, mgoConn, cmd, ... )    
    --获取集合
    local coll = db:get_col(self.collName)
    if not coll then
        return nil, "get collection occur error"
    end

    local result, err = coll[cmd](coll, ... )

    mgoConn:set_keepalive(self.keepalive, self.poolSize)
    
    return result, err
end

for k,v in pairs(collCmd) do
    
    _M[k] =
            function (self, ...)
                   --获取连接客户端
                local db, mgoConn, err = getDB(self)
                if not db then
                    return nil, "get current database occur error " .. err
                end
                return packCollCmd(self, db, mgoConn, v, ...)
            end
end

function _M.sort( self, fields )
    self._sort = fields
    return self
end

function _M.limit( self, num )
    self._limit = num 
    return self
end

function _M.skip( self, num )
    self._skip = num
    return self
end

--[[
find 方法要在用户每次操作后直接帮他关闭链接，所以没有返回游标给用户，而是自己内部遍历了下，直接返回结果
]] 
function _M.find( self, ... )
       --获取连接客户端
    local db, mgoConn, err = getDB(self)
    if not db then
        return nil, "get current database occur error " .. err
    end
    --获取集合
    local coll, err = db:get_col(self.collName)
    
    if not coll then
        return nil, "get collection occur error"
    end

    local cursor, err = coll["find"](coll, ... )
    if not cursor then
        return nil, string.format("%s %s %s", self.collName, "find occur error", err)
    end

    cursor:limit(self._limit)
    cursor:skip(self._skip)
    cursor = cursor:sort(self._sort)
    local results = {}
    for index, item in cursor:pairs() do
        table.insert(results, item)
    end
    mgoConn:set_keepalive(self.keepalive, self.poolSize)
    return results, err    
end

return _M