mongo 2.19.1 → 2.21.0

data/lib/mongo/collection/view/aggregation.rb

@@ -15,6 +15,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+require 'mongo/collection/view/aggregation/behavior'
+
 module Mongo
   class Collection
     class View
@@ -23,46 +25,11 @@ module Mongo
       #
       # @since 2.0.0
       class Aggregation
-        extend Forwardable
-        include Enumerable
-        include Immutable
-        include Iterable
-        include Explainable
-        include Loggable
-        include Retryable
+        include Behavior
 
-        # @return [ View ] view The collection view.
-        attr_reader :view
         # @return [ Array<Hash> ] pipeline The aggregation pipeline.
         attr_reader :pipeline
 
-        # Delegate necessary operations to the view.
-        def_delegators :view, :collection, :read, :cluster
-
-        # Delegate necessary operations to the collection.
-        def_delegators :collection, :database, :client
-
-        # The reroute message.
-        #
-        # @since 2.1.0
-        # @deprecated
-        REROUTE = 'Rerouting the Aggregation operation to the primary server.'.freeze
-
-        # Set to true if disk usage is allowed during the aggregation.
-        #
-        # @example Set disk usage flag.
-        #   aggregation.allow_disk_use(true)
-        #
-        # @param [ true, false ] value The flag value.
-        #
-        # @return [ true, false, Aggregation ] The aggregation if a value was
-        #   set or the value if used as a getter.
-        #
-        # @since 2.0.0
-        def allow_disk_use(value = nil)
-          configure(:allow_disk_use, value)
-        end
-
         # Initialize the aggregation for the provided collection view, pipeline
         # and options.
         #
@@ -86,59 +53,29 @@ module Mongo
         # @option options [ Hash ] :let Mapping of variables to use in the pipeline.
         #   See the server documentation for details.
         # @option options [ Integer ] :max_time_ms The maximum amount of time in
-        #   milliseconds to allow the aggregation to run.
-        # @option options [ true, false ] :use_cursor Indicates whether the command
-        #   will request that the server provide results using a cursor. Note that
-        #   as of server version 3.6, aggregations always provide results using a
-        #   cursor and this option is therefore not valid.
+        #   milliseconds to allow the aggregation to run. This option is deprecated, use
+        #   :timeout_ms instead.
         # @option options [ Session ] :session The session to use.
+        # @option options [ :cursor_lifetime | :iteration ] :timeout_mode How to interpret
+        #   :timeout_ms (whether it applies to the lifetime of the cursor, or per
+        #   iteration).
+        # @option options [ Integer ] :timeout_ms The operation timeout in milliseconds.
+        #    Must be a non-negative integer. An explicit value of 0 means infinite.
+        #    The default value is unset which means the value is inherited from
+        #    the collection or the database or the client.
         #
         # @since 2.0.0
         def initialize(view, pipeline, options = {})
-          @view = view
-          @pipeline = pipeline.dup
-          unless Mongo.broken_view_aggregate || view.filter.empty?
-            @pipeline.unshift(:$match => view.filter)
+          perform_setup(view, options) do
+            @pipeline = pipeline.dup
+            unless Mongo.broken_view_aggregate || view.filter.empty?
+              @pipeline.unshift(:$match => view.filter)
+            end
           end
-          @options = BSON::Document.new(options).freeze
-        end
-
-        # Get the explain plan for the aggregation.
-        #
-        # @example Get the explain plan for the aggregation.
-        #   aggregation.explain
-        #
-        # @return [ Hash ] The explain plan.
-        #
-        # @since 2.0.0
-        def explain
-          self.class.new(view, pipeline, options.merge(explain: true)).first
-        end
-
-        # Whether this aggregation will write its result to a database collection.
-        #
-        # @return [ Boolean ] Whether the aggregation will write its result
-        #   to a collection.
-        #
-        # @api private
-        def write?
-          pipeline.any? { |op| op.key?('$out') || op.key?(:$out) || op.key?('$merge') || op.key?(:$merge) }
         end
 
         private
 
-        def server_selector
-          @view.send(:server_selector)
-        end
-
-        def aggregate_spec(session, read_preference)
-          Builder::Aggregation.new(
-            pipeline,
-            view,
-            options.merge(session: session, read_preference: read_preference)
-          ).specification
-        end
-
         def new(options)
           Aggregation.new(view, pipeline, options)
         end
@@ -180,32 +117,29 @@ module Mongo
 
         end
 
-        def send_initial_query(server, session)
-          server.with_connection do |connection|
+        def send_initial_query(server, context)
+          if server.load_balancer?
+            # Connection will be checked in when cursor is drained.
+            connection = server.pool.check_out(context: context)
             initial_query_op(
-              session,
+              context.session,
               effective_read_preference(connection)
             ).execute_with_connection(
               connection,
-              context: Operation::Context.new(client: client, session: session)
+              context: context
             )
+          else
+            server.with_connection do |connection|
+              initial_query_op(
+                context.session,
+                effective_read_preference(connection)
+              ).execute_with_connection(
+                connection,
+                context: context
+              )
+            end
           end
         end
-
-        # Skip, sort, limit, projection are specified as pipeline stages
-        # rather than as options.
-        def cache_options
-          {
-            namespace: collection.namespace,
-            selector: pipeline,
-            read_concern: view.read_concern,
-            read_preference: view.read_preference,
-            collation: options[:collation],
-            # Aggregations can read documents from more than one collection,
-            # so they will be cleared on every write operation.
-            multi_collection: true,
-          }
-        end
       end
     end
   end

data/lib/mongo/collection/view/builder/aggregation.rb

@@ -113,17 +113,11 @@ module Mongo
               command[:readConcern] = Options::Mapper.transform_values_to_strings(
                 read_concern)
             end
-            command[:cursor] = cursor if cursor
+            command[:cursor] = batch_size_doc
             command.merge!(Options::Mapper.transform_documents(options, MAPPINGS))
             command
           end
 
-          def cursor
-            if options[:use_cursor] == true || options[:use_cursor].nil?
-              batch_size_doc
-            end
-          end
-
           def batch_size_doc
             value = options[:batch_size] || view.batch_size
             if value == 0 && write?

data/lib/mongo/collection/view/change_stream.rb

@@ -15,6 +15,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+require 'mongo/collection/view/aggregation/behavior'
 require 'mongo/collection/view/change_stream/retryable'
 
 module Mongo
@@ -35,7 +36,8 @@ module Mongo
       #
       #
       # @since 2.5.0
-      class ChangeStream < Aggregation
+      class ChangeStream
+        include Aggregation::Behavior
         include Retryable
 
         # @return [ String ] The fullDocument option default value.
@@ -60,6 +62,10 @@ module Mongo
         # @since 2.5.0
         attr_reader :options
 
+        # @return [ Cursor ] the underlying cursor for this operation
+        # @api private
+        attr_reader :cursor
+
         # Initialize the change stream for the provided collection view, pipeline
         # and options.
         #
@@ -125,11 +131,13 @@ module Mongo
         #
         # @since 2.5.0
         def initialize(view, pipeline, changes_for, options = {})
-          @view = view
-          @changes_for = changes_for
-          @change_stream_filters = pipeline && pipeline.dup
-          @options = options && options.dup.freeze
-          @start_after = @options[:start_after]
+          # change stream cursors can only be :iterable, so we don't allow
+          # timeout_mode to be specified.
+          perform_setup(view, options, forbid: %i[ timeout_mode ]) do
+            @changes_for = changes_for
+            @change_stream_filters = pipeline && pipeline.dup
+            @start_after = @options[:start_after]
+          end
 
           # The resume token tracked by the change stream, used only
           # when there is no cursor, or no cursor resume token
@@ -181,24 +189,30 @@ module Mongo
         # @return [ BSON::Document | nil ] A change stream document.
         # @since 2.6.0
         def try_next
+          recreate_cursor! if @timed_out
+
           raise StopIteration.new if closed?
+
           begin
             doc = @cursor.try_next
           rescue Mongo::Error => e
-            if !e.change_stream_resumable?
-              raise
-            end
-
-            # Rerun initial aggregation.
-            # Any errors here will stop iteration and break out of this
-            # method.
+            # "If a next call fails with a timeout error, drivers MUST NOT
+            # invalidate the change stream. The subsequent next call MUST
+            # perform a resume attempt to establish a new change stream on the
+            # server..."
+            #
+            # However, SocketTimeoutErrors are TimeoutErrors, but are also
+            # change-stream-resumable. To preserve existing (specified) behavior,
+            # We only count timeouts when the error is not also
+            # change-stream-resumable.
+            @timed_out = e.is_a?(Mongo::Error::TimeoutError) && !e.change_stream_resumable?
+
+            raise unless @timed_out || e.change_stream_resumable?
 
-            # Save cursor's resume token so we can use it
-            # to create a new cursor
             @resume_token = @cursor.resume_token
+            raise e if @timed_out
 
-            close
-            create_cursor!
+            recreate_cursor!(@cursor.context)
             retry
           end
 
@@ -231,14 +245,17 @@ module Mongo
         #   This method ignores any errors that occur when closing the
         #   server-side cursor.
         #
+        # @params [ Hash ] opts Options to be passed to the cursor close
+        #   command.
+        #
         # @return [ nil ] Always nil.
         #
         # @since 2.5.0
-        def close
+        def close(opts = {})
           unless closed?
             begin
-              @cursor.close
-            rescue Error::OperationFailure, Error::SocketError, Error::SocketTimeoutError, Error::MissingConnection
+              @cursor.close(opts)
+            rescue Error::OperationFailure::Family, Error::SocketError, Error::SocketTimeoutError, Error::MissingConnection
               # ignore
             end
             @cursor = nil
@@ -284,6 +301,28 @@ module Mongo
           cursor_resume_token || @resume_token
         end
 
+        # "change streams are an abstraction around tailable-awaitData cursors..."
+        #
+        # @return :tailable_await
+        def cursor_type
+          :tailable_await
+        end
+
+        # "change streams...implicitly use ITERATION mode"
+        #
+        # @return :iteration
+        def timeout_mode
+          :iteration
+        end
+
+        # Returns the value of the max_await_time_ms option that was
+        # passed to this change stream.
+        #
+        # @return [ Integer | nil ] the max_await_time_ms value
+        def max_await_time_ms
+          options[:max_await_time_ms]
+        end
+
         private
 
         def for_cluster?
@@ -298,19 +337,23 @@ module Mongo
           !for_cluster? && !for_database?
         end
 
-        def create_cursor!
+        def create_cursor!(timeout_ms = nil)
           # clear the cache because we may get a newer or an older server
           # (rolling upgrades)
           @start_at_operation_time_supported = nil
 
-          session = client.send(:get_session, @options)
+          session = client.get_session(@options)
+          context = Operation::Context.new(client: client, session: session, view: self, operation_timeouts: timeout_ms ? { operation_timeout_ms: timeout_ms } : operation_timeouts)
+
           start_at_operation_time = nil
           start_at_operation_time_supported = nil
-          @cursor = read_with_retry_cursor(session, server_selector, view) do |server|
+
+          @cursor = read_with_retry_cursor(session, server_selector, self, context: context) do |server|
             server.with_connection do |connection|
               start_at_operation_time_supported = connection.description.server_version_gte?('4.0')
 
-              result = send_initial_query(connection, session)
+              result = send_initial_query(connection, context)
+
               if doc = result.replies.first && result.replies.first.documents.first
                 start_at_operation_time = doc['operationTime']
               else
@@ -324,6 +367,7 @@ module Mongo
               result
             end
           end
+
           @start_at_operation_time = start_at_operation_time
           @start_at_operation_time_supported = start_at_operation_time_supported
         end
@@ -390,11 +434,11 @@ module Mongo
           end
         end
 
-        def send_initial_query(connection, session)
-          initial_query_op(session, view.read_preference)
+        def send_initial_query(connection, context)
+          initial_query_op(context.session, view.read_preference)
             .execute_with_connection(
               connection,
-              context: Operation::Context.new(client: client, session: session),
+              context: context,
             )
         end
 
@@ -412,6 +456,15 @@ module Mongo
         def resuming?
           !!@resuming
         end
+
+        # Recreates the current cursor (typically as a consequence of attempting
+        # to resume the change stream)
+        def recreate_cursor!(context = nil)
+          @timed_out = false
+
+          close
+          create_cursor!(context&.remaining_timeout_ms)
+        end
       end
     end
   end

data/lib/mongo/collection/view/iterable.rb

@@ -15,6 +15,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+require 'mongo/cursor_host'
+
 module Mongo
   class Collection
     class View
@@ -24,13 +26,7 @@ module Mongo
       #
       # @since 2.0.0
       module Iterable
-
-        # Returns the cursor associated with this view, if any.
-        #
-        # @return [ nil | Cursor ] The cursor, if any.
-        #
-        # @api private
-        attr_reader :cursor
+        include Mongo::CursorHost
 
         # Iterate through documents returned by a query with this +View+.
         #
@@ -45,48 +41,21 @@ module Mongo
         #
         # @yieldparam [ Hash ] Each matching document.
         def each
-          # If the caching cursor is closed and was not fully iterated,
-          # the documents we have in it are not the complete result set and
-          # we have no way of completing that iteration.
-          # Therefore, discard that cursor and start iteration again.
-          # The case of the caching cursor not being closed and not having
-          # been fully iterated isn't tested - see RUBY-2773.
-          @cursor = if use_query_cache? && cached_cursor && (
-            cached_cursor.fully_iterated? || !cached_cursor.closed?
-          )
-            cached_cursor
-          else
-            session = client.send(:get_session, @options)
-            select_cursor(session).tap do |cursor|
-              if use_query_cache?
-                # No need to store the cursor in the query cache if there is
-                # already a cached cursor stored at this key.
-                QueryCache.set(cursor, **cache_options)
-              end
-            end
-          end
+          @cursor = prefer_cached_cursor? ? cached_cursor : new_cursor_for_iteration
+          return @cursor.to_enum unless block_given?
 
-          if use_query_cache?
-            # If a query with a limit is performed, the query cache will
-            # re-use results from an earlier query with the same or larger
-            # limit, and then impose the lower limit during iteration.
-            limit_for_cached_query = respond_to?(:limit) ? QueryCache.normalized_limit(limit) : nil
-          end
+          limit_for_cached_query = compute_limit_for_cached_query
 
-          if block_given?
-            # Ruby versions 2.5 and older do not support arr[0..nil] syntax, so
-            # this must be a separate conditional.
-            cursor_to_iterate = if limit_for_cached_query
-              @cursor.to_a[0...limit_for_cached_query]
-            else
-              @cursor
-            end
-
-            cursor_to_iterate.each do |doc|
-              yield doc
-            end
+          # Ruby versions 2.5 and older do not support arr[0..nil] syntax, so
+          # this must be a separate conditional.
+          cursor_to_iterate = if limit_for_cached_query
+            @cursor.to_a[0...limit_for_cached_query]
           else
-            @cursor.to_enum
+            @cursor
+          end
+
+          cursor_to_iterate.each do |doc|
+            yield doc
           end
         end
 
@@ -100,7 +69,7 @@ module Mongo
         #
         # @return [ nil ] Always nil.
         #
-        # @raise [ Error::OperationFailure ] If the server cursor close fails.
+        # @raise [ Error::OperationFailure::Family ] If the server cursor close fails.
         #
         # @since 2.1.0
         def close_query
@@ -113,18 +82,25 @@ module Mongo
         private
 
         def select_cursor(session)
+          context = Operation::Context.new(
+            client: client,
+            session: session,
+            operation_timeouts: operation_timeouts,
+            view: self
+          )
+
           if respond_to?(:write?, true) && write?
             server = server_selector.select_server(cluster, nil, session, write_aggregation: true)
-            result = send_initial_query(server, session)
+            result = send_initial_query(server, context)
 
             if use_query_cache?
-              CachingCursor.new(view, result, server, session: session)
+              CachingCursor.new(view, result, server, session: session, context: context)
             else
-              Cursor.new(view, result, server, session: session)
+              Cursor.new(view, result, server, session: session, context: context)
             end
           else
-            read_with_retry_cursor(session, server_selector, view) do |server|
-              send_initial_query(server, session)
+            read_with_retry_cursor(session, server_selector, view, context: context) do |server|
+              send_initial_query(server, context)
             end
           end
         end
@@ -162,29 +138,27 @@ module Mongo
             let: options[:let],
             limit: limit,
             allow_disk_use: options[:allow_disk_use],
+            allow_partial_results: options[:allow_partial_results],
             read: read,
             read_concern: options[:read_concern] || read_concern,
             batch_size: batch_size,
             hint: options[:hint],
             max_scan: options[:max_scan],
-            max_time_ms: options[:max_time_ms],
             max_value: options[:max_value],
             min_value: options[:min_value],
             no_cursor_timeout: options[:no_cursor_timeout],
             return_key: options[:return_key],
             show_disk_loc: options[:show_disk_loc],
             comment: options[:comment],
-            oplog_replay: if (v = options[:oplog_replay]).nil?
-              collection.options[:oplog_replay]
-            else
-              v
-            end,
+            oplog_replay: oplog_replay
           }
 
           if spec[:oplog_replay]
             collection.client.log_warn("The :oplog_replay option is deprecated and ignored by MongoDB 4.4 and later")
           end
 
+          maybe_set_tailable_options(spec)
+
           if explained?
             spec[:explain] = options[:explain]
             Operation::Explain.new(spec)
@@ -193,13 +167,71 @@ module Mongo
           end
         end
 
-        def send_initial_query(server, session = nil)
-          initial_query_op(session).execute(server, context: Operation::Context.new(client: client, session: session))
+        def send_initial_query(server, context)
+          operation = initial_query_op(context.session)
+          if server.load_balancer?
+            # Connection will be checked in when cursor is drained.
+            connection = server.pool.check_out(context: context)
+            operation.execute_with_connection(connection, context: context)
+          else
+            operation.execute(server, context: context)
+          end
         end
 
         def use_query_cache?
           QueryCache.enabled? && !collection.system_collection?
         end
+
+        # If the caching cursor is closed and was not fully iterated,
+        # the documents we have in it are not the complete result set and
+        # we have no way of completing that iteration.
+        # Therefore, discard that cursor and start iteration again.
+        def prefer_cached_cursor?
+          use_query_cache? &&
+            cached_cursor &&
+            (cached_cursor.fully_iterated? || !cached_cursor.closed?)
+        end
+
+        # Start a new cursor for use when iterating (via #each).
+        def new_cursor_for_iteration
+          session = client.get_session(@options)
+          select_cursor(session).tap do |cursor|
+            if use_query_cache?
+              # No need to store the cursor in the query cache if there is
+              # already a cached cursor stored at this key.
+              QueryCache.set(cursor, **cache_options)
+            end
+          end
+        end
+
+        def compute_limit_for_cached_query
+          return nil unless use_query_cache? && respond_to?(:limit)
+
+          # If a query with a limit is performed, the query cache will
+          # re-use results from an earlier query with the same or larger
+          # limit, and then impose the lower limit during iteration.
+          return QueryCache.normalized_limit(limit)
+        end
+
+        # Add tailable cusror options to the command specifiction if needed.
+        #
+        # @param [ Hash ] spec The command specification.
+        def maybe_set_tailable_options(spec)
+          case cursor_type
+          when :tailable
+            spec[:tailable] = true
+          when :tailable_await
+            spec[:tailable] = true
+            spec[:await_data] = true
+          end
+        end
+
+        # @return [ true | false | nil ] options[:oplog_replay], if
+        #    set, otherwise the same option from the collection.
+        def oplog_replay
+          v = options[:oplog_replay]
+          v.nil? ? collection.options[:oplog_replay] : v
+        end
       end
     end
   end