HDDS-15881. Support setting storagePolicy via SetBucketProperty

hadoop-hdds/common/src/main/java/org/apache/hadoop/ozone/OzoneConsts.java

@@ -300,6 +300,7 @@ public final class OzoneConsts {
   public static final String STORAGE_TYPE = "storageType";
   public static final String STORAGE_POLICY = "storagePolicy";
   public static final String ALLOW_FALLBACK_STORAGE_POLICY = "allowFallbackStoragePolicy";
+  public static final String UNSET_STORAGE_POLICY = "unSetStoragePolicy";
   public static final String RESOURCE_TYPE = "resourceType";
   public static final String IS_VERSION_ENABLED = "isVersionEnabled";
   public static final String CREATION_TIME = "creationTime";

hadoop-ozone/cli-shell/src/main/java/org/apache/hadoop/ozone/shell/bucket/CreateBucketHandler.java

@@ -21,6 +21,8 @@
 import java.io.IOException;
 import org.apache.hadoop.hdds.client.DefaultReplicationConfig;
 import org.apache.hadoop.hdds.client.OzoneQuota;
+import org.apache.hadoop.hdds.client.OzoneStoragePolicy;
+import org.apache.hadoop.hdds.client.StoragePolicy;
 import org.apache.hadoop.ozone.OzoneConsts;
 import org.apache.hadoop.ozone.client.BucketArgs;
 import org.apache.hadoop.ozone.client.OzoneBucket;
@@ -67,6 +69,19 @@ public class CreateBucketHandler extends BucketHandler {
   @CommandLine.Mixin
   private SetSpaceQuotaOptions quotaOptions;
 
+  @Option(names = {"--storage-policy", "-s"},
+      description = "Bucket StoragePolicy. Allowed values: HOT, WARM, COLD, null. Default: WARM.",
+      defaultValue = "WARM")
+  private String storagePolicyStr;
+
+  @Option(names = {"--allow-fallback-storage-policy", "-a"},
+      description = "When true, allocation may fall back to the StoragePolicy's " +
+          "fallback tier if the creation tier is unavailable. Default: true.",
+      defaultValue = "true")
+  private String allowFallBackStoragePolicyStr;
+
+  private static final String NULL_STORAGE_POLICY = "null";
+
   /**
    * Executes create bucket.
    */
@@ -78,8 +93,14 @@ public void execute(OzoneClient client, OzoneAddress address)
       ownerName = UserGroupInformation.getCurrentUser().getShortUserName();
     }
 
+    StoragePolicy storagePolicy = parseStoragePolicy(storagePolicyStr);
+    Boolean allowFallBackStoragePolicy =
+        Boolean.valueOf(allowFallBackStoragePolicyStr);
+
     BucketArgs.Builder bb =
         new BucketArgs.Builder()
+            .setStoragePolicy(storagePolicy)
+            .setAllowFallbackStoragePolicy(allowFallBackStoragePolicy)
             .setVersioning(false).setOwner(ownerName);
     if (allowedBucketLayout != null) {
       bb.setBucketLayout(allowedBucketLayout);
@@ -127,6 +148,28 @@ public void execute(OzoneClient client, OzoneAddress address)
     }
   }
 
+  /**
+   * Parse a user-supplied {@code --storage-policy} value into a
+   * {@link StoragePolicy}. Returns {@code null} when the caller passed
+   * "null" (any case) or an empty value, to leave the bucket without an
+   * explicit StoragePolicy (the server defaults it to WARM on create).
+   *
+   * @throws IllegalArgumentException if the value is not one of HOT, WARM,
+   *   COLD, or "null".
+   */
+  private static StoragePolicy parseStoragePolicy(String value) {
+    if (Strings.isNullOrEmpty(value)
+        || NULL_STORAGE_POLICY.equalsIgnoreCase(value)) {
+      return null;
+    }
+    try {
+      return OzoneStoragePolicy.valueOf(value.toUpperCase());
+    } catch (IllegalArgumentException e) {
+      throw new IllegalArgumentException("Invalid storage policy: " + value
+          + ". Allowed String values are: HOT, WARM, COLD, or null.");
+    }
+  }
+
   private static class LayoutConverter implements CommandLine.ITypeConverter<BucketLayout> {
     @Override
     public BucketLayout convert(String value) {

hadoop-ozone/cli-shell/src/main/java/org/apache/hadoop/ozone/shell/bucket/UpdateBucketHandler.java

@@ -17,9 +17,13 @@
 
 package org.apache.hadoop.ozone.shell.bucket;
 
+import com.google.common.base.Strings;
 import java.io.IOException;
+import org.apache.hadoop.hdds.client.OzoneStoragePolicy;
+import org.apache.hadoop.hdds.client.StoragePolicy;
 import org.apache.hadoop.ozone.client.OzoneBucket;
 import org.apache.hadoop.ozone.client.OzoneClient;
+import org.apache.hadoop.ozone.om.helpers.OmBucketArgs;
 import org.apache.hadoop.ozone.shell.OzoneAddress;
 import picocli.CommandLine.Command;
 import picocli.CommandLine.Option;
@@ -35,6 +39,20 @@ public class UpdateBucketHandler extends BucketHandler {
       description = "Owner of the bucket to set")
   private String ownerName;
 
+  @Option(names = {"--storage-policy", "-s"},
+      description = "Bucket StoragePolicy. Allowed values: HOT, WARM, COLD, null "
+          + "(null clears the bucket's StoragePolicy). Leave unset to keep the "
+          + "current value.")
+  private String storagePolicyStr;
+
+  @Option(names = {"--allow-fallback-storage-policy", "-a"},
+      description = "When true, allocation may fall back to the StoragePolicy's "
+          + "fallback tier if the creation tier is unavailable. Leave unset to "
+          + "keep the current value.")
+  private String allowFallBackStoragePolicyStr;
+
+  private static final String NULL_STORAGE_POLICY = "null";
+
   @Override
   protected void execute(OzoneClient client, OzoneAddress address)
       throws IOException {
@@ -52,8 +70,64 @@ protected void execute(OzoneClient client, OzoneAddress address)
       }
     }
 
+    // Update StoragePolicy / allowFallback if requested.
+    boolean shouldSetStoragePolicyProperty = false;
+    OmBucketArgs.Builder bucketArgsBuilder = OmBucketArgs.newBuilder()
+        .setVolumeName(volumeName)
+        .setBucketName(bucketName);
+
+    if (hasStoragePolicy()) {
+      shouldSetStoragePolicyProperty = true;
+      StoragePolicy storagePolicy = getStoragePolicy();
+      if (storagePolicy == null) {
+        bucketArgsBuilder.setUnSetStoragePolicy(true);
+      } else {
+        bucketArgsBuilder.setStoragePolicy(storagePolicy);
+      }
+    }
+
+    Boolean allowFallBackStoragePolicy = getAllowFallBackStoragePolicy();
+    if (allowFallBackStoragePolicy != null) {
+      shouldSetStoragePolicyProperty = true;
+      bucketArgsBuilder.setAllowFallbackStoragePolicy(allowFallBackStoragePolicy);
+    }
+    if (shouldSetStoragePolicyProperty) {
+      bucket.setStoragePolicyProperty(bucketArgsBuilder.build());
+    }
+
     OzoneBucket updatedBucket = client.getObjectStore().getVolume(volumeName)
         .getBucket(bucketName);
     printObjectAsJson(updatedBucket);
   }
+
+  private boolean hasStoragePolicy() {
+    return !Strings.isNullOrEmpty(storagePolicyStr);
+  }
+
+  /**
+   * Parse the {@code --storagepolicy} value. Returns {@code null} when the
+   * user passed "null" (any case), meaning the policy should be cleared;
+   * otherwise the matching {@link StoragePolicy}.
+   *
+   * @throws IllegalArgumentException if the value is not HOT, WARM, COLD, or
+   *   "null".
+   */
+  private StoragePolicy getStoragePolicy() {
+    if (Strings.isNullOrEmpty(storagePolicyStr)
+        || NULL_STORAGE_POLICY.equalsIgnoreCase(storagePolicyStr)) {
+      return null;
+    }
+    try {
+      return OzoneStoragePolicy.valueOf(storagePolicyStr.toUpperCase());
+    } catch (IllegalArgumentException e) {
+      throw new IllegalArgumentException("Invalid storage policy: "
+          + storagePolicyStr
+          + ". Allowed String values are: HOT, WARM, COLD, or null.");
+    }
+  }
+
+  private Boolean getAllowFallBackStoragePolicy() {
+    return allowFallBackStoragePolicyStr == null ? null
+        : Boolean.valueOf(allowFallBackStoragePolicyStr);
+  }
 }

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/OzoneBucket.java

@@ -55,6 +55,7 @@
 import org.apache.hadoop.ozone.om.helpers.BasicOmKeyInfo;
 import org.apache.hadoop.ozone.om.helpers.BucketLayout;
 import org.apache.hadoop.ozone.om.helpers.ErrorInfo;
+import org.apache.hadoop.ozone.om.helpers.OmBucketArgs;
 import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
 import org.apache.hadoop.ozone.om.helpers.OmLifecycleConfiguration;
 import org.apache.hadoop.ozone.om.helpers.OmMultipartInfo;
@@ -377,6 +378,17 @@ public void setStoragePolicy(StoragePolicy newStoragePolicy) throws IOException
     storagePolicy = newStoragePolicy;
   }
 
+  /**
+   * Sets the bucket's storage-policy properties carried in the given
+   * {@link OmBucketArgs} (storage policy, allowFallback, or unset). Used by the
+   * update path, where the policy may be absent or explicitly cleared.
+   * @param args Bucket arguments carrying the properties to update.
+   * @throws IOException
+   */
+  public void setStoragePolicyProperty(OmBucketArgs args) throws IOException {
+    proxy.setBucketStoragePolicy(args);
+  }
+
   /**
    * Enable/Disable versioning of the bucket.
    * @param newVersioning

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/protocol/ClientProtocol.java

@@ -50,6 +50,7 @@
 import org.apache.hadoop.ozone.om.helpers.DeleteTenantState;
 import org.apache.hadoop.ozone.om.helpers.ErrorInfo;
 import org.apache.hadoop.ozone.om.helpers.LeaseKeyInfo;
+import org.apache.hadoop.ozone.om.helpers.OmBucketArgs;
 import org.apache.hadoop.ozone.om.helpers.OmKeyArgs;
 import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
 import org.apache.hadoop.ozone.om.helpers.OmKeyLocationInfo;
@@ -293,6 +294,16 @@ void setBucketStoragePolicy(String volumeName, String bucketName,
                               StoragePolicy storagePolicy)
       throws IOException;
 
+  /**
+   * Sets the storage-policy related properties of a bucket carried in the
+   * given {@link OmBucketArgs} (storage policy, allowFallback, or unset).
+   * Unlike the three-arg overload, the policy may be absent or explicitly
+   * cleared via {@link OmBucketArgs.Builder#setUnSetStoragePolicy}.
+   * @param args Bucket arguments carrying the properties to update.
+   * @throws IOException
+   */
+  void setBucketStoragePolicy(OmBucketArgs args) throws IOException;
+
   /**
    * Deletes a bucket if it is empty.
    * @param volumeName Name of the Volume

hadoop-ozone/client/src/main/java/org/apache/hadoop/ozone/client/rpc/RpcClient.java

@@ -1236,6 +1236,14 @@ public void setBucketStoragePolicy(
     ozoneManagerClient.setBucketProperty(builder.build());
   }
 
+  @Override
+  public void setBucketStoragePolicy(OmBucketArgs args) throws IOException {
+    Objects.requireNonNull(args, "args == null");
+    verifyVolumeName(args.getVolumeName());
+    verifyBucketName(args.getBucketName());
+    ozoneManagerClient.setBucketProperty(args);
+  }
+
   @Override
   public void setBucketQuota(String volumeName, String bucketName,
       long quotaInNamespace, long quotaInBytes) throws IOException {

hadoop-ozone/common/src/main/java/org/apache/hadoop/ozone/om/helpers/OmBucketArgs.java

@@ -55,6 +55,11 @@ public final class OmBucketArgs extends WithMetadata implements Auditable {
    * the flag was not set.
    */
   private final Boolean allowFallbackStoragePolicy;
+  /**
+   * Whether to clear (unset) the bucket's storage policy. {@code null} or
+   * {@code false} leaves the existing policy unchanged.
+   */
+  private final Boolean unSetStoragePolicy;
 
   /**
    * Bucket encryption key info if encryption is enabled.
@@ -81,6 +86,7 @@ private OmBucketArgs(Builder b) {
     this.isVersionEnabled = b.isVersionEnabled;
     this.storagePolicy = b.storagePolicy;
     this.allowFallbackStoragePolicy = b.allowFallbackStoragePolicy;
+    this.unSetStoragePolicy = b.unSetStoragePolicy;
     this.ownerName = b.ownerName;
     this.defaultReplicationConfig = b.defaultReplicationConfig;
     this.quotaInBytesSet = b.quotaInBytesSet;
@@ -132,6 +138,14 @@ public Boolean getAllowFallbackStoragePolicy() {
     return allowFallbackStoragePolicy;
   }
 
+  /**
+   * Returns whether the bucket's storage policy should be cleared (unset).
+   * @return unSetStoragePolicy (may be {@code null} when not set).
+   */
+  public Boolean getUnSetStoragePolicy() {
+    return unSetStoragePolicy;
+  }
+
   /**
    * Returns Bucket Quota in bytes.
    * @return quotaInBytes.
@@ -213,6 +227,10 @@ public Map<String, String> toAuditMap() {
       auditMap.put(OzoneConsts.ALLOW_FALLBACK_STORAGE_POLICY,
           String.valueOf(this.allowFallbackStoragePolicy));
     }
+    if (this.unSetStoragePolicy != null) {
+      auditMap.put(OzoneConsts.UNSET_STORAGE_POLICY,
+          String.valueOf(this.unSetStoragePolicy));
+    }
     if (this.ownerName != null) {
       auditMap.put(OzoneConsts.OWNER, this.ownerName);
     }
@@ -249,6 +267,7 @@ public static class Builder extends WithMetadata.Builder {
     private Boolean isVersionEnabled;
     private StoragePolicy storagePolicy;
     private Boolean allowFallbackStoragePolicy;
+    private Boolean unSetStoragePolicy;
     private boolean quotaInBytesSet = false;
     private long quotaInBytes;
     private boolean quotaInNamespaceSet = false;
@@ -306,6 +325,11 @@ public Builder setAllowFallbackStoragePolicy(Boolean allowFallback) {
       return this;
     }
 
+    public Builder setUnSetStoragePolicy(Boolean unSet) {
+      this.unSetStoragePolicy = unSet;
+      return this;
+    }
+
     public Builder setQuotaInBytes(long quota) {
       quotaInBytesSet = true;
       quotaInBytes = quota;
@@ -371,6 +395,9 @@ public BucketArgs getProtobuf() {
     if (allowFallbackStoragePolicy != null) {
       builder.setAllowFallbackStoragePolicy(allowFallbackStoragePolicy);
     }
+    if (unSetStoragePolicy != null) {
+      builder.setUnSetStoragePolicy(unSetStoragePolicy);
+    }
     if (quotaInBytesSet && (
         quotaInBytes > 0 || quotaInBytes == OzoneConsts.QUOTA_RESET)) {
       builder.setQuotaInBytes(quotaInBytes);
@@ -416,6 +443,9 @@ public static Builder builderFromProtobuf(BucketArgs bucketArgs) {
     if (bucketArgs.hasAllowFallbackStoragePolicy()) {
       builder.setAllowFallbackStoragePolicy(bucketArgs.getAllowFallbackStoragePolicy());
     }
+    if (bucketArgs.hasUnSetStoragePolicy()) {
+      builder.setUnSetStoragePolicy(bucketArgs.getUnSetStoragePolicy());
+    }
     if (bucketArgs.hasOwnerName()) {
       builder.setOwnerName(bucketArgs.getOwnerName());
     }