able to configure the quota for a bucket

1 files changed, 93 insertions, 0 deletions

diff --git a/weed/shell/command_s3_bucket_quota.go b/weed/shell/command_s3_bucket_quota.go
new file mode 100644
index 000000000..17be1d5b0
--- /dev/null
+++ b/weed/shell/command_s3_bucket_quota.go
@@ -0,0 +1,93 @@
+package shell
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
+	"io"
+)
+
+func init() {
+	Commands = append(Commands, &commandS3BucketQuota{})
+}
+
+type commandS3BucketQuota struct {
+}
+
+func (c *commandS3BucketQuota) Name() string {
+	return "s3.bucket.quota"
+}
+
+func (c *commandS3BucketQuota) Help() string {
+	return `set/remove/enable/disable quota for a bucket
+
+	Example:
+		s3.bucket.quota -name=<bucket_name> -operation=set -sizeMB=1024
+`
+}
+
+func (c *commandS3BucketQuota) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {
+
+	bucketCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
+	bucketName := bucketCommand.String("name", "", "bucket name")
+	operationName := bucketCommand.String("op", "set", "operation name [set|remove|enable|disable]")
+	sizeMB := bucketCommand.Int64("sizeMB", 0, "bucket quota size in MiB")
+	if err = bucketCommand.Parse(args); err != nil {
+		return nil
+	}
+
+	if *bucketName == "" {
+		return fmt.Errorf("empty bucket name")
+	}
+
+	err = commandEnv.WithFilerClient(false, func(client filer_pb.SeaweedFilerClient) error {
+
+		ctx := context.Background()
+
+		resp, err := client.GetFilerConfiguration(ctx, &filer_pb.GetFilerConfigurationRequest{})
+		if err != nil {
+			return fmt.Errorf("get filer configuration: %v", err)
+		}
+		filerBucketsPath := resp.DirBuckets
+
+		lookupResp, err := client.LookupDirectoryEntry(ctx, &filer_pb.LookupDirectoryEntryRequest{
+			Directory: filerBucketsPath,
+			Name:      *bucketName,
+		})
+		if err != nil {
+			return fmt.Errorf("did not find bucket %s: %v", *bucketName, err)
+		}
+		bucketEntry := lookupResp.Entry
+
+		switch *operationName {
+		case "set":
+			bucketEntry.Quota = *sizeMB * 1024 * 1024
+		case "remove":
+			bucketEntry.Quota = 0
+		case "enable":
+			if bucketEntry.Quota < 0 {
+				bucketEntry.Quota = -bucketEntry.Quota
+			}
+		case "disable":
+			if bucketEntry.Quota > 0 {
+				bucketEntry.Quota = -bucketEntry.Quota
+			}
+		}
+
+		if err := filer_pb.UpdateEntry(client, &filer_pb.UpdateEntryRequest{
+			Directory: filerBucketsPath,
+			Entry:     bucketEntry,
+		}); err != nil {
+			return err
+		}
+
+		println("updated quota for bucket", *bucketName)
+
+		return nil
+
+	})
+
+	return err
+
+}