add marshal/unmarshal wrapper

Author: BeryJu

pkg/roles/api/auth/provider.go

@@ -14,7 +14,6 @@ import (
 	"go.uber.org/zap"
 	"golang.org/x/crypto/bcrypt"
 	"golang.org/x/oauth2"
-	"google.golang.org/protobuf/proto"
 )
 
 type AuthProvider struct {
@@ -79,19 +78,15 @@ func (ap *AuthProvider) CreateUser(ctx context.Context, username, password strin
 			},
 		},
 	}
-	userProto, err := proto.Marshal(&user)
-	if err != nil {
-		return err
-	}
 
-	_, err = ap.inst.KV().Put(
+	_, err = ap.inst.KV().PutObj(
 		ctx,
 		ap.inst.KV().Key(
 			types.KeyRole,
 			types.KeyUsers,
 			username,
 		).String(),
-		string(userProto),
+		&user,
 	)
 	if err != nil {
 		return err

pkg/roles/api/auth/user.go

@@ -7,7 +7,6 @@ import (
 	"beryju.io/gravity/pkg/roles/api/types"
 	"go.etcd.io/etcd/api/v3/mvccpb"
 	clientv3 "go.etcd.io/etcd/client/v3"
-	"google.golang.org/protobuf/proto"
 )
 
 func (ap *AuthProvider) userFromKV(raw *mvccpb.KeyValue) (*types.User, error) {
@@ -17,23 +16,19 @@ func (ap *AuthProvider) userFromKV(raw *mvccpb.KeyValue) (*types.User, error) {
 		types.KeyUsers,
 	).Prefix(true).String()
 	user.Username = strings.TrimPrefix(string(raw.Key), prefix)
-	err := proto.Unmarshal(raw.Value, user)
+	err := ap.inst.KV().Unmarshal(raw.Value, user)
 	if err != nil {
 		return user, err
 	}
 	return user, nil
 }
 
 func (ap *AuthProvider) putUser(u *types.User, ctx context.Context, opts ...clientv3.OpOption) error {
-	raw, err := proto.Marshal(u)
-	if err != nil {
-		return err
-	}
 	fullKey := ap.inst.KV().Key(
 		types.KeyRole,
 		types.KeyUsers,
 		u.Username,
 	).String()
-	_, err = ap.inst.KV().Put(ctx, fullKey, string(raw), opts...)
+	_, err := ap.inst.KV().PutObj(ctx, fullKey, u, opts...)
 	return err
 }

pkg/roles/tsdb/api.go

@@ -13,7 +13,6 @@ import (
 	"github.com/swaggest/usecase/status"
 	clientv3 "go.etcd.io/etcd/client/v3"
 	"go.uber.org/zap"
-	"google.golang.org/protobuf/proto"
 )
 
 func (r *Role) APIMetrics() usecase.Interactor {
@@ -52,7 +51,7 @@ func (r *Role) APIMetrics() usecase.Interactor {
 				continue
 			}
 			v := types.MetricsRecord{}
-			err = proto.Unmarshal(kv.Value, &v)
+			err = r.i.KV().Unmarshal(kv.Value, &v)
 			if err != nil {
 				value, err := strconv.ParseInt(string(kv.Value), 10, 0)
 				if err != nil {

pkg/roles/tsdb/role.go

@@ -20,7 +20,6 @@ import (
 	"github.com/swaggest/rest/web"
 	clientv3 "go.etcd.io/etcd/client/v3"
 	"go.uber.org/zap"
-	"google.golang.org/protobuf/proto"
 )
 
 type Role struct {
@@ -170,14 +169,10 @@ func (r *Role) write(ctx context.Context) {
 		v := types.MetricsRecord{
 			Value: int64(value.Value),
 		}
-		rv, err := proto.Marshal(&v)
-		if err != nil {
-			r.log.Warn("failed to marshal message", zap.Error(err))
-		}
-		_, err = r.i.KV().Put(
+		_, err = r.i.KV().PutObj(
 			ks.Context(),
 			key,
-			string(rv),
+			&v,
 			clientv3.WithLease(lease.ID),
 		)
 		if err != nil {

pkg/storage/client.go

@@ -2,9 +2,11 @@ package storage
 
 import (
 	"context"
+	"encoding/json"
 	"time"
 
 	"beryju.io/gravity/pkg/storage/trace"
+	"github.com/gogo/protobuf/proto"
 	"go.uber.org/zap"
 
 	clientv3 "go.etcd.io/etcd/client/v3"
@@ -157,3 +159,28 @@ func (c *Client) Delete(ctx context.Context, key string, opts ...clientv3.OpOpti
 	}
 	return res, err
 }
+
+func (c *Client) PutObj(ctx context.Context, key string, val any, opts ...clientv3.OpOption) (*clientv3.PutResponse, error) {
+	marshalled, err := c.Marshal(val)
+	if err != nil {
+		return nil, err
+	}
+	return c.Put(ctx, key, string(marshalled), opts...)
+}
+
+func (c *Client) Unmarshal(raw []byte, out any) error {
+	err := json.Unmarshal(raw, out)
+	if err != nil {
+		if p, ok := out.(proto.Message); ok {
+			return proto.Unmarshal(raw, p)
+		}
+	}
+	return err
+}
+
+func (c *Client) Marshal(in any) ([]byte, error) {
+	if p, ok := in.(proto.Message); ok {
+		return proto.Marshal(p)
+	}
+	return json.Marshal(in)
+}