Referal api (#2792)

## Description

## Tests

server/cmd/museum/main.go

@@ -612,6 +612,7 @@ func main() {
 	}
 
 	privateAPI.GET("/storage-bonus/details", storageBonusHandler.GetStorageBonusDetails)
+	privateAPI.POST("/storage-bonus/change-code", storageBonusHandler.UpdateReferralCode)
 	privateAPI.GET("/storage-bonus/referral-view", storageBonusHandler.GetReferralView)
 	privateAPI.POST("/storage-bonus/referral-claim", storageBonusHandler.ClaimReferral)
 

server/ente/storagebonus/referral.go

@@ -18,6 +18,10 @@ type UserReferralPlanStat struct {
 	UpgradedCount int      `json:"upgradedCount"`
 }
 
+type UpdateReferralCodeRequest struct {
+	Code string `json:"code" binding:"required"`
+}
+
 // PlanInfo represents the referral plan metadata
 type PlanInfo struct {
 	// IsEnabled indicates if the referral plan is enabled for given user

server/pkg/api/admin.go

@@ -243,7 +243,7 @@ func (h *AdminHandler) UpdateReferral(c *gin.Context) {
 	}
 	go h.DiscordController.NotifyAdminAction(
 		fmt.Sprintf("Admin (%d) updating referral code for %d to %s", auth.GetUserID(c.Request.Header), request.UserID, request.Code))
-	err := h.StorageBonusCtl.UpdateReferralCode(c, request.UserID, request.Code)
+	err := h.StorageBonusCtl.UpdateReferralCode(c, request.UserID, request.Code, true)
 	if err != nil {
 		logrus.WithError(err).Error("Failed to disable 2FA")
 		handler.Error(c, stacktrace.Propagate(err, ""))

server/pkg/api/storage_bonus.go

@@ -4,6 +4,7 @@ import (
 	"net/http"
 
 	"github.com/ente-io/museum/ente"
+	entity "github.com/ente-io/museum/ente/storagebonus"
 	"github.com/ente-io/museum/pkg/controller/storagebonus"
 	"github.com/ente-io/museum/pkg/utils/auth"
 	"github.com/ente-io/museum/pkg/utils/handler"
@@ -24,6 +25,21 @@ func (h StorageBonusHandler) GetReferralView(context *gin.Context) {
 	context.JSON(http.StatusOK, response)
 }
 
+func (h StorageBonusHandler) UpdateReferralCode(context *gin.Context) {
+	var request entity.UpdateReferralCodeRequest
+	if err := context.ShouldBindJSON(&request); err != nil {
+		handler.Error(context, stacktrace.Propagate(ente.NewBadRequestWithMessage(err.Error()), ""))
+		return
+	}
+	userID := auth.GetUserID(context.Request.Header)
+	err := h.Controller.UpdateReferralCode(context, userID, request.Code, false)
+	if err != nil {
+		handler.Error(context, stacktrace.Propagate(err, ""))
+		return
+	}
+	context.JSON(http.StatusOK, gin.H{})
+}
+
 func (h StorageBonusHandler) GetStorageBonusDetails(context *gin.Context) {
 	response, err := h.Controller.GetStorageBonusDetailResponse(context, auth.GetUserID(context.Request.Header))
 	if err != nil {

server/pkg/controller/storagebonus/referral.go

@@ -133,15 +133,15 @@ func (c *Controller) GetOrCreateReferralCode(ctx *gin.Context, userID int64) (*s
 	return referralCode, nil
 }
 
-func (c *Controller) UpdateReferralCode(ctx *gin.Context, userID int64, code string) error {
+func (c *Controller) UpdateReferralCode(ctx *gin.Context, userID int64, code string, isAdminEdit bool) error {
 	code = strings.ToUpper(code)
 	if !random.IsAlphanumeric(code) {
 		return stacktrace.Propagate(ente.NewBadRequestWithMessage("code is not alphanumeric"), "")
 	}
 	if len(code) < 4 || len(code) > 8 {
 		return stacktrace.Propagate(ente.NewBadRequestWithMessage("code length should be between 4 and 8"), "")
 	}
-	err := c.StorageBonus.AddNewCode(ctx, userID, code)
+	err := c.StorageBonus.AddNewCode(ctx, userID, code, isAdminEdit)
 	if err != nil {
 		return stacktrace.Propagate(err, "failed to update referral code")
 	}

server/pkg/repo/storagebonus/referral_codes.go

@@ -39,23 +39,32 @@ func (r *Repository) InsertCode(ctx context.Context, userID int64, code string)
 // AddNewCode and mark the old one as inactive for a given userID.
 // Note: This method is not being used in the initial MVP as we don't allow user to change the storagebonus
 // code
-func (r *Repository) AddNewCode(ctx context.Context, userID int64, code string) error {
+func (r *Repository) AddNewCode(ctx context.Context, userID int64, code string, isAdminEdit bool) error {
 	// check current referral code count
 	var count int
 	err := r.DB.QueryRowContext(ctx, "SELECT COALESCE(COUNT(*),0) FROM referral_codes WHERE user_id = $1", userID).Scan(&count)
 	if err != nil {
 		return stacktrace.Propagate(err, "failed to get storagebonus code count for user %d", userID)
 	}
-	if count > maxReferralChangeAllowed {
+	if !isAdminEdit && count > maxReferralChangeAllowed {
 		return stacktrace.Propagate(&ente.ApiError{
 			Code:           "REFERRAL_CHANGE_LIMIT_REACHED",
 			Message:        fmt.Sprintf("max referral code change limit %d reached", maxReferralChangeAllowed),
 			HttpStatusCode: http.StatusTooManyRequests,
 		}, "max referral code change limit reached for user %d", userID)
 	}
+	// check if code already exists
+	var existCount int
+	err = r.DB.QueryRowContext(ctx, "SELECT COALESCE(COUNT(*),0) FROM referral_codes WHERE code = $1", code).Scan(&existCount)
+	if err != nil {
+		return stacktrace.Propagate(err, "failed to check if code already exists for user %d", userID)
+	}
+	if existCount > 0 {
+		return stacktrace.Propagate(entity.CodeAlreadyExistsErr, "storagebonus code %s already exists", code)
+	}
 	_, err = r.DB.ExecContext(ctx, "UPDATE referral_codes SET is_active = FALSE WHERE user_id = $1", userID)
 	if err != nil {
-		return stacktrace.Propagate(err, "failed to update storagebonus code for user %d", userID)
+		return stacktrace.Propagate(err, "failed to update remove existing code code for user %d", userID)
 	}
 	return r.InsertCode(ctx, userID, code)
 }

server/pkg/repo/storagebonus/referral_codes_test.go

@@ -65,7 +65,7 @@ func TestAddNewReferralCode(t *testing.T) {
 	assert.Nil(t, err)
 
 	newCode := "C22222"
-	err = repo.AddNewCode(context.Background(), userID, newCode)
+	err = repo.AddNewCode(context.Background(), userID, newCode, false)
 	assert.Nil(t, err)
 
 	referralCode, err := repo.GetCode(context.Background(), userID)