3390

Author: comandeo-mongo

lib/mongo/collection/view/iterable.rb

@@ -99,7 +99,7 @@ def select_cursor(session)
               Cursor.new(view, result, server, session: session, context: context)
             end
           else
-            read_with_retry_cursor(session, server_selector, view) do |server|
+            read_with_retry_cursor(session, server_selector, view, context: context) do |server|
               send_initial_query(server, context)
             end
           end

lib/mongo/collection/view/readable.rb

@@ -187,7 +187,7 @@ def count(opts = {})
               session: session,
               operation_timeouts: operation_timeouts(opts)
             )
-            read_with_retry(session, selector) do |server|
+            read_with_retry(session, selector, context) do |server|
               Operation::Count.new(
                 selector: cmd,
                 db_name: database.name,
@@ -237,7 +237,7 @@ def count_documents(opts = {})
           pipeline << { :'$limit' => opts[:limit] } if opts[:limit]
           pipeline << { :'$group' => { _id: 1, n: { :'$sum' => 1 } } }
 
-          opts = opts.slice(:hint, :max_time_ms, :read, :collation, :session, :comment)
+          opts = opts.slice(:hint, :max_time_ms, :read, :collation, :session, :comment, :timeout_ms)
           opts[:collation] ||= collation
 
           first = aggregate(pipeline, opts).first
@@ -279,12 +279,12 @@ def estimated_document_count(opts = {})
           read_pref = opts[:read] || read_preference
           selector = ServerSelector.get(read_pref || server_selector)
           with_session(opts) do |session|
-            read_with_retry(session, selector) do |server|
-              context = Operation::Context.new(
-                client: client,
-                session: session,
-                operation_timeouts: operation_timeouts(opts)
-              )
+            context = Operation::Context.new(
+              client: client,
+              session: session,
+              operation_timeouts: operation_timeouts(opts)
+            )
+            read_with_retry(session, selector, context) do |server|
               cmd = { count: collection.name }
               cmd[:maxTimeMS] = opts[:max_time_ms] if opts[:max_time_ms]
               if read_concern
@@ -347,7 +347,12 @@ def distinct(field_name, opts = {})
           read_pref = opts[:read] || read_preference
           selector = ServerSelector.get(read_pref || server_selector)
           with_session(opts) do |session|
-            read_with_retry(session, selector) do |server|
+            context = Operation::Context.new(
+              client: client,
+              session: session,
+              operation_timeouts: operation_timeouts(opts)
+            )
+            read_with_retry(session, selector, context) do |server|
               Operation::Distinct.new(
                 selector: cmd,
                 db_name: database.name,
@@ -360,11 +365,7 @@ def distinct(field_name, opts = {})
                 collation: opts[:collation] || opts['collation'] || collation,
               ).execute(
                 server,
-                context: Operation::Context.new(
-                  client: client,
-                  session: session,
-                  operation_timeouts: operation_timeouts(opts)
-                )
+                context: context
               )
             end.first['values']
           end

lib/mongo/index/view.rb

@@ -273,9 +273,15 @@ def get(keys_or_name)
       #
       # @since 2.0.0
       def each(&block)
-        session = client.send(:get_session, @options)
-        cursor = read_with_retry_cursor(session, ServerSelector.primary, self) do |server|
-          send_initial_query(server, session)
+        session = client.get_session(@options)
+        context = Operation::Context.new(
+          client: client,
+          session: session,
+          operation_timeouts: operation_timeouts(@options)
+        )
+
+        cursor = read_with_retry_cursor(session, ServerSelector.primary, self, context: context) do |server|
+          send_initial_query(server, session, context)
         end
         if block_given?
           cursor.each do |doc|
@@ -386,8 +392,8 @@ def normalize_models(models, server)
         end
       end
 
-      def send_initial_query(server, session)
-        initial_query_op(session).execute(server, context: Operation::Context.new(client: client, session: session))
+      def send_initial_query(server, session, context)
+        initial_query_op(session).execute(server, context: context)
       end
     end
   end

lib/mongo/operation/find/op_msg.rb

@@ -71,7 +71,9 @@ def apply_find_timeouts_to(spec, timeout_ms)
             end
           end
 
-          spec
+          spec.tap do |spc|
+            spc.delete(:maxTimeMS) if spc[:maxTimeMS].nil?
+          end
         end
 
         def selector(connection)

spec/runners/unified/crud_operations.rb

@@ -47,6 +47,23 @@ def get_find_view(op)
       end
     end
 
+    def count(op)
+      collection = entities.get(:collection, op.use!('object'))
+      use_arguments(op) do |args|
+        opts = {}
+        if session = args.use('session')
+          opts[:session] = entities.get(:session, session)
+        end
+        if comment = args.use('comment')
+          opts[:comment] = comment
+        end
+        if timeout_ms = args.use('timeoutMS')
+          opts[:timeout_ms] = timeout_ms
+        end
+        collection.count(args.use!('filter'), **opts)
+      end
+    end
+
     def count_documents(op)
       collection = entities.get(:collection, op.use!('object'))
       use_arguments(op) do |args|
@@ -57,6 +74,9 @@ def count_documents(op)
         if comment = args.use('comment')
           opts[:comment] = comment
         end
+        if timeout_ms = args.use('timeoutMS')
+          opts[:timeout_ms] = timeout_ms
+        end
         collection.find(args.use!('filter')).count_documents(**opts)
       end
     end
@@ -71,6 +91,9 @@ def estimated_document_count(op)
         if comment = args.use('comment')
           opts[:comment] = comment
         end
+        if timeout_ms = args.use('timeoutMS')
+          opts[:timeout_ms] = timeout_ms
+        end
         collection.estimated_document_count(**opts)
       end
     end
@@ -85,6 +108,9 @@ def distinct(op)
         if comment = args.use('comment')
           opts[:comment] = comment
         end
+        if timeout_ms = args.use('timeoutMS')
+          opts[:timeout_ms] = timeout_ms
+        end
         req = collection.find(args.use!('filter'), **opts).distinct(args.use!('fieldName'), **opts)
         result = req.to_a
       end

spec/runners/unified/ddl_operations.rb

@@ -137,6 +137,9 @@ def list_indexes(op)
         if session = args.use('session')
           opts[:session] = entities.get(:session, session)
         end
+        if timeout_ms = args.use('timeoutMS')
+          opts[:timeout_ms] = timeout_ms
+        end
         collection.indexes(**opts).to_a
       end
     end