github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/utils/jsondb/jsondb.go (about)

     1  // Copyright (c) Tailscale Inc & AUTHORS
     2  // SPDX-License-Identifier: BSD-3-Clause
     3  
     4  // Package jsondb provides a trivial "database": a Go object saved to
     5  // disk as JSON.
     6  package jsondb
     7  
     8  import (
     9  	"os"
    10  	"reflect"
    11  
    12  	"github.com/Asutorufa/yuhaiin/pkg/log"
    13  	"google.golang.org/protobuf/encoding/protojson"
    14  	"google.golang.org/protobuf/proto"
    15  	"google.golang.org/protobuf/reflect/protoreflect"
    16  )
    17  
    18  // DB is a database backed by a JSON file.
    19  type DB[T proto.Message] struct {
    20  	// Data is the contents of the database.
    21  	Data T
    22  
    23  	path string
    24  }
    25  
    26  // Open opens the database at path, creating it with a zero value if
    27  // necessary.
    28  func Open[T proto.Message](path string, defaultValue T) *DB[T] {
    29  	val := reflect.New(reflect.TypeOf(defaultValue).Elem()).Interface().(T)
    30  
    31  	bs, err := os.ReadFile(path)
    32  	if err == nil {
    33  		err = protojson.UnmarshalOptions{DiscardUnknown: true}.Unmarshal(bs, val)
    34  		if err != nil {
    35  			log.Warn("proto json unmarshal failed", "err", err)
    36  		}
    37  	} else {
    38  		log.Warn("open jsonDB file failed", "path", path, "err", err)
    39  	}
    40  
    41  	MergeDefault(val.ProtoReflect(), defaultValue.ProtoReflect())
    42  
    43  	return &DB[T]{
    44  		Data: val,
    45  		path: path,
    46  	}
    47  }
    48  
    49  func MergeDefault(src, def protoreflect.Message) {
    50  	def.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool {
    51  
    52  		vi, vok := v.Interface().(protoreflect.Message)
    53  
    54  		if src.IsValid() && !src.Has(fd) && vok {
    55  			src.Set(fd, v)
    56  		}
    57  
    58  		sv := src.Get(fd)
    59  
    60  		svi, sok := sv.Interface().(protoreflect.Message)
    61  
    62  		if sok && vok {
    63  			MergeDefault(svi, vi)
    64  		}
    65  
    66  		return true
    67  	})
    68  }
    69  
    70  func New[T proto.Message](t T, path string) *DB[T] { return &DB[T]{t, path} }
    71  
    72  // Save writes db.Data back to disk.
    73  func (db *DB[T]) Save() error {
    74  	bs, err := protojson.MarshalOptions{Multiline: true, Indent: "\t", EmitUnpopulated: true}.Marshal(db.Data)
    75  	if err != nil {
    76  		return err
    77  	}
    78  
    79  	return os.WriteFile(db.path, bs, 0600)
    80  }