diff --git a/blueprints/data-solutions/data-platform-spark/demo/orchestrate_pyspark.py b/blueprints/data-solutions/data-platform-spark/demo/orchestrate_pyspark.py new file mode 100644 index 00000000..a8e0f2d3 --- /dev/null +++ b/blueprints/data-solutions/data-platform-spark/demo/orchestrate_pyspark.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python + +# Copyright 2019 Google LLC +# +# 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. + +import datetime +import datetime +import os + +from airflow import models +from airflow.providers.google.cloud.operators.dataproc import ( + DataprocCreateBatchOperator, DataprocDeleteBatchOperator, DataprocGetBatchOperator, DataprocListBatchesOperator + +) +from airflow.utils.dates import days_ago + +# -------------------------------------------------------------------------------- +# Get variables +# -------------------------------------------------------------------------------- +BQ_LOCATION = os.environ.get("BQ_LOCATION") +CURATED_BQ_DATASET = os.environ.get("CURATED_BQ_DATASET") +CURATED_GCS = os.environ.get("CURATED_GCS") +CURATED_PRJ = os.environ.get("CURATED_PRJ") +DP_KMS_KEY = os.environ.get("DP_KMS_KEY", "") +DP_REGION = os.environ.get("DP_REGION") +GCP_REGION = os.environ.get("GCP_REGION") +LAND_PRJ = os.environ.get("LAND_PRJ") +LAND_BQ_DATASET = os.environ.get("LAND_BQ_DATASET") +LAND_GCS = os.environ.get("LAND_GCS") +PHS_NAME = os.environ.get("PHS_NAME") +PROCESSING_GCS = os.environ.get("PROCESSING_GCS") +PROCESSING_PRJ = os.environ.get("PROCESSING_PRJ") +PROCESSING_SA_DP = os.environ.get("PROCESSING_SA_DP") +PROCESSING_SA_SUBNET = os.environ.get("PROCESSING_SUBNET") +PROCESSING_SA_VPC = os.environ.get("PROCESSING_VPC") + +PYTHON_FILE_LOCATION = "gs://"+PROCESSING_GCS+"/pyspark_sort.py" +PHS_CLUSTER_PATH = "projects/"+PROCESSING_PRJ+"/regions/"+DP_REGION+"/clusters/"+PHS_NAME + +default_args = { + # Tell airflow to start one day ago, so that it runs as soon as you upload it + "start_date": days_ago(1), + "region": DP_REGION, +} +with models.DAG( + "dataproc_batch_operators", # The id you will see in the DAG airflow page + default_args=default_args, # The interval with which to schedule the DAG + schedule_interval=None, # Override to match your needs +) as dag: + + create_batch = DataprocCreateBatchOperator( + task_id="batch_create", + project_id=PROCESSING_PRJ, + batch={ + "environment_config": { + "execution_config": { + "service_account": PROCESSING_SA_DP, + "subnetwork_uri": PROCESSING_SA_SUBNET + } + }, + "pyspark_batch": { + "main_python_file_uri": PYTHON_FILE_LOCATION, + }, + "history_server_cluster": PHS_NAME, + }, + batch_id="batch-create-phs", + ) + + list_batches = DataprocListBatchesOperator( + task_id="list-all-batches", + ) + + get_batch = DataprocGetBatchOperator( + task_id="get_batch", + batch_id="batch-create-phs", + ) + + create_batch >> list_batches >> get_batch \ No newline at end of file