<commit_before>package service

import (
	"github.com/go-kit/kit/endpoint"
	"github.com/kolide/kolide-ose/server/kolide"
	"golang.org/x/net/context"
)

////////////////////////////////////////////////////////////////////////////////
// Get Pack
////////////////////////////////////////////////////////////////////////////////

type getPackRequest struct {
	ID uint
}

type packResponse struct {
	kolide.Pack
	QueryCount   uint `json:"query_count"`
	TargetsCount uint `json:"targets_count"`
}

type getPackResponse struct {
	Pack packResponse `json:"pack,omitempty"`
	Err  error        `json:"error,omitempty"`
}

func (r getPackResponse) error() error { return r.Err }

func makeGetPackEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(getPackRequest)

		pack, err := svc.GetPack(ctx, req.ID)
		if err != nil {
			return getPackResponse{Err: err}, nil
		}

		queries, err := svc.GetScheduledQueriesInPack(ctx, pack.ID, kolide.ListOptions{})
		if err != nil {
			return getPackResponse{Err: err}, nil
		}

		hosts, err := svc.ListHostsInPack(ctx, pack.ID, kolide.ListOptions{})
		if err != nil {
			return getPackResponse{Err: err}, nil
		}

		return getPackResponse{
			Pack: packResponse{
				Pack:         *pack,
				QueryCount:   uint(len(queries)),
				TargetsCount: uint(len(hosts)),
			},
		}, nil
	}
}

////////////////////////////////////////////////////////////////////////////////
// List Packs
////////////////////////////////////////////////////////////////////////////////

type listPacksRequest struct {
	ListOptions kolide.ListOptions
}

type listPacksResponse struct {
	Packs []packResponse `json:"packs"`
	Err   error          `json:"error,omitempty"`
}

func (r listPacksResponse) error() error { return r.Err }

func makeListPacksEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(listPacksRequest)
		packs, err := svc.ListPacks(ctx, req.ListOptions)
		if err != nil {
			return getPackResponse{Err: err}, nil
		}

		resp := listPacksResponse{Packs: []packResponse{}}
		for _, pack := range packs {
			queries, err := svc.GetScheduledQueriesInPack(ctx, pack.ID, kolide.ListOptions{})
			if err != nil {
				return getPackResponse{Err: err}, nil
			}
			hosts, err := svc.ListHostsInPack(ctx, pack.ID, kolide.ListOptions{})
			if err != nil {
				return getPackResponse{Err: err}, nil
			}
			resp.Packs = append(resp.Packs, packResponse{
				Pack:         *pack,
				QueryCount:   uint(len(queries)),
				TargetsCount: uint(len(hosts)),
			})
		}
		return resp, nil
	}
}

////////////////////////////////////////////////////////////////////////////////
// Create Pack
////////////////////////////////////////////////////////////////////////////////

type createPackRequest struct {
	payload kolide.PackPayload
}

type createPackResponse struct {
	Pack *kolide.Pack `json:"pack,omitempty"`
	Err  error        `json:"error,omitempty"`
}

func (r createPackResponse) error() error { return r.Err }

func makeCreatePackEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(createPackRequest)
		pack, err := svc.NewPack(ctx, req.payload)
		if err != nil {
			return createPackResponse{Err: err}, nil
		}
		return createPackResponse{pack, nil}, nil
	}
}

////////////////////////////////////////////////////////////////////////////////
// Modify Pack
////////////////////////////////////////////////////////////////////////////////

type modifyPackRequest struct {
	ID      uint
	payload kolide.PackPayload
}

type modifyPackResponse struct {
	Pack *kolide.Pack `json:"pack,omitempty"`
	Err  error        `json:"error,omitempty"`
}

func (r modifyPackResponse) error() error { return r.Err }

func makeModifyPackEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(modifyPackRequest)
		pack, err := svc.ModifyPack(ctx, req.ID, req.payload)
		if err != nil {
			return modifyPackResponse{Err: err}, nil
		}
		return modifyPackResponse{pack, nil}, nil
	}
}

////////////////////////////////////////////////////////////////////////////////
// Delete Pack
////////////////////////////////////////////////////////////////////////////////

type deletePackRequest struct {
	ID uint
}

type deletePackResponse struct {
	Err error `json:"error,omitempty"`
}

func (r deletePackResponse) error() error { return r.Err }

func makeDeletePackEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(deletePackRequest)
		err := svc.DeletePack(ctx, req.ID)
		if err != nil {
			return deletePackResponse{Err: err}, nil
		}
		return deletePackResponse{}, nil
	}
}

////////////////////////////////////////////////////////////////////////////////
// Add Label To Pack
////////////////////////////////////////////////////////////////////////////////

type addLabelToPackRequest struct {
	PackID  uint
	LabelID uint
}

type addLabelToPackResponse struct {
	Err error `json:"error,omitempty"`
}

func (r addLabelToPackResponse) error() error { return r.Err }

func makeAddLabelToPackEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(addLabelToPackRequest)
		err := svc.AddLabelToPack(ctx, req.LabelID, req.PackID)
		if err != nil {
			return addLabelToPackResponse{Err: err}, nil
		}
		return addLabelToPackResponse{}, nil
	}
}

////////////////////////////////////////////////////////////////////////////////
// Get Labels For Pack
////////////////////////////////////////////////////////////////////////////////

type getLabelsForPackRequest struct {
	PackID uint
}

type getLabelsForPackResponse struct {
	Labels []kolide.Label `json:"labels"`
	Err    error          `json:"error,omitempty"`
}

func (r getLabelsForPackResponse) error() error { return r.Err }

func makeGetLabelsForPackEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(getLabelsForPackRequest)
		labels, err := svc.ListLabelsForPack(ctx, req.PackID)
		if err != nil {
			return getLabelsForPackResponse{Err: err}, nil
		}

		var resp getLabelsForPackResponse
		for _, label := range labels {
			resp.Labels = append(resp.Labels, *label)
		}
		return resp, nil
	}
}

////////////////////////////////////////////////////////////////////////////////
// Delete Label From Pack
////////////////////////////////////////////////////////////////////////////////

type deleteLabelFromPackRequest struct {
	LabelID uint
	PackID  uint
}

type deleteLabelFromPackResponse struct {
	Err error `json:"error,omitempty"`
}

func (r deleteLabelFromPackResponse) error() error { return r.Err }

func makeDeleteLabelFromPackEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(deleteLabelFromPackRequest)
		err := svc.RemoveLabelFromPack(ctx, req.LabelID, req.PackID)
		if err != nil {
			return deleteLabelFromPackResponse{Err: err}, nil
		}
		return deleteLabelFromPackResponse{}, nil
	}
}
<commit_msg>Return query_count and targets_count in all pack responses (#731)<commit_after>package service

import (
	"github.com/go-kit/kit/endpoint"
	"github.com/kolide/kolide-ose/server/kolide"
	"golang.org/x/net/context"
)

////////////////////////////////////////////////////////////////////////////////
// Get Pack
////////////////////////////////////////////////////////////////////////////////

type getPackRequest struct {
	ID uint
}

type packResponse struct {
	kolide.Pack
	QueryCount   uint `json:"query_count"`
	TargetsCount uint `json:"targets_count"`
}

type getPackResponse struct {
	Pack packResponse `json:"pack,omitempty"`
	Err  error        `json:"error,omitempty"`
}

func (r getPackResponse) error() error { return r.Err }

func makeGetPackEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(getPackRequest)

		pack, err := svc.GetPack(ctx, req.ID)
		if err != nil {
			return getPackResponse{Err: err}, nil
		}

		queries, err := svc.GetScheduledQueriesInPack(ctx, pack.ID, kolide.ListOptions{})
		if err != nil {
			return getPackResponse{Err: err}, nil
		}

		hosts, err := svc.ListHostsInPack(ctx, pack.ID, kolide.ListOptions{})
		if err != nil {
			return getPackResponse{Err: err}, nil
		}

		return getPackResponse{
			Pack: packResponse{
				Pack:         *pack,
				QueryCount:   uint(len(queries)),
				TargetsCount: uint(len(hosts)),
			},
		}, nil
	}
}

////////////////////////////////////////////////////////////////////////////////
// List Packs
////////////////////////////////////////////////////////////////////////////////

type listPacksRequest struct {
	ListOptions kolide.ListOptions
}

type listPacksResponse struct {
	Packs []packResponse `json:"packs"`
	Err   error          `json:"error,omitempty"`
}

func (r listPacksResponse) error() error { return r.Err }

func makeListPacksEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(listPacksRequest)
		packs, err := svc.ListPacks(ctx, req.ListOptions)
		if err != nil {
			return getPackResponse{Err: err}, nil
		}

		resp := listPacksResponse{Packs: []packResponse{}}
		for _, pack := range packs {
			queries, err := svc.GetScheduledQueriesInPack(ctx, pack.ID, kolide.ListOptions{})
			if err != nil {
				return getPackResponse{Err: err}, nil
			}
			hosts, err := svc.ListHostsInPack(ctx, pack.ID, kolide.ListOptions{})
			if err != nil {
				return getPackResponse{Err: err}, nil
			}
			resp.Packs = append(resp.Packs, packResponse{
				Pack:         *pack,
				QueryCount:   uint(len(queries)),
				TargetsCount: uint(len(hosts)),
			})
		}
		return resp, nil
	}
}

////////////////////////////////////////////////////////////////////////////////
// Create Pack
////////////////////////////////////////////////////////////////////////////////

type createPackRequest struct {
	payload kolide.PackPayload
}

type createPackResponse struct {
	Pack packResponse `json:"pack,omitempty"`
	Err  error        `json:"error,omitempty"`
}

func (r createPackResponse) error() error { return r.Err }

func makeCreatePackEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(createPackRequest)
		pack, err := svc.NewPack(ctx, req.payload)
		if err != nil {
			return createPackResponse{Err: err}, nil
		}

		queries, err := svc.GetScheduledQueriesInPack(ctx, pack.ID, kolide.ListOptions{})
		if err != nil {
			return createPackResponse{Err: err}, nil
		}

		hosts, err := svc.ListHostsInPack(ctx, pack.ID, kolide.ListOptions{})
		if err != nil {
			return createPackResponse{Err: err}, nil
		}

		return createPackResponse{
			Pack: packResponse{
				Pack:         *pack,
				QueryCount:   uint(len(queries)),
				TargetsCount: uint(len(hosts)),
			},
			Err: nil,
		}, nil
	}
}

////////////////////////////////////////////////////////////////////////////////
// Modify Pack
////////////////////////////////////////////////////////////////////////////////

type modifyPackRequest struct {
	ID      uint
	payload kolide.PackPayload
}

type modifyPackResponse struct {
	Pack packResponse `json:"pack,omitempty"`
	Err  error        `json:"error,omitempty"`
}

func (r modifyPackResponse) error() error { return r.Err }

func makeModifyPackEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(modifyPackRequest)
		pack, err := svc.ModifyPack(ctx, req.ID, req.payload)
		if err != nil {
			return modifyPackResponse{Err: err}, nil
		}

		queries, err := svc.GetScheduledQueriesInPack(ctx, pack.ID, kolide.ListOptions{})
		if err != nil {
			return modifyPackResponse{Err: err}, nil
		}

		hosts, err := svc.ListHostsInPack(ctx, pack.ID, kolide.ListOptions{})
		if err != nil {
			return modifyPackResponse{Err: err}, nil
		}

		return modifyPackResponse{
			Pack: packResponse{
				Pack:         *pack,
				QueryCount:   uint(len(queries)),
				TargetsCount: uint(len(hosts)),
			},
			Err: nil,
		}, nil
	}
}

////////////////////////////////////////////////////////////////////////////////
// Delete Pack
////////////////////////////////////////////////////////////////////////////////

type deletePackRequest struct {
	ID uint
}

type deletePackResponse struct {
	Err error `json:"error,omitempty"`
}

func (r deletePackResponse) error() error { return r.Err }

func makeDeletePackEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(deletePackRequest)
		err := svc.DeletePack(ctx, req.ID)
		if err != nil {
			return deletePackResponse{Err: err}, nil
		}
		return deletePackResponse{}, nil
	}
}

////////////////////////////////////////////////////////////////////////////////
// Add Label To Pack
////////////////////////////////////////////////////////////////////////////////

type addLabelToPackRequest struct {
	PackID  uint
	LabelID uint
}

type addLabelToPackResponse struct {
	Err error `json:"error,omitempty"`
}

func (r addLabelToPackResponse) error() error { return r.Err }

func makeAddLabelToPackEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(addLabelToPackRequest)
		err := svc.AddLabelToPack(ctx, req.LabelID, req.PackID)
		if err != nil {
			return addLabelToPackResponse{Err: err}, nil
		}
		return addLabelToPackResponse{}, nil
	}
}

////////////////////////////////////////////////////////////////////////////////
// Get Labels For Pack
////////////////////////////////////////////////////////////////////////////////

type getLabelsForPackRequest struct {
	PackID uint
}

type getLabelsForPackResponse struct {
	Labels []kolide.Label `json:"labels"`
	Err    error          `json:"error,omitempty"`
}

func (r getLabelsForPackResponse) error() error { return r.Err }

func makeGetLabelsForPackEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(getLabelsForPackRequest)
		labels, err := svc.ListLabelsForPack(ctx, req.PackID)
		if err != nil {
			return getLabelsForPackResponse{Err: err}, nil
		}

		var resp getLabelsForPackResponse
		for _, label := range labels {
			resp.Labels = append(resp.Labels, *label)
		}
		return resp, nil
	}
}

////////////////////////////////////////////////////////////////////////////////
// Delete Label From Pack
////////////////////////////////////////////////////////////////////////////////

type deleteLabelFromPackRequest struct {
	LabelID uint
	PackID  uint
}

type deleteLabelFromPackResponse struct {
	Err error `json:"error,omitempty"`
}

func (r deleteLabelFromPackResponse) error() error { return r.Err }

func makeDeleteLabelFromPackEndpoint(svc kolide.Service) endpoint.Endpoint {
	return func(ctx context.Context, request interface{}) (interface{}, error) {
		req := request.(deleteLabelFromPackRequest)
		err := svc.RemoveLabelFromPack(ctx, req.LabelID, req.PackID)
		if err != nil {
			return deleteLabelFromPackResponse{Err: err}, nil
		}
		return deleteLabelFromPackResponse{}, nil
	}
}
<|endoftext|>


<commit_before>package client

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"

	kitlog "github.com/go-kit/kit/log"
	retryablehttp "github.com/hashicorp/go-retryablehttp"
)

// NomadServer is connection parameters to a nomad server
type NomadServer struct {
	Address string
	Port    int
}

// Details represents a json object in a nomad job
type Details struct {
	Running int `json:"running"`
}

// Summary represents a json object in a nomad job
type Summary struct {
	Details *Details `json:"server"`
}

// JobSummary represents a json object in a nomad job
type JobSummary struct {
	Summary *Summary `json:"Summary"`
}

// Job is a representation of nomad job
type Job struct {
	Name       string      `json:"Name"`
	Priority   int         `json:"Priority"`
	Status     string      `json:"status"`
	JobSummary *JobSummary `json:"JobSummary"`
}

// Task represents an allocated nomad task
type Task struct {
	State string `json:"State"`
}

// Alloc is a representation of a nomad allocation
type Alloc struct {
	ID            string          `json:"ID"`
	JobID         string          `json:"JobID"`
	NodeID        string          `json:"NodeID"`
	Name          string          `json:"Name"`
	ClientStatus  string          `json:"ClientStatus"`
	DesiredStatus string          `json:"DesiredStatus"`
	Tasks         map[string]Task `json:"TaskStates"`
}

// Host is a representation of a nomad client node
type Host struct {
	ID        string         `json:"ID"`
	Name      string         `json:"Name"`
	Drain     bool           `json:"Drain"`
	Resources *NodeResources `json:"Resources,omitempty"`
}

// NodeResources represents the resources of a nomad node
type NodeResources struct {
	Networks *Networks `json:"Networks"`
}

// Networks represents the network available on a nomad node
type Networks struct {
	IP string `json:"IP"`
}

// AllocNotFound indicates a missing allocation for a Job
type AllocNotFound struct {
	Jobname  string
	Hostname string
}

func log(keyvals ...interface{}) *bytes.Buffer {
	w := new(bytes.Buffer)
	kitlog.NewJSONLogger(w).Log(keyvals)
	return w
}

// Jobs will parse the json representation from the nomad rest api
// /v1/jobs
func Jobs(nomad *NomadServer) ([]Job, int, error) {
	jobs := make([]Job, 0)
	status, err := decodeJSON(url(nomad)+"/v1/jobs", &jobs)
	return jobs, status, err
}

// FindJob will parse the json representation and find the supplied job name
func FindJob(nomad *NomadServer, name string) (*Job, error) {
	jobs, _, _ := Jobs(nomad)
	for _, job := range jobs {
		if job.Name == name {
			return &job, nil
		}
	}
	return &Job{}, errors.New("job not found")
}

// Hosts will parse the json representation from the nomad rest api
// /v1/nodes
func Hosts(nomad *NomadServer) ([]Host, int, error) {
	hosts := make([]Host, 0)
	status, err := decodeJSON(url(nomad)+"/v1/nodes", &hosts)
	return hosts, status, err
}

// Drain will inform nomad to add/remove all allocations from that host
// depending on the value of enable
// Returns the http status code or an error
func Drain(nomad *NomadServer, id string, enable bool) (int, error) {
	resp, err := retryablehttp.Post(url(nomad)+"/v1/node/"+id+"/drain?enable="+strconv.FormatBool(enable), "application/json", nil)
	if resp != nil && resp.Body != nil {
		defer resp.Body.Close()
		return resp.StatusCode, err
	}
	return http.StatusInternalServerError, err
}

// StopJob will send a delete request to a nomad server to stop the provided job
func StopJob(nomad *NomadServer, job *Job) (int, error) {
	url := fmt.Sprintf("http://%v:%v/v1/job/%v", nomad.Address, nomad.Port, job.Name)
	req, err := retryablehttp.NewRequest(http.MethodDelete, url, bytes.NewReader([]byte{}))
	if err != nil {
		buf := log("event", "stop_job_request_error", "jobname", job.Name, "error", err.Error())
		return http.StatusInternalServerError, errors.New(buf.String())
	}
	resp, err := retryablehttp.NewClient().Do(req)
	if err != nil {
		buf := log("event", "stop_job_client_error", "jobname", job.Name, "error", err.Error())
		return resp.StatusCode, errors.New(buf.String())
	}
	return resp.StatusCode, nil
}

// SubmitJob requests that nomade run a json launch file located at the supplied path
// Returns an http status code or error
func SubmitJob(nomad *NomadServer, launchFilePath string) (int, error) {
	file, err := ioutil.ReadFile(launchFilePath)
	if err != nil {
		return http.StatusBadRequest, err
	}
	resp, err := retryablehttp.Post(url(nomad)+"/v1/jobs", "application/json", bytes.NewReader(file))
	defer resp.Body.Close()
	return resp.StatusCode, err
}

// Allocs will parse the json representation from the nomad rest api
// /v1/allocations
func Allocs(nomad *NomadServer) []Alloc {
	allocs := make([]Alloc, 0)
	decodeJSON(url(nomad)+"/v1/allocations", &allocs)
	return allocs
}

// FindAlloc will search through the Allocs on a provided Host to look for the
// allocations that match the provided Job
func FindAlloc(nomad *NomadServer, job *Job, host *Host) (*Alloc, error) {
	allocs := Allocs(nomad)
	for _, alloc := range allocs {
		if alloc.NodeID == host.ID && strings.Contains(alloc.Name, job.Name) {
			// We may be looking at a stale allocation and a newer one exists
			if alloc.DesiredStatus == "stop" && len(allocs) > 1 {
				continue
			}
			return &alloc, nil
		}
	}
	return &Alloc{}, &AllocNotFound{Hostname: host.Name, Jobname: job.Name}
}

// CheckTaskStates will return whether all of the tasks for an Alloc are
// in the state provided
func (alloc *Alloc) CheckTaskStates(state string) bool {
	for k := range alloc.Tasks {
		task := alloc.Tasks[k]
		if task.State != state {
			return false
		}
	}
	return true
}

// HostID finds a nomad Host that matches the provided hostname
func HostID(nomad *NomadServer, hostname *string) (*Host, error) {
	hosts, _, err := Hosts(nomad)
	if err != nil {
		return &Host{}, err
	}
	for _, host := range hosts {
		if *hostname == host.Name {
			return &host, nil
		}
	}
	buf := log("event", "node_not_found", "hostname", hostname)
	return &Host{}, errors.New(buf.String())

}

func (e *AllocNotFound) Error() string {
	return fmt.Sprintf("Unable to find '%v' job on '%v' host.", e.Jobname, e.Hostname)
}

func url(nomad *NomadServer) string {
	return fmt.Sprintf("http://%v:%v", nomad.Address, nomad.Port)
}

func decodeJSON(url string, target interface{}) (int, error) {
	r, err := retryablehttp.Get(url)
	if err != nil && r != nil && r.Body != nil {
		defer r.Body.Close()
		return r.StatusCode, err
	} else if err != nil {
		return http.StatusInternalServerError, err
	}
	return http.StatusOK, json.NewDecoder(r.Body).Decode(target)
}
<commit_msg>Added drew's changes to client<commit_after>package client

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"

	kitlog "github.com/go-kit/kit/log"
	retryablehttp "github.com/hashicorp/go-retryablehttp"
)

// NomadServer is connection parameters to a nomad server
type NomadServer struct {
	Address string
	Port    int
}

// Details represents a json object in a nomad job
type Details struct {
	Running int `json:"running"`
}

// Summary represents a json object in a nomad job
type Summary struct {
	Details *Details `json:"server"`
}

// JobSummary represents a json object in a nomad job
type JobSummary struct {
	JobID   string             `json:"JobID"`
	Summary map[string]Details `json:"Summary"`
}

// Job is a representation of nomad job
type Job struct {
	Name       string      `json:"Name"`
	Priority   int         `json:"Priority"`
	Status     string      `json:"status"`
	JobSummary *JobSummary `json:"JobSummary"`
}

// Task represents an allocated nomad task
type Task struct {
	State string `json:"State"`
}

// Alloc is a representation of a nomad allocation
type Alloc struct {
	ID            string          `json:"ID"`
	JobID         string          `json:"JobID"`
	NodeID        string          `json:"NodeID"`
	Name          string          `json:"Name"`
	ClientStatus  string          `json:"ClientStatus"`
	DesiredStatus string          `json:"DesiredStatus"`
	TaskGroup     string          `json:"TaskGroup"`
	Tasks         map[string]Task `json:"TaskStates"`
}

// Host is a representation of a nomad client node
type Host struct {
	ID        string         `json:"ID"`
	Name      string         `json:"Name"`
	Drain     bool           `json:"Drain"`
	Resources *NodeResources `json:"Resources,omitempty"`
}

// NodeResources represents the resources of a nomad node
type NodeResources struct {
	Networks *Networks `json:"Networks"`
}

// Networks represents the network available on a nomad node
type Networks struct {
	IP string `json:"IP"`
}

// AllocNotFound indicates a missing allocation for a Job
type AllocNotFound struct {
	Jobname  string
	Hostname string
}

func log(keyvals ...interface{}) *bytes.Buffer {
	w := new(bytes.Buffer)
	kitlog.NewJSONLogger(w).Log(keyvals)
	return w
}

// Jobs will parse the json representation from the nomad rest api
// /v1/jobs
func Jobs(nomad *NomadServer) ([]Job, int, error) {
	jobs := make([]Job, 0)
	status, err := decodeJSON(url(nomad)+"/v1/jobs", &jobs)
	return jobs, status, err
}

// FindJob will parse the json representation and find the supplied job name
func FindJob(nomad *NomadServer, name string) (*Job, error) {
	jobs, _, _ := Jobs(nomad)
	for _, job := range jobs {
		if job.Name == name {
			return &job, nil
		}
	}
	return &Job{}, errors.New("job not found")
}

// Hosts will parse the json representation from the nomad rest api
// /v1/nodes
func Hosts(nomad *NomadServer) ([]Host, int, error) {
	hosts := make([]Host, 0)
	status, err := decodeJSON(url(nomad)+"/v1/nodes", &hosts)
	return hosts, status, err
}

// Drain will inform nomad to add/remove all allocations from that host
// depending on the value of enable
// Returns the http status code or an error
func Drain(nomad *NomadServer, id string, enable bool) (int, error) {
	resp, err := retryablehttp.Post(url(nomad)+"/v1/node/"+id+"/drain?enable="+strconv.FormatBool(enable), "application/json", nil)
	if resp != nil && resp.Body != nil {
		defer resp.Body.Close()
		return resp.StatusCode, err
	}
	return http.StatusInternalServerError, err
}

// StopJob will send a delete request to a nomad server to stop the provided job
func StopJob(nomad *NomadServer, job *Job) (int, error) {
	url := fmt.Sprintf("http://%v:%v/v1/job/%v", nomad.Address, nomad.Port, job.Name)
	req, err := retryablehttp.NewRequest(http.MethodDelete, url, bytes.NewReader([]byte{}))
	if err != nil {
		buf := log("event", "stop_job_request_error", "jobname", job.Name, "error", err.Error())
		return http.StatusInternalServerError, errors.New(buf.String())
	}
	resp, err := retryablehttp.NewClient().Do(req)
	if err != nil {
		buf := log("event", "stop_job_client_error", "jobname", job.Name, "error", err.Error())
		return resp.StatusCode, errors.New(buf.String())
	}
	return resp.StatusCode, nil
}

// SubmitJob requests that nomade run a json launch file located at the supplied path
// Returns an http status code or error
func SubmitJob(nomad *NomadServer, launchFilePath string) (int, error) {
	file, err := ioutil.ReadFile(launchFilePath)
	if err != nil {
		return http.StatusBadRequest, err
	}
	resp, err := retryablehttp.Post(url(nomad)+"/v1/jobs", "application/json", bytes.NewReader(file))
	defer resp.Body.Close()
	return resp.StatusCode, err
}

// Allocs will parse the json representation from the nomad rest api
// /v1/allocations
func Allocs(nomad *NomadServer) []Alloc {
	allocs := make([]Alloc, 0)
	decodeJSON(url(nomad)+"/v1/allocations", &allocs)
	return allocs
}

// FindAlloc will search through the Allocs on a provided Host to look for the
// allocations that match the provided Job
func FindAlloc(nomad *NomadServer, job *Job, host *Host) (*Alloc, error) {
	allocs := Allocs(nomad)
	for _, alloc := range allocs {
		if alloc.NodeID == host.ID && strings.Contains(alloc.Name, job.Name) {
			// We may be looking at a stale allocation and a newer one exists
			if alloc.DesiredStatus == "stop" && len(allocs) > 1 {
				continue
			}
			return &alloc, nil
		}
	}
	return &Alloc{}, &AllocNotFound{Hostname: host.Name, Jobname: job.Name}
}

//Active determines whether a given host is active.  This is true when we find a worker
//on a given host that isn't in a "stop" desired status
func Active(nomad *NomadServer, job *Job, host *Host) bool {
	allocs := Allocs(nomad)
	for _, alloc := range allocs {
		if alloc.NodeID == host.ID && strings.Contains(alloc.Name, job.Name) {
			if alloc.DesiredStatus != "stop" && strings.Contains(alloc.Name, "worker") {
				return true
			}
		}
	}
	return false
}

//Services returns a list of services present on a given host
func Services(nomad *NomadServer, job *Job, host *Host) []string {
	services := make([]string, 0)
	allocs := Allocs(nomad)
	for _, alloc := range allocs {
		if alloc.NodeID == host.ID && strings.Contains(alloc.Name, job.Name) {
			if alloc.DesiredStatus != "stop" {
				services = append(services, alloc.TaskGroup)
			}
		}
	}
	return services
}

// CheckTaskStates will return whether all of the tasks for an Alloc are
// in the state provided
func (alloc *Alloc) CheckTaskStates(state string) bool {
	for k := range alloc.Tasks {
		task := alloc.Tasks[k]
		if task.State != state {
			return false
		}
	}
	return true
}

// HostID finds a nomad Host that matches the provided hostname
func HostID(nomad *NomadServer, hostname *string) (*Host, error) {
	hosts, _, err := Hosts(nomad)
	if err != nil {
		return &Host{}, err
	}
	for _, host := range hosts {
		if *hostname == host.Name {
			return &host, nil
		}
	}
	buf := log("event", "node_not_found", "hostname", hostname)
	return &Host{}, errors.New(buf.String())

}

func (e *AllocNotFound) Error() string {
	return fmt.Sprintf("Unable to find '%v' job on '%v' host.", e.Jobname, e.Hostname)
}

func url(nomad *NomadServer) string {
	return fmt.Sprintf("http://%v:%v", nomad.Address, nomad.Port)
}

func decodeJSON(url string, target interface{}) (int, error) {
	r, err := retryablehttp.Get(url)
	if err != nil && r != nil && r.Body != nil {
		defer r.Body.Close()
		return r.StatusCode, err
	} else if err != nil {
		return http.StatusInternalServerError, err
	}
	return http.StatusOK, json.NewDecoder(r.Body).Decode(target)
}
<|endoftext|>


<commit_before>package client

import (
	"errors"
	"strings"

	imap "github.com/emersion/imap/common"
	"github.com/emersion/imap/commands"
	"github.com/emersion/imap/responses"
)

// Requests a checkpoint of the currently selected mailbox. A checkpoint refers
// to any implementation-dependent housekeeping associated with the mailbox that
// is not normally executed as part of each command.
func (c *Client) Check() (err error) {
	if c.State != imap.SelectedState {
		err = errors.New("No mailbox selected")
		return
	}

	cmd := &commands.Check{}

	status, err := c.execute(cmd, nil)
	if err != nil {
		return
	}

	err = status.Err()
	return
}

// Permanently removes all messages that have the \Deleted flag set from the
// currently selected mailbox, and returns to the authenticated state from the
// selected state.
func (c *Client) Close() (err error) {
	if c.State != imap.SelectedState {
		err = errors.New("No mailbox selected")
		return
	}

	cmd := &commands.Close{}

	status, err := c.execute(cmd, nil)
	if err != nil {
		return
	}

	err = status.Err()
	if err != nil {
		return
	}

	c.State = imap.AuthenticatedState
	c.Mailbox = nil
	return
}

// Permanently removes all messages that have the \Deleted flag set from the
// currently selected mailbox.
// If ch is not nil, sends sequence IDs of each deleted message to this channel.
func (c *Client) Expunge(ch chan<- uint32) (err error) {
	defer close(ch)

	if c.State != imap.SelectedState {
		err = errors.New("No mailbox selected")
		return
	}

	cmd := &commands.Expunge{}

	var res *responses.Expunge
	if ch != nil {
		res = &responses.Expunge{SeqIds: ch}
	}

	status, err := c.execute(cmd, res)
	if err != nil {
		return
	}

	err = status.Err()
	return
}

func (c *Client) search(uid bool, criteria *imap.SearchCriteria) (ids []uint32, err error) {
	if c.State != imap.SelectedState {
		err = errors.New("No mailbox selected")
		return
	}

	var cmd imap.Commander
	cmd = &commands.Search{
		Charset: "UTF-8",
		Criteria: criteria,
	}
	if uid {
		cmd = &commands.Uid{Cmd: cmd}
	}

	res := &responses.Search{}

	status, err := c.execute(cmd, res)
	if err != nil {
		return
	}

	err = status.Err()
	ids = res.Ids
	return
}

// Searches the mailbox for messages that match the given searching criteria.
// Searching criteria consist of one or more search keys. The response contains
// a list of message sequence IDs corresponding to those messages that match the
// searching criteria.
// When multiple keys are specified, the result is the intersection (AND
// function) of all the messages that match those keys.
// Criteria must be UTF-8 encoded.
// See RFC 3501 section 6.4.4 for a list of searching criteria.
func (c *Client) Search(criteria *imap.SearchCriteria) (seqids []uint32, err error) {
	return c.search(false, criteria)
}

// Identical to Search, but unique identifiers are returned instead of message
// sequence numbers.
func (c *Client) UidSearch(criteria *imap.SearchCriteria) (uids []uint32, err error) {
	return c.search(true, criteria)
}

func (c *Client) fetch(uid bool, seqset *imap.SeqSet, items []string, ch chan *imap.Message) (err error) {
	defer close(ch)

	if c.State != imap.SelectedState {
		err = errors.New("No mailbox selected")
		return
	}

	var cmd imap.Commander
	cmd = &commands.Fetch{
		SeqSet: seqset,
		Items: items,
	}
	if uid {
		cmd = &commands.Uid{Cmd: cmd}
	}

	res := &responses.Fetch{Messages: ch}

	status, err := c.execute(cmd, res)
	if err != nil {
		return
	}

	err = status.Err()
	return
}

// Retrieves data associated with a message in the mailbox.
// See RFC 3501 section 6.4.5 for a list of items that can be requested.
func (c *Client) Fetch(seqset *imap.SeqSet, items []string, ch chan *imap.Message) (err error) {
	return c.fetch(false, seqset, items, ch)
}

// Identical to Fetch, but seqset is interpreted as containing unique
// identifiers instead of message sequence numbers.
func (c *Client) UidFetch(seqset *imap.SeqSet, items []string, ch chan *imap.Message) (err error) {
	return c.fetch(true, seqset, items, ch)
}

func (c *Client) store(uid bool, seqset *imap.SeqSet, item string, value interface{}, ch chan *imap.Message) (err error) {
	defer (func () {
		if ch != nil {
			close(ch)
		}
	})()

	if c.State != imap.SelectedState {
		err = errors.New("No mailbox selected")
		return
	}

	// If ch is nil, the updated values are data which will be lost, so don't
	// retrieve it.
	if ch == nil && !strings.HasSuffix(item, imap.SilentOp) {
		item += imap.SilentOp
	}

	var cmd imap.Commander
	cmd = &commands.Store{
		SeqSet: seqset,
		Item: item,
		Value: value,
	}
	if uid {
		cmd = &commands.Uid{Cmd: cmd}
	}

	var res *responses.Fetch
	if ch != nil {
		res = &responses.Fetch{Messages: ch}
	}

	status, err := c.execute(cmd, res)
	if err != nil {
		return
	}

	err = status.Err()
	return
}

// Alters data associated with a message in the mailbox. If ch is not nil, the
// updated value of the data will be sent to this channel.
// See RFC 3501 section 6.4.6 for a list of items that can be updated.
func (c *Client) Store(seqset *imap.SeqSet, item string, value interface{}, ch chan *imap.Message) (err error) {
	return c.store(false, seqset, item, value, ch)
}

// Identical to Store, but seqset is interpreted as containing unique
// identifiers instead of message sequence numbers.
func (c *Client) UidStore(seqset *imap.SeqSet, item string, value interface{}, ch chan *imap.Message) (err error) {
	return c.store(true, seqset, item, value, ch)
}

func (c *Client) copy(uid bool, seqset *imap.SeqSet, dest string) (err error) {
	if c.State != imap.SelectedState {
		err = errors.New("No mailbox selected")
		return
	}

	var cmd imap.Commander
	cmd = &commands.Copy{
		SeqSet: seqset,
		Mailbox: dest,
	}
	if uid {
		cmd = &commands.Uid{Cmd: cmd}
	}

	status, err := c.execute(cmd, nil)
	if err != nil {
		return
	}

	err = status.Err()
	return
}

// Copies the specified message(s) to the end of the specified destination
// mailbox.
func (c *Client) Copy(seqset *imap.SeqSet, dest string) (err error) {
	return c.copy(false, seqset, dest)
}

// Identical to Copy, but seqset is interpreted as containing unique
// identifiers instead of message sequence numbers.
func (c *Client) UidCopy(seqset *imap.SeqSet, dest string) (err error) {
	return c.copy(true, seqset, dest)
}
<commit_msg>Fixes client build<commit_after>package client

import (
	"errors"
	"strings"

	imap "github.com/emersion/imap/common"
	"github.com/emersion/imap/commands"
	"github.com/emersion/imap/responses"
)

// Requests a checkpoint of the currently selected mailbox. A checkpoint refers
// to any implementation-dependent housekeeping associated with the mailbox that
// is not normally executed as part of each command.
func (c *Client) Check() (err error) {
	if c.State != imap.SelectedState {
		err = errors.New("No mailbox selected")
		return
	}

	cmd := &commands.Check{}

	status, err := c.execute(cmd, nil)
	if err != nil {
		return
	}

	err = status.Err()
	return
}

// Permanently removes all messages that have the \Deleted flag set from the
// currently selected mailbox, and returns to the authenticated state from the
// selected state.
func (c *Client) Close() (err error) {
	if c.State != imap.SelectedState {
		err = errors.New("No mailbox selected")
		return
	}

	cmd := &commands.Close{}

	status, err := c.execute(cmd, nil)
	if err != nil {
		return
	}

	err = status.Err()
	if err != nil {
		return
	}

	c.State = imap.AuthenticatedState
	c.Mailbox = nil
	return
}

// Permanently removes all messages that have the \Deleted flag set from the
// currently selected mailbox.
// If ch is not nil, sends sequence IDs of each deleted message to this channel.
func (c *Client) Expunge(ch chan uint32) (err error) {
	defer close(ch)

	if c.State != imap.SelectedState {
		err = errors.New("No mailbox selected")
		return
	}

	cmd := &commands.Expunge{}

	var res *responses.Expunge
	if ch != nil {
		res = &responses.Expunge{SeqIds: ch}
	}

	status, err := c.execute(cmd, res)
	if err != nil {
		return
	}

	err = status.Err()
	return
}

func (c *Client) search(uid bool, criteria *imap.SearchCriteria) (ids []uint32, err error) {
	if c.State != imap.SelectedState {
		err = errors.New("No mailbox selected")
		return
	}

	var cmd imap.Commander
	cmd = &commands.Search{
		Charset: "UTF-8",
		Criteria: criteria,
	}
	if uid {
		cmd = &commands.Uid{Cmd: cmd}
	}

	res := &responses.Search{}

	status, err := c.execute(cmd, res)
	if err != nil {
		return
	}

	err = status.Err()
	ids = res.Ids
	return
}

// Searches the mailbox for messages that match the given searching criteria.
// Searching criteria consist of one or more search keys. The response contains
// a list of message sequence IDs corresponding to those messages that match the
// searching criteria.
// When multiple keys are specified, the result is the intersection (AND
// function) of all the messages that match those keys.
// Criteria must be UTF-8 encoded.
// See RFC 3501 section 6.4.4 for a list of searching criteria.
func (c *Client) Search(criteria *imap.SearchCriteria) (seqids []uint32, err error) {
	return c.search(false, criteria)
}

// Identical to Search, but unique identifiers are returned instead of message
// sequence numbers.
func (c *Client) UidSearch(criteria *imap.SearchCriteria) (uids []uint32, err error) {
	return c.search(true, criteria)
}

func (c *Client) fetch(uid bool, seqset *imap.SeqSet, items []string, ch chan *imap.Message) (err error) {
	defer close(ch)

	if c.State != imap.SelectedState {
		err = errors.New("No mailbox selected")
		return
	}

	var cmd imap.Commander
	cmd = &commands.Fetch{
		SeqSet: seqset,
		Items: items,
	}
	if uid {
		cmd = &commands.Uid{Cmd: cmd}
	}

	res := &responses.Fetch{Messages: ch}

	status, err := c.execute(cmd, res)
	if err != nil {
		return
	}

	err = status.Err()
	return
}

// Retrieves data associated with a message in the mailbox.
// See RFC 3501 section 6.4.5 for a list of items that can be requested.
func (c *Client) Fetch(seqset *imap.SeqSet, items []string, ch chan *imap.Message) (err error) {
	return c.fetch(false, seqset, items, ch)
}

// Identical to Fetch, but seqset is interpreted as containing unique
// identifiers instead of message sequence numbers.
func (c *Client) UidFetch(seqset *imap.SeqSet, items []string, ch chan *imap.Message) (err error) {
	return c.fetch(true, seqset, items, ch)
}

func (c *Client) store(uid bool, seqset *imap.SeqSet, item string, value interface{}, ch chan *imap.Message) (err error) {
	defer (func () {
		if ch != nil {
			close(ch)
		}
	})()

	if c.State != imap.SelectedState {
		err = errors.New("No mailbox selected")
		return
	}

	// If ch is nil, the updated values are data which will be lost, so don't
	// retrieve it.
	if ch == nil && !strings.HasSuffix(item, imap.SilentOp) {
		item += imap.SilentOp
	}

	var cmd imap.Commander
	cmd = &commands.Store{
		SeqSet: seqset,
		Item: item,
		Value: value,
	}
	if uid {
		cmd = &commands.Uid{Cmd: cmd}
	}

	var res *responses.Fetch
	if ch != nil {
		res = &responses.Fetch{Messages: ch}
	}

	status, err := c.execute(cmd, res)
	if err != nil {
		return
	}

	err = status.Err()
	return
}

// Alters data associated with a message in the mailbox. If ch is not nil, the
// updated value of the data will be sent to this channel.
// See RFC 3501 section 6.4.6 for a list of items that can be updated.
func (c *Client) Store(seqset *imap.SeqSet, item string, value interface{}, ch chan *imap.Message) (err error) {
	return c.store(false, seqset, item, value, ch)
}

// Identical to Store, but seqset is interpreted as containing unique
// identifiers instead of message sequence numbers.
func (c *Client) UidStore(seqset *imap.SeqSet, item string, value interface{}, ch chan *imap.Message) (err error) {
	return c.store(true, seqset, item, value, ch)
}

func (c *Client) copy(uid bool, seqset *imap.SeqSet, dest string) (err error) {
	if c.State != imap.SelectedState {
		err = errors.New("No mailbox selected")
		return
	}

	var cmd imap.Commander
	cmd = &commands.Copy{
		SeqSet: seqset,
		Mailbox: dest,
	}
	if uid {
		cmd = &commands.Uid{Cmd: cmd}
	}

	status, err := c.execute(cmd, nil)
	if err != nil {
		return
	}

	err = status.Err()
	return
}

// Copies the specified message(s) to the end of the specified destination
// mailbox.
func (c *Client) Copy(seqset *imap.SeqSet, dest string) (err error) {
	return c.copy(false, seqset, dest)
}

// Identical to Copy, but seqset is interpreted as containing unique
// identifiers instead of message sequence numbers.
func (c *Client) UidCopy(seqset *imap.SeqSet, dest string) (err error) {
	return c.copy(true, seqset, dest)
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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
//
// https://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.

/*
This package is for loading different mailing list data types into Cloud Storage.
*/

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/google/project-OCEAN/1-raw-data/gcs"
	"github.com/google/project-OCEAN/1-raw-data/mailinglists/googlegroups"
	"github.com/google/project-OCEAN/1-raw-data/mailinglists/mailman"
	"github.com/google/project-OCEAN/1-raw-data/mailinglists/pipermail"
	"github.com/google/project-OCEAN/1-raw-data/utils"
)

var (
	//Below variables used if build run
	buildListRun = flag.Bool("build-list-run", false, "Use flag to run build list run vs manual command line run.")
	allListRun   = flag.Bool("all-list-run", false, "Use flag to get variables from command-line or do a full mailing list run or to do simple build test run of one mailing list.")
	allDateRun   = flag.Bool("all-date-run", false, "Use flag to get variables from command-line or do a full run")
	projectID    = flag.String("project-id", "", "GCP Project id.")

	//Below variables used if manual run
	bucketName   = flag.String("bucket-name", "mailinglists", "Bucket name to store files.")
	subDirectory = flag.String("subdirectory", "", "Subdirectory to store files. Enter 1 or more and use spaces to identify. CAUTION also enter the groupNames to load to in the same order.")
	mailingList  = flag.String("mailinglist", "", "Choose which mailing list to process either pipermail (default), mailman, googlegroups")
	groupNames   = flag.String("groupname", "", "Mailing list group name. Enter 1 or more and use spaces to identify. CAUTION also enter the buckets to load to in the same order.")
	startDate    = flag.String("start-date", "", "Start date in format of year-month-date and 4dig-2dig-2dig.")
	endDate      = flag.String("end-date", "", "End date in format of year-month-date and 4dig-2dig-2dig.")
	workerNum    = flag.Int("workers", 1, "Number of workers to use for goroutines.")
	subNames     []string
)

func getData(ctx context.Context, storage gcs.Connection, httpToDom utils.HttpDomResponse, workerNum, numMonths int, mailingList, groupName, startDateString, endDateString string) {
	switch mailingList {
	//TODO add start and end dates to pipermail
	case "pipermail":
		if err := pipermail.GetPipermailData(ctx, storage, groupName, startDateString, endDateString, httpToDom); err != nil {
			log.Fatalf("Mailman load failed: %v", err)
		}
	case "mailman":
		if err := mailman.GetMailmanData(ctx, storage, groupName, startDateString, endDateString, numMonths); err != nil {
			log.Fatalf("Mailman load failed: %v", err)
		}
		//TODO add start and end dates to google groups
	case "gg":
		if err := googlegroups.GetGoogleGroupsData(ctx, "", groupName, startDateString, endDateString, storage, workerNum); err != nil {
			log.Fatalf("GoogleGroups load failed: %v", err)
		}
	default:
		log.Fatalf("Mailing list %v is not an option. Change the option submitted.", mailingList)
	}
}

func main() {
	var err error
	numMonths := 1
	httpToDom := utils.DomResponse
	flag.Parse()
	fmt.Printf(*projectID)

	//Setup Storage connection
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	storageConn := gcs.StorageConnection{
		ProjectID:  *projectID,
		BucketName: *bucketName,
	}
	if err = storageConn.ConnectClient(ctx); err != nil {
		log.Fatalf("Connect GCS failes: %v", err)
	}

	if *buildListRun {
		//Build run to load mailing list data
		now := time.Now()
		//Set variables in build that aren't coming in on command line
		groupName := ""

		// Setup bucket connection whether new or not
		storageConn.BucketName = *bucketName
		if err := storageConn.CreateBucket(ctx); err != nil {
			log.Fatalf("Create GCS Bucket failed: %v", err)
		}

		// Run Build to load all mailing lists
		if *allListRun {
			log.Printf("Build all lists ")
			mailListMap := map[string]string{"gg-angular": "2009-09-01", "gg-golang-announce": "2011-05-01", "gg-golang-checkins": "2009-11-01", " gg-golang-codereviews": "2013-12-01", "gg-golang-dev": "2009-11-01", "gg-golang-nuts": "2009-11-01", "gg-nodejs": "2009-06-01", "mailman-python-announce-list": "1999-04-01", "mailman-python-dev": "1999-04-01", "mailman-python-ideas": "2006-12-01", "pipermail-python-announce-list": "1999-04-01", "pipermail-python-dev": "1995-03-01", "pipermail-python-ideas": "2006-12-01", "pipermail-python-list": "1999-02-01"}

			for subName, origStartDate := range mailListMap {
				startDateResult, endDateResult := "", ""
				storageConn.SubDirectory = subName
				*mailingList = strings.SplitN(subName, "-", 2)[0]
				groupName = strings.SplitN(subName, "-", 2)[1]

				if *allDateRun {
					//Load all months
					log.Printf("All Date Cloud Run")
					//Set start and end dates with first mailing list date and current end date
					if startDateResult, endDateResult, err = utils.FixDate(origStartDate, *endDate); err != nil {
						log.Fatalf("Date error: %v", err)
					}
				} else {
					//Set start and end dates split by one month
					if startDateResult, endDateResult, err = utils.SplitDatesByMonth(*startDate, *endDate, 1); err != nil {
						log.Fatalf("Date error: %v", err)
					}
					log.Printf("One Month Run All MailingLists")
					startDateResult = now.AddDate(0, -1, 0).Format("2006-01-02")
					endDateResult = now.Format("2006-01-02")
				}

				//Get mailing list data and store
				getData(ctx, &storageConn, httpToDom, *workerNum, numMonths, *mailingList, groupName, startDateResult, endDateResult)
			}
		} else {
			log.Printf("Build test run with mailman")
			storageConn.SubDirectory = "mailman-python-announce-list"
			groupName = "python-announce-list"
			*startDate = now.AddDate(0, -1, 0).Format("2006-01-02")
			*endDate = now.AddDate(0, -1, 1).Format("2006-01-02")

			if err := mailman.GetMailmanData(ctx, &storageConn, groupName, *startDate, *endDate, numMonths); err != nil {
				log.Fatalf("Mailman test build load failed: %v", err)
			}
		}
	} else {

		//Check and create bucket if needed
		if err := storageConn.CreateBucket(ctx); err != nil {
			log.Fatalf("Create GCS Bucket failed: %v", err)
		}

		if *subDirectory != "" {
			subNames = strings.Split(*subDirectory, " ")
		}

		for idx, groupName := range strings.Split(*groupNames, " ") {
			//Apply sub directory name to storageConn if it exists
			if *subDirectory != "" {
				storageConn.SubDirectory = subNames[idx]
			}

			//Get mailing list data and store
			getData(ctx, &storageConn, httpToDom, *workerNum, numMonths, *mailingList, groupName, *startDate, *endDate)
		}
	}
}
<commit_msg>Cleaning up storage code. Bucket will be created in either branch, so this reduces code.<commit_after>// Copyright 2020 Google LLC
//
// 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
//
// https://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.

/*
This package is for loading different mailing list data types into Cloud Storage.
*/

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/google/project-OCEAN/1-raw-data/gcs"
	"github.com/google/project-OCEAN/1-raw-data/mailinglists/googlegroups"
	"github.com/google/project-OCEAN/1-raw-data/mailinglists/mailman"
	"github.com/google/project-OCEAN/1-raw-data/mailinglists/pipermail"
	"github.com/google/project-OCEAN/1-raw-data/utils"
)

var (
	//Below variables used if build run
	buildListRun = flag.Bool("build-list-run", false, "Use flag to run build list run vs manual command line run.")
	allListRun   = flag.Bool("all-list-run", false, "Use flag to get variables from command-line or do a full mailing list run or to do simple build test run of one mailing list.")
	allDateRun   = flag.Bool("all-date-run", false, "Use flag to get variables from command-line or do a full run")
	projectID    = flag.String("project-id", "", "GCP Project id.")

	//Below variables used if manual run
	bucketName   = flag.String("bucket-name", "mailinglists", "Bucket name to store files.")
	subDirectory = flag.String("subdirectory", "", "Subdirectory to store files. Enter 1 or more and use spaces to identify. CAUTION also enter the groupNames to load to in the same order.")
	mailingList  = flag.String("mailinglist", "", "Choose which mailing list to process either pipermail (default), mailman, googlegroups")
	groupNames   = flag.String("groupname", "", "Mailing list group name. Enter 1 or more and use spaces to identify. CAUTION also enter the buckets to load to in the same order.")
	startDate    = flag.String("start-date", "", "Start date in format of year-month-date and 4dig-2dig-2dig.")
	endDate      = flag.String("end-date", "", "End date in format of year-month-date and 4dig-2dig-2dig.")
	workerNum    = flag.Int("workers", 1, "Number of workers to use for goroutines.")
	subNames     []string
)

func getData(ctx context.Context, storage gcs.Connection, httpToDom utils.HttpDomResponse, workerNum, numMonths int, mailingList, groupName, startDateString, endDateString string) {
	switch mailingList {
	//TODO add start and end dates to pipermail
	case "pipermail":
		if err := pipermail.GetPipermailData(ctx, storage, groupName, startDateString, endDateString, httpToDom); err != nil {
			log.Fatalf("Mailman load failed: %v", err)
		}
	case "mailman":
		if err := mailman.GetMailmanData(ctx, storage, groupName, startDateString, endDateString, numMonths); err != nil {
			log.Fatalf("Mailman load failed: %v", err)
		}
		//TODO add start and end dates to google groups
	case "gg":
		if err := googlegroups.GetGoogleGroupsData(ctx, "", groupName, startDateString, endDateString, storage, workerNum); err != nil {
			log.Fatalf("GoogleGroups load failed: %v", err)
		}
	default:
		log.Fatalf("Mailing list %v is not an option. Change the option submitted.", mailingList)
	}
}

func main() {
	var err error
	numMonths := 1
	httpToDom := utils.DomResponse
	flag.Parse()
	fmt.Printf(*projectID)

	//Setup Storage connection
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	storageConn := gcs.StorageConnection{
		ProjectID:  *projectID,
		BucketName: *bucketName,
	}
	if err := storageConn.ConnectClient(ctx); err != nil {
		log.Fatalf("Connect GCS failes: %v", err)
	}
	if err := storageConn.CreateBucket(ctx); err != nil {
		log.Fatalf("Create GCS Bucket failed: %v", err)
	}

	if *buildListRun {
		//Build run to load mailing list data
		now := time.Now()
		//Set variables in build that aren't coming in on command line
		groupName := ""

		// Run Build to load all mailing lists
		if *allListRun {
			log.Printf("Build all lists ")
			mailListMap := map[string]string{"gg-angular": "2009-09-01", "gg-golang-announce": "2011-05-01", "gg-golang-checkins": "2009-11-01", " gg-golang-codereviews": "2013-12-01", "gg-golang-dev": "2009-11-01", "gg-golang-nuts": "2009-11-01", "gg-nodejs": "2009-06-01", "mailman-python-announce-list": "1999-04-01", "mailman-python-dev": "1999-04-01", "mailman-python-ideas": "2006-12-01", "pipermail-python-announce-list": "1999-04-01", "pipermail-python-dev": "1995-03-01", "pipermail-python-ideas": "2006-12-01", "pipermail-python-list": "1999-02-01"}

			for subName, origStartDate := range mailListMap {
				startDateResult, endDateResult := "", ""
				storageConn.SubDirectory = subName
				*mailingList = strings.SplitN(subName, "-", 2)[0]
				groupName = strings.SplitN(subName, "-", 2)[1]

				if *allDateRun {
					//Load all months
					log.Printf("All Date Cloud Run")
					//Set start and end dates with first mailing list date and current end date
					if startDateResult, endDateResult, err = utils.FixDate(origStartDate, *endDate); err != nil {
						log.Fatalf("Date error: %v", err)
					}
				} else {
					//Set start and end dates split by one month
					if startDateResult, endDateResult, err = utils.SplitDatesByMonth(*startDate, *endDate, 1); err != nil {
						log.Fatalf("Date error: %v", err)
					}
					log.Printf("One Month Run All MailingLists")
					startDateResult = now.AddDate(0, -1, 0).Format("2006-01-02")
					endDateResult = now.Format("2006-01-02")
				}

				//Get mailing list data and store
				getData(ctx, &storageConn, httpToDom, *workerNum, numMonths, *mailingList, groupName, startDateResult, endDateResult)
			}
		} else {
			log.Printf("Build test run with mailman")
			storageConn.SubDirectory = "mailman-python-announce-list"
			groupName = "python-announce-list"
			*startDate = now.AddDate(0, -1, 0).Format("2006-01-02")
			*endDate = now.AddDate(0, -1, 1).Format("2006-01-02")

			if err := mailman.GetMailmanData(ctx, &storageConn, groupName, *startDate, *endDate, numMonths); err != nil {
				log.Fatalf("Mailman test build load failed: %v", err)
			}
		}
	} else {

		if *subDirectory != "" {
			subNames = strings.Split(*subDirectory, " ")
		}

		for idx, groupName := range strings.Split(*groupNames, " ") {
			//Apply sub directory name to storageConn if it exists
			if *subDirectory != "" {
				storageConn.SubDirectory = subNames[idx]
			}

			//Get mailing list data and store
			getData(ctx, &storageConn, httpToDom, *workerNum, numMonths, *mailingList, groupName, *startDate, *endDate)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

// Git remote helper for the Keybase file system.

package main

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

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/kbfs/env"
	"github.com/keybase/kbfs/kbfsgit"
	"github.com/keybase/kbfs/libfs"
	"github.com/keybase/kbfs/libgit"
	"github.com/keybase/kbfs/libkbfs"
	"github.com/pkg/errors"
)

var version = flag.Bool("version", false, "Print version")

const usageFormatStr = `Usage:
  git-remote-keybase -version

To run against remote KBFS servers:
  git-remote-keybase %s <remote> [keybase://<repo>]

To run in a local testing environment:
  git-remote-keybase %s <remote> [keybase://<repo>]

Defaults:
%s
`

func getUsageString(ctx libkbfs.Context) string {
	remoteUsageStr := libkbfs.GetRemoteUsageString()
	localUsageStr := libkbfs.GetLocalUsageString()
	defaultUsageStr := libkbfs.GetDefaultsUsageString(ctx)
	return fmt.Sprintf(
		usageFormatStr, remoteUsageStr, localUsageStr, defaultUsageStr)
}

func getLocalGitDir() (gitDir string, err error) {
	gitDir = os.Getenv("GIT_DIR")
	fi, err := os.Stat(gitDir)
	if err != nil {
		return "", err
	}
	if !fi.IsDir() {
		return "", errors.Errorf("GIT_DIR=%s, but is not a dir", gitDir)
	}
	// On Windows, git annoyingly puts normal slashes in the
	// environment variable.
	return filepath.FromSlash(gitDir), nil
}

func start() (startErr *libfs.Error) {
	kbCtx := env.NewContext()

	switch kbCtx.GetRunMode() {
	case libkb.StagingRunMode:
		fmt.Fprintf(os.Stderr, "Running in staging mode\n")
	case libkb.DevelRunMode:
		fallthrough
	case libkb.NoRunMode:
		fmt.Fprintf(os.Stderr, "Running in devel mode\n")
	}

	defaultParams, storageRoot, err := libgit.Params(kbCtx, kbCtx.GetDataDir())
	if err != nil {
		return libfs.InitError(err.Error())
	}
	defer func() {
		rmErr := os.RemoveAll(storageRoot)
		if rmErr != nil {
			fmt.Fprintf(os.Stderr,
				"Error cleaning storage dir %s: %+v\n", storageRoot, rmErr)
		}
	}()
	defaultLogPath := filepath.Join(kbCtx.GetLogDir(), libkb.GitLogFileName)

	// Duplicate the stderr fd, so that when the logger closes it when
	// redirecting log messages to a file, we will still be able to
	// write status updates back to the git process.
	stderrFile, err := dupStderr()
	if err != nil {
		return libfs.InitError(err.Error())
	}
	defer stderrFile.Close()

	defer func() {
		// Now that the stderr has been duplicated, print all errors
		// to the duplicate as well as to the new `os.Stderr` down in
		// `main()`, so that the error shows up both in the log and to
		// the user.
		if startErr != nil {
			fmt.Fprintf(stderrFile, "git-remote-keybase error: (%d) %s\n",
				startErr.Code, startErr.Message)
		}
	}()

	kbfsParams := libkbfs.AddFlagsWithDefaults(
		flag.CommandLine, defaultParams, defaultLogPath)
	flag.Parse()

	if *version {
		fmt.Printf("%s\n", libkbfs.VersionString())
		return nil
	}

	if len(flag.Args()) < 1 {
		fmt.Print(getUsageString(kbCtx))
		return libfs.InitError("no remote repo specified")
	}

	remote := flag.Arg(0)
	var repo string
	if len(flag.Args()) > 1 {
		repo = flag.Arg(1)
	}

	if len(flag.Args()) > 2 {
		fmt.Print(getUsageString(kbCtx))
		return libfs.InitError("extra arguments specified (flags go before the first argument)")
	}

	gitDir, err := getLocalGitDir()
	if err != nil {
		return libfs.InitError(err.Error())
	}

	options := kbfsgit.StartOptions{
		KbfsParams: *kbfsParams,
		Remote:     remote,
		Repo:       repo,
		GitDir:     gitDir,
	}

	ctx := context.Background()
	return kbfsgit.Start(
		ctx, options, kbCtx, defaultLogPath, os.Stdin, os.Stdout, stderrFile)
}

func main() {
	runMode := os.Getenv("KEYBASE_RUN_MODE")
	if len(runMode) == 0 {
		// Default to prod.
		os.Setenv("KEYBASE_RUN_MODE", "prod")
	}

	err := start()
	if err != nil {
		fmt.Fprintf(os.Stderr, "git-remote-keybase error: (%d) %s\n",
			err.Code, err.Message)
		os.Exit(err.Code)
	}
	os.Exit(0)
}
<commit_msg>git-remote-keybase: panic on NoRunMode<commit_after>// Copyright 2017 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

// Git remote helper for the Keybase file system.

package main

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

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/kbfs/env"
	"github.com/keybase/kbfs/kbfsgit"
	"github.com/keybase/kbfs/libfs"
	"github.com/keybase/kbfs/libgit"
	"github.com/keybase/kbfs/libkbfs"
	"github.com/pkg/errors"
)

var version = flag.Bool("version", false, "Print version")

const usageFormatStr = `Usage:
  git-remote-keybase -version

To run against remote KBFS servers:
  git-remote-keybase %s <remote> [keybase://<repo>]

To run in a local testing environment:
  git-remote-keybase %s <remote> [keybase://<repo>]

Defaults:
%s
`

func getUsageString(ctx libkbfs.Context) string {
	remoteUsageStr := libkbfs.GetRemoteUsageString()
	localUsageStr := libkbfs.GetLocalUsageString()
	defaultUsageStr := libkbfs.GetDefaultsUsageString(ctx)
	return fmt.Sprintf(
		usageFormatStr, remoteUsageStr, localUsageStr, defaultUsageStr)
}

func getLocalGitDir() (gitDir string, err error) {
	gitDir = os.Getenv("GIT_DIR")
	fi, err := os.Stat(gitDir)
	if err != nil {
		return "", err
	}
	if !fi.IsDir() {
		return "", errors.Errorf("GIT_DIR=%s, but is not a dir", gitDir)
	}
	// On Windows, git annoyingly puts normal slashes in the
	// environment variable.
	return filepath.FromSlash(gitDir), nil
}

func start() (startErr *libfs.Error) {
	kbCtx := env.NewContext()

	switch kbCtx.GetRunMode() {
	case libkb.ProductionRunMode:
	case libkb.StagingRunMode:
		fmt.Fprintf(os.Stderr, "Running in staging mode\n")
	case libkb.DevelRunMode:
		fmt.Fprintf(os.Stderr, "Running in devel mode\n")
	default:
		panic(fmt.Sprintf("Unexpected run mode: %s", kbCtx.GetRunMode()))
	}

	defaultParams, storageRoot, err := libgit.Params(kbCtx, kbCtx.GetDataDir())
	if err != nil {
		return libfs.InitError(err.Error())
	}
	defer func() {
		rmErr := os.RemoveAll(storageRoot)
		if rmErr != nil {
			fmt.Fprintf(os.Stderr,
				"Error cleaning storage dir %s: %+v\n", storageRoot, rmErr)
		}
	}()
	defaultLogPath := filepath.Join(kbCtx.GetLogDir(), libkb.GitLogFileName)

	// Duplicate the stderr fd, so that when the logger closes it when
	// redirecting log messages to a file, we will still be able to
	// write status updates back to the git process.
	stderrFile, err := dupStderr()
	if err != nil {
		return libfs.InitError(err.Error())
	}
	defer stderrFile.Close()

	defer func() {
		// Now that the stderr has been duplicated, print all errors
		// to the duplicate as well as to the new `os.Stderr` down in
		// `main()`, so that the error shows up both in the log and to
		// the user.
		if startErr != nil {
			fmt.Fprintf(stderrFile, "git-remote-keybase error: (%d) %s\n",
				startErr.Code, startErr.Message)
		}
	}()

	kbfsParams := libkbfs.AddFlagsWithDefaults(
		flag.CommandLine, defaultParams, defaultLogPath)
	flag.Parse()

	if *version {
		fmt.Printf("%s\n", libkbfs.VersionString())
		return nil
	}

	if len(flag.Args()) < 1 {
		fmt.Print(getUsageString(kbCtx))
		return libfs.InitError("no remote repo specified")
	}

	remote := flag.Arg(0)
	var repo string
	if len(flag.Args()) > 1 {
		repo = flag.Arg(1)
	}

	if len(flag.Args()) > 2 {
		fmt.Print(getUsageString(kbCtx))
		return libfs.InitError("extra arguments specified (flags go before the first argument)")
	}

	gitDir, err := getLocalGitDir()
	if err != nil {
		return libfs.InitError(err.Error())
	}

	options := kbfsgit.StartOptions{
		KbfsParams: *kbfsParams,
		Remote:     remote,
		Repo:       repo,
		GitDir:     gitDir,
	}

	ctx := context.Background()
	return kbfsgit.Start(
		ctx, options, kbCtx, defaultLogPath, os.Stdin, os.Stdout, stderrFile)
}

func main() {
	runMode := os.Getenv("KEYBASE_RUN_MODE")
	if len(runMode) == 0 {
		// Default to prod.
		os.Setenv("KEYBASE_RUN_MODE", "prod")
	}

	err := start()
	if err != nil {
		fmt.Fprintf(os.Stderr, "git-remote-keybase error: (%d) %s\n",
			err.Code, err.Message)
		os.Exit(err.Code)
	}
	os.Exit(0)
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 main

import (
	"fmt"
	"github.com/jacobsa/aws/exp/sdb"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"math/rand"
	"sync"
)

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type integrationTest struct {
	db sdb.SimpleDB
}

func (t *integrationTest) SetUp(i *TestInfo) {
	var err error

	// Open a connection.
	t.db, err = sdb.NewSimpleDB(g_region, g_accessKey)
	AssertEq(nil, err)
}

// Generate an item name likely to be unique.
func (t *integrationTest) makeItemName() sdb.ItemName {
	return sdb.ItemName(fmt.Sprintf("item.%16x", uint64(rand.Int63())))
}

////////////////////////////////////////////////////////////////////////
// Domains
////////////////////////////////////////////////////////////////////////

var g_domainsTestDb sdb.SimpleDB
var g_domainsTestDomain0 sdb.Domain
var g_domainsTestDomain1 sdb.Domain

type DomainsTest struct {
	integrationTest

	mutex           sync.Mutex
	domainsToDelete []sdb.Domain  // Protected by mutex
}

func init() { RegisterTestSuite(&DomainsTest{}) }

func (t *DomainsTest) SetUpTestSuite() {
	var err error

	// Open a connection.
	g_domainsTestDb, err = sdb.NewSimpleDB(g_region, g_accessKey)
	if err != nil {
		panic(err)
	}

	// Create domain 0.
	g_domainsTestDomain0, err = g_domainsTestDb.OpenDomain("DomainsTest.domain0")
	if err != nil {
		panic(err)
	}

	// Create domain 1.
	g_domainsTestDomain1, err = g_domainsTestDb.OpenDomain("DomainsTest.domain1")
	if err != nil {
		panic(err)
	}
}

func (t *DomainsTest) TearDownTestSuite() {
	// Delete both domains.
	if err := g_domainsTestDb.DeleteDomain(g_domainsTestDomain0); err != nil {
		panic(err)
	}

	if err := g_domainsTestDb.DeleteDomain(g_domainsTestDomain1); err != nil {
		panic(err)
	}

	// Clear variables.
	g_domainsTestDb = nil
	g_domainsTestDomain0 = nil
	g_domainsTestDomain1 = nil
}

func (t *DomainsTest) TearDown() {
	t.mutex.Lock()
	defer t.mutex.Unlock()

	// Delete each of the domains created during the test.
	for _, d := range t.domainsToDelete {
		ExpectEq(nil, t.db.DeleteDomain(d), "Domain: %s", d.Name())
	}
}

func (t *DomainsTest) InvalidAccessKey() {
	// Open a connection with an unknown key ID.
	wrongKey := g_accessKey
	wrongKey.Id += "taco"

	db, err := sdb.NewSimpleDB(g_region, wrongKey)
	AssertEq(nil, err)

	// Attempt to create a domain.
	_, err = db.OpenDomain("some_domain")

	ExpectThat(err, Error(HasSubstr("403")))
	ExpectThat(err, Error(HasSubstr("Key Id")))
	ExpectThat(err, Error(HasSubstr("exist")))
}

func (t *DomainsTest) SeparatelyNamedDomainsHaveIndependentItems() {
	var err error

	// Set up an item in the first domain.
	itemName := t.makeItemName()
	err = g_domainsTestDomain0.PutAttributes(
		itemName,
		[]sdb.PutUpdate{
			sdb.PutUpdate{Name: "enchilada", Value: "queso"},
		},
		[]sdb.Precondition{},
	)

	AssertEq(nil, err)

	// Get attributes for the same name in the other domain. There should be
	// none.
	attrs, err := g_domainsTestDomain1.GetAttributes(itemName, true, []string{})
	AssertEq(nil, err)

	ExpectThat(attrs, ElementsAre())
}

func (t *DomainsTest) IdenticallyNamedDomainsHaveIdenticalItems() {
	var err error

	// Set up an item in the first domain.
	itemName := t.makeItemName()
	err = g_domainsTestDomain0.PutAttributes(
		itemName,
		[]sdb.PutUpdate{
			sdb.PutUpdate{Name: "enchilada", Value: "queso"},
		},
		[]sdb.Precondition{},
	)

	AssertEq(nil, err)

	// Get attributes for the same name in another domain object opened with the
	// same name.
	domain1, err := t.db.OpenDomain(g_domainsTestDomain0.Name())
	AssertEq(nil, err)

	attrs, err := domain1.GetAttributes(itemName, true, []string{})
	AssertEq(nil, err)

	ExpectThat(
		attrs,
		ElementsAre(
			DeepEquals(sdb.Attribute{Name: "enchilada", Value: "queso"}),
		),
	)
}

func (t *DomainsTest) Delete() {
	var err error
	domainName := "DomainsTest.Delete"

	// Create a domain, then delete it.
	domain, err := t.db.OpenDomain(domainName)
	AssertEq(nil, err)

	err = t.db.DeleteDomain(domain)
	AssertEq(nil, err)

	// Delete again; nothing should go wrong.
	err = t.db.DeleteDomain(domain)
	AssertEq(nil, err)

	// Attempt to write to the domain.
	err = domain.PutAttributes(
		"some_item",
		[]sdb.PutUpdate{
			sdb.PutUpdate{Name: "foo", Value: "bar"},
		},
		[]sdb.Precondition{},
	)

	ExpectThat(err, Error(HasSubstr("NoSuchDomain")))
	ExpectThat(err, Error(HasSubstr("domain")))
	ExpectThat(err, Error(HasSubstr("exist")))
}

////////////////////////////////////////////////////////////////////////
// Items
////////////////////////////////////////////////////////////////////////

type ItemsTest struct {
	integrationTest
}

func init() { RegisterTestSuite(&ItemsTest{}) }

func (t *ItemsTest) WrongAccessKeySecret() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidUtf8ItemName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidUtf8AttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidUtf8AttributeValue() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) LongItemName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) LongAttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) LongAttributeValue() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) PutThenGet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchPutThenGet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchPutThenBatchGet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) GetForNonExistentItem() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) GetParticularAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchGetParticularAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchGetForNonExistentItems() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) GetNonExistentAttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchGetNonExistentAttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) FailedValuePrecondition() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) FailedExistencePrecondition() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) FailedNonExistencePrecondition() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SuccessfulPreconditions() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) DeleteParticularAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) DeleteAllAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchDelete() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidSelectQuery() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectAll() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectItemName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectCount() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectWithPredicates() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectWithSortOrder() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectWithLimit() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectEmptyResultSet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectLargeResultSet() {
	ExpectEq("TODO", "")
}
<commit_msg>Added static setup for ItemsTest.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 main

import (
	"fmt"
	"github.com/jacobsa/aws/exp/sdb"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"math/rand"
	"sync"
)

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type integrationTest struct {
	db sdb.SimpleDB
}

func (t *integrationTest) SetUp(i *TestInfo) {
	var err error

	// Open a connection.
	t.db, err = sdb.NewSimpleDB(g_region, g_accessKey)
	AssertEq(nil, err)
}

// Generate an item name likely to be unique.
func (t *integrationTest) makeItemName() sdb.ItemName {
	return sdb.ItemName(fmt.Sprintf("item.%16x", uint64(rand.Int63())))
}

////////////////////////////////////////////////////////////////////////
// Domains
////////////////////////////////////////////////////////////////////////

var g_domainsTestDb sdb.SimpleDB
var g_domainsTestDomain0 sdb.Domain
var g_domainsTestDomain1 sdb.Domain

type DomainsTest struct {
	integrationTest

	mutex           sync.Mutex
	domainsToDelete []sdb.Domain  // Protected by mutex
}

func init() { RegisterTestSuite(&DomainsTest{}) }

func (t *DomainsTest) SetUpTestSuite() {
	var err error

	// Open a connection.
	g_domainsTestDb, err = sdb.NewSimpleDB(g_region, g_accessKey)
	if err != nil {
		panic(err)
	}

	// Create domain 0.
	g_domainsTestDomain0, err = g_domainsTestDb.OpenDomain("DomainsTest.domain0")
	if err != nil {
		panic(err)
	}

	// Create domain 1.
	g_domainsTestDomain1, err = g_domainsTestDb.OpenDomain("DomainsTest.domain1")
	if err != nil {
		panic(err)
	}
}

func (t *DomainsTest) TearDownTestSuite() {
	// Delete both domains.
	if err := g_domainsTestDb.DeleteDomain(g_domainsTestDomain0); err != nil {
		panic(err)
	}

	if err := g_domainsTestDb.DeleteDomain(g_domainsTestDomain1); err != nil {
		panic(err)
	}

	// Clear variables.
	g_domainsTestDb = nil
	g_domainsTestDomain0 = nil
	g_domainsTestDomain1 = nil
}

func (t *DomainsTest) TearDown() {
	t.mutex.Lock()
	defer t.mutex.Unlock()

	// Delete each of the domains created during the test.
	for _, d := range t.domainsToDelete {
		ExpectEq(nil, t.db.DeleteDomain(d), "Domain: %s", d.Name())
	}
}

func (t *DomainsTest) InvalidAccessKey() {
	// Open a connection with an unknown key ID.
	wrongKey := g_accessKey
	wrongKey.Id += "taco"

	db, err := sdb.NewSimpleDB(g_region, wrongKey)
	AssertEq(nil, err)

	// Attempt to create a domain.
	_, err = db.OpenDomain("some_domain")

	ExpectThat(err, Error(HasSubstr("403")))
	ExpectThat(err, Error(HasSubstr("Key Id")))
	ExpectThat(err, Error(HasSubstr("exist")))
}

func (t *DomainsTest) SeparatelyNamedDomainsHaveIndependentItems() {
	var err error

	// Set up an item in the first domain.
	itemName := t.makeItemName()
	err = g_domainsTestDomain0.PutAttributes(
		itemName,
		[]sdb.PutUpdate{
			sdb.PutUpdate{Name: "enchilada", Value: "queso"},
		},
		[]sdb.Precondition{},
	)

	AssertEq(nil, err)

	// Get attributes for the same name in the other domain. There should be
	// none.
	attrs, err := g_domainsTestDomain1.GetAttributes(itemName, true, []string{})
	AssertEq(nil, err)

	ExpectThat(attrs, ElementsAre())
}

func (t *DomainsTest) IdenticallyNamedDomainsHaveIdenticalItems() {
	var err error

	// Set up an item in the first domain.
	itemName := t.makeItemName()
	err = g_domainsTestDomain0.PutAttributes(
		itemName,
		[]sdb.PutUpdate{
			sdb.PutUpdate{Name: "enchilada", Value: "queso"},
		},
		[]sdb.Precondition{},
	)

	AssertEq(nil, err)

	// Get attributes for the same name in another domain object opened with the
	// same name.
	domain1, err := t.db.OpenDomain(g_domainsTestDomain0.Name())
	AssertEq(nil, err)

	attrs, err := domain1.GetAttributes(itemName, true, []string{})
	AssertEq(nil, err)

	ExpectThat(
		attrs,
		ElementsAre(
			DeepEquals(sdb.Attribute{Name: "enchilada", Value: "queso"}),
		),
	)
}

func (t *DomainsTest) Delete() {
	var err error
	domainName := "DomainsTest.Delete"

	// Create a domain, then delete it.
	domain, err := t.db.OpenDomain(domainName)
	AssertEq(nil, err)

	err = t.db.DeleteDomain(domain)
	AssertEq(nil, err)

	// Delete again; nothing should go wrong.
	err = t.db.DeleteDomain(domain)
	AssertEq(nil, err)

	// Attempt to write to the domain.
	err = domain.PutAttributes(
		"some_item",
		[]sdb.PutUpdate{
			sdb.PutUpdate{Name: "foo", Value: "bar"},
		},
		[]sdb.Precondition{},
	)

	ExpectThat(err, Error(HasSubstr("NoSuchDomain")))
	ExpectThat(err, Error(HasSubstr("domain")))
	ExpectThat(err, Error(HasSubstr("exist")))
}

////////////////////////////////////////////////////////////////////////
// Items
////////////////////////////////////////////////////////////////////////

var g_itemsTestDb     sdb.SimpleDB
var g_itemsTestDomain sdb.Domain

type ItemsTest struct {
	integrationTest
}

func init() { RegisterTestSuite(&ItemsTest{}) }

func (t *ItemsTest) SetUpTestSuite() {
	var err error

	// Open a connection.
	g_itemsTestDb, err = sdb.NewSimpleDB(g_region, g_accessKey)
	if err != nil {
		panic(err)
	}

	// Create a domain.
	g_itemsTestDomain, err = g_itemsTestDb.OpenDomain("ItemsTest.domain")
	if err != nil {
		panic(err)
	}
}

func (t *ItemsTest) TearDownTestSuite() {
	// Delete the domain.
	if err := g_itemsTestDb.DeleteDomain(g_itemsTestDomain); err != nil {
		panic(err)
	}

	// Clear variables.
	g_itemsTestDb = nil
	g_itemsTestDomain = nil
}

func (t *ItemsTest) WrongAccessKeySecret() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidUtf8ItemName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidUtf8AttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidUtf8AttributeValue() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) LongItemName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) LongAttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) LongAttributeValue() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) PutThenGet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchPutThenGet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchPutThenBatchGet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) GetForNonExistentItem() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) GetParticularAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchGetParticularAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchGetForNonExistentItems() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) GetNonExistentAttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchGetNonExistentAttributeName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) FailedValuePrecondition() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) FailedExistencePrecondition() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) FailedNonExistencePrecondition() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SuccessfulPreconditions() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) DeleteParticularAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) DeleteAllAttributes() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) BatchDelete() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) InvalidSelectQuery() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectAll() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectItemName() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectCount() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectWithPredicates() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectWithSortOrder() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectWithLimit() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectEmptyResultSet() {
	ExpectEq("TODO", "")
}

func (t *ItemsTest) SelectLargeResultSet() {
	ExpectEq("TODO", "")
}
<|endoftext|>