refactor: header validation and add custom error

pyadav · pyadav · commit 23895f16fb93 · 2024-01-31T23:21:53.000+05:30
diff --git a/mobius/internal/providers/azure/base.go b/mobius/internal/providers/azure/base.go
@@ -3,7 +3,6 @@ package azure
 import (
 	"net/http"
 
-	"github.com/go-playground/validator/v10"
 	"github.com/missingstudio/studio/backend/internal/providers/base"
 	"github.com/missingstudio/studio/backend/pkg/utils"
 	"github.com/missingstudio/studio/common/errors"
@@ -12,10 +11,10 @@ import (
 type AzureProviderFactory struct{}
 
 type AzureHeaders struct {
-	APIKey       string `validate:"required" json:"Authorization"`
-	ResourceName string `validate:"required" json:"X-Ms-Azure-Resource-Name"`
-	DeploymentID string `validate:"required" json:"X-Ms-Deployment-ID"`
-	APIVersion   string `validate:"required" json:"X-Ms-API-Version"`
+	APIKey       string `validate:"required" json:"Authorization" error:"API key is required"`
+	ResourceName string `validate:"required" json:"X-Ms-Azure-Resource-Name" error:"Resource Name is required"`
+	DeploymentID string `validate:"required" json:"X-Ms-Deployment-ID" error:"Deployment ID is required"`
+	APIVersion   string `validate:"required" json:"X-Ms-API-Version" error:"API Version is required"`
 }
 
 func (azf AzureProviderFactory) Validate(headers http.Header) (*AzureHeaders, error) {
@@ -24,9 +23,8 @@ func (azf AzureProviderFactory) Validate(headers http.Header) (*AzureHeaders, er
 		return nil, errors.New(err)
 	}
 
-	validate := validator.New()
-	if err := validate.Struct(azHeaders); err != nil {
-		return nil, errors.NewBadRequest("provider's required headers are missing")
+	if err := utils.ValidateHeaders(azHeaders); err != nil {
+		return nil, err
 	}
 
 	return &azHeaders, nil
diff --git a/mobius/internal/providers/openai/base.go b/mobius/internal/providers/openai/base.go
@@ -4,7 +4,6 @@ import (
 	"net/http"
 	"strings"
 
-	"github.com/go-playground/validator/v10"
 	"github.com/missingstudio/studio/backend/internal/providers/base"
 	"github.com/missingstudio/studio/backend/pkg/utils"
 	"github.com/missingstudio/studio/common/errors"
@@ -13,18 +12,17 @@ import (
 type OpenAIProviderFactory struct{}
 
 type OpenAIHeaders struct {
-	APIKey string `validate:"required" json:"Authorization"`
+	APIKey string `validate:"required" json:"Authorization" error:"API key is required"`
 }
 
-func (oaif OpenAIProviderFactory) Validate(headers http.Header) (*OpenAIHeaders, error) {
+func (azf OpenAIProviderFactory) Validate(headers http.Header) (*OpenAIHeaders, error) {
 	var oaiHeaders OpenAIHeaders
 	if err := utils.UnmarshalHeader(headers, &oaiHeaders); err != nil {
 		return nil, errors.New(err)
 	}
 
-	validate := validator.New()
-	if err := validate.Struct(oaiHeaders); err != nil {
-		return nil, errors.NewBadRequest("provider's required headers are missing")
+	if err := utils.ValidateHeaders(oaiHeaders); err != nil {
+		return nil, err
 	}
 
 	return &oaiHeaders, nil
diff --git a/mobius/pkg/utils/headers.go b/mobius/pkg/utils/headers.go
@@ -1,8 +1,13 @@
 package utils
 
 import (
+	"fmt"
 	"net/http"
 	"reflect"
+	"strings"
+
+	"github.com/go-playground/validator/v10"
+	"github.com/missingstudio/studio/common/errors"
 )
 
 // UnmarshalHeader unmarshals an http.Header into a struct
@@ -22,3 +27,28 @@ func UnmarshalHeader(header http.Header, v interface{}) error {
 
 	return nil
 }
+
+// ValidateHeaders is a generic function to validate any structure with the `validate` struct tag.
+func ValidateHeaders(data interface{}) error {
+	validate := validator.New()
+	if err := validate.Struct(data); err != nil {
+		errorMessages := []string{}
+
+		// Collect all validation errors
+		for _, e := range err.(validator.ValidationErrors) {
+			fieldName := e.Field()
+			field, _ := reflect.TypeOf(data).FieldByName(fieldName)
+			customMessage := field.Tag.Get("error")
+
+			if customMessage == "" {
+				errorMessages = append(errorMessages, fmt.Sprintf("Validation error on field %s", fieldName))
+			} else {
+				errorMessages = append(errorMessages, customMessage)
+			}
+		}
+
+		return errors.New(fmt.Errorf("Validation failed: %v", strings.Join(errorMessages, ", ")))
+	}
+
+	return nil
+}
