* Added 'count' method to the connection handle class

* Added 'remove' method to the connection handle class * Added 'update' method to the connection handle class * Added 'rows' method to the cursor handle class * Added more type checking for some methods * Cleaned up Makefile git-svn-id: http://luamongo.googlecode.com/svn/trunk@3 17e306f6-9439-11de-a13a-71b910e68cc8
moai · Apr 17, 2010 · c01d251 · c01d251
1 parent edd74e6
commit c01d251
Show file tree

Hide file tree

Showing 4 changed files with 161 additions and 12 deletions.
diff --git a/Makefile b/Makefile
@@ -1,9 +1,9 @@
 CC= g++
-CFLAGS= -g -O2 -shared -fpic -I /usr/include/lua5.1/ -I/usr/local/include/mongo/
+CFLAGS= -g -O2 -shared -fpic -I /usr/include/lua5.1/ -I/usr/include/mongo/
 AR= ar rcu
 RANLIB= ranlib
 RM= rm -f
-LIBS=-L /usr/local/lib -lmongoclient -lboost_thread-mt -lboost_filesystem-mt
+LIBS=-lmongoclient -lboost_thread-mt -lboost_filesystem-mt
 OUTLIB=mongo.so
 
 LDFLAGS= $(LIBS)

diff --git a/common.h b/common.h
@@ -1,9 +1,13 @@
-#define LUAMONGO_CONNECTION	"mongo.Connection"
-#define LUAMONGO_CURSOR		"mongo.Cursor"
+#define	LUAMONGO_CONNECTION	"mongo.Connection"
+#define	LUAMONGO_CURSOR		"mongo.Cursor"
 
-#define LUAMONGO_ERR_CONNECTION_FAILED	"Connection failed: %s"
-#define LUAMONGO_ERR_QUERY_FAILED	"Query failed: %s"
-#define LUAMONGO_ERR_INSERT_FAILED	"Insert failed: %s"
-#define LUAMONGO_ERR_CONNECT_FAILED	"Connection to %s failed: %s"
-#define LUAMONGO_UNSUPPORTED_BSON_TYPE	"Unsupported BSON type"
+#define	LUAMONGO_ERR_CONNECTION_FAILED	"Connection failed: %s"
+#define	LUAMONGO_ERR_QUERY_FAILED	"Query failed: %s"
+#define	LUAMONGO_ERR_INSERT_FAILED	"Insert failed: %s"
+#define	LUAMONGO_ERR_CONNECT_FAILED	"Connection to %s failed: %s"
+#define	LUAMONGO_ERR_COUNT_FAILED	"Count failed: %s"
+#define	LUAMONGO_ERR_REMOVE_FAILED	"Remove failed: %s"
+#define	LUAMONGO_ERR_UPDATE_FAILED	"Update failed: %s"
+#define	LUAMONGO_UNSUPPORTED_BSON_TYPE	"Unsupported BSON type"
 #define	LUAMONGO_UNSUPPORTED_LUA_TYPE	"Unsupported Lua type `%s'"
+#define	LUAMONGO_REQUIRES_JSON_OR_TABLE	"JSON string or Lua table required"
diff --git a/mongo_connection.cpp b/mongo_connection.cpp
@@ -59,6 +59,28 @@ static int connection_connect(lua_State *L) {
     return 1;
 }
 
+static int connection_count(lua_State *L) {
+    void *ud = 0;
+
+    ud = luaL_checkudata(L, 1, LUAMONGO_CONNECTION);
+    DBClientConnection *connection = *((DBClientConnection **)ud); 
+
+    const char *ns = luaL_checkstring(L, 2);
+
+    int count = 0;
+
+    try {
+        count = connection->count(ns);
+    } catch (std::exception &e) {
+	lua_pushnil(L);
+	lua_pushfstring(L, LUAMONGO_ERR_COUNT_FAILED, e.what());
+	return 2;
+    }
+
+    lua_pushinteger(L, count);
+    return 1;
+}
+
 static int connection_insert(lua_State *L) {
     void *ud = 0;
 
@@ -77,11 +99,17 @@ static int connection_insert(lua_State *L) {
 	    lua_to_bson(L, 3, data);
 
 	    connection->insert(ns, data);
+	} else {
+	    throw(LUAMONGO_REQUIRES_JSON_OR_TABLE);
 	}
     } catch (std::exception &e) {
 	lua_pushnil(L);
 	lua_pushfstring(L, LUAMONGO_ERR_INSERT_FAILED, e.what());
 	return 2;
+    } catch (const char *err) {
+        lua_pushnil(L);
+	lua_pushstring(L, err);
+	return 2;
     }
 
     lua_pushboolean(L, 1);
@@ -107,17 +135,111 @@ static int connection_query(lua_State *L) {
 		BSONObj obj;
 		lua_to_bson(L, 3, obj);
 		query = obj;
+	    } else {
+		throw(LUAMONGO_REQUIRES_JSON_OR_TABLE);
 	    }
 	} catch (std::exception &e) {
 	    lua_pushnil(L);
 	    lua_pushfstring(L, LUAMONGO_ERR_QUERY_FAILED, e.what());
 	    return 2;
+	} catch (const char *err) {
+	    lua_pushnil(L);
+	    lua_pushstring(L, err);
+	    return 2;
 	}
     }
 
     return cursor_create(L, connection, ns, query);
 }
 
+static int connection_remove(lua_State *L) {
+    void *ud = 0;
+
+    ud = luaL_checkudata(L, 1, LUAMONGO_CONNECTION);
+    DBClientConnection *connection = *((DBClientConnection **)ud); 
+
+    const char *ns = luaL_checkstring(L, 2);
+
+    try {
+	int type = lua_type(L, 3);
+	bool justOne = lua_toboolean(L, 4);
+
+	if (type == LUA_TSTRING) {
+	    const char *jsonstr = luaL_checkstring(L, 3);
+	    connection->remove(ns, fromjson(jsonstr), justOne);
+	} else if (type == LUA_TTABLE) {
+	    BSONObj data;
+	    lua_to_bson(L, 3, data);
+
+	    connection->remove(ns, data, justOne);
+	} else {
+	    throw(LUAMONGO_REQUIRES_JSON_OR_TABLE);
+	}
+    } catch (std::exception &e) {
+	lua_pushnil(L);
+	lua_pushfstring(L, LUAMONGO_ERR_REMOVE_FAILED, e.what());
+	return 2;
+    } catch (const char *err) {
+        lua_pushnil(L);
+	lua_pushstring(L, err);
+	return 2;
+    }
+
+    lua_pushboolean(L, 1);
+    return 1;
+}
+
+static int connection_update(lua_State *L) {
+    void *ud = 0;
+
+    ud = luaL_checkudata(L, 1, LUAMONGO_CONNECTION);
+    DBClientConnection *connection = *((DBClientConnection **)ud); 
+
+    const char *ns = luaL_checkstring(L, 2);
+
+    try {
+	int type_query = lua_type(L, 3);
+	int type_obj = lua_type(L, 4);
+
+	bool upsert = lua_toboolean(L, 5);
+	bool multi = lua_toboolean(L, 6);
+
+	BSONObj query;
+	BSONObj obj;
+
+	if (type_query == LUA_TSTRING) {
+	    const char *jsonstr = luaL_checkstring(L, 3);
+	    query = fromjson(jsonstr);
+	} else if (type_query == LUA_TTABLE) {
+	    lua_to_bson(L, 3, query);
+	} else {
+	    throw(LUAMONGO_REQUIRES_JSON_OR_TABLE);
+	}
+
+	if (type_obj == LUA_TSTRING) {
+	    const char *jsonstr = luaL_checkstring(L, 4);
+	    obj = fromjson(jsonstr);
+	} else if (type_obj == LUA_TTABLE) {
+	    lua_to_bson(L, 4, obj);
+	} else {
+	    throw(LUAMONGO_REQUIRES_JSON_OR_TABLE);
+	}
+
+	connection->update(ns, query, obj, upsert, multi);
+    } catch (std::exception &e) {
+	lua_pushnil(L);
+	lua_pushfstring(L, LUAMONGO_ERR_UPDATE_FAILED, e.what());
+	return 2;
+    } catch (const char *err) {
+        lua_pushnil(L);
+	lua_pushstring(L, err);
+	return 2;
+    }
+
+    lua_pushboolean(L, 1);
+    return 1;
+}
+
 static int connection_gc(lua_State *L) {
     void *ud = 0;
 
@@ -143,8 +265,11 @@ static int connection_tostring(lua_State *L) {
 int mongo_connection_register(lua_State *L) {
     static const luaL_Reg connection_methods[] = {
 	{"connect", connection_connect},
-        {"insert", connection_insert},
-        {"query", connection_query},
+	{"count", connection_count},
+	{"insert", connection_insert},
+	{"query", connection_query},
+	{"remove", connection_remove},
+	{"update", connection_update},
         {NULL, NULL}
     };
 

diff --git a/mongo_cursor.cpp b/mongo_cursor.cpp
@@ -45,7 +45,6 @@ static int cursor_next(lua_State *L) {
     DBClientCursor *cursor = *((DBClientCursor **)ud);
 
     if (cursor->more()) {
-	//lua_pushstring(L, cursor->next().jsonString().c_str());
 	bson_to_lua(L, cursor->next());	
     } else {
 	lua_pushnil(L);
@@ -54,6 +53,26 @@ static int cursor_next(lua_State *L) {
     return 1;
 }
 
+static int row_iterator(lua_State *L) {
+    void *ud = 0;
+
+    ud = luaL_checkudata(L, lua_upvalueindex(1), LUAMONGO_CURSOR);
+    DBClientCursor *cursor = *((DBClientCursor **)ud);
+
+    if (cursor->more()) {
+	bson_to_lua(L, cursor->next());	
+    } else {
+	lua_pushnil(L);
+    }
+
+    return 1;
+}
+
+static int cursor_rows(lua_State *L) {
+    lua_pushcclosure(L, row_iterator, 1);
+    return 1;
+}
+
 static int cursor_gc(lua_State *L) {
     void *ud = 0;
 
@@ -72,6 +91,7 @@ static int cursor_tostring(lua_State *L) {
 int mongo_cursor_register(lua_State *L) {
     static const luaL_Reg cursor_methods[] = {
 	{"next", cursor_next},
+	{"rows", cursor_rows},
         {NULL, NULL}
     };