Add support to instance selection disk config and confidential compute to Dataproc

providers/google/src/airflow/providers/google/cloud/operators/dataproc.py

@@ -84,6 +84,15 @@ class PreemptibilityType(Enum):
     NON_PREEMPTIBLE = "NON_PREEMPTIBLE"
 
 
+class ConfidentialInstanceType(Enum):
+    """Contains possible Type values of Confidential Instance applicable for every secondary worker of Cluster."""
+
+    CONFIDENTIAL_INSTANCE_TYPE_UNSPECIFIED = "CONFIDENTIAL_INSTANCE_TYPE_UNSPECIFIED"
+    SEV = "SEV"
+    SEV_SNP = "SEV_SNP"
+    TDX = "TDX"
+
+
 @dataclass
 class InstanceSelection:
     """
@@ -97,10 +106,12 @@ class InstanceSelection:
         Dataproc will first try to create a VM based on the machine-type with priority rank and fallback
         to next rank based on availability. Machine types and instance selections with the same priority have
         the same preference.
+    :param disk_config: Disk config for the instance selection.
     """
 
     machine_types: list[str]
     rank: int = 0
+    disk_config: dict | None = None
 
 
 @dataclass
@@ -117,6 +128,21 @@ class InstanceFlexibilityPolicy:
     instance_selection_list: list[InstanceSelection]
 
 
+@dataclass
+class ConfidentialInstanceConfig:
+    """
+    Confidential instance configuration.
+
+    Representation for google.cloud.dataproc.v1#google.cloud.dataproc.v1.ConfidentialInstanceConfig.
+
+    :param enable_confidential_compute: Whether to enable confidential compute.
+    :param confidential_instance_type: Confidential instance type.
+    """
+
+    enable_confidential_compute: bool
+    confidential_instance_type: str = ConfidentialInstanceType.CONFIDENTIAL_INSTANCE_TYPE_UNSPECIFIED.value
+
+
 class ClusterGenerator:
     """
     Create a new Dataproc Cluster.
@@ -194,6 +220,7 @@ class ClusterGenerator:
     :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
     :param service_account: The service account of the dataproc instances.
     :param service_account_scopes: The URIs of service account scopes to be included.
+    :param confidential_instance_config: Confidential instance config for the cluster.
     :param idle_delete_ttl: The longest duration that cluster would keep alive while
         staying idle. Passing this threshold will cause cluster to be auto-deleted.
         A duration in seconds.
@@ -212,6 +239,7 @@ class ClusterGenerator:
         ``projects/[PROJECT_STORING_KEYS]/locations/[LOCATION]/keyRings/[KEY_RING_NAME]/cryptoKeys/[KEY_NAME]`` # noqa
     :param enable_component_gateway: Provides access to the web interfaces of default and selected optional
         components on the cluster.
+    :param confidential_instance_config: Confidential compute instance config for the cluster.
     :param driver_pool_size: The number of driver nodes in the node group.
     :param driver_pool_id: The ID for the driver pool. Must be unique within the cluster. Use this ID to
         identify the driver group in future operations, such as resizing the node group.
@@ -269,6 +297,7 @@ def __init__(
         idle_stop_ttl: int | None = None,
         auto_stop_time: datetime | None = None,
         auto_stop_ttl: int | None = None,
+        confidential_instance_config: dict | ConfidentialInstanceConfig | None = None,
         customer_managed_key: str | None = None,
         enable_component_gateway: bool | None = False,
         driver_pool_size: int = 0,
@@ -324,6 +353,7 @@ def __init__(
         self.idle_stop_ttl = idle_stop_ttl
         self.auto_stop_time = auto_stop_time
         self.auto_stop_ttl = auto_stop_ttl
+        self.confidential_instance_config = confidential_instance_config
         self.customer_managed_key = customer_managed_key
         self.enable_component_gateway = enable_component_gateway
         self.single_node = num_workers == 0
@@ -407,6 +437,11 @@ def _build_gce_cluster_config(self, cluster_data):
         if self.service_account_scopes:
             cluster_data[config]["service_account_scopes"] = self.service_account_scopes
 
+        if self.confidential_instance_config:
+            if isinstance(self.confidential_instance_config, ConfidentialInstanceConfig):
+                cluster_data[config]["confidential_instance_config"] = vars(self.confidential_instance_config)
+            else:
+                cluster_data[config]["confidential_instance_config"] = self.confidential_instance_config
         return cluster_data
 
     def _build_lifecycle_config(self, cluster_data):
@@ -532,7 +567,7 @@ def _build_cluster_data(self):
             if self.secondary_worker_instance_flexibility_policy:
                 cluster_data["secondary_worker_config"]["instance_flexibility_policy"] = {
                     "instance_selection_list": [
-                        vars(s)
+                        {k: v for k, v in vars(s).items() if v is not None}
                         for s in self.secondary_worker_instance_flexibility_policy.instance_selection_list
                     ]
                 }

providers/google/tests/unit/google/cloud/operators/test_dataproc.py

@@ -103,6 +103,7 @@
         "tags": ["tags"],
         "service_account": "service_account",
         "service_account_scopes": ["service_account_scopes"],
+        "confidential_instance_config": {"confidential_instance_type": "SEV"},
     },
     "master_config": {
         "num_instances": 2,
@@ -256,6 +257,7 @@
                         "projects/project_id/zones/zone/machineTypes/machine2",
                     ],
                     "rank": 0,
+                    "disk_config": {"boot_disk_type": "pd-standard", "boot_disk_size_gb": 100},
                 },
                 {"machine_types": ["projects/project_id/zones/zone/machineTypes/machine3"], "rank": 1},
             ],
@@ -350,6 +352,7 @@
                         "projects/project_id/zones/zone/machineTypes/machine2",
                     ],
                     "rank": 0,
+                    "disk_config": {"boot_disk_type": "pd-standard", "boot_disk_size_gb": 100},
                 },
                 {"machine_types": ["projects/project_id/zones/zone/machineTypes/machine3"], "rank": 1},
             ],
@@ -369,6 +372,7 @@
                         "projects/project_id/zones/zone/machineTypes/machine2",
                     ],
                     "rank": 0,
+                    "disk_config": {"boot_disk_type": "pd-standard", "boot_disk_size_gb": 30},
                 },
                 {"machine_types": ["projects/project_id/zones/zone/machineTypes/machine3"], "rank": 1},
             ],
@@ -720,6 +724,7 @@ def test_build(self):
             cluster_tier="CLUSTER_TIER_STANDARD",
             cluster_type="STANDARD",
             engine="DEFAULT",
+            confidential_instance_config={"enable_confidential_compute": True},
         )
         cluster = generator.make()
         assert cluster == CONFIG
@@ -805,6 +810,7 @@ def test_build_with_flex_migs(self):
                             "projects/project_id/zones/zone/machineTypes/machine2",
                         ],
                         0,
+                        {"boot_disk_type": "pd-standard", "boot_disk_size_gb": 30},
                     ),
                     InstanceSelection(["projects/project_id/zones/zone/machineTypes/machine3"], 1),
                 ]
@@ -852,6 +858,7 @@ def test_build_with_master_and_worker_flex_migs(self):
                             "projects/project_id/zones/zone/machineTypes/machine2",
                         ],
                         0,
+                        {"boot_disk_type": "pd-standard", "boot_disk_size_gb": 100},
                     ),
                     InstanceSelection(["projects/project_id/zones/zone/machineTypes/machine3"], 1),
                 ]
@@ -864,6 +871,7 @@ def test_build_with_master_and_worker_flex_migs(self):
                             "projects/project_id/zones/zone/machineTypes/machine2",
                         ],
                         0,
+                        {"boot_disk_type": "pd-standard", "boot_disk_size_gb": 100},
                     ),
                     InstanceSelection(["projects/project_id/zones/zone/machineTypes/machine3"], 1),
                 ]