Repository: incubator-airflow Updated Branches: refs/heads/master f64224204 -> ec80f9441
[AIRFLOW-2124] Upload Python file to a bucket for Dataproc If the Python Dataproc file is on local storage, we want to upload this to google cloud storage before submitting it to the dataproc cluster Closes #3130 from Fokko/airflow-stash-files-on-gcs Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/ec80f944 Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/ec80f944 Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/ec80f944 Branch: refs/heads/master Commit: ec80f944183b744ff7cab7b72f5350240a8ac4ae Parents: f642242 Author: Fokko Driesprong <[email protected]> Authored: Wed Mar 21 11:46:37 2018 +0100 Committer: Fokko Driesprong <[email protected]> Committed: Wed Mar 21 11:46:37 2018 +0100 ---------------------------------------------------------------------- airflow/contrib/hooks/gcp_dataproc_hook.py | 7 +++ airflow/contrib/operators/dataproc_operator.py | 51 +++++++++++++++++++-- airflow/www/views.py | 1 - 3 files changed, 53 insertions(+), 6 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/ec80f944/airflow/contrib/hooks/gcp_dataproc_hook.py ---------------------------------------------------------------------- diff --git a/airflow/contrib/hooks/gcp_dataproc_hook.py b/airflow/contrib/hooks/gcp_dataproc_hook.py index bc5aa83..2973f34 100644 --- a/airflow/contrib/hooks/gcp_dataproc_hook.py +++ b/airflow/contrib/hooks/gcp_dataproc_hook.py @@ -201,6 +201,13 @@ class DataProcHook(GoogleCloudBaseHook): http_authorized = self._authorize() return build('dataproc', self.api_version, http=http_authorized) + def get_cluster(self, project_id, region, cluster_name): + return self.get_conn().projects().regions().clusters().get( + projectId=project_id, + region=region, + clusterName=cluster_name + ).execute(num_retries=5) + def submit(self, project_id, job, region='global'): submitted = _DataProcJob(self.get_conn(), project_id, job, region) if not submitted.wait_for_done(): http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/ec80f944/airflow/contrib/operators/dataproc_operator.py ---------------------------------------------------------------------- diff --git a/airflow/contrib/operators/dataproc_operator.py b/airflow/contrib/operators/dataproc_operator.py index 1de456a..b1164c4 100644 --- a/airflow/contrib/operators/dataproc_operator.py +++ b/airflow/contrib/operators/dataproc_operator.py @@ -13,11 +13,13 @@ # limitations under the License. # +import ntpath +import os import time import uuid - from airflow.contrib.hooks.gcp_dataproc_hook import DataProcHook +from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook from airflow.exceptions import AirflowException from airflow.models import BaseOperator from airflow.utils.decorators import apply_defaults @@ -884,6 +886,33 @@ class DataProcPySparkOperator(BaseOperator): template_fields = ['arguments', 'job_name', 'cluster_name'] ui_color = '#0273d4' + @staticmethod + def _generate_temp_filename(filename): + dt = time.strftime('%Y%m%d%H%M%S') + return "{}_{}_{}".format(dt, str(uuid.uuid1())[:8], ntpath.basename(filename)) + + """ + Upload a local file to a Google Cloud Storage bucket + """ + def _upload_file_temp(self, bucket, local_file): + temp_filename = self._generate_temp_filename(local_file) + if not bucket: + raise AirflowException( + "If you want Airflow to upload the local file to a temporary bucket, set " + "the 'temp_bucket' key in the connection string") + + self.log.info("Uploading %s to %s", local_file, temp_filename) + + GoogleCloudStorageHook( + google_cloud_storage_conn_id=self.gcp_conn_id + ).upload( + bucket=bucket, + object=temp_filename, + mime_type='application/x-python', + filename=local_file + ) + return "gs://{}/{}".format(bucket, temp_filename) + @apply_defaults def __init__( self, @@ -917,12 +946,24 @@ class DataProcPySparkOperator(BaseOperator): self.region = region def execute(self, context): - hook = DataProcHook(gcp_conn_id=self.gcp_conn_id, - delegate_to=self.delegate_to) - job = hook.create_job_template(self.task_id, self.cluster_name, "pysparkJob", - self.dataproc_properties) + hook = DataProcHook( + gcp_conn_id=self.gcp_conn_id, + delegate_to=self.delegate_to + ) + job = hook.create_job_template( + self.task_id, self.cluster_name, "pysparkJob", self.dataproc_properties) + # Check if the file is local, if that is the case, upload it to a bucket + if os.path.isfile(self.main): + cluster_info = hook.get_cluster( + project_id=hook.project_id, + region=self.region, + cluster_name=self.cluster_name + ) + bucket = cluster_info['config']['configBucket'] + self.main = self._upload_file_temp(bucket, self.main) job.set_python_main(self.main) + job.add_args(self.arguments) job.add_jar_file_uris(self.dataproc_jars) job.add_archive_uris(self.archives) http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/ec80f944/airflow/www/views.py ---------------------------------------------------------------------- diff --git a/airflow/www/views.py b/airflow/www/views.py index 44ef8e0..2f56175 100644 --- a/airflow/www/views.py +++ b/airflow/www/views.py @@ -2672,7 +2672,6 @@ class ConnectionModelView(wwwutils.SuperUserMixin, AirflowModelView): 'extra__google_cloud_platform__key_path': StringField('Keyfile Path'), 'extra__google_cloud_platform__keyfile_dict': PasswordField('Keyfile JSON'), 'extra__google_cloud_platform__scope': StringField('Scopes (comma separated)'), - } form_choices = { 'conn_type': models.Connection._types
