Merge pull request yatish27#47 from IntractableQuery/master

Add throttling support to reduce API calls

Author: yatish27

README.md

@@ -50,7 +50,7 @@ OR
 	salesforce = SalesforceBulkApi::Api.new(client)
 
 
-Sample operations:
+### Sample operations:
 
     # Insert/Create
     # Add as many fields per record as needed.
@@ -82,21 +82,26 @@ Sample operations:
     # Query
 	res = salesforce.query("Account", "select id, name, createddate from Account limit 3") # We just need to pass the sobject name and the query string
 
-Helpful methods:
+### Helpful methods:
 
     # Check status of a job via #job_from_id
 	job = salesforce.job_from_id('a00A0001009zA2m') # Returns a SalesforceBulkApi::Job instance
 	puts "status is: #{job.check_job_status.inspect}"
 
-
-Listening to events:
+### Listening to events:
 
     # A job is created
     # Useful when you need to store the job_id before any work begins, then if you fail during a complex load scenario, you can wait for your
     # previous job(s) to finish.
-  salesforce.on_job_created do |job|
-    puts "Job #{job.job_id} created!"
-  end
+    salesforce.on_job_created do |job|
+      puts "Job #{job.job_id} created!"
+    end
+
+### Throttling API calls:
+
+    # By default, this gem (and maybe your app driving it) will query job/batch statuses at an unbounded rate.  We
+    # can fix that, e.g.:
+    salesforce.connection.set_status_throttle(30) # only check status of individual jobs/batches every 30 seconds
 
 ## Installation
 

lib/salesforce_bulk_api.rb

@@ -1,16 +1,18 @@
 require 'rubygems'
 require 'bundler'
 Bundler.require()
-require "salesforce_bulk_api/version"
+require 'salesforce_bulk_api/version'
 require 'net/https'
 require 'xmlsimple'
 require 'csv'
+require 'salesforce_bulk_api/concerns/throttling'
 require 'salesforce_bulk_api/job'
 require 'salesforce_bulk_api/connection'
 
 module SalesforceBulkApi
 
   class Api
+    attr_reader :connection
 
     @@SALESFORCE_API_VERSION = '32.0'
 

lib/salesforce_bulk_api/concerns/throttling.rb

@@ -0,0 +1,62 @@
+module SalesforceBulkApi::Concerns
+  module Throttling
+
+    def throttles
+      @throttles.dup
+    end
+
+    def add_throttle(&throttling_callback)
+      @throttles ||= []
+      @throttles << throttling_callback
+    end
+
+    def set_status_throttle(limit_seconds)
+      set_throttle_limit_in_seconds(limit_seconds,
+        throttle_by_keys: %i(http_method path),
+        only_if: ->(details) { details[:path] == :get })
+    end
+
+    def set_throttle_limit_in_seconds(limit_seconds, throttle_by_keys: %i(http_method path), only_if: Proc.new{true})
+      add_throttle do |details|
+        limit_log = get_limit_log(Time.now - limit_seconds)
+        key = extract_constraint_key_from(details, throttle_by_keys)
+        last_request = limit_log[key]
+
+        if !last_request.nil? && only_if.call(details)
+          seconds_since_last_request = Time.now.to_f - last_request.to_f
+          need_to_wait_seconds = limit_seconds - seconds_since_last_request
+          sleep(need_to_wait_seconds) if need_to_wait_seconds > 0
+        end
+
+        limit_log[key] = Time.now
+      end
+    end
+
+    private
+
+    def extract_constraint_key_from(details, throttle_by_keys)
+      hash = {}
+      throttle_by_keys.each { |k| hash[k] = details[k] }
+      hash
+    end
+
+    def get_limit_log(prune_older_than)
+      @limits ||= Hash.new(0)
+
+      @limits.delete_if do |k, v|
+        v < prune_older_than
+      end
+
+      @limits
+    end
+
+    def throttle(details={})
+      (@throttles || []).each do |callback|
+        args = [details]
+        args = args[0..callback.arity]
+        callback.call(*args)
+      end
+    end
+
+  end
+end

lib/salesforce_bulk_api/connection.rb

@@ -2,6 +2,7 @@ module SalesforceBulkApi
 require 'timeout'
 
   class Connection
+    include Concerns::Throttling
 
     @@XML_HEADER = '<?xml version="1.0" encoding="utf-8" ?>'
     @@API_VERSION = nil
@@ -35,14 +36,15 @@ def login()
     end
 
     def post_xml(host, path, xml, headers)
-      count :post
       host = host || @@INSTANCE_HOST
       if host != @@LOGIN_HOST # Not login, need to add session id to header
-        headers['X-SFDC-Session'] = @session_id;
+        headers['X-SFDC-Session'] = @session_id
         path = "#{@@PATH_PREFIX}#{path}"
       end
       i = 0
       begin
+        count :post
+        throttle(http_method: :post, path: path)
         https(host).post(path, xml, headers).body
       rescue
         i += 1
@@ -57,12 +59,14 @@ def post_xml(host, path, xml, headers)
     end
 
     def get_request(host, path, headers)
-      count :get
       host = host || @@INSTANCE_HOST
       path = "#{@@PATH_PREFIX}#{path}"
       if host != @@LOGIN_HOST # Not login, need to add session id to header
         headers['X-SFDC-Session'] = @session_id;
       end
+
+      count :get
+      throttle(http_method: :get, path: path)
       https(host).get(path, headers).body
     end
 
@@ -74,23 +78,20 @@ def https(host)
     end
 
     def parse_instance()
-      @instance=@server_url.match(/https:\/\/[a-z]{2}[0-9]{1,2}/).to_s.gsub("https://","")
+      @instance = @server_url.match(/https:\/\/[a-z]{2}[0-9]{1,2}/).to_s.gsub("https://","")
       @instance = @server_url.split(".salesforce.com")[0].split("://")[1] if @instance.nil? || @instance.empty?
       return @instance
     end
 
     def counters
-      {
-        get: @counters[:get],
-        post: @counters[:post]
-      }
+      @counters.dup
     end
 
     private
 
-    def count(name)
+    def count(http_method)
       @counters ||= Hash.new(0)
-      @counters[name] += 1
+      @counters[http_method] += 1
     end
 
   end