mongo 2.15.0 → 2.18.1

data/lib/mongo/operation/write_command/op_msg.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+# encoding: utf-8
+
+# Copyright (C) 2021 MongoDB Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+module Mongo
+  module Operation
+    class WriteCommand
+
+      # A MongoDB write command operation sent as an op message.
+      #
+      # @api private
+      class OpMsg < OpMsgBase
+        include Validatable
+
+        private
+
+        def selector(connection)
+          super.tap do |selector|
+            if selector.key?(:findAndModify)
+              validate_find_options(connection, selector)
+            end
+            if wc = spec[:write_concern]
+              selector[:writeConcern] = wc.options
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/mongo/operation/write_command.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+# encoding: utf-8
+
+# Copyright (C) 2021 MongoDB Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require 'mongo/operation/write_command/command'
+require 'mongo/operation/write_command/op_msg'
+
+module Mongo
+  module Operation
+
+    # A MongoDB general command operation.
+    #
+    # @api private
+    class WriteCommand
+      include Specifiable
+      include OpMsgOrCommand
+    end
+  end
+end

data/lib/mongo/operation.rb

@@ -20,12 +20,14 @@ require 'mongo/operation/shared/causal_consistency_supported'
 require 'mongo/operation/shared/write'
 require 'mongo/operation/shared/idable'
 require 'mongo/operation/shared/specifiable'
+require 'mongo/operation/shared/validatable'
 require 'mongo/operation/shared/object_id_generator'
 require 'mongo/operation/shared/op_msg_or_command'
 require 'mongo/operation/shared/op_msg_or_find_command'
 
 require 'mongo/operation/op_msg_base'
 require 'mongo/operation/command'
+require 'mongo/operation/write_command'
 require 'mongo/operation/aggregate'
 require 'mongo/operation/result'
 require 'mongo/operation/collections_info'
@@ -54,6 +56,14 @@ require 'mongo/operation/drop_index'
 
 module Mongo
 
+  # This module encapsulates all of the operation classes defined by the driver.
+  #
+  # The operation classes take Ruby options as constructor parameters.
+  # For example, :read contains read preference and :read_concern contains read
+  # concern, whereas server commands use readConcern field for the read
+  # concern and read preference is passed as $readPreference or secondaryOk
+  # wire protocol flag bit.
+  #
   # @api private
   module Operation
 

data/lib/mongo/protocol/caching_hash.rb

@@ -0,0 +1,69 @@
+# frozen_string_literal: true
+# encoding: utf-8
+
+# Copyright (C) 2014-2022 MongoDB Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+module Mongo
+  module Protocol
+
+    # A Hash that caches the results of #to_bson.
+    #
+    # @api private
+    class CachingHash
+
+      def initialize(hash)
+        @hash = hash
+      end
+
+      def bson_type
+        Hash::BSON_TYPE
+      end
+
+      # Caches the result of to_bson and writes it to the given buffer on subsequent
+      # calls to this method. If this method is originally called without validation,
+      # and then is subsequently called with validation, we will want to recalculate
+      # the to_bson to trigger the validations.
+      #
+      # @param [ BSON::ByteBuffer ] buffer The encoded BSON buffer to append to.
+      # @param [ true, false ] validating_keys Whether keys should be validated when serializing.
+      #
+      # @return [ BSON::ByteBuffer ] The buffer with the encoded object.
+      def to_bson(buffer = BSON::ByteBuffer.new, validating_keys = BSON::Config.validating_keys?)
+        if !@bytes
+          @bytes = @hash.to_bson(BSON::ByteBuffer.new, validating_keys).to_s
+        elsif needs_validation?(validating_keys)
+          @validated = true
+          return @hash.to_bson(buffer, validating_keys)
+        end
+        @validated ||= validating_keys
+        buffer.put_bytes(@bytes)
+      end
+
+      private
+
+      # Checks the current value for validating keys, and whether or not this
+      # bson has been validated in the past, and decides if we need to recalculate
+      # the to_bson to check the validations.
+      #
+      # @param [ true, false ] validating_keys Whether keys should be validated when serializing.
+      #
+      # @return [ true, false ] Whether or not the bson needs to be recalculated
+      #   with validation.
+      def needs_validation?(validating_keys)
+        !@validated && validating_keys
+      end
+    end
+  end
+end

data/lib/mongo/protocol/msg.rb

@@ -87,7 +87,9 @@ module Mongo
         ] + @sequences.map do |section|
           {type: 1, payload: {
             identifier: section.identifier,
-            sequence: section.documents,
+            sequence: section.documents.map do |doc|
+              CachingHash.new(doc)
+            end,
           }}
         end
         @request_id = nil
@@ -298,6 +300,24 @@ module Mongo
         Msg.new(@flags, @options, main_document, *@sequences)
       end
 
+      # Returns the number of documents returned from the server.
+      #
+      # The Msg instance must be for a server reply and the reply must return
+      # an active cursor (either a newly created one or one whose iteration is
+      # continuing via getMore).
+      #
+      # @return [ Integer ] Number of returned documents.
+      def number_returned
+        if doc = documents.first
+          if cursor = doc['cursor']
+            if batch = cursor['firstBatch'] || cursor['nextBatch']
+              return batch.length
+            end
+          end
+        end
+        raise NotImplementedError, "number_returned is only defined for cursor replies"
+      end
+
       private
 
       # Validate that the documents in this message are all smaller than the

data/lib/mongo/protocol/query.rb

@@ -64,13 +64,21 @@ module Mongo
       def initialize(database, collection, selector, options = {})
         @database = database
         @namespace = "#{database}.#{collection}"
+        if selector.nil?
+          raise ArgumentError, 'Selector cannot be nil'
+        end
         @selector = selector
         @options = options
         @project = options[:project]
         @limit = determine_limit
         @skip = options[:skip]  || 0
         @flags = options[:flags] || []
-        @upconverter = Upconverter.new(collection, selector, options, flags)
+        @upconverter = Upconverter.new(
+          collection,
+          BSON::Document.new(selector),
+          BSON::Document.new(options),
+          flags,
+        )
         super
       end
 
@@ -224,7 +232,7 @@ module Mongo
         }.freeze
 
         SPECIAL_FIELD_MAPPINGS = {
-          :$readPreference => 'readPreference',
+          :$readPreference => '$readPreference',
           :$orderby => 'sort',
           :$hint => 'hint',
           :$comment => 'comment',
@@ -249,16 +257,6 @@ module Mongo
           :partial => 'allowPartialResults'
         }.freeze
 
-        # Find command constant.
-        #
-        # @since 2.1.0
-        FIND = 'find'.freeze
-
-        # Filter attribute constant.
-        #
-        # @since 2.1.0
-        FILTER = 'filter'.freeze
-
         # @return [ String ] collection The name of the collection.
         attr_reader :collection
 
@@ -283,6 +281,15 @@ module Mongo
         #
         # @since 2.1.0
         def initialize(collection, filter, options, flags)
+          # Although the docstring claims both hashes and BSON::Documents
+          # are acceptable, this class expects the filter and options to
+          # contain symbol keys which isn't what the operation layer produces.
+          unless BSON::Document === filter
+            raise ArgumentError, 'Filter must provide indifferent access'
+          end
+          unless BSON::Document === options
+            raise ArgumentError, 'Options must provide indifferent access'
+          end
           @collection = collection
           @filter = filter
           @options = options
@@ -311,7 +318,7 @@ module Mongo
         #
         # @since 2.1.0
         def command_name
-          ((filter[:$query] || !command?) ? FIND : filter.keys.first).to_s
+          ((filter[:$query] || !command?) ? :find : filter.keys.first).to_s
         end
 
         private
@@ -333,15 +340,25 @@ module Mongo
         end
 
         def find_command
-          document = BSON::Document.new
-          document.store(FIND, collection)
-          document.store(FILTER, query_filter)
+          document = BSON::Document.new(
+            find: collection,
+            filter: query_filter,
+          )
           OPTION_MAPPINGS.each do |legacy, option|
             document.store(option, options[legacy]) unless options[legacy].nil?
           end
-          Mongo::Lint.validate_camel_case_read_preference(filter['readPreference'])
+          if Lint.enabled?
+            filter.each do |k, v|
+              unless String === k
+                raise Error::LintError, "All keys in filter must be strings: #{filter.inspect}"
+              end
+            end
+          end
+          Lint.validate_camel_case_read_preference(filter['readPreference'])
           SPECIAL_FIELD_MAPPINGS.each do |special, normal|
-            document.store(normal, filter[special]) unless filter[special].nil?
+            unless (v = filter[special]).nil?
+              document.store(normal, v)
+            end
           end
           FLAG_MAPPINGS.each do |legacy, flag|
             document.store(flag, true) if flags.include?(legacy)

data/lib/mongo/protocol.rb

@@ -6,6 +6,7 @@ require 'mongo/protocol/serializers'
 require 'mongo/protocol/registry'
 require 'mongo/protocol/bit_vector'
 require 'mongo/protocol/message'
+require 'mongo/protocol/caching_hash'
 
 # Client Requests
 require 'mongo/protocol/compressed'

data/lib/mongo/query_cache.rb

@@ -179,7 +179,8 @@ module Mongo
       #
       # @api private
       def get(**opts)
-        limit = opts[:limit]
+        limit = normalized_limit(opts[:limit])
+
         _namespace_key = namespace_key(**opts)
         _cache_key = cache_key(**opts)
 
@@ -189,7 +190,7 @@ module Mongo
         caching_cursor = namespace_hash[_cache_key]
         return nil unless caching_cursor
 
-        caching_cursor_limit = caching_cursor.view.limit
+        caching_cursor_limit = normalized_limit(caching_cursor.view.limit)
 
         # There are two scenarios in which a caching cursor could fulfill the
         # query:
@@ -199,6 +200,7 @@ module Mongo
         #
         # Otherwise, return nil because the stored cursor will not satisfy
         # the query.
+
         if limit && (caching_cursor_limit.nil? || caching_cursor_limit >= limit)
           caching_cursor
         elsif limit.nil? && caching_cursor_limit.nil?
@@ -208,6 +210,14 @@ module Mongo
         end
       end
 
+      def normalized_limit(limit)
+        return nil unless limit
+        # For the purposes of caching, a limit of 0 means no limit, as mongo treats it as such.
+        return nil if limit == 0
+        # For the purposes of caching, a negative limit is the same as as a positive limit.
+        limit.abs
+      end
+
       private
 
       def cache_key(**opts)
@@ -270,6 +280,21 @@ module Mongo
       ensure
         QueryCache.clear
       end
+
+      # ActiveJob middleware that activates the query cache for each job.
+      module ActiveJob
+        def self.included(base)
+          base.class_eval do
+            around_perform do |_job, block|
+              QueryCache.cache do
+                block.call
+              end
+            ensure
+              QueryCache.clear
+            end
+          end
+        end
+      end
     end
   end
 end

data/lib/mongo/retryable.rb

@@ -191,25 +191,29 @@ module Mongo
     # @note This only retries operations on not master failures, since it is
     #   the only case we can be sure a partial write did not already occur.
     #
-    # @param [ nil | Session ] session Optional session to use with the operation.
     # @param [ nil | Hash | WriteConcern::Base ] write_concern The write concern.
-    # @param [ true | false ] ending_transaction True if the write operation is abortTransaction or
-    #   commitTransaction, false otherwise.
+    # @param [ true | false ] ending_transaction True if the write operation is
+    #   abortTransaction or commitTransaction, false otherwise.
+    # @param [ Context ] context The context for the operation.
     # @param [ Proc ] block The block to execute.
     #
-    # @yieldparam [ Server ] server The server to which the write should be sent.
+    # @yieldparam [ Connection ] connection The connection through which the
+    #   write should be sent.
     # @yieldparam [ Integer ] txn_num Transaction number (NOT the ACID kind).
+    # @yieldparam [ Operation::Context ] context The operation context.
     #
     # @return [ Result ] The result of the operation.
     #
     # @since 2.1.0
-    def write_with_retry(session, write_concern, ending_transaction = false, &block)
+    def write_with_retry(write_concern, ending_transaction: false, context:, &block)
+      session = context.session
+
       if ending_transaction && !session
         raise ArgumentError, 'Cannot end a transaction without a session'
       end
 
       unless ending_transaction || retry_write_allowed?(session, write_concern)
-        return legacy_write_with_retry(nil, session, &block)
+        return legacy_write_with_retry(nil, context: context, &block)
       end
 
       # If we are here, session is not nil. A session being nil would have
@@ -218,23 +222,46 @@ module Mongo
       server = select_server(cluster, ServerSelector.primary, session)
 
       unless ending_transaction || server.retry_writes?
-        return legacy_write_with_retry(server, session, &block)
+        return legacy_write_with_retry(server, context: context, &block)
       end
 
-      txn_num = if session.in_transaction?
-        session.txn_num
-      else
-        session.next_txn_num
-      end
+      txn_num = nil
+
       begin
-        yield(server, txn_num, false)
-      rescue Error::SocketError, Error::SocketTimeoutError => e
+        connection_succeeded = false
+        server.with_connection(connection_global_id: context.connection_global_id) do |connection|
+          connection_succeeded = true
+
+          session.materialize_if_needed
+          txn_num = if session.in_transaction?
+            session.txn_num
+          else
+            session.next_txn_num
+          end
+
+          # The context needs to be duplicated here because we will be using
+          # it later for the retry as well.
+          yield(connection, txn_num, context.dup)
+        end
+      rescue Error::SocketError, Error::SocketTimeoutError, Auth::Unauthorized => e
         e.add_note('modern retry')
         e.add_note("attempt 1")
         if !e.label?('RetryableWriteError')
-          raise e
+          # If we get an auth error, it was raised when connecting the connection
+          # and therefore we didn't have the connection yet to add labels.
+          # Therefore, check if it is retryable, and if it is, add the label
+          # and retry it. We also want to retry this if there was a Socket error
+          # when trying to create the connection.
+          if !connection_succeeded && !session.in_transaction? && e.write_retryable?
+            e.add_label('RetryableWriteError')
+          else
+            raise e
+          end
         end
-        retry_write(e, session, txn_num, &block)
+
+        # Context#with creates a new context, which is not necessary here
+        # but the API is less prone to misuse this way.
+        retry_write(e, txn_num, context: context.with(is_retry: true), &block)
       rescue Error::OperationFailure => e
         e.add_note('modern retry')
         e.add_note("attempt 1")
@@ -244,7 +271,9 @@ module Mongo
           raise e
         end
 
-        retry_write(e, session, txn_num, &block)
+        # Context#with creates a new context, which is not necessary here
+        # but the API is less prone to misuse this way.
+        retry_write(e, txn_num, context: context.with(is_retry: true), &block)
       end
     end
 
@@ -258,23 +287,30 @@ module Mongo
     # delegates to legacy_write_with_retry which performs write retries using
     # legacy logic.
     #
-    # @param [ nil | Session ] session Optional session to use with the operation.
     # @param [ nil | Hash | WriteConcern::Base ] write_concern The write concern.
+    # @param [ Context ] context The context for the operation.
     #
-    # @yieldparam [ Server ] server The server to which the write should be sent.
+    # @yieldparam [ Connection ] connection The connection through which the
+    #   write should be sent.
+    # @yieldparam [ nil ] txn_num nil as transaction number.
+    # @yieldparam [ Operation::Context ] context The operation context.
     #
     # @api private
-    def nro_write_with_retry(session, write_concern, &block)
+    def nro_write_with_retry(write_concern, context:, &block)
+      session = context.session
+
+      server = select_server(cluster, ServerSelector.primary, session)
       if session && session.client.options[:retry_writes]
-        server = select_server(cluster, ServerSelector.primary, session)
         begin
-          yield server
+          server.with_connection(connection_global_id: context.connection_global_id) do |connection|
+            yield connection, nil, context
+          end
         rescue Error::SocketError, Error::SocketTimeoutError, Error::OperationFailure => e
           e.add_note('retries disabled')
           raise e
         end
       else
-        legacy_write_with_retry(nil, session, &block)
+        legacy_write_with_retry(server, context: context, &block)
       end
     end
 
@@ -287,12 +323,17 @@ module Mongo
     # @param [ Server ] server The server which should be used for the
     #   operation. If not provided, the current primary will be retrieved from
     #   the cluster.
-    # @param [ nil | Session ] session Optional session to use with the operation.
+    # @param [ Context ] context The context for the operation.
     #
-    # @yieldparam [ Server ] server The server to which the write should be sent.
+    # @yieldparam [ Connection ] connection The connection through which the
+    #    write should be sent.
+    # @yieldparam [ nil ] txn_num nil as transaction number.
+    # @yieldparam [ Operation::Context ] context The operation context.
     #
     # @api private
-    def legacy_write_with_retry(server = nil, session = nil)
+    def legacy_write_with_retry(server = nil, context:)
+      session = context.session
+
       # This is the pre-session retry logic, and is not subject to
       # current retryable write specifications.
       # In particular it does not retry on SocketError and SocketTimeoutError.
@@ -300,7 +341,10 @@ module Mongo
       begin
         attempt += 1
         server ||= select_server(cluster, ServerSelector.primary, session)
-        yield server
+        server.with_connection(connection_global_id: context.connection_global_id) do |connection|
+          # Legacy retries do not use txn_num
+          yield connection, nil, context.dup
+        end
       rescue Error::OperationFailure => e
         e.add_note('legacy retry')
         e.add_note("attempt #{attempt}")
@@ -331,7 +375,7 @@ module Mongo
           raise e
         end
         retry_read(e, server_selector, session, &block)
-      rescue Error::OperationFailure => e
+      rescue Error::OperationFailure, Auth::Unauthorized => e
         e.add_note('modern retry')
         e.add_note("attempt 1")
         if session.in_transaction? || !e.write_retryable?
@@ -423,7 +467,9 @@ module Mongo
       raise original_error
     end
 
-    def retry_write(original_error, session, txn_num, &block)
+    def retry_write(original_error, txn_num, context:, &block)
+      session = context.session
+
       # We do not request a scan of the cluster here, because error handling
       # for the error which triggered the retry should have updated the
       # server description and/or topology as necessary (specifically,
@@ -444,7 +490,9 @@ module Mongo
         raise Error::RaiseOriginalError
       end
       log_retry(original_error, message: 'Write retry')
-      yield(server, txn_num, true)
+      server.with_connection(connection_global_id: context.connection_global_id) do |connection|
+        yield(connection, txn_num, context)
+      end
     rescue Error::SocketError, Error::SocketTimeoutError => e
       e.add_note('modern retry')
       e.add_note('attempt 2')