def test_execute(self, mock_hook): op = DataprocDeleteClusterOperator( task_id=TASK_ID, region=GCP_LOCATION, project_id=GCP_PROJECT, cluster_name=CLUSTER_NAME, request_id=REQUEST_ID, gcp_conn_id=GCP_CONN_ID, retry=RETRY, timeout=TIMEOUT, metadata=METADATA, impersonation_chain=IMPERSONATION_CHAIN, ) op.execute(context={}) mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN) mock_hook.return_value.delete_cluster.assert_called_once_with( region=GCP_LOCATION, project_id=GCP_PROJECT, cluster_name=CLUSTER_NAME, cluster_uuid=None, request_id=REQUEST_ID, retry=RETRY, timeout=TIMEOUT, metadata=METADATA, )
pyspark_task = DataprocSubmitJobOperator(task_id="pyspark_task", job=PYSPARK_JOB, location=REGION, project_id=PROJECT_ID) hive_task = DataprocSubmitJobOperator(task_id="hive_task", job=HIVE_JOB, location=REGION, project_id=PROJECT_ID) hadoop_task = DataprocSubmitJobOperator(task_id="hadoop_task", job=HADOOP_JOB, location=REGION, project_id=PROJECT_ID) delete_cluster = DataprocDeleteClusterOperator( task_id="delete_cluster", project_id=PROJECT_ID, cluster_name=CLUSTER_NAME, region=REGION, ) create_cluster >> scale_cluster scale_cluster >> hive_task >> delete_cluster scale_cluster >> pig_task >> delete_cluster scale_cluster >> spark_sql_task >> delete_cluster scale_cluster >> spark_task >> delete_cluster scale_cluster >> pyspark_task >> delete_cluster scale_cluster >> hadoop_task >> delete_cluster
cluster_name=CLUSTER_NAME, cluster_config=CLUSTER_CONFIG, labels={"tenant": TENANT, "created-by": USER, }, ) pr_producer_pyspark_task = DataprocSubmitJobOperator( task_id="run_cee_pr_producer", impersonation_chain=CONNECT_SA, job=PR_PYSPARK_JOB, location=REGION, cluster_name=CLUSTER_NAME, cluster_config=CLUSTER_CONFIG, labels={"tenant": TENANT, "created-by": USER, }, ) # Delete Dataproc cluster. delete_dataproc_cluster = DataprocDeleteClusterOperator( task_id="delete_dataproc_cluster", impersonation_chain=CONNECT_SA, region=REGION, cluster_name=CLUSTER_NAME, labels={"tenant": TENANT, "created-by": USER, }, ) ( create_dataproc_cluster >> assign_permissions >> [cg_producer_pyspark_task, pr_producer_pyspark_task] >> delete_dataproc_cluster )
default_args = { 'owner': 'airflow', 'depends_on_past': True, 'start_date': datetime(2021, 3, 9, tzinfo=local_tz), 'email': ['*****@*****.**'], 'email_on_failure': True, 'email_on_retry': True, 'retries': 1, 'retry_delay': timedelta(minutes=1), 'project_id': gcp_config['project_id'], 'region': gcp_config['region'], 'gcp_conn_id': gcp_config['conn_id'] } with DAG( 'delete_dataproc', default_args=default_args, description='delete_dataproc', schedule_interval='@once' ) as dag: delete_dataproc = DataprocDeleteClusterOperator( task_id='delete_dataproc', cluster_name=dataproc_config['cluster_name'] ) delete_dataproc
) pyspark_job_task_predict = DataprocSubmitJobOperator( task_id="pyspark_job_task", job=get_pyspark_job_config( predict_config_script['script_path'], predict_config_script['config_path_daily'], CLUSTER_NAME_DAILY ), location=REGION, project_id=PROJECT_ID ) delete_cluster = DataprocDeleteClusterOperator( task_id="delete_cluster", project_id=PROJECT_ID, cluster_name=CLUSTER_NAME_DAILY, region=REGION ) delete_cluster_on_failure = DataprocDeleteClusterOperator( task_id="delete_cluster", project_id=PROJECT_ID, cluster_name=CLUSTER_NAME_DAILY, region=REGION, trigger_rule=TriggerRule.ONE_FAILED ) bq_delete = BigQueryOperator( task_id="delete_record_of_execution_date", sql=delete_query, use_legacy_sql=False,
start_pipelines.append( PythonOperator( task_id=str(x) + "_" + DF_PIPELINES[x], python_callable=start_pipeline_function, dag=dag, pool=DF_PIPELINE_POOL, op_kwargs={"pipeline": DF_PIPELINES[x]}, trigger_rule="all_done", )) delete_firewall_rule = PythonOperator( task_id="delete_firewall_rule", provide_context=True, python_callable=remove_firewall_function, dag=dag, trigger_rule="all_done", ) delete_dataproc_cluster = DataprocDeleteClusterOperator( task_id="delete_dataproc_cluster", project_id="data-analytics-webinar", cluster_name="etl-cluster", region="us-central1", trigger_rule="all_done", ) # Define DAG dependencies. for x in range(len(DF_PIPELINES)): create_dataproc_cluster >> create_firewall_rule >> start_pipelines[ x] >> delete_firewall_rule >> delete_dataproc_cluster