Repository: apache/airflow
Pull request: #15245 - passing custom image family name to the DataProcClusterCreateoperator
Closes: #15245

This is my first PR. I tried to follow the CONTRIBUTING guide step by step. It passed all the pre-commit tests.

My PR adds a new argument to  **DataprocCreateClusterOperator**  which allows user to pass **custom image family** while creating **GCP DataProc** cluster. any advice or suggestions are welcome.




Problem statement:
Passing Custom Image Family Name to the DataprocClusterCreateOperator()
**Description**
Currently, we can only pass custom Image name to  **DataprocClusterCreateOperator(),**
as the custom image expires after 60 days, we either need to update the image or we need to pass the expiration token.

Functionality is already available in **gcloud** and **REST**.

`gcloud dataproc clusters test_cluster ......
--image projects/<custom_image_project_id>/global/images/family/<family_name>
.......
`


**Use case / motivation**
The user should be able to pass either Custom Image or Custom Image family name,
so we don't have to update the image up on expiration or use expiration token.


**Are you willing to submit a PR?**
Yes

**Related Issues**
None

Hints from issue discussion:
Thanks for opening your first issue here! Be sure to follow the issue template!

I will work on this, please let me know any suggestions or comments.
There is workaround to this, we can just pass **family/<family_name>** in custom_image parameter to DataprocClusterCreateOperator(), taking advantage of because the way **custom_image_url** is created.
we will get the exact URL for **custom_image_family_name**.

![image](https://user-images.githubusercontent.com/16856802/113821798-3717da80-979a-11eb-87d4-1e9f48f711cc.png)

But I really think there should be a dedicated parameter **custom_image_family** for avoiding any confusion, and this substitution might not work in the future.

Commit message: [Airflow-15245] - sanity check added for image_version and custom_image_family.

Patch to review:
diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py
index d578565b06526..144ccad9bdc20 100644
--- a/airflow/providers/google/cloud/operators/dataproc.py
+++ b/airflow/providers/google/cloud/operators/dataproc.py
@@ -75,6 +75,10 @@ class ClusterGenerator:
     :param custom_image_project_id: project id for the custom Dataproc image, for more info see
         https://cloud.google.com/dataproc/docs/guides/dataproc-images
     :type custom_image_project_id: str
+    :param custom_image_family: family for the custom Dataproc image,
+        family name can be provide using --family flag while creating custom image, for more info see
+        https://cloud.google.com/dataproc/docs/guides/dataproc-images
+    :type custom_image_family: str
     :param autoscaling_policy: The autoscaling policy used by the cluster. Only resource names
         including projectid and location (region) are valid. Example:
         ``projects/[projectId]/locations/[dataproc_region]/autoscalingPolicies/[policy_id]``
@@ -163,6 +167,7 @@ def __init__(
         metadata: Optional[Dict] = None,
         custom_image: Optional[str] = None,
         custom_image_project_id: Optional[str] = None,
+        custom_image_family: Optional[str] = None,
         image_version: Optional[str] = None,
         autoscaling_policy: Optional[str] = None,
         properties: Optional[Dict] = None,
@@ -194,6 +199,7 @@ def __init__(
         self.metadata = metadata
         self.custom_image = custom_image
         self.custom_image_project_id = custom_image_project_id
+        self.custom_image_family = custom_image_family
         self.image_version = image_version
         self.properties = properties or {}
         self.optional_components = optional_components
@@ -220,6 +226,12 @@ def __init__(
         if self.custom_image and self.image_version:
             raise ValueError("The custom_image and image_version can't be both set")
 
+        if self.custom_image_family and self.image_version:
+            raise ValueError("The image_version and custom_image_family can't be both set")
+
+        if self.custom_image_family and self.custom_image:
+            raise ValueError("The custom_image and custom_image_family can't be both set")
+
         if self.single_node and self.num_preemptible_workers > 0:
             raise ValueError("Single node cannot have preemptible workers.")
 
@@ -346,6 +358,16 @@ def _build_cluster_data(self):
             if not self.single_node:
                 cluster_data['worker_config']['image_uri'] = custom_image_url
 
+        elif self.custom_image_family:
+            project_id = self.custom_image_project_id or self.project_id
+            custom_image_url = (
+                'https://www.googleapis.com/compute/beta/projects/'
+                '{}/global/images/family/{}'.format(project_id, self.custom_image_family)
+            )
+            cluster_data['master_config']['image_uri'] = custom_image_url
+            if not self.single_node:
+                cluster_data['worker_config']['image_uri'] = custom_image_url
+
         cluster_data = self._build_gce_cluster_config(cluster_data)
 
         if self.single_node:
diff --git a/tests/providers/google/cloud/operators/test_dataproc.py b/tests/providers/google/cloud/operators/test_dataproc.py
index e66acb4afb0d3..01a907f4bd2cd 100644
--- a/tests/providers/google/cloud/operators/test_dataproc.py
+++ b/tests/providers/google/cloud/operators/test_dataproc.py
@@ -144,6 +144,26 @@ def test_image_version(self):
             )
             assert "custom_image and image_version" in str(ctx.value)
 
+    def test_custom_image_family_1(self):
+        with pytest.raises(ValueError) as ctx:
+            ClusterGenerator(
+                image_version="image_version",
+                custom_image_family="custom_image_family",
+                project_id=GCP_PROJECT,
+                cluster_name=CLUSTER_NAME,
+            )
+            assert "image_version and custom_image_family" in str(ctx.value)
+
+    def test_custom_image_family_2(self):
+        with pytest.raises(ValueError) as ctx:
+            ClusterGenerator(
+                custom_image="custom_image",
+                custom_image_family="custom_image_family",
+                project_id=GCP_PROJECT,
+                cluster_name=CLUSTER_NAME,
+            )
+            assert "custom_image and custom_image_family" in str(ctx.value)
+
     def test_nodes_number(self):
         with pytest.raises(AssertionError) as ctx:
             ClusterGenerator(


Code review comments:
diff_hunk: @@ -346,6 +358,16 @@ def _build_cluster_data(self):
             if not self.single_node:
                 cluster_data['worker_config']['image_uri'] = custom_image_url
 
+        elif self.custom_image_family:
+            project_id = self.custom_image_project_id or self.project_id
+            custom_image_url = (
+                'https://www.googleapis.com/compute/beta/projects/'
+                '{}/global/images/family/{}'.format(project_id, self.custom_image_family)
path: airflow/providers/google/cloud/operators/dataproc.py
text: @user1:
Let's use f-string here
```suggestion
                f'{project_id}/global/images/family/{self.custom_image_family}'
```