#50 - query results as Enumerable

4 files changed, 85 insertions, 13 deletions

diff --git a/rbejdb/src/rbbson.c b/rbejdb/src/rbbson.c
index d8fb203..7acf287 100644
--- a/rbejdb/src/rbbson.c
+++ b/rbejdb/src/rbbson.c
@@ -20,9 +20,18 @@ VALUE iterate_array_callback(VALUE val, VALUE bsonWrap);
 
 VALUE bson_array_to_ruby(bson_iterator* it);
 
+VALUE bsonWrapClass = Qnil;
+
+
+void init_ruby_to_bson() {
+    bsonWrapClass = rb_define_class(BSON_RUBY_CLASS, rb_cObject);
+}
+
 
 VALUE createBsonWrap(bson* bsonval, VALUE rbobj, int flags) {
-    VALUE bsonWrapClass = rb_define_class(BSON_RUBY_CLASS, rb_cObject);
+    if (NIL_P(bsonWrapClass)) {
+        rb_raise(rb_eRuntimeError, "Ruby to BSON library must be initialized");
+    }
     VALUE bsonWrap = Data_Wrap_Struct(bsonWrapClass, NULL, NULL, ruby_xmalloc(sizeof(RBBSON)));
 
     RBBSON* rbbson;
diff --git a/rbejdb/src/rbbson.h b/rbejdb/src/rbbson.h
index f0fd5ee..e04e84f 100644
--- a/rbejdb/src/rbbson.h
+++ b/rbejdb/src/rbbson.h
@@ -17,6 +17,8 @@
 extern "C" {
 #endif
 
+    void init_ruby_to_bson();
+
     void ruby_to_bson(VALUE rbobj, bson** bsonbuf, int flags);
 
     VALUE bson_to_ruby(bson* bsonval);
diff --git a/rbejdb/src/rbejdb.c b/rbejdb/src/rbejdb.c
index f9275f2..9eea7e4 100644
--- a/rbejdb/src/rbejdb.c
+++ b/rbejdb/src/rbejdb.c
@@ -25,6 +25,18 @@ typedef struct {
     EJDB* ejdb;
 } RBEJDB;
 
+typedef struct {
+    TCLIST* results;
+} RBEJDB_RESULTS;
+
+
+VALUE create_EJDB_query_results(TCLIST* qres);
+
+
+VALUE ejdbClass;
+VALUE ejdbResultsClass;
+
+
 static int raise_ejdb_error(EJDB *ejdb) {
     int ecode = ejdbecode(ejdb);
     const char *emsg = ejdberrmsg(ecode);
@@ -39,7 +51,7 @@ EJDB* getEJDB(VALUE self) {
 }
 
 void EJDB_free(RBEJDB* rejdb) {
-    if (!rejdb->ejdb) {
+    if (rejdb->ejdb) {
         ejdbclose(rejdb->ejdb);
         ejdbdel(rejdb->ejdb);
     }
@@ -80,13 +92,15 @@ void EJDB_dropCollection(VALUE self, VALUE collName, VALUE prune) {
     Check_SafeStr(collName);
 
     EJDB* ejdb = getEJDB(self);
-    if (!ejdbrmcoll(ejdb, StringValuePtr(collName), TYPE(prune) == T_TRUE)) {
+    if (!ejdbrmcoll(ejdb, StringValuePtr(collName), RTEST(prune))) {
         raise_ejdb_error(ejdb);        
     }
 }
 
 void EJDB_ensureCollection(int argc, VALUE* argv, VALUE self) {
-    VALUE collName, copts;
+    VALUE collName;
+    VALUE copts;
+
     rb_scan_args(argc, argv, "11", &collName, &copts);
 
     Check_SafeStr(collName);
@@ -101,8 +115,8 @@ void EJDB_ensureCollection(int argc, VALUE* argv, VALUE self) {
         VALUE records = rb_hash_aref(copts, rb_str_new2("records"));
 
         jcopts.cachedrecords = !NIL_P(cachedrecords) ? NUM2INT(cachedrecords) : 0;
-        jcopts.compressed = TYPE(compressed) == T_TRUE;
-        jcopts.large = TYPE(large) == T_TRUE;
+        jcopts.compressed = RTEST(compressed);
+        jcopts.large = RTEST(large);
         jcopts.records = !NIL_P(records) ? NUM2INT(records) : 0;
     }
 
@@ -143,8 +157,14 @@ VALUE EJDB_save(int argc, VALUE *argv, VALUE self) {
     return Qnil;
 }
 
-VALUE EJDB_find(VALUE self, VALUE collName, VALUE q) {
+VALUE EJDB_find(int argc, VALUE* argv, VALUE self) {
+    VALUE collName;
+    VALUE q;
+
+    rb_scan_args(argc, argv, "11", &collName, &q);
+
     Check_SafeStr(collName);
+    q = !NIL_P(q) ? q :rb_hash_new();
 
     EJDB* ejdb = getEJDB(self);
 
@@ -162,6 +182,35 @@ VALUE EJDB_find(VALUE self, VALUE collName, VALUE q) {
     int qflags = 0;
     TCLIST* qres = ejdbqryexecute(coll, ejq, &count, qflags, NULL);
 
+    return create_EJDB_query_results(qres);
+}
+
+
+void EJDB_results_free(RBEJDB_RESULTS* rbres) {
+    if (rbres->results) {
+        tclistdel(rbres->results);
+    }
+    ruby_xfree(rbres);
+}
+
+VALUE create_EJDB_query_results(TCLIST* qres) {
+    VALUE results = Data_Wrap_Struct(ejdbResultsClass, NULL, EJDB_results_free, ruby_xmalloc(sizeof(RBEJDB_RESULTS)));
+    RBEJDB_RESULTS* rbresults;
+    Data_Get_Struct(results, RBEJDB_RESULTS, rbresults);
+
+    rbresults->results = qres;
+    return results;
+}
+
+VALUE EJDB_results_each(VALUE self) {
+    RBEJDB_RESULTS* rbresults;
+    Data_Get_Struct(self, RBEJDB_RESULTS, rbresults);
+
+    if (!rbresults || !rbresults->results) {
+        rb_raise(rb_eRuntimeError, "Each() method called on invalid ejdb query results");
+    }
+
+    TCLIST* qres = rbresults->results;
     int i;
     for (i = 0; i < TCLISTNUM(qres); i++) {
         char* bsrawdata = TCLISTVALPTR(qres, i);
@@ -170,12 +219,14 @@ VALUE EJDB_find(VALUE self, VALUE collName, VALUE q) {
         rb_yield(bson_to_ruby(&bsonval));
     }
 
-    tclistdel(qres);
-    ejdbquerydel(ejq);
+    return Qnil;
 }
 
+
 Init_rbejdb() {
-    VALUE ejdbClass = rb_define_class("EJDB", rb_cObject);
+    init_ruby_to_bson();
+
+    ejdbClass = rb_define_class("EJDB", rb_cObject);
     rb_define_alloc_func(ejdbClass, EJDB_alloc);
     rb_define_private_method(ejdbClass, "initialize", RUBY_METHOD_FUNC(EJDB_init), 0);
 
@@ -184,8 +235,13 @@ Init_rbejdb() {
     rb_define_method(ejdbClass, "open", RUBY_METHOD_FUNC(EJDB_open), 2);
     rb_define_method(ejdbClass, "is_open?", RUBY_METHOD_FUNC(EJDB_is_open), 0);
     rb_define_method(ejdbClass, "save", RUBY_METHOD_FUNC(EJDB_save), -1);
-    rb_define_method(ejdbClass, "find", RUBY_METHOD_FUNC(EJDB_find), 2);
+    rb_define_method(ejdbClass, "find", RUBY_METHOD_FUNC(EJDB_find), -1);
 
     rb_define_method(ejdbClass, "dropCollection", RUBY_METHOD_FUNC(EJDB_dropCollection), 2);
     rb_define_method(ejdbClass, "ensureCollection", RUBY_METHOD_FUNC(EJDB_ensureCollection), -1);
+
+
+    ejdbResultsClass = rb_define_class("EJDBResults", rb_cObject);
+    rb_include_module(ejdbResultsClass, rb_mEnumerable);
+    rb_define_method(ejdbResultsClass, "each", RUBY_METHOD_FUNC(EJDB_results_each), 0);
 }
\ No newline at end of file
diff --git a/rbejdb/test/t1.rb b/rbejdb/test/t1.rb
index 072bff2..ad21d1d 100644
--- a/rbejdb/test/t1.rb
+++ b/rbejdb/test/t1.rb
@@ -6,6 +6,7 @@ ejdb.open("zoo", EJDB::DEFAULT_OPEN_MODE)
 raise "Failed to open ejdb" unless ejdb.is_open?
 
 ejdb.dropCollection("parrots", true)
+ejdb.dropCollection("cows", true)
 
 ejdb.ensureCollection("parrots")
 ejdb.ensureCollection("parrots", {"large" => true, "records" => 200000})
@@ -23,13 +24,17 @@ ejdb.save("parrots", parrot1)
 parrot2 = {:name => "Mamadoo", :size => 666, "likes" => ["green color", "night", ["toys", "joys"], parrot1]}
 ejdb.save("parrots", parrot2)
 
+ejdb.save("cows", :name => "moo")
 
-ejdb.find("parrots", {:name => "Cacadoo"}) { |res|
+
+ejdb.find("parrots", {:name => "Cacadoo"}).each { |res|
   puts res.inspect
 }
 
-ejdb.find("parrots", {:name => {"$in" => ["Mamadoo", "Sauron"]}}) { |res|
+ejdb.find("parrots", {:name => {"$in" => ["Mamadoo", "Sauron"]}}).each { |res|
   puts res.inspect
 }
 
+puts ejdb.find("cows", {}).to_a.inspect
+
 puts "CONGRATULATIONS!!! EJDB tests have passed completely!"