Style

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
2025-04-23 11:15:23 -07:00
parent b3160366a4
commit b65524d60b
3 changed files with 147 additions and 76 deletions
--- a/pkg/coredata/business_impact.go
+++ b/pkg/coredata/business_impact.go
@@ -16,45 +16,90 @@ package coredata

 import (
 	"database/sql/driver"
+	"encoding/json"
 	"fmt"
 )

-func (bi BusinessImpact) MarshalText() ([]byte, error) {
-	return []byte(bi.String()), nil
+type BusinessImpact string
+
+const (
+	BusinessImpactLow      BusinessImpact = "LOW"
+	BusinessImpactMedium   BusinessImpact = "MEDIUM"
+	BusinessImpactHigh     BusinessImpact = "HIGH"
+	BusinessImpactCritical BusinessImpact = "CRITICAL"
+)
+
+func (i BusinessImpact) String() string {
+	return string(i)
 }

-func (bi *BusinessImpact) UnmarshalText(data []byte) error {
-	val := string(data)
-
-	switch val {
-	case BusinessImpactLow.String():
-		*bi = BusinessImpactLow
-	case BusinessImpactMedium.String():
-		*bi = BusinessImpactMedium
-	case BusinessImpactHigh.String():
-		*bi = BusinessImpactHigh
-	case BusinessImpactCritical.String():
-		*bi = BusinessImpactCritical
+func (i *BusinessImpact) Scan(value interface{}) error {
+	switch v := value.(type) {
+	case string:
+		switch v {
+		case "LOW":
+			*i = BusinessImpactLow
+		case "MEDIUM":
+			*i = BusinessImpactMedium
+		case "HIGH":
+			*i = BusinessImpactHigh
+		case "CRITICAL":
+			*i = BusinessImpactCritical
+		default:
+			return fmt.Errorf("invalid BusinessImpact value: %q", v)
+		}
 	default:
-		return fmt.Errorf("invalid BusinessImpact value: %q", val)
+		return fmt.Errorf("unsupported type for BusinessImpact: %T", value)
 	}
-
 	return nil
 }

-func (bi BusinessImpact) String() string {
-	return string(bi)
+func (i BusinessImpact) Value() (driver.Value, error) {
+	return i.String(), nil
 }

-func (bi *BusinessImpact) Scan(value any) error {
-	val, ok := value.(string)
-	if !ok {
-		return fmt.Errorf("invalid scan source for BusinessImpact, expected string got %T", value)
+func (i BusinessImpact) MarshalJSON() ([]byte, error) {
+	return json.Marshal(i.String())
+}
+
+func (i *BusinessImpact) UnmarshalJSON(data []byte) error {
+	var s string
+	if err := json.Unmarshal(data, &s); err != nil {
+		return err
 	}

-	return bi.UnmarshalText([]byte(val))
+	switch s {
+	case "LOW":
+		*i = BusinessImpactLow
+	case "MEDIUM":
+		*i = BusinessImpactMedium
+	case "HIGH":
+		*i = BusinessImpactHigh
+	case "CRITICAL":
+		*i = BusinessImpactCritical
+	default:
+		return fmt.Errorf("invalid BusinessImpact value: %q", s)
+	}
+	return nil
 }

-func (bi BusinessImpact) Value() (driver.Value, error) {
-	return bi.String(), nil
+func (i *BusinessImpact) UnmarshalText(text []byte) error {
+	var s string
+	if err := json.Unmarshal(text, &s); err != nil {
+		return err
+	}
+
+	switch s {
+	case "LOW":
+		*i = BusinessImpactLow
+	case "MEDIUM":
+		*i = BusinessImpactMedium
+	case "HIGH":
+		*i = BusinessImpactHigh
+	case "CRITICAL":
+		*i = BusinessImpactCritical
+	default:
+		return fmt.Errorf("invalid BusinessImpact value: %q", s)
+	}
+	return nil
 }
--- a/pkg/coredata/data_sensitivity.go
+++ b/pkg/coredata/data_sensitivity.go
@@ -16,47 +16,97 @@ package coredata

 import (
 	"database/sql/driver"
+	"encoding/json"
 	"fmt"
 )

-func (ds DataSensitivity) MarshalText() ([]byte, error) {
-	return []byte(ds.String()), nil
+type DataSensitivity string
+
+const (
+	DataSensitivityNone     DataSensitivity = "NONE"
+	DataSensitivityLow      DataSensitivity = "LOW"
+	DataSensitivityMedium   DataSensitivity = "MEDIUM"
+	DataSensitivityHigh     DataSensitivity = "HIGH"
+	DataSensitivityCritical DataSensitivity = "CRITICAL"
+)
+
+func (i DataSensitivity) String() string {
+	return string(i)
 }

-func (ds *DataSensitivity) UnmarshalText(data []byte) error {
-	val := string(data)
-
-	switch val {
-	case DataSensitivityNone.String():
-		*ds = DataSensitivityNone
-	case DataSensitivityLow.String():
-		*ds = DataSensitivityLow
-	case DataSensitivityMedium.String():
-		*ds = DataSensitivityMedium
-	case DataSensitivityHigh.String():
-		*ds = DataSensitivityHigh
-	case DataSensitivityCritical.String():
-		*ds = DataSensitivityCritical
+func (i *DataSensitivity) Scan(value interface{}) error {
+	switch v := value.(type) {
+	case string:
+		switch v {
+		case "NONE":
+			*i = DataSensitivityNone
+		case "LOW":
+			*i = DataSensitivityLow
+		case "MEDIUM":
+			*i = DataSensitivityMedium
+		case "HIGH":
+			*i = DataSensitivityHigh
+		case "CRITICAL":
+			*i = DataSensitivityCritical
+		default:
+			return fmt.Errorf("invalid DataSensitivity value: %q", v)
+		}
 	default:
-		return fmt.Errorf("invalid DataSensitivity value: %q", val)
+		return fmt.Errorf("unsupported type for DataSensitivity: %T", value)
 	}
-
 	return nil
 }

-func (ds DataSensitivity) String() string {
-	return string(ds)
+func (i DataSensitivity) Value() (driver.Value, error) {
+	return i.String(), nil
 }

-func (ds *DataSensitivity) Scan(value any) error {
-	val, ok := value.(string)
-	if !ok {
-		return fmt.Errorf("invalid scan source for DataSensitivity, expected string got %T", value)
+func (i DataSensitivity) MarshalJSON() ([]byte, error) {
+	return json.Marshal(i.String())
+}
+
+func (i *DataSensitivity) UnmarshalJSON(data []byte) error {
+	var s string
+	if err := json.Unmarshal(data, &s); err != nil {
+		return err
 	}

-	return ds.UnmarshalText([]byte(val))
+	switch s {
+	case "NONE":
+		*i = DataSensitivityNone
+	case "LOW":
+		*i = DataSensitivityLow
+	case "MEDIUM":
+		*i = DataSensitivityMedium
+	case "HIGH":
+		*i = DataSensitivityHigh
+	case "CRITICAL":
+		*i = DataSensitivityCritical
+	default:
+		return fmt.Errorf("invalid DataSensitivity value: %q", s)
+	}
+	return nil
 }

-func (ds DataSensitivity) Value() (driver.Value, error) {
-	return ds.String(), nil
+func (i *DataSensitivity) UnmarshalText(text []byte) error {
+	var s string
+	if err := json.Unmarshal(text, &s); err != nil {
+		return err
+	}
+
+	switch s {
+	case "NONE":
+		*i = DataSensitivityNone
+	case "LOW":
+		*i = DataSensitivityLow
+	case "MEDIUM":
+		*i = DataSensitivityMedium
+	case "HIGH":
+		*i = DataSensitivityHigh
+	case "CRITICAL":
+		*i = DataSensitivityCritical
+	default:
+		return fmt.Errorf("invalid DataSensitivity value: %q", s)
+	}
+	return nil
 }
--- a/pkg/coredata/vendor_risk_assessment.go
+++ b/pkg/coredata/vendor_risk_assessment.go
@@ -43,33 +43,9 @@ type (
 		UpdatedAt       time.Time       `db:"updated_at"`
 	}

-	// DataSensitivity represents the level of data sensitivity handled by a vendor
-	DataSensitivity string
-
-	// BusinessImpact represents the impact level the vendor has on business operations
-	BusinessImpact string
-
-	// RiskAssessments is a collection of RiskAssessment objects
 	VendorRiskAssessments []*VendorRiskAssessment
 )

-// Constants for DataSensitivity
-const (
-	DataSensitivityNone     DataSensitivity = "NONE"     // No sensitive data
-	DataSensitivityLow      DataSensitivity = "LOW"      // Public or non-sensitive data
-	DataSensitivityMedium   DataSensitivity = "MEDIUM"   // Internal/restricted data
-	DataSensitivityHigh     DataSensitivity = "HIGH"     // Confidential data
-	DataSensitivityCritical DataSensitivity = "CRITICAL" // Regulated/PII/financial data
-)
-
-// Constants for BusinessImpact
-const (
-	BusinessImpactLow      BusinessImpact = "LOW"      // Minimal impact on business
-	BusinessImpactMedium   BusinessImpact = "MEDIUM"   // Moderate impact on business
-	BusinessImpactHigh     BusinessImpact = "HIGH"     // Significant business impact
-	BusinessImpactCritical BusinessImpact = "CRITICAL" // Critical to business operations
-)
-
 func (v VendorRiskAssessment) CursorKey(orderBy VendorRiskAssessmentOrderField) page.CursorKey {
 	switch orderBy {
 	case VendorRiskAssessmentOrderFieldCreatedAt: