<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
)

func (i *Installer) generateCrashReport(sourceError error) {
	// TODO: Use the live endpoint
	//reportURL := "https://metriton.datawire.io/crash-report"
	reportURL := "https://metriton.datawire.io/beta/crash-report"

	report := &crashReportCreationRequest{
		Product:         "edgectl",
		Command:         "install",
		ProductVersion:  displayVersion,
		Error:           sourceError.Error(),
		AESVersion:      i.version,
		Address:         i.address,
		Hostname:        i.hostname,
		ClusterID:       i.clusterID,
		InstallID:       i.scout.installID,
		TraceID:         fmt.Sprintf("%v", i.scout.metadata["trace_id"]),
		ClusterInfo:     fmt.Sprintf("%v", i.scout.metadata["cluster_info"]),
		Managed:         fmt.Sprintf("%v", i.scout.metadata["managed"]),
		KubectlVersion:  i.k8sVersion.Client.GitVersion,
		KubectlPlatform: i.k8sVersion.Client.Platform,
		K8sVersion:      i.k8sVersion.Server.GitVersion,
		K8sPlatform:     i.k8sVersion.Server.Platform,
	}
	buf := new(bytes.Buffer)
	_ = json.NewEncoder(buf).Encode(report)
	resp, err := http.Post(reportURL, "application/json", buf)
	if err != nil {
		i.log.Printf("failed to initiate anonymous crash report due to error: %v", err.Error())
		return
	}
	defer resp.Body.Close()
	content, err := ioutil.ReadAll(resp.Body)
	if resp.StatusCode != 201 {
		i.log.Print("skipping anonymous crash report and log submission for this failure")
		i.log.Printf("%v: %q", resp.StatusCode, string(content))
		return
	}
	crashReport := crashReportCreationResponse{}
	err = json.Unmarshal(content, &crashReport)
	if err != nil {
		i.log.Printf("failed to generate anonymous crash report due to error: %v", err.Error())
		return
	}
	i.log.Printf("uploading anonymous crash report and logs under report ID: %v", crashReport.ReportId)
	i.Report("crash_report", ScoutMeta{"crash_report_id", crashReport.ReportId})
	i.uploadCrashReportData(crashReport, i.gatherCrashReportData())
}

func (i *Installer) gatherCrashReportData() []byte {
	buffer := bytes.NewBuffer([]byte{})

	buffer.WriteString("========== edgectl logs ==========\n")
	fileContent, err := ioutil.ReadFile(i.logName)
	if err != nil {
		i.log.Printf("failed to read log file %v: %v", i.logName, err.Error())
	}
	buffer.Write(fileContent)

	buffer.WriteString("\n========== kubectl describe ==========\n")
	describe, err := i.SilentCaptureKubectl("describe ambassador namespace", "", "-n", "ambassador", "describe", "all")
	if err != nil {
		i.log.Printf("failed to describe ambassador resources: %v", err.Error())
	}
	buffer.WriteString(describe)

	buffer.WriteString("\n========== kubectl logs ==========\n")
	ambassadorLogs, err := i.SilentCaptureKubectl("read ambassador logs", "", "-n", "ambassador", "logs", "deployments/ambassador", "--tail=1000")
	if err != nil {
		i.log.Printf("failed to read ambassador logs: %v", err.Error())
	}
	buffer.WriteString(ambassadorLogs)

	return buffer.Bytes()
}

func (i *Installer) uploadCrashReportData(crashReport crashReportCreationResponse, uploadContent []byte) {
	client := &http.Client{}
	req, err := http.NewRequest(crashReport.Method, crashReport.UploadURL, bytes.NewReader(uploadContent))
	if err != nil {
		i.log.Print(err.Error())
		return
	}

	res, err := client.Do(req)
	if err != nil {
		i.log.Print(err.Error())
		return
	}
	defer res.Body.Close()
	_, err = ioutil.ReadAll(res.Body)
	if err != nil {
		i.log.Print(err.Error())
		return
	}
}

// crashReportCreationRequest is used to initiate a crash report request
type crashReportCreationRequest struct {
	Product         string
	ProductVersion  string
	Command         string
	Error           string
	AESVersion      string
	Address         string
	Hostname        string
	ClusterID       string
	InstallID       string
	TraceID         string
	ClusterInfo     string
	Managed         string
	KubectlVersion  string
	KubectlPlatform string
	K8sVersion      string
	K8sPlatform     string
}

// crashReportCreationResponse is used to receive a crash report creation response
type crashReportCreationResponse struct {
	ReportId  string
	Method    string
	UploadURL string
}
<commit_msg>use the live metriton endpoint<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
)

func (i *Installer) generateCrashReport(sourceError error) {
	reportURL := "https://metriton.datawire.io/crash-report"

	report := &crashReportCreationRequest{
		Product:         "edgectl",
		Command:         "install",
		ProductVersion:  displayVersion,
		Error:           sourceError.Error(),
		AESVersion:      i.version,
		Address:         i.address,
		Hostname:        i.hostname,
		ClusterID:       i.clusterID,
		InstallID:       i.scout.installID,
		TraceID:         fmt.Sprintf("%v", i.scout.metadata["trace_id"]),
		ClusterInfo:     fmt.Sprintf("%v", i.scout.metadata["cluster_info"]),
		Managed:         fmt.Sprintf("%v", i.scout.metadata["managed"]),
		KubectlVersion:  i.k8sVersion.Client.GitVersion,
		KubectlPlatform: i.k8sVersion.Client.Platform,
		K8sVersion:      i.k8sVersion.Server.GitVersion,
		K8sPlatform:     i.k8sVersion.Server.Platform,
	}
	buf := new(bytes.Buffer)
	_ = json.NewEncoder(buf).Encode(report)
	resp, err := http.Post(reportURL, "application/json", buf)
	if err != nil {
		i.log.Printf("failed to initiate anonymous crash report due to error: %v", err.Error())
		return
	}
	defer resp.Body.Close()
	content, err := ioutil.ReadAll(resp.Body)
	if resp.StatusCode != 201 {
		i.log.Print("skipping anonymous crash report and log submission for this failure")
		i.log.Printf("%v: %q", resp.StatusCode, string(content))
		return
	}
	crashReport := crashReportCreationResponse{}
	err = json.Unmarshal(content, &crashReport)
	if err != nil {
		i.log.Printf("failed to generate anonymous crash report due to error: %v", err.Error())
		return
	}
	i.log.Printf("uploading anonymous crash report and logs under report ID: %v", crashReport.ReportId)
	i.Report("crash_report", ScoutMeta{"crash_report_id", crashReport.ReportId})
	i.uploadCrashReportData(crashReport, i.gatherCrashReportData())
}

func (i *Installer) gatherCrashReportData() []byte {
	buffer := bytes.NewBuffer([]byte{})

	buffer.WriteString("========== edgectl logs ==========\n")
	fileContent, err := ioutil.ReadFile(i.logName)
	if err != nil {
		i.log.Printf("failed to read log file %v: %v", i.logName, err.Error())
	}
	buffer.Write(fileContent)

	buffer.WriteString("\n========== kubectl describe ==========\n")
	describe, err := i.SilentCaptureKubectl("describe ambassador namespace", "", "-n", "ambassador", "describe", "all")
	if err != nil {
		i.log.Printf("failed to describe ambassador resources: %v", err.Error())
	}
	buffer.WriteString(describe)

	buffer.WriteString("\n========== kubectl logs ==========\n")
	ambassadorLogs, err := i.SilentCaptureKubectl("read ambassador logs", "", "-n", "ambassador", "logs", "deployments/ambassador", "--tail=1000")
	if err != nil {
		i.log.Printf("failed to read ambassador logs: %v", err.Error())
	}
	buffer.WriteString(ambassadorLogs)

	return buffer.Bytes()
}

func (i *Installer) uploadCrashReportData(crashReport crashReportCreationResponse, uploadContent []byte) {
	client := &http.Client{}
	req, err := http.NewRequest(crashReport.Method, crashReport.UploadURL, bytes.NewReader(uploadContent))
	if err != nil {
		i.log.Print(err.Error())
		return
	}

	res, err := client.Do(req)
	if err != nil {
		i.log.Print(err.Error())
		return
	}
	defer res.Body.Close()
	_, err = ioutil.ReadAll(res.Body)
	if err != nil {
		i.log.Print(err.Error())
		return
	}
}

// crashReportCreationRequest is used to initiate a crash report request
type crashReportCreationRequest struct {
	Product         string
	ProductVersion  string
	Command         string
	Error           string
	AESVersion      string
	Address         string
	Hostname        string
	ClusterID       string
	InstallID       string
	TraceID         string
	ClusterInfo     string
	Managed         string
	KubectlVersion  string
	KubectlPlatform string
	K8sVersion      string
	K8sPlatform     string
}

// crashReportCreationResponse is used to receive a crash report creation response
type crashReportCreationResponse struct {
	ReportId  string
	Method    string
	UploadURL string
}
<|endoftext|>


<commit_before>package app

import (
	"errors"
	"fmt"
	"io"
	"log"
	"math/rand"
	"strconv"
	"time"

	"github.com/mesos/mesos-go"
	"github.com/mesos/mesos-go/backoff"
	xmetrics "github.com/mesos/mesos-go/extras/metrics"
	"github.com/mesos/mesos-go/extras/scheduler/controller"
	"github.com/mesos/mesos-go/httpcli/httpsched"
	"github.com/mesos/mesos-go/scheduler"
	"github.com/mesos/mesos-go/scheduler/calls"
	"github.com/mesos/mesos-go/scheduler/events"
)

func Run(cfg Config) error {
	log.Printf("scheduler running with configuration: %+v", cfg)

	state, err := newInternalState(cfg)
	if err != nil {
		return err
	}
	var (
		frameworkInfo  = buildFrameworkInfo(cfg)
		controlContext = &controllerContext{
			events: events.Decorators{
				eventMetrics(state.metricsAPI, time.Now, state.config.summaryMetrics),
				events.Decorator(logAllEvents).If(state.config.verbose),
			}.Apply(buildEventHandler(state)),

			callerChanged: httpsched.Decorators{
				callMetrics(state.metricsAPI, time.Now, state.config.summaryMetrics),
				logCalls(map[scheduler.Call_Type]string{scheduler.Call_SUBSCRIBE: "connecting..."}),
				// the next decorator must be last since it tracks the subscribed caller we'll use
				httpsched.CallerTracker(func(c httpsched.Caller) { state.cli = c }),
			}.Combine(),

			errHandler: func(err error) {
				if err != nil && err != io.EOF {
					log.Println(err)
				} else {
					log.Println("disconnected")
				}
			},

			state:              state,
			registrationTokens: backoff.Notifier(1*time.Second, 15*time.Second, nil),
		}
	)
	return controller.Run(controlContext, frameworkInfo, state.cli)
}

// buildEventHandler generates and returns a handler to process events received from the subscription.
func buildEventHandler(state *internalState) events.Handler {
	callOptions := scheduler.CallOptions{} // should be applied to every outgoing call
	return events.NewMux(
		events.Handle(scheduler.Event_FAILURE, events.HandlerFunc(func(e *scheduler.Event) error {
			log.Println("received a FAILURE event")
			f := e.GetFailure()
			failure(f.ExecutorID, f.AgentID, f.Status)
			return nil
		})),
		events.Handle(scheduler.Event_OFFERS, events.HandlerFunc(func(e *scheduler.Event) error {
			if state.config.verbose {
				log.Println("received an OFFERS event")
			}
			offers := e.GetOffers().GetOffers()
			state.metricsAPI.offersReceived.Int(len(offers))
			resourceOffers(state, callOptions[:], offers)
			return nil
		})),
		events.Handle(scheduler.Event_UPDATE, events.HandlerFunc(func(e *scheduler.Event) error {
			statusUpdate(state, callOptions[:], e.GetUpdate().GetStatus())
			return nil
		})),
		events.Handle(scheduler.Event_ERROR, events.HandlerFunc(func(e *scheduler.Event) error {
			// it's recommended that we abort and re-try subscribing; setting
			// err here will cause the event loop to terminate and the connection
			// will be reset.
			return fmt.Errorf("ERROR: " + e.GetError().GetMessage())
		})),
		events.Handle(scheduler.Event_SUBSCRIBED, events.HandlerFunc(func(e *scheduler.Event) (err error) {
			log.Println("received a SUBSCRIBED event")
			if state.frameworkID == "" {
				state.frameworkID = e.GetSubscribed().GetFrameworkID().GetValue()
				if state.frameworkID == "" {
					// sanity check
					err = errors.New("mesos gave us an empty frameworkID")
				} else {
					callOptions = append(callOptions, calls.Framework(state.frameworkID))
				}
			}
			return
		})),
	)
} // buildEventHandler

func failure(eid *mesos.ExecutorID, aid *mesos.AgentID, stat *int32) {
	if eid != nil {
		// executor failed..
		msg := "executor '" + eid.Value + "' terminated"
		if aid != nil {
			msg += "on agent '" + aid.Value + "'"
		}
		if stat != nil {
			msg += ", with status=" + strconv.Itoa(int(*stat))
		}
		log.Println(msg)
	} else if aid != nil {
		// agent failed..
		log.Println("agent '" + aid.Value + "' terminated")
	}
}

func refuseSecondsWithJitter(d time.Duration) scheduler.CallOpt {
	return calls.Filters(func(f *mesos.Filters) {
		s := time.Duration(rand.Int63n(int64(d))).Seconds()
		f.RefuseSeconds = &s
	})
}

func resourceOffers(state *internalState, callOptions scheduler.CallOptions, offers []mesos.Offer) {
	callOptions = append(callOptions, refuseSecondsWithJitter(state.config.maxRefuseSeconds))
	tasksLaunchedThisCycle := 0
	offersDeclined := 0
	for i := range offers {
		var (
			remaining = mesos.Resources(offers[i].Resources)
			tasks     = []mesos.TaskInfo{}
		)

		if state.config.verbose {
			log.Println("received offer id '" + offers[i].ID.Value + "' with resources " + remaining.String())
		}

		var wantsExecutorResources mesos.Resources
		if len(offers[i].ExecutorIDs) == 0 {
			wantsExecutorResources = mesos.Resources(state.executor.Resources)
		}

		flattened := remaining.Flatten()

		// avoid the expense of computing these if we can...
		if state.config.summaryMetrics && state.config.resourceTypeMetrics {
			for name, restype := range flattened.Types() {
				if restype == mesos.SCALAR {
					sum := flattened.SumScalars(mesos.NamedResources(name))
					state.metricsAPI.offeredResources(sum.GetValue(), name)
				}
			}
		}

		taskWantsResources := state.wantsTaskResources.Plus(wantsExecutorResources...)
		for state.tasksLaunched < state.totalTasks && flattened.ContainsAll(taskWantsResources) {
			state.tasksLaunched++
			taskID := state.tasksLaunched

			if state.config.verbose {
				log.Println("launching task " + strconv.Itoa(taskID) + " using offer " + offers[i].ID.Value)
			}

			task := mesos.TaskInfo{TaskID: mesos.TaskID{Value: strconv.Itoa(taskID)}}
			task.Name = "Task " + task.TaskID.Value
			task.AgentID = offers[i].AgentID
			task.Executor = state.executor
			task.Resources = remaining.Find(state.wantsTaskResources.Flatten(mesos.Role(state.role).Assign()))

			remaining.Subtract(task.Resources...)
			tasks = append(tasks, task)

			flattened = remaining.Flatten()
		}

		// build Accept call to launch all of the tasks we've assembled
		accept := calls.Accept(
			calls.OfferWithOperations(
				offers[i].ID,
				calls.OpLaunch(tasks...),
			),
		).With(callOptions...)

		// send Accept call to mesos
		err := httpsched.CallNoData(state.cli, accept)
		if err != nil {
			log.Printf("failed to launch tasks: %+v", err)
		} else {
			if n := len(tasks); n > 0 {
				tasksLaunchedThisCycle += n
			} else {
				offersDeclined++
			}
		}
	}
	state.metricsAPI.offersDeclined.Int(offersDeclined)
	state.metricsAPI.tasksLaunched.Int(tasksLaunchedThisCycle)
	if state.config.summaryMetrics {
		state.metricsAPI.launchesPerOfferCycle(float64(tasksLaunchedThisCycle))
	}
}

func statusUpdate(state *internalState, callOptions scheduler.CallOptions, s mesos.TaskStatus) {
	msg := "Task " + s.TaskID.Value + " is in state " + s.GetState().String()
	if m := s.GetMessage(); m != "" {
		msg += " with message '" + m + "'"
	}
	if state.config.verbose {
		log.Println(msg)
	}

	if uuid := s.GetUUID(); len(uuid) > 0 {
		ack := calls.Acknowledge(
			s.GetAgentID().GetValue(),
			s.TaskID.Value,
			uuid,
		).With(callOptions...)

		// send Ack call to mesos
		err := httpsched.CallNoData(state.cli, ack)
		if err != nil {
			log.Printf("failed to ack status update for task: %+v", err)
			return
		}
	}
	switch st := s.GetState(); st {
	case mesos.TASK_FINISHED:
		state.tasksFinished++
		state.metricsAPI.tasksFinished()

	case mesos.TASK_LOST, mesos.TASK_KILLED, mesos.TASK_FAILED, mesos.TASK_ERROR:
		state.err = errors.New("Exiting because task " + s.GetTaskID().Value +
			" is in an unexpected state " + st.String() +
			" with reason " + s.GetReason().String() +
			" from source " + s.GetSource().String() +
			" with message '" + s.GetMessage() + "'")
		state.done = true
		return
	}

	if state.tasksFinished == state.totalTasks {
		log.Println("mission accomplished, terminating")
		state.done = true
	} else {
		tryReviveOffers(state, callOptions)
	}
}

func tryReviveOffers(state *internalState, callOptions scheduler.CallOptions) {
	// limit the rate at which we request offer revival
	select {
	case <-state.reviveTokens:
		// not done yet, revive offers!
		err := httpsched.CallNoData(state.cli, calls.Revive().With(callOptions...))
		if err != nil {
			log.Printf("failed to revive offers: %+v", err)
			return
		}
	default:
		// noop
	}
}

// logAllEvents logs every observed event; this is somewhat expensive to do
func logAllEvents(h events.Handler) events.Handler {
	return events.HandlerFunc(func(e *scheduler.Event) error {
		log.Printf("%+v\n", *e)
		return h.HandleEvent(e)
	})
}

// eventMetrics logs metrics for every processed API event
func eventMetrics(metricsAPI *metricsAPI, clock func() time.Time, timingMetrics bool) events.Decorator {
	timed := metricsAPI.eventReceivedLatency
	if !timingMetrics {
		timed = nil
	}
	harness := xmetrics.NewHarness(metricsAPI.eventReceivedCount, metricsAPI.eventErrorCount, timed, clock)
	return events.Metrics(harness)
}

// callMetrics logs metrics for every outgoing Mesos call
func callMetrics(metricsAPI *metricsAPI, clock func() time.Time, timingMetrics bool) httpsched.Decorator {
	timed := metricsAPI.callLatency
	if !timingMetrics {
		timed = nil
	}
	harness := xmetrics.NewHarness(metricsAPI.callCount, metricsAPI.callErrorCount, timed, clock)
	return httpsched.CallerMetrics(harness)
}

func logCalls(messages map[scheduler.Call_Type]string) httpsched.Decorator {
	return func(caller httpsched.Caller) httpsched.Caller {
		return &httpsched.CallerAdapter{
			CallFunc: func(c *scheduler.Call) (mesos.Response, httpsched.Caller, error) {
				if message, ok := messages[c.GetType()]; ok {
					log.Println(message)
				}
				return caller.Call(c)
			},
		}
	}
}

type controllerContext struct {
	state              *internalState
	events             events.Handler
	registrationTokens <-chan struct{}
	callerChanged      httpsched.Decorator
	errHandler         func(error)
}

func (ci *controllerContext) Handler() events.Handler                    { return ci.events }
func (ci *controllerContext) Done() bool                                 { return ci.state.done }
func (ci *controllerContext) FrameworkID() string                        { return ci.state.frameworkID }
func (ci *controllerContext) RegistrationTokens() <-chan struct{}        { return ci.registrationTokens }
func (ci *controllerContext) Caller(c httpsched.Caller) httpsched.Caller { return ci.callerChanged(c) }
func (ci *controllerContext) Error(err error) {
	if ci.errHandler != nil {
		ci.errHandler(err)
	}
}
<commit_msg>refactor example-scheduler controller-context generation into factory func<commit_after>package app

import (
	"errors"
	"fmt"
	"io"
	"log"
	"math/rand"
	"strconv"
	"time"

	"github.com/mesos/mesos-go"
	"github.com/mesos/mesos-go/backoff"
	xmetrics "github.com/mesos/mesos-go/extras/metrics"
	"github.com/mesos/mesos-go/extras/scheduler/controller"
	"github.com/mesos/mesos-go/httpcli/httpsched"
	"github.com/mesos/mesos-go/scheduler"
	"github.com/mesos/mesos-go/scheduler/calls"
	"github.com/mesos/mesos-go/scheduler/events"
)

func Run(cfg Config) error {
	log.Printf("scheduler running with configuration: %+v", cfg)

	state, err := newInternalState(cfg)
	if err != nil {
		return err
	}
	var (
		frameworkInfo  = buildFrameworkInfo(cfg)
		controlContext = buildControllerContext(state, frameworkInfo)
	)
	return controller.Run(controlContext, frameworkInfo, state.cli)
}

func buildControllerContext(state *internalState, frameworkInfo *mesos.FrameworkInfo) controller.Context {
	return &controllerContext{
		events: events.Decorators{
			eventMetrics(state.metricsAPI, time.Now, state.config.summaryMetrics),
			events.Decorator(logAllEvents).If(state.config.verbose),
		}.Apply(buildEventHandler(state)),

		callerChanged: httpsched.Decorators{
			callMetrics(state.metricsAPI, time.Now, state.config.summaryMetrics),
			logCalls(map[scheduler.Call_Type]string{scheduler.Call_SUBSCRIBE: "connecting..."}),
			// the next decorator must be last since it tracks the subscribed caller we'll use
			httpsched.CallerTracker(func(c httpsched.Caller) { state.cli = c }),
		}.Combine(),

		errHandler: func(err error) {
			if err != nil && err != io.EOF {
				log.Println(err)
			} else {
				log.Println("disconnected")
			}
		},

		state:              state,
		registrationTokens: backoff.Notifier(1*time.Second, 15*time.Second, nil),
	}
}

// buildEventHandler generates and returns a handler to process events received from the subscription.
func buildEventHandler(state *internalState) events.Handler {
	callOptions := scheduler.CallOptions{} // should be applied to every outgoing call
	return events.NewMux(
		events.Handle(scheduler.Event_FAILURE, events.HandlerFunc(func(e *scheduler.Event) error {
			log.Println("received a FAILURE event")
			f := e.GetFailure()
			failure(f.ExecutorID, f.AgentID, f.Status)
			return nil
		})),
		events.Handle(scheduler.Event_OFFERS, events.HandlerFunc(func(e *scheduler.Event) error {
			if state.config.verbose {
				log.Println("received an OFFERS event")
			}
			offers := e.GetOffers().GetOffers()
			state.metricsAPI.offersReceived.Int(len(offers))
			resourceOffers(state, callOptions[:], offers)
			return nil
		})),
		events.Handle(scheduler.Event_UPDATE, events.HandlerFunc(func(e *scheduler.Event) error {
			statusUpdate(state, callOptions[:], e.GetUpdate().GetStatus())
			return nil
		})),
		events.Handle(scheduler.Event_ERROR, events.HandlerFunc(func(e *scheduler.Event) error {
			// it's recommended that we abort and re-try subscribing; setting
			// err here will cause the event loop to terminate and the connection
			// will be reset.
			return fmt.Errorf("ERROR: " + e.GetError().GetMessage())
		})),
		events.Handle(scheduler.Event_SUBSCRIBED, events.HandlerFunc(func(e *scheduler.Event) (err error) {
			log.Println("received a SUBSCRIBED event")
			if state.frameworkID == "" {
				state.frameworkID = e.GetSubscribed().GetFrameworkID().GetValue()
				if state.frameworkID == "" {
					// sanity check
					err = errors.New("mesos gave us an empty frameworkID")
				} else {
					callOptions = append(callOptions, calls.Framework(state.frameworkID))
				}
			}
			return
		})),
	)
} // buildEventHandler

func failure(eid *mesos.ExecutorID, aid *mesos.AgentID, stat *int32) {
	if eid != nil {
		// executor failed..
		msg := "executor '" + eid.Value + "' terminated"
		if aid != nil {
			msg += "on agent '" + aid.Value + "'"
		}
		if stat != nil {
			msg += ", with status=" + strconv.Itoa(int(*stat))
		}
		log.Println(msg)
	} else if aid != nil {
		// agent failed..
		log.Println("agent '" + aid.Value + "' terminated")
	}
}

func refuseSecondsWithJitter(d time.Duration) scheduler.CallOpt {
	return calls.Filters(func(f *mesos.Filters) {
		s := time.Duration(rand.Int63n(int64(d))).Seconds()
		f.RefuseSeconds = &s
	})
}

func resourceOffers(state *internalState, callOptions scheduler.CallOptions, offers []mesos.Offer) {
	callOptions = append(callOptions, refuseSecondsWithJitter(state.config.maxRefuseSeconds))
	tasksLaunchedThisCycle := 0
	offersDeclined := 0
	for i := range offers {
		var (
			remaining = mesos.Resources(offers[i].Resources)
			tasks     = []mesos.TaskInfo{}
		)

		if state.config.verbose {
			log.Println("received offer id '" + offers[i].ID.Value + "' with resources " + remaining.String())
		}

		var wantsExecutorResources mesos.Resources
		if len(offers[i].ExecutorIDs) == 0 {
			wantsExecutorResources = mesos.Resources(state.executor.Resources)
		}

		flattened := remaining.Flatten()

		// avoid the expense of computing these if we can...
		if state.config.summaryMetrics && state.config.resourceTypeMetrics {
			for name, restype := range flattened.Types() {
				if restype == mesos.SCALAR {
					sum := flattened.SumScalars(mesos.NamedResources(name))
					state.metricsAPI.offeredResources(sum.GetValue(), name)
				}
			}
		}

		taskWantsResources := state.wantsTaskResources.Plus(wantsExecutorResources...)
		for state.tasksLaunched < state.totalTasks && flattened.ContainsAll(taskWantsResources) {
			state.tasksLaunched++
			taskID := state.tasksLaunched

			if state.config.verbose {
				log.Println("launching task " + strconv.Itoa(taskID) + " using offer " + offers[i].ID.Value)
			}

			task := mesos.TaskInfo{TaskID: mesos.TaskID{Value: strconv.Itoa(taskID)}}
			task.Name = "Task " + task.TaskID.Value
			task.AgentID = offers[i].AgentID
			task.Executor = state.executor
			task.Resources = remaining.Find(state.wantsTaskResources.Flatten(mesos.Role(state.role).Assign()))

			remaining.Subtract(task.Resources...)
			tasks = append(tasks, task)

			flattened = remaining.Flatten()
		}

		// build Accept call to launch all of the tasks we've assembled
		accept := calls.Accept(
			calls.OfferWithOperations(
				offers[i].ID,
				calls.OpLaunch(tasks...),
			),
		).With(callOptions...)

		// send Accept call to mesos
		err := httpsched.CallNoData(state.cli, accept)
		if err != nil {
			log.Printf("failed to launch tasks: %+v", err)
		} else {
			if n := len(tasks); n > 0 {
				tasksLaunchedThisCycle += n
			} else {
				offersDeclined++
			}
		}
	}
	state.metricsAPI.offersDeclined.Int(offersDeclined)
	state.metricsAPI.tasksLaunched.Int(tasksLaunchedThisCycle)
	if state.config.summaryMetrics {
		state.metricsAPI.launchesPerOfferCycle(float64(tasksLaunchedThisCycle))
	}
}

func statusUpdate(state *internalState, callOptions scheduler.CallOptions, s mesos.TaskStatus) {
	msg := "Task " + s.TaskID.Value + " is in state " + s.GetState().String()
	if m := s.GetMessage(); m != "" {
		msg += " with message '" + m + "'"
	}
	if state.config.verbose {
		log.Println(msg)
	}

	if uuid := s.GetUUID(); len(uuid) > 0 {
		ack := calls.Acknowledge(
			s.GetAgentID().GetValue(),
			s.TaskID.Value,
			uuid,
		).With(callOptions...)

		// send Ack call to mesos
		err := httpsched.CallNoData(state.cli, ack)
		if err != nil {
			log.Printf("failed to ack status update for task: %+v", err)
			return
		}
	}
	switch st := s.GetState(); st {
	case mesos.TASK_FINISHED:
		state.tasksFinished++
		state.metricsAPI.tasksFinished()

	case mesos.TASK_LOST, mesos.TASK_KILLED, mesos.TASK_FAILED, mesos.TASK_ERROR:
		state.err = errors.New("Exiting because task " + s.GetTaskID().Value +
			" is in an unexpected state " + st.String() +
			" with reason " + s.GetReason().String() +
			" from source " + s.GetSource().String() +
			" with message '" + s.GetMessage() + "'")
		state.done = true
		return
	}

	if state.tasksFinished == state.totalTasks {
		log.Println("mission accomplished, terminating")
		state.done = true
	} else {
		tryReviveOffers(state, callOptions)
	}
}

func tryReviveOffers(state *internalState, callOptions scheduler.CallOptions) {
	// limit the rate at which we request offer revival
	select {
	case <-state.reviveTokens:
		// not done yet, revive offers!
		err := httpsched.CallNoData(state.cli, calls.Revive().With(callOptions...))
		if err != nil {
			log.Printf("failed to revive offers: %+v", err)
			return
		}
	default:
		// noop
	}
}

// logAllEvents logs every observed event; this is somewhat expensive to do
func logAllEvents(h events.Handler) events.Handler {
	return events.HandlerFunc(func(e *scheduler.Event) error {
		log.Printf("%+v\n", *e)
		return h.HandleEvent(e)
	})
}

// eventMetrics logs metrics for every processed API event
func eventMetrics(metricsAPI *metricsAPI, clock func() time.Time, timingMetrics bool) events.Decorator {
	timed := metricsAPI.eventReceivedLatency
	if !timingMetrics {
		timed = nil
	}
	harness := xmetrics.NewHarness(metricsAPI.eventReceivedCount, metricsAPI.eventErrorCount, timed, clock)
	return events.Metrics(harness)
}

// callMetrics logs metrics for every outgoing Mesos call
func callMetrics(metricsAPI *metricsAPI, clock func() time.Time, timingMetrics bool) httpsched.Decorator {
	timed := metricsAPI.callLatency
	if !timingMetrics {
		timed = nil
	}
	harness := xmetrics.NewHarness(metricsAPI.callCount, metricsAPI.callErrorCount, timed, clock)
	return httpsched.CallerMetrics(harness)
}

func logCalls(messages map[scheduler.Call_Type]string) httpsched.Decorator {
	return func(caller httpsched.Caller) httpsched.Caller {
		return &httpsched.CallerAdapter{
			CallFunc: func(c *scheduler.Call) (mesos.Response, httpsched.Caller, error) {
				if message, ok := messages[c.GetType()]; ok {
					log.Println(message)
				}
				return caller.Call(c)
			},
		}
	}
}

type controllerContext struct {
	state              *internalState
	events             events.Handler
	registrationTokens <-chan struct{}
	callerChanged      httpsched.Decorator
	errHandler         func(error)
}

func (ci *controllerContext) Handler() events.Handler                    { return ci.events }
func (ci *controllerContext) Done() bool                                 { return ci.state.done }
func (ci *controllerContext) FrameworkID() string                        { return ci.state.frameworkID }
func (ci *controllerContext) RegistrationTokens() <-chan struct{}        { return ci.registrationTokens }
func (ci *controllerContext) Caller(c httpsched.Caller) httpsched.Caller { return ci.callerChanged(c) }
func (ci *controllerContext) Error(err error) {
	if ci.errHandler != nil {
		ci.errHandler(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package service

import (
	"fmt"
	"os"
	"path/filepath"

	"github.com/juju/cmd"
	"github.com/juju/errors"
	"github.com/juju/names"
	"gopkg.in/juju/charm.v6-unstable"
	charmresource "gopkg.in/juju/charm.v6-unstable/resource"
	"gopkg.in/juju/charmrepo.v2-unstable"
	"gopkg.in/macaroon.v1"
	"launchpad.net/gnuflag"

	"github.com/juju/juju/api"
	apiservice "github.com/juju/juju/api/service"
	"github.com/juju/juju/cmd/juju/block"
	"github.com/juju/juju/cmd/modelcmd"
	"github.com/juju/juju/resource"
	"github.com/juju/juju/resource/resourceadapters"
)

// NewUpgradeCharmCommand returns a command which upgrades service's charm.
func NewUpgradeCharmCommand() cmd.Command {
	return modelcmd.Wrap(&upgradeCharmCommand{})
}

// UpgradeCharm is responsible for upgrading a service's charm.
type upgradeCharmCommand struct {
	modelcmd.ModelCommandBase
	ServiceName string
	ForceUnits  bool
	ForceSeries bool
	RepoPath    string // defaults to JUJU_REPOSITORY
	SwitchURL   string
	CharmPath   string
	Revision    int // defaults to -1 (latest)
	// Resources is a map of resource name to filename to be uploaded on upgrade.
	Resources map[string]string
}

const upgradeCharmDoc = `
When no flags are set, the service's charm will be upgraded to the latest
revision available in the repository from which it was originally deployed. An
explicit revision can be chosen with the --revision flag.

If the charm was not originally deployed from a repository, but from a path,
then a path will need to be supplied to allow an updated copy of the charm
to be located.

If the charm came from a local repository, its path will be assumed to be
$JUJU_REPOSITORY unless overridden by --repository. Note that deploying from
a local repository is deprecated in favour of deploying from a path.

Deploying from a path or local repository is intended to suit the workflow of a charm
author working on a single client machine; use of this deployment method from
multiple clients is not supported and may lead to confusing behaviour. Each
local charm gets uploaded with the revision specified in the charm, if possible,
otherwise it gets a unique revision (highest in state + 1).

When deploying from a path, the --path flag is used to specify the location from
which to load the updated charm. Note that the directory containing the charm must
match what was originally used to deploy the charm as a superficial check that the
updated charm is compatible.

Resources may be uploaded at upgrade time by specifying the --resource flag.
Following the resource flag should be name=filepath pair.  This flag may be
repeated more than once to upload more than one resource.

  juju upgrade-charm foo --resource bar=/some/file.tgz --resource baz=./docs/cfg.xml

Where bar and baz are resources named in the metadata for the foo charm.

If the new version of a charm does not explicitly support the service's series, the
upgrade is disallowed unless the --force-series flag is used. This option should be
used with caution since using a charm on a machine running an unsupported series may
cause unexpected behavior.

When using a local repository, the --switch flag allows you to replace the charm
with an entirely different one. The new charm's URL and revision are inferred as
they would be when running a deploy command.

Please note that --switch is dangerous, because juju only has limited
information with which to determine compatibility; the operation will succeed,
regardless of potential havoc, so long as the following conditions hold:

- The new charm must declare all relations that the service is currently
participating in.
- All config settings shared by the old and new charms must
have the same types.

The new charm may add new relations and configuration settings.

--switch and --path are mutually exclusive.

--path and --revision are mutually exclusive. The revision of the updated charm
is determined by the contents of the charm at the specified path.

--switch and --revision are mutually exclusive. To specify a given revision
number with --switch, give it in the charm URL, for instance "cs:wordpress-5"
would specify revision number 5 of the wordpress charm.

Use of the --force-units flag is not generally recommended; units upgraded while in an
error state will not have upgrade-charm hooks executed, and may cause unexpected
behavior.
`

func (c *upgradeCharmCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "upgrade-charm",
		Args:    "<service>",
		Purpose: "upgrade a service's charm",
		Doc:     upgradeCharmDoc,
	}
}

func (c *upgradeCharmCommand) SetFlags(f *gnuflag.FlagSet) {
	f.BoolVar(&c.ForceUnits, "force-units", false, "upgrade all units immediately, even if in error state")
	f.BoolVar(&c.ForceSeries, "force-series", false, "upgrade even if series of deployed services are not supported by the new charm")
	f.StringVar(&c.RepoPath, "repository", os.Getenv("JUJU_REPOSITORY"), "local charm repository path")
	f.StringVar(&c.SwitchURL, "switch", "", "crossgrade to a different charm")
	f.StringVar(&c.CharmPath, "path", "", "upgrade to a charm located at path")
	f.IntVar(&c.Revision, "revision", -1, "explicit revision of current charm")
	f.Var(stringMap{&c.Resources}, "resource", "resource to be uploaded to the controller")
}

func (c *upgradeCharmCommand) Init(args []string) error {
	switch len(args) {
	case 1:
		if !names.IsValidService(args[0]) {
			return fmt.Errorf("invalid service name %q", args[0])
		}
		c.ServiceName = args[0]
	case 0:
		return fmt.Errorf("no service specified")
	default:
		return cmd.CheckEmpty(args[1:])
	}
	if c.SwitchURL != "" && c.Revision != -1 {
		return fmt.Errorf("--switch and --revision are mutually exclusive")
	}
	if c.CharmPath != "" && c.Revision != -1 {
		return fmt.Errorf("--path and --revision are mutually exclusive")
	}
	if c.SwitchURL != "" && c.CharmPath != "" {
		return fmt.Errorf("--switch and --path are mutually exclusive")
	}
	return nil
}

func (c *upgradeCharmCommand) newServiceAPIClient() (*apiservice.Client, error) {
	root, err := c.NewAPIRoot()
	if err != nil {
		return nil, errors.Trace(err)
	}
	return apiservice.NewClient(root), nil
}

// Run connects to the specified environment and starts the charm
// upgrade process.
func (c *upgradeCharmCommand) Run(ctx *cmd.Context) error {
	client, err := c.NewAPIClient()
	if err != nil {
		return err
	}
	defer client.Close()

	serviceClient, err := c.newServiceAPIClient()
	if err != nil {
		return err
	}

	oldURL, err := serviceClient.GetCharmURL(c.ServiceName)
	if err != nil {
		return err
	}

	newRef := c.SwitchURL
	if newRef == "" {
		newRef = c.CharmPath
	}
	if c.SwitchURL == "" && c.CharmPath == "" {
		// No new URL specified, but revision might have been.
		newRef = oldURL.WithRevision(c.Revision).String()
	}

	httpClient, err := c.HTTPClient()
	if err != nil {
		return errors.Trace(err)
	}
	csClient := newCharmStoreClient(ctx, httpClient)

	addedURL, csMac, err := c.addCharm(oldURL, newRef, ctx, client, csClient)
	if err != nil {
		return block.ProcessBlockedError(err, block.BlockChange)
	}

	ids, err := c.upgradeResources(client, addedURL, csMac)
	if err != nil {
		return errors.Trace(err)
	}

	cfg := apiservice.SetCharmConfig{
		ServiceName: c.ServiceName,
		CharmUrl:    addedURL.String(),
		ForceSeries: c.ForceSeries,
		ForceUnits:  c.ForceUnits,
		ResourceIDs: ids,
	}

	return block.ProcessBlockedError(serviceClient.SetCharm(cfg), block.BlockChange)
}

// upgradeResources pushes metadata up to the server for each resource defined
// in the new charm's metadata and returns a map of resource names to pending
// IDs to include in the upgrage-charm call.
func (c *upgradeCharmCommand) upgradeResources(client *api.Client, cURL *charm.URL, csMac *macaroon.Macaroon) (map[string]string, error) {
	meta, err := getMetaResources(cURL, client)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if len(meta) == 0 {
		return nil, nil
	}

	current, err := getResources(c.ServiceName, c.NewAPIRoot)
	if err != nil {
		return nil, errors.Trace(err)
	}
	filtered := filterResources(meta, current, c.Resources)

	// Note: the validity of user-supplied resources to be uploaded will be
	// checked further down the stack.
	return handleResources(c, c.Resources, c.ServiceName, cURL, csMac, filtered)
}

// TODO(ericsnow) Move these helpers into handleResources()?

func getMetaResources(cURL *charm.URL, client *api.Client) (map[string]charmresource.Meta, error) {
	// this gets the charm info that was added to the controller using addcharm.
	charmInfo, err := client.CharmInfo(cURL.String())
	if err != nil {
		return nil, errors.Trace(err)
	}
	return charmInfo.Meta.Resources, nil
}

func getResources(serviceID string, newAPIRoot func() (api.Connection, error)) (map[string]resource.Resource, error) {
	resclient, err := resourceadapters.NewAPIClient(newAPIRoot)
	if err != nil {
		return nil, errors.Trace(err)
	}
	svcs, err := resclient.ListResources([]string{serviceID})
	if err != nil {
		return nil, errors.Trace(err)
	}
	// ListResources guarantees a number of values returned == number of
	// services passed in.
	return resource.AsMap(svcs[0].Resources), nil
}

// TODO(ericsnow) Move filterResources() and shouldUploadMeta()
// somewhere more general under the "resource" package?

func filterResources(meta map[string]charmresource.Meta, current map[string]resource.Resource, uploads map[string]string) map[string]charmresource.Meta {
	filtered := make(map[string]charmresource.Meta)
	for name, res := range meta {
		if shouldUpgradeResource(res, uploads, current) {
			filtered[name] = res
		}
	}
	return filtered
}

// shouldUpgradeResource reports whether we should upload the metadata for the given
// resource.  This is always true for resources we're adding with the --resource
// flag. For resources we're not adding with --resource, we only upload metadata
// for charmstore resources.  Previously uploaded resources stay pinned to the
// data the user uploaded.
func shouldUpgradeResource(res charmresource.Meta, uploads map[string]string, current map[string]resource.Resource) bool {
	// Always upload metadata for resources the user is uploading during
	// upgrade-charm.
	if _, ok := uploads[res.Name]; ok {
		return true
	}
	cur, ok := current[res.Name]
	if !ok {
		// If there's no information on the server, there should be.
		return true
	}
	// Never override existing resources a user has already uploaded.
	if cur.Origin == charmresource.OriginUpload {
		return false
	}
	return true
}

// addCharm interprets the new charmRef and adds the specified charm if the new charm is different
// to what's already deployed as specified by oldURL.
func (c *upgradeCharmCommand) addCharm(oldURL *charm.URL, charmRef string, ctx *cmd.Context,
	client *api.Client, csClient *csClient,
) (*charm.URL, *macaroon.Macaroon, error) {
	// Charm may have been supplied via a path reference.
	ch, newURL, err := charmrepo.NewCharmAtPathForceSeries(charmRef, oldURL.Series, c.ForceSeries)
	if err == nil {
		_, newName := filepath.Split(charmRef)
		if newName != oldURL.Name {
			return nil, nil, fmt.Errorf("cannot upgrade %q to %q", oldURL.Name, newName)
		}
		addedURL, err := client.AddLocalCharm(newURL, ch)
		return addedURL, nil, err
	}
	if _, ok := err.(*charmrepo.NotFoundError); ok {
		return nil, nil, errors.Errorf("no charm found at %q", charmRef)
	}
	// If we get a "not exists" or invalid path error then we attempt to interpret
	// the supplied charm reference as a URL below, otherwise we return the error.
	if err != os.ErrNotExist && !charmrepo.IsInvalidPathError(err) {
		return nil, nil, err
	}

	// Charm has been supplied as a URL so we resolve and deploy using the store.
	conf, err := getClientConfig(client)
	if err != nil {
		return nil, nil, err
	}

	newURL, supportedSeries, repo, err := resolveCharmStoreEntityURL(resolveCharmStoreEntityParams{
		urlStr:          charmRef,
		requestedSeries: oldURL.Series,
		forceSeries:     c.ForceSeries,
		csParams:        csClient.params,
		repoPath:        ctx.AbsPath(c.RepoPath),
		conf:            conf,
	})
	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	if !c.ForceSeries && oldURL.Series != "" && newURL.Series == "" && !isSeriesSupported(oldURL.Series, supportedSeries) {
		series := []string{"no series"}
		if len(supportedSeries) > 0 {
			series = supportedSeries
		}
		return nil, nil, errors.Errorf(
			"cannot upgrade from single series %q charm to a charm supporting %q. Use --force-series to override.",
			oldURL.Series, series,
		)
	}
	// If no explicit revision was set with either SwitchURL
	// or Revision flags, discover the latest.
	if *newURL == *oldURL {
		newRef, _ := charm.ParseURL(charmRef)
		if newRef.Revision != -1 {
			return nil, nil, fmt.Errorf("already running specified charm %q", newURL)
		}
		if newURL.Schema == "cs" {
			// No point in trying to upgrade a charm store charm when
			// we just determined that's the latest revision
			// available.
			return nil, nil, fmt.Errorf("already running latest charm %q", newURL)
		}
	}

	addedURL, csMac, err := addCharmFromURL(client, newURL, repo, csClient)
	if err != nil {
		return nil, nil, err
	}
	ctx.Infof("Added charm %q to the model.", addedURL)
	return addedURL, csMac, nil
}
<commit_msg>Factor out getUpgradeResources().<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package service

import (
	"fmt"
	"os"
	"path/filepath"

	"github.com/juju/cmd"
	"github.com/juju/errors"
	"github.com/juju/names"
	"gopkg.in/juju/charm.v6-unstable"
	charmresource "gopkg.in/juju/charm.v6-unstable/resource"
	"gopkg.in/juju/charmrepo.v2-unstable"
	"gopkg.in/macaroon.v1"
	"launchpad.net/gnuflag"

	"github.com/juju/juju/api"
	apiservice "github.com/juju/juju/api/service"
	"github.com/juju/juju/cmd/juju/block"
	"github.com/juju/juju/cmd/modelcmd"
	"github.com/juju/juju/resource"
	"github.com/juju/juju/resource/resourceadapters"
)

// NewUpgradeCharmCommand returns a command which upgrades service's charm.
func NewUpgradeCharmCommand() cmd.Command {
	return modelcmd.Wrap(&upgradeCharmCommand{})
}

// UpgradeCharm is responsible for upgrading a service's charm.
type upgradeCharmCommand struct {
	modelcmd.ModelCommandBase
	ServiceName string
	ForceUnits  bool
	ForceSeries bool
	RepoPath    string // defaults to JUJU_REPOSITORY
	SwitchURL   string
	CharmPath   string
	Revision    int // defaults to -1 (latest)
	// Resources is a map of resource name to filename to be uploaded on upgrade.
	Resources map[string]string
}

const upgradeCharmDoc = `
When no flags are set, the service's charm will be upgraded to the latest
revision available in the repository from which it was originally deployed. An
explicit revision can be chosen with the --revision flag.

If the charm was not originally deployed from a repository, but from a path,
then a path will need to be supplied to allow an updated copy of the charm
to be located.

If the charm came from a local repository, its path will be assumed to be
$JUJU_REPOSITORY unless overridden by --repository. Note that deploying from
a local repository is deprecated in favour of deploying from a path.

Deploying from a path or local repository is intended to suit the workflow of a charm
author working on a single client machine; use of this deployment method from
multiple clients is not supported and may lead to confusing behaviour. Each
local charm gets uploaded with the revision specified in the charm, if possible,
otherwise it gets a unique revision (highest in state + 1).

When deploying from a path, the --path flag is used to specify the location from
which to load the updated charm. Note that the directory containing the charm must
match what was originally used to deploy the charm as a superficial check that the
updated charm is compatible.

Resources may be uploaded at upgrade time by specifying the --resource flag.
Following the resource flag should be name=filepath pair.  This flag may be
repeated more than once to upload more than one resource.

  juju upgrade-charm foo --resource bar=/some/file.tgz --resource baz=./docs/cfg.xml

Where bar and baz are resources named in the metadata for the foo charm.

If the new version of a charm does not explicitly support the service's series, the
upgrade is disallowed unless the --force-series flag is used. This option should be
used with caution since using a charm on a machine running an unsupported series may
cause unexpected behavior.

When using a local repository, the --switch flag allows you to replace the charm
with an entirely different one. The new charm's URL and revision are inferred as
they would be when running a deploy command.

Please note that --switch is dangerous, because juju only has limited
information with which to determine compatibility; the operation will succeed,
regardless of potential havoc, so long as the following conditions hold:

- The new charm must declare all relations that the service is currently
participating in.
- All config settings shared by the old and new charms must
have the same types.

The new charm may add new relations and configuration settings.

--switch and --path are mutually exclusive.

--path and --revision are mutually exclusive. The revision of the updated charm
is determined by the contents of the charm at the specified path.

--switch and --revision are mutually exclusive. To specify a given revision
number with --switch, give it in the charm URL, for instance "cs:wordpress-5"
would specify revision number 5 of the wordpress charm.

Use of the --force-units flag is not generally recommended; units upgraded while in an
error state will not have upgrade-charm hooks executed, and may cause unexpected
behavior.
`

func (c *upgradeCharmCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "upgrade-charm",
		Args:    "<service>",
		Purpose: "upgrade a service's charm",
		Doc:     upgradeCharmDoc,
	}
}

func (c *upgradeCharmCommand) SetFlags(f *gnuflag.FlagSet) {
	f.BoolVar(&c.ForceUnits, "force-units", false, "upgrade all units immediately, even if in error state")
	f.BoolVar(&c.ForceSeries, "force-series", false, "upgrade even if series of deployed services are not supported by the new charm")
	f.StringVar(&c.RepoPath, "repository", os.Getenv("JUJU_REPOSITORY"), "local charm repository path")
	f.StringVar(&c.SwitchURL, "switch", "", "crossgrade to a different charm")
	f.StringVar(&c.CharmPath, "path", "", "upgrade to a charm located at path")
	f.IntVar(&c.Revision, "revision", -1, "explicit revision of current charm")
	f.Var(stringMap{&c.Resources}, "resource", "resource to be uploaded to the controller")
}

func (c *upgradeCharmCommand) Init(args []string) error {
	switch len(args) {
	case 1:
		if !names.IsValidService(args[0]) {
			return fmt.Errorf("invalid service name %q", args[0])
		}
		c.ServiceName = args[0]
	case 0:
		return fmt.Errorf("no service specified")
	default:
		return cmd.CheckEmpty(args[1:])
	}
	if c.SwitchURL != "" && c.Revision != -1 {
		return fmt.Errorf("--switch and --revision are mutually exclusive")
	}
	if c.CharmPath != "" && c.Revision != -1 {
		return fmt.Errorf("--path and --revision are mutually exclusive")
	}
	if c.SwitchURL != "" && c.CharmPath != "" {
		return fmt.Errorf("--switch and --path are mutually exclusive")
	}
	return nil
}

func (c *upgradeCharmCommand) newServiceAPIClient() (*apiservice.Client, error) {
	root, err := c.NewAPIRoot()
	if err != nil {
		return nil, errors.Trace(err)
	}
	return apiservice.NewClient(root), nil
}

// Run connects to the specified environment and starts the charm
// upgrade process.
func (c *upgradeCharmCommand) Run(ctx *cmd.Context) error {
	client, err := c.NewAPIClient()
	if err != nil {
		return err
	}
	defer client.Close()

	serviceClient, err := c.newServiceAPIClient()
	if err != nil {
		return err
	}

	oldURL, err := serviceClient.GetCharmURL(c.ServiceName)
	if err != nil {
		return err
	}

	newRef := c.SwitchURL
	if newRef == "" {
		newRef = c.CharmPath
	}
	if c.SwitchURL == "" && c.CharmPath == "" {
		// No new URL specified, but revision might have been.
		newRef = oldURL.WithRevision(c.Revision).String()
	}

	httpClient, err := c.HTTPClient()
	if err != nil {
		return errors.Trace(err)
	}
	csClient := newCharmStoreClient(ctx, httpClient)

	addedURL, csMac, err := c.addCharm(oldURL, newRef, ctx, client, csClient)
	if err != nil {
		return block.ProcessBlockedError(err, block.BlockChange)
	}

	ids, err := c.upgradeResources(client, addedURL, csMac)
	if err != nil {
		return errors.Trace(err)
	}

	cfg := apiservice.SetCharmConfig{
		ServiceName: c.ServiceName,
		CharmUrl:    addedURL.String(),
		ForceSeries: c.ForceSeries,
		ForceUnits:  c.ForceUnits,
		ResourceIDs: ids,
	}

	return block.ProcessBlockedError(serviceClient.SetCharm(cfg), block.BlockChange)
}

// upgradeResources pushes metadata up to the server for each resource defined
// in the new charm's metadata and returns a map of resource names to pending
// IDs to include in the upgrage-charm call.
func (c *upgradeCharmCommand) upgradeResources(client *api.Client, cURL *charm.URL, csMac *macaroon.Macaroon) (map[string]string, error) {
	filtered, err := getUpgradeResources(c, c.ServiceName, cURL, client, c.Resources)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if len(filtered) == 0 {
		return nil, nil
	}

	// Note: the validity of user-supplied resources to be uploaded will be
	// checked further down the stack.
	return handleResources(c, c.Resources, c.ServiceName, cURL, csMac, filtered)
}

// TODO(ericsnow) Move these helpers into handleResources()?

func getUpgradeResources(c APICmd, serviceID string, cURL *charm.URL, client *api.Client, cliResources map[string]string) (map[string]charmresource.Meta, error) {
	meta, err := getMetaResources(cURL, client)
	if err != nil {
		return nil, errors.Trace(err)
	}
	if len(meta) == 0 {
		return nil, nil
	}

	current, err := getResources(serviceID, c.NewAPIRoot)
	if err != nil {
		return nil, errors.Trace(err)
	}
	filtered := filterResources(meta, current, cliResources)
	return filtered, nil
}

func getMetaResources(cURL *charm.URL, client *api.Client) (map[string]charmresource.Meta, error) {
	// this gets the charm info that was added to the controller using addcharm.
	charmInfo, err := client.CharmInfo(cURL.String())
	if err != nil {
		return nil, errors.Trace(err)
	}
	return charmInfo.Meta.Resources, nil
}

func getResources(serviceID string, newAPIRoot func() (api.Connection, error)) (map[string]resource.Resource, error) {
	resclient, err := resourceadapters.NewAPIClient(newAPIRoot)
	if err != nil {
		return nil, errors.Trace(err)
	}
	svcs, err := resclient.ListResources([]string{serviceID})
	if err != nil {
		return nil, errors.Trace(err)
	}
	// ListResources guarantees a number of values returned == number of
	// services passed in.
	return resource.AsMap(svcs[0].Resources), nil
}

// TODO(ericsnow) Move filterResources() and shouldUploadMeta()
// somewhere more general under the "resource" package?

func filterResources(meta map[string]charmresource.Meta, current map[string]resource.Resource, uploads map[string]string) map[string]charmresource.Meta {
	filtered := make(map[string]charmresource.Meta)
	for name, res := range meta {
		if shouldUpgradeResource(res, uploads, current) {
			filtered[name] = res
		}
	}
	return filtered
}

// shouldUpgradeResource reports whether we should upload the metadata for the given
// resource.  This is always true for resources we're adding with the --resource
// flag. For resources we're not adding with --resource, we only upload metadata
// for charmstore resources.  Previously uploaded resources stay pinned to the
// data the user uploaded.
func shouldUpgradeResource(res charmresource.Meta, uploads map[string]string, current map[string]resource.Resource) bool {
	// Always upload metadata for resources the user is uploading during
	// upgrade-charm.
	if _, ok := uploads[res.Name]; ok {
		return true
	}
	cur, ok := current[res.Name]
	if !ok {
		// If there's no information on the server, there should be.
		return true
	}
	// Never override existing resources a user has already uploaded.
	if cur.Origin == charmresource.OriginUpload {
		return false
	}
	return true
}

// addCharm interprets the new charmRef and adds the specified charm if the new charm is different
// to what's already deployed as specified by oldURL.
func (c *upgradeCharmCommand) addCharm(oldURL *charm.URL, charmRef string, ctx *cmd.Context,
	client *api.Client, csClient *csClient,
) (*charm.URL, *macaroon.Macaroon, error) {
	// Charm may have been supplied via a path reference.
	ch, newURL, err := charmrepo.NewCharmAtPathForceSeries(charmRef, oldURL.Series, c.ForceSeries)
	if err == nil {
		_, newName := filepath.Split(charmRef)
		if newName != oldURL.Name {
			return nil, nil, fmt.Errorf("cannot upgrade %q to %q", oldURL.Name, newName)
		}
		addedURL, err := client.AddLocalCharm(newURL, ch)
		return addedURL, nil, err
	}
	if _, ok := err.(*charmrepo.NotFoundError); ok {
		return nil, nil, errors.Errorf("no charm found at %q", charmRef)
	}
	// If we get a "not exists" or invalid path error then we attempt to interpret
	// the supplied charm reference as a URL below, otherwise we return the error.
	if err != os.ErrNotExist && !charmrepo.IsInvalidPathError(err) {
		return nil, nil, err
	}

	// Charm has been supplied as a URL so we resolve and deploy using the store.
	conf, err := getClientConfig(client)
	if err != nil {
		return nil, nil, err
	}

	newURL, supportedSeries, repo, err := resolveCharmStoreEntityURL(resolveCharmStoreEntityParams{
		urlStr:          charmRef,
		requestedSeries: oldURL.Series,
		forceSeries:     c.ForceSeries,
		csParams:        csClient.params,
		repoPath:        ctx.AbsPath(c.RepoPath),
		conf:            conf,
	})
	if err != nil {
		return nil, nil, errors.Trace(err)
	}
	if !c.ForceSeries && oldURL.Series != "" && newURL.Series == "" && !isSeriesSupported(oldURL.Series, supportedSeries) {
		series := []string{"no series"}
		if len(supportedSeries) > 0 {
			series = supportedSeries
		}
		return nil, nil, errors.Errorf(
			"cannot upgrade from single series %q charm to a charm supporting %q. Use --force-series to override.",
			oldURL.Series, series,
		)
	}
	// If no explicit revision was set with either SwitchURL
	// or Revision flags, discover the latest.
	if *newURL == *oldURL {
		newRef, _ := charm.ParseURL(charmRef)
		if newRef.Revision != -1 {
			return nil, nil, fmt.Errorf("already running specified charm %q", newURL)
		}
		if newURL.Schema == "cs" {
			// No point in trying to upgrade a charm store charm when
			// we just determined that's the latest revision
			// available.
			return nil, nil, fmt.Errorf("already running latest charm %q", newURL)
		}
	}

	addedURL, csMac, err := addCharmFromURL(client, newURL, repo, csClient)
	if err != nil {
		return nil, nil, err
	}
	ctx.Infof("Added charm %q to the model.", addedURL)
	return addedURL, csMac, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	minhashlsh "github.com/ekzhu/minhash-lsh"
)

var (
	setFilename    string
	minhashSeed    int64
	minhashSize    int
	threshold      float64
	outputSelfPair bool
	hasID          bool
)

func main() {
	flag.StringVar(&setFilename, "input", "", "The set file as input")
	flag.Int64Var(&minhashSeed, "seed", 42, "The Minhash seed")
	flag.IntVar(&minhashSize, "sigsize", 128,
		"The Minhash signature size in number of hash functions")
	flag.Float64Var(&threshold, "threshold", 0.9, "The Jaccard similarity threshold")
	flag.BoolVar(&outputSelfPair, "selfpair", false, "Allow self-pair in results")
	flag.BoolVar(&hasID, "hasIDfield", true, "The input set file has ID field in the beginning of each line")
	flag.Parse()

	// Create Minhash signatures
	start := time.Now()
	sets := readSets(setFilename, hasID)
	setSigs := make([]setSig, 0)
	for setSig := range createSigantures(sets) {
		setSigs = append(setSigs, setSig)
	}
	signatureCreationTime := time.Now().Sub(start)

	// Indexing
	start = time.Now()
	lsh := minhashlsh.NewMinhashLSH(minhashSize, threshold)
	for _, s := range setSigs {
		lsh.Add(s.ID, s.signature)
	}
	lsh.Index()
	indexingTime := time.Now().Sub(start)

	// Querying and output results
	start = time.Now()
	pairs := make(chan pair)
	go func() {
		defer close(pairs)
		for _, s := range setSigs {
			for _, candidateID := range lsh.Query(s.signature) {
				if !outputSelfPair && candidateID == s.ID {
					continue
				}
				pairs <- pair{s.ID, candidateID}
			}
		}
	}()
	w := bufio.NewWriter(os.Stdout)
	for pair := range pairs {
		w.WriteString(pair.String() + "\n")
	}
	if err := w.Flush(); err != nil {
		panic(err)
	}
	searchTime := time.Now().Sub(start)

	fmt.Printf("Creating Minhash signature time: %.2f seconds\n", signatureCreationTime.Seconds())
	fmt.Printf("Indexing time: %.2f seconds\n", indexingTime.Seconds())
	fmt.Printf("All pair search time: %.2f seconds\n", searchTime.Seconds())
}

func pointquery() {
	panic("Not implemented")
}

type valueCountPair struct {
	value string
	count int
}

var valueCountRegex = regexp.MustCompile(`^(?P<value>.*)____(?P<count>[0-9]+)$`)

func (p *valueCountPair) Parse(str string) error {
	indexes := valueCountRegex.FindStringSubmatchIndex(str)
	if indexes == nil || len(indexes) != 6 {
		return errors.New("Incorrect value count pair detected: " + str)
	}
	p.value = str[indexes[2]:indexes[3]]
	var err error
	p.count, err = strconv.Atoi(str[indexes[4]:indexes[5]])
	if err != nil {
		panic(str + "\n" + err.Error())
	}
	return nil
}

type set struct {
	ID     string
	values []string
}

// readSets takes a set file having the following format:
// 1. One set per line
// 2. Each set, all items are separated by whitespaces
// 3. If the parameter firstItemIsID is set to true,
//    the first itme is the unique ID of the set.
// 4. The rest of the items with the following format:
//    <value>____<frequency>
//    * value is an unique element of the set
//    * frequency is an integer count of the occurance of value
//    * ____ (4 underscores) is the separator
func readSets(setFilename string, firstItemIsID bool) <-chan set {
	sets := make(chan set)
	go func() {
		defer close(sets)
		file, err := os.Open(setFilename)
		if err != nil {
			panic(err)
		}
		defer file.Close()
		scanner := bufio.NewScanner(file)
		scanner.Buffer(nil, 4096*1024*1024*8)
		var count int
		for scanner.Scan() {
			items := strings.Split(scanner.Text(), " ")
			var ID string
			if firstItemIsID {
				ID = items[0]
				items = items[1:]
			} else {
				ID = strconv.Itoa(count)
			}
			values := make([]string, len(items))
			for i, item := range items {
				var pair valueCountPair
				if err := pair.Parse(item); err != nil {
					fmt.Println(items)
					panic(err)
				}
				values[i] = pair.value
			}
			sets <- set{ID, values}
			count++
		}
		if err := scanner.Err(); err != nil {
			panic(err)
		}
	}()
	return sets
}

type setSig struct {
	ID        string
	size      int
	signature []uint64
}

func createSigantures(sets <-chan set) <-chan setSig {
	out := make(chan setSig)
	go func() {
		defer close(out)
		for set := range sets {
			mh := minhashlsh.NewMinhash(minhashSeed, minhashSize)
			for _, v := range set.values {
				mh.Push([]byte(v))
			}
			out <- setSig{set.ID, len(set.values), mh.Signature()}
		}
	}()
	return out
}

type pair struct {
	ID1 string
	ID2 string
}

func (p *pair) String() string {
	if p.ID1 <= p.ID2 {
		return fmt.Sprintf("%s, %s", p.ID1, p.ID2)
	}
	return fmt.Sprintf("%s, %s", p.ID2, p.ID1)
}
<commit_msg>output timing to stdout<commit_after>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	minhashlsh "github.com/ekzhu/minhash-lsh"
)

var (
	setFilename    string
	minhashSeed    int64
	minhashSize    int
	threshold      float64
	outputSelfPair bool
	hasID          bool
)

func main() {
	flag.StringVar(&setFilename, "input", "", "The set file as input")
	flag.Int64Var(&minhashSeed, "seed", 42, "The Minhash seed")
	flag.IntVar(&minhashSize, "sigsize", 128,
		"The Minhash signature size in number of hash functions")
	flag.Float64Var(&threshold, "threshold", 0.9, "The Jaccard similarity threshold")
	flag.BoolVar(&outputSelfPair, "selfpair", false, "Allow self-pair in results")
	flag.BoolVar(&hasID, "hasIDfield", true, "The input set file has ID field in the beginning of each line")
	flag.Parse()

	// Create Minhash signatures
	start := time.Now()
	sets := readSets(setFilename, hasID)
	setSigs := make([]setSig, 0)
	for setSig := range createSigantures(sets) {
		setSigs = append(setSigs, setSig)
	}
	signatureCreationTime := time.Now().Sub(start)
	fmt.Fprintf(os.Stderr, "Creating Minhash signature time: %.2f seconds\n", signatureCreationTime.Seconds())

	// Indexing
	start = time.Now()
	lsh := minhashlsh.NewMinhashLSH(minhashSize, threshold)
	for _, s := range setSigs {
		lsh.Add(s.ID, s.signature)
	}
	lsh.Index()
	indexingTime := time.Now().Sub(start)
	fmt.Fprintf(os.Stderr, "Indexing time: %.2f seconds\n", indexingTime.Seconds())

	// Querying and output results
	start = time.Now()
	pairs := make(chan pair)
	go func() {
		defer close(pairs)
		for _, s := range setSigs {
			for _, candidateID := range lsh.Query(s.signature) {
				if !outputSelfPair && candidateID == s.ID {
					continue
				}
				pairs <- pair{s.ID, candidateID}
			}
		}
	}()
	w := bufio.NewWriter(os.Stdout)
	for pair := range pairs {
		w.WriteString(pair.String() + "\n")
	}
	if err := w.Flush(); err != nil {
		panic(err)
	}
	searchTime := time.Now().Sub(start)
	fmt.Fprintf(os.Stderr, "All pair search time: %.2f seconds\n", searchTime.Seconds())
}

func pointquery() {
	panic("Not implemented")
}

type valueCountPair struct {
	value string
	count int
}

var valueCountRegex = regexp.MustCompile(`^(?P<value>.*)____(?P<count>[0-9]+)$`)

func (p *valueCountPair) Parse(str string) error {
	indexes := valueCountRegex.FindStringSubmatchIndex(str)
	if indexes == nil || len(indexes) != 6 {
		return errors.New("Incorrect value count pair detected: " + str)
	}
	p.value = str[indexes[2]:indexes[3]]
	var err error
	p.count, err = strconv.Atoi(str[indexes[4]:indexes[5]])
	if err != nil {
		panic(str + "\n" + err.Error())
	}
	return nil
}

type set struct {
	ID     string
	values []string
}

// readSets takes a set file having the following format:
// 1. One set per line
// 2. Each set, all items are separated by whitespaces
// 3. If the parameter firstItemIsID is set to true,
//    the first itme is the unique ID of the set.
// 4. The rest of the items with the following format:
//    <value>____<frequency>
//    * value is an unique element of the set
//    * frequency is an integer count of the occurance of value
//    * ____ (4 underscores) is the separator
func readSets(setFilename string, firstItemIsID bool) <-chan set {
	sets := make(chan set)
	go func() {
		defer close(sets)
		file, err := os.Open(setFilename)
		if err != nil {
			panic(err)
		}
		defer file.Close()
		scanner := bufio.NewScanner(file)
		scanner.Buffer(nil, 4096*1024*1024*8)
		var count int
		for scanner.Scan() {
			items := strings.Split(scanner.Text(), " ")
			var ID string
			if firstItemIsID {
				ID = items[0]
				items = items[1:]
			} else {
				ID = strconv.Itoa(count)
			}
			values := make([]string, len(items))
			for i, item := range items {
				var pair valueCountPair
				if err := pair.Parse(item); err != nil {
					fmt.Println(items)
					panic(err)
				}
				values[i] = pair.value
			}
			sets <- set{ID, values}
			count++
		}
		if err := scanner.Err(); err != nil {
			panic(err)
		}
	}()
	return sets
}

type setSig struct {
	ID        string
	size      int
	signature []uint64
}

func createSigantures(sets <-chan set) <-chan setSig {
	out := make(chan setSig)
	go func() {
		defer close(out)
		for set := range sets {
			mh := minhashlsh.NewMinhash(minhashSeed, minhashSize)
			for _, v := range set.values {
				mh.Push([]byte(v))
			}
			out <- setSig{set.ID, len(set.values), mh.Signature()}
		}
	}()
	return out
}

type pair struct {
	ID1 string
	ID2 string
}

func (p *pair) String() string {
	if p.ID1 <= p.ID2 {
		return fmt.Sprintf("%s, %s", p.ID1, p.ID2)
	}
	return fmt.Sprintf("%s, %s", p.ID2, p.ID1)
}
<|endoftext|>


<commit_before>// Copyright 2019 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package spantest implements creation/destruction of a temporary Spanner
// database.
package spantest

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"io/ioutil"
	"os/user"
	"regexp"
	"strings"

	"cloud.google.com/go/spanner"
	spandb "cloud.google.com/go/spanner/admin/database/apiv1"
	"golang.org/x/oauth2"
	"google.golang.org/api/option"
	dbpb "google.golang.org/genproto/googleapis/spanner/admin/database/v1"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/hardcoded/chromeinfra"
)

// TempDBConfig specifies how to create a temporary database.
type TempDBConfig struct {
	// InstanceName is the name of Spannner instance where to create the
	// temporary database.
	// Format: projects/{project}/instances/{instance}.
	// Defaults to chromeinfra.TestSpannerInstance.
	InstanceName string

	// TokenSource will be used to authenticate to Spanner.
	// If nil, auth.Authenticator with SilentLogin and chrome-infra auth options
	// will be used.
	// This means that that the user may have to login with luci-auth tool.
	TokenSource oauth2.TokenSource

	// InitScriptPath is a path to a DDL script to initialize the database.
	//
	// In lieu of a proper DDL parser, it is parsed using regexes.
	// Therefore the script MUST:
	//   - Use `#`` and/or `--`` for comments. No block comments.
	//   - Separate DDL statements with `;\n`.
	//
	// If empty, the database is created with no tables.
	InitScriptPath string
}

func (cfg *TempDBConfig) tokenSource(ctx context.Context) (oauth2.TokenSource, error) {
	if cfg.TokenSource != nil {
		return cfg.TokenSource, nil
	}

	opts := chromeinfra.DefaultAuthOptions()
	opts.Scopes = spandb.DefaultAuthScopes()
	a := auth.NewAuthenticator(ctx, auth.SilentLogin, opts)
	if err := a.CheckLoginRequired(); err != nil {
		return nil, errors.Annotate(err, "please login with `luci-auth login -scopes %q`", strings.Join(opts.Scopes, " ")).Err()
	}
	return a.TokenSource()
}

var ddlStatementSepRe = regexp.MustCompile(`;\s*\n`)
var commentRe = regexp.MustCompile(`(--|#)[^\n]*`)

// readDDLStatements read the file at cfg.InitScriptPath as a sequence of DDL
// statements. If the path is empty, returns (nil, nil).
func (cfg *TempDBConfig) readDDLStatements() ([]string, error) {
	if cfg.InitScriptPath == "" {
		return nil, nil
	}

	contents, err := ioutil.ReadFile(cfg.InitScriptPath)
	if err != nil {
		return nil, err
	}

	statements := ddlStatementSepRe.Split(string(contents), -1)
	ret := statements[:0]
	for _, stmt := range statements {
		stmt = commentRe.ReplaceAllString(stmt, "")
		stmt = strings.TrimSpace(stmt)
		if stmt != "" {
			ret = append(ret, stmt)
		}
	}
	return ret, nil
}

// TempDB is a temporary Spanner database.
type TempDB struct {
	Name string
	ts   oauth2.TokenSource
}

// Client returns a spanner client connected to the database.
func (db *TempDB) Client(ctx context.Context) (*spanner.Client, error) {
	return spanner.NewClient(ctx, db.Name, option.WithTokenSource(db.ts))
}

// Drop deletes the database.
func (db *TempDB) Drop(ctx context.Context) error {
	client, err := spandb.NewDatabaseAdminClient(ctx, option.WithTokenSource(db.ts))
	if err != nil {
		return err
	}
	defer client.Close()

	return client.DropDatabase(ctx, &dbpb.DropDatabaseRequest{
		Database: db.Name,
	})
}

// NewTempDB creates a temporary database with a random name.
// The caller is responsible for calling Drop on the returned TempDB to
// cleanup resources after usage.
func NewTempDB(ctx context.Context, cfg TempDBConfig) (*TempDB, error) {
	instanceName := cfg.InstanceName
	if instanceName == "" {
		instanceName = chromeinfra.TestSpannerInstance
	}

	ts, err := cfg.tokenSource(ctx)
	if err != nil {
		return nil, err
	}

	initStatements, err := cfg.readDDLStatements()
	if err != nil {
		return nil, errors.Annotate(err, "failed to read %q", cfg.InitScriptPath).Err()
	}

	client, err := spandb.NewDatabaseAdminClient(ctx, option.WithTokenSource(ts))
	if err != nil {
		return nil, err
	}
	defer client.Close()

	// Generate a random database name.
	var random uint32
	if err := binary.Read(rand.Reader, binary.LittleEndian, &random); err != nil {
		panic(err)
	}
	dbName := fmt.Sprintf("test%d", random)
	if u, err := user.Current(); err == nil && u.Username != "" {
		dbName += "_by_" + u.Username
	}

	dbOp, err := client.CreateDatabase(ctx, &dbpb.CreateDatabaseRequest{
		Parent:          instanceName,
		CreateStatement: "CREATE DATABASE " + dbName,
		ExtraStatements: initStatements,
	})
	if err != nil {
		return nil, errors.Annotate(err, "failed to create database").Err()
	}
	db, err := dbOp.Wait(ctx)
	if err != nil {
		return nil, errors.Annotate(err, "failed to create database").Err()
	}

	return &TempDB{
		Name: db.Name,
		ts:   ts,
	}, nil
}
<commit_msg>Reland "[spantest] Sanitize username in db name"<commit_after>// Copyright 2019 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package spantest implements creation/destruction of a temporary Spanner
// database.
package spantest

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"io/ioutil"
	"os/user"
	"regexp"
	"strings"

	"cloud.google.com/go/spanner"
	spandb "cloud.google.com/go/spanner/admin/database/apiv1"
	"golang.org/x/oauth2"
	"google.golang.org/api/option"
	dbpb "google.golang.org/genproto/googleapis/spanner/admin/database/v1"

	"go.chromium.org/luci/auth"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/hardcoded/chromeinfra"
)

// TempDBConfig specifies how to create a temporary database.
type TempDBConfig struct {
	// InstanceName is the name of Spannner instance where to create the
	// temporary database.
	// Format: projects/{project}/instances/{instance}.
	// Defaults to chromeinfra.TestSpannerInstance.
	InstanceName string

	// TokenSource will be used to authenticate to Spanner.
	// If nil, auth.Authenticator with SilentLogin and chrome-infra auth options
	// will be used.
	// This means that that the user may have to login with luci-auth tool.
	TokenSource oauth2.TokenSource

	// InitScriptPath is a path to a DDL script to initialize the database.
	//
	// In lieu of a proper DDL parser, it is parsed using regexes.
	// Therefore the script MUST:
	//   - Use `#`` and/or `--`` for comments. No block comments.
	//   - Separate DDL statements with `;\n`.
	//
	// If empty, the database is created with no tables.
	InitScriptPath string
}

func (cfg *TempDBConfig) tokenSource(ctx context.Context) (oauth2.TokenSource, error) {
	if cfg.TokenSource != nil {
		return cfg.TokenSource, nil
	}

	opts := chromeinfra.DefaultAuthOptions()
	opts.Scopes = spandb.DefaultAuthScopes()
	a := auth.NewAuthenticator(ctx, auth.SilentLogin, opts)
	if err := a.CheckLoginRequired(); err != nil {
		return nil, errors.Annotate(err, "please login with `luci-auth login -scopes %q`", strings.Join(opts.Scopes, " ")).Err()
	}
	return a.TokenSource()
}

var ddlStatementSepRe = regexp.MustCompile(`;\s*\n`)
var commentRe = regexp.MustCompile(`(--|#)[^\n]*`)

// readDDLStatements read the file at cfg.InitScriptPath as a sequence of DDL
// statements. If the path is empty, returns (nil, nil).
func (cfg *TempDBConfig) readDDLStatements() ([]string, error) {
	if cfg.InitScriptPath == "" {
		return nil, nil
	}

	contents, err := ioutil.ReadFile(cfg.InitScriptPath)
	if err != nil {
		return nil, err
	}

	statements := ddlStatementSepRe.Split(string(contents), -1)
	ret := statements[:0]
	for _, stmt := range statements {
		stmt = commentRe.ReplaceAllString(stmt, "")
		stmt = strings.TrimSpace(stmt)
		if stmt != "" {
			ret = append(ret, stmt)
		}
	}
	return ret, nil
}

// TempDB is a temporary Spanner database.
type TempDB struct {
	Name string
	ts   oauth2.TokenSource
}

// Client returns a spanner client connected to the database.
func (db *TempDB) Client(ctx context.Context) (*spanner.Client, error) {
	return spanner.NewClient(ctx, db.Name, option.WithTokenSource(db.ts))
}

// Drop deletes the database.
func (db *TempDB) Drop(ctx context.Context) error {
	client, err := spandb.NewDatabaseAdminClient(ctx, option.WithTokenSource(db.ts))
	if err != nil {
		return err
	}
	defer client.Close()

	return client.DropDatabase(ctx, &dbpb.DropDatabaseRequest{
		Database: db.Name,
	})
}

var dbNameAlphabetInversedRe = regexp.MustCompile(`[^\w]+`)

// NewTempDB creates a temporary database with a random name.
// The caller is responsible for calling Drop on the returned TempDB to
// cleanup resources after usage.
func NewTempDB(ctx context.Context, cfg TempDBConfig) (*TempDB, error) {
	instanceName := cfg.InstanceName
	if instanceName == "" {
		instanceName = chromeinfra.TestSpannerInstance
	}

	ts, err := cfg.tokenSource(ctx)
	if err != nil {
		return nil, err
	}

	initStatements, err := cfg.readDDLStatements()
	if err != nil {
		return nil, errors.Annotate(err, "failed to read %q", cfg.InitScriptPath).Err()
	}

	client, err := spandb.NewDatabaseAdminClient(ctx, option.WithTokenSource(ts))
	if err != nil {
		return nil, err
	}
	defer client.Close()

	// Generate a random database name.
	var random uint32
	if err := binary.Read(rand.Reader, binary.LittleEndian, &random); err != nil {
		panic(err)
	}
	dbName := fmt.Sprintf("test%d", random)
	if u, err := user.Current(); err == nil && u.Username != "" {
		dbName += "_by_" + dbNameAlphabetInversedRe.ReplaceAllLiteralString(u.Username, "_")
	}

	dbOp, err := client.CreateDatabase(ctx, &dbpb.CreateDatabaseRequest{
		Parent:          instanceName,
		CreateStatement: "CREATE DATABASE " + dbName,
		ExtraStatements: initStatements,
	})
	if err != nil {
		return nil, errors.Annotate(err, "failed to create database").Err()
	}
	db, err := dbOp.Wait(ctx)
	if err != nil {
		return nil, errors.Annotate(err, "failed to create database").Err()
	}

	return &TempDB{
		Name: db.Name,
		ts:   ts,
	}, nil
}
<|endoftext|>


<commit_before>//
// Provides unified interface to access needed flags when you are testing.
//
// Property file
//
// "NewTestFlags()" has two way to load property file:
//
// 	OWL_TEST_PROPS_FILE - Environment Variable
// 	"-owl.test.propfile" - flag of "go test"
//
// Entry Environment Variables
//
// Following environment variables are supported when calling "NewTestFlags()".
//
// 	OWL_TEST_PROPS - As same as "-owl.test"
// 	OWL_TEST_PROPS_SEP - As same as "-owl.test.sep"
//
// Entry Flags
//
// There are only two flags of Golang needed:
//
// 	-owl.test=<properties>
// 	-owl.test.sep=<separator for properties>
//
// The format used by "owl.test" is property file:
//
// https://en.wikipedia.org/wiki/.properties
//
// In order to separate properties, "owl.test.sep"(as regular expression) would be used to
// recognize a record of property file.
//
// See "DEFAULT_SEPARATOR" constant for default separator.
//
// Loading priority
//
// When you execute go test with viable values for both of the environment variables and the flags,
// the priority is:
//
// 	1. Load file from environment variable(OWL_TEST_PROPS_FILE)
// 	2. Load properties of environment variable(OWL_TEST_PROPS)
// 	3. Load file from flag("-owl.test.propfile")
// 	4. Load properties from flag("-owl.test")
//
// Pre-defined Properties - Features
//
// There are some pre-defined properties:
//
// 	mysql - MySql connection
//
// 	client.http.host - HTTP client
// 	client.http.port - HTTP client
// 	client.http.ssl - HTTP client
// 	client.http.resource - HTTP client
//
// 	client.jsonrpc.host - JSONRPC Client
// 	client.jsonrpc.port - JSONRPC Client
//
// 	it.web.enable - IT to Web
//
// The object of "*TestFlags" provides various functions to check
// whether or not some configuration for testing are enabled.
//
// For example, "HasMySql()" would let you know whether "mysql=<conn>" is viable.
//
// Pre-defined Properties - Owl Databases of MySql
//
// Following list shows build-in supporting databases of Owl Database:
//
// 	mysql.owl_portal - MySql connection on OWL-Portal
// 	mysql.owl_graph - MySql connection on OWL-Graph
// 	mysql.owl_uic - MySql connection on OWL-Uic
// 	mysql.owl_links - MySql connection on OWL-Links
// 	mysql.owl_grafana - MySql connection on OWL-Grafana
// 	mysql.owl_dashboard - MySql connection on OWL-Dashboard
// 	mysql.owl_boss - MySql connection on OWL-Boss
//
// You could use "HasMySqlOfOwlDb(int)" or "GetMysqlOfOwlDb(int)" to retrieve value of properties.
//
// Constraint
//
// The empty string of property value would be considered as non-viable.
package flag

import (
	"flag"
	"fmt"
	"regexp"
	"strings"

	"github.com/juju/errors"
	"github.com/spf13/viper"
)

const (
	// Default separator
	DEFAULT_SEPARATOR = "\\s+"
)

/*
Bit reservation principals:

	Bits (0~7): For clients of various protocols
	Bits (8~15): For databases
	Bits (16~23): For misc(e.x. mocking server)
*/
const (
	// Feature of HTTP client
	F_HttpClient = 0x01
	// Feature of JSONRPC client
	F_JsonRpcClient = 0x02
	// Feature of MySql
	F_MySql = 0x100
	// Feature of IT web
	F_ItWeb = 0x10000
)

const (
	ENV_OWL_TEST_PROPS      = "OWL_TEST_PROPS"
	ENV_OWL_TEST_PROPS_SEP  = "OWL_TEST_PROPS_SEP"
	ENV_OWL_TEST_PROPS_FILE = "OWL_TEST_PROPS_FILE"
)

const (
	OWL_DB_PORTAL    = 0x01
	OWL_DB_GRAPH     = 0x02
	OWL_DB_UIC       = 0x04
	OWL_DB_LINKS     = 0x8
	OWL_DB_GRAFANA   = 0x10
	OWL_DB_DASHBOARD = 0x20
	OWL_DB_BOSS      = 0x40
)

var owlDbMap = map[int]string{
	OWL_DB_PORTAL:    "mysql.owl_portal",
	OWL_DB_GRAPH:     "mysql.owl_graph",
	OWL_DB_UIC:       "mysql.owl_uic",
	OWL_DB_LINKS:     "mysql.owl_links",
	OWL_DB_GRAFANA:   "mysql.owl_grafana",
	OWL_DB_DASHBOARD: "mysql.owl_dashboard",
	OWL_DB_BOSS:      "mysql.owl_boss",
}

var (
	owlTest         = flag.String("owl.test", "", "Owl typedFlags for testing properties")
	owlTestSep      = flag.String("owl.test.sep", DEFAULT_SEPARATOR, "Owl typedFlags for separator of properties")
	owlTestPropFile = flag.String("owl.test.propfile", "", "Owl property file for testing")
)

// Initializes the object of "*TestFlags" by parsing flag automatically.
//
// This function parses os.Args every time it is get called.
func NewTestFlags() *TestFlags {
	viperLoader := newMultiPropLoader()
	viperLoader.loadFromEnv()
	viperLoader.loadFromFlag()

	/**
	 * Setup Flags of testing
	 */
	newFlags := newTestFlags(viperLoader.viperObj)
	// :~)

	return newFlags
}

// Convenient type used to access specific testing environment of OWL.
type TestFlags struct {
	typedFlags map[string]interface{}

	viperObj *viper.Viper
}

func (f *TestFlags) GetViper() *viper.Viper {
	return f.viperObj
}

// Gets property value of "mysql"
func (f *TestFlags) GetMySql() string {
	if f.HasMySql() {
		return f.typedFlags["mysql"].(string)
	}

	return ""
}

func (f *TestFlags) GetMysqlOfOwlDb(owlDb int) string {
	propName, ok := owlDbMap[owlDb]
	if !ok {
		panic(fmt.Sprintf("Unsupported OWL Db: %v", owlDb))
	}

	return strings.TrimSpace(f.viperObj.GetString(propName))
}

// Gets property values of:
// 	client.http.host
// 	client.http.port
// 	client.http.resource
// 	client.http.ssl
func (f *TestFlags) GetHttpClient() (string, uint16, string, bool) {
	if f.HasHttpClient() {
		return f.typedFlags["client.http.host"].(string), f.typedFlags["client.http.port"].(uint16),
			f.typedFlags["client.http.resource"].(string), f.typedFlags["client.http.ssl"].(bool)
	}

	return "", 0, "", false
}

// Gets property values of "client.jsonrpc.host" and "client.jsonrpc.port"
func (f *TestFlags) GetJsonRpcClient() (string, uint16) {
	if f.HasJsonRpcClient() {
		return f.typedFlags["client.jsonrpc.host"].(string), f.typedFlags["client.jsonrpc.port"].(uint16)
	}

	return "", 0
}

// Gives "true" if and only if following properties are viable:
//
// 	client.jsonrpc.host=
// 	client.jsonrpc.port=
//
// Example:
// 	"-owl.flag=client.jsonrpc.host=127.0.0.1 client.jsonrpc.port=3396"
func (f *TestFlags) HasJsonRpcClient() bool {
	_, hostOk := f.typedFlags["client.jsonrpc.host"]
	_, portOk := f.typedFlags["client.jsonrpc.port"]

	return hostOk && portOk
}

// Gives "true" if and only if following properties are viable:
//
// 	client.http.host=
// 	client.http.port=
//
// Example:
// 	"-owl.flag=client.http.host=127.0.0.1 client.http.port=3396"
func (f *TestFlags) HasHttpClient() bool {
	_, hostOk := f.typedFlags["client.http.host"]
	_, portOk := f.typedFlags["client.http.port"]

	return hostOk && portOk
}

// Gives "true" if and only if "mysql" property is non-empty
//
// Example:
// 	"-owl.flag=mysql=root:cepave@tcp(192.168.20.50:3306)/falcon_portal_test?parseTime=True&loc=Local"
func (f *TestFlags) HasMySql() bool {
	_, ok := f.typedFlags["mysql"]
	return ok
}

// Gives "true" if and only if "mysql.<db>" property is non-empty
//
// Example:
// 	"-owl.flag=mysql.portal=root:cepave@tcp(192.168.20.50:3306)/falcon_portal_test?parseTime=True&loc=Local"
func (f *TestFlags) HasMySqlOfOwlDb(owlDb int) bool {
	return f.GetMysqlOfOwlDb(owlDb) != ""
}

// Gives "true" if and only if "it.web.enable" property is true
//
// Example:
// 	"-owl.flag=it.web.enable=true"
func (f *TestFlags) HasItWeb() bool {
	hasItWeb, ok := f.typedFlags["it.web.enable"].(bool)
	return ok && hasItWeb
}

func (f *TestFlags) setupByViper() {
	viperObj := f.viperObj

	/**
	 * MySql
	 */
	if viperObj.IsSet("mysql") {
		setNonEmptyString(f.typedFlags, "mysql", viperObj)
	}
	// :~)

	/**
	 * HTTP client
	 */
	if viperObj.IsSet("client.http.host") {
		setNonEmptyString(f.typedFlags, "client.http.host", viperObj)
	}
	if viperObj.IsSet("client.http.port") {
		setValidPort(f.typedFlags, "client.http.port", viperObj)
	}
	viperObj.SetDefault("client.http.ssl", false)
	f.typedFlags["client.http.ssl"] = viperObj.GetBool("client.http.ssl")
	viperObj.SetDefault("client.http.resource", "")
	f.typedFlags["client.http.resource"] = viperObj.GetString("client.http.resource")
	// :~)

	/**
	 * JSONRPC Client
	 */
	if viperObj.IsSet("client.jsonrpc.host") {
		setNonEmptyString(f.typedFlags, "client.jsonrpc.host", viperObj)
	}
	if viperObj.IsSet("client.jsonrpc.port") {
		setValidPort(f.typedFlags, "client.jsonrpc.port", viperObj)
	}
	// :~)

	/**
	 * Start web for integration test
	 */
	if viperObj.IsSet("it.web.enable") {
		f.typedFlags["it.web.enable"] = viperObj.GetBool("it.web.enable")
	}
	// :~)
}

func setValidPort(props map[string]interface{}, key string, viperObj *viper.Viper) {
	v := viperObj.GetInt(key)
	if v > 0 {
		props[key] = uint16(v)
	}
}
func setNonEmptyString(props map[string]interface{}, key string, viperObj *viper.Viper) {
	value := strings.TrimSpace(viperObj.GetString(key))
	if len(value) > 0 {
		props[key] = value
	}
}

func newTestFlags(viperObj *viper.Viper) *TestFlags {
	testFlag := &TestFlags{
		make(map[string]interface{}),
		viperObj,
	}
	testFlag.setupByViper()
	return testFlag
}

func newMultiPropLoader() *multiPropLoader {
	viperObj := viper.New()
	viperObj.SetConfigType("properties")

	return &multiPropLoader{viperObj}
}

type multiPropLoader struct {
	viperObj *viper.Viper
}

func (l *multiPropLoader) loadFromEnv() {
	envLoader := viper.New()
	envLoader.SetDefault(ENV_OWL_TEST_PROPS_SEP, DEFAULT_SEPARATOR)
	envLoader.BindEnv(ENV_OWL_TEST_PROPS)
	envLoader.BindEnv(ENV_OWL_TEST_PROPS_SEP)
	envLoader.BindEnv(ENV_OWL_TEST_PROPS_FILE)

	l.loadProperties(
		envLoader.GetString(ENV_OWL_TEST_PROPS_FILE),
		envLoader.GetString(ENV_OWL_TEST_PROPS),
		envLoader.GetString(ENV_OWL_TEST_PROPS_SEP),
	)
}
func (l *multiPropLoader) loadFromFlag() {
	if !flag.Parsed() {
		flag.Parse()
	}

	l.loadProperties(*owlTestPropFile, *owlTest, *owlTestSep)
}

func (l *multiPropLoader) loadProperties(filename string, propertiesString string, separator string) {
	/**
	 * Loads property file into viper object
	 */
	if filename != "" {
		l.viperObj.SetConfigFile(filename)
		if err := l.viperObj.MergeInConfig(); err != nil {
			logger.Warnf("Load property file[%s] has error: %v", filename, err)
		}
	}
	// :~)

	propertiesString = strings.TrimSpace(propertiesString)
	if propertiesString == "" {
		return
	}

	splitRegExp := regexp.MustCompile(separator)
	properties := splitRegExp.ReplaceAllString(propertiesString, "\n")

	/**
	 * Loads properties into viper object
	 */
	if err := errors.Annotate(
		l.viperObj.MergeConfig(strings.NewReader(properties)),
		"Read owl.test as format of property file has error",
	); err != nil {
		panic(errors.Details(err))
	}
	// :~)
}
<commit_msg>[OWL-1950] Refactoring the implementation of NewTestFlags() to singleton pattern<commit_after>//
// Provides unified interface to access needed flags when you are testing.
//
// Property file
//
// "NewTestFlags()" has two way to load property file:
//
// 	OWL_TEST_PROPS_FILE - Environment Variable
// 	"-owl.test.propfile" - flag of "go test"
//
// Entry Environment Variables
//
// Following environment variables are supported when calling "NewTestFlags()".
//
// 	OWL_TEST_PROPS - As same as "-owl.test"
// 	OWL_TEST_PROPS_SEP - As same as "-owl.test.sep"
//
// Entry Flags
//
// There are only two flags of Golang needed:
//
// 	-owl.test=<properties>
// 	-owl.test.sep=<separator for properties>
//
// The format used by "owl.test" is property file:
//
// https://en.wikipedia.org/wiki/.properties
//
// In order to separate properties, "owl.test.sep"(as regular expression) would be used to
// recognize a record of property file.
//
// See "DEFAULT_SEPARATOR" constant for default separator.
//
// Loading priority
//
// When you execute go test with viable values for both of the environment variables and the flags,
// the priority is:
//
// 	1. Load file from environment variable(OWL_TEST_PROPS_FILE)
// 	2. Load properties of environment variable(OWL_TEST_PROPS)
// 	3. Load file from flag("-owl.test.propfile")
// 	4. Load properties from flag("-owl.test")
//
// Pre-defined Properties - Features
//
// There are some pre-defined properties:
//
// 	mysql - MySql connection
//
// 	client.http.host - HTTP client
// 	client.http.port - HTTP client
// 	client.http.ssl - HTTP client
// 	client.http.resource - HTTP client
//
// 	client.jsonrpc.host - JSONRPC Client
// 	client.jsonrpc.port - JSONRPC Client
//
// 	it.web.enable - IT to Web
//
// The object of "*TestFlags" provides various functions to check
// whether or not some configuration for testing are enabled.
//
// For example, "HasMySql()" would let you know whether "mysql=<conn>" is viable.
//
// Pre-defined Properties - Owl Databases of MySql
//
// Following list shows build-in supporting databases of Owl Database:
//
// 	mysql.owl_portal - MySql connection on OWL-Portal
// 	mysql.owl_graph - MySql connection on OWL-Graph
// 	mysql.owl_uic - MySql connection on OWL-Uic
// 	mysql.owl_links - MySql connection on OWL-Links
// 	mysql.owl_grafana - MySql connection on OWL-Grafana
// 	mysql.owl_dashboard - MySql connection on OWL-Dashboard
// 	mysql.owl_boss - MySql connection on OWL-Boss
//
// You could use "HasMySqlOfOwlDb(int)" or "GetMysqlOfOwlDb(int)" to retrieve value of properties.
//
// Constraint
//
// The empty string of property value would be considered as non-viable.
package flag

import (
	"flag"
	"fmt"
	"regexp"
	"strings"

	"github.com/juju/errors"
	"github.com/spf13/viper"
)

const (
	// Default separator
	DEFAULT_SEPARATOR = "\\s+"
)

/*
Bit reservation principals:

	Bits (0~7): For clients of various protocols
	Bits (8~15): For databases
	Bits (16~23): For misc(e.x. mocking server)
*/
const (
	// Feature of HTTP client
	F_HttpClient = 0x01
	// Feature of JSONRPC client
	F_JsonRpcClient = 0x02
	// Feature of MySql
	F_MySql = 0x100
	// Feature of IT web
	F_ItWeb = 0x10000
)

const (
	ENV_OWL_TEST_PROPS      = "OWL_TEST_PROPS"
	ENV_OWL_TEST_PROPS_SEP  = "OWL_TEST_PROPS_SEP"
	ENV_OWL_TEST_PROPS_FILE = "OWL_TEST_PROPS_FILE"
)

const (
	OWL_DB_PORTAL    = 0x01
	OWL_DB_GRAPH     = 0x02
	OWL_DB_UIC       = 0x04
	OWL_DB_LINKS     = 0x8
	OWL_DB_GRAFANA   = 0x10
	OWL_DB_DASHBOARD = 0x20
	OWL_DB_BOSS      = 0x40
)

var owlDbMap = map[int]string{
	OWL_DB_PORTAL:    "mysql.owl_portal",
	OWL_DB_GRAPH:     "mysql.owl_graph",
	OWL_DB_UIC:       "mysql.owl_uic",
	OWL_DB_LINKS:     "mysql.owl_links",
	OWL_DB_GRAFANA:   "mysql.owl_grafana",
	OWL_DB_DASHBOARD: "mysql.owl_dashboard",
	OWL_DB_BOSS:      "mysql.owl_boss",
}

var (
	owlTest         = flag.String("owl.test", "", "Owl typedFlags for testing properties")
	owlTestSep      = flag.String("owl.test.sep", DEFAULT_SEPARATOR, "Owl typedFlags for separator of properties")
	owlTestPropFile = flag.String("owl.test.propfile", "", "Owl property file for testing")
)

var singletonTestFlags *TestFlags

// Initializes the object of "*TestFlags" by parsing flag automatically.
//
// This function parses os.Args every time it is get called.
func NewTestFlags() *TestFlags {
	if singletonTestFlags != nil {
		return singletonTestFlags.clone()
	}

	viperLoader := newMultiPropLoader()
	viperLoader.loadFromEnv()
	viperLoader.loadFromFlag()

	/**
	 * Setup Flags of testing
	 */
	newFlags := newTestFlags(viperLoader.viperObj)
	// :~)

	singletonTestFlags = newFlags

	return singletonTestFlags.clone()
}

// Convenient type used to access specific testing environment of OWL.
type TestFlags struct {
	typedFlags map[string]interface{}

	viperObj *viper.Viper
}

func (f *TestFlags) GetViper() *viper.Viper {
	return f.viperObj
}

// Gets property value of "mysql"
func (f *TestFlags) GetMySql() string {
	if f.HasMySql() {
		return f.typedFlags["mysql"].(string)
	}

	return ""
}

func (f *TestFlags) GetMysqlOfOwlDb(owlDb int) string {
	propName, ok := owlDbMap[owlDb]
	if !ok {
		panic(fmt.Sprintf("Unsupported OWL Db: %v", owlDb))
	}

	return strings.TrimSpace(f.viperObj.GetString(propName))
}

// Gets property values of:
// 	client.http.host
// 	client.http.port
// 	client.http.resource
// 	client.http.ssl
func (f *TestFlags) GetHttpClient() (string, uint16, string, bool) {
	if f.HasHttpClient() {
		return f.typedFlags["client.http.host"].(string), f.typedFlags["client.http.port"].(uint16),
			f.typedFlags["client.http.resource"].(string), f.typedFlags["client.http.ssl"].(bool)
	}

	return "", 0, "", false
}

// Gets property values of "client.jsonrpc.host" and "client.jsonrpc.port"
func (f *TestFlags) GetJsonRpcClient() (string, uint16) {
	if f.HasJsonRpcClient() {
		return f.typedFlags["client.jsonrpc.host"].(string), f.typedFlags["client.jsonrpc.port"].(uint16)
	}

	return "", 0
}

// Gives "true" if and only if following properties are viable:
//
// 	client.jsonrpc.host=
// 	client.jsonrpc.port=
//
// Example:
// 	"-owl.flag=client.jsonrpc.host=127.0.0.1 client.jsonrpc.port=3396"
func (f *TestFlags) HasJsonRpcClient() bool {
	_, hostOk := f.typedFlags["client.jsonrpc.host"]
	_, portOk := f.typedFlags["client.jsonrpc.port"]

	return hostOk && portOk
}

// Gives "true" if and only if following properties are viable:
//
// 	client.http.host=
// 	client.http.port=
//
// Example:
// 	"-owl.flag=client.http.host=127.0.0.1 client.http.port=3396"
func (f *TestFlags) HasHttpClient() bool {
	_, hostOk := f.typedFlags["client.http.host"]
	_, portOk := f.typedFlags["client.http.port"]

	return hostOk && portOk
}

// Gives "true" if and only if "mysql" property is non-empty
//
// Example:
// 	"-owl.flag=mysql=root:cepave@tcp(192.168.20.50:3306)/falcon_portal_test?parseTime=True&loc=Local"
func (f *TestFlags) HasMySql() bool {
	_, ok := f.typedFlags["mysql"]
	return ok
}

// Gives "true" if and only if "mysql.<db>" property is non-empty
//
// Example:
// 	"-owl.flag=mysql.portal=root:cepave@tcp(192.168.20.50:3306)/falcon_portal_test?parseTime=True&loc=Local"
func (f *TestFlags) HasMySqlOfOwlDb(owlDb int) bool {
	return f.GetMysqlOfOwlDb(owlDb) != ""
}

// Gives "true" if and only if "it.web.enable" property is true
//
// Example:
// 	"-owl.flag=it.web.enable=true"
func (f *TestFlags) HasItWeb() bool {
	hasItWeb, ok := f.typedFlags["it.web.enable"].(bool)
	return ok && hasItWeb
}

func (f *TestFlags) clone() *TestFlags {
	newViper := viper.New()

	for k, v := range f.viperObj.AllSettings() {
		newViper.Set(k, v)
	}

	newTypedFlags := make(map[string]interface{})
	for k, v := range f.typedFlags {
		newTypedFlags[k] = v
	}

	return &TestFlags{
		newTypedFlags, newViper,
	}
}

func (f *TestFlags) setupByViper() {
	viperObj := f.viperObj

	/**
	 * MySql
	 */
	if viperObj.IsSet("mysql") {
		setNonEmptyString(f.typedFlags, "mysql", viperObj)
	}
	// :~)

	/**
	 * HTTP client
	 */
	if viperObj.IsSet("client.http.host") {
		setNonEmptyString(f.typedFlags, "client.http.host", viperObj)
	}
	if viperObj.IsSet("client.http.port") {
		setValidPort(f.typedFlags, "client.http.port", viperObj)
	}
	viperObj.SetDefault("client.http.ssl", false)
	f.typedFlags["client.http.ssl"] = viperObj.GetBool("client.http.ssl")
	viperObj.SetDefault("client.http.resource", "")
	f.typedFlags["client.http.resource"] = viperObj.GetString("client.http.resource")
	// :~)

	/**
	 * JSONRPC Client
	 */
	if viperObj.IsSet("client.jsonrpc.host") {
		setNonEmptyString(f.typedFlags, "client.jsonrpc.host", viperObj)
	}
	if viperObj.IsSet("client.jsonrpc.port") {
		setValidPort(f.typedFlags, "client.jsonrpc.port", viperObj)
	}
	// :~)

	/**
	 * Start web for integration test
	 */
	if viperObj.IsSet("it.web.enable") {
		f.typedFlags["it.web.enable"] = viperObj.GetBool("it.web.enable")
	}
	// :~)
}

func setValidPort(props map[string]interface{}, key string, viperObj *viper.Viper) {
	v := viperObj.GetInt(key)
	if v > 0 {
		props[key] = uint16(v)
	}
}
func setNonEmptyString(props map[string]interface{}, key string, viperObj *viper.Viper) {
	value := strings.TrimSpace(viperObj.GetString(key))
	if len(value) > 0 {
		props[key] = value
	}
}

func newTestFlags(viperObj *viper.Viper) *TestFlags {
	testFlag := &TestFlags{
		make(map[string]interface{}),
		viperObj,
	}
	testFlag.setupByViper()
	return testFlag
}

func newMultiPropLoader() *multiPropLoader {
	viperObj := viper.New()
	viperObj.SetConfigType("properties")

	return &multiPropLoader{viperObj}
}

type multiPropLoader struct {
	viperObj *viper.Viper
}

func (l *multiPropLoader) loadFromEnv() {
	envLoader := viper.New()
	envLoader.SetDefault(ENV_OWL_TEST_PROPS_SEP, DEFAULT_SEPARATOR)
	envLoader.BindEnv(ENV_OWL_TEST_PROPS)
	envLoader.BindEnv(ENV_OWL_TEST_PROPS_SEP)
	envLoader.BindEnv(ENV_OWL_TEST_PROPS_FILE)

	l.loadProperties(
		envLoader.GetString(ENV_OWL_TEST_PROPS_FILE),
		envLoader.GetString(ENV_OWL_TEST_PROPS),
		envLoader.GetString(ENV_OWL_TEST_PROPS_SEP),
	)
}
func (l *multiPropLoader) loadFromFlag() {
	if !flag.Parsed() {
		flag.Parse()
	}

	l.loadProperties(*owlTestPropFile, *owlTest, *owlTestSep)
}

func (l *multiPropLoader) loadProperties(filename string, propertiesString string, separator string) {
	/**
	 * Loads property file into viper object
	 */
	if filename != "" {
		l.viperObj.SetConfigFile(filename)
		if err := l.viperObj.MergeInConfig(); err != nil {
			logger.Warnf("Load property file[%s] has error: %v", filename, err)
		}
	}
	// :~)

	propertiesString = strings.TrimSpace(propertiesString)
	if propertiesString == "" {
		return
	}

	splitRegExp := regexp.MustCompile(separator)
	properties := splitRegExp.ReplaceAllString(propertiesString, "\n")

	/**
	 * Loads properties into viper object
	 */
	if err := errors.Annotate(
		l.viperObj.MergeConfig(strings.NewReader(properties)),
		"Read owl.test as format of property file has error",
	); err != nil {
		panic(errors.Details(err))
	}
	// :~)
}
<|endoftext|>