lihao/skynet
1
0
Fork 0
mirror of https://github.com/cloudwu/skynet.git synced 2026-07-25 04:33:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

space to tab

Browse Source
This commit is contained in:
Cloud Wu
2014-07-21 14:27:50 +08:00
parent b933046fb9
commit 39ff6eb2a6
1 changed files with 114 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files

228
lualib/mongo.lua
View File

@@ -1,123 +1,123 @@
local bson = require "bson" local bson = require "bson"
local socket = require "socket" local socket = require "socket"
local socketchannel = require "socketchannel" local socketchannel = require "socketchannel"
local skynet = require "skynet" local skynet = require "skynet"
local driver = require "mongo.driver" local driver = require "mongo.driver"
local md5 = require "md5" local md5 = require "md5"
local rawget = rawget local rawget = rawget
local assert = assert local assert = assert
local bson_encode = bson.encode local bson_encode = bson.encode
local bson_encode_order = bson.encode_order local bson_encode_order = bson.encode_order
local bson_decode = bson.decode local bson_decode = bson.decode
local empty_bson = bson_encode {} local empty_bson = bson_encode {}
local mongo = {} local mongo = {}
mongo.null = assert(bson.null) mongo.null = assert(bson.null)
mongo.maxkey = assert(bson.maxkey) mongo.maxkey = assert(bson.maxkey)
mongo.minkey = assert(bson.minkey) mongo.minkey = assert(bson.minkey)
mongo.type = assert(bson.type) mongo.type = assert(bson.type)
local mongo_cursor = {} local mongo_cursor = {}
local cursor_meta = { local cursor_meta = {
__index = mongo_cursor, __index = mongo_cursor,
} }
local mongo_client = {} local mongo_client = {}
local client_meta = { local client_meta = {
__index = function(self, key) __index = function(self, key)
return rawget(mongo_client, key) or self:getDB(key) return rawget(mongo_client, key) or self:getDB(key)
end, end,
__tostring = function (self) __tostring = function (self)
local port_string local port_string
if self.port then if self.port then
port_string = ":" .. tostring(self.port) port_string = ":" .. tostring(self.port)
else else
port_string = "" port_string = ""
end end
return "[mongo client : " .. self.host .. port_string .."]" return "[mongo client : " .. self.host .. port_string .."]"
end, end,
-- DO NOT need disconnect, because channel will shutdown during gc -- DO NOT need disconnect, because channel will shutdown during gc
} }
local mongo_db = {} local mongo_db = {}
local db_meta = { local db_meta = {
__index = function (self, key) __index = function (self, key)
return rawget(mongo_db, key) or self:getCollection(key) return rawget(mongo_db, key) or self:getCollection(key)
end, end,
__tostring = function (self) __tostring = function (self)
return "[mongo db : " .. self.name .. "]" return "[mongo db : " .. self.name .. "]"
end end
} }
local mongo_collection = {} local mongo_collection = {}
local collection_meta = { local collection_meta = {
__index = function(self, key) __index = function(self, key)
return rawget(mongo_collection, key) or self:getCollection(key) return rawget(mongo_collection, key) or self:getCollection(key)
end , end ,
__tostring = function (self) __tostring = function (self)
return "[mongo collection : " .. self.full_name .. "]" return "[mongo collection : " .. self.full_name .. "]"
end end
} }
local function dispatch_reply(so) local function dispatch_reply(so)
local len_reply = so:read(4) local len_reply = so:read(4)
local reply = so:read(driver.length(len_reply)) local reply = so:read(driver.length(len_reply))
local result = { result = {} } local result = { result = {} }
local succ, reply_id, document, cursor_id, startfrom = driver.reply(reply, result.result) local succ, reply_id, document, cursor_id, startfrom = driver.reply(reply, result.result)
result.document = document result.document = document
result.cursor_id = cursor_id result.cursor_id = cursor_id
result.startfrom = startfrom result.startfrom = startfrom
result.data = reply result.data = reply
return reply_id, succ, result return reply_id, succ, result
end end
local function __parse_addr(addr) local function __parse_addr(addr)
local host, port = string.match(addr, "([^:]+):(.+)") local host, port = string.match(addr, "([^:]+):(.+)")
return host, tonumber(port) return host, tonumber(port)
end end
local function mongo_auth(mongoc) local function mongo_auth(mongoc)
local user = rawget(mongoc, "username") local user = rawget(mongoc, "username")
local pass = rawget(mongoc, "password") local pass = rawget(mongoc, "password")
return function() return function()
if user ~= nil and pass ~= nil then if user ~= nil and pass ~= nil then
assert(mongoc:auth(user, pass)) assert(mongoc:auth(user, pass))
end end
local rs_data = mongoc:runCommand("ismaster") local rs_data = mongoc:runCommand("ismaster")
if rs_data.ok == 1 then if rs_data.ok == 1 then
if rs_data.hosts then if rs_data.hosts then
local backup = {} local backup = {}
for _, v in ipairs(rs_data.hosts) do for _, v in ipairs(rs_data.hosts) do
local host, port = __parse_addr(v) local host, port = __parse_addr(v)
table.insert(backup, {host = host, port = port}) table.insert(backup, {host = host, port = port})
end end
mongoc.__sock:changebackup(backup) mongoc.__sock:changebackup(backup)
end end
if rs_data.ismaster then if rs_data.ismaster then
return return
else else
local host, port = __parse_addr(rs_data.primary) local host, port = __parse_addr(rs_data.primary)
mongoc.host = host mongoc.host = host
mongoc.port = port mongoc.port = port
mongoc.__sock:changehost(host, port) mongoc.__sock:changehost(host, port)
end end
end end
end end
end end
function mongo.client( conf ) function mongo.client( conf )
local first = conf local first = conf
local backup = nil local backup = nil
if conf.rs then if conf.rs then
first = conf.rs[1] first = conf.rs[1]
backup = conf.rs backup = conf.rs
end end
local obj = { local obj = {
host = first.host, host = first.host,
port = first.port or 27017, port = first.port or 27017,
} }
@@ -128,10 +128,10 @@ function mongo.client( conf )
port = obj.port, port = obj.port,
response = dispatch_reply, response = dispatch_reply,
auth = mongo_auth(obj), auth = mongo_auth(obj),
backup = backup, backup = backup,
} }
setmetatable(obj, client_meta) setmetatable(obj, client_meta)
obj.__sock:connect(true) -- try connect only once obj.__sock:connect(true) -- try connect only once
return obj return obj
end end
@@ -139,32 +139,32 @@ function mongo_client:getDB(dbname)
local db = { local db = {
connection = self, connection = self,
name = dbname, name = dbname,
full_name = dbname, full_name = dbname,
database = false, database = false,
__cmd = dbname .. "." .. "$cmd", __cmd = dbname .. "." .. "$cmd",
} }
db.database = db db.database = db
return setmetatable(db, db_meta) return setmetatable(db, db_meta)
end end
function mongo_client:disconnect() function mongo_client:disconnect()
if self.__sock then if self.__sock then
local so = self.__sock local so = self.__sock
self.__sock = false self.__sock = false
so:close() so:close()
end end
end end
function mongo_client:genId() function mongo_client:genId()
local id = self.__id + 1 local id = self.__id + 1
self.__id = id self.__id = id
return id return id
end end
function mongo_client:runCommand(...) function mongo_client:runCommand(...)
if not self.admin then if not self.admin then
self.admin = self:getDB "admin" self.admin = self:getDB "admin"
end end
return self.admin:runCommand(...) return self.admin:runCommand(...)
end end
@@ -176,14 +176,14 @@ function mongo_client:auth(user,password)
return false return false
end end
local key = md5.sumhexa(string.format("%s%s%s",result.nonce,user,password)) local key = md5.sumhexa(string.format("%s%s%s",result.nonce,user,password))
local result= self:runCommand ("authenticate",1,"user",user,"nonce",result.nonce,"key",key) local result= self:runCommand ("authenticate",1,"user",user,"nonce",result.nonce,"key",key)
return result.ok == 1 return result.ok == 1
end end
function mongo_client:logout() function mongo_client:logout()
local result = self:runCommand "logout" local result = self:runCommand "logout"
return result.ok == 1 return result.ok == 1
end end
function mongo_db:runCommand(cmd,cmd_v,...) function mongo_db:runCommand(cmd,cmd_v,...)
@@ -196,18 +196,18 @@ function mongo_db:runCommand(cmd,cmd_v,...)
else else
bson_cmd = bson_encode_order(cmd,cmd_v,...) bson_cmd = bson_encode_order(cmd,cmd_v,...)
end end
local pack = driver.query(request_id, 0, self.__cmd, 0, 1, bson_cmd) local pack = driver.query(request_id, 0, self.__cmd, 0, 1, bson_cmd)
-- we must hold req (req.data), because req.document is a lightuserdata, it's a pointer to the string (req.data) -- we must hold req (req.data), because req.document is a lightuserdata, it's a pointer to the string (req.data)
local req = sock:request(pack, request_id) local req = sock:request(pack, request_id)
local doc = req.document local doc = req.document
return bson_decode(doc) return bson_decode(doc)
end end
function mongo_db:getCollection(collection) function mongo_db:getCollection(collection)
local col = { local col = {
connection = self.connection, connection = self.connection,
name = collection, name = collection,
full_name = self.full_name .. "." .. collection, full_name = self.full_name .. "." .. collection,
database = self.database, database = self.database,
} }
self[collection] = setmetatable(col, collection_meta) self[collection] = setmetatable(col, collection_meta)
@@ -218,20 +218,20 @@ mongo_collection.getCollection = mongo_db.getCollection
function mongo_collection:insert(doc) function mongo_collection:insert(doc)
if doc._id == nil then if doc._id == nil then
doc._id = bson.objectid() doc._id = bson.objectid()
end end
local sock = self.connection.__sock local sock = self.connection.__sock
local pack = driver.insert(0, self.full_name, bson_encode(doc)) local pack = driver.insert(0, self.full_name, bson_encode(doc))
-- flags support 1: ContinueOnError -- flags support 1: ContinueOnError
sock:request(pack) sock:request(pack)
end end
function mongo_collection:batch_insert(docs) function mongo_collection:batch_insert(docs)
for i=1,#docs do for i=1,#docs do
if docs[i]._id == nil then if docs[i]._id == nil then
docs[i]._id = bson.objectid() docs[i]._id = bson.objectid()
end end
docs[i] = bson_encode(docs[i]) docs[i] = bson_encode(docs[i])
end end
local sock = self.connection.__sock local sock = self.connection.__sock
local pack = driver.insert(0, self.full_name, docs) local pack = driver.insert(0, self.full_name, docs)
@@ -239,7 +239,7 @@ function mongo_collection:batch_insert(docs)
end end
function mongo_collection:update(selector,update,upsert,multi) function mongo_collection:update(selector,update,upsert,multi)
local flags = (upsert and 1 or 0) + (multi and 2 or 0) local flags = (upsert and 1 or 0) + (multi and 2 or 0)
local sock = self.connection.__sock local sock = self.connection.__sock
local pack = driver.update(self.full_name, flags, bson_encode(selector), bson_encode(update)) local pack = driver.update(self.full_name, flags, bson_encode(selector), bson_encode(update))
sock:request(pack) sock:request(pack)
@@ -255,24 +255,24 @@ function mongo_collection:findOne(query, selector)
local conn = self.connection local conn = self.connection
local request_id = conn:genId() local request_id = conn:genId()
local sock = conn.__sock local sock = conn.__sock
local pack = driver.query(request_id, 0, self.full_name, 0, 1, query and bson_encode(query) or empty_bson, selector and bson_encode(selector)) local pack = driver.query(request_id, 0, self.full_name, 0, 1, query and bson_encode(query) or empty_bson, selector and bson_encode(selector))
-- we must hold req (req.data), because req.document is a lightuserdata, it's a pointer to the string (req.data) -- we must hold req (req.data), because req.document is a lightuserdata, it's a pointer to the string (req.data)
local req = sock:request(pack, request_id) local req = sock:request(pack, request_id)
local doc = req.document local doc = req.document
return bson_decode(doc) return bson_decode(doc)
end end
function mongo_collection:find(query, selector) function mongo_collection:find(query, selector)
return setmetatable( { return setmetatable( {
__collection = self, __collection = self,
__query = query and bson_encode(query) or empty_bson, __query = query and bson_encode(query) or empty_bson,
__selector = selector and bson_encode(selector), __selector = selector and bson_encode(selector),
__ptr = nil, __ptr = nil,
__data = nil, __data = nil,
__cursor = nil, __cursor = nil,
__document = {}, __document = {},
__flags = 0, __flags = 0,
} , cursor_meta) } , cursor_meta)
end end
function mongo_cursor:hasNext() function mongo_cursor:hasNext()
@@ -285,42 +285,42 @@ function mongo_cursor:hasNext()
local sock = conn.__sock local sock = conn.__sock
local pack local pack
if self.__data == nil then if self.__data == nil then
pack = driver.query(request_id, self.__flags, self.__collection.full_name,0,0,self.__query,self.__selector) pack = driver.query(request_id, self.__flags, self.__collection.full_name,0,0,self.__query,self.__selector)
else else
if self.__cursor then if self.__cursor then
pack = driver.more(request_id, self.__collection.full_name,0,self.__cursor) pack = driver.more(request_id, self.__collection.full_name,0,self.__cursor)
else else
-- no more -- no more
self.__document = nil self.__document = nil
self.__data = nil self.__data = nil
return false return false
end end
end end
local ok, result = pcall(sock.request,sock,pack, request_id) local ok, result = pcall(sock.request,sock,pack, request_id)
local doc = result.document local doc = result.document
local cursor = result.cursor_id local cursor = result.cursor_id
if ok then if ok then
if doc then if doc then
self.__document = result.result self.__document = result.result
self.__data = result.data self.__data = result.data
self.__ptr = 1 self.__ptr = 1
self.__cursor = cursor self.__cursor = cursor
return true return true
else else
self.__document = nil self.__document = nil
self.__data = nil self.__data = nil
self.__cursor = nil self.__cursor = nil
return false return false
end end
else else
self.__document = nil self.__document = nil
self.__data = nil self.__data = nil
self.__cursor = nil self.__cursor = nil
if doc then if doc then
local err = bson_decode(doc) local err = bson_decode(doc)
error(err["$err"]) error(err["$err"])
else else
error("Reply from mongod error") error("Reply from mongod error")
@@ -333,11 +333,11 @@ end
function mongo_cursor:next() function mongo_cursor:next()
if self.__ptr == nil then if self.__ptr == nil then
error "Call hasNext first" error "Call hasNext first"
end end
local r = bson_decode(self.__document[self.__ptr]) local r = bson_decode(self.__document[self.__ptr])
self.__ptr = self.__ptr + 1 self.__ptr = self.__ptr + 1
if self.__ptr > #self.__document then if self.__ptr > #self.__document then
self.__ptr = nil self.__ptr = nil
end end