From 3df52cc927b58a56408f8e6f068e832377dd86f9 Mon Sep 17 00:00:00 2001
From: Blythe Dunham <blythe@snowgiraffe.com>
Date: Sun, 6 Dec 2009 15:45:42 -0800
Subject: [PATCH] Add support for the finalize option of the group command

---
 lib/mongo/collection.rb | 32 +++++++++++++++++++++++---------
 test/test_collection.rb | 10 ++++++++++
 test/test_db_api.rb     |  6 ++++++
 3 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/lib/mongo/collection.rb b/lib/mongo/collection.rb
index 9a913a2..b0f62d7 100644
--- a/lib/mongo/collection.rb
+++ b/lib/mongo/collection.rb
@@ -346,22 +346,34 @@ module Mongo
     # :command :: if true, run the group as a command instead of in an
     #             eval - it is likely that this option will eventually be
     #             deprecated and all groups will be run as commands
-    def group(keys, condition, initial, reduce, command=false)
+    def group(keys, condition, initial, reduce, command=false, finalize=nil)
+
       if command
         hash = {}
         keys.each do |k|
           hash[k] = 1
         end
 
-        reduce = Code.new(reduce) unless reduce.is_a?(Code)
+        reduce =   Code.new(reduce)   unless reduce.is_a?(Code)
+
+        group_command = {
+          "group" => {
+            "ns"      => @name,
+            "$reduce" => reduce,
+            "key"     => hash,
+            "cond"    => condition,
+            "initial" => initial
+          }
+        }
+
+        # only add finalize if specified
+        if finalize
+          finalize = Code.new(finalize) unless finalize.is_a?(Code)
+          group_command['group']['finalize'] = finalize
+        end
+
+        result = @db.command group_command
 
-        result = @db.command({"group" =>
-                                  {
-                                    "ns" => @name,
-                                    "$reduce" => reduce,
-                                    "key" => hash,
-                                    "cond" => condition,
-                                    "initial" => initial}})
         if result["ok"] == 1
           return result["retval"]
         else
@@ -369,6 +381,8 @@ module Mongo
         end
       end
 
+      raise OperationFailure, "finalize is only supported with the group command" if finalize
+
       case reduce
       when Code
         scope = reduce.scope
diff --git a/test/test_collection.rb b/test/test_collection.rb
index bc0eea5..136c166 100644
--- a/test/test_collection.rb
+++ b/test/test_collection.rb
@@ -377,6 +377,16 @@ class TestCollection < Test::Unit::TestCase
 #     assert_equal 1, @@test.group([], {}, {"count" => 0},
 #                                  Code.new(reduce_function,
 #                                           {"inc_value" => 0.5}), true)[0]["count"]
+
+    # test finalize
+    #assert_equal( 3, 
+    #  @@test.group(
+    #    [], {}, {"count" => 0}, 
+    #    Code.new(reduce_function,{"inc_value" => 2}), true, 
+    #    Code.new("function (o) { o.final_count = o.count - 1; }")
+    #  )[0]["final_count"]
+    #)
+
     end
 
   context "A collection with two records" do 
diff --git a/test/test_db_api.rb b/test/test_db_api.rb
index fe5333f..368ac7a 100644
--- a/test/test_db_api.rb
+++ b/test/test_db_api.rb
@@ -553,6 +553,12 @@ class DBAPITest < Test::Unit::TestCase
     assert_equal 1, test.group([], {"a" => {"$gt" => 1}}, {"count" => 0}, "function (obj, prev) { prev.count++; }")[0]["count"]
     assert_equal 1, test.group([], {"a" => {"$gt" => 1}}, {"count" => 0}, "function (obj, prev) { prev.count++; }", true)[0]["count"]
 
+    finalize = "function (obj) { obj.f = obj.count - 1; }"
+    assert_equal 2, test.group([], {}, {"count" => 0}, "function (obj, prev) { prev.count++; }", true, finalize)[0]["f"]
+    assert_raise OperationFailure do
+      test.group([], {}, {"count" => 0}, "function (obj, prev) { prev.count++; }", false, finalize)[0]["f"]
+    end
+
     test.insert("a" => 2, "b" => 3)
     expected = [{"a" => 2, "count" => 2},
                 {"a" => nil, "count" => 1},