From f62e9e6732e5be2e4e9bd60fc3b9683519d7fbc4 Mon Sep 17 00:00:00 2001
From: Rimas Silkaitis <neovintage@gmail.com>
Date: Sat, 3 Jul 2010 11:24:24 -0700
Subject: [PATCH] Added helper functions to create and remove stored js
 functions in system.js

---
 lib/mongo/db.rb | 31 +++++++++++++++++++++++++++++++
 test/db_test.rb |  8 ++++++++
 2 files changed, 39 insertions(+)

diff --git a/lib/mongo/db.rb b/lib/mongo/db.rb
index e9aea98fbd..8649d57ada 100644
--- a/lib/mongo/db.rb
+++ b/lib/mongo/db.rb
@@ -29,6 +29,7 @@ class DB
     SYSTEM_INDEX_COLLECTION = "system.indexes"
     SYSTEM_PROFILE_COLLECTION = "system.profile"
     SYSTEM_USER_COLLECTION = "system.users"
+    SYSTEM_JS_COLLECTION = "system.js"
     SYSTEM_COMMAND_COLLECTION = "$cmd"
 
     # Counter for generating unique request ids.
@@ -106,6 +107,36 @@ def authenticate(username, password, save_auth=true)
       end
     end
 
+    # Adds a stored Javascript function to the database which can executed  
+    # server-side in map_reduce, db.eval and $where clauses.
+    # 
+    # @param [String] function_name
+    # @param [String] code
+    #
+    # @return [String] the function name saved to the database
+    def add_stored_function(function_name, code)
+      self[SYSTEM_JS_COLLECTION].save(
+        {
+          "_id" => function_name, 
+          :value => BSON::Code.new(code)
+        }
+      )
+    end
+    
+    # Removes stored Javascript function from the database.  Returns
+    # false if the function does not exist
+    #
+    # @param [String] function_name
+    #
+    # @return [Boolean]
+    def remove_stored_function(function_name)
+      if self[SYSTEM_JS_COLLECTION].find_one({"_id" => function_name})
+        self[SYSTEM_JS_COLLECTION].remove({"_id" => function_name}, :safe => true)
+      else
+        return false
+      end
+    end
+    
     # Adds a user to this database for use with authentication. If the user already
     # exists in the system, the password will be updated.
     #
diff --git a/test/db_test.rb b/test/db_test.rb
index 62cf462bcb..dfeacaa828 100644
--- a/test/db_test.rb
+++ b/test/db_test.rb
@@ -247,6 +247,14 @@ def test_remove_non_existant_user
     assert !@@db.remove_user("joe")
   end
 
+  def test_stored_function_management
+    @@db.add_stored_function("sum", "function (x, y) { return x + y; }")
+    assert_equal @@db.eval("return sum(2,3);"), 5
+    assert @@db.remove_stored_function("sum")
+    assert_raise OperationFailure do 
+      @@db.eval("return sum(2,3);")
+    end
+  end
 
   if @@version >= "1.3.5"
     def test_db_stats