Create secondary custom disk Sample

Updated tests approach

Author: Thoughtseize1

compute/client_library/snippets/disks/create_secondary_custom.py

@@ -0,0 +1,140 @@
+#  Copyright 2024 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.
+# flake8: noqa
+
+from __future__ import annotations
+
+import sys
+from typing import Any
+
+from google.api_core.extended_operation import ExtendedOperation
+from google.cloud import compute_v1
+
+
+def wait_for_extended_operation(
+    operation: ExtendedOperation, verbose_name: str = "operation", timeout: int = 300
+) -> Any:
+    """
+    Waits for the extended (long-running) operation to complete.
+    If the operation is successful, it will return its result.
+    If the operation ends with an error, an exception will be raised.
+    If there were any warnings during the execution of the operation
+    they will be printed to sys.stderr.
+    Args:
+        operation: a long-running operation you want to wait on.
+        verbose_name: (optional) a more verbose name of the operation,
+            used only during error and warning reporting.
+        timeout: how long (in seconds) to wait for operation to finish.
+            If None, wait indefinitely.
+    Returns:
+        Whatever the operation.result() returns.
+    Raises:
+        This method will raise the exception received from `operation.exception()`
+        or RuntimeError if there is no exception set, but there is an `error_code`
+        set for the `operation`.
+        In case of an operation taking longer than `timeout` seconds to complete,
+        a `concurrent.futures.TimeoutError` will be raised.
+    """
+    result = operation.result(timeout=timeout)
+
+    if operation.error_code:
+        print(
+            f"Error during {verbose_name}: [Code: {operation.error_code}]: {operation.error_message}",
+            file=sys.stderr,
+            flush=True,
+        )
+        print(f"Operation ID: {operation.name}", file=sys.stderr, flush=True)
+        raise operation.exception() or RuntimeError(operation.error_message)
+
+    if operation.warnings:
+        print(f"Warnings during {verbose_name}:\n", file=sys.stderr, flush=True)
+        for warning in operation.warnings:
+            print(f" - {warning.code}: {warning.message}", file=sys.stderr, flush=True)
+
+    return result
+
+
+def create_secondary_custom_disk(
+    primary_disk_name: str,
+    primary_disk_project: str,
+    primary_disk_zone: str,
+    secondary_disk_name: str,
+    secondary_disk_project: str,
+    secondary_disk_zone: str,
+    disk_size_gb: int,
+    disk_type: str = "pd-ssd",
+
+) -> compute_v1.Disk:
+    """Creates a custom secondary disk whose properties differ from the primary disk.
+    Args:
+        primary_disk_name (str): The name of the primary disk.
+        primary_disk_project (str): The project of the primary disk.
+        primary_disk_zone (str): The location of the primary disk.
+        secondary_disk_name (str): The name of the secondary disk.
+        secondary_disk_project (str): The project of the secondary disk.
+        secondary_disk_zone (str): The location of the secondary disk.
+        disk_size_gb (int): The size of the disk in GB. Should be the same as the primary disk.
+        disk_type (str): The type of the disk. Must be one of pd-ssd or pd-balanced.
+    """
+    disk_client = compute_v1.DisksClient()
+    disk = compute_v1.Disk()
+    disk.name = secondary_disk_name
+    disk.size_gb = disk_size_gb
+    disk.type = f"zones/{primary_disk_zone}/diskTypes/{disk_type}"
+    disk.async_primary_disk = compute_v1.DiskAsyncReplication(
+        disk=f"projects/{primary_disk_project}/zones/{primary_disk_zone}/disks/{primary_disk_name}"
+    )
+
+    # Specify additional guest OS features for the secondary disk
+    # Set to one or more of the following values: - VIRTIO_SCSI_MULTIQUEUE - WINDOWS - MULTI_IP_SUBNET
+    # - UEFI_COMPATIBLE - GVNIC - SEV_CAPABLE - SUSPEND_RESUME_COMPATIBLE
+    # - SEV_LIVE_MIGRATABLE_V2 - SEV_SNP_CAPABLE - TDX_CAPABLE - IDPF
+    disk.guest_os_features = [compute_v1.GuestOsFeature(type="MULTI_IP_SUBNET")]
+
+    # Assign additional labels to the secondary disk
+    disk.labels = {
+        "source-disk": primary_disk_name,
+        "secondary-disk-for-replication": "true",
+    }
+
+    operation = disk_client.insert(
+        project=secondary_disk_project, zone=secondary_disk_zone, disk_resource=disk
+    )
+    wait_for_extended_operation(operation, "create_secondary_disk")
+
+    secondary_disk = disk_client.get(
+        project=secondary_disk_project,
+        zone=secondary_disk_zone,
+        disk=secondary_disk_name,
+    )
+    return secondary_disk
+
+
+if __name__ == "__main__":
+    import os
+
+    PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+    a = create_secondary_custom_disk(
+        primary_disk_name="disk-1",
+        primary_disk_project=PROJECT_ID,
+        primary_disk_zone="europe-west2-c",
+        secondary_disk_name="secondary-disk-name6",
+        secondary_disk_project=PROJECT_ID,
+        secondary_disk_zone="europe-west1-c",
+        disk_size_gb=100,
+        disk_type="pd-ssd",
+    )
+    print(a.labels)
+    print(a.labels["secondary-disk-for-replication"])
+    print(a.labels["source-disk"])

compute/client_library/snippets/tests/test_disks.py

@@ -28,6 +28,7 @@
 from ..disks.create_hyperdisk_from_pool import create_hyperdisk_from_pool
 from ..disks.create_hyperdisk_storage_pool import create_hyperdisk_storage_pool
 from ..disks.create_kms_encrypted_disk import create_kms_encrypted_disk
+from ..disks.create_secondary_custom import create_secondary_custom_disk
 from ..disks.create_secondary_disk import create_secondary_disk
 from ..disks.create_secondary_region_disk import create_secondary_region_disk
 from ..disks.delete import delete_disk
@@ -45,10 +46,12 @@
 
 PROJECT = google.auth.default()[1]
 ZONE = "europe-west2-c"
+ZONE_SECONDARY = "europe-west1-c"
 REGION = "europe-west2"
-SECOND_REGION = "europe-central2"
+REGION_SECONDARY = "europe-central2"
 KMS_KEYRING_NAME = "compute-test-keyring"
 KMS_KEY_NAME = "compute-test-key"
+DISK_SIZE = 11
 
 
 @pytest.fixture()
@@ -96,6 +99,23 @@ def test_disk():
     delete_disk(PROJECT, ZONE, test_disk_name)
 
 
+@pytest.fixture
+def test_empty_pd_balanced_disk():
+    """
+    Creates and deletes a pd_balanced disk in secondary zone.
+    """
+    disk_name = "test-pd-balanced-disk" + uuid.uuid4().hex[:4]
+    disk = create_empty_disk(
+        PROJECT,
+        ZONE_SECONDARY,
+        disk_name,
+        f"zones/{ZONE_SECONDARY}/diskTypes/pd-balanced",
+        disk_size_gb=DISK_SIZE,
+    )
+    yield disk
+    delete_disk(PROJECT, ZONE_SECONDARY, disk_name)
+
+
 @pytest.fixture
 def test_snapshot(test_disk):
     """
@@ -110,22 +130,22 @@ def test_snapshot(test_disk):
 
 
 @pytest.fixture()
-def autodelete_disk_name():
-    disk_name = "test-disk-" + uuid.uuid4().hex[:10]
+def autodelete_regional_disk_name():
+    disk_name = "secondary-region-disk" + uuid.uuid4().hex[:4]
     yield disk_name
     try:
-        delete_disk(PROJECT, ZONE, disk_name)
+        delete_regional_disk(PROJECT, REGION_SECONDARY, disk_name)
     except NotFound:
         # The disk was already deleted
         pass
 
 
 @pytest.fixture()
-def autodelete_regional_disk_name():
-    disk_name = "secondary-region-disk" + uuid.uuid4().hex[:4]
+def autodelete_disk_name():
+    disk_name = "test-disk-" + uuid.uuid4().hex[:10]
     yield disk_name
     try:
-        delete_regional_disk(PROJECT, SECOND_REGION, disk_name)
+        delete_disk(PROJECT, ZONE, disk_name)
     except NotFound:
         # The disk was already deleted
         pass
@@ -165,7 +185,7 @@ def autodelete_regional_blank_disk():
     disk_type = f"regions/{REGION}/diskTypes/pd-balanced"
 
     disk = create_regional_disk(
-        PROJECT, REGION, replica_zones, disk_name, disk_type, 11
+        PROJECT, REGION, replica_zones, disk_name, disk_type, DISK_SIZE
     )
 
     yield disk
@@ -379,31 +399,6 @@ def test_create_hyperdisk(autodelete_disk_name):
     assert "hyperdisk" in disk.type_.lower()
 
 
-def test_create_secondary(autodelete_disk_name):
-    disk_size = 12
-    primary_disk = create_empty_disk(
-        PROJECT,
-        ZONE,
-        autodelete_disk_name,
-        f"zones/{ZONE}/diskTypes/pd-balanced",
-        disk_size_gb=disk_size,
-    )
-    disk_name = "secondary-region-disk" + uuid.uuid4().hex[:4]
-    disk = create_secondary_disk(
-        autodelete_disk_name,
-        PROJECT,
-        ZONE,
-        disk_name,
-        PROJECT,
-        "europe-west3-c",
-        disk_size,
-    )
-    try:
-        assert disk.async_primary_disk.disk == primary_disk.self_link
-    finally:
-        delete_disk(PROJECT, "europe-west3-c", disk_name)
-
-
 def test_create_secondary_region(
     autodelete_regional_blank_disk, autodelete_regional_disk_name
 ):
@@ -413,7 +408,38 @@ def test_create_secondary_region(
         REGION,
         autodelete_regional_disk_name,
         PROJECT,
-        SECOND_REGION,
-        11,
+        REGION_SECONDARY,
+        DISK_SIZE,
     )
     assert disk.async_primary_disk.disk == autodelete_regional_blank_disk.self_link
+
+
+def test_create_secondary(test_empty_pd_balanced_disk, autodelete_disk_name):
+    disk = create_secondary_disk(
+        primary_disk_name=test_empty_pd_balanced_disk.name,
+        primary_disk_project=PROJECT,
+        primary_disk_zone=ZONE_SECONDARY,
+        secondary_disk_name=autodelete_disk_name,
+        secondary_disk_project=PROJECT,
+        secondary_disk_zone=ZONE,
+        disk_size_gb=DISK_SIZE,
+        disk_type="pd-ssd",
+    )
+    assert disk.async_primary_disk.disk == test_empty_pd_balanced_disk.self_link
+
+
+def test_create_custom_secondary_disk(
+    test_empty_pd_balanced_disk, autodelete_disk_name
+):
+    disk = create_secondary_custom_disk(
+        primary_disk_name=test_empty_pd_balanced_disk.name,
+        primary_disk_project=PROJECT,
+        primary_disk_zone=ZONE_SECONDARY,
+        secondary_disk_name=autodelete_disk_name,
+        secondary_disk_project=PROJECT,
+        secondary_disk_zone=ZONE,
+        disk_size_gb=DISK_SIZE,
+        disk_type="pd-ssd",
+    )
+    assert disk.labels["secondary-disk-for-replication"] == "true"
+    assert disk.labels["source-disk"] == test_empty_pd_balanced_disk.name