From 047fc60c91d5f42234ee097b6e85402752bfd60a Mon Sep 17 00:00:00 2001
From: Mike Dirolf <mike@10gen.com>
Date: Mon, 17 Aug 2009 11:11:03 -0400
Subject: [PATCH] allow [] on DB, Collection instances to get
 collections/sub-collections

---
 lib/mongo/collection.rb  | 11 +++++++++++
 lib/mongo/db.rb          |  1 +
 tests/test_collection.rb | 14 ++++++++++++++
 3 files changed, 26 insertions(+)

diff --git a/lib/mongo/collection.rb b/lib/mongo/collection.rb
index a803312..eee7756 100644
--- a/lib/mongo/collection.rb
+++ b/lib/mongo/collection.rb
@@ -48,6 +48,17 @@ module XGen
           @hint = nil
         end
 
+        # Get a sub-collection of this collection by name.
+        #
+        # Raises InvalidName if an invalid collection name is used.
+        #
+        # :name :: the name of the collection to get
+        def [](name)
+          name = "#{self.name}.#{name}"
+          return Collection.new(self, name) if !db.strict? || db.collection_names.include?(name)
+          raise "Collection #{name} doesn't exist. Currently in strict mode."
+        end
+
         # Set hint fields to use and return +self+. hint may be a single field
         # name, array of field names, or a hash (preferably an OrderedHash).
         # May be +nil+.
diff --git a/lib/mongo/db.rb b/lib/mongo/db.rb
index 1053155..efe2c07 100644
--- a/lib/mongo/db.rb
+++ b/lib/mongo/db.rb
@@ -246,6 +246,7 @@ module XGen
           return Collection.new(self, name) if !strict? || collection_names.include?(name)
           raise "Collection #{name} doesn't exist. Currently in strict mode."
         end
+        alias_method :[], :collection
 
         # Drop collection +name+. Returns +true+ on success or if the
         # collection does not exist, +false+ otherwise.
diff --git a/tests/test_collection.rb b/tests/test_collection.rb
index 2549c69..698fbd7 100644
--- a/tests/test_collection.rb
+++ b/tests/test_collection.rb
@@ -31,6 +31,20 @@ class TestCollection < Test::Unit::TestCase
     @@test.drop()
   end
 
+  def test_collection
+    assert_raise InvalidName do
+      @@db["te$t"]
+    end
+
+    assert_kind_of Collection, @@db["test"]
+    assert_equal @@db["test"].name(), @@db.collection("test").name()
+    assert_equal @@db["test"].name(), @@db[:test].name()
+
+    assert_kind_of Collection, @@db["test"]["foo"]
+    assert_equal @@db["test"]["foo"].name(), @@db.collection("test.foo").name()
+    assert_equal @@db["test"]["foo"].name(), @@db["test.foo"].name()
+  end
+
   def test_safe_insert
     a = {"hello" => "world"}
     @@test.insert(a)