android_compile/go/android_compile/datastore.go

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

package main

import (
	"context"
	"fmt"
	"sort"
	"time"

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

type CompileTask struct {
	Issue    int    `json:"issue"`
	PatchSet int    `json:"patchset"`
	Hash     string `json:"hash"`

	LunchTarget string `json:"lunch_target"`
	MMMATargets string `json:"mmma_targets"`

	Checkout string `json:"checkout"`

	Created   time.Time `json:"created"`
	Completed time.Time `json:"completed"`

	WithPatchSucceeded bool `json:"withpatch_success"`
	NoPatchSucceeded   bool `json:"nopatch_success"`

	WithPatchLog string `json:"withpatch_log"`
	NoPatchLog   string `json:"nopatch_log"`

	IsMasterBranch bool `json:"is_master_branch"`
	Done           bool `json:"done"`
	InfraFailure   bool `json:"infra_failure"`
}

type CompileTaskAndKey struct {
	task *CompileTask
	key  *datastore.Key
}
type sortTasks []*CompileTaskAndKey

func (a sortTasks) Len() int      { return len(a) }
func (a sortTasks) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a sortTasks) Less(i, j int) bool {
	return a[i].task.Created.Before(a[j].task.Created)
}

func GetCompileTasksAndKeys() ([]*CompileTaskAndKey, []*CompileTaskAndKey, error) {
	waitingTasksAndKeys := []*CompileTaskAndKey{}
	runningTasksAndKeys := []*CompileTaskAndKey{}

	it := GetPendingTasks()
	for {
		t := &CompileTask{}
		datastoreKey, err := it.Next(t)
		if err == iterator.Done {
			break
		} else if err != nil {
			return nil, nil, fmt.Errorf("Failed to retrieve list of tasks: %s", err)
		}
		taskAndKey := &CompileTaskAndKey{task: t, key: datastoreKey}
		if t.Checkout == "" {
			waitingTasksAndKeys = append(waitingTasksAndKeys, taskAndKey)
		} else {
			runningTasksAndKeys = append(runningTasksAndKeys, taskAndKey)
		}
	}
	sort.Sort(sortTasks(waitingTasksAndKeys))
	sort.Sort(sortTasks(runningTasksAndKeys))

	return waitingTasksAndKeys, runningTasksAndKeys, nil
}

func GetCompileTasks() ([]*CompileTask, []*CompileTask, error) {
	waitingTasksAndKeys, runningTasksAndKeys, err := GetCompileTasksAndKeys()
	if err != nil {
		return nil, nil, err
	}
	waitingTasks := []*CompileTask{}
	for _, taskAndKey := range waitingTasksAndKeys {
		waitingTasks = append(waitingTasks, taskAndKey.task)
	}
	runningTasks := []*CompileTask{}
	for _, taskAndKey := range runningTasksAndKeys {
		runningTasks = append(runningTasks, taskAndKey.task)
	}
	return waitingTasks, runningTasks, nil
}

func DatastoreInit(project string, ns string, ts oauth2.TokenSource) error {
	return ds.InitWithOpt(project, ns, option.WithTokenSource(ts))
}

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

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

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

	task := &CompileTask{}
	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(ctx context.Context, k *datastore.Key, t *CompileTask) (*datastore.Key, error) {
	return ds.DS.Put(ctx, k, t)
}

func UpdateDSTask(ctx context.Context, k *datastore.Key, t *CompileTask) (*datastore.Key, error) {
	return ds.DS.Put(ctx, k, t)
}