Add SalesforceToGcsOperator (#10760)

Adds SalesforceToGcsOperator that allows users to transfer data from Salesforce to GCS bucket. Co-authored-by: Tomek Urbaszek <[email protected]>
apache · Oct 30, 2020 · 2f703df · 2f703df
1 parent 9687b3b
commit 2f703df
Show file tree

Hide file tree

Showing 9 changed files with 493 additions and 1 deletion.
diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst
@@ -584,7 +584,7 @@ apache.hive amazon,microsoft.mssql,mysql,presto,samba,vertica
 apache.livy http
 dingding http
 discord http
-google amazon,apache.cassandra,cncf.kubernetes,facebook,microsoft.azure,microsoft.mssql,mysql,postgres,presto,sftp
+google amazon,apache.cassandra,cncf.kubernetes,facebook,microsoft.azure,microsoft.mssql,mysql,postgres,presto,salesforce,sftp
 hashicorp google
 microsoft.azure google,oracle
 microsoft.mssql odbc

diff --git a/airflow/providers/dependencies.json b/airflow/providers/dependencies.json
@@ -38,6 +38,7 @@
  "mysql",
  "postgres",
  "presto",
+ "salesforce",
  "sftp"
  ],
  "hashicorp": [

diff --git a/airflow/providers/google/cloud/example_dags/example_salesforce_to_gcs.py b/airflow/providers/google/cloud/example_dags/example_salesforce_to_gcs.py
@@ -0,0 +1,123 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you 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.
+
+"""
+Example Airflow DAG that shows how to use SalesforceToGcsOperator.
+"""
+import os
+
+from airflow import models
+from airflow.providers.google.cloud.operators.bigquery import (
+ BigQueryCreateEmptyDatasetOperator,
+ BigQueryCreateEmptyTableOperator,
+ BigQueryDeleteDatasetOperator,
+ BigQueryExecuteQueryOperator,
+)
+from airflow.providers.google.cloud.operators.gcs import GCSCreateBucketOperator, GCSDeleteBucketOperator
+from airflow.providers.google.cloud.transfers.gcs_to_bigquery import GCSToBigQueryOperator
+from airflow.providers.google.cloud.transfers.salesforce_to_gcs import SalesforceToGcsOperator
+from airflow.utils.dates import days_ago
+
+GCP_PROJECT_ID = os.environ.get("GCP_PROJECT_ID", "example-project")
+GCS_BUCKET = os.environ.get("GCS_BUCKET", "airflow-salesforce-bucket")
+DATASET_NAME = os.environ.get("SALESFORCE_DATASET_NAME", "salesforce_test_dataset")
+TABLE_NAME = os.environ.get("SALESFORCE_TABLE_NAME", "salesforce_test_datatable")
+GCS_OBJ_PATH = os.environ.get("GCS_OBJ_PATH", "results.csv")
+QUERY = "SELECT Id, Name, Company, Phone, Email, CreatedDate, LastModifiedDate, IsDeleted FROM Lead"
+GCS_CONN_ID = os.environ.get("GCS_CONN_ID", "google_cloud_default")
+SALESFORCE_CONN_ID = os.environ.get("SALESFORCE_CONN_ID", "salesforce_default")
+
+
+with models.DAG(
+ "example_salesforce_to_gcs",
+ schedule_interval=None, # Override to match your needs
+ start_date=days_ago(1),
+) as dag:
+ create_bucket = GCSCreateBucketOperator(
+ task_id="create_bucket",
+ bucket_name=GCS_BUCKET,
+ project_id=GCP_PROJECT_ID,
+ gcp_conn_id=GCS_CONN_ID,
+ )
+
+ # [START howto_operator_salesforce_to_gcs]
+ gcs_upload_task = SalesforceToGcsOperator(
+ query=QUERY,
+ include_deleted=True,
+ bucket_name=GCS_BUCKET,
+ object_name=GCS_OBJ_PATH,
+ salesforce_conn_id=SALESFORCE_CONN_ID,
+ export_format='csv',
+ coerce_to_timestamp=False,
+ record_time_added=False,
+ gcp_conn_id=GCS_CONN_ID,
+ task_id="upload_to_gcs",
+ dag=dag,
+ )
+ # [END howto_operator_salesforce_to_gcs]
+
+ create_dataset = BigQueryCreateEmptyDatasetOperator(
+ task_id="create_dataset", dataset_id=DATASET_NAME, project_id=GCP_PROJECT_ID, gcp_conn_id=GCS_CONN_ID
+ )
+
+ create_table = BigQueryCreateEmptyTableOperator(
+ task_id="create_table",
+ dataset_id=DATASET_NAME,
+ table_id=TABLE_NAME,
+ schema_fields=[
+ {'name': 'id', 'type': 'STRING', 'mode': 'NULLABLE'},
+ {'name': 'name', 'type': 'STRING', 'mode': 'NULLABLE'},
+ {'name': 'company', 'type': 'STRING', 'mode': 'NULLABLE'},
+ {'name': 'phone', 'type': 'STRING', 'mode': 'NULLABLE'},
+ {'name': 'email', 'type': 'STRING', 'mode': 'NULLABLE'},
+ {'name': 'createddate', 'type': 'STRING', 'mode': 'NULLABLE'},
+ {'name': 'lastmodifieddate', 'type': 'STRING', 'mode': 'NULLABLE'},
+ {'name': 'isdeleted', 'type': 'BOOL', 'mode': 'NULLABLE'},
+ ],
+ )
+
+ load_csv = GCSToBigQueryOperator(
+ task_id='gcs_to_bq',
+ bucket=GCS_BUCKET,
+ source_objects=[GCS_OBJ_PATH],
+ destination_project_dataset_table=f"{DATASET_NAME}.{TABLE_NAME}",
+ write_disposition='WRITE_TRUNCATE',
+ )
+
+ read_data_from_gcs = BigQueryExecuteQueryOperator(
+ task_id="read_data_from_gcs",
+ sql=f"SELECT COUNT(*) FROM `{GCP_PROJECT_ID}.{DATASET_NAME}.{TABLE_NAME}`",
+ use_legacy_sql=False,
+ )
+
+ delete_bucket = GCSDeleteBucketOperator(
+ task_id="delete_bucket",
+ bucket_name=GCS_BUCKET,
+ )
+
+ delete_dataset = BigQueryDeleteDatasetOperator(
+ task_id="delete_dataset",
+ project_id=GCP_PROJECT_ID,
+ dataset_id=DATASET_NAME,
+ delete_contents=True,
+ )
+
+ create_bucket >> gcs_upload_task >> load_csv
+ create_dataset >> create_table >> load_csv
+ load_csv >> read_data_from_gcs
+ read_data_from_gcs >> delete_bucket
+ read_data_from_gcs >> delete_dataset
diff --git a/airflow/providers/google/cloud/transfers/salesforce_to_gcs.py b/airflow/providers/google/cloud/transfers/salesforce_to_gcs.py
@@ -0,0 +1,125 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you 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.
+
+import os
+import tempfile
+from typing import Dict, Optional
+
+from airflow.models import BaseOperator
+from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.salesforce.hooks.salesforce import SalesforceHook
+
+
+class SalesforceToGcsOperator(BaseOperator):
+ """
+ Submits Salesforce query and uploads results to Google Cloud Storage
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:SalesforceToGcsOperator`
+
+ :param query: The query to make to Salesforce.
+ :type query: str
+ :param bucket_name: The bucket to upload to.
+ :type bucket_name: str
+ :param object_name: The object name to set when uploading the file.
+ :type object_name: str
+ :param salesforce_conn_id: the name of the connection that has the parameters
+ we need to connect to Salesforce.
+ :type conn_id: str
+ :param include_deleted: True if the query should include deleted records.
+ :type include_deleted: bool
+ :param query_params: Additional optional arguments
+ :type query_params: dict
+ :param export_format: Desired format of files to be exported.
+ :type export_format: str
+ :param coerce_to_timestamp: True if you want all datetime fields to be converted into Unix timestamps.
+ False if you want them to be left in the same format as they were in Salesforce.
+ Leaving the value as False will result in datetimes being strings. Default: False
+ :type coerce_to_timestamp: bool
+ :param record_time_added: 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
+ :type record_time_added: bool
+ :param gzip: Option to compress local file or file data for upload
+ :type gzip: bool
+ :param gcp_conn_id: the name of the connection that has the parameters we need to connect to GCS.
+ :type conn_id: str
+ """
+
+ template_fields = (
+ 'query',
+ 'bucket_name',
+ 'object_name',
+ )
+ template_ext = ('.sql',)
+
+ def __init__(
+ self,
+ *,
+ query: str,
+ bucket_name: str,
+ object_name: str,
+ salesforce_conn_id: str,
+ include_deleted: bool = False,
+ query_params: Optional[dict] = None,
+ export_format: str = "csv",
+ coerce_to_timestamp: bool = False,
+ record_time_added: bool = False,
+ gzip: bool = False,
+ gcp_conn_id: str = "google_cloud_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.query = query
+ self.bucket_name = bucket_name
+ self.object_name = object_name
+ self.salesforce_conn_id = salesforce_conn_id
+ self.export_format = export_format
+ self.coerce_to_timestamp = coerce_to_timestamp
+ self.record_time_added = record_time_added
+ self.gzip = gzip
+ self.gcp_conn_id = gcp_conn_id
+ self.include_deleted = include_deleted
+ self.query_params = query_params
+
+ def execute(self, context: Dict):
+ salesforce = SalesforceHook(conn_id=self.salesforce_conn_id)
+ response = salesforce.make_query(
+ query=self.query, include_deleted=self.include_deleted, query_params=self.query_params
+ )
+
+ with tempfile.TemporaryDirectory() as tmp:
+ path = os.path.join(tmp, "salesforce_temp_file")
+ salesforce.write_object_to_file(
+ query_results=response["records"],
+ filename=path,
+ fmt=self.export_format,
+ coerce_to_timestamp=self.coerce_to_timestamp,
+ record_time_added=self.record_time_added,
+ )
+
+ hook = GCSHook(gcp_conn_id=self.gcp_conn_id)
+ hook.upload(
+ bucket_name=self.bucket_name,
+ object_name=self.object_name,
+ filename=path,
+ gzip=self.gzip,
+ )
+
+ gcs_uri = "gs://{}/{}".format(self.bucket_name, self.object_name)
+ self.log.info("%s uploaded to GCS", gcs_uri)
+ return gcs_uri
diff --git a/docs/howto/operator/google/transfer/salesforce_to_gcs.rst b/docs/howto/operator/google/transfer/salesforce_to_gcs.rst
@@ -0,0 +1,49 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you 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.
+
+Salesforce To GCS Operators
+==============================
+
+.. contents::
+ :depth: 1
+ :local:
+
+Prerequisite Tasks
+^^^^^^^^^^^^^^^^^^
+
+.. include:: /howto/operator/google/_partials/prerequisite_tasks.rst
+
+.. _howto/operator:SalesforceToGcsOperator:
+
+SalesforceToGcsOperator
+------------------------------
+
+Use the
+:class:`~airflow.providers.google.cloud.transfers.salesforce_to_gcs.SalesforceToGcsOperator`
+to execute a Salesforce query to fetch data and load it to GCS.
+
+.. exampleinclude:: /../airflow/providers/google/cloud/example_dags/example_salesforce_to_gcs.py
+ :language: python
+ :start-after: [START howto_operator_salesforce_to_gcs]
+ :end-before: [END howto_operator_salesforce_to_gcs]
+
+Reference
+^^^^^^^^^
+
+For further information, look at:
+
+* `Simple Salesforce Documentation <https://simple-salesforce.readthedocs.io/en/latest/>`__
diff --git a/docs/operators-and-hooks-ref.rst b/docs/operators-and-hooks-ref.rst
@@ -992,6 +992,11 @@ These integrations allow you to copy data from/to Google Cloud.
  - :doc:`How to use <howto/operator/google/transfer/facebook_ads_to_gcs>`
  - :mod:`airflow.providers.google.cloud.transfers.facebook_ads_to_gcs`
 
+ * - `Salesforce <https://www.salesforce.com/>`__
+ - `Google Cloud Storage (GCS) <https://cloud.google.com/gcs/>`__
+ - :doc:`How to use <howto/operator/google/transfer/salesforce_to_gcs>`
+ - :mod:`airflow.providers.google.cloud.transfers.salesforce_to_gcs`
+
 
  * - `Google Ads <https://ads.google.com/>`__
  - `Google Cloud Storage (GCS) <https://cloud.google.com/gcs/>`__