From e5d6d324dba00523f8852a35d1490eaf1a159ff9 Mon Sep 17 00:00:00 2001
From: Neeraj Gupta <254676+ua741@users.noreply.github.com>
Date: Wed, 21 Aug 2024 16:28:24 +0530
Subject: [PATCH] Add API to update referral code

---
 server/cmd/museum/main.go                      |  1 +
 server/ente/storagebonus/referral.go           |  4 ++++
 server/pkg/api/storage_bonus.go                | 16 ++++++++++++++++
 server/pkg/repo/storagebonus/referral_codes.go |  2 +-
 4 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/server/cmd/museum/main.go b/server/cmd/museum/main.go
index c0e25e4e1c..c6de4f5c56 100644
--- a/server/cmd/museum/main.go
+++ b/server/cmd/museum/main.go
@@ -612,6 +612,7 @@ func main() {
 	}
 
 	privateAPI.GET("/storage-bonus/details", storageBonusHandler.GetStorageBonusDetails)
+	privateAPI.GET("/storage-bonus/change-code", storageBonusHandler.UpdateReferralCode)
 	privateAPI.GET("/storage-bonus/referral-view", storageBonusHandler.GetReferralView)
 	privateAPI.POST("/storage-bonus/referral-claim", storageBonusHandler.ClaimReferral)
 
diff --git a/server/ente/storagebonus/referral.go b/server/ente/storagebonus/referral.go
index 816e0ca4e4..be356012e2 100644
--- a/server/ente/storagebonus/referral.go
+++ b/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
diff --git a/server/pkg/api/storage_bonus.go b/server/pkg/api/storage_bonus.go
index 4f2ba1046b..04d0aff898 100644
--- a/server/pkg/api/storage_bonus.go
+++ b/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)
+	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 {
diff --git a/server/pkg/repo/storagebonus/referral_codes.go b/server/pkg/repo/storagebonus/referral_codes.go
index d3f49ec3d0..d28bc2e0ae 100644
--- a/server/pkg/repo/storagebonus/referral_codes.go
+++ b/server/pkg/repo/storagebonus/referral_codes.go
@@ -55,7 +55,7 @@ func (r *Repository) AddNewCode(ctx context.Context, userID int64, code string)
 	}
 	_, 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)
 }