feat(compute): add compute disk create with snapshot schedule (#9788)

TetyanaYahodska · web-flow · commit c16bb4c143ce · 2024-12-30T23:07:17.000-05:00
* Implemented compute_disk_create_with_snapshot_schedule sample, created test

* Fixed code

* Fixed code

* Fixed test

* Fixed code

* Fixed code as requested in the comments

* Fixed lint issue

* Fixed lint issue

* Deleted redundant code
diff --git a/compute/cloud-client/src/main/java/compute/disks/CreateDiskWithSnapshotSchedule.java b/compute/cloud-client/src/main/java/compute/disks/CreateDiskWithSnapshotSchedule.java
@@ -0,0 +1,73 @@
+/*
+ * 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.
+ */
+
+package compute.disks;
+
+// [START compute_disk_create_with_snapshot_schedule]
+import com.google.cloud.compute.v1.Disk;
+import com.google.cloud.compute.v1.DisksClient;
+import com.google.cloud.compute.v1.Operation;
+import com.google.cloud.compute.v1.Operation.Status;
+import java.io.IOException;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+public class CreateDiskWithSnapshotSchedule {
+  public static void main(String[] args)
+          throws IOException, ExecutionException, InterruptedException, TimeoutException {
+    // TODO(developer): Replace these variables before running the sample.
+    // Project ID or project number of the Cloud project you want to use.
+    String projectId = "YOUR_PROJECT_ID";
+    // Name of the zone in which you want to create the disk.
+    String zone = "us-central1-a";
+    // Name of the disk you want to create.
+    String diskName = "YOUR_DISK_NAME";
+    // Name of the schedule you want to link to the disk.
+    String snapshotScheduleName = "YOUR_SCHEDULE_NAME";
+
+    createDiskWithSnapshotSchedule(projectId, zone, diskName, snapshotScheduleName);
+  }
+
+  // Creates disk with linked snapshot schedule.
+  public static Status createDiskWithSnapshotSchedule(
+      String projectId, String zone, String diskName, String snapshotScheduleName)
+      throws IOException, ExecutionException, InterruptedException, TimeoutException {
+    // Initialize client that will be used to send requests. This client only needs to be created
+    // once, and can be reused for multiple requests.
+    try (DisksClient disksClient = DisksClient.create()) {
+      String region = zone.substring(0, zone.lastIndexOf('-'));
+      // Get the resource policy to link to the disk
+      String resourcePolicyLink = String.format("projects/%s/regions/%s/resourcePolicies/%s",
+              projectId, region, snapshotScheduleName);
+
+      Disk disk = Disk.newBuilder()
+              .setName(diskName)
+              .setZone(zone)
+              .addAllResourcePolicies(List.of(resourcePolicyLink))
+              .build();
+
+      Operation response = disksClient.insertAsync(projectId, zone, disk).get(3, TimeUnit.MINUTES);
+
+      if (response.hasError()) {
+        throw new Error("Disk creation failed! " + response.getError());
+      }
+      return response.getStatus();
+    }
+  }
+}
+// [END compute_disk_create_with_snapshot_schedule]
diff --git a/compute/cloud-client/src/test/java/compute/Util.java b/compute/cloud-client/src/test/java/compute/Util.java
@@ -30,6 +30,8 @@
 import com.google.cloud.compute.v1.RegionInstanceTemplatesClient;
 import com.google.cloud.compute.v1.Reservation;
 import com.google.cloud.compute.v1.ReservationsClient;
+import com.google.cloud.compute.v1.ResourcePoliciesClient;
+import com.google.cloud.compute.v1.ResourcePolicy;
 import com.google.cloud.compute.v1.Snapshot;
 import com.google.cloud.compute.v1.SnapshotsClient;
 import com.google.cloud.compute.v1.StoragePool;
@@ -39,7 +41,9 @@
 import compute.disks.DeleteSnapshot;
 import compute.disks.RegionalDelete;
 import compute.reservation.DeleteReservation;
+import compute.snapshotschedule.DeleteSnapshotSchedule;
 import java.io.IOException;
+import java.lang.Error;
 import java.nio.charset.StandardCharsets;
 import java.security.SecureRandom;
 import java.time.Instant;
@@ -283,6 +287,21 @@ public static void deleteStoragePool(String project, String zone, String storage
     }
   }
 
+  // Delete snapshot schedule which starts with the given prefixToDelete and
+  // has creation timestamp >24 hours.
+  public static void cleanUpExistingSnapshotSchedule(
+          String prefixToDelete, String projectId, String region)
+          throws IOException, ExecutionException, InterruptedException, TimeoutException {
+    try (ResourcePoliciesClient resourcePoliciesClient = ResourcePoliciesClient.create()) {
+      for (ResourcePolicy resource : resourcePoliciesClient.list(projectId, region).iterateAll()) {
+        if (containPrefixToDeleteAndZone(resource, prefixToDelete, region)
+                && isCreatedBeforeThresholdTime(resource.getCreationTimestamp())) {
+          DeleteSnapshotSchedule.deleteSnapshotSchedule(projectId, region, resource.getName());
+        }
+      }
+    }
+  }
+
   public static boolean containPrefixToDeleteAndZone(
       Object resource, String prefixToDelete, String zone) {
     boolean containPrefixAndZone = false;
@@ -308,6 +327,11 @@ public static boolean containPrefixToDeleteAndZone(
         containPrefixAndZone = ((StoragePool) resource).getName().contains(prefixToDelete)
             && ((StoragePool) resource).getZone().contains(zone);
       }
+      if (resource instanceof ResourcePolicy) {
+        containPrefixAndZone = ((ResourcePolicy) resource).getName().contains(prefixToDelete)
+                && ((ResourcePolicy) resource).getRegion()
+                .contains(zone.substring(0, zone.lastIndexOf('-')));
+      }
     } catch (NullPointerException e) {
       System.out.println("Resource not found, skipping deletion:");
     }
diff --git a/compute/cloud-client/src/test/java/compute/disks/DisksIT.java b/compute/cloud-client/src/test/java/compute/disks/DisksIT.java
@@ -19,6 +19,7 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 import static org.junit.Assert.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 
 import com.google.cloud.compute.v1.AttachedDisk;
 import com.google.cloud.compute.v1.AttachedDiskInitializeParams;
@@ -37,6 +38,8 @@
 import com.google.cloud.compute.v1.SnapshotsClient;
 import compute.DeleteInstance;
 import compute.Util;
+import compute.snapshotschedule.CreateSnapshotSchedule;
+import compute.snapshotschedule.DeleteSnapshotSchedule;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.PrintStream;
@@ -81,7 +84,8 @@ public class DisksIT {
   private static String SECONDARY_DISK;
   private static final long DISK_SIZE = 10L;
   private static String SECONDARY_CUSTOM_DISK;
-
+  private static String DISK_WITH_SNAPSHOT_SCHEDULE;
+  private static String SNAPSHOT_SCHEDULE;
   private ByteArrayOutputStream stdOut;
 
   // Check if the required environment variables are set.
@@ -113,6 +117,8 @@ public static void setup()
     SECONDARY_REGIONAL_DISK = "gcloud-test-disk-secondary-regional-" + uuid;
     SECONDARY_DISK = "gcloud-test-disk-secondary-" + uuid;
     SECONDARY_CUSTOM_DISK = "gcloud-test-disk-custom-" + uuid;
+    DISK_WITH_SNAPSHOT_SCHEDULE = "gcloud-test-disk-shapshot-" + uuid;
+    SNAPSHOT_SCHEDULE = "gcloud-test-snapshot-schedule-" + uuid;
 
     // Cleanup existing stale resources.
     Util.cleanUpExistingInstances("test-disks", PROJECT_ID, ZONE);
@@ -122,6 +128,7 @@ public static void setup()
             "gcloud-test-disk-secondary-regional-", PROJECT_ID, "us-central1");
     Util.cleanUpExistingRegionalDisks("gcloud-test-disk-", PROJECT_ID, REGION);
     Util.cleanUpExistingSnapshots("gcloud-test-snapshot-", PROJECT_ID);
+    Util.cleanUpExistingSnapshotSchedule("gcloud-test-snapshot-schedule-", PROJECT_ID, REGION);
 
     // Create disk from image.
     Image debianImage = null;
@@ -154,7 +161,8 @@ public static void setup()
     TimeUnit.SECONDS.sleep(10);
     SetDiskAutodelete.setDiskAutodelete(PROJECT_ID, ZONE, INSTANCE_NAME, DISK_NAME, true);
     assertThat(stdOut.toString()).contains("Disk autodelete field updated.");
-
+    CreateSnapshotSchedule.createSnapshotSchedule(PROJECT_ID, REGION, SNAPSHOT_SCHEDULE,
+            "description", 10, "US");
     // Create zonal and regional blank disks for testing attach and resize.
     createZonalDisk();
     createRegionalDisk();
@@ -191,6 +199,8 @@ public static void cleanUp()
     RegionalDelete.deleteRegionalDisk(PROJECT_ID, "us-central1", SECONDARY_REGIONAL_DISK);
     DeleteDisk.deleteDisk(PROJECT_ID, "us-central1-c", SECONDARY_DISK);
     DeleteDisk.deleteDisk(PROJECT_ID, "us-central1-c", SECONDARY_CUSTOM_DISK);
+    DeleteDisk.deleteDisk(PROJECT_ID, ZONE, DISK_WITH_SNAPSHOT_SCHEDULE);
+    DeleteSnapshotSchedule.deleteSnapshotSchedule(PROJECT_ID, REGION, SNAPSHOT_SCHEDULE);
 
     stdOut.close();
     System.setOut(out);
@@ -325,24 +335,28 @@ public void testCreateReplicatedDisk()
           throws IOException, ExecutionException, InterruptedException, TimeoutException {
     Status status = CreateReplicatedDisk.createReplicatedDisk(PROJECT_ID, REGION,
             replicaZones, REGIONAL_REPLICATED_DISK, 100, DISK_TYPE);
-    Disk disk = Util.getRegionalDisk(PROJECT_ID, REGION, REGIONAL_REPLICATED_DISK);
 
     assertThat(status).isEqualTo(Status.DONE);
-    assertEquals(REGIONAL_REPLICATED_DISK, disk.getName());
+    assertDoesNotThrow(() -> {
+      Disk disk = Util.getRegionalDisk(PROJECT_ID, REGION, REGIONAL_REPLICATED_DISK);
+      assertEquals(REGIONAL_REPLICATED_DISK, disk.getName());
+    });
   }
 
   @Test
   public void testCreateDiskSecondaryRegional()
       throws IOException, ExecutionException, InterruptedException, TimeoutException {
-    String diskType =  String.format(
+    String diskType = String.format(
         "projects/%s/regions/%s/diskTypes/pd-balanced", PROJECT_ID, REGION);
     Status status = CreateDiskSecondaryRegional.createDiskSecondaryRegional(
         PROJECT_ID, PROJECT_ID, REGIONAL_BLANK_DISK, SECONDARY_REGIONAL_DISK,
         REGION, "us-central1", DISK_SIZE,  diskType);
-    Disk disk = Util.getRegionalDisk(PROJECT_ID, "us-central1", SECONDARY_REGIONAL_DISK);
 
     assertThat(status).isEqualTo(Status.DONE);
-    assertEquals(SECONDARY_REGIONAL_DISK, disk.getName());
+    assertDoesNotThrow(() -> {
+      Disk disk = Util.getRegionalDisk(PROJECT_ID, "us-central1", SECONDARY_REGIONAL_DISK);
+      assertEquals(SECONDARY_REGIONAL_DISK, disk.getName());
+    });
   }
 
   @Test
@@ -353,10 +367,12 @@ public void testCreateDiskSecondaryZonal()
     Status status = CreateDiskSecondaryZonal.createDiskSecondaryZonal(
         PROJECT_ID, PROJECT_ID, EMPTY_DISK_NAME, SECONDARY_DISK, ZONE,
         "us-central1-c", DISK_SIZE, diskType);
-    Disk disk = Util.getDisk(PROJECT_ID, "us-central1-c", SECONDARY_DISK);
 
     assertThat(status).isEqualTo(Status.DONE);
-    assertEquals(SECONDARY_DISK, disk.getName());
+    assertDoesNotThrow(() -> {
+      Disk disk = Util.getDisk(PROJECT_ID, "us-central1-c", SECONDARY_DISK);
+      assertEquals(SECONDARY_DISK, disk.getName());
+    });
   }
 
   @Test
@@ -367,9 +383,24 @@ public void testCreateSecondaryCustomDisk()
     Status status = CreateSecondaryCustomDisk.createSecondaryCustomDisk(
         PROJECT_ID, PROJECT_ID, EMPTY_DISK_NAME, SECONDARY_CUSTOM_DISK, ZONE,
         "us-central1-c", DISK_SIZE,  diskType);
-    Disk disk = Util.getDisk(PROJECT_ID, "us-central1-c", SECONDARY_CUSTOM_DISK);
 
     assertThat(status).isEqualTo(Status.DONE);
-    assertEquals(SECONDARY_CUSTOM_DISK, disk.getName());
+    assertDoesNotThrow(() -> {
+      Disk disk = Util.getDisk(PROJECT_ID, "us-central1-c", SECONDARY_CUSTOM_DISK);
+      assertEquals(SECONDARY_CUSTOM_DISK, disk.getName());
+    });
+  }
+
+  @Test
+  void testCreateDiskWithSnapshotSchedule()
+          throws IOException, ExecutionException, InterruptedException, TimeoutException {
+    Status status = CreateDiskWithSnapshotSchedule.createDiskWithSnapshotSchedule(
+            PROJECT_ID, ZONE, DISK_WITH_SNAPSHOT_SCHEDULE, SNAPSHOT_SCHEDULE);
+
+    assertThat(status).isEqualTo(Status.DONE);
+    assertDoesNotThrow(() -> {
+      Disk disk = Util.getDisk(PROJECT_ID, ZONE, DISK_WITH_SNAPSHOT_SCHEDULE);
+      assertEquals(DISK_WITH_SNAPSHOT_SCHEDULE, disk.getName());
+    });
   }
 }
