From b603a2b22cec5ea406ea5911e277d30a03df4c85 Mon Sep 17 00:00:00 2001
From: Jamis Buck <jamis.buck@mongodb.com>
Date: Tue, 29 Aug 2023 14:46:26 -0600
Subject: [PATCH] MONGOID-5647 Allow #count to be used with #for_js (#5693)
 (#5696)

* MONGOID-5647 Allow #count to be used with #for_js

* look for $where in nested hashes

* fix typo in method name
---
 lib/mongoid/contextual/mongo.rb       | 27 ++++++++++++++++++++++++++-
 spec/mongoid/contextual/mongo_spec.rb | 10 ++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/lib/mongoid/contextual/mongo.rb b/lib/mongoid/contextual/mongo.rb
index 35d564705..598945aa0 100644
--- a/lib/mongoid/contextual/mongo.rb
+++ b/lib/mongoid/contextual/mongo.rb
@@ -65,7 +65,14 @@ def cached?
       # @return [ Integer ] The number of matches.
       def count(options = {}, &block)
         return super(&block) if block_given?
-        try_cache(:count) { view.count_documents(options) }
+
+        try_cache(:count) do
+          if valid_for_count_documents?
+            view.count_documents(options)
+          else
+            view.count(options)
+          end
+        end
       end
 
       # Get the estimated number of documents matching the query.
@@ -902,6 +909,24 @@ def process_raw_docs(raw_docs, limit)
         docs = eager_load(docs)
         limit ? docs : docs.first
       end
+
+      # Queries whether the current context is valid for use with
+      # the #count_documents? predicate. A context is valid if it
+      # does not include a `$where` operator.
+      #
+      # @return [ true | false ] whether or not the current context
+      #   excludes a `$where` operator.
+      def valid_for_count_documents?(hash = view.filter)
+        # Note that `view.filter` is a BSON::Document, and all keys in a
+        # BSON::Document are strings; we don't need to worry about symbol
+        # representations of `$where`.
+        hash.keys.each do |key|
+          return false if key == '$where'
+          return false if hash[key].is_a?(Hash) && !valid_for_count_documents?(hash[key])
+        end
+
+        true
+      end
     end
   end
 end
diff --git a/spec/mongoid/contextual/mongo_spec.rb b/spec/mongoid/contextual/mongo_spec.rb
index 441d0e271..f7b9c3218 100644
--- a/spec/mongoid/contextual/mongo_spec.rb
+++ b/spec/mongoid/contextual/mongo_spec.rb
@@ -189,6 +189,16 @@
         end
       end
     end
+
+    context 'when for_js is present' do
+      let(:context) do
+        Band.for_js('this.name == "Depeche Mode"')
+      end
+
+      it 'counts the expected records' do
+        expect(context.count).to eq(1)
+      end
+    end
   end
 
   describe "#estimated_count" do