RUBY-205 RUBY-150 Support new connection URI options

lib/mongo/connection.rb

@@ -138,15 +138,19 @@ module Mongo
     #
     # @param opts Any of the options available for Connection.new
     #
-    # @return [Mongo::Connection]
-    def self.from_uri(uri, opts={})
-      nodes, auths = Mongo::URIParser.parse(uri)
-      opts.merge!({:auths => auths})
-      if nodes.length == 1
-        Connection.new(nodes[0][0], nodes[0][1], opts)
-      elsif nodes.length > 1
-        nodes << opts
-        ReplSetConnection.new(*nodes)
+    # @return [Mongo::Connection, Mongo::ReplSetConnection]
+    def self.from_uri(string, extra_opts={})
+      uri = URIParser.new(string)
+      opts = uri.connection_options
+      opts.merge!(extra_opts)
+
+      if uri.nodes.length == 1
+        opts.merge!({:auths => uri.auths})
+        Connection.new(uri.nodes[0][0], uri.nodes[0][1], opts)
+      elsif uri.nodes.length > 1
+        nodes = uri.nodes.clone
+        nodes_with_opts = nodes << opts
+        ReplSetConnection.new(*nodes_with_opts)
       else
         raise MongoArgumentError, "No nodes specified. Please ensure that you've provided at least one node."
       end

lib/mongo/repl_set_connection.rb

@@ -32,6 +32,8 @@ module Mongo
     # @param [Array] args A list of host-port pairs ending with a hash containing any options. See
     #   the examples below for exactly how to use the constructor.
     #
+    # @option options [String] :rs_name (nil) The name of the replica set to connect to. You
+    #   can use this option to verify that you're connecting to the right replica set.
     # @option options [Boolean, Hash] :safe (false) Set the default safe-mode options
     #   propogated to DB objects instantiated off of this Connection. This
     #   default can be overridden upon instantiation of any DB by explicity setting a :safe value

lib/mongo/util/uri_parser.rb

@@ -17,28 +17,99 @@
 # ++
 
 module Mongo
-  module URIParser
+  class URIParser
 
     DEFAULT_PORT = 27017
     MONGODB_URI_MATCHER = /(([-_.\w\d]+):([-_\w\d]+)@)?([-.\w\d]+)(:([\w\d]+))?(\/([-\d\w]+))?/
     MONGODB_URI_SPEC = "mongodb://[username:password@]host1[:port1][,host2[:port2],...[,hostN[:portN]]][/database]"
+    SPEC_ATTRS = [:nodes, :auths]
+    OPT_ATTRS  = [:connect, :replicaset, :slaveok, :safe, :w, :wtimeout, :fsync]
 
-    extend self
+    OPT_VALID  = {:connect    => lambda {|arg| ['direct', 'replicaset'].include?(arg)},
+                  :replicaset => lambda {|arg| arg.length > 0},
+                  :slaveok    => lambda {|arg| ['true', 'false'].include?(arg)},
+                  :safe       => lambda {|arg| ['true', 'false'].include?(arg)},
+                  :w          => lambda {|arg| arg =~ /^\d+$/ },
+                  :wtimeout   => lambda {|arg| arg =~ /^\d+$/ },
+                  :fsync      => lambda {|arg| ['true', 'false'].include?(arg)}
+                 }
+
+    OPT_ERR    = {:connect    => "must be 'direct' or 'replicaset'",
+                  :replicaset => "must be a string containing the name of the replica set to connect to",
+                  :slaveok    => "must be 'true' or 'false'",
+                  :safe       => "must be 'true' or 'false'",
+                  :w          => "must be an integer specifying number of nodes to replica to",
+                  :wtimeout   => "must be an integer specifying milliseconds",
+                  :fsync      => "must be 'true' or 'false'"
+                 }
+
+    OPT_CONV   = {:connect    => lambda {|arg| arg},
+                  :replicaset => lambda {|arg| arg},
+                  :slaveok    => lambda {|arg| arg == 'true' ? true : false},
+                  :safe       => lambda {|arg| arg == 'true' ? true : false},
+                  :w          => lambda {|arg| arg.to_i},
+                  :wtimeout   => lambda {|arg| arg.to_i},
+                  :fsync      => lambda {|arg| arg == 'true' ? true : false}
+                 }
+
+    ATTRS      = SPEC_ATTRS + OPT_ATTRS
+
+    attr_reader *ATTRS
 
     # Parse a MongoDB URI. This method is used by Connection.from_uri.
     # Returns an array of nodes and an array of db authorizations, if applicable.
-    #
-    # @private
-    def parse(string)
+    def initialize(string)
       if string =~ /^mongodb:\/\//
         string = string[10..-1]
       else
         raise MongoArgumentError, "MongoDB URI must match this spec: #{MONGODB_URI_SPEC}"
       end
 
-      nodes = []
-      auths = []
-      specs = string.split(',')
+      hosts, opts = string.split('?')
+      parse_hosts(hosts)
+      parse_options(opts)
+      configure_connect
+    end
+
+    def connection_options
+      opts = {}
+
+      if (@w || @wtimeout || @fsync) && !@safe
+        raise MongoArgumentError, "Safe must be true if w, wtimeout, or fsync is specified"
+      end
+
+      if @safe
+        if @w || @wtimeout || @fsync
+          safe_opts = {}
+          safe_opts[:w] = @w if @w
+          safe_opts[:wtimeout] = @wtimeout if @wtimeout
+          safe_opts[:fsync] = @fsync if @fsync
+        else
+          safe_opts = true
+        end
+
+        opts[:safe] = safe_opts
+      end
+
+      if @slaveok
+        if @connect == 'direct'
+          opts[:slave_ok] = true
+        else
+          opts[:read_secondary] = true
+        end
+      end
+
+      opts[:rs_name] = @replicaset if @replicaset
+
+      opts
+    end
+
+    private
+
+    def parse_hosts(hosts)
+      @nodes = []
+      @auths = []
+      specs = hosts.split(',')
       specs.each do |spec|
         matches  = MONGODB_URI_MATCHER.match(spec)
         if !matches
@@ -52,8 +123,8 @@ module Mongo
         if !(port.to_s =~ /^\d+$/)
           raise MongoArgumentError, "Invalid port #{port}; port must be specified as digits."
         end
-        port  = port.to_i
-        db    = matches[8]
+        port    = port.to_i
+        db      = matches[8]
 
         if uname && pwd && db
           auths << {'db_name' => db, 'username' => uname, 'password' => pwd}
@@ -62,10 +133,47 @@ module Mongo
             "and db if any one of these is specified."
         end
 
-        nodes << [host, port]
+        @nodes << [host, port]
+      end
+    end
+
+    # This method uses the lambdas defined in OPT_VALID and OPT_CONV to validate
+    # and convert the given options.
+    def parse_options(opts)
+      return unless opts
+      separator = opts.include?('&') ? '&' : ';'
+      opts.split(separator).each do |attr|
+        key, value = attr.split('=')
+        key   = key.to_sym
+        value = value.strip.downcase
+        if !OPT_ATTRS.include?(key)
+          raise MongoArgumentError, "Invalid Mongo URI option #{key}"
+        end
+
+        if OPT_VALID[key].call(value)
+          instance_variable_set("@#{key}", OPT_CONV[key].call(value))
+        else
+          raise MongoArgumentError, "Invalid value for #{key}: #{OPT_ERR[key]}"
+        end
+      end
+    end
+
+    def configure_connect
+      if @nodes.length > 1 && !@connect
+        @connect = 'replicaset'
       end
 
-      [nodes, auths]
+      if !@connect
+        if @nodes.length > 1
+          @connect = 'replicaset'
+        else
+          @connect = 'direct'
+        end
+      end
+
+      if @connect == 'direct' && @replicaset
+        raise MongoArgumentError, "If specifying a replica set name, please also specify that connect=replicaset"
+      end
     end
   end
 end

test/replica_sets/connect_test.rb

@@ -16,8 +16,8 @@ class ConnectTest < Test::Unit::TestCase
 
   def test_connect_with_deprecated_multi
     @conn = Connection.multi([[RS.host, RS.ports[0]], [RS.host, RS.ports[1]]], :name => RS.name)
-    assert @conn.connected?
     assert @conn.is_a?(ReplSetConnection)
+    assert @conn.connected?
   end
 
   def test_connect_bad_name

test/replica_sets/connection_string_test.rb

@@ -0,0 +1,32 @@
+$:.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
+require './test/replica_sets/rs_test_helper'
+
+# NOTE: This test expects a replica set of three nodes to be running on RS.host,
+# on ports TEST_PORT, RS.ports[1], and TEST + 2.
+class ConnectionStringTest < Test::Unit::TestCase
+  include Mongo
+
+  def setup
+    RS.restart_killed_nodes
+  end
+
+  def teardown
+    RS.restart_killed_nodes
+  end
+
+  def test_connect_with_connection_string
+    @conn = Connection.from_uri("mongodb://#{RS.host}:#{RS.ports[0]},#{RS.host}:#{RS.ports[1]}?replicaset=#{RS.name}")
+    assert @conn.is_a?(ReplSetConnection)
+    assert @conn.connected?
+  end
+
+  def test_connect_with_full_connection_string
+    @conn = Connection.from_uri("mongodb://#{RS.host}:#{RS.ports[0]},#{RS.host}:#{RS.ports[1]}?replicaset=#{RS.name};safe=true;w=2;fsync=true;slaveok=true")
+    assert @conn.is_a?(ReplSetConnection)
+    assert @conn.connected?
+    assert_equal 2, @conn.safe[:w]
+    assert @conn.safe[:fsync]
+    assert @conn.read_pool
+  end
+
+end

test/unit/connection_test.rb

@@ -22,8 +22,8 @@ class ConnectionTest < Test::Unit::TestCase
         TCPSocket.stubs(:new).returns(new_mock_socket)
 
         admin_db = new_mock_db
-        admin_db.expects(:command).returns({'ok' => 1, 'ismaster' => 1})
-        @conn.expects(:[]).with('admin').returns(admin_db)
+        admin_db.expects(:command).returns({'ok' => 1, 'ismaster' => 1}).twice
+        @conn.expects(:[]).with('admin').returns(admin_db).twice
         @conn.connect
       end
 
@@ -65,8 +65,8 @@ class ConnectionTest < Test::Unit::TestCase
         @conn = Connection.from_uri("mongodb://localhost", :connect => false)
 
         admin_db = new_mock_db
-        admin_db.expects(:command).returns({'ok' => 1, 'ismaster' => 1})
-        @conn.expects(:[]).with('admin').returns(admin_db)
+        admin_db.expects(:command).returns({'ok' => 1, 'ismaster' => 1}).twice
+        @conn.expects(:[]).with('admin').returns(admin_db).twice
         @conn.expects(:apply_saved_authentication)
         @conn.connect
       end

test/unit/repl_set_connection_test.rb

@@ -67,16 +67,6 @@ class ReplSetConnectionTest < Test::Unit::TestCase
         assert_equal ['localhost', 27017], @conn.nodes[0]
         assert_equal ['mydb.com', 27018], @conn.nodes[1]
       end
-
-      should "parse a uri specifying multiple nodes with auth" do
-        @conn = Connection.from_uri("mongodb://kyle:s3cr3t@localhost:27017/app,mickey:m0u5e@mydb.com:27018/dsny", :connect => false)
-        assert_equal ['localhost', 27017], @conn.nodes[0]
-        assert_equal ['mydb.com', 27018], @conn.nodes[1]
-        auth_hash = {'username' => 'kyle', 'password' => 's3cr3t', 'db_name' => 'app'}
-        assert_equal auth_hash, @conn.auths[0]
-        auth_hash = {'username' => 'mickey', 'password' => 'm0u5e', 'db_name' => 'dsny'}
-        assert_equal auth_hash, @conn.auths[1]
-      end
     end
   end
 end

test/uri_test.rb

@@ -0,0 +1,75 @@
+require './test/test_helper'
+
+class TestThreading < Test::Unit::TestCase
+  include Mongo
+
+  def test_uri_without_port
+    parser = Mongo::URIParser.new('mongodb://localhost')
+    assert_equal 1, parser.nodes.length
+    assert_equal 'localhost', parser.nodes[0][0]
+    assert_equal 27017, parser.nodes[0][1]
+  end
+
+  def test_basic_uri
+    parser = Mongo::URIParser.new('mongodb://localhost:27018')
+    assert_equal 1, parser.nodes.length
+    assert_equal 'localhost', parser.nodes[0][0]
+    assert_equal 27018, parser.nodes[0][1]
+  end
+
+  def test_multiple_uris
+    parser = Mongo::URIParser.new('mongodb://a.example.com:27018,b.example.com')
+    assert_equal 2, parser.nodes.length
+    assert_equal 'a.example.com', parser.nodes[0][0]
+    assert_equal 27018, parser.nodes[0][1]
+    assert_equal 'b.example.com', parser.nodes[1][0]
+    assert_equal 27017, parser.nodes[1][1]
+  end
+
+  def test_multiple_uris_with_auths
+    parser = Mongo::URIParser.new('mongodb://bob:secret@a.example.com:27018/test,joe:secret2@b.example.com/test2')
+    assert_equal 2, parser.nodes.length
+    assert_equal 'a.example.com', parser.nodes[0][0]
+    assert_equal 27018, parser.nodes[0][1]
+    assert_equal 'b.example.com', parser.nodes[1][0]
+    assert_equal 27017, parser.nodes[1][1]
+    assert_equal 2, parser.auths.length
+    assert_equal "bob", parser.auths[0]["username"]
+    assert_equal "secret", parser.auths[0]["password"]
+    assert_equal "test", parser.auths[0]["db_name"]
+    assert_equal "joe", parser.auths[1]["username"]
+    assert_equal "secret2", parser.auths[1]["password"]
+    assert_equal "test2", parser.auths[1]["db_name"]
+  end
+
+  def test_opts_basic
+    parser = Mongo::URIParser.new('mongodb://localhost:27018?connect=direct;slaveok=true;safe=true')
+    assert_equal 'direct', parser.connect
+    assert parser.slaveok
+    assert parser.safe
+  end
+
+  def test_opts_with_amp_separator
+    parser = Mongo::URIParser.new('mongodb://localhost:27018?connect=direct&slaveok=true&safe=true')
+    assert_equal 'direct', parser.connect
+    assert parser.slaveok
+    assert parser.safe
+  end
+
+  def test_opts_safe
+    parser = Mongo::URIParser.new('mongodb://localhost:27018?safe=true;w=2;wtimeout=200;fsync=true')
+    assert parser.safe
+    assert_equal 2, parser.w
+    assert_equal 200, parser.wtimeout
+    assert parser.fsync
+  end
+
+  def test_opts_replica_set
+    assert_raise_error MongoArgumentError, "specify that connect=replicaset" do
+      Mongo::URIParser.new('mongodb://localhost:27018?replicaset=foo')
+    end
+    parser = Mongo::URIParser.new('mongodb://localhost:27018?connect=replicaset;replicaset=foo')
+    assert_equal 'foo', parser.replicaset
+    assert_equal 'replicaset', parser.connect
+  end
+end