From 16df743c451597698593b617b18d369fcf139ae0 Mon Sep 17 00:00:00 2001
From: Kharitonov Maksim Andreevich <makharitonov@ozon.ru>
Date: Wed, 27 Jul 2022 13:47:26 +0300
Subject: [PATCH] rgw/admin: Individual bucket quota support

---
 rgw/admin/bucket_quota.go      | 28 ++++++++++++++++++++
 rgw/admin/bucket_quota_test.go | 48 ++++++++++++++++++++++++++++++++++
 rgw/admin/errors.go            |  1 +
 rgw/admin/quota.go             |  1 +
 4 files changed, 78 insertions(+)
 create mode 100644 rgw/admin/bucket_quota.go
 create mode 100644 rgw/admin/bucket_quota_test.go

diff --git a/rgw/admin/bucket_quota.go b/rgw/admin/bucket_quota.go
new file mode 100644
index 0000000..87b225c
--- /dev/null
+++ b/rgw/admin/bucket_quota.go
@@ -0,0 +1,28 @@
+//go:build ceph_preview
+// +build ceph_preview
+
+package admin
+
+import (
+	"context"
+	"net/http"
+)
+
+// SetIndividualBucketQuota sets quota to a specific bucket
+// https://docs.ceph.com/en/latest/radosgw/adminops/#set-quota-for-an-individual-bucket
+func (api *API) SetIndividualBucketQuota(ctx context.Context, quota QuotaSpec) error {
+	if quota.UID == "" {
+		return errMissingUserID
+	}
+
+	if quota.Bucket == "" {
+		return errMissingUserBucket
+	}
+
+	_, err := api.call(ctx, http.MethodPut, "/bucket?quota", valueToURLParams(quota, []string{"bucket", "uid", "enabled", "max-size", "max-size-kb", "max-objects"}))
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
diff --git a/rgw/admin/bucket_quota_test.go b/rgw/admin/bucket_quota_test.go
new file mode 100644
index 0000000..70bd383
--- /dev/null
+++ b/rgw/admin/bucket_quota_test.go
@@ -0,0 +1,48 @@
+//go:build ceph_preview
+// +build ceph_preview
+
+package admin
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+var testBucketQuota = 1000000
+
+func (suite *RadosGWTestSuite) TestBucketQuota() {
+	suite.SetupConnection()
+	co, err := New(suite.endpoint, suite.accessKey, suite.secretKey, newDebugHTTPClient(http.DefaultClient))
+	assert.NoError(suite.T(), err)
+
+	s3, err := newS3Agent(suite.accessKey, suite.secretKey, suite.endpoint, true)
+	assert.NoError(suite.T(), err)
+
+	err = s3.createBucket(suite.bucketTestName)
+	assert.NoError(suite.T(), err)
+
+	suite.T().Run("set bucket quota but no user is specified", func(t *testing.T) {
+		err := co.SetIndividualBucketQuota(context.Background(), QuotaSpec{})
+		assert.Error(suite.T(), err)
+		assert.EqualError(suite.T(), err, errMissingUserID.Error())
+	})
+
+	suite.T().Run("set bucket quota but no bucket is specified", func(t *testing.T) {
+		err := co.SetIndividualBucketQuota(context.Background(), QuotaSpec{UID: "admin"})
+		assert.Error(suite.T(), err)
+		assert.EqualError(suite.T(), err, errMissingUserBucket.Error())
+	})
+
+	suite.T().Run("set bucket quota", func(t *testing.T) {
+		err := co.SetIndividualBucketQuota(context.Background(), QuotaSpec{UID: "admin", Bucket: suite.bucketTestName, MaxSizeKb: &testBucketQuota})
+		assert.NoError(suite.T(), err)
+
+		bucketInfo, err := co.GetBucketInfo(context.Background(), Bucket{Bucket: suite.bucketTestName})
+		assert.NoError(suite.T(), err)
+
+		assert.Equal(suite.T(), &testBucketQuota, bucketInfo.BucketQuota.MaxSizeKb)
+	})
+}
diff --git a/rgw/admin/errors.go b/rgw/admin/errors.go
index 7ce3c64..58e6641 100644
--- a/rgw/admin/errors.go
+++ b/rgw/admin/errors.go
@@ -93,6 +93,7 @@ var (
 	errMissingUserCap         = errors.New("missing user capabilities")
 	errMissingBucketID        = errors.New("missing bucket ID")
 	errMissingBucket          = errors.New("missing bucket")
+	errMissingUserBucket      = errors.New("missing bucket")
 )
 
 // errorReason is the reason of the error
diff --git a/rgw/admin/quota.go b/rgw/admin/quota.go
index 5ebac86..9184ee5 100644
--- a/rgw/admin/quota.go
+++ b/rgw/admin/quota.go
@@ -11,6 +11,7 @@ import (
 // Only user's quota are supported
 type QuotaSpec struct {
 	UID        string `json:"user_id" url:"uid"`
+	Bucket     string `json:"bucket" url:"bucket"`
 	QuotaType  string `url:"quota-type"`
 	Enabled    *bool  `json:"enabled" url:"enabled"`
 	CheckOnRaw bool   `json:"check_on_raw"`