├── LICENSE ├── Makefile ├── README.md ├── lua-mongo.c ├── lua-socket.c ├── mongo.lua └── test.lua /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2013 codingnow.com 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 4 | 5 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 6 | 7 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 8 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | LUAINCLUDE=-I/usr/local/include 2 | LUALIB=-L/usr/local/bin -llua52 3 | SOCKETLIB=-lws2_32 4 | 5 | .PHONY: all win linux 6 | 7 | all : 8 | @echo Please do \'make PLATFORM\' where PLATFORM is one of these: 9 | @echo win linux 10 | 11 | win: mongo.dll 12 | 13 | linux: mongo.so 14 | 15 | mongo.dll : lua-mongo.c lua-socket.c 16 | gcc --shared -Wall -g $^ -o$@ $(LUAINCLUDE) $(LUALIB) $(SOCKETLIB) 17 | 18 | mongo.so : lua-mongo.c lua-socket.c 19 | gcc --shared -Wall -fPIC -g $^ $(LUAINCLUDE) -o$@ 20 | 21 | clean: 22 | rm -f mongo.dll mongo.so 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## Introduction 2 | 3 | This is a simple lua mongo driver, work in progress. 4 | 5 | ## Building 6 | 7 | Install lua-bson first. 8 | https://github.com/cloudwu/lua-bson 9 | 10 | ``` 11 | make win 12 | ``` 13 | or 14 | ``` 15 | make linux 16 | ``` 17 | 18 | ## Features 19 | 20 | * connect to mongod 21 | * runCommand 22 | * insert document 23 | * update document 24 | * find return basic cursor, cursor hasNext() and next() 25 | * findone 26 | 27 | ## Todo list 28 | 29 | * write concern 30 | * replica set 31 | * gridFS 32 | * tailable cursor 33 | * more options for cursor 34 | * more command 35 | * and more ... 36 | 37 | ## Getting started 38 | 39 | See test.lua 40 | -------------------------------------------------------------------------------- /lua-mongo.c: -------------------------------------------------------------------------------- 1 | #define LUA_LIB 2 | 3 | #include 4 | #include 5 | 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | #define OP_REPLY 1 12 | #define OP_MSG 1000 13 | #define OP_UPDATE 2001 14 | #define OP_INSERT 2002 15 | #define OP_QUERY 2004 16 | #define OP_GET_MORE 2005 17 | #define OP_DELETE 2006 18 | #define OP_KILL_CURSORS 2007 19 | 20 | #define REPLY_CURSORNOTFOUND 1 21 | #define REPLY_QUERYFAILURE 2 22 | #define REPLY_AWAITCAPABLE 8 // ignore because mongo 1.6+ always set it 23 | 24 | #define DEFAULT_CAP 128 25 | 26 | struct connection { 27 | int sock; 28 | int id; 29 | }; 30 | 31 | struct response { 32 | int flags; 33 | int32_t cursor_id[2]; 34 | int starting_from; 35 | int number; 36 | }; 37 | 38 | struct buffer { 39 | int size; 40 | int cap; 41 | uint8_t * ptr; 42 | uint8_t buffer[DEFAULT_CAP]; 43 | }; 44 | 45 | static inline uint32_t 46 | to_little_endian(uint32_t v) { 47 | union { 48 | uint32_t v; 49 | uint8_t b[4]; 50 | } u; 51 | u.v = v; 52 | return u.b[0] | u.b[1] << 8 | u.b[2] << 16 | u.b[3] << 24; 53 | } 54 | 55 | typedef void * document; 56 | 57 | static inline uint32_t 58 | get_length(const document buffer) { 59 | union { 60 | uint32_t v; 61 | uint8_t b[4]; 62 | } u; 63 | memcpy(&u.v, buffer, 4); 64 | return u.b[0] | u.b[1] << 8 | u.b[2] << 16 | u.b[3] << 24; 65 | } 66 | 67 | static inline void 68 | buffer_destroy(struct buffer *b) { 69 | if (b->ptr != b->buffer) { 70 | free(b->ptr); 71 | } 72 | } 73 | 74 | static inline void 75 | buffer_create(struct buffer *b) { 76 | b->size = 0; 77 | b->cap = DEFAULT_CAP; 78 | b->ptr = b->buffer; 79 | } 80 | 81 | static inline void 82 | buffer_reserve(struct buffer *b, int sz) { 83 | if (b->size + sz <= b->cap) 84 | return; 85 | do { 86 | b->cap *= 2; 87 | } while (b->cap <= b->size + sz); 88 | 89 | if (b->ptr == b->buffer) { 90 | b->ptr = malloc(b->cap); 91 | memcpy(b->ptr, b->buffer, b->size); 92 | } else { 93 | b->ptr = realloc(b->ptr, b->cap); 94 | } 95 | } 96 | 97 | static inline void 98 | write_int32(struct buffer *b, int32_t v) { 99 | uint32_t uv = (uint32_t)v; 100 | buffer_reserve(b,4); 101 | b->ptr[b->size++] = uv & 0xff; 102 | b->ptr[b->size++] = (uv >> 8)&0xff; 103 | b->ptr[b->size++] = (uv >> 16)&0xff; 104 | b->ptr[b->size++] = (uv >> 24)&0xff; 105 | } 106 | 107 | static inline void 108 | write_bytes(struct buffer *b, const void * buf, int sz) { 109 | buffer_reserve(b,sz); 110 | memcpy(b->ptr + b->size, buf, sz); 111 | b->size += sz; 112 | } 113 | 114 | static void 115 | write_string(struct buffer *b, const char *key, size_t sz) { 116 | buffer_reserve(b,sz+1); 117 | memcpy(b->ptr + b->size, key, sz); 118 | b->ptr[b->size+sz] = '\0'; 119 | b->size+=sz+1; 120 | } 121 | 122 | static inline int 123 | reserve_length(struct buffer *b) { 124 | int sz = b->size; 125 | buffer_reserve(b,4); 126 | b->size +=4; 127 | return sz; 128 | } 129 | 130 | static inline void 131 | write_length(struct buffer *b, int32_t v, int off) { 132 | uint32_t uv = (uint32_t)v; 133 | b->ptr[off++] = uv & 0xff; 134 | b->ptr[off++] = (uv >> 8)&0xff; 135 | b->ptr[off++] = (uv >> 16)&0xff; 136 | b->ptr[off++] = (uv >> 24)&0xff; 137 | } 138 | 139 | // 1 integer id 140 | // 2 integer flags 141 | // 3 string collection name 142 | // 4 integer skip 143 | // 5 integer return number 144 | // 6 document query 145 | // 7 document selector (optional) 146 | // return string package 147 | static int 148 | op_query(lua_State *L) { 149 | int id = luaL_checkinteger(L,1); 150 | document query = lua_touserdata(L,6); 151 | if (query == NULL) { 152 | return luaL_error(L, "require query document"); 153 | } 154 | document selector = lua_touserdata(L,7); 155 | int flags = luaL_checkinteger(L, 2); 156 | size_t nsz = 0; 157 | const char *name = luaL_checklstring(L,3,&nsz); 158 | int skip = luaL_checkinteger(L, 4); 159 | int number = luaL_checkinteger(L, 5); 160 | 161 | luaL_Buffer b; 162 | luaL_buffinit(L,&b); 163 | 164 | struct buffer buf; 165 | buffer_create(&buf); 166 | int len = reserve_length(&buf); 167 | write_int32(&buf, id); 168 | write_int32(&buf, 0); 169 | write_int32(&buf, OP_QUERY); 170 | write_int32(&buf, flags); 171 | write_string(&buf, name, nsz); 172 | write_int32(&buf, skip); 173 | write_int32(&buf, number); 174 | 175 | int32_t query_len = get_length(query); 176 | int total = buf.size + query_len; 177 | int32_t selector_len = 0; 178 | if (selector) { 179 | selector_len = get_length(selector); 180 | total += selector_len; 181 | } 182 | 183 | write_length(&buf, total, len); 184 | luaL_addlstring(&b, (const char *)buf.ptr, buf.size); 185 | buffer_destroy(&buf); 186 | 187 | luaL_addlstring(&b, (const char *)query, query_len); 188 | 189 | if (selector) { 190 | luaL_addlstring(&b, (const char *)selector, selector_len); 191 | } 192 | 193 | luaL_pushresult(&b); 194 | 195 | return 1; 196 | } 197 | 198 | // 1 string data 199 | // 2 result document table 200 | // return boolean succ (false -> request id, error document) 201 | // number request_id 202 | // document first 203 | // string cursor_id 204 | // integer startfrom 205 | static int 206 | op_reply(lua_State *L) { 207 | size_t data_len = 0; 208 | const char * data = luaL_checklstring(L,1,&data_len); 209 | struct { 210 | // int32_t length; // total message size, including this 211 | int32_t request_id; // identifier for this message 212 | int32_t response_id; // requestID from the original request 213 | // (used in reponses from db) 214 | int32_t opcode; // request type 215 | int32_t flags; 216 | int32_t cursor_id[2]; 217 | int32_t starting; 218 | int32_t number; 219 | } const *reply = (const void *)data; 220 | 221 | if (data_len < sizeof(*reply)) { 222 | lua_pushboolean(L, 0); 223 | return 1; 224 | } 225 | 226 | int id = to_little_endian(reply->response_id); 227 | int flags = to_little_endian(reply->flags); 228 | if (flags & REPLY_QUERYFAILURE) { 229 | lua_pushboolean(L,0); 230 | lua_pushinteger(L, id); 231 | lua_pushlightuserdata(L, (void *)(reply+1)); 232 | return 3; 233 | } 234 | 235 | int starting_from = to_little_endian(reply->starting); 236 | int number = to_little_endian(reply->number); 237 | int sz = (int)data_len - sizeof(*reply); 238 | const uint8_t * doc = (const uint8_t *)(reply+1); 239 | 240 | if (lua_istable(L,2)) { 241 | int i = 1; 242 | while (sz > 4) { 243 | lua_pushlightuserdata(L, (void *)doc); 244 | lua_rawseti(L, 2, i); 245 | 246 | int32_t doc_len = get_length((const document)doc); 247 | 248 | doc += doc_len; 249 | sz -= doc_len; 250 | 251 | ++i; 252 | } 253 | if (i != number + 1) { 254 | lua_pushboolean(L,0); 255 | lua_pushinteger(L, id); 256 | return 2; 257 | } 258 | int c = lua_rawlen(L, 2); 259 | for (;i<=c;i++) { 260 | lua_pushnil(L); 261 | lua_rawseti(L, 2, i); 262 | } 263 | } 264 | lua_pushboolean(L,1); 265 | lua_pushinteger(L, id); 266 | if (number == 0) 267 | lua_pushnil(L); 268 | else 269 | lua_pushlightuserdata(L, (void *)(reply+1)); 270 | if (reply->cursor_id[0] == 0 && reply->cursor_id[1]==0) { 271 | // closed cursor 272 | lua_pushnil(L); 273 | } else { 274 | lua_pushlstring(L, (const char *)(reply->cursor_id), 8); 275 | } 276 | lua_pushinteger(L, starting_from); 277 | 278 | return 5; 279 | } 280 | 281 | /* 282 | 1 string cursor_id 283 | return string package 284 | */ 285 | static int 286 | op_kill(lua_State *L) { 287 | size_t cursor_len = 0; 288 | const char * cursor_id = luaL_tolstring(L, 1, &cursor_len); 289 | if (cursor_len != 8) { 290 | return luaL_error(L, "Invalid cursor id"); 291 | } 292 | 293 | struct buffer buf; 294 | buffer_create(&buf); 295 | 296 | int len = reserve_length(&buf); 297 | write_int32(&buf, 0); 298 | write_int32(&buf, 0); 299 | write_int32(&buf, OP_KILL_CURSORS); 300 | 301 | write_int32(&buf, 0); 302 | write_int32(&buf, 1); 303 | write_bytes(&buf, cursor_id, 8); 304 | 305 | write_length(&buf, buf.size, len); 306 | 307 | lua_pushlstring(L, (const char *)buf.ptr, buf.size); 308 | buffer_destroy(&buf); 309 | 310 | return 1; 311 | } 312 | 313 | /* 314 | 1 string collection 315 | 2 integer single remove 316 | 3 document selector 317 | 318 | return string package 319 | */ 320 | static int 321 | op_delete(lua_State *L) { 322 | document selector = lua_touserdata(L,3); 323 | if (selector == NULL) { 324 | luaL_error(L, "Invalid param"); 325 | } 326 | size_t sz = 0; 327 | const char * name = luaL_checklstring(L,1,&sz); 328 | 329 | luaL_Buffer b; 330 | luaL_buffinit(L,&b); 331 | 332 | struct buffer buf; 333 | buffer_create(&buf); 334 | int len = reserve_length(&buf); 335 | write_int32(&buf, 0); 336 | write_int32(&buf, 0); 337 | write_int32(&buf, OP_DELETE); 338 | write_int32(&buf, 0); 339 | write_string(&buf, name, sz); 340 | write_int32(&buf, lua_tointeger(L,2)); 341 | 342 | int32_t selector_len = get_length(selector); 343 | int total = buf.size + selector_len; 344 | write_length(&buf, total, len); 345 | 346 | luaL_addlstring(&b, (const char *)buf.ptr, buf.size); 347 | buffer_destroy(&buf); 348 | 349 | luaL_addlstring(&b, (const char *)selector, selector_len); 350 | luaL_pushresult(&b); 351 | 352 | return 1; 353 | } 354 | 355 | /* 356 | 1 integer id 357 | 2 string collection 358 | 3 integer number 359 | 4 cursor_id (8 bytes string/ 64bit) 360 | 361 | return string package 362 | */ 363 | static int 364 | op_get_more(lua_State *L) { 365 | int id = luaL_checkinteger(L, 1); 366 | size_t sz = 0; 367 | const char * name = luaL_checklstring(L,2,&sz); 368 | int number = luaL_checkinteger(L, 3); 369 | size_t cursor_len = 0; 370 | const char * cursor_id = luaL_tolstring(L, 4, &cursor_len); 371 | if (cursor_len != 8) { 372 | return luaL_error(L, "Invalid cursor id"); 373 | } 374 | 375 | struct buffer buf; 376 | buffer_create(&buf); 377 | int len = reserve_length(&buf); 378 | write_int32(&buf, id); 379 | write_int32(&buf, 0); 380 | write_int32(&buf, OP_GET_MORE); 381 | write_int32(&buf, 0); 382 | write_string(&buf, name, sz); 383 | write_int32(&buf, number); 384 | write_bytes(&buf, cursor_id, 8); 385 | write_length(&buf, buf.size, len); 386 | 387 | lua_pushlstring(L, (const char *)buf.ptr, buf.size); 388 | buffer_destroy(&buf); 389 | 390 | return 1; 391 | } 392 | 393 | // 1 string collection 394 | // 2 integer flags 395 | // 3 document selector 396 | // 4 document update 397 | // return string package 398 | static int 399 | op_update(lua_State *L) { 400 | document selector = lua_touserdata(L,3); 401 | document update = lua_touserdata(L,4); 402 | if (selector == NULL || update == NULL) { 403 | luaL_error(L, "Invalid param"); 404 | } 405 | size_t sz = 0; 406 | const char * name = luaL_checklstring(L,1,&sz); 407 | 408 | luaL_Buffer b; 409 | luaL_buffinit(L, &b); 410 | 411 | struct buffer buf; 412 | buffer_create(&buf); 413 | // make package header, don't raise L error 414 | int len = reserve_length(&buf); 415 | write_int32(&buf, 0); 416 | write_int32(&buf, 0); 417 | write_int32(&buf, OP_UPDATE); 418 | write_int32(&buf, 0); 419 | write_string(&buf, name, sz); 420 | write_int32(&buf, lua_tointeger(L,2)); 421 | 422 | int32_t selector_len = get_length(selector); 423 | int32_t update_len = get_length(update); 424 | 425 | int total = buf.size + selector_len + update_len; 426 | write_length(&buf, total, len); 427 | 428 | luaL_addlstring(&b, (const char *)buf.ptr, buf.size); 429 | buffer_destroy(&buf); 430 | 431 | luaL_addlstring(&b, (const char *)selector, selector_len); 432 | luaL_addlstring(&b, (const char *)update, update_len); 433 | 434 | luaL_pushresult(&b); 435 | 436 | return 1; 437 | } 438 | 439 | static int 440 | document_length(lua_State *L) { 441 | if (lua_isuserdata(L, 3)) { 442 | document doc = lua_touserdata(L,3); 443 | return get_length(doc); 444 | } 445 | if (lua_istable(L,3)) { 446 | int total = 0; 447 | int s = lua_rawlen(L,3); 448 | int i; 449 | for (i=1;i<=s;i++) { 450 | lua_rawgeti(L, 3, i); 451 | document doc = lua_touserdata(L,-1); 452 | if (doc == NULL) { 453 | lua_pop(L,1); 454 | return luaL_error(L, "Invalid document at %d", i); 455 | } else { 456 | total += get_length(doc); 457 | lua_pop(L,1); 458 | } 459 | } 460 | return total; 461 | } 462 | return luaL_error(L, "Insert need documents"); 463 | } 464 | 465 | // 1 integer flags 466 | // 2 string collection 467 | // 3 documents 468 | // return string package 469 | static int 470 | op_insert(lua_State *L) { 471 | size_t sz = 0; 472 | const char * name = luaL_checklstring(L,2,&sz); 473 | int dsz = document_length(L); 474 | 475 | luaL_Buffer b; 476 | luaL_buffinit(L, &b); 477 | 478 | struct buffer buf; 479 | buffer_create(&buf); 480 | // make package header, don't raise L error 481 | int len = reserve_length(&buf); 482 | write_int32(&buf, 0); 483 | write_int32(&buf, 0); 484 | write_int32(&buf, OP_INSERT); 485 | write_int32(&buf, lua_tointeger(L,1)); 486 | write_string(&buf, name, sz); 487 | 488 | int total = buf.size + dsz; 489 | write_length(&buf, total, len); 490 | 491 | luaL_addlstring(&b, (const char *)buf.ptr, buf.size); 492 | buffer_destroy(&buf); 493 | 494 | if (lua_isuserdata(L,3)) { 495 | document doc = lua_touserdata(L,3); 496 | luaL_addlstring(&b, (const char *)doc, get_length(doc)); 497 | } else { 498 | int s = lua_rawlen(L, 3); 499 | int i; 500 | for (i=1;i<=s;i++) { 501 | lua_rawgeti(L,3,i); 502 | document doc = lua_touserdata(L,3); 503 | lua_pop(L,1); // must call lua_pop before luaL_addlstring, because addlstring may change stack top 504 | luaL_addlstring(&b, (const char *)doc, get_length(doc)); 505 | } 506 | } 507 | 508 | luaL_pushresult(&b); 509 | 510 | return 1; 511 | } 512 | 513 | // string 4 bytes length 514 | // return integer 515 | static int 516 | reply_length(lua_State *L) { 517 | const char * rawlen_str = luaL_checkstring(L, 1); 518 | int rawlen = 0; 519 | memcpy(&rawlen, rawlen_str, sizeof(int)); 520 | int length = to_little_endian(rawlen); 521 | lua_pushinteger(L, length - 4); 522 | return 1; 523 | } 524 | 525 | LUAMOD_API int 526 | luaopen_mongo_driver(lua_State *L) { 527 | luaL_checkversion(L); 528 | luaL_Reg l[] ={ 529 | { "query", op_query }, 530 | { "reply", op_reply }, 531 | { "kill", op_kill }, 532 | { "delete", op_delete }, 533 | { "more", op_get_more }, 534 | { "update", op_update }, 535 | { "insert", op_insert }, 536 | { "length", reply_length }, 537 | { NULL, NULL }, 538 | }; 539 | 540 | luaL_newlib(L,l); 541 | return 1; 542 | } 543 | -------------------------------------------------------------------------------- /lua-socket.c: -------------------------------------------------------------------------------- 1 | #define LUA_LIB 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | #include 8 | 9 | #if defined(_WIN32) || defined(_WIN64) 10 | 11 | #define _WIN32_WINNT 0x0501 12 | 13 | #include 14 | #include 15 | 16 | static void 17 | init_winsock() { 18 | WSADATA wsaData; 19 | WSAStartup(MAKEWORD(2,2), &wsaData); 20 | } 21 | 22 | #define close(fd) closesocket(fd) 23 | 24 | #else 25 | 26 | #include 27 | #include 28 | #define INVALID_SOCKET (-1) 29 | 30 | static void 31 | init_winsock() { 32 | } 33 | 34 | #endif 35 | 36 | #ifndef MSG_WAITALL 37 | #define MSG_WAITALL 0 38 | #endif 39 | 40 | #define LOCALBUFFER 8192 41 | 42 | static int 43 | lopen(lua_State *L) { 44 | const char * host = luaL_checkstring(L,1); 45 | int port = luaL_checkinteger(L,2); 46 | 47 | char port_str[32]; 48 | int status; 49 | 50 | struct addrinfo ai_hints; 51 | struct addrinfo *ai_list = NULL; 52 | struct addrinfo *ai_ptr = NULL; 53 | 54 | sprintf( port_str, "%d", port ); 55 | 56 | memset( &ai_hints, 0, sizeof( ai_hints ) ); 57 | ai_hints.ai_family = AF_UNSPEC; 58 | ai_hints.ai_socktype = SOCK_STREAM; 59 | ai_hints.ai_protocol = IPPROTO_TCP; 60 | 61 | status = getaddrinfo( host, port_str, &ai_hints, &ai_list ); 62 | if ( status != 0 ) { 63 | return 0; 64 | } 65 | int sock=INVALID_SOCKET; 66 | for ( ai_ptr = ai_list; ai_ptr != NULL; ai_ptr = ai_ptr->ai_next ) { 67 | sock = socket( ai_ptr->ai_family, ai_ptr->ai_socktype, ai_ptr->ai_protocol ); 68 | if ( sock == INVALID_SOCKET ) { 69 | continue; 70 | } 71 | 72 | status = connect( sock, ai_ptr->ai_addr, ai_ptr->ai_addrlen ); 73 | if ( status != 0 ) { 74 | close(sock); 75 | sock = INVALID_SOCKET; 76 | continue; 77 | } 78 | break; 79 | } 80 | 81 | freeaddrinfo( ai_list ); 82 | 83 | if (sock != INVALID_SOCKET) { 84 | lua_pushinteger(L, sock); 85 | return 1; 86 | } 87 | 88 | return 0; 89 | } 90 | 91 | static int 92 | lclose(lua_State *L) { 93 | int sock = luaL_checkinteger(L, 1); 94 | close(sock); 95 | 96 | return 0; 97 | } 98 | 99 | static int 100 | lread(lua_State *L) { 101 | int fd = luaL_checkinteger(L,1); 102 | int sz = luaL_checkinteger(L,2); 103 | char tmp[LOCALBUFFER]; 104 | void * buffer = tmp; 105 | if (sz > LOCALBUFFER) { 106 | buffer = lua_newuserdata(L, sz); 107 | } 108 | char * ptr = buffer; 109 | for (;;) { 110 | int bytes = recv(fd, ptr, sz, MSG_WAITALL); 111 | if (bytes < 0) { 112 | switch (errno) { 113 | case EAGAIN: 114 | case EINTR: 115 | continue; 116 | } 117 | return 0; 118 | } 119 | if (bytes == 0) { 120 | return 0; 121 | } 122 | if (bytes < sz) { 123 | ptr += bytes; 124 | sz -= bytes; 125 | return 0; 126 | } 127 | lua_pushlstring(L, buffer, sz); 128 | return 1; 129 | } 130 | } 131 | 132 | static int 133 | lwrite(lua_State *L) { 134 | int sock = luaL_checkinteger(L,1); 135 | size_t sz = 0; 136 | const char *buffer = luaL_checklstring(L, 2, &sz); 137 | for (;;) { 138 | int err = send(sock, buffer, sz, 0); 139 | if (err < 0) { 140 | switch (errno) { 141 | case EAGAIN: 142 | case EINTR: 143 | continue; 144 | } 145 | } 146 | if (err != sz) { 147 | lua_pushboolean(L,0); 148 | return 1; 149 | } 150 | break; 151 | } 152 | lua_pushboolean(L,1); 153 | return 1; 154 | } 155 | 156 | LUAMOD_API int 157 | luaopen_mongo_socket(lua_State *L) { 158 | init_winsock(); 159 | luaL_checkversion(L); 160 | luaL_Reg l[] ={ 161 | { "open", lopen }, 162 | { "close", lclose }, 163 | { "read", lread }, 164 | { "write", lwrite }, 165 | { NULL, NULL }, 166 | }; 167 | 168 | luaL_newlib(L,l); 169 | 170 | return 1; 171 | } 172 | -------------------------------------------------------------------------------- /mongo.lua: -------------------------------------------------------------------------------- 1 | local bson = require "bson" 2 | local socket = require "mongo.socket" 3 | local driver = require "mongo.driver" 4 | local rawget = rawget 5 | local assert = assert 6 | 7 | local bson_encode = bson.encode 8 | local bson_encode_order = bson.encode_order 9 | local bson_decode = bson.decode 10 | local empty_bson = bson_encode {} 11 | 12 | local mongo = {} 13 | mongo.null = assert(bson.null) 14 | mongo.maxkey = assert(bson.maxkey) 15 | mongo.minkey = assert(bson.minkey) 16 | mongo.type = assert(bson.type) 17 | 18 | local mongo_cursor = {} 19 | local cursor_meta = { 20 | __index = mongo_cursor, 21 | } 22 | 23 | local mongo_client = {} 24 | 25 | local client_meta = { 26 | __index = function(self, key) 27 | return rawget(mongo_client, key) or self:getDB(key) 28 | end, 29 | __tostring = function (self) 30 | local port_string 31 | if self.port then 32 | port_string = ":" .. tostring(self.port) 33 | else 34 | port_string = "" 35 | end 36 | 37 | return "[mongo client : " .. self.host .. port_string .."]" 38 | end, 39 | __gc = function(self) 40 | self:disconnect() 41 | end 42 | } 43 | 44 | local mongo_db = {} 45 | 46 | local db_meta = { 47 | __index = function (self, key) 48 | return rawget(mongo_db, key) or self:getCollection(key) 49 | end, 50 | __tostring = function (self) 51 | return "[mongo db : " .. self.name .. "]" 52 | end 53 | } 54 | 55 | local mongo_collection = {} 56 | local collection_meta = { 57 | __index = function(self, key) 58 | return rawget(mongo_collection, key) or self:getCollection(key) 59 | end , 60 | __tostring = function (self) 61 | return "[mongo collection : " .. self.full_name .. "]" 62 | end 63 | } 64 | 65 | function mongo.client( obj ) 66 | obj.port = obj.port or 27017 67 | obj.__id = 0 68 | obj.__sock = assert(socket.open(obj.host, obj.port),"Connect failed") 69 | return setmetatable(obj, client_meta) 70 | end 71 | 72 | function mongo_client:getDB(dbname) 73 | local db = { 74 | connection = self, 75 | name = dbname, 76 | full_name = dbname, 77 | database = false, 78 | __cmd = dbname .. "." .. "$cmd", 79 | } 80 | 81 | db.database = db 82 | 83 | return setmetatable(db, db_meta) 84 | end 85 | 86 | function mongo_client:disconnect() 87 | if self.__sock then 88 | socket.close(self.__sock) 89 | self.__sock = nil 90 | end 91 | end 92 | 93 | function mongo_client:genId() 94 | local id = self.__id + 1 95 | self.__id = id 96 | return id 97 | end 98 | 99 | function mongo_client:runCommand(...) 100 | if not self.admin then 101 | self.admin = self:getDB "admin" 102 | end 103 | return self.admin:runCommand(...) 104 | end 105 | 106 | local function get_reply(sock, result) 107 | local length = driver.length(socket.read(sock, 4)) 108 | local reply = socket.read(sock, length) 109 | return reply, driver.reply(reply, result) 110 | end 111 | 112 | function mongo_db:runCommand(cmd,cmd_v,...) 113 | local request_id = self.connection:genId() 114 | local sock = self.connection.__sock 115 | local bson_cmd 116 | if not cmd_v then 117 | bson_cmd = bson_encode_order(cmd,1) 118 | else 119 | bson_cmd = bson_encode_order(cmd,cmd_v,...) 120 | end 121 | local pack = driver.query(request_id, 0, self.__cmd, 0, 1, bson_cmd) 122 | -- todo: check send 123 | socket.write(sock, pack) 124 | 125 | local _, succ, reply_id, doc = get_reply(sock) 126 | assert(request_id == reply_id, "Reply from mongod error") 127 | -- todo: check succ 128 | return bson_decode(doc) 129 | end 130 | 131 | function mongo_db:getCollection(collection) 132 | local col = { 133 | connection = self.connection, 134 | name = collection, 135 | full_name = self.full_name .. "." .. collection, 136 | database = self.database, 137 | } 138 | self[collection] = setmetatable(col, collection_meta) 139 | return col 140 | end 141 | 142 | mongo_collection.getCollection = mongo_db.getCollection 143 | 144 | function mongo_collection:insert(doc) 145 | if doc._id == nil then 146 | doc._id = bson.objectid() 147 | end 148 | local sock = self.connection.__sock 149 | local pack = driver.insert(0, self.full_name, bson_encode(doc)) 150 | -- todo: check send 151 | -- flags support 1: ContinueOnError 152 | socket.write(sock, pack) 153 | end 154 | 155 | function mongo_collection:batch_insert(docs) 156 | for i=1,#docs do 157 | if docs[i]._id == nil then 158 | docs[i]._id = bson.objectid() 159 | end 160 | docs[i] = bson_encode(docs[i]) 161 | end 162 | local sock = self.connection.__sock 163 | local pack = driver.insert(0, self.full_name, docs) 164 | -- todo: check send 165 | socket.write(sock, pack) 166 | end 167 | 168 | function mongo_collection:update(selector,update,upsert,multi) 169 | local flags = (upsert and 1 or 0) + (multi and 2 or 0) 170 | local sock = self.connection.__sock 171 | local pack = driver.update(self.full_name, flags, bson_encode(selector), bson_encode(update)) 172 | -- todo: check send 173 | socket.write(sock, pack) 174 | end 175 | 176 | function mongo_collection:delete(selector, single) 177 | local sock = self.connection.__sock 178 | local pack = driver.delete(self.full_name, single, bson_encode(selector)) 179 | -- todo: check send 180 | socket.write(sock, pack) 181 | end 182 | 183 | function mongo_collection:findOne(query, selector) 184 | local request_id = self.connection:genId() 185 | local sock = self.connection.__sock 186 | 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)) 187 | 188 | -- todo: check send 189 | socket.write(sock, pack) 190 | 191 | local _, succ, reply_id, doc = get_reply(sock) 192 | assert(request_id == reply_id, "Reply from mongod error") 193 | -- todo: check succ 194 | return bson_decode(doc) 195 | end 196 | 197 | function mongo_collection:find(query, selector) 198 | return setmetatable( { 199 | __collection = self, 200 | __query = query and bson_encode(query) or empty_bson, 201 | __selector = selector and bson_encode(selector), 202 | __ptr = nil, 203 | __data = nil, 204 | __cursor = nil, 205 | __document = {}, 206 | __flags = 0, 207 | } , cursor_meta) 208 | end 209 | 210 | function mongo_cursor:hasNext() 211 | if self.__ptr == nil then 212 | if self.__document == nil then 213 | return false 214 | end 215 | local conn = self.__collection.connection 216 | local request_id = conn:genId() 217 | local sock = conn.__sock 218 | local pack 219 | if self.__data == nil then 220 | pack = driver.query(request_id, self.__flags, self.__collection.full_name,0,0,self.__query,self.__selector) 221 | else 222 | if self.__cursor then 223 | pack = driver.more(request_id, self.__collection.full_name,0,self.__cursor) 224 | else 225 | -- no more 226 | self.__document = nil 227 | self.__data = nil 228 | return false 229 | end 230 | end 231 | 232 | --todo: check send 233 | socket.write(sock, pack) 234 | 235 | local data, succ, reply_id, doc, cursor = get_reply(sock, self.__document) 236 | assert(request_id == reply_id, "Reply from mongod error") 237 | if succ then 238 | if doc then 239 | self.__data = data 240 | self.__ptr = 1 241 | self.__cursor = cursor 242 | return true 243 | else 244 | self.__document = nil 245 | self.__data = nil 246 | self.__cursor = nil 247 | return false 248 | end 249 | else 250 | self.__document = nil 251 | self.__data = nil 252 | self.__cursor = nil 253 | if doc then 254 | local err = bson_decode(doc) 255 | error(err["$err"]) 256 | else 257 | error("Reply from mongod error") 258 | end 259 | end 260 | end 261 | 262 | return true 263 | end 264 | 265 | function mongo_cursor:next() 266 | if self.__ptr == nil then 267 | error "Call hasNext first" 268 | end 269 | local r = bson_decode(self.__document[self.__ptr]) 270 | self.__ptr = self.__ptr + 1 271 | if self.__ptr > #self.__document then 272 | self.__ptr = nil 273 | end 274 | 275 | return r 276 | end 277 | 278 | function mongo_cursor:close() 279 | -- todo: warning hasNext after close 280 | if self.__cursor then 281 | local sock = self.__collection.connection.__sock 282 | local pack = driver.kill(self.__cursor) 283 | -- todo: check send 284 | socket.write(sock, pack) 285 | end 286 | end 287 | 288 | return mongo -------------------------------------------------------------------------------- /test.lua: -------------------------------------------------------------------------------- 1 | mongo = require "mongo" 2 | 3 | db = mongo.client { host = "localhost" } 4 | 5 | local r = db:runCommand "listDatabases" 6 | 7 | for k,v in ipairs(r.databases) do 8 | print(v.name) 9 | end 10 | 11 | 12 | local loc = db:getDB "hello" 13 | local c = loc.system.namespaces:find() 14 | 15 | while c:hasNext() do 16 | local r = c:next() 17 | print(r.name) 18 | end 19 | 20 | print "===============" 21 | 22 | db.hello.world:insert {} 23 | local r = db:runCommand ("getLastError",1,"w",1) 24 | print(r.ok) 25 | 26 | local c = db.hello.world:find() 27 | 28 | while c:hasNext() do 29 | local r = c:next() 30 | print(mongo.type(r._id)) 31 | end 32 | --------------------------------------------------------------------------------