add SalesforceToS3 Operator

Author: benjamingregory

operators/salesforce_to_s3_operator.py

@@ -1,4 +1,13 @@
+from tempfile import NamedTemporaryFile
+import logging
+import json
+
+from airflow.utils.decorators import apply_defaults
 from airflow.models import BaseOperator
+from airflow.hooks.S3_hook import S3Hook
+
+from airflow.contrib.hooks.salesforce_hook import SalesforceHook
+
 
 class SalesforceBulkQueryToS3Operator(BaseOperator):
     """
@@ -30,7 +39,6 @@ def __init__(self, sf_conn_id, soql, object_type, #sf config
 
     def execute(self, context):
         sf_conn = SalesforceHook(self.sf_conn_id).get_conn()
-        s3_conn = S3Hook(self.s3_conn_id)
 
         logging.info(self.soql)
         query_results = sf_conn.bulk.__getattr__(self.object).query(self.soql)
@@ -40,4 +48,172 @@ def execute(self, context):
         query_results = [json.dumps(result, ensure_ascii=False) for result in query_results]
         query_results = '\n'.join(query_results)
 
-        s3.load_string(query_results, self.s3_key, bucket_name=self.s3_bucket, replace=True)
+        s3.load_string(query_results, self.s3_key, bucket_name=self.s3_bucket, replace=True)
+
+
+class SalesforceToS3Operator(BaseOperator):
+    """
+    Make a query against Salesforce and write the resulting data to a file.
+    """
+    template_fields = ("s3_key",
+                       "query")
+
+    @apply_defaults
+    def __init__(
+        self,
+        sf_conn_id,
+        sf_obj,
+        s3_conn_id,
+        s3_bucket,
+        s3_key,
+        sf_fields=None,
+        fmt="csv",
+        query=None,
+        relationship_object=None,
+        record_time_added=False,
+        coerce_to_timestamp=False,
+        *args,
+        **kwargs
+    ):
+        """
+        Initialize the operator
+        :param sf_conn_id:          Name of the Airflow connection that has
+                                    the following information:
+                                        - username
+                                        - password
+                                        - security_token
+        :param sf_obj:              Name of the relevant Salesforce object
+        :param s3_conn_id:          The destination s3 connection id.
+        :type s3_conn_id:           string
+        :param s3_bucket:           The destination s3 bucket.
+        :type s3_bucket:            string
+        :param s3_key:              The destination s3 key.
+        :type s3_key:               string
+        :param sf_fields:           *(optional)* list of fields that you want
+                                    to get from the object.
+                                    If *None*, then this will get all fields
+                                    for the object
+        :param fmt:                 *(optional)* format that the s3_key of the
+                                    data should be in. Possible values include:
+                                        - csv
+                                        - json
+                                        - ndjson
+                                    *Default: csv*
+        :param query:               *(optional)* A specific query to run for
+                                    the given object.  This will override
+                                    default query creation.
+                                    *Default: None*
+        :param relationship_object: *(optional)* Some queries require
+                                    relationship objects to work, and
+                                    these are not the same names as
+                                    the SF object.  Specify that
+                                    relationship object here.
+                                    *Default: None*
+        :param record_time_added:   *(optional)* True if you want to add a
+                                    Unix timestamp field to the resulting data
+                                    that marks when the data was
+                                    fetched from Salesforce.
+                                    *Default: False*.
+        :param coerce_to_timestamp: *(optional)* True if you want to convert
+                                    all fields with dates and datetimes
+                                    into Unix timestamp (UTC).
+                                    *Default: False*.
+        """
+
+        super(SalesforceToS3Operator, self).__init__(*args, **kwargs)
+
+        self.sf_conn_id = sf_conn_id
+        self.object = sf_obj
+        self.fields = sf_fields
+        self.s3_conn_id = s3_conn_id
+        self.s3_bucket = s3_bucket
+        self.s3_key = s3_key
+        self.fmt = fmt.lower()
+        self.query = query
+        self.relationship_object = relationship_object
+        self.record_time_added = record_time_added
+        self.coerce_to_timestamp = coerce_to_timestamp
+
+    def special_query(self, query, sf_hook, relationship_object=None):
+        if not query:
+            raise ValueError("Query is None.  Cannot query nothing")
+
+        sf_hook.sign_in()
+
+        results = sf_hook.make_query(query)
+        if relationship_object:
+            records = []
+            for r in results['records']:
+                if r.get(relationship_object, None):
+                    records.extend(r[relationship_object]['records'])
+            results['records'] = records
+
+        return results
+
+    def execute(self, context):
+        """
+        Execute the operator.
+        This will get all the data for a particular Salesforce model
+        and write it to a file.
+        """
+        logging.info("Prepping to gather data from Salesforce")
+
+        # Open a name temporary file to store output file until S3 upload
+        with NamedTemporaryFile("w") as tmp:
+
+            # Load the SalesforceHook
+            hook = SalesforceHook(conn_id=self.sf_conn_id, output=tmp.name)
+
+            # Attempt to login to Salesforce
+            # If this process fails, it will raise an error and die.
+            try:
+                hook.sign_in()
+            except:
+                logging.debug('Unable to login.')
+
+            # Get object from Salesforce
+            # If fields were not defined, all fields are pulled.
+            if not self.fields:
+                self.fields = hook.get_available_fields(self.object)
+
+            logging.info(
+                "Making request for "
+                "{0} fields from {1}".format(len(self.fields), self.object)
+            )
+
+            if self.query:
+                query = self.special_query(self.query,
+                                           hook,
+                                           relationship_object=self.relationship_object
+                                           )
+            else:
+                query = hook.get_object_from_salesforce(self.object,
+                                                        self.fields)
+
+            # output the records from the query to a file
+            # the list of records is stored under the "records" key
+            logging.info("Writing query results to: {0}".format(tmp.name))
+
+            hook.write_object_to_file(query['records'],
+                                      filename=tmp.name,
+                                      fmt=self.fmt,
+                                      coerce_to_timestamp=self.coerce_to_timestamp,
+                                      record_time_added=self.record_time_added)
+
+            # Flush the temp file and upload temp file to S3
+            tmp.flush()
+
+            dest_s3 = S3Hook(s3_conn_id=self.s3_conn_id)
+
+            dest_s3.load_file(
+                filename=tmp.name,
+                key=self.output,
+                bucket_name=self.s3_bucket,
+                replace=True
+            )
+
+            dest_s3.connection.close()
+
+            tmp.close()
+
+        logging.info("Query finished!")