Skip to content

Commit

Permalink
Allow the setting of safe mode globally on the Connection,
Browse files Browse the repository at this point in the history
DB, and Collection levels. The safe mode setting will
automatically be inherited down the hierarchy Connection ->
DB -> Collection -> (insert, update, remove). This default
can be overridden at any time. Connection#safe, DB#safe, and
Collection#safe will yield the current default value.
  • Loading branch information
banker committed Nov 3, 2010
1 parent f7d151c commit 68af3db
Show file tree
Hide file tree
Showing 7 changed files with 224 additions and 20 deletions.
40 changes: 27 additions & 13 deletions lib/mongo/collection.rb
Original file line number Diff line number Diff line change
Expand Up @@ -21,13 +21,18 @@ module Mongo
# A named collection of documents in a database.
class Collection

attr_reader :db, :name, :pk_factory, :hint
attr_reader :db, :name, :pk_factory, :hint, :safe

# Initialize a collection object.
#
# @param [DB] db a MongoDB database instance.
# @param [String, Symbol] name the name of the collection.
#
# @option options [Boolean, Hash] :safe (false) Set the default safe-mode options
# for insert, update, and remove method called on this Collection instance. If no
# value is provided, the default value set on this instance's DB will be used. This
# default can be overridden for any invocation of insert, update, or remove.
#
# @raise [InvalidNSName]
# if collection name is empty, contains '$', or starts or ends with '.'
#
Expand All @@ -37,7 +42,7 @@ class Collection
# @return [Collection]
#
# @core collections constructor_details
def initialize(db, name, pk_factory=nil)
def initialize(db, name, pk_factory=nil, options={})
case name
when Symbol, String
else
Expand All @@ -60,6 +65,7 @@ def initialize(db, name, pk_factory=nil)
@connection = @db.connection
@logger = @connection.logger
@pk_factory = pk_factory || BSON::ObjectId
@safe = options.has_key?(:safe) ? options[:safe] : @db.safe
@hint = nil
end

Expand Down Expand Up @@ -245,16 +251,19 @@ def save(doc, opts={})
# @option opts [Boolean, Hash] :safe (+false+)
# run the operation in safe mode, which run a getlasterror command on the
# database to report any assertion. In addition, a hash can be provided to
# run an fsync and/or wait for replication of the insert (>= 1.5.1). See the options
# for DB#error.
# run an fsync and/or wait for replication of the insert (>= 1.5.1). Safe
# options provided here will override any safe options set on this collection,
# its database object, or the current connection. See the options on
# for DB#get_last_error.
#
# @see DB#remove for options that can be passed to :safe.
#
# @core insert insert-instance_method
def insert(doc_or_docs, options={})
doc_or_docs = [doc_or_docs] unless doc_or_docs.is_a?(Array)
doc_or_docs.collect! { |doc| @pk_factory.create_pk(doc) }
result = insert_documents(doc_or_docs, @name, true, options[:safe])
safe = options.has_key?(:safe) ? options[:safe] : @safe
result = insert_documents(doc_or_docs, @name, true, safe)
result.size > 1 ? result : result.first
end
alias_method :<<, :insert
Expand All @@ -265,10 +274,11 @@ def insert(doc_or_docs, options={})
# If specified, only matching documents will be removed.
#
# @option opts [Boolean, Hash] :safe (+false+)
# run the operation in safe mode, which run a getlasterror command on the
# run the operation in safe mode, which will run a getlasterror command on the
# database to report any assertion. In addition, a hash can be provided to
# run an fsync and/or wait for replication of the remove (>= 1.5.1). See the options
# for DB#get_last_error.
# run an fsync and/or wait for replication of the remove (>= 1.5.1). Safe
# options provided here will override any safe options set on this collection,
# its database, or the current connection. See the options for DB#get_last_error for more details.
#
# @example remove all documents from the 'users' collection:
# users.remove
Expand All @@ -287,14 +297,15 @@ def insert(doc_or_docs, options={})
# @core remove remove-instance_method
def remove(selector={}, opts={})
# Initial byte is 0.
safe = opts.has_key?(:safe) ? opts[:safe] : @safe
message = BSON::ByteBuffer.new("\0\0\0\0")
BSON::BSON_RUBY.serialize_cstr(message, "#{@db.name}.#{@name}")
message.put_int(0)
message.put_binary(BSON::BSON_CODER.serialize(selector, false, true).to_s)

@logger.debug("MONGODB #{@db.name}['#{@name}'].remove(#{selector.inspect})") if @logger
if opts[:safe]
@connection.send_message_with_safe_check(Mongo::Constants::OP_DELETE, message, @db.name, nil, opts[:safe])
if safe
@connection.send_message_with_safe_check(Mongo::Constants::OP_DELETE, message, @db.name, nil, safe)
# the return value of send_message_with_safe_check isn't actually meaningful --
# only the fact that it didn't raise an error is -- so just return true
true
Expand All @@ -320,11 +331,14 @@ def remove(selector={}, opts={})
# @option opts [Boolean] :safe (+false+)
# If true, check that the save succeeded. OperationFailure
# will be raised on an error. Note that a safe check requires an extra
# round-trip to the database.
# round-trip to the database. Safe options provided here will override any safe
# options set on this collection, its database object, or the current collection.
# See the options for DB#get_last_error for details.
#
# @core update update-instance_method
def update(selector, document, options={})
# Initial byte is 0.
safe = options.has_key?(:safe) ? options[:safe] : @safe
message = BSON::ByteBuffer.new("\0\0\0\0")
BSON::BSON_RUBY.serialize_cstr(message, "#{@db.name}.#{@name}")
update_options = 0
Expand All @@ -334,8 +348,8 @@ def update(selector, document, options={})
message.put_binary(BSON::BSON_CODER.serialize(selector, false, true).to_s)
message.put_binary(BSON::BSON_CODER.serialize(document, false, true).to_s)
@logger.debug("MONGODB #{@db.name}['#{@name}'].update(#{selector.inspect}, #{document.inspect})") if @logger
if options[:safe]
@connection.send_message_with_safe_check(Mongo::Constants::OP_UPDATE, message, @db.name, nil, options[:safe])
if safe
@connection.send_message_with_safe_check(Mongo::Constants::OP_UPDATE, message, @db.name, nil, safe)
else
@connection.send_message(Mongo::Constants::OP_UPDATE, message, nil)
end
Expand Down
12 changes: 10 additions & 2 deletions lib/mongo/connection.rb
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,8 @@ class Connection
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]"

attr_reader :logger, :size, :host, :port, :nodes, :auths, :sockets, :checked_out, :primary, :secondaries, :arbiters
attr_reader :logger, :size, :host, :port, :nodes, :auths, :sockets, :checked_out, :primary, :secondaries, :arbiters,
:safe

# Counter for generating unique request ids.
@@current_request_id = 0
Expand All @@ -61,6 +62,10 @@ class Connection
# @param [String, Hash] host.
# @param [Integer] port specify a port number here if only one host is being specified.
#
# @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
# on initialization.
# @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.
Expand Down Expand Up @@ -114,6 +119,9 @@ def initialize(host=nil, port=nil, options={})
# Mutex for synchronizing pool access
@connection_mutex = Mutex.new

# Global safe option. This is false by default.
@safe = options[:safe] || false

# Create a mutex when a new key, in this case a socket,
# is added to the hash.
@safe_mutexes = Hash.new { |h, k| h[k] = Mutex.new }
Expand Down Expand Up @@ -312,7 +320,7 @@ def db(db_name, options={})
#
# @core databases []-instance_method
def [](db_name)
DB.new(db_name, self)
DB.new(db_name, self, :safe => @safe)
end

# Drop a database.
Expand Down
21 changes: 16 additions & 5 deletions lib/mongo/db.rb
Original file line number Diff line number Diff line change
Expand Up @@ -45,8 +45,8 @@ class DB
# Returns the value of the +strict+ flag.
def strict?; @strict; end

# The name of the database.
attr_reader :name
# The name of the database and the local safe option.
attr_reader :name, :safe

# The Mongo::Connection instance connecting to the MongoDB server.
attr_reader :connection
Expand All @@ -65,12 +65,19 @@ def strict?; @strict; end
# fields the factory wishes to inject. (NOTE: if the object already has a primary key,
# the factory should not inject a new key).
#
# @option options [Boolean, Hash] :safe (false) Set the default safe-mode options
# propogated to Collection objects instantiated off of this DB. If no
# value is provided, the default value set on this instance's Connection object will be used. This
# default can be overridden upon instantiation of any collection by explicity setting a :safe value
# on initialization
#
# @core databases constructor_details
def initialize(db_name, connection, options={})
@name = Mongo::Support.validate_db_name(db_name)
@connection = connection
@strict = options[:strict]
@pk_factory = options[:pk]
@safe = options.has_key?(:safe) ? options[:safe] : @connection.safe
end

# Authenticate with the given username and password. Note that mongod
Expand Down Expand Up @@ -259,9 +266,13 @@ def create_collection(name, options={})
# @raise [MongoDBError] if collection does not already exist and we're in +strict+ mode.
#
# @return [Mongo::Collection]
def collection(name)
return Collection.new(self, name, @pk_factory) if !strict? || collection_names.include?(name)
raise Mongo::MongoDBError, "Collection #{name} doesn't exist. Currently in strict mode."
def collection(name, options={})
if strict? && !collection_names.include?(name)
raise Mongo::MongoDBError, "Collection #{name} doesn't exist. Currently in strict mode."
else
options[:safe] = options.has_key?(:safe) ? options[:safe] : @safe
Collection.new(self, name, @pk_factory, options)
end
end
alias_method :[], :collection

Expand Down
42 changes: 42 additions & 0 deletions test/safe_test.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
require './test/test_helper'
include Mongo

class SafeTest < Test::Unit::TestCase
context "Safe tests: " do
setup do
@con = standard_connection(:safe => {:w => 1})
@db = @con[MONGO_TEST_DB]
@col = @db['test-safe']
@col.create_index([[:a, 1]], :unique => true)
@col.remove
end

should "propogate safe option on insert" do
@col.insert({:a => 1})

assert_raise_error(OperationFailure, "duplicate key") do
@col.insert({:a => 1})
end
end

should "allow safe override on insert" do
@col.insert({:a => 1})
@col.insert({:a => 1}, :safe => false)
end

should "propogate safe option on update" do
@col.insert({:a => 1})
@col.insert({:a => 2})

assert_raise_error(OperationFailure, "duplicate key") do
@col.update({:a => 2}, {:a => 1})
end
end

should "allow safe override on update" do
@col.insert({:a => 1})
@col.insert({:a => 2})
@col.update({:a => 2}, {:a => 1}, :safe => false)
end
end
end
2 changes: 2 additions & 0 deletions test/unit/db_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,9 @@ def insert_message(db, documents)
context "DB commands" do
setup do
@conn = stub()
@conn.stubs(:safe)
@db = DB.new("testing", @conn)
@db.stubs(:safe)
@collection = mock()
@db.stubs(:system_command_collection).returns(@collection)
end
Expand Down
2 changes: 2 additions & 0 deletions test/unit/grid_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,14 @@ class GridTest < Test::Unit::TestCase
context "GridFS: " do
setup do
@conn = stub()
@conn.stubs(:safe)
@db = DB.new("testing", @conn)
@files = mock()
@chunks = mock()

@db.expects(:[]).with('fs.files').returns(@files)
@db.expects(:[]).with('fs.chunks').returns(@chunks)
@db.stubs(:safe)
end

context "Grid classe with standard connections" do
Expand Down
Loading

0 comments on commit 68af3db

Please sign in to comment.