From 606667f205774943fc64884bc86b2befddfd75c3 Mon Sep 17 00:00:00 2001
From: chrislu <chris.lu@gmail.com>
Date: Fri, 21 Jan 2022 01:42:20 -0800
Subject: [PATCH] able to configure the quota for a bucket

---
 weed/filer/entry.go                   |  4 ++
 weed/filer/entry_codec.go             |  3 +
 weed/shell/command_s3_bucket_quota.go | 93 +++++++++++++++++++++++++++
 3 files changed, 100 insertions(+)
 create mode 100644 weed/shell/command_s3_bucket_quota.go

diff --git a/weed/filer/entry.go b/weed/filer/entry.go
index 8fa75fe6b..9f83da4aa 100644
--- a/weed/filer/entry.go
+++ b/weed/filer/entry.go
@@ -43,6 +43,7 @@ type Entry struct {
 	HardLinkCounter int32
 	Content         []byte
 	Remote          *filer_pb.RemoteEntry
+	Quota           int64
 }
 
 func (entry *Entry) Size() uint64 {
@@ -70,6 +71,7 @@ func (entry *Entry) ShallowClone() *Entry {
 	newEntry.HardLinkCounter = entry.HardLinkCounter
 	newEntry.Content = entry.Content
 	newEntry.Remote = entry.Remote
+	newEntry.Quota = entry.Quota
 
 	return newEntry
 }
@@ -96,6 +98,7 @@ func (entry *Entry) ToExistingProtoEntry(message *filer_pb.Entry) {
 	message.HardLinkCounter = entry.HardLinkCounter
 	message.Content = entry.Content
 	message.RemoteEntry = entry.Remote
+	message.Quota = entry.Quota
 }
 
 func FromPbEntryToExistingEntry(message *filer_pb.Entry, fsEntry *Entry) {
@@ -106,6 +109,7 @@ func FromPbEntryToExistingEntry(message *filer_pb.Entry, fsEntry *Entry) {
 	fsEntry.HardLinkCounter = message.HardLinkCounter
 	fsEntry.Content = message.Content
 	fsEntry.Remote = message.RemoteEntry
+	fsEntry.Quota = message.Quota
 }
 
 func (entry *Entry) ToProtoFullEntry() *filer_pb.FullEntry {
diff --git a/weed/filer/entry_codec.go b/weed/filer/entry_codec.go
index 55c937b39..0a917bea9 100644
--- a/weed/filer/entry_codec.go
+++ b/weed/filer/entry_codec.go
@@ -118,6 +118,9 @@ func EqualEntry(a, b *Entry) bool {
 	if !proto.Equal(a.Remote, b.Remote) {
 		return false
 	}
+	if a.Quota != b.Quota {
+		return false
+	}
 	return true
 }
 
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
+
+}