Add ability to query using Composite Batch Api

lib/active_force.rb

@@ -9,8 +9,16 @@ module ActiveForce
 
   class << self
     attr_accessor :sfdc_client
+    attr_writer :composite_batch_query_threshold
+
+    def composite_batch_query_threshold
+      @composite_batch_query_threshold ||= 100_000
+
+      return @composite_batch_query_threshold.call if @composite_batch_query_threshold.respond_to?(:call)
+
+      @composite_batch_query_threshold
+    end
   end
 
   self.sfdc_client = Restforce.new
-
 end

lib/active_force/active_query.rb

@@ -1,6 +1,7 @@
 require 'active_support/all'
 require 'active_force/query'
 require 'active_force/select_builder'
+require 'active_force/composite_batch_query'
 require 'forwardable'
 
 module ActiveForce
@@ -245,7 +246,13 @@ def quote_string(s)
     end
 
     def result
-      sfdc_client.query(self.to_s)
+      soql = self.to_s
+
+      if soql.length >= ActiveForce.composite_batch_query_threshold
+        CompositeBatchQuery.call(soql, sfdc_client)
+      else
+        sfdc_client.query(soql)
+      end
     end
 
     def build_order_by(args)

lib/active_force/composite_batch_query.rb

@@ -0,0 +1,42 @@
+require 'active_support/all'
+
+module ActiveForce
+  class CompositeBatchQuery
+    def self.call(soql, sfdc_client = ActiveForce.sfdc_client)
+      new(soql, sfdc_client).call
+    end
+
+    attr_reader :sfdc_client, :soql
+
+    def initialize(soql, sfdc_client)
+      @sfdc_client = sfdc_client
+      @soql = soql
+    end
+
+    def call
+      results = sfdc_client.batch do |subrequests|
+        subrequests.requests << {
+          method: "GET",
+          url: "v#{subrequests.options[:api_version]}/query?" + {q: soql}.to_query
+        }
+      end
+
+      r = results.first
+
+      process_composite_batch_error_response(r) if r.statusCode >= 300
+
+      r.result
+    end
+
+    private
+
+    def process_composite_batch_error_response(r)
+      response_value = {status: r.statusCode, body: r.result.as_json}
+      error_code = r.result.dig(0, "errorCode")
+      message = "#{error_code}: #{r.result.dig(0, "message")}"
+      message << "\nRESPONSE: #{r.result.to_json}"
+
+      raise Restforce::ErrorCode.get_exception_class(error_code).new(message, response_value)
+    end
+  end
+end

spec/active_force/active_query_spec.rb

@@ -557,6 +557,26 @@ def check_ranges(base_query, start, finish, &format_block)
     end
   end
 
+  describe '#result routing' do
+    after do
+      ActiveForce.instance_variable_set(:@composite_batch_query_threshold, nil)
+    end
+
+    it 'uses sfdc_client.query when SOQL is at or below the threshold' do
+      ActiveForce.composite_batch_query_threshold = 100_000
+      expect(client).to receive(:query).and_return([])
+      expect(ActiveForce::CompositeBatchQuery).not_to receive(:call)
+      active_query.where("Id = 'foo'").to_a
+    end
+
+    it 'uses CompositeBatchQuery.call when SOQL exceeds the threshold' do
+      ActiveForce.composite_batch_query_threshold = 0
+      expect(ActiveForce::CompositeBatchQuery).to receive(:call).and_return([])
+      expect(client).not_to receive(:query)
+      active_query.where("Id = 'foo'").to_a
+    end
+  end
+
   describe "#order" do
     context 'when it is symbol' do 
       it "should add an order condition with actual SF field name" do

spec/active_force/composite_batch_query_spec.rb

@@ -0,0 +1,101 @@
+require 'spec_helper'
+
+describe ActiveForce::CompositeBatchQuery do
+  let(:soql) { "SELECT Id FROM Account WHERE Name = 'Test'" }
+  let(:api_version) { '58.0' }
+  let(:sfdc_client) { double(Restforce) }
+  let(:subrequests) do
+    double('subrequests', requests: [], options: { api_version: api_version })
+  end
+
+  describe '.call' do
+    it 'delegates to new(...).call' do
+      result_record = double('result_record')
+      instance = instance_double(described_class)
+      allow(described_class).to receive(:new).with(soql, sfdc_client).and_return(instance)
+      allow(instance).to receive(:call).and_return(result_record)
+
+      expect(described_class.call(soql, sfdc_client)).to eq(result_record)
+    end
+
+    it 'defaults to ActiveForce.sfdc_client when no client is provided' do
+      default_client = double(Restforce)
+      allow(ActiveForce).to receive(:sfdc_client).and_return(default_client)
+
+      instance = instance_double(described_class)
+      allow(described_class).to receive(:new).with(soql, default_client).and_return(instance)
+      allow(instance).to receive(:call).and_return(double('result'))
+
+      expect { described_class.call(soql) }.not_to raise_error
+      expect(described_class).to have_received(:new).with(soql, default_client)
+    end
+  end
+
+  describe '#call' do
+    context 'when the response is successful' do
+      let(:result_record) { double(Restforce::Collection) }
+      let(:batch_result) do
+        double(Restforce::Mash, statusCode: 200, result: result_record)
+      end
+
+      before do
+        allow(sfdc_client).to receive(:batch).and_yield(subrequests).and_return([batch_result])
+      end
+
+      it 'sends a GET batch subrequest with the correct URL' do
+        described_class.call(soql, sfdc_client)
+
+        expected_url = "v#{api_version}/query?" + { q: soql }.to_query
+        expect(subrequests.requests).to include(
+          hash_including(method: 'GET', url: expected_url)
+        )
+      end
+
+      it 'returns the result from the batch response' do
+        expect(described_class.call(soql, sfdc_client)).to eq(result_record)
+      end
+    end
+
+    context 'when the response statusCode is exactly 300' do
+      let(:error_body) do
+        [{ 'errorCode' => 'MULTIPLE_CHOICES', 'message' => 'multiple records found' }]
+      end
+      let(:batch_result) do
+        double('batch_result', statusCode: 300, result: error_body)
+      end
+
+      before do
+        allow(sfdc_client).to receive(:batch).and_yield(subrequests).and_return([batch_result])
+      end
+
+      it 'treats it as an error' do
+        expect {
+          described_class.call(soql, sfdc_client)
+        }.to raise_error(Restforce::ResponseError)
+      end
+    end
+
+    context 'when the response is an error' do
+      let(:error_body) do
+        [{ 'errorCode' => 'MALFORMED_QUERY', 'message' => 'unexpected token' }]
+      end
+      let(:batch_result) do
+        double(Restforce::Mash, statusCode: 400, result: error_body)
+      end
+
+      before do
+        allow(sfdc_client).to receive(:batch).and_yield(subrequests).and_return([batch_result])
+      end
+
+      it 'raises a Restforce error with the correct error code and message' do
+        expect {
+          described_class.call(soql, sfdc_client)
+        }.to raise_error(Restforce::ErrorCode::MalformedQuery) do |error|
+          expect(error.message).to include('MALFORMED_QUERY')
+          expect(error.message).to include('unexpected token')
+          expect(error.message).to include('RESPONSE:')
+        end
+      end
+    end
+  end
+end

spec/active_force_spec.rb

@@ -4,4 +4,24 @@
   it 'should have a version number' do
     expect(ActiveForce::VERSION).to_not be_nil
   end
+
+  describe '.composite_batch_query_threshold' do
+    after do
+      ActiveForce.instance_variable_set(:@composite_batch_query_threshold, nil)
+    end
+
+    it 'returns 100_000 by default' do
+      expect(ActiveForce.composite_batch_query_threshold).to eq(100_000)
+    end
+
+    it 'allows setting a custom integer value' do
+      ActiveForce.composite_batch_query_threshold = 25_000
+      expect(ActiveForce.composite_batch_query_threshold).to eq(25_000)
+    end
+
+    it 'supports callable values (proc/lambda)' do
+      ActiveForce.composite_batch_query_threshold = -> { 50_000 }
+      expect(ActiveForce.composite_batch_query_threshold).to eq(50_000)
+    end
+  end
 end