leasing/go/leasing/datastore.go

/*
	Used by the Leasing Server to interact with the cloud datastore.
*/

package main

import (
	"context"
	"fmt"

	"cloud.google.com/go/datastore"
	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/ds"
	"google.golang.org/api/option"
)

func DatastoreInit(project string, ns string) error {
	ts, err := auth.NewDefaultTokenSource(*local, "https://www.googleapis.com/auth/datastore")
	if err != nil {
		return fmt.Errorf("Problem setting up default token source: %s", err)
	}
	return ds.InitWithOpt(project, ns, option.WithTokenSource(ts))
}

func GetRunningDSTasks() *datastore.Iterator {
	q := ds.NewQuery(ds.TASK).EventualConsistency().Filter("Done =", false)
	return ds.DS.Run(context.TODO(), q)
}

func GetAllDSTasks(filterUser string) *datastore.Iterator {
	q := ds.NewQuery(ds.TASK).EventualConsistency()
	if filterUser != "" {
		q = q.Filter("Requester =", filterUser)
	}
	return ds.DS.Run(context.TODO(), q)
}

func GetNewDSKey() *datastore.Key {
	return ds.NewKey(ds.TASK)
}

func GetDSTask(taskID int64) (*datastore.Key, *Task, error) {
	key := ds.NewKey(ds.TASK)
	key.ID = taskID

	task := &Task{}
	if err := ds.DS.Get(context.TODO(), key, task); err != nil {
		return nil, nil, fmt.Errorf("Error retrieving task from Datastore: %v", err)
	}
	return key, task, nil
}

func PutDSTask(k *datastore.Key, t *Task) (*datastore.Key, error) {
	return ds.DS.Put(context.Background(), k, t)
}

func UpdateDSTask(k *datastore.Key, t *Task) (*datastore.Key, error) {
	return ds.DS.Put(context.Background(), k, t)
}