[Dataproc] Optimize Example DAG And Push Job ID in XCOM (#399)

Author: pankajastro

astronomer/providers/google/cloud/example_dags/example_dataproc.py

@@ -5,13 +5,8 @@
 
 from airflow import models
 from airflow.providers.google.cloud.operators.dataproc import (
-    ClusterGenerator,
     DataprocCreateClusterOperator,
-    DataprocCreateWorkflowTemplateOperator,
     DataprocDeleteClusterOperator,
-    DataprocInstantiateInlineWorkflowTemplateOperator,
-    DataprocInstantiateWorkflowTemplateOperator,
-    DataprocUpdateClusterOperator,
 )
 from airflow.providers.google.cloud.operators.gcs import (
     GCSCreateBucketOperator,
@@ -40,49 +35,10 @@
         "machine_type_uri": "n1-standard-4",
         "disk_config": {"boot_disk_type": "pd-standard", "boot_disk_size_gb": 1024},
     },
-    "worker_config": {
-        "num_instances": 2,
-        "machine_type_uri": "n1-standard-4",
-        "disk_config": {"boot_disk_type": "pd-standard", "boot_disk_size_gb": 1024},
-    },
 }
 
 # [END how_to_cloud_dataproc_create_cluster]
 
-# Cluster definition: Generating Cluster Config for DataprocCreateClusterOperator
-# [START how_to_cloud_dataproc_create_cluster_generate_cluster_config]
-path = "gs://goog-dataproc-initialization-actions-us-central1/python/pip-install.sh"
-
-CLUSTER_GENERATOR_CONFIG = ClusterGenerator(
-    project_id="test",
-    zone="us-central1-a",
-    master_machine_type="n1-standard-4",
-    worker_machine_type="n1-standard-4",
-    num_workers=2,
-    storage_bucket="test",
-    init_actions_uris=[path],
-    metadata={"PIP_PACKAGES": "pyyaml requests pandas openpyxl"},
-).make()
-
-create_cluster_operator = DataprocCreateClusterOperator(
-    task_id="create_dataproc_cluster",
-    cluster_name="test",
-    project_id="test",
-    region="us-central1",
-    cluster_config=CLUSTER_GENERATOR_CONFIG,
-)
-# [END how_to_cloud_dataproc_create_cluster_generate_cluster_config]
-
-# Update options
-# [START how_to_cloud_dataproc_updatemask_cluster_operator]
-CLUSTER_UPDATE = {
-    "config": {"worker_config": {"num_instances": 3}, "secondary_worker_config": {"num_instances": 3}}
-}
-UPDATE_MASK = {
-    "paths": ["config.worker_config.num_instances", "config.secondary_worker_config.num_instances"]
-}
-# [END how_to_cloud_dataproc_updatemask_cluster_operator]
-
 TIMEOUT = {"seconds": 1 * 24 * 60 * 60}
 
 # Jobs definitions
@@ -183,39 +139,6 @@
     )
     # [END howto_create_bucket_task]
 
-    # [START how_to_cloud_dataproc_update_cluster_operator]
-    scale_cluster = DataprocUpdateClusterOperator(
-        task_id="scale_cluster",
-        cluster_name=CLUSTER_NAME,
-        cluster=CLUSTER_UPDATE,
-        update_mask=UPDATE_MASK,
-        graceful_decommission_timeout=TIMEOUT,
-        project_id=PROJECT_ID,
-        region=REGION,
-    )
-    # [END how_to_cloud_dataproc_update_cluster_operator]
-
-    # [START how_to_cloud_dataproc_create_workflow_template]
-    create_workflow_template = DataprocCreateWorkflowTemplateOperator(
-        task_id="create_workflow_template",
-        template=WORKFLOW_TEMPLATE,
-        project_id=PROJECT_ID,
-        region=REGION,
-    )
-    # [END how_to_cloud_dataproc_create_workflow_template]
-
-    # [START how_to_cloud_dataproc_trigger_workflow_template]
-    trigger_workflow = DataprocInstantiateWorkflowTemplateOperator(
-        task_id="trigger_workflow", region=REGION, project_id=PROJECT_ID, template_id=WORKFLOW_NAME
-    )
-    # [END how_to_cloud_dataproc_trigger_workflow_template]
-
-    # [START how_to_cloud_dataproc_instantiate_inline_workflow_template]
-    instantiate_inline_workflow_template = DataprocInstantiateInlineWorkflowTemplateOperator(
-        task_id="instantiate_inline_workflow_template", template=WORKFLOW_TEMPLATE, region=REGION
-    )
-    # [END how_to_cloud_dataproc_instantiate_inline_workflow_template]
-
     # [START howto_DataprocSubmitJobOperatorAsync]
     pig_task = DataprocSubmitJobOperatorAsync(
         task_id="pig_task", job=PIG_JOB, region=REGION, project_id=PROJECT_ID
@@ -243,20 +166,21 @@
     # [END howto_DataprocSubmitJobOperatorAsync]
     # [START how_to_cloud_dataproc_delete_cluster_operator]
     delete_cluster = DataprocDeleteClusterOperator(
-        task_id="delete_cluster", project_id=PROJECT_ID, cluster_name=CLUSTER_NAME, region=REGION
+        task_id="delete_cluster",
+        project_id=PROJECT_ID,
+        cluster_name=CLUSTER_NAME,
+        region=REGION,
+        trigger_rule="all_done",
     )
     # [END how_to_cloud_dataproc_delete_cluster_operator]
     # [START howto_delete_buckettask]
     delete_bucket = GCSDeleteBucketOperator(
         task_id="delete_bucket",
         bucket_name=BUCKET,
+        trigger_rule="all_done",
     )
     # [END howto_delete_buckettask]
 
-    create_cluster >> scale_cluster >> create_bucket
-    scale_cluster >> create_workflow_template >> trigger_workflow >> delete_cluster
-    scale_cluster >> hive_task >> delete_cluster >> delete_bucket
-    scale_cluster >> pig_task >> delete_cluster >> delete_bucket
-    scale_cluster >> spark_sql_task >> delete_cluster >> delete_bucket
-    scale_cluster >> spark_task >> delete_cluster >> delete_bucket
-    scale_cluster >> hadoop_task >> delete_cluster >> delete_bucket
+    create_cluster >> create_bucket
+    create_cluster >> pig_task >> hive_task >> delete_cluster >> delete_bucket
+    create_cluster >> spark_task >> spark_sql_task >> hadoop_task >> delete_cluster >> delete_bucket

astronomer/providers/google/cloud/operators/dataproc.py

@@ -81,7 +81,7 @@ def execute_complete(self, context: Dict[str, Any], event: Optional[Dict[str, st
         """
         if event:
             if event["status"] == "success":
-                self.log.debug("Job %s completed successfully.", self.job_id)
-                return event["message"]
+                self.log.info("Job %s completed successfully.", event["job_id"])
+                return event["job_id"]
             raise AirflowException(event["message"])
         raise AirflowException("No event received in trigger callback")

astronomer/providers/google/cloud/triggers/dataproc.py

@@ -28,9 +28,8 @@ def __init__(
         project_id: Optional[str] = None,
         gcp_conn_id: str = "google_cloud_default",
         polling_interval: float = 5.0,
-        **kwargs: Any,
     ) -> None:
-        super().__init__(**kwargs)
+        super().__init__()
         self.project_id = project_id
         self.gcp_conn_id = gcp_conn_id
         self.dataproc_job_id = dataproc_job_id
@@ -70,15 +69,19 @@ async def _get_job_status(self, hook: DataprocHookAsync) -> Dict[str, str]:
         job = await hook.get_job(job_id=self.dataproc_job_id, region=self.region, project_id=self.project_id)
         state = job.status.state
         if state == JobStatus.State.ERROR:
-            return {"status": "error", "message": "Job Failed"}
+            return {"status": "error", "message": "Job Failed", "job_id": self.dataproc_job_id}
         elif state in {
             JobStatus.State.CANCELLED,
             JobStatus.State.CANCEL_PENDING,
             JobStatus.State.CANCEL_STARTED,
         }:
-            return {"status": "error", "message": "Job got cancelled"}
+            return {"status": "error", "message": "Job got cancelled", "job_id": self.dataproc_job_id}
         elif JobStatus.State.DONE == state:
-            return {"status": "success", "message": "Job completed successfully"}
+            return {
+                "status": "success",
+                "message": "Job completed successfully",
+                "job_id": self.dataproc_job_id,
+            }
         elif JobStatus.State.ATTEMPT_FAILURE == state:
-            return {"status": "pending", "message": "Job is in pending state"}
-        return {"status": "pending", "message": "Job is in pending state"}
+            return {"status": "pending", "message": "Job is in pending state", "job_id": self.dataproc_job_id}
+        return {"status": "pending", "message": "Job is in pending state", "job_id": self.dataproc_job_id}

tests/google/cloud/operators/test_dataproc.py

@@ -52,7 +52,7 @@ def test_dataproc_operator_execute_async(mock_submit_job):
     "event",
     [
         ({"status": "error", "message": "test failure message"}),
-        (None),
+        None,
     ],
 )
 @mock.patch("airflow.providers.google.cloud.operators.dataproc.DataprocHook.submit_job")
@@ -73,4 +73,6 @@ def test_dataproc_operator_execute_success_async(mock_submit_job):
     task = DataprocSubmitJobOperatorAsync(
         task_id="task-id", job=SPARK_JOB, region=TEST_REGION, project_id=TEST_PROJECT_ID
     )
-    assert task.execute_complete(context=None, event={"status": "success", "message": "success"})
+    assert task.execute_complete(
+        context=None, event={"status": "success", "message": "success", "job_id": TEST_JOB_ID}
+    )

tests/google/cloud/triggers/test_dataproc.py

@@ -106,11 +106,23 @@ async def test_dataproc_submit_return_exception(mock_get_job_status):
 @pytest.mark.parametrize(
     "state, response",
     [
-        (JobStatus.State.DONE, {"status": "success", "message": "Job completed successfully"}),
-        (JobStatus.State.ERROR, {"status": "error", "message": "Job Failed"}),
-        (JobStatus.State.CANCELLED, {"status": "error", "message": "Job got cancelled"}),
-        (JobStatus.State.ATTEMPT_FAILURE, {"status": "pending", "message": "Job is in pending state"}),
-        (JobStatus.State.SETUP_DONE, {"status": "pending", "message": "Job is in pending state"}),
+        (
+            JobStatus.State.DONE,
+            {"status": "success", "message": "Job completed successfully", "job_id": TEST_JOB_ID},
+        ),
+        (JobStatus.State.ERROR, {"status": "error", "message": "Job Failed", "job_id": TEST_JOB_ID}),
+        (
+            JobStatus.State.CANCELLED,
+            {"status": "error", "message": "Job got cancelled", "job_id": TEST_JOB_ID},
+        ),
+        (
+            JobStatus.State.ATTEMPT_FAILURE,
+            {"status": "pending", "message": "Job is in pending state", "job_id": TEST_JOB_ID},
+        ),
+        (
+            JobStatus.State.SETUP_DONE,
+            {"status": "pending", "message": "Job is in pending state", "job_id": TEST_JOB_ID},
+        ),
     ],
 )
 async def test_dataproc_get_job_status(state, response):