From 103d7e5c62eac58466b33a24426f39183e1062a2 Mon Sep 17 00:00:00 2001
From: Jim Menard <jim@10gen.com>
Date: Tue, 16 Dec 2008 17:08:15 -0500
Subject: [PATCH] New optional "strict" mode for databases.

---
 README.rdoc          |  4 ++++
 lib/mongo/db.rb      | 31 +++++++++++++++++++++++++++----
 tests/test_db_api.rb | 43 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 74 insertions(+), 4 deletions(-)

diff --git a/README.rdoc b/README.rdoc
index f283ce0..8f4922b 100644
--- a/README.rdoc
+++ b/README.rdoc
@@ -39,6 +39,8 @@ http://github.com/geir/mongo-java-driver/tree/master.
 
 == Release Notes
 
+Added "strict" db attribute.
+
 I plan to remove the auto-generation of _id primary keys.
 
 If you ran tests using code before release
@@ -52,6 +54,8 @@ type
 
 = To Do
 
+* Add way to ask db if it is master or slave.
+
 * Tests that prove that this driver's ObjectID and Geir's Java version do the
   same thing. (I've done so manually.)
 
diff --git a/lib/mongo/db.rb b/lib/mongo/db.rb
index 88a5fb7..0971b34 100644
--- a/lib/mongo/db.rb
+++ b/lib/mongo/db.rb
@@ -27,12 +27,20 @@ module XGen
         SYSTEM_INDEX_COLLECTION = "system.indexes"
         SYSTEM_COMMAND_COLLECTION = "$cmd"
 
+        # Strict mode means that trying to access a collection that does not
+        # exist will raise an error. Strict mode is off (false) by default.
+        attr_writer :strict
+
+        # Returns the value of the +strict+ flag.
+        def strict?; @strict; end
+
         attr_reader :name, :socket
 
         def initialize(db_name, host, port)
           raise "Invalid DB name" if !db_name || (db_name && db_name.length > 0 && db_name.include?("."))
           @name, @host, @port = db_name, host, port
           @socket = TCPSocket.new(@host, @port)
+          @strict = false
         end
 
         def collection_names
@@ -47,9 +55,18 @@ module XGen
           query(SYSTEM_NAMESPACE_COLLECTION, Query.new(selector))
         end
 
+        # Create a collection. If +strict+ is false, will return existing or
+        # new collection. If +strict+ is true, will raise an error if
+        # collection +name+ does not already exist.
         def create_collection(name, options={})
           # First check existence
-          return Collection.new(self, name) if collection_names.include?(name)
+          if collection_names.include?(full_coll_name(name))
+            if strict?
+              raise "Collection #{name} already exists. Currently in strict mode."
+            else
+              return Collection.new(self, name)
+            end
+          end
 
           # Create new collection
           oh = OrderedHash.new
@@ -66,10 +83,16 @@ module XGen
           Admin.new(self)
         end
 
+        # Return a collection. If +strict+ is false, will return existing or
+        # new collection. If +strict+ is true, will raise an error if
+        # collection +name+ already exists.
         def collection(name)
-          # We do not implement the Java driver's optional strict mode, which
-          # throws an exception if the collection does not exist.
-          create_collection(name)
+          return Collection.new(self, name) if collection_names.include?(full_coll_name(name))
+          if strict?
+            raise "Collection #{name} doesn't exist. Currently in strict mode."
+          else
+            create_collection(name)
+          end
         end
 
         def drop_collection(name)
diff --git a/tests/test_db_api.rb b/tests/test_db_api.rb
index ae6d818..6a0da23 100644
--- a/tests/test_db_api.rb
+++ b/tests/test_db_api.rb
@@ -203,4 +203,47 @@ class DBAPITest < Test::Unit::TestCase
     assert_equal 1, rows.length
     assert_equal regex, rows[0]['b']
   end
+
+  def test_strict
+    assert !@db.strict?
+    @db.strict = true
+    assert @db.strict?
+  end
+
+  def test_strict_access_collection
+    @db.strict = true
+    begin
+      @db.collection('does-not-exist')
+      fail "expected exception"
+    rescue => ex
+      assert_equal "Collection does-not-exist doesn't exist. Currently in strict mode.", ex.to_s
+    ensure
+      @db.strict = false
+      @db.drop_collection('does-not-exist')
+    end
+  end
+
+  def test_strict_create_collection
+    @db.drop_collection('foobar')
+    @db.strict = true
+
+    begin
+      @db.create_collection('foobar')
+      assert true
+    rescue => ex
+      fail "did not expect exception \"#{ex}\""
+    end
+
+    # Now the collection exists. This time we should see an exception.
+    begin
+      @db.create_collection('foobar')
+      fail "expected exception"
+    rescue => ex
+      assert_equal "Collection foobar already exists. Currently in strict mode.", ex.to_s
+    ensure
+      @db.strict = false
+      @db.drop_collection('foobar')
+    end
+  end
+
 end