You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

97 lines
2.3 KiB

3 years ago
3 years ago
  1. package shell
  2. import (
  3. "context"
  4. "flag"
  5. "fmt"
  6. "io"
  7. "github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
  8. )
  9. func init() {
  10. Commands = append(Commands, &commandS3BucketQuota{})
  11. }
  12. type commandS3BucketQuota struct {
  13. }
  14. func (c *commandS3BucketQuota) Name() string {
  15. return "s3.bucket.quota"
  16. }
  17. func (c *commandS3BucketQuota) Help() string {
  18. return `set/remove/enable/disable quota for a bucket
  19. Example:
  20. s3.bucket.quota -name=<bucket_name> -op=set -sizeMB=1024
  21. `
  22. }
  23. func (c *commandS3BucketQuota) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {
  24. bucketCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
  25. bucketName := bucketCommand.String("name", "", "bucket name")
  26. operationName := bucketCommand.String("op", "set", "operation name [set|get|remove|enable|disable]")
  27. sizeMB := bucketCommand.Int64("sizeMB", 0, "bucket quota size in MiB")
  28. if err = bucketCommand.Parse(args); err != nil {
  29. return nil
  30. }
  31. if *bucketName == "" {
  32. return fmt.Errorf("empty bucket name")
  33. }
  34. err = commandEnv.WithFilerClient(false, func(client filer_pb.SeaweedFilerClient) error {
  35. ctx := context.Background()
  36. resp, err := client.GetFilerConfiguration(ctx, &filer_pb.GetFilerConfigurationRequest{})
  37. if err != nil {
  38. return fmt.Errorf("get filer configuration: %v", err)
  39. }
  40. filerBucketsPath := resp.DirBuckets
  41. lookupResp, err := client.LookupDirectoryEntry(ctx, &filer_pb.LookupDirectoryEntryRequest{
  42. Directory: filerBucketsPath,
  43. Name: *bucketName,
  44. })
  45. if err != nil {
  46. return fmt.Errorf("did not find bucket %s: %v", *bucketName, err)
  47. }
  48. bucketEntry := lookupResp.Entry
  49. switch *operationName {
  50. case "set":
  51. bucketEntry.Quota = *sizeMB * 1024 * 1024
  52. case "get":
  53. fmt.Fprintf(writer, "bucket quota: %dMiB \n", bucketEntry.Quota/1024/1024)
  54. return nil
  55. case "remove":
  56. bucketEntry.Quota = 0
  57. case "enable":
  58. if bucketEntry.Quota < 0 {
  59. bucketEntry.Quota = -bucketEntry.Quota
  60. }
  61. case "disable":
  62. if bucketEntry.Quota > 0 {
  63. bucketEntry.Quota = -bucketEntry.Quota
  64. }
  65. }
  66. if err := filer_pb.UpdateEntry(client, &filer_pb.UpdateEntryRequest{
  67. Directory: filerBucketsPath,
  68. Entry: bucketEntry,
  69. }); err != nil {
  70. return err
  71. }
  72. println("updated quota for bucket", *bucketName)
  73. return nil
  74. })
  75. return err
  76. }