Add support for multiple results in salesforce batch jobs (#1686)

* #1685 Add support for multiple results in batch jobs
Adds a new method 'get_batch_result_ids' and maintains backwards compatibility
with the old method, adding a warning.

* Add the ability to merge result sets and basic testing.
Note that this only supports CSV since this is what the library does by default.

luigi/contrib/salesforce.py

@@ -17,6 +17,7 @@
 import time
 import abc
 import logging
+import warnings
 import xml.etree.ElementTree as ET
 from collections import OrderedDict
 import re
@@ -139,6 +140,11 @@ def is_soql_file(self):
  """Override to True if soql property is a file path."""
  return False
 
+ @property
+ def content_type(self):
+ """Override to use a different content type. (e.g. XML)"""
+ return "CSV"
+
  def run(self):
  if self.use_sandbox and not self.sandbox_name:
  raise Exception("Parameter sf_sandbox_name must be provided when uploading to a Salesforce Sandbox")
@@ -170,11 +176,23 @@ def run(self):
  if 'foreign key relationships not supported' not in status['state_message'].lower():
  raise Exception(msg)
  else:
- result_id = sf.get_batch_results(job_id, batch_id)
- data = sf.get_batch_result(job_id, batch_id, result_id)
-
- with open(self.output().fn, 'w') as outfile:
- outfile.write(data)
+ result_ids = sf.get_batch_result_ids(job_id, batch_id)
+
+ # If there's only one result, just download it, otherwise we need to merge the resulting downloads
+ if len(result_ids) == 1:
+ data = sf.get_batch_result(job_id, batch_id, result_ids[0])
+ with open(self.output().path, 'w') as outfile:
+ outfile.write(data)
+ else:
+ # Download each file to disk, and then merge into one.
+ # Preferring to do it this way so as to minimize memory consumption.
+ for i, result_id in enumerate(result_ids):
+ logger.info("Downloading batch result %s for batch: %s and job: %s" % (result_id, batch_id, job_id))
+ with open("%s.%d" % (self.output().path, i), 'w') as outfile:
+ outfile.write(sf.get_batch_result(job_id, batch_id, result_id))
+
+ logger.info("Merging results of batch %s" % batch_id)
+ self.merge_batch_results(result_ids)
  finally:
  logger.info("Closing job %s" % job_id)
  sf.close_job(job_id)
@@ -184,11 +202,30 @@ def run(self):
  data_file = sf.query_all(self.soql)
 
  reader = csv.reader(data_file)
- with open(self.output().fn, 'w') as outfile:
+ with open(self.output().path, 'w') as outfile:
  writer = csv.writer(outfile, dialect='excel')
  for row in reader:
  writer.writerow(row)
 
+ def merge_batch_results(self, result_ids):
+ """
+ Merges the resulting files of a multi-result batch bulk query.
+ """
+ outfile = open(self.output().path, 'w')
+
+ if self.content_type == 'CSV':
+ for i, result_id in enumerate(result_ids):
+ with open("%s.%d" % (self.output().path, i), 'r') as f:
+ header = f.readline()
+ if i == 0:
+ outfile.write(header)
+ for line in f:
+ outfile.write(line)
+ else:
+ raise Exception("Batch result merging not implemented for %s" % self.content_type)
+
+ outfile.close()
+
 
 class SalesforceAPI(object):
  """
@@ -353,15 +390,17 @@ def restful(self, path, params):
  else:
  return json_result
 
- def create_operation_job(self, operation, obj, external_id_field_name=None, content_type='CSV'):
+ def create_operation_job(self, operation, obj, external_id_field_name=None, content_type=None):
  """
  Creates a new SF job that for doing any operation (insert, upsert, update, delete, query)
 
  :param operation: delete, insert, query, upsert, update, hardDelete. Must be lowercase.
  :param obj: Parent SF object
  :param external_id_field_name: Optional.
- :param content_type: XML, CSV, ZIP_CSV, or ZIP_XML. Defaults to CSV
  """
+ if content_type is None:
+ content_type = self.content_type
+
  if not self.has_active_session():
  self.start_session()
 
@@ -419,7 +458,7 @@ def close_job(self, job_id):
 
  return response
 
- def create_batch(self, job_id, data, file_type='csv'):
+ def create_batch(self, job_id, data, file_type=None):
  """
  Creates a batch with either a string of data or a file containing data.
 
@@ -429,13 +468,15 @@ def create_batch(self, job_id, data, file_type='csv'):
 
  :param job_id: job_id as returned by 'create_operation_job(...)'
  :param data:
- :param file_type:
 
  :return: Returns batch_id
  """
  if not job_id or not self.has_active_session():
  raise Exception("Can not create a batch without a valid job_id and an active session.")
 
+ if file_type is None:
+ file_type = self.content_type.lower()
+
  headers = self._get_create_batch_content_headers(file_type)
  headers['Content-Length'] = len(data)
 
@@ -473,22 +514,27 @@ def block_on_batch(self, job_id, batch_id, sleep_time_seconds=5, max_wait_time_s
 
  def get_batch_results(self, job_id, batch_id):
  """
- Get results of a batch that has completed processing.
- If the batch is a CSV file, the response is in CSV format.
- If the batch is an XML file, the response is in XML format.
+ DEPRECATED: Use `get_batch_result_ids`
+ """
+ warnings.warn("get_batch_results is deprecated and only returns one batch result. Please use get_batch_result_ids")
+ return self.get_batch_result_ids(job_id, batch_id)[0]
+
+ def get_batch_result_ids(self, job_id, batch_id):
+ """
+ Get result IDs of a batch that has completed processing.
 
  :param job_id: job_id as returned by 'create_operation_job(...)'
  :param batch_id: batch_id as returned by 'create_batch(...)'
- :return: batch result response as either CSV or XML, dependent on the batch
+ :return: list of batch result IDs to be used in 'get_batch_result(...)'
  """
  response = requests.get(self._get_batch_results_url(job_id, batch_id),
  headers=self._get_batch_info_headers())
  response.raise_for_status()
 
  root = ET.fromstring(response.text)
- result = root.find('%sresult' % self.API_NS).text
+ result_ids = [r.text for r in root.findall('%sresult' % self.API_NS)]
 
- return result
+ return result_ids
 
  def get_batch_result(self, job_id, batch_id, result_id):
  """

test/contrib/salesforce_test.py

@@ -0,0 +1,117 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright (c) 2016 Simply Measured
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not
+# use this file except in compliance with the License. You may obtain a copy of
+# the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations under
+# the License.
+#
+# This method will be used by the mock to replace requests.get
+
+"""
+Unit test for the Salesforce contrib package
+"""
+
+from luigi.contrib.salesforce import SalesforceAPI, QuerySalesforce
+
+from helpers import unittest
+import mock
+from luigi.mock import MockTarget
+from luigi.six import PY3
+import re
+
+
+def mocked_requests_get(*args, **kwargs):
+ class MockResponse:
+ def __init__(self, body, status_code):
+ self.body = body
+ self.status_code = status_code
+
+ @property
+ def text(self):
+ return self.body
+
+ def raise_for_status(self):
+ return None
+
+ result_list = (
+ '<result-list xmlns="http://www.force.com/2009/06/asyncapi/dataload">'
+ '<result>1234</result><result>1235</result><result>1236</result>'
+ '</result-list>'
+ )
+ return MockResponse(result_list, 200)
+
+# Keep open around so we can use it in the mock responses
+old__open = open
+
+
+def mocked_open(*args, **kwargs):
+ if re.match("job_data", args[0]):
+ return MockTarget(args[0]).open(args[1])
+ else:
+ return old__open(*args)
+
+
+class TestSalesforceAPI(unittest.TestCase):
+ # We patch 'requests.get' with our own method. The mock object is passed in to our test case method.
+ @mock.patch('requests.get', side_effect=mocked_requests_get)
+ def test_deprecated_results(self, mock_get):
+ sf = SalesforceAPI('xx', 'xx', 'xx')
+ result_id = sf.get_batch_results('job_id', 'batch_id')
+ self.assertEqual('1234', result_id)
+
+ @mock.patch('requests.get', side_effect=mocked_requests_get)
+ def test_result_ids(self, mock_get):
+ sf = SalesforceAPI('xx', 'xx', 'xx')
+ result_ids = sf.get_batch_result_ids('job_id', 'batch_id')
+ self.assertEqual(['1234', '1235', '1236'], result_ids)
+
+
+class TestQuerySalesforce(QuerySalesforce):
+ def output(self):
+ return MockTarget('job_data.csv')
+
+ @property
+ def object_name(self):
+ return 'dual'
+
+ @property
+ def soql(self):
+ return "SELECT * FROM %s" % self.object_name
+
+
+class TestSalesforceQuery(unittest.TestCase):
+ patch_name = '__builtin__.open'
+ if PY3:
+ patch_name = 'builtins.open'
+
+ @mock.patch(patch_name, side_effect=mocked_open)
+ def setUp(self, mock_open):
+ MockTarget.fs.clear()
+ self.result_ids = ['a', 'b', 'c']
+
+ counter = 1
+ self.all_lines = "Lines\n"
+ self.header = "Lines"
+ for i, id in enumerate(self.result_ids):
+ filename = "%s.%d" % ('job_data.csv', i)
+ with MockTarget(filename).open('w') as f:
+ line = "%d line\n%d line" % ((counter), (counter+1))
+ f.write(self.header + "\n" + line + "\n")
+ self.all_lines += line+"\n"
+ counter += 2
+
+ @mock.patch(patch_name, side_effect=mocked_open)
+ def test_multi_csv_download(self, mock_open):
+ qsf = TestQuerySalesforce()
+
+ qsf.merge_batch_results(self.result_ids)
+ self.assertEqual(MockTarget(qsf.output().path).open('r').read(), self.all_lines)