autoroll/go/status/firestore.go

package status

import (
	"context"
	"time"

	fs "cloud.google.com/go/firestore"
	"go.skia.org/infra/go/firestore"
	"go.skia.org/infra/go/skerr"
	"go.skia.org/infra/go/util"
	"golang.org/x/oauth2"
)

const (
	// App name used for Firestore.
	fsApp = "autoroll"

	// Collection name for AutoRollStatus.
	collectionStatus = "status"

	// Firestore-related constants.
	defaultAttempts = 3
	defaultTimeout  = 10 * time.Second
)

// FirestoreDB implements DB using Firestore.
type FirestoreDB struct {
	client   *firestore.Client
	coll     *fs.CollectionRef
	exportDB DB
}

// NewFirestoreDBWithParams returns a FirestoreDB instance using the given params.
func NewFirestoreDBWithParams(ctx context.Context, project, namespace, instance string, ts oauth2.TokenSource) (*FirestoreDB, error) {
	client, err := firestore.NewClient(ctx, project, namespace, instance, ts)
	if err != nil {
		return nil, skerr.Wrap(err)
	}
	return NewFirestoreDB(ctx, client)
}

// NewFirestoreDB returns a FirestoreDB instance using the given Client.
func NewFirestoreDB(ctx context.Context, client *firestore.Client) (*FirestoreDB, error) {
	return &FirestoreDB{
		client: client,
		coll:   client.Collection(collectionStatus),
	}, nil
}

// Close the DB.
func (d *FirestoreDB) Close() error {
	return skerr.Wrap(d.client.Close())
}

// Get implements DB.
func (d *FirestoreDB) Get(ctx context.Context, rollerID string) (*AutoRollStatus, error) {
	ref := d.coll.Doc(rollerID)
	doc, err := d.client.Get(ctx, ref, defaultAttempts, defaultTimeout)
	if err != nil {
		return nil, skerr.Wrapf(err, "retrieving config for %s", rollerID)
	}
	rv := new(AutoRollStatus)
	if err := doc.DataTo(rv); err != nil {
		return nil, skerr.Wrapf(err, "decoding config for %s", rollerID)
	}
	return rv, nil
}

// Set implements DB.
func (d *FirestoreDB) Set(ctx context.Context, rollerID string, st *AutoRollStatus) error {
	ref := d.coll.Doc(rollerID)
	if _, err := ref.Set(ctx, st); err != nil {
		return skerr.Wrap(err)
	}
	if d.exportDB != nil && util.In(rollerID, exportRollers) {
		exportStatus := &AutoRollStatus{
			AutoRollMiniStatus: st.AutoRollMiniStatus,
		}
		if err := d.exportDB.Set(ctx, rollerID, exportStatus); err != nil {
			return skerr.Wrap(err)
		}
	}
	return nil
}

var _ DB = &FirestoreDB{}