edge/pkg/storage/rpc/client/document_store.go

package client

import (
	"context"
	"encoding/gob"
	"time"

	"github.com/keegancsmith/rpc"
	"github.com/pkg/errors"

	"forge.cadoles.com/arcad/edge/pkg/storage"
	"forge.cadoles.com/arcad/edge/pkg/storage/filter"
	"forge.cadoles.com/arcad/edge/pkg/storage/rpc/server/document"
)

func init() {
	gob.Register(storage.Document{})
	gob.Register(storage.DocumentID(""))
	gob.Register(time.Time{})
	gob.Register(map[string]interface{}{})
	gob.Register([]interface{}{})
	gob.Register([]map[string]interface{}{})
}

type DocumentStore struct {
	client *rpc.Client
}

// Delete implements storage.DocumentStore.
func (s *DocumentStore) Delete(ctx context.Context, collection string, id storage.DocumentID) error {
	args := document.DeleteDocumentArgs{
		Collection: collection,
		DocumentID: id,
	}

	reply := document.DeleteDocumentReply{}

	if err := s.client.Call(ctx, "Service.DeleteDocument", args, &reply); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// Get implements storage.DocumentStore.
func (s *DocumentStore) Get(ctx context.Context, collection string, id storage.DocumentID) (storage.Document, error) {
	args := document.GetDocumentArgs{
		Collection: collection,
		DocumentID: id,
	}

	reply := document.GetDocumentReply{}

	if err := s.client.Call(ctx, "Service.GetDocument", args, &reply); err != nil {
		return nil, errors.WithStack(err)
	}

	return reply.Document, nil
}

// Query implements storage.DocumentStore.
func (s *DocumentStore) Query(ctx context.Context, collection string, filter *filter.Filter, funcs ...storage.QueryOptionFunc) ([]storage.Document, error) {
	opts := &storage.QueryOptions{}
	for _, fn := range funcs {
		fn(opts)
	}

	args := document.QueryDocumentsArgs{
		Collection: collection,
		Filter:     nil,
		Options:    opts,
	}

	if filter != nil {
		args.Filter = filter.AsMap()
	}

	reply := document.QueryDocumentsReply{}

	if err := s.client.Call(ctx, "Service.QueryDocuments", args, &reply); err != nil {
		return nil, errors.WithStack(err)
	}

	return reply.Documents, nil
}

// Upsert implements storage.DocumentStore.
func (s *DocumentStore) Upsert(ctx context.Context, collection string, doc storage.Document) (storage.Document, error) {
	args := document.UpsertDocumentArgs{
		Collection: collection,
		Document:   doc,
	}

	reply := document.UpsertDocumentReply{}

	if err := s.client.Call(ctx, "Service.UpsertDocument", args, &reply); err != nil {
		return nil, errors.WithStack(err)
	}

	return reply.Document, nil
}

func NewDocumentStore(client *rpc.Client) *DocumentStore {
	return &DocumentStore{client}
}

var _ storage.DocumentStore = &DocumentStore{}