github.com/tickoalcantara12/micro/v3@v3.0.0-20221007104245-9d75b9bcbab9/client/cli/store/snapshot/restore.go (about)

     1  package snapshot
     2  
     3  import (
     4  	"encoding/gob"
     5  	"io"
     6  	"net/url"
     7  	"os"
     8  	"time"
     9  
    10  	"github.com/tickoalcantara12/micro/v3/service/store"
    11  	"github.com/pkg/errors"
    12  )
    13  
    14  // Restore emits records from a go-micro store snapshot
    15  type Restore interface {
    16  	// Init validates the RestoreOptions and returns an error if they are invalid.
    17  	// Init must be called before a Restore is used
    18  	Init(opts ...RestoreOption) error
    19  	// Start opens a channel over which records from the snapshot are retrieved.
    20  	// The channel will be closed when the entire snapshot has been read.
    21  	Start() (<-chan *store.Record, error)
    22  }
    23  
    24  // RestoreOptions configure a Restore
    25  type RestoreOptions struct {
    26  	Source string
    27  }
    28  
    29  // RestoreOption is an individual option
    30  type RestoreOption func(r *RestoreOptions)
    31  
    32  // Source is the source URL of a snapshot, e.g. file:///path/to/file
    33  func Source(source string) RestoreOption {
    34  	return func(r *RestoreOptions) {
    35  		r.Source = source
    36  	}
    37  }
    38  
    39  // FileRestore reads records from a file
    40  type FileRestore struct {
    41  	Options RestoreOptions
    42  
    43  	path string
    44  }
    45  
    46  func NewFileRestore(opts ...RestoreOption) Restore {
    47  	r := &FileRestore{}
    48  	for _, o := range opts {
    49  		o(&r.Options)
    50  	}
    51  	return r
    52  }
    53  
    54  func (f *FileRestore) Init(opts ...RestoreOption) error {
    55  	for _, o := range opts {
    56  		o(&f.Options)
    57  	}
    58  	u, err := url.Parse(f.Options.Source)
    59  	if err != nil {
    60  		return errors.Wrap(err, "source is invalid")
    61  	}
    62  	if u.Scheme != "file" {
    63  		return errors.Errorf("unsupported scheme %s (wanted file)", u.Scheme)
    64  	}
    65  	f.path = u.Path
    66  	return nil
    67  }
    68  
    69  // Start starts reading records from a file. The returned channel is closed when complete
    70  func (f *FileRestore) Start() (<-chan *store.Record, error) {
    71  	fi, err := os.Open(f.path)
    72  	if err != nil {
    73  		return nil, errors.Wrapf(err, "Couldn't open file %s", f.path)
    74  	}
    75  	recordChan := make(chan *store.Record)
    76  	go func(records chan<- *store.Record, reader io.ReadCloser) {
    77  		defer close(recordChan)
    78  		defer reader.Close()
    79  		dec := gob.NewDecoder(fi)
    80  		var r record
    81  		for {
    82  			err := dec.Decode(&r)
    83  			if err == io.EOF {
    84  				break
    85  			}
    86  			if err != nil {
    87  				panic(err)
    88  			}
    89  			rec := &store.Record{
    90  				Key: r.Key,
    91  			}
    92  			rec.Value = make([]byte, len(r.Value))
    93  			copy(rec.Value, r.Value)
    94  			if !r.ExpiresAt.IsZero() {
    95  				rec.Expiry = time.Until(r.ExpiresAt)
    96  			}
    97  			recordChan <- rec
    98  		}
    99  	}(recordChan, fi)
   100  	return recordChan, nil
   101  }