mirror of
https://github.com/cloudwu/skynet.git
synced 2026-07-24 20:23:06 +00:00
mongodb 交互协议调整为 OP_MSG (#1649)
* mongodb 交互协议调整为 OP_MSG * Update Co-authored-by: yy151474 <yy151474@alibaba-inc.com>
This commit is contained in:
@@ -10,21 +10,17 @@
|
||||
#include <stdio.h>
|
||||
#include <string.h>
|
||||
|
||||
#define OP_REPLY 1
|
||||
#define OP_MSG 1000
|
||||
#define OP_UPDATE 2001
|
||||
#define OP_INSERT 2002
|
||||
#define OP_QUERY 2004
|
||||
#define OP_GET_MORE 2005
|
||||
#define OP_DELETE 2006
|
||||
#define OP_KILL_CURSORS 2007
|
||||
#define OP_COMPRESSED 2012
|
||||
#define OP_MSG 2013
|
||||
|
||||
typedef enum {
|
||||
MSG_CHECKSUM_PRESENT = 1 << 0,
|
||||
MSG_MORE_TO_COME = 1 << 1,
|
||||
MSG_EXHAUST_ALLOWED = 1 << 16,
|
||||
} msg_flags_t;
|
||||
|
||||
#define REPLY_CURSORNOTFOUND 1
|
||||
#define REPLY_QUERYFAILURE 2
|
||||
#define REPLY_AWAITCAPABLE 8 // ignore because mongo 1.6+ always set it
|
||||
|
||||
#define DEFAULT_CAP 128
|
||||
|
||||
struct connection {
|
||||
int sock;
|
||||
int id;
|
||||
@@ -106,6 +102,13 @@ write_int32(struct buffer *b, int32_t v) {
|
||||
b->ptr[b->size++] = (uv >> 24)&0xff;
|
||||
}
|
||||
|
||||
static inline void
|
||||
write_int8(struct buffer *b, int8_t v) {
|
||||
uint8_t uv = (uint8_t)v;
|
||||
buffer_reserve(b, 1);
|
||||
b->ptr[b->size++] = uv;
|
||||
}
|
||||
|
||||
static inline void
|
||||
write_bytes(struct buffer *b, const void * buf, int sz) {
|
||||
buffer_reserve(b,sz);
|
||||
@@ -113,6 +116,7 @@ write_bytes(struct buffer *b, const void * buf, int sz) {
|
||||
b->size += sz;
|
||||
}
|
||||
|
||||
/*
|
||||
static void
|
||||
write_string(struct buffer *b, const char *key, size_t sz) {
|
||||
buffer_reserve(b,sz+1);
|
||||
@@ -120,6 +124,7 @@ write_string(struct buffer *b, const char *key, size_t sz) {
|
||||
b->ptr[b->size+sz] = '\0';
|
||||
b->size+=sz+1;
|
||||
}
|
||||
*/
|
||||
|
||||
static inline int
|
||||
reserve_length(struct buffer *b) {
|
||||
@@ -138,389 +143,69 @@ write_length(struct buffer *b, int32_t v, int off) {
|
||||
b->ptr[off++] = (uv >> 24)&0xff;
|
||||
}
|
||||
|
||||
// 1 integer id
|
||||
// 2 integer flags
|
||||
// 3 string collection name
|
||||
// 4 integer skip
|
||||
// 5 integer return number
|
||||
// 6 document query
|
||||
// 7 document selector (optional)
|
||||
// return string package
|
||||
static int
|
||||
op_query(lua_State *L) {
|
||||
int id = luaL_checkinteger(L,1);
|
||||
document query = lua_touserdata(L,6);
|
||||
if (query == NULL) {
|
||||
return luaL_error(L, "require query document");
|
||||
}
|
||||
document selector = lua_touserdata(L,7);
|
||||
int flags = luaL_checkinteger(L, 2);
|
||||
size_t nsz = 0;
|
||||
const char *name = luaL_checklstring(L,3,&nsz);
|
||||
int skip = luaL_checkinteger(L, 4);
|
||||
int number = luaL_checkinteger(L, 5);
|
||||
struct header_t {
|
||||
//int32_t message_length; // total message size, include this
|
||||
int32_t request_id; // identifier for this message
|
||||
int32_t response_to; // requestID from the original request(used in responses from the database)
|
||||
int32_t opcode; // message type
|
||||
|
||||
luaL_Buffer b;
|
||||
luaL_buffinit(L,&b);
|
||||
|
||||
struct buffer buf;
|
||||
buffer_create(&buf);
|
||||
int len = reserve_length(&buf);
|
||||
write_int32(&buf, id);
|
||||
write_int32(&buf, 0);
|
||||
write_int32(&buf, OP_QUERY);
|
||||
write_int32(&buf, flags);
|
||||
write_string(&buf, name, nsz);
|
||||
write_int32(&buf, skip);
|
||||
write_int32(&buf, number);
|
||||
|
||||
int32_t query_len = get_length(query);
|
||||
int total = buf.size + query_len;
|
||||
int32_t selector_len = 0;
|
||||
if (selector) {
|
||||
selector_len = get_length(selector);
|
||||
total += selector_len;
|
||||
}
|
||||
|
||||
write_length(&buf, total, len);
|
||||
luaL_addlstring(&b, (const char *)buf.ptr, buf.size);
|
||||
buffer_destroy(&buf);
|
||||
|
||||
luaL_addlstring(&b, (const char *)query, query_len);
|
||||
|
||||
if (selector) {
|
||||
luaL_addlstring(&b, (const char *)selector, selector_len);
|
||||
}
|
||||
|
||||
luaL_pushresult(&b);
|
||||
|
||||
return 1;
|
||||
}
|
||||
int32_t flags;
|
||||
};
|
||||
|
||||
// 1 string data
|
||||
// 2 result document table
|
||||
// return boolean succ (false -> request id, error document)
|
||||
// number request_id
|
||||
// document first
|
||||
// string cursor_id
|
||||
// integer startfrom
|
||||
static int
|
||||
op_reply(lua_State *L) {
|
||||
unpack_reply(lua_State *L) {
|
||||
size_t data_len = 0;
|
||||
const char * data = luaL_checklstring(L,1,&data_len);
|
||||
struct reply_type {
|
||||
// int32_t length; // total message size, including this
|
||||
int32_t request_id; // identifier for this message
|
||||
int32_t response_id; // requestID from the original request
|
||||
// (used in reponses from db)
|
||||
int32_t opcode; // request type
|
||||
int32_t flags;
|
||||
int32_t cursor_id[2];
|
||||
int32_t starting;
|
||||
int32_t number;
|
||||
};
|
||||
const struct reply_type* reply = (const struct reply_type*)data;
|
||||
const struct header_t* h = (const struct header_t*)data;
|
||||
|
||||
if (data_len < sizeof(*reply)) {
|
||||
if (data_len < sizeof(*h)) {
|
||||
lua_pushboolean(L, 0);
|
||||
return 1;
|
||||
}
|
||||
|
||||
int id = little_endian(reply->response_id);
|
||||
int flags = little_endian(reply->flags);
|
||||
if (flags & REPLY_QUERYFAILURE) {
|
||||
lua_pushboolean(L,0);
|
||||
lua_pushinteger(L, id);
|
||||
lua_pushlightuserdata(L, (void *)(reply+1));
|
||||
return 3;
|
||||
int opcode = little_endian(h->opcode);
|
||||
if (opcode != OP_MSG) {
|
||||
return luaL_error(L, "Unsupported opcode:%d", opcode);
|
||||
}
|
||||
|
||||
int starting_from = little_endian(reply->starting);
|
||||
int number = little_endian(reply->number);
|
||||
int sz = (int)data_len - sizeof(*reply);
|
||||
const uint8_t * doc = (const uint8_t *)(reply+1);
|
||||
int id = little_endian(h->response_to);
|
||||
int flags = little_endian(h->flags);
|
||||
|
||||
if (lua_istable(L,2)) {
|
||||
int i = 1;
|
||||
while (sz > 4) {
|
||||
lua_pushlightuserdata(L, (void *)doc);
|
||||
lua_rawseti(L, 2, i);
|
||||
|
||||
int32_t doc_len = get_length((document)doc);
|
||||
if (doc_len <= 0) {
|
||||
return luaL_error(L, "Invalid result bson document");
|
||||
if (flags != 0) {
|
||||
if ((flags & MSG_CHECKSUM_PRESENT) != 0) {
|
||||
return luaL_error(L, "Unsupported OP_MSG flag checksumPresent");
|
||||
}
|
||||
|
||||
doc += doc_len;
|
||||
sz -= doc_len;
|
||||
if ((flags ^ MSG_MORE_TO_COME) != 0) {
|
||||
return luaL_error(L, "Unsupported OP_MSG flag:%d", flags);
|
||||
}
|
||||
}
|
||||
|
||||
++i;
|
||||
int sz = (int)data_len - sizeof(*h);
|
||||
|
||||
const uint8_t * section = (const uint8_t *)(h+1);
|
||||
|
||||
uint8_t payload_type = *section;
|
||||
const uint8_t * doc = section+1;
|
||||
|
||||
if (payload_type != 0) {
|
||||
return luaL_error(L, "Unsupported OP_MSG payload type: %d", payload_type);
|
||||
}
|
||||
if (i != number + 1) {
|
||||
lua_pushboolean(L,0);
|
||||
lua_pushinteger(L, id);
|
||||
return 2;
|
||||
}
|
||||
int c = lua_rawlen(L, 2);
|
||||
for (;i<=c;i++) {
|
||||
lua_pushnil(L);
|
||||
lua_rawseti(L, 2, i);
|
||||
}
|
||||
} else {
|
||||
if (sz >= 4) {
|
||||
sz -= get_length((document)doc);
|
||||
}
|
||||
}
|
||||
if (sz != 0) {
|
||||
return luaL_error(L, "Invalid result bson document");
|
||||
|
||||
int32_t doc_sz = get_length((document)(doc));
|
||||
if ((sz - 1) != doc_sz) {
|
||||
return luaL_error(L, "Unsupported OP_MSG reply: >1 section");
|
||||
}
|
||||
|
||||
lua_pushboolean(L, 1);
|
||||
lua_pushinteger(L, id);
|
||||
if (number == 0)
|
||||
lua_pushnil(L);
|
||||
else
|
||||
lua_pushlightuserdata(L, (void *)(reply+1));
|
||||
if (reply->cursor_id[0] == 0 && reply->cursor_id[1]==0) {
|
||||
// closed cursor
|
||||
lua_pushnil(L);
|
||||
} else {
|
||||
lua_pushlstring(L, (const char *)(reply->cursor_id), 8);
|
||||
}
|
||||
lua_pushinteger(L, starting_from);
|
||||
|
||||
return 5;
|
||||
}
|
||||
|
||||
/*
|
||||
1 string cursor_id
|
||||
return string package
|
||||
*/
|
||||
static int
|
||||
op_kill(lua_State *L) {
|
||||
size_t cursor_len = 0;
|
||||
const char * cursor_id = luaL_tolstring(L, 1, &cursor_len);
|
||||
if (cursor_len != 8) {
|
||||
return luaL_error(L, "Invalid cursor id");
|
||||
}
|
||||
|
||||
struct buffer buf;
|
||||
buffer_create(&buf);
|
||||
|
||||
int len = reserve_length(&buf);
|
||||
write_int32(&buf, 0);
|
||||
write_int32(&buf, 0);
|
||||
write_int32(&buf, OP_KILL_CURSORS);
|
||||
|
||||
write_int32(&buf, 0);
|
||||
write_int32(&buf, 1);
|
||||
write_bytes(&buf, cursor_id, 8);
|
||||
|
||||
write_length(&buf, buf.size, len);
|
||||
|
||||
lua_pushlstring(L, (const char *)buf.ptr, buf.size);
|
||||
buffer_destroy(&buf);
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
/*
|
||||
1 string collection
|
||||
2 integer single remove
|
||||
3 document selector
|
||||
|
||||
return string package
|
||||
*/
|
||||
static int
|
||||
op_delete(lua_State *L) {
|
||||
document selector = lua_touserdata(L,3);
|
||||
if (selector == NULL) {
|
||||
luaL_error(L, "Invalid param");
|
||||
}
|
||||
size_t sz = 0;
|
||||
const char * name = luaL_checklstring(L,1,&sz);
|
||||
|
||||
luaL_Buffer b;
|
||||
luaL_buffinit(L,&b);
|
||||
|
||||
struct buffer buf;
|
||||
buffer_create(&buf);
|
||||
int len = reserve_length(&buf);
|
||||
write_int32(&buf, 0);
|
||||
write_int32(&buf, 0);
|
||||
write_int32(&buf, OP_DELETE);
|
||||
write_int32(&buf, 0);
|
||||
write_string(&buf, name, sz);
|
||||
write_int32(&buf, lua_tointeger(L,2));
|
||||
|
||||
int32_t selector_len = get_length(selector);
|
||||
int total = buf.size + selector_len;
|
||||
write_length(&buf, total, len);
|
||||
|
||||
luaL_addlstring(&b, (const char *)buf.ptr, buf.size);
|
||||
buffer_destroy(&buf);
|
||||
|
||||
luaL_addlstring(&b, (const char *)selector, selector_len);
|
||||
luaL_pushresult(&b);
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
/*
|
||||
1 integer id
|
||||
2 string collection
|
||||
3 integer number
|
||||
4 cursor_id (8 bytes string/ 64bit)
|
||||
|
||||
return string package
|
||||
*/
|
||||
static int
|
||||
op_get_more(lua_State *L) {
|
||||
int id = luaL_checkinteger(L, 1);
|
||||
size_t sz = 0;
|
||||
const char * name = luaL_checklstring(L,2,&sz);
|
||||
int number = luaL_checkinteger(L, 3);
|
||||
size_t cursor_len = 0;
|
||||
const char * cursor_id = luaL_tolstring(L, 4, &cursor_len);
|
||||
if (cursor_len != 8) {
|
||||
return luaL_error(L, "Invalid cursor id");
|
||||
}
|
||||
|
||||
struct buffer buf;
|
||||
buffer_create(&buf);
|
||||
int len = reserve_length(&buf);
|
||||
write_int32(&buf, id);
|
||||
write_int32(&buf, 0);
|
||||
write_int32(&buf, OP_GET_MORE);
|
||||
write_int32(&buf, 0);
|
||||
write_string(&buf, name, sz);
|
||||
write_int32(&buf, number);
|
||||
write_bytes(&buf, cursor_id, 8);
|
||||
write_length(&buf, buf.size, len);
|
||||
|
||||
lua_pushlstring(L, (const char *)buf.ptr, buf.size);
|
||||
buffer_destroy(&buf);
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
// 1 string collection
|
||||
// 2 integer flags
|
||||
// 3 document selector
|
||||
// 4 document update
|
||||
// return string package
|
||||
static int
|
||||
op_update(lua_State *L) {
|
||||
document selector = lua_touserdata(L,3);
|
||||
document update = lua_touserdata(L,4);
|
||||
if (selector == NULL || update == NULL) {
|
||||
luaL_error(L, "Invalid param");
|
||||
}
|
||||
size_t sz = 0;
|
||||
const char * name = luaL_checklstring(L,1,&sz);
|
||||
|
||||
luaL_Buffer b;
|
||||
luaL_buffinit(L, &b);
|
||||
|
||||
struct buffer buf;
|
||||
buffer_create(&buf);
|
||||
// make package header, don't raise L error
|
||||
int len = reserve_length(&buf);
|
||||
write_int32(&buf, 0);
|
||||
write_int32(&buf, 0);
|
||||
write_int32(&buf, OP_UPDATE);
|
||||
write_int32(&buf, 0);
|
||||
write_string(&buf, name, sz);
|
||||
write_int32(&buf, lua_tointeger(L,2));
|
||||
|
||||
int32_t selector_len = get_length(selector);
|
||||
int32_t update_len = get_length(update);
|
||||
|
||||
int total = buf.size + selector_len + update_len;
|
||||
write_length(&buf, total, len);
|
||||
|
||||
luaL_addlstring(&b, (const char *)buf.ptr, buf.size);
|
||||
buffer_destroy(&buf);
|
||||
|
||||
luaL_addlstring(&b, (const char *)selector, selector_len);
|
||||
luaL_addlstring(&b, (const char *)update, update_len);
|
||||
|
||||
luaL_pushresult(&b);
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
static int
|
||||
document_length(lua_State *L) {
|
||||
if (lua_isuserdata(L, 3)) {
|
||||
document doc = lua_touserdata(L,3);
|
||||
return get_length(doc);
|
||||
}
|
||||
if (lua_istable(L,3)) {
|
||||
int total = 0;
|
||||
int s = lua_rawlen(L,3);
|
||||
int i;
|
||||
for (i=1;i<=s;i++) {
|
||||
lua_rawgeti(L, 3, i);
|
||||
document doc = lua_touserdata(L,-1);
|
||||
if (doc == NULL) {
|
||||
lua_pop(L,1);
|
||||
return luaL_error(L, "Invalid document at %d", i);
|
||||
} else {
|
||||
total += get_length(doc);
|
||||
lua_pop(L,1);
|
||||
}
|
||||
}
|
||||
return total;
|
||||
}
|
||||
return luaL_error(L, "Insert need documents");
|
||||
}
|
||||
|
||||
// 1 integer flags
|
||||
// 2 string collection
|
||||
// 3 documents
|
||||
// return string package
|
||||
static int
|
||||
op_insert(lua_State *L) {
|
||||
size_t sz = 0;
|
||||
const char * name = luaL_checklstring(L,2,&sz);
|
||||
int dsz = document_length(L);
|
||||
|
||||
luaL_Buffer b;
|
||||
luaL_buffinit(L, &b);
|
||||
|
||||
struct buffer buf;
|
||||
buffer_create(&buf);
|
||||
// make package header, don't raise L error
|
||||
int len = reserve_length(&buf);
|
||||
write_int32(&buf, 0);
|
||||
write_int32(&buf, 0);
|
||||
write_int32(&buf, OP_INSERT);
|
||||
write_int32(&buf, lua_tointeger(L,1));
|
||||
write_string(&buf, name, sz);
|
||||
|
||||
int total = buf.size + dsz;
|
||||
write_length(&buf, total, len);
|
||||
|
||||
luaL_addlstring(&b, (const char *)buf.ptr, buf.size);
|
||||
buffer_destroy(&buf);
|
||||
|
||||
if (lua_isuserdata(L,3)) {
|
||||
document doc = lua_touserdata(L,3);
|
||||
luaL_addlstring(&b, (const char *)doc, get_length(doc));
|
||||
} else {
|
||||
int s = lua_rawlen(L, 3);
|
||||
int i;
|
||||
for (i=1;i<=s;i++) {
|
||||
lua_rawgeti(L,3,i);
|
||||
document doc = lua_touserdata(L,-1);
|
||||
lua_pop(L,1); // must call lua_pop before luaL_addlstring, because addlstring may change stack top
|
||||
luaL_addlstring(&b, (const char *)doc, get_length(doc));
|
||||
}
|
||||
}
|
||||
|
||||
luaL_pushresult(&b);
|
||||
|
||||
return 1;
|
||||
lua_pushlightuserdata(L, (void *)(doc));
|
||||
return 3;
|
||||
}
|
||||
|
||||
// string 4 bytes length
|
||||
@@ -535,18 +220,52 @@ reply_length(lua_State *L) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
// @param 1 request_id int
|
||||
// @param 2 flags int
|
||||
// @param 3 command bson document
|
||||
// @return
|
||||
static int
|
||||
op_msg(lua_State *L) {
|
||||
int id = luaL_checkinteger(L, 1);
|
||||
int flags = luaL_checkinteger(L, 2);
|
||||
document cmd = lua_touserdata(L, 3);
|
||||
|
||||
if (cmd == NULL) {
|
||||
return luaL_error(L, "opmsg require cmd document");
|
||||
}
|
||||
|
||||
luaL_Buffer b;
|
||||
luaL_buffinit(L, &b);
|
||||
|
||||
struct buffer buf;
|
||||
buffer_create(&buf);
|
||||
int len = reserve_length(&buf);
|
||||
write_int32(&buf, id);
|
||||
write_int32(&buf, 0);
|
||||
write_int32(&buf, OP_MSG);
|
||||
write_int32(&buf, flags);
|
||||
write_int8(&buf, 0);
|
||||
|
||||
int32_t cmd_len = get_length(cmd);
|
||||
int total = buf.size + cmd_len;
|
||||
|
||||
write_length(&buf, total, len);
|
||||
luaL_addlstring(&b, (const char *)buf.ptr, buf.size);
|
||||
buffer_destroy(&buf);
|
||||
|
||||
luaL_addlstring(&b, (const char *)cmd, cmd_len);
|
||||
luaL_pushresult(&b);
|
||||
return 1;
|
||||
}
|
||||
|
||||
|
||||
LUAMOD_API int
|
||||
luaopen_skynet_mongo_driver(lua_State *L) {
|
||||
luaL_checkversion(L);
|
||||
luaL_Reg l[] ={
|
||||
{ "query", op_query },
|
||||
{ "reply", op_reply },
|
||||
{ "kill", op_kill },
|
||||
{ "delete", op_delete },
|
||||
{ "more", op_get_more },
|
||||
{ "update", op_update },
|
||||
{ "insert", op_insert },
|
||||
{ "reply", unpack_reply }, // 接收响应
|
||||
{ "length", reply_length },
|
||||
{ "op_msg", op_msg},
|
||||
{ NULL, NULL },
|
||||
};
|
||||
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
local bson = require "bson"
|
||||
local socket = require "skynet.socket"
|
||||
|
||||
require "skynet.socket"
|
||||
|
||||
local socketchannel = require "skynet.socketchannel"
|
||||
local skynet = require "skynet"
|
||||
local driver = require "skynet.mongo.driver"
|
||||
@@ -68,11 +70,9 @@ local collection_meta = {
|
||||
local function dispatch_reply(so)
|
||||
local len_reply = so:read(4)
|
||||
local reply = so:read(driver.length(len_reply))
|
||||
local result = { result = {} }
|
||||
local succ, reply_id, document, cursor_id, startfrom = driver.reply(reply, result.result)
|
||||
local result = {}
|
||||
local succ, reply_id, document = driver.reply(reply)
|
||||
result.document = document
|
||||
result.cursor_id = cursor_id
|
||||
result.startfrom = startfrom
|
||||
result.data = reply
|
||||
return reply_id, succ, result
|
||||
end
|
||||
@@ -298,17 +298,37 @@ function mongo_db:runCommand(cmd,cmd_v,...)
|
||||
local sock = conn.__sock
|
||||
local bson_cmd
|
||||
if not cmd_v then
|
||||
bson_cmd = bson_encode_order(cmd,1)
|
||||
-- ensure cmd remains in first place
|
||||
bson_cmd = bson_encode_order(cmd,1, "$db", self.name)
|
||||
else
|
||||
bson_cmd = bson_encode_order(cmd,cmd_v,...)
|
||||
bson_cmd = bson_encode_order(cmd,cmd_v, "$db", self.name, ...)
|
||||
end
|
||||
local pack = driver.query(request_id, 0, self.__cmd, 0, 1, bson_cmd)
|
||||
|
||||
local pack = driver.op_msg(request_id, 0, bson_cmd)
|
||||
-- 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 doc = req.document
|
||||
return bson_decode(doc)
|
||||
end
|
||||
|
||||
--- send command without response
|
||||
function mongo_db:send_command(cmd, cmd_v, ...)
|
||||
local conn = self.connection
|
||||
local request_id = conn:genId()
|
||||
local sock = conn.__sock
|
||||
local bson_cmd
|
||||
if not cmd_v then
|
||||
-- ensure cmd remains in first place
|
||||
bson_cmd = bson_encode_order(cmd, 1, "$db", self.name, "writeConcern", {w=0})
|
||||
else
|
||||
bson_cmd = bson_encode_order(cmd, cmd_v, "$db", self.name, "writeConcern", {w=0}, ...)
|
||||
end
|
||||
|
||||
local pack = driver.op_msg(request_id, 2, bson_cmd)
|
||||
sock:request(pack)
|
||||
return {ok=1} -- fake successful response
|
||||
end
|
||||
|
||||
function mongo_db:getCollection(collection)
|
||||
local col = {
|
||||
connection = self.connection,
|
||||
@@ -326,10 +346,7 @@ function mongo_collection:insert(doc)
|
||||
if doc._id == nil then
|
||||
doc._id = bson.objectid()
|
||||
end
|
||||
local sock = self.connection.__sock
|
||||
local pack = driver.insert(0, self.full_name, bson_encode(doc))
|
||||
-- flags support 1: ContinueOnError
|
||||
sock:request(pack)
|
||||
self.database:send_command("insert", self.name, "documents", {bson_encode(doc)})
|
||||
end
|
||||
|
||||
local function werror(r)
|
||||
@@ -353,6 +370,11 @@ function mongo_collection:safe_insert(doc)
|
||||
return werror(r)
|
||||
end
|
||||
|
||||
function mongo_collection:raw_safe_insert(doc)
|
||||
local r = self.database:runCommand("insert", self.name, "documents", {doc})
|
||||
return werror(r)
|
||||
end
|
||||
|
||||
function mongo_collection:batch_insert(docs)
|
||||
for i=1,#docs do
|
||||
if docs[i]._id == nil then
|
||||
@@ -360,11 +382,12 @@ function mongo_collection:batch_insert(docs)
|
||||
end
|
||||
docs[i] = bson_encode(docs[i])
|
||||
end
|
||||
local sock = self.connection.__sock
|
||||
local pack = driver.insert(0, self.full_name, docs)
|
||||
sock:request(pack)
|
||||
|
||||
self.database:send_command("insert", self.name, "documents", docs)
|
||||
end
|
||||
|
||||
mongo_collection.insert_many = mongo_collection.batch_insert
|
||||
|
||||
function mongo_collection:safe_batch_insert(docs)
|
||||
for i = 1, #docs do
|
||||
if docs[i]._id == nil then
|
||||
@@ -377,16 +400,20 @@ function mongo_collection:safe_batch_insert(docs)
|
||||
return werror(r)
|
||||
end
|
||||
|
||||
function mongo_collection:update(selector,update,upsert,multi)
|
||||
local flags = (upsert and 1 or 0) + (multi and 2 or 0)
|
||||
local sock = self.connection.__sock
|
||||
local pack = driver.update(self.full_name, flags, bson_encode(selector), bson_encode(update))
|
||||
sock:request(pack)
|
||||
mongo_collection.safe_insert_many = mongo_collection.safe_batch_insert
|
||||
|
||||
function mongo_collection:update(query,update,upsert,multi)
|
||||
self.database:send_command("update", self.name, "updates", {bson_encode({
|
||||
q = query,
|
||||
u = update,
|
||||
upsert = upsert,
|
||||
multi = multi
|
||||
})})
|
||||
end
|
||||
|
||||
function mongo_collection:safe_update(selector, update, upsert, multi)
|
||||
function mongo_collection:safe_update(query, update, upsert, multi)
|
||||
local r = self.database:runCommand("update", self.name, "updates", {bson_encode({
|
||||
q = selector,
|
||||
q = query,
|
||||
u = update,
|
||||
upsert = upsert,
|
||||
multi = multi,
|
||||
@@ -394,6 +421,20 @@ function mongo_collection:safe_update(selector, update, upsert, multi)
|
||||
return werror(r)
|
||||
end
|
||||
|
||||
function mongo_collection:batch_update(updates)
|
||||
local updates_tb = {}
|
||||
for i = 1, #updates do
|
||||
updates_tb[i] = bson_encode({
|
||||
q = updates[i].query,
|
||||
u = updates[i].update,
|
||||
upsert = updates[i].upsert,
|
||||
multi = updates[i].multi,
|
||||
})
|
||||
end
|
||||
|
||||
self.database:send_command("update", self.name, "updates", updates_tb)
|
||||
end
|
||||
|
||||
function mongo_collection:safe_batch_update(updates)
|
||||
local updates_tb = {}
|
||||
for i = 1, #updates do
|
||||
@@ -409,25 +450,31 @@ function mongo_collection:safe_batch_update(updates)
|
||||
return werror(r)
|
||||
end
|
||||
|
||||
function mongo_collection:delete(selector, single)
|
||||
local sock = self.connection.__sock
|
||||
local pack = driver.delete(self.full_name, single, bson_encode(selector))
|
||||
sock:request(pack)
|
||||
function mongo_collection:raw_safe_update(update)
|
||||
local r = self.database:runCommand("update", self.name, "updates", {update})
|
||||
return werror(r)
|
||||
end
|
||||
|
||||
function mongo_collection:safe_delete(selector, single)
|
||||
function mongo_collection:delete(query, single)
|
||||
self.database:runCommand("delete", self.name, "deletes", {bson_encode({
|
||||
q = query,
|
||||
limit = single and 1 or 0,
|
||||
})})
|
||||
end
|
||||
|
||||
function mongo_collection:safe_delete(query, single)
|
||||
local r = self.database:runCommand("delete", self.name, "deletes", {bson_encode({
|
||||
q = selector,
|
||||
q = query,
|
||||
limit = single and 1 or 0,
|
||||
})})
|
||||
return werror(r)
|
||||
end
|
||||
|
||||
function mongo_collection:safe_batch_delete(selectors, single)
|
||||
function mongo_collection:safe_batch_delete(deletes, single)
|
||||
local delete_tb = {}
|
||||
for i = 1, #selectors do
|
||||
for i = 1, #deletes do
|
||||
delete_tb[i] = bson_encode({
|
||||
q = selectors[i],
|
||||
q = deletes[i],
|
||||
limit = single and 1 or 0,
|
||||
})
|
||||
end
|
||||
@@ -435,30 +482,29 @@ function mongo_collection:safe_batch_delete(selectors, single)
|
||||
return werror(r)
|
||||
end
|
||||
|
||||
function mongo_collection:findOne(query, selector)
|
||||
local conn = self.connection
|
||||
local request_id = conn:genId()
|
||||
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))
|
||||
-- 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 doc = req.document
|
||||
return bson_decode(doc)
|
||||
function mongo_collection:raw_safe_delete(delete)
|
||||
local r = self.database:runCommand("delete", self.name, "deletes", {delete})
|
||||
return werror(r)
|
||||
end
|
||||
|
||||
function mongo_collection:find(query, selector)
|
||||
function mongo_collection:findOne(query, projection)
|
||||
local cursor = self:find(query, projection)
|
||||
return cursor:hasNext() and cursor:next()
|
||||
end
|
||||
|
||||
function mongo_collection:find(query, projection)
|
||||
return setmetatable( {
|
||||
__collection = self,
|
||||
__query = query and bson_encode(query) or empty_bson,
|
||||
__selector = selector and bson_encode(selector),
|
||||
__projection = projection and bson_encode(projection) or empty_bson,
|
||||
__ptr = nil,
|
||||
__data = nil,
|
||||
__cursor = nil,
|
||||
__document = {},
|
||||
__flags = 0,
|
||||
__skip = 0,
|
||||
__sortquery = nil,
|
||||
__limit = 0,
|
||||
__sort = empty_bson,
|
||||
} , cursor_meta)
|
||||
end
|
||||
|
||||
@@ -479,7 +525,7 @@ function mongo_cursor:sort(key, key_v, ...)
|
||||
local key_list = unfold({}, key, key_v , ...)
|
||||
key = bson_encode_order(table.unpack(key_list))
|
||||
end
|
||||
self.__sortquery = bson_encode {['$query'] = self.__query, ['$orderby'] = key}
|
||||
self.__sort = key
|
||||
return self
|
||||
end
|
||||
|
||||
@@ -610,21 +656,37 @@ function mongo_collection:findAndModify(doc)
|
||||
return self.database:runCommand(table.unpack(cmd))
|
||||
end
|
||||
|
||||
-- https://docs.mongodb.com/manual/reference/command/aggregate/
|
||||
-- collection:aggregate({ { ["$project"] = {tags = 1} } }, {cursor={}})
|
||||
-- @param pipeline: array
|
||||
-- @param options: map
|
||||
-- @return
|
||||
function mongo_collection:aggregate(pipeline, options)
|
||||
assert(pipeline)
|
||||
local cmd = {"aggregate", self.name, "pipeline", pipeline}
|
||||
for k, v in pairs(options) do
|
||||
table.insert(cmd, k)
|
||||
table.insert(cmd, v)
|
||||
end
|
||||
return self.database:runCommand(table.unpack(cmd))
|
||||
end
|
||||
|
||||
function mongo_cursor:hasNext()
|
||||
if self.__ptr == nil then
|
||||
if self.__document == nil then
|
||||
return false
|
||||
end
|
||||
local conn = self.__collection.connection
|
||||
local request_id = conn:genId()
|
||||
local sock = conn.__sock
|
||||
local pack
|
||||
local response
|
||||
|
||||
local database = self.__collection.database
|
||||
if self.__data == nil then
|
||||
local query = self.__sortquery or self.__query
|
||||
pack = driver.query(request_id, self.__flags, self.__collection.full_name, self.__skip, self.__limit, query, self.__selector)
|
||||
local name = self.__collection.name
|
||||
response = database:runCommand("find", name, "filter", self.__query, "sort", self.__sort,
|
||||
"skip", self.__skip, "limit", self.__limit, "projection", self.__projection)
|
||||
else
|
||||
if self.__cursor then
|
||||
pack = driver.more(request_id, self.__collection.full_name, self.__limit, self.__cursor)
|
||||
if self.__cursor and self.__cursor > 0 then
|
||||
local name = self.__collection.name
|
||||
response = database:runCommand("getMore", bson.int64(self.__cursor), "collection", name, "batchSize", self.__limit)
|
||||
else
|
||||
-- no more
|
||||
self.__document = nil
|
||||
@@ -633,45 +695,31 @@ function mongo_cursor:hasNext()
|
||||
end
|
||||
end
|
||||
|
||||
local ok, result = pcall(sock.request,sock,pack, request_id)
|
||||
if response.ok ~= 1 then
|
||||
self.__document = nil
|
||||
self.__data = nil
|
||||
self.__cursor = nil
|
||||
error(response["$err"] or "Reply from mongod error")
|
||||
end
|
||||
|
||||
local doc = result.document
|
||||
local cursor = result.cursor_id
|
||||
|
||||
if ok then
|
||||
if doc then
|
||||
local doc = result.result
|
||||
self.__document = doc
|
||||
self.__data = result.data
|
||||
local cursor = response.cursor
|
||||
self.__document = cursor.firstBatch or cursor.nextBatch
|
||||
self.__data = response
|
||||
self.__ptr = 1
|
||||
self.__cursor = cursor
|
||||
self.__cursor = cursor.id
|
||||
|
||||
local limit = self.__limit
|
||||
if cursor and limit > 0 then
|
||||
limit = limit - #doc
|
||||
if cursor.id > 0 and limit > 0 then
|
||||
limit = limit - #self.__document
|
||||
if limit <= 0 then
|
||||
-- reach limit
|
||||
self:close()
|
||||
end
|
||||
|
||||
self.__limit = limit
|
||||
end
|
||||
|
||||
return true
|
||||
else
|
||||
self.__document = nil
|
||||
self.__data = nil
|
||||
self.__cursor = nil
|
||||
return false
|
||||
end
|
||||
else
|
||||
self.__document = nil
|
||||
self.__data = nil
|
||||
self.__cursor = nil
|
||||
if doc then
|
||||
local err = bson_decode(doc)
|
||||
error(err["$err"])
|
||||
else
|
||||
error("Reply from mongod error")
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
return true
|
||||
@@ -681,7 +729,7 @@ function mongo_cursor:next()
|
||||
if self.__ptr == nil then
|
||||
error "Call hasNext first"
|
||||
end
|
||||
local r = bson_decode(self.__document[self.__ptr])
|
||||
local r = self.__document[self.__ptr]
|
||||
self.__ptr = self.__ptr + 1
|
||||
if self.__ptr > #self.__document then
|
||||
self.__ptr = nil
|
||||
@@ -692,9 +740,8 @@ end
|
||||
|
||||
function mongo_cursor:close()
|
||||
if self.__cursor then
|
||||
local sock = self.__collection.connection.__sock
|
||||
local pack = driver.kill(self.__cursor)
|
||||
sock:request(pack)
|
||||
local coll = self.__collection
|
||||
coll.database:send_command("killCursors", coll.name, "cursors", {self.__cursor})
|
||||
self.__cursor = nil
|
||||
end
|
||||
end
|
||||
|
||||
@@ -14,7 +14,7 @@ local function _create_client()
|
||||
{
|
||||
host = host, port = port,
|
||||
username = username, password = password,
|
||||
authdb = db_name,
|
||||
authdb = "admin",
|
||||
}
|
||||
)
|
||||
end
|
||||
@@ -209,7 +209,7 @@ local function test_safe_batch_delete()
|
||||
db.testcoll:safe_batch_delete(docs)
|
||||
|
||||
local ret = db.testcoll:find()
|
||||
assert(length == ret:count(), "test safe batch delete failed")
|
||||
assert((length - del_num) == ret:count(), "test safe batch delete failed")
|
||||
end
|
||||
|
||||
skynet.start(function()
|
||||
|
||||
Reference in New Issue
Block a user