Allow specification of replica set name on connect.

Raise ReplicaSetConnectionError if expected name doesn't match set.

lib/mongo/connection.rb

@@ -64,6 +64,8 @@ module Mongo
     # @option options [Boolean] :slave_ok (false) Must be set to +true+ when connecting
     #   to a single, slave node.
     # @option options [Logger, #debug] :logger (nil) Logger instance to receive driver operation log.
+    # @option options [String] :name (nil) The name of the replica set to connect to. An exception will be
+    #   raised if unable to connect to a replica set with this name.
     # @option options [Integer] :pool_size (1) The maximum number of socket connections that can be
     #   opened to the database.
     # @option options [Float] :timeout (5.0) When all of the connections to the pool are checked out,
@@ -83,6 +85,9 @@ module Mongo
     #
     # @see http://www.mongodb.org/display/DOCS/Replica+Pairs+in+Ruby Replica pairs in Ruby
     #
+    # @raise [ReplicaSetConnectionError] This is raised if a replica set name is specified and the
+    #   driver fails to connect to a replica set with that name.
+    #
     # @core connections
     def initialize(host=nil, port=nil, options={})
       @auths        = []
@@ -96,6 +101,9 @@ module Mongo
       # Host and port of current master.
       @host = @port = nil
 
+      # Replica set name
+      @replica_set_name = options[:name]
+
       # Lock for request ids.
       @id_lock = Mutex.new
 
@@ -600,6 +608,7 @@ module Mongo
 
         config = self['admin'].command({:ismaster => 1}, :sock => socket)
 
+        check_set_name(config, socket)
       rescue OperationFailure, SocketError, SystemCallError, IOError => ex
         close unless connected?
       ensure
@@ -617,6 +626,21 @@ module Mongo
       config
     end
 
+    # Make sure that we're connected to the expected replica set.
+    def check_set_name(config, socket)
+      if @replica_set_name
+        config = self['admin'].command({:replSetGetStatus => 1},
+                   :sock => socket, :check_response => false)
+
+        if !Mongo::Support.ok?(config)
+          raise ReplicaSetConnectionError, config['errmsg']
+        elsif config['set'] != @replica_set_name
+          raise ReplicaSetConnectionError,
+            "Attempting to connect to replica set '#{config['set']}' but expected '#{@replica_set_name}'"
+        end
+      end
+    end
+
     # Set the specified node as primary, and
     # apply any saved authentication credentials.
     def set_primary(node)

lib/mongo/exceptions.rb

@@ -42,6 +42,9 @@ module Mongo
   # Raised on failures in connection to the database server.
   class ConnectionError < MongoRubyError; end
 
+  # Raised on failures in connection to the database server.
+  class ReplicaSetConnectionError < ConnectionError; end
+
   # Raised on failures in connection to the database server.
   class ConnectionTimeoutError < MongoRubyError; end
 

test/connection_test.rb

@@ -43,6 +43,12 @@ class TestConnection < Test::Unit::TestCase
     assert_raise Mongo::InvalidNSName do @conn.db('te st') end
   end
 
+  def test_replica_set_connection_name
+    assert_raise_error(Mongo::ReplicaSetConnectionError, "replSet") do
+      standard_connection(:name => "replica-set-foo")
+    end
+  end
+
   def test_options_passed_to_db
     @pk_mock = Object.new
     db = @conn.db('test', :pk => @pk_mock, :strict => true)

test/replica_sets/connect_test.rb

@@ -7,8 +7,16 @@ require './test/test_helper'
 class ConnectTest < Test::Unit::TestCase
   include Mongo
 
+  def test_connect_bad_name
+    assert_raise_error(ReplicaSetConnectionError, "expected 'wrong-repl-set-name'") do
+      Mongo::Connection.multi([['localhost', 27017], ['localhost', 27018], ['localhost', 27019]],
+        :name => "wrong-repl-set-name")
+    end
+  end
+
   def test_connect
-    @conn = Mongo::Connection.multi([['localhost', 27017], ['localhost', 27018], ['localhost', 27019]])
+    @conn = Mongo::Connection.multi([['localhost', 27017], ['localhost', 27018], ['localhost', 27019]],
+                                    :name => "foo")
     assert @conn.connected?
   end