<commit_before>// Copyright 2018 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
//
//      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 proxy provides a record/replay HTTP proxy. It is designed to support
// both an in-memory API (cloud.google.com/go/httpreplay) and a standalone server
// (cloud.google.com/go/httpreplay/cmd/httpr).
package proxy

// See github.com/google/martian/cmd/proxy/main.go for the origin of much of this.

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/google/martian"
	"github.com/google/martian/fifo"
	"github.com/google/martian/httpspec"
	"github.com/google/martian/martianlog"
	"github.com/google/martian/mitm"
)

// A Proxy is an HTTP proxy that supports recording or replaying requests.
type Proxy struct {
	// The certificate that the proxy uses to participate in TLS.
	CACert *x509.Certificate

	// The URL of the proxy.
	URL *url.URL

	// Initial state of the client.
	Initial []byte

	mproxy        *martian.Proxy
	filename      string          // for log
	logger        *Logger         // for recording only
	ignoreHeaders map[string]bool // headers the user has asked to ignore
}

// ForRecording returns a Proxy configured to record.
func ForRecording(filename string, port int) (*Proxy, error) {
	p, err := newProxy(filename)
	if err != nil {
		return nil, err
	}

	// Construct a group that performs the standard proxy stack of request/response
	// modifications.
	stack, _ := httpspec.NewStack("httpr") // second arg is an internal group that we don't need
	p.mproxy.SetRequestModifier(stack)
	p.mproxy.SetResponseModifier(stack)

	// Make a group for logging requests and responses.
	logGroup := fifo.NewGroup()
	skipAuth := skipLoggingByHost("accounts.google.com")
	logGroup.AddRequestModifier(skipAuth)
	logGroup.AddResponseModifier(skipAuth)
	p.logger = newLogger()
	logGroup.AddRequestModifier(p.logger)
	logGroup.AddResponseModifier(p.logger)

	stack.AddRequestModifier(logGroup)
	stack.AddResponseModifier(logGroup)

	// Ordinary debug logging.
	logger := martianlog.NewLogger()
	logger.SetDecode(true)
	stack.AddRequestModifier(logger)
	stack.AddResponseModifier(logger)

	if err := p.start(port); err != nil {
		return nil, err
	}
	return p, nil
}

type hideTransport http.Transport

func (t *hideTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	return (*http.Transport)(t).RoundTrip(req)
}

func newProxy(filename string) (*Proxy, error) {
	mproxy := martian.NewProxy()
	// Set up a man-in-the-middle configuration with a CA certificate so the proxy can
	// participate in TLS.
	x509c, priv, err := mitm.NewAuthority("cloud.google.com/go/httpreplay", "HTTPReplay Authority", time.Hour)
	if err != nil {
		return nil, err
	}
	mc, err := mitm.NewConfig(x509c, priv)
	if err != nil {
		return nil, err
	}
	mc.SetValidity(time.Hour)
	mc.SetOrganization("cloud.google.com/go/httpreplay")
	mc.SkipTLSVerify(false)
	if err != nil {
		return nil, err
	}
	mproxy.SetMITM(mc)
	return &Proxy{
		mproxy:        mproxy,
		CACert:        x509c,
		filename:      filename,
		ignoreHeaders: map[string]bool{},
	}, nil
}

func (p *Proxy) start(port int) error {
	l, err := net.Listen("tcp4", fmt.Sprintf(":%d", port))
	if err != nil {
		return err
	}
	p.URL = &url.URL{Scheme: "http", Host: l.Addr().String()}
	go p.mproxy.Serve(l)
	return nil
}

// Transport returns an http.Transport for clients who want to talk to the proxy.
func (p *Proxy) Transport() *http.Transport {
	caCertPool := x509.NewCertPool()
	caCertPool.AddCert(p.CACert)
	return &http.Transport{
		TLSClientConfig: &tls.Config{RootCAs: caCertPool},
		Proxy:           func(*http.Request) (*url.URL, error) { return p.URL, nil },
	}
}

// RemoveRequestHeaders will remove request headers matching patterns from the log,
// and skip matching them. Pattern is taken literally except for *, which matches any
// sequence of characters.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) RemoveRequestHeaders(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerRemoveRequestHeaders(pat)
	}
}

// ClearHeaders will replace matching headers with CLEARED.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) ClearHeaders(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerClearHeaders(pat)
	}
}

// RemoveQueryParams will remove query parameters matching patterns from the request
// URL before logging, and skip matching them. Pattern is taken literally except for
// *, which matches any sequence of characters.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) RemoveQueryParams(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerRemoveParams(pat)
	}
}

// ClearQueryParams will replace matching query params in the request URL with CLEARED.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) ClearQueryParams(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerClearParams(pat)
	}
}

// IgnoreHeader will cause h to be ignored during matching on replay.
// Deprecated: use RemoveRequestHeaders instead.
func (p *Proxy) IgnoreHeader(h string) {
	p.ignoreHeaders[http.CanonicalHeaderKey(h)] = true
}

// Close closes the proxy. If the proxy is recording, it also writes the log.
func (p *Proxy) Close() error {
	p.mproxy.Close()
	if p.logger != nil {
		return p.writeLog()
	}
	return nil
}

func (p *Proxy) writeLog() error {
	lg := p.logger.Extract()
	lg.Initial = p.Initial
	bytes, err := json.MarshalIndent(lg, "", "  ")
	if err != nil {
		return err
	}
	return ioutil.WriteFile(p.filename, bytes, 0600) // only accessible by owner
}

// skipLoggingByHost disables logging for traffic to a particular host.
type skipLoggingByHost string

func (s skipLoggingByHost) ModifyRequest(req *http.Request) error {
	if strings.HasPrefix(req.Host, string(s)) {
		martian.NewContext(req).SkipLogging()
	}
	return nil
}

func (s skipLoggingByHost) ModifyResponse(res *http.Response) error {
	return s.ModifyRequest(res.Request)
}
<commit_msg>httpreplay: set up MITM config only once<commit_after>// Copyright 2018 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
//
//      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 proxy provides a record/replay HTTP proxy. It is designed to support
// both an in-memory API (cloud.google.com/go/httpreplay) and a standalone server
// (cloud.google.com/go/httpreplay/cmd/httpr).
package proxy

// See github.com/google/martian/cmd/proxy/main.go for the origin of much of this.

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/google/martian"
	"github.com/google/martian/fifo"
	"github.com/google/martian/httpspec"
	"github.com/google/martian/martianlog"
	"github.com/google/martian/mitm"
)

// A Proxy is an HTTP proxy that supports recording or replaying requests.
type Proxy struct {
	// The certificate that the proxy uses to participate in TLS.
	CACert *x509.Certificate

	// The URL of the proxy.
	URL *url.URL

	// Initial state of the client.
	Initial []byte

	mproxy        *martian.Proxy
	filename      string          // for log
	logger        *Logger         // for recording only
	ignoreHeaders map[string]bool // headers the user has asked to ignore
}

// ForRecording returns a Proxy configured to record.
func ForRecording(filename string, port int) (*Proxy, error) {
	p, err := newProxy(filename)
	if err != nil {
		return nil, err
	}

	// Construct a group that performs the standard proxy stack of request/response
	// modifications.
	stack, _ := httpspec.NewStack("httpr") // second arg is an internal group that we don't need
	p.mproxy.SetRequestModifier(stack)
	p.mproxy.SetResponseModifier(stack)

	// Make a group for logging requests and responses.
	logGroup := fifo.NewGroup()
	skipAuth := skipLoggingByHost("accounts.google.com")
	logGroup.AddRequestModifier(skipAuth)
	logGroup.AddResponseModifier(skipAuth)
	p.logger = newLogger()
	logGroup.AddRequestModifier(p.logger)
	logGroup.AddResponseModifier(p.logger)

	stack.AddRequestModifier(logGroup)
	stack.AddResponseModifier(logGroup)

	// Ordinary debug logging.
	logger := martianlog.NewLogger()
	logger.SetDecode(true)
	stack.AddRequestModifier(logger)
	stack.AddResponseModifier(logger)

	if err := p.start(port); err != nil {
		return nil, err
	}
	return p, nil
}

type hideTransport http.Transport

func (t *hideTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	return (*http.Transport)(t).RoundTrip(req)
}

var (
	configOnce sync.Once
	cert       *x509.Certificate
	config     *mitm.Config
	configErr  error
)

func newProxy(filename string) (*Proxy, error) {
	configOnce.Do(func() {
		// Set up a man-in-the-middle configuration with a CA certificate so the proxy can
		// participate in TLS.
		x509c, priv, err := mitm.NewAuthority("cloud.google.com/go/httpreplay", "HTTPReplay Authority", 100*time.Hour)
		if err != nil {
			configErr = err
			return
		}
		cert = x509c
		config, configErr = mitm.NewConfig(x509c, priv)
		if config != nil {
			config.SetValidity(100 * time.Hour)
			config.SetOrganization("cloud.google.com/go/httpreplay")
			config.SkipTLSVerify(false)
		}
	})
	if configErr != nil {
		return nil, configErr
	}
	mproxy := martian.NewProxy()
	mproxy.SetMITM(config)
	return &Proxy{
		mproxy:        mproxy,
		CACert:        cert,
		filename:      filename,
		ignoreHeaders: map[string]bool{},
	}, nil
}

func (p *Proxy) start(port int) error {
	l, err := net.Listen("tcp4", fmt.Sprintf(":%d", port))
	if err != nil {
		return err
	}
	p.URL = &url.URL{Scheme: "http", Host: l.Addr().String()}
	go p.mproxy.Serve(l)
	return nil
}

// Transport returns an http.Transport for clients who want to talk to the proxy.
func (p *Proxy) Transport() *http.Transport {
	caCertPool := x509.NewCertPool()
	caCertPool.AddCert(p.CACert)
	return &http.Transport{
		TLSClientConfig: &tls.Config{RootCAs: caCertPool},
		Proxy:           func(*http.Request) (*url.URL, error) { return p.URL, nil },
	}
}

// RemoveRequestHeaders will remove request headers matching patterns from the log,
// and skip matching them. Pattern is taken literally except for *, which matches any
// sequence of characters.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) RemoveRequestHeaders(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerRemoveRequestHeaders(pat)
	}
}

// ClearHeaders will replace matching headers with CLEARED.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) ClearHeaders(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerClearHeaders(pat)
	}
}

// RemoveQueryParams will remove query parameters matching patterns from the request
// URL before logging, and skip matching them. Pattern is taken literally except for
// *, which matches any sequence of characters.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) RemoveQueryParams(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerRemoveParams(pat)
	}
}

// ClearQueryParams will replace matching query params in the request URL with CLEARED.
//
// This only needs to be called during recording; the patterns will be saved to the
// log for replay.
func (p *Proxy) ClearQueryParams(patterns []string) {
	for _, pat := range patterns {
		p.logger.log.Converter.registerClearParams(pat)
	}
}

// IgnoreHeader will cause h to be ignored during matching on replay.
// Deprecated: use RemoveRequestHeaders instead.
func (p *Proxy) IgnoreHeader(h string) {
	p.ignoreHeaders[http.CanonicalHeaderKey(h)] = true
}

// Close closes the proxy. If the proxy is recording, it also writes the log.
func (p *Proxy) Close() error {
	p.mproxy.Close()
	if p.logger != nil {
		return p.writeLog()
	}
	return nil
}

func (p *Proxy) writeLog() error {
	lg := p.logger.Extract()
	lg.Initial = p.Initial
	bytes, err := json.MarshalIndent(lg, "", "  ")
	if err != nil {
		return err
	}
	return ioutil.WriteFile(p.filename, bytes, 0600) // only accessible by owner
}

// skipLoggingByHost disables logging for traffic to a particular host.
type skipLoggingByHost string

func (s skipLoggingByHost) ModifyRequest(req *http.Request) error {
	if strings.HasPrefix(req.Host, string(s)) {
		martian.NewContext(req).SkipLogging()
	}
	return nil
}

func (s skipLoggingByHost) ModifyResponse(res *http.Response) error {
	return s.ModifyRequest(res.Request)
}
<|endoftext|>


<commit_before>package grpc

import (
	"net/http"

	"github.com/micro/go-micro/v3/errors"
	"google.golang.org/grpc/codes"
)

func microError(err *errors.Error) codes.Code {
	switch err {
	case nil:
		return codes.OK
	}

	switch err.Code {
	case http.StatusOK:
		return codes.OK
	case http.StatusBadRequest:
		return codes.InvalidArgument
	case http.StatusRequestTimeout:
		return codes.DeadlineExceeded
	case http.StatusNotFound:
		return codes.NotFound
	case http.StatusConflict:
		return codes.AlreadyExists
	case http.StatusForbidden:
		return codes.PermissionDenied
	case http.StatusUnauthorized:
		return codes.Unauthenticated
	case http.StatusPreconditionFailed:
		return codes.FailedPrecondition
	case http.StatusNotImplemented:
		return codes.Unimplemented
	case http.StatusInternalServerError:
		return codes.Internal
	case http.StatusServiceUnavailable:
		return codes.Unavailable
	}

	return codes.Unknown
}
<commit_msg>simplifies code (#1934)<commit_after>package grpc

import (
	"net/http"

	"github.com/micro/go-micro/v3/errors"
	"google.golang.org/grpc/codes"
)

var errMapping = map[int32]codes.Code{
	http.StatusOK:                  codes.OK,
	http.StatusBadRequest:          codes.InvalidArgument,
	http.StatusRequestTimeout:      codes.DeadlineExceeded,
	http.StatusNotFound:            codes.NotFound,
	http.StatusConflict:            codes.AlreadyExists,
	http.StatusForbidden:           codes.PermissionDenied,
	http.StatusUnauthorized:        codes.Unauthenticated,
	http.StatusPreconditionFailed:  codes.FailedPrecondition,
	http.StatusNotImplemented:      codes.Unimplemented,
	http.StatusInternalServerError: codes.Internal,
	http.StatusServiceUnavailable:  codes.Unavailable,
}

func microError(err *errors.Error) codes.Code {
	if err == nil {
		return codes.OK
	}

	if code, ok := errMapping[err.Code]; ok {
		return code
	}
	return codes.Unknown
}
<|endoftext|>


<commit_before>//   Copyright 2017 MSolution.IO
//
//   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 (
	"context"
	"database/sql"
	"errors"
	"flag"
	"math/rand"
	"strconv"
	"time"

	"github.com/trackit/jsonlog"

	"github.com/trackit/trackit-server/aws"
	"github.com/trackit/trackit-server/aws/s3"
	"github.com/trackit/trackit-server/db"
	"github.com/aws/aws-sdk-go/aws/awserr"
)

// taskIngest ingests billing data for a given BillRepository and AwsAccount.
func taskIngest(ctx context.Context) error {
	args := flag.Args()
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	logger.Debug("Running task 'ingest'.", map[string]interface{}{
		"args": args,
	})
	if len(args) != 2 {
		return errors.New("taskIngest requires two integer arguments")
	} else if aa, err := strconv.Atoi(args[0]); err != nil {
		return err
	} else if br, err := strconv.Atoi(args[1]); err != nil {
		return err
	} else {
		return ingestBillingDataForBillRepository(ctx, aa, br)
	}
}

// ingestBillingDataForBillRepository ingests the billing data for a
// BillRepository.
func ingestBillingDataForBillRepository(ctx context.Context, aaId, brId int) (err error) {
	var tx *sql.Tx
	var aa aws.AwsAccount
	var br s3.BillRepository
	var updateId int64
	var latestManifest time.Time
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	defer func() {
		if tx != nil {
			if err != nil {
				tx.Rollback()
			} else {
				tx.Commit()
			}
		}
	}()
	if tx, err = db.Db.BeginTx(ctx, nil); err != nil {
	} else if aa, err = aws.GetAwsAccountWithId(aaId, tx); err != nil {
	} else if br, err = s3.GetBillRepositoryForAwsAccountById(aa, brId, tx); err != nil {
	} else if updateId, err = registerUpdate(db.Db, br); err != nil {
	} else if latestManifest, err = s3.UpdateReport(ctx, aa, br); err != nil {
		if billError, castok := err.(awserr.Error); castok {
			br.Error = billError.Message()
			s3.UpdateBillRepositoryWithoutContext(br, db.Db)
		}
	} else {
		br.Error = ""
		err = updateBillRepositoryForNextUpdate(ctx, tx, br, latestManifest)
	}
	if err != nil {
		logger.Error("Failed to ingest billing data.", map[string]interface{}{
			"awsAccountId":     aaId,
			"billRepositoryId": brId,
			"error":            err.Error(),
		})
	}
	updateCompletion(ctx, aaId, brId, db.Db, updateId, err)
	updateSubAccounts(ctx, aa)
	return
}

func updateSubAccounts(ctx context.Context, aa aws.AwsAccount) {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	var tx *sql.Tx
	var err error
	defer func() {
		if tx != nil {
			if err != nil {
				tx.Rollback()
			} else {
				tx.Commit()
			}
		}
	}()
	if tx, err = db.Db.BeginTx(ctx, nil); err == nil {
		err = aws.PutSubAccounts(ctx, aa, tx)
	}
	if err != nil {
		logger.Error("Failed to update sub accounts.", map[string]interface{}{
			"awsAccountId": aa.Id,
			"error":        err.Error(),
		})
	} else {
		logger.Info("Sub accounts updated.", map[string]interface{}{
			"awsAccountId": aa.Id,
		})
	}
}

func registerUpdate(db *sql.DB, br s3.BillRepository) (int64, error) {
	const sqlstr = `INSERT INTO aws_bill_update_job(
		aws_bill_repository_id,
		worker_id,
		error
	) VALUES (?, ?, "")`
	res, err := db.Exec(sqlstr, br.Id, backendId)
	if err != nil {
		return 0, err
	}
	return res.LastInsertId()
}

func updateCompletion(ctx context.Context, aaId, brId int, db *sql.DB, updateId int64, err error) {
	rErr := registerUpdateCompletion(db, updateId, err)
	if rErr != nil {
		logger := jsonlog.LoggerFromContextOrDefault(ctx)
		logger.Error("Failed to register ingestion completion.", map[string]interface{}{
			"awsAccountId":     aaId,
			"billRepositoryId": brId,
			"error":            rErr.Error(),
			"updateId":         updateId,
		})
	}
}

func registerUpdateCompletion(db *sql.DB, updateId int64, err error) error {
	const sqlstr = `UPDATE aws_bill_update_job SET
		completed=?,
		error=?
	WHERE id=?`
	var errorValue string
	var now = time.Now()
	if err != nil {
		errorValue = err.Error()
	}
	_, err = db.Exec(sqlstr, now, errorValue, updateId)
	return err
}

const (
	UpdateIntervalMinutes = 6 * 60
	UpdateIntervalWindow  = 2 * 60
)

// updateBillRepositoryForNextUpdate plans the next update for a
// BillRepository.
func updateBillRepositoryForNextUpdate(ctx context.Context, tx *sql.Tx, br s3.BillRepository, latestManifest time.Time) error {
	if latestManifest.After(br.LastImportedManifest) {
		br.LastImportedManifest = latestManifest
	}
	updateDeltaMinutes := time.Duration(UpdateIntervalMinutes-UpdateIntervalWindow/2+rand.Int63n(UpdateIntervalWindow)) * time.Minute
	br.NextUpdate = time.Now().Add(updateDeltaMinutes)
	return s3.UpdateBillRepository(br, tx)
}
<commit_msg>change sub account logs<commit_after>//   Copyright 2017 MSolution.IO
//
//   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 (
	"context"
	"database/sql"
	"errors"
	"flag"
	"math/rand"
	"strconv"
	"time"

	"github.com/trackit/jsonlog"

	"github.com/trackit/trackit-server/aws"
	"github.com/trackit/trackit-server/aws/s3"
	"github.com/trackit/trackit-server/db"
	"github.com/aws/aws-sdk-go/aws/awserr"
)

// taskIngest ingests billing data for a given BillRepository and AwsAccount.
func taskIngest(ctx context.Context) error {
	args := flag.Args()
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	logger.Debug("Running task 'ingest'.", map[string]interface{}{
		"args": args,
	})
	if len(args) != 2 {
		return errors.New("taskIngest requires two integer arguments")
	} else if aa, err := strconv.Atoi(args[0]); err != nil {
		return err
	} else if br, err := strconv.Atoi(args[1]); err != nil {
		return err
	} else {
		return ingestBillingDataForBillRepository(ctx, aa, br)
	}
}

// ingestBillingDataForBillRepository ingests the billing data for a
// BillRepository.
func ingestBillingDataForBillRepository(ctx context.Context, aaId, brId int) (err error) {
	var tx *sql.Tx
	var aa aws.AwsAccount
	var br s3.BillRepository
	var updateId int64
	var latestManifest time.Time
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	defer func() {
		if tx != nil {
			if err != nil {
				tx.Rollback()
			} else {
				tx.Commit()
			}
		}
	}()
	if tx, err = db.Db.BeginTx(ctx, nil); err != nil {
	} else if aa, err = aws.GetAwsAccountWithId(aaId, tx); err != nil {
	} else if br, err = s3.GetBillRepositoryForAwsAccountById(aa, brId, tx); err != nil {
	} else if updateId, err = registerUpdate(db.Db, br); err != nil {
	} else if latestManifest, err = s3.UpdateReport(ctx, aa, br); err != nil {
		if billError, castok := err.(awserr.Error); castok {
			br.Error = billError.Message()
			s3.UpdateBillRepositoryWithoutContext(br, db.Db)
		}
	} else {
		br.Error = ""
		err = updateBillRepositoryForNextUpdate(ctx, tx, br, latestManifest)
	}
	if err != nil {
		logger.Error("Failed to ingest billing data.", map[string]interface{}{
			"awsAccountId":     aaId,
			"billRepositoryId": brId,
			"error":            err.Error(),
		})
	}
	updateCompletion(ctx, aaId, brId, db.Db, updateId, err)
	updateSubAccounts(ctx, aa)
	return
}

func updateSubAccounts(ctx context.Context, aa aws.AwsAccount) {
	logger := jsonlog.LoggerFromContextOrDefault(ctx)
	var tx *sql.Tx
	var err error
	defer func() {
		if tx != nil {
			if err != nil {
				tx.Rollback()
			} else {
				tx.Commit()
			}
		}
	}()
	if tx, err = db.Db.BeginTx(ctx, nil); err != nil {
		logger.Error("Failed to get DB Tx", err.Error())
	} else if err = aws.PutSubAccounts(ctx, aa, tx); err != nil {
		logger.Warning("Failed to update sub accounts.", map[string]interface{}{
			"awsAccountId": aa.Id,
			"error":        err.Error(),
		})
	} else {
		logger.Info("Sub accounts updated.", map[string]interface{}{
			"awsAccountId": aa.Id,
		})
	}
}

func registerUpdate(db *sql.DB, br s3.BillRepository) (int64, error) {
	const sqlstr = `INSERT INTO aws_bill_update_job(
		aws_bill_repository_id,
		worker_id,
		error
	) VALUES (?, ?, "")`
	res, err := db.Exec(sqlstr, br.Id, backendId)
	if err != nil {
		return 0, err
	}
	return res.LastInsertId()
}

func updateCompletion(ctx context.Context, aaId, brId int, db *sql.DB, updateId int64, err error) {
	rErr := registerUpdateCompletion(db, updateId, err)
	if rErr != nil {
		logger := jsonlog.LoggerFromContextOrDefault(ctx)
		logger.Error("Failed to register ingestion completion.", map[string]interface{}{
			"awsAccountId":     aaId,
			"billRepositoryId": brId,
			"error":            rErr.Error(),
			"updateId":         updateId,
		})
	}
}

func registerUpdateCompletion(db *sql.DB, updateId int64, err error) error {
	const sqlstr = `UPDATE aws_bill_update_job SET
		completed=?,
		error=?
	WHERE id=?`
	var errorValue string
	var now = time.Now()
	if err != nil {
		errorValue = err.Error()
	}
	_, err = db.Exec(sqlstr, now, errorValue, updateId)
	return err
}

const (
	UpdateIntervalMinutes = 6 * 60
	UpdateIntervalWindow  = 2 * 60
)

// updateBillRepositoryForNextUpdate plans the next update for a
// BillRepository.
func updateBillRepositoryForNextUpdate(ctx context.Context, tx *sql.Tx, br s3.BillRepository, latestManifest time.Time) error {
	if latestManifest.After(br.LastImportedManifest) {
		br.LastImportedManifest = latestManifest
	}
	updateDeltaMinutes := time.Duration(UpdateIntervalMinutes-UpdateIntervalWindow/2+rand.Int63n(UpdateIntervalWindow)) * time.Minute
	br.NextUpdate = time.Now().Add(updateDeltaMinutes)
	return s3.UpdateBillRepository(br, tx)
}
<|endoftext|>


<commit_before>/*******************************************************************************
* Copyright (C) Zenoss, Inc. 2013, all rights reserved.
*
* This content is made available according to terms specified in
* License.zenoss under the directory where your Zenoss product is installed.
*
*******************************************************************************/

package main

// This is here the command line arguments are parsed and executed.

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/zenoss/serviced"
	clientlib "github.com/zenoss/serviced/client"
	"log"
	"reflect"
	"strconv"
	"strings"
)

// A type to represent the CLI. All the command will have the same signature.
// This makes it easy to call them arbitrarily.
type ServicedCli struct{}

// A helper function that creates a subcommand
func Subcmd(name, signature, description string) *flag.FlagSet {
	flags := flag.NewFlagSet(name, flag.ContinueOnError)
	flags.Usage = func() {
		fmt.Printf("\nUsage: serviced %s %s\n\n%s\n\n", name, signature, description)
		flags.PrintDefaults()
	}
	return flags
}

// Use reflection to aquire the give method by name. For example to get method
// CmdFoo, pass 'foo'. A method is returned. The second return argument
// indicates if the argument was found.
func (cli *ServicedCli) getMethod(name string) (reflect.Method, bool) {

	// Contruct the method name to be CmdFoo, where foo was passed
	methodName := "Cmd"
	for _, part := range strings.Split(name, "-") {
		methodName = methodName + strings.ToUpper(part[:1]) + strings.ToLower(part[1:])
	}
	return reflect.TypeOf(cli).MethodByName(methodName)
}

// Construct a new command line parsing object.
func NewServicedCli() (s *ServicedCli) {
	return &ServicedCli{}
}

// Show usage of serviced command line options.
func (cli *ServicedCli) CmdHelp(args ...string) error {
	if len(args) > 0 {
		method, exists := cli.getMethod(args[0])
		if !exists {
			fmt.Println("Error: Command not found:", args[0])
		} else {
			method.Func.CallSlice([]reflect.Value{
				reflect.ValueOf(cli),
				reflect.ValueOf([]string{"--help"}),
			})[0].Interface()
			return nil
		}
	}
	help := fmt.Sprintf("Usage: serviced [OPTIONS] COMMAND [arg...]\n\nA container based service management system.\n\nCommands:\n")
	for _, command := range [][2]string{
		{"hosts", "Display hosts"},
		{"update-host", "Update a host"},
		{"add-host", "Add a host"},
		{"remove-host", "Remove a host"},
		{"pools", "Show pools"},
		{"add-pool", "Add pool"},
		{"services", "Show services"},
		{"add-service", "Add a service"},
		{"remove-service", "Remote a service"},
		{"start-service", "Start a service"},
		{"stop-service", "Stop a service"},
	} {
		help += fmt.Sprintf("    %-30.30s%s\n", command[0], command[1])
	}
	fmt.Println(help)
	return nil
}

// Attempt to find the command give on the CLI by looking up the method on the
// CLI interface. If found, execute it. Otherwise show usage.
func ParseCommands(args ...string) error {
	cli := NewServicedCli()

	if len(args) > 0 {
		method, exists := cli.getMethod(args[0])
		if !exists {
			fmt.Println("Error: Command not found:", args[0])
			return cli.CmdHelp(args[1:]...)
		}
		ret := method.Func.CallSlice([]reflect.Value{
			reflect.ValueOf(cli),
			reflect.ValueOf(args[1:]),
		})[0].Interface()
		if ret == nil {
			return nil
		}
		return ret.(error)
	}
	return cli.CmdHelp(args...)
}

// Create a client to the control plane.
func getClient() (c serviced.ControlPlane) {
	// setup the client
	c, err := clientlib.NewControlClient(options.port)
	if err != nil {
		log.Fatalf("Could not create acontrol plane client %v", err)
	}
	return c
}

// List the hosts associated with the control plane.
func (cli *ServicedCli) CmdHosts(args ...string) error {

	cmd := Subcmd("hosts", "[OPTIONS]", "List hosts")
	if err := cmd.Parse(args); err != nil {
		return err
	}

	client := getClient()

	var hosts map[string]*serviced.Host
	request := serviced.EntityRequest{}
	err := client.GetHosts(request, &hosts)
	if err != nil {
		log.Fatalf("Could not get hosts %v", err)
	}
	hostsJson, err := json.MarshalIndent(hosts, " ", "  ")
	if err == nil {
		fmt.Printf("%s\n", hostsJson)
	}
	return err
}

// Add a host to the control plane given the host:port.
func (cli *ServicedCli) CmdAddHost(args ...string) error {

	cmd := Subcmd("add-host", "HOST:PORT RESOURCE_POOL", "Add host")
	if err := cmd.Parse(args); err != nil {
		return nil
	}

	if len(cmd.Args()) != 2 {
		cmd.Usage()
		return nil
	}

	client, err := clientlib.NewAgentClient(cmd.Arg(0))
	if err != nil {
		log.Fatalf("Could not create connection to host %s: %v", args[0], err)
	}

	var remoteHost serviced.Host
	err = client.GetInfo(0, &remoteHost)
	if err != nil {
		log.Fatalf("Could not get remote host info: %v", err)
	}
	remoteHost.PoolId = cmd.Arg(1)
	log.Printf("Got host info: %v", remoteHost)

	controlPlane := getClient()
	var unused int

	err = controlPlane.AddHost(remoteHost, &unused)
	if err != nil {
		log.Fatalf("Could not add host: %v", err)
	}
	fmt.Println(remoteHost.Id)
	return err
}

// Update the host information. This method contacts the agent running on
// HOST:PORT to update the information assoicated with the host.
func (cli *ServicedCli) CmdUpdateHost(args ...string) error {

	cmd := Subcmd("update-host", "HOST:PORT", "Update the host information.")
	if err := cmd.Parse(args); err != nil {
		return nil
	}

	if len(cmd.Args()) != 1 {
		cmd.Usage()
		return nil
	}

	client, err := clientlib.NewAgentClient(cmd.Arg(0))
	if err != nil {
		log.Fatalf("Could not create connection to host %s: %v", args[0], err)
	}

	var remoteHost serviced.Host
	err = client.GetInfo(0, &remoteHost)
	if err != nil {
		log.Fatalf("Could not get remote host info: %v", err)
	}
	log.Printf("Got host info: %v", remoteHost)

	controlPlane := getClient()
	var unused int

	err = controlPlane.UpdateHost(remoteHost, &unused)
	if err != nil {
		log.Fatalf("Could not update host: %v", err)
	}
	return err
}

// This method removes the given host (by HOSTID) from the system.
func (cli *ServicedCli) CmdRemoveHost(args ...string) error {
	cmd := Subcmd("remove-host", "HOSTID", "Remove the host.")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	if len(cmd.Args()) != 1 {
		cmd.Usage()
		return nil
	}

	controlPlane := getClient()
	var unused int
	err := controlPlane.RemoveHost(cmd.Arg(0), &unused)
	if err != nil {
		log.Fatalf("Could not remove host: %v", err)
	}
	log.Printf("Host %s removed.", cmd.Arg(0))
	return err
}

// A convinience struct for printing to command line
type poolWithHost struct {
	serviced.ResourcePool
	Hosts []string
}

// Print a list of pools. Args are ignored.
func (cli *ServicedCli) CmdPools(args ...string) error {
	cmd := Subcmd("pools", "[OPTIONS]", "Display pools")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	controlPlane := getClient()
	request := serviced.EntityRequest{}
	var pools map[string]*serviced.ResourcePool
	err := controlPlane.GetResourcePools(request, &pools)
	if err != nil {
		log.Fatalf("Could not get resource pools: %v", err)
	}
	poolsWithHost := make(map[string]poolWithHost)
	for _, pool := range pools {

		// get pool hosts
		var poolHosts []*serviced.PoolHost
		err = controlPlane.GetHostsForResourcePool(pool.Id, &poolHosts)
		if err != nil {
			log.Fatalf("Could not get hosts for Pool %s: %v", pool.Id, err)
		}
		hosts := make([]string, len(poolHosts))
		for i, hostPool := range poolHosts {
			hosts[i] = hostPool.HostId
		}
		poolsWithHost[pool.Id] = poolWithHost{*pool, hosts}
	}
	poolsWithHostJson, err := json.MarshalIndent(poolsWithHost, " ", "  ")
	if err == nil {
		fmt.Printf("%s\n", poolsWithHostJson)
	}
	return err
}

// Add a new pool given some parameters.
func (cli *ServicedCli) CmdAddPool(args ...string) error {
	cmd := Subcmd("add-pool", "[options] POOLID CORE_LIMIT MEMORY_LIMIT PRIORITY", "Add resource pool")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	if len(cmd.Args()) < 4 {
		cmd.Usage()
		return nil
	}
	pool, _ := serviced.NewResourcePool(cmd.Arg(0))
	coreLimit, err := strconv.Atoi(cmd.Arg(1))
	if err != nil {
		log.Fatal("Bad core limit %s: %v", cmd.Arg(1), err)
	}
	pool.CoreLimit = coreLimit
	memoryLimit, err := strconv.Atoi(cmd.Arg(2))
	if err != nil {
		log.Fatal("Bad memory limit %s: %v", cmd.Arg(2), err)
	}
	pool.MemoryLimit = uint64(memoryLimit)
	controlPlane := getClient()
	var unused int
	err = controlPlane.AddResourcePool(*pool, &unused)
	if err != nil {
		log.Fatalf("Could not add resource pool: %v", err)
	}
	fmt.Printf("%s\n", pool.Id)
	return err
}

// Print the list of available services.
func (cli *ServicedCli) CmdServices(args ...string) error {
	cmd := Subcmd("services", "[CMD]", "Show services")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	controlPlane := getClient()
	request := serviced.EntityRequest{}
	var services []*serviced.Service
	err := controlPlane.GetServices(request, &services)
	if err != nil {
		log.Fatalf("Could not get services: %v", err)
	}
	servicesJson, err := json.MarshalIndent(services, " ", " ")
	if err != nil {
		log.Fatalf("Problem marshaling services object: %s", err)
	}
	fmt.Printf("%s\n", servicesJson)
	return err
}

// Add a service given a set of paramters.
func (cli *ServicedCli) CmdAddService(args ...string) error {
	cmd := Subcmd("add-service", "NAME POOLID IMAGEID COMMAND", "Add service.")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	if len(cmd.Args()) < 4 {
		cmd.Usage()
		return nil
	}
	controlPlane := getClient()

	service, err := serviced.NewService()
	if err != nil {
		log.Fatal("Could not create service:%v\n", err)
	}
	service.Name = cmd.Arg(0)
	service.PoolId = cmd.Arg(1)
	service.ImageId = cmd.Arg(2)
	startup := cmd.Arg(3)
	for i := 4; i < len(cmd.Args()); i++ {
		startup = startup + " " + cmd.Arg(i)
	}
	service.Startup = startup

	log.Printf("Calling AddService.\n")
	var unused int
	err = controlPlane.AddService(*service, &unused)
	if err != nil {
		log.Fatalf("Could not add services: %v", err)
	}
	fmt.Println(service.Id)
	return err
}

// Remove a service given the SERVICEID.
func (cli *ServicedCli) CmdRemoveService(args ...string) error {
	cmd := Subcmd("remove-service", "SERVICEID", "Remove a service.")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	if len(cmd.Args()) != 1 {
		cmd.Usage()
		return nil
	}
	controlPlane := getClient()

	var unused int
	err := controlPlane.RemoveService(cmd.Arg(0), &unused)
	if err != nil {
		log.Fatalf("Could not remove service: %v", err)
	}
	return err
}

// Schedule a service to start given a service id.
func (cli *ServicedCli) CmdStartService(args ...string) error {
	cmd := Subcmd("start-service", "SERVICEID", "Start a service.")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	if len(cmd.Args()) != 1 {
		cmd.Usage()
		return nil
	}
	controlPlane := getClient()
	var hostId string
	err := controlPlane.StartService(cmd.Arg(0), &hostId)
	if err != nil {
		log.Fatalf("Could not start service: %v", err)
	}
	log.Printf("Sevice scheduled to start on host %s\n", hostId)
	return err
}

// Schedule a service to stop given a service id.
func (cli *ServicedCli) CmdStopService(args ...string) error {
	cmd := Subcmd("stop-service", "SERVICEID", "Stop a service.")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	if len(cmd.Args()) != 1 {
		cmd.Usage()
		return nil
	}
	controlPlane := getClient()
	var unused int
	err := controlPlane.StopService(cmd.Arg(0), &unused)
	if err != nil {
		log.Fatalf("Could not stop service: %v", err)
	}
	log.Printf("Sevice scheduled to stop.")
	return err
}
<commit_msg>removed useless update-host command.<commit_after>/*******************************************************************************
* Copyright (C) Zenoss, Inc. 2013, all rights reserved.
*
* This content is made available according to terms specified in
* License.zenoss under the directory where your Zenoss product is installed.
*
*******************************************************************************/

package main

// This is here the command line arguments are parsed and executed.

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/zenoss/serviced"
	clientlib "github.com/zenoss/serviced/client"
	"log"
	"reflect"
	"strconv"
	"strings"
)

// A type to represent the CLI. All the command will have the same signature.
// This makes it easy to call them arbitrarily.
type ServicedCli struct{}

// A helper function that creates a subcommand
func Subcmd(name, signature, description string) *flag.FlagSet {
	flags := flag.NewFlagSet(name, flag.ContinueOnError)
	flags.Usage = func() {
		fmt.Printf("\nUsage: serviced %s %s\n\n%s\n\n", name, signature, description)
		flags.PrintDefaults()
	}
	return flags
}

// Use reflection to aquire the give method by name. For example to get method
// CmdFoo, pass 'foo'. A method is returned. The second return argument
// indicates if the argument was found.
func (cli *ServicedCli) getMethod(name string) (reflect.Method, bool) {

	// Contruct the method name to be CmdFoo, where foo was passed
	methodName := "Cmd"
	for _, part := range strings.Split(name, "-") {
		methodName = methodName + strings.ToUpper(part[:1]) + strings.ToLower(part[1:])
	}
	return reflect.TypeOf(cli).MethodByName(methodName)
}

// Construct a new command line parsing object.
func NewServicedCli() (s *ServicedCli) {
	return &ServicedCli{}
}

// Show usage of serviced command line options.
func (cli *ServicedCli) CmdHelp(args ...string) error {
	if len(args) > 0 {
		method, exists := cli.getMethod(args[0])
		if !exists {
			fmt.Println("Error: Command not found:", args[0])
		} else {
			method.Func.CallSlice([]reflect.Value{
				reflect.ValueOf(cli),
				reflect.ValueOf([]string{"--help"}),
			})[0].Interface()
			return nil
		}
	}
	help := fmt.Sprintf("Usage: serviced [OPTIONS] COMMAND [arg...]\n\nA container based service management system.\n\nCommands:\n")
	for _, command := range [][2]string{
		{"hosts", "Display hosts"},
		{"add-host", "Add a host"},
		{"remove-host", "Remove a host"},
		{"pools", "Show pools"},
		{"add-pool", "Add pool"},
		{"services", "Show services"},
		{"add-service", "Add a service"},
		{"remove-service", "Remote a service"},
		{"start-service", "Start a service"},
		{"stop-service", "Stop a service"},
	} {
		help += fmt.Sprintf("    %-30.30s%s\n", command[0], command[1])
	}
	fmt.Println(help)
	return nil
}

// Attempt to find the command give on the CLI by looking up the method on the
// CLI interface. If found, execute it. Otherwise show usage.
func ParseCommands(args ...string) error {
	cli := NewServicedCli()

	if len(args) > 0 {
		method, exists := cli.getMethod(args[0])
		if !exists {
			fmt.Println("Error: Command not found:", args[0])
			return cli.CmdHelp(args[1:]...)
		}
		ret := method.Func.CallSlice([]reflect.Value{
			reflect.ValueOf(cli),
			reflect.ValueOf(args[1:]),
		})[0].Interface()
		if ret == nil {
			return nil
		}
		return ret.(error)
	}
	return cli.CmdHelp(args...)
}

// Create a client to the control plane.
func getClient() (c serviced.ControlPlane) {
	// setup the client
	c, err := clientlib.NewControlClient(options.port)
	if err != nil {
		log.Fatalf("Could not create acontrol plane client %v", err)
	}
	return c
}

// List the hosts associated with the control plane.
func (cli *ServicedCli) CmdHosts(args ...string) error {

	cmd := Subcmd("hosts", "[OPTIONS]", "List hosts")
	if err := cmd.Parse(args); err != nil {
		return err
	}

	client := getClient()

	var hosts map[string]*serviced.Host
	request := serviced.EntityRequest{}
	err := client.GetHosts(request, &hosts)
	if err != nil {
		log.Fatalf("Could not get hosts %v", err)
	}
	hostsJson, err := json.MarshalIndent(hosts, " ", "  ")
	if err == nil {
		fmt.Printf("%s\n", hostsJson)
	}
	return err
}

// Add a host to the control plane given the host:port.
func (cli *ServicedCli) CmdAddHost(args ...string) error {

	cmd := Subcmd("add-host", "HOST:PORT RESOURCE_POOL", "Add host")
	if err := cmd.Parse(args); err != nil {
		return nil
	}

	if len(cmd.Args()) != 2 {
		cmd.Usage()
		return nil
	}

	client, err := clientlib.NewAgentClient(cmd.Arg(0))
	if err != nil {
		log.Fatalf("Could not create connection to host %s: %v", args[0], err)
	}

	var remoteHost serviced.Host
	err = client.GetInfo(0, &remoteHost)
	if err != nil {
		log.Fatalf("Could not get remote host info: %v", err)
	}
	remoteHost.PoolId = cmd.Arg(1)
	log.Printf("Got host info: %v", remoteHost)

	controlPlane := getClient()
	var unused int

	err = controlPlane.AddHost(remoteHost, &unused)
	if err != nil {
		log.Fatalf("Could not add host: %v", err)
	}
	fmt.Println(remoteHost.Id)
	return err
}

// This method removes the given host (by HOSTID) from the system.
func (cli *ServicedCli) CmdRemoveHost(args ...string) error {
	cmd := Subcmd("remove-host", "HOSTID", "Remove the host.")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	if len(cmd.Args()) != 1 {
		cmd.Usage()
		return nil
	}

	controlPlane := getClient()
	var unused int
	err := controlPlane.RemoveHost(cmd.Arg(0), &unused)
	if err != nil {
		log.Fatalf("Could not remove host: %v", err)
	}
	log.Printf("Host %s removed.", cmd.Arg(0))
	return err
}

// A convinience struct for printing to command line
type poolWithHost struct {
	serviced.ResourcePool
	Hosts []string
}

// Print a list of pools. Args are ignored.
func (cli *ServicedCli) CmdPools(args ...string) error {
	cmd := Subcmd("pools", "[OPTIONS]", "Display pools")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	controlPlane := getClient()
	request := serviced.EntityRequest{}
	var pools map[string]*serviced.ResourcePool
	err := controlPlane.GetResourcePools(request, &pools)
	if err != nil {
		log.Fatalf("Could not get resource pools: %v", err)
	}
	poolsWithHost := make(map[string]poolWithHost)
	for _, pool := range pools {

		// get pool hosts
		var poolHosts []*serviced.PoolHost
		err = controlPlane.GetHostsForResourcePool(pool.Id, &poolHosts)
		if err != nil {
			log.Fatalf("Could not get hosts for Pool %s: %v", pool.Id, err)
		}
		hosts := make([]string, len(poolHosts))
		for i, hostPool := range poolHosts {
			hosts[i] = hostPool.HostId
		}
		poolsWithHost[pool.Id] = poolWithHost{*pool, hosts}
	}
	poolsWithHostJson, err := json.MarshalIndent(poolsWithHost, " ", "  ")
	if err == nil {
		fmt.Printf("%s\n", poolsWithHostJson)
	}
	return err
}

// Add a new pool given some parameters.
func (cli *ServicedCli) CmdAddPool(args ...string) error {
	cmd := Subcmd("add-pool", "[options] POOLID CORE_LIMIT MEMORY_LIMIT PRIORITY", "Add resource pool")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	if len(cmd.Args()) < 4 {
		cmd.Usage()
		return nil
	}
	pool, _ := serviced.NewResourcePool(cmd.Arg(0))
	coreLimit, err := strconv.Atoi(cmd.Arg(1))
	if err != nil {
		log.Fatal("Bad core limit %s: %v", cmd.Arg(1), err)
	}
	pool.CoreLimit = coreLimit
	memoryLimit, err := strconv.Atoi(cmd.Arg(2))
	if err != nil {
		log.Fatal("Bad memory limit %s: %v", cmd.Arg(2), err)
	}
	pool.MemoryLimit = uint64(memoryLimit)
	controlPlane := getClient()
	var unused int
	err = controlPlane.AddResourcePool(*pool, &unused)
	if err != nil {
		log.Fatalf("Could not add resource pool: %v", err)
	}
	fmt.Printf("%s\n", pool.Id)
	return err
}

// Print the list of available services.
func (cli *ServicedCli) CmdServices(args ...string) error {
	cmd := Subcmd("services", "[CMD]", "Show services")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	controlPlane := getClient()
	request := serviced.EntityRequest{}
	var services []*serviced.Service
	err := controlPlane.GetServices(request, &services)
	if err != nil {
		log.Fatalf("Could not get services: %v", err)
	}
	servicesJson, err := json.MarshalIndent(services, " ", " ")
	if err != nil {
		log.Fatalf("Problem marshaling services object: %s", err)
	}
	fmt.Printf("%s\n", servicesJson)
	return err
}

// Add a service given a set of paramters.
func (cli *ServicedCli) CmdAddService(args ...string) error {
	cmd := Subcmd("add-service", "NAME POOLID IMAGEID COMMAND", "Add service.")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	if len(cmd.Args()) < 4 {
		cmd.Usage()
		return nil
	}
	controlPlane := getClient()

	service, err := serviced.NewService()
	if err != nil {
		log.Fatal("Could not create service:%v\n", err)
	}
	service.Name = cmd.Arg(0)
	service.PoolId = cmd.Arg(1)
	service.ImageId = cmd.Arg(2)
	startup := cmd.Arg(3)
	for i := 4; i < len(cmd.Args()); i++ {
		startup = startup + " " + cmd.Arg(i)
	}
	service.Startup = startup

	log.Printf("Calling AddService.\n")
	var unused int
	err = controlPlane.AddService(*service, &unused)
	if err != nil {
		log.Fatalf("Could not add services: %v", err)
	}
	fmt.Println(service.Id)
	return err
}

// Remove a service given the SERVICEID.
func (cli *ServicedCli) CmdRemoveService(args ...string) error {
	cmd := Subcmd("remove-service", "SERVICEID", "Remove a service.")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	if len(cmd.Args()) != 1 {
		cmd.Usage()
		return nil
	}
	controlPlane := getClient()

	var unused int
	err := controlPlane.RemoveService(cmd.Arg(0), &unused)
	if err != nil {
		log.Fatalf("Could not remove service: %v", err)
	}
	return err
}

// Schedule a service to start given a service id.
func (cli *ServicedCli) CmdStartService(args ...string) error {
	cmd := Subcmd("start-service", "SERVICEID", "Start a service.")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	if len(cmd.Args()) != 1 {
		cmd.Usage()
		return nil
	}
	controlPlane := getClient()
	var hostId string
	err := controlPlane.StartService(cmd.Arg(0), &hostId)
	if err != nil {
		log.Fatalf("Could not start service: %v", err)
	}
	log.Printf("Sevice scheduled to start on host %s\n", hostId)
	return err
}

// Schedule a service to stop given a service id.
func (cli *ServicedCli) CmdStopService(args ...string) error {
	cmd := Subcmd("stop-service", "SERVICEID", "Stop a service.")
	if err := cmd.Parse(args); err != nil {
		return nil
	}
	if len(cmd.Args()) != 1 {
		cmd.Usage()
		return nil
	}
	controlPlane := getClient()
	var unused int
	err := controlPlane.StopService(cmd.Arg(0), &unused)
	if err != nil {
		log.Fatalf("Could not stop service: %v", err)
	}
	log.Printf("Sevice scheduled to stop.")
	return err
}
<|endoftext|>


<commit_before>package backrestservice

/*
Copyright 2018 Crunchy Data Solutions, Inc.
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.
*/

import (
	"errors"
	log "github.com/Sirupsen/logrus"
	crv1 "github.com/crunchydata/postgres-operator/apis/cr/v1"
	"github.com/crunchydata/postgres-operator/apiserver"
	msgs "github.com/crunchydata/postgres-operator/apiservermsgs"
	"github.com/crunchydata/postgres-operator/kubeapi"
	"github.com/crunchydata/postgres-operator/util"
	"k8s.io/api/core/v1"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const backrestCommand = "pgbackrest"
const backrestStanza = "--stanza=db"
const backrestInfoCommand = "info"
const containername = "database"

//  CreateBackup ...
// pgo backrest mycluster
// pgo backrest --selector=name=mycluster
func CreateBackup(request *msgs.CreateBackrestBackupRequest) msgs.CreateBackrestBackupResponse {
	resp := msgs.CreateBackrestBackupResponse{}
	resp.Status.Code = msgs.Ok
	resp.Status.Msg = ""
	resp.Results = make([]string, 0)

	if request.Selector != "" {
		//use the selector instead of an argument list to filter on

		clusterList := crv1.PgclusterList{}

		err := kubeapi.GetpgclustersBySelector(apiserver.RESTClient, &clusterList, request.Selector, apiserver.Namespace)
		if err != nil {
			resp.Status.Code = msgs.Error
			resp.Status.Msg = err.Error()
			return resp
		}

		if len(clusterList.Items) == 0 {
			log.Debug("no clusters found")
			resp.Results = append(resp.Results, "no clusters found with that selector")
			return resp
		} else {
			newargs := make([]string, 0)
			for _, cluster := range clusterList.Items {
				newargs = append(newargs, cluster.Spec.Name)
			}
			request.Args = newargs
		}

	}

	for _, clusterName := range request.Args {
		log.Debug("create backrestbackup called for " + clusterName)
		taskName := clusterName + "-backrest-backup"

		cluster := crv1.Pgcluster{}
		found, err := kubeapi.Getpgcluster(apiserver.RESTClient, &cluster, clusterName, apiserver.Namespace)
		if !found {
			resp.Status.Code = msgs.Error
			resp.Status.Msg = clusterName + " was not found, verify cluster name"
			return resp
		} else if err != nil {
			resp.Status.Code = msgs.Error
			resp.Status.Msg = err.Error()
			return resp
		}

		result := crv1.Pgtask{}

		// error if it already exists
		found, err = kubeapi.Getpgtask(apiserver.RESTClient, &result, taskName, apiserver.Namespace)
		if !found {
			log.Debug("backrest backup pgtask " + taskName + " not found so we create it")
		} else if err != nil {
			resp.Results = append(resp.Results, "error getting pgtask for "+taskName)
			break
		} else {
			log.Debug("pgtask " + taskName + " was found so we recreate it")
			//remove the existing pgtask
			err := kubeapi.Deletepgtask(apiserver.RESTClient, taskName, apiserver.Namespace)
			if err != nil {
				resp.Status.Code = msgs.Error
				resp.Status.Msg = err.Error()
				return resp
			}

			//remove any previous backup job
			removeBackupJob(taskName)
		}

		//get pod name from cluster
		var podname string
		podname, err = getPrimaryPodName(&cluster)

		if err != nil {
			log.Error(err)
			resp.Status.Code = msgs.Error
			resp.Status.Msg = err.Error()
			return resp
		}

		err = kubeapi.Createpgtask(apiserver.RESTClient, getBackupParams(clusterName, taskName, crv1.PgtaskBackrestBackup, podname, "database"), apiserver.Namespace)
		if err != nil {
			resp.Status.Code = msgs.Error
			resp.Status.Msg = err.Error()
			return resp
		}
		resp.Results = append(resp.Results, "created Pgtask "+taskName)

	}

	return resp
}

func getBackupParams(clusterName, taskName, action, podName, containerName string) *crv1.Pgtask {
	var newInstance *crv1.Pgtask

	spec := crv1.PgtaskSpec{}
	spec.Name = taskName
	spec.TaskType = crv1.PgtaskBackrest
	spec.Parameters = make(map[string]string)
	spec.Parameters[util.LABEL_PG_CLUSTER] = clusterName
	spec.Parameters[util.LABEL_POD_NAME] = podName
	spec.Parameters[util.LABEL_CONTAINER_NAME] = containerName
	spec.Parameters[util.LABEL_BACKREST_COMMAND] = action

	newInstance = &crv1.Pgtask{
		ObjectMeta: meta_v1.ObjectMeta{
			Name: taskName,
		},
		Spec: spec,
	}
	return newInstance
}

func removeBackupJob(name string) {

	_, found := kubeapi.GetJob(apiserver.Clientset, name, apiserver.Namespace)
	if !found {
		return
	}

	log.Debugf("found backrest backup job %s will remove\n", name)

	kubeapi.DeleteJob(apiserver.Clientset, name, apiserver.Namespace)
}

func getPrimaryPodName(cluster *crv1.Pgcluster) (string, error) {
	var podname string

	selector := util.LABEL_PGPOOL + "!=true," + util.LABEL_PG_CLUSTER + "=" + cluster.Spec.Name + "," + util.LABEL_PRIMARY + "=true"

	pods, err := kubeapi.GetPods(apiserver.Clientset, selector, apiserver.Namespace)
	if err != nil {
		return podname, err
	}

	for _, p := range pods.Items {
		if isPrimary(&p) && isReady(&p) {
			return p.Name, err
		}
	}

	return podname, errors.New("primary pod is not in Ready state")
}

func isPrimary(pod *v1.Pod) bool {
	if pod.ObjectMeta.Labels[util.LABEL_PRIMARY] == "true" {
		return true
	}
	return false

}

func isReady(pod *v1.Pod) bool {
	readyCount := 0
	containerCount := 0
	for _, stat := range pod.Status.ContainerStatuses {
		containerCount++
		if stat.Ready {
			readyCount++
		}
	}
	if readyCount != containerCount {
		return false
	}
	return true

}

// ShowBackrest ...
func ShowBackrest(name, selector string) msgs.ShowBackrestResponse {
	var err error

	response := msgs.ShowBackrestResponse{}
	response.Status = msgs.Status{Code: msgs.Ok, Msg: ""}
	response.Items = make([]msgs.ShowBackrestDetail, 0)

	if selector == "" && name == "all" {
	} else {
		if selector == "" {
			selector = "name=" + name
		}
	}

	clusterList := crv1.PgclusterList{}

	//get a list of all clusters
	err = kubeapi.GetpgclustersBySelector(apiserver.RESTClient,
		&clusterList, selector, apiserver.Namespace)
	if err != nil {
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	log.Debug("clusters found len is %d\n", len(clusterList.Items))

	for _, c := range clusterList.Items {
		detail := msgs.ShowBackrestDetail{}
		detail.Name = c.Name

		podname, err := getPrimaryPodName(&c)

		if err != nil {
			log.Error(err)
			response.Status.Code = msgs.Error
			response.Status.Msg = err.Error()
			return response
		}

		//here is where we would exec to get the backrest info
		info, err := getInfo(c.Name, podname)
		if err != nil {
			detail.Info = err.Error()
		} else {
			detail.Info = info
		}

		response.Items = append(response.Items, detail)
	}

	return response

}

func getInfo(clusterName, podname string) (string, error) {

	var err error

	cmd := make([]string, 0)

	log.Info("backrest info command requested")
	//pgbackrest --stanza=db info
	cmd = append(cmd, backrestCommand)
	cmd = append(cmd, backrestStanza)
	cmd = append(cmd, backrestInfoCommand)

	log.Infof("command is %v ", cmd)
	output, stderr, err := kubeapi.ExecToPodThroughAPI(apiserver.RESTConfig, apiserver.Clientset, cmd, containername, podname, apiserver.Namespace, nil)
	log.Info("output=[" + output + "]")
	log.Info("stderr=[" + stderr + "]")

	if err != nil {
		log.Error(err)
		return "", err
	}
	log.Debug("backrest info ends")
	return output, err

}

//  Restore ...
// pgo restore mycluster --to-cluster=restored
func Restore(request *msgs.RestoreRequest) msgs.RestoreResponse {
	resp := msgs.RestoreResponse{}
	resp.Status.Code = msgs.Ok
	resp.Status.Msg = ""
	resp.Results = make([]string, 0)

	log.Debugf("Restore %v\n", request)

	cluster := crv1.Pgcluster{}
	found, err := kubeapi.Getpgcluster(apiserver.RESTClient, &cluster, request.FromCluster, apiserver.Namespace)
	if !found {
		resp.Status.Code = msgs.Error
		resp.Status.Msg = request.FromCluster + " was not found, verify cluster name"
		return resp
	} else if err != nil {
		resp.Status.Code = msgs.Error
		resp.Status.Msg = err.Error()
		return resp
	}

	//verify that the cluster we are restoring from has backrest enabled
	if cluster.Spec.UserLabels[util.LABEL_BACKREST] != "true" {
		resp.Status.Code = msgs.Error
		resp.Status.Msg = "can't restore, cluster restoring from does not have backrest enabled"
		return resp
	}

	pgtask := getRestoreParams(request)
	existingTask := crv1.Pgtask{}

	//delete any existing pgtask with the same name
	found, err = kubeapi.Getpgtask(apiserver.RESTClient,
		&existingTask,
		pgtask.Name,
		apiserver.Namespace)
	if found {
		log.Debug("deleting prior pgtask " + pgtask.Name)
		err = kubeapi.Deletepgtask(apiserver.RESTClient,
			pgtask.Name,
			apiserver.Namespace)
		if err != nil {
			resp.Status.Code = msgs.Error
			resp.Status.Msg = err.Error()
			return resp
		}
	}

	//create a pgtask for the restore workflow
	err = kubeapi.Createpgtask(apiserver.RESTClient,
		pgtask,
		apiserver.Namespace)
	if err != nil {
		resp.Status.Code = msgs.Error
		resp.Status.Msg = err.Error()
		return resp
	}

	resp.Results = append(resp.Results, "restore performed on "+request.FromCluster+" to "+request.ToCluster+" type="+request.RestoreType)

	return resp
}

func getRestoreParams(request *msgs.RestoreRequest) *crv1.Pgtask {
	var newInstance *crv1.Pgtask

	spec := crv1.PgtaskSpec{}
	spec.Name = "backrest-restore-" + request.FromCluster + "-to-" + request.ToCluster
	spec.TaskType = crv1.PgtaskBackrestRestore
	spec.Parameters = make(map[string]string)
	spec.Parameters[util.LABEL_BACKREST_RESTORE_FROM_CLUSTER] = request.FromCluster
	spec.Parameters[util.LABEL_BACKREST_RESTORE_TO_CLUSTER] = request.ToCluster
	spec.Parameters[util.LABEL_BACKREST_RESTORE_TYPE] = request.RestoreType
	spec.Parameters[util.LABEL_BACKREST_RESTORE_PITR_TARGET] = request.PITRTarget

	newInstance = &crv1.Pgtask{
		ObjectMeta: meta_v1.ObjectMeta{
			Name: spec.Name,
		},
		Spec: spec,
	}
	return newInstance
}
<commit_msg>fix backrest job removal<commit_after>package backrestservice

/*
Copyright 2018 Crunchy Data Solutions, Inc.
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.
*/

import (
	"errors"
	log "github.com/Sirupsen/logrus"
	crv1 "github.com/crunchydata/postgres-operator/apis/cr/v1"
	"github.com/crunchydata/postgres-operator/apiserver"
	msgs "github.com/crunchydata/postgres-operator/apiservermsgs"
	"github.com/crunchydata/postgres-operator/kubeapi"
	"github.com/crunchydata/postgres-operator/util"
	"k8s.io/api/core/v1"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	//"time"
)

const backrestCommand = "pgbackrest"
const backrestStanza = "--stanza=db"
const backrestInfoCommand = "info"
const containername = "database"

//  CreateBackup ...
// pgo backrest mycluster
// pgo backrest --selector=name=mycluster
func CreateBackup(request *msgs.CreateBackrestBackupRequest) msgs.CreateBackrestBackupResponse {
	resp := msgs.CreateBackrestBackupResponse{}
	resp.Status.Code = msgs.Ok
	resp.Status.Msg = ""
	resp.Results = make([]string, 0)

	if request.Selector != "" {
		//use the selector instead of an argument list to filter on

		clusterList := crv1.PgclusterList{}

		err := kubeapi.GetpgclustersBySelector(apiserver.RESTClient, &clusterList, request.Selector, apiserver.Namespace)
		if err != nil {
			resp.Status.Code = msgs.Error
			resp.Status.Msg = err.Error()
			return resp
		}

		if len(clusterList.Items) == 0 {
			log.Debug("no clusters found")
			resp.Results = append(resp.Results, "no clusters found with that selector")
			return resp
		} else {
			newargs := make([]string, 0)
			for _, cluster := range clusterList.Items {
				newargs = append(newargs, cluster.Spec.Name)
			}
			request.Args = newargs
		}

	}

	for _, clusterName := range request.Args {
		log.Debug("create backrestbackup called for " + clusterName)
		taskName := clusterName + "-backrest-backup"

		cluster := crv1.Pgcluster{}
		found, err := kubeapi.Getpgcluster(apiserver.RESTClient, &cluster, clusterName, apiserver.Namespace)
		if !found {
			resp.Status.Code = msgs.Error
			resp.Status.Msg = clusterName + " was not found, verify cluster name"
			return resp
		} else if err != nil {
			resp.Status.Code = msgs.Error
			resp.Status.Msg = err.Error()
			return resp
		}

		result := crv1.Pgtask{}

		// error if it already exists
		found, err = kubeapi.Getpgtask(apiserver.RESTClient, &result, taskName, apiserver.Namespace)
		if !found {
			log.Debug("backrest backup pgtask " + taskName + " not found so we create it")
		} else if err != nil {

			resp.Results = append(resp.Results, "error getting pgtask for "+taskName)
			break
		} else {

			log.Debug("pgtask " + taskName + " was found so we recreate it")
			//remove the existing pgtask
			err := kubeapi.Deletepgtask(apiserver.RESTClient, taskName, apiserver.Namespace)
			if err != nil {
				resp.Status.Code = msgs.Error
				resp.Status.Msg = err.Error()
				return resp
			}

			//remove any previous backup job

			kubeapi.DeleteJobs(apiserver.Clientset, util.LABEL_PG_CLUSTER+"="+clusterName+","+util.LABEL_BACKREST+"=true", apiserver.Namespace)
			//time.Sleep(time.Seconds * 2)
		}

		//get pod name from cluster
		var podname string
		podname, err = getPrimaryPodName(&cluster)

		if err != nil {
			log.Error(err)
			resp.Status.Code = msgs.Error
			resp.Status.Msg = err.Error()
			return resp
		}

		err = kubeapi.Createpgtask(apiserver.RESTClient, getBackupParams(clusterName, taskName, crv1.PgtaskBackrestBackup, podname, "database"), apiserver.Namespace)
		if err != nil {
			resp.Status.Code = msgs.Error
			resp.Status.Msg = err.Error()
			return resp
		}
		resp.Results = append(resp.Results, "created Pgtask "+taskName)

	}

	return resp
}

func getBackupParams(clusterName, taskName, action, podName, containerName string) *crv1.Pgtask {
	var newInstance *crv1.Pgtask

	spec := crv1.PgtaskSpec{}
	spec.Name = taskName
	spec.TaskType = crv1.PgtaskBackrest
	spec.Parameters = make(map[string]string)
	spec.Parameters[util.LABEL_PG_CLUSTER] = clusterName
	spec.Parameters[util.LABEL_POD_NAME] = podName
	spec.Parameters[util.LABEL_CONTAINER_NAME] = containerName
	spec.Parameters[util.LABEL_BACKREST_COMMAND] = action

	newInstance = &crv1.Pgtask{
		ObjectMeta: meta_v1.ObjectMeta{
			Name: taskName,
		},
		Spec: spec,
	}
	return newInstance
}

func removeBackupJob(clusterName string) {

}

func getPrimaryPodName(cluster *crv1.Pgcluster) (string, error) {
	var podname string

	selector := util.LABEL_PGPOOL + "!=true," + util.LABEL_PG_CLUSTER + "=" + cluster.Spec.Name + "," + util.LABEL_PRIMARY + "=true"

	pods, err := kubeapi.GetPods(apiserver.Clientset, selector, apiserver.Namespace)
	if err != nil {
		return podname, err
	}

	for _, p := range pods.Items {
		if isPrimary(&p) && isReady(&p) {
			return p.Name, err
		}
	}

	return podname, errors.New("primary pod is not in Ready state")
}

func isPrimary(pod *v1.Pod) bool {
	if pod.ObjectMeta.Labels[util.LABEL_PRIMARY] == "true" {
		return true
	}
	return false

}

func isReady(pod *v1.Pod) bool {
	readyCount := 0
	containerCount := 0
	for _, stat := range pod.Status.ContainerStatuses {
		containerCount++
		if stat.Ready {
			readyCount++
		}
	}
	if readyCount != containerCount {
		return false
	}
	return true

}

// ShowBackrest ...
func ShowBackrest(name, selector string) msgs.ShowBackrestResponse {
	var err error

	response := msgs.ShowBackrestResponse{}
	response.Status = msgs.Status{Code: msgs.Ok, Msg: ""}
	response.Items = make([]msgs.ShowBackrestDetail, 0)

	if selector == "" && name == "all" {
	} else {
		if selector == "" {
			selector = "name=" + name
		}
	}

	clusterList := crv1.PgclusterList{}

	//get a list of all clusters
	err = kubeapi.GetpgclustersBySelector(apiserver.RESTClient,
		&clusterList, selector, apiserver.Namespace)
	if err != nil {
		response.Status.Code = msgs.Error
		response.Status.Msg = err.Error()
		return response
	}

	log.Debug("clusters found len is %d\n", len(clusterList.Items))

	for _, c := range clusterList.Items {
		detail := msgs.ShowBackrestDetail{}
		detail.Name = c.Name

		podname, err := getPrimaryPodName(&c)

		if err != nil {
			log.Error(err)
			response.Status.Code = msgs.Error
			response.Status.Msg = err.Error()
			return response
		}

		//here is where we would exec to get the backrest info
		info, err := getInfo(c.Name, podname)
		if err != nil {
			detail.Info = err.Error()
		} else {
			detail.Info = info
		}

		response.Items = append(response.Items, detail)
	}

	return response

}

func getInfo(clusterName, podname string) (string, error) {

	var err error

	cmd := make([]string, 0)

	log.Info("backrest info command requested")
	//pgbackrest --stanza=db info
	cmd = append(cmd, backrestCommand)
	cmd = append(cmd, backrestStanza)
	cmd = append(cmd, backrestInfoCommand)

	log.Infof("command is %v ", cmd)
	output, stderr, err := kubeapi.ExecToPodThroughAPI(apiserver.RESTConfig, apiserver.Clientset, cmd, containername, podname, apiserver.Namespace, nil)
	log.Info("output=[" + output + "]")
	log.Info("stderr=[" + stderr + "]")

	if err != nil {
		log.Error(err)
		return "", err
	}
	log.Debug("backrest info ends")
	return output, err

}

//  Restore ...
// pgo restore mycluster --to-cluster=restored
func Restore(request *msgs.RestoreRequest) msgs.RestoreResponse {
	resp := msgs.RestoreResponse{}
	resp.Status.Code = msgs.Ok
	resp.Status.Msg = ""
	resp.Results = make([]string, 0)

	log.Debugf("Restore %v\n", request)

	cluster := crv1.Pgcluster{}
	found, err := kubeapi.Getpgcluster(apiserver.RESTClient, &cluster, request.FromCluster, apiserver.Namespace)
	if !found {
		resp.Status.Code = msgs.Error
		resp.Status.Msg = request.FromCluster + " was not found, verify cluster name"
		return resp
	} else if err != nil {
		resp.Status.Code = msgs.Error
		resp.Status.Msg = err.Error()
		return resp
	}

	//verify that the cluster we are restoring from has backrest enabled
	if cluster.Spec.UserLabels[util.LABEL_BACKREST] != "true" {
		resp.Status.Code = msgs.Error
		resp.Status.Msg = "can't restore, cluster restoring from does not have backrest enabled"
		return resp
	}

	pgtask := getRestoreParams(request)
	existingTask := crv1.Pgtask{}

	//delete any existing pgtask with the same name
	found, err = kubeapi.Getpgtask(apiserver.RESTClient,
		&existingTask,
		pgtask.Name,
		apiserver.Namespace)
	if found {
		log.Debug("deleting prior pgtask " + pgtask.Name)
		err = kubeapi.Deletepgtask(apiserver.RESTClient,
			pgtask.Name,
			apiserver.Namespace)
		if err != nil {
			resp.Status.Code = msgs.Error
			resp.Status.Msg = err.Error()
			return resp
		}
	}

	//create a pgtask for the restore workflow
	err = kubeapi.Createpgtask(apiserver.RESTClient,
		pgtask,
		apiserver.Namespace)
	if err != nil {
		resp.Status.Code = msgs.Error
		resp.Status.Msg = err.Error()
		return resp
	}

	resp.Results = append(resp.Results, "restore performed on "+request.FromCluster+" to "+request.ToCluster+" type="+request.RestoreType)

	return resp
}

func getRestoreParams(request *msgs.RestoreRequest) *crv1.Pgtask {
	var newInstance *crv1.Pgtask

	spec := crv1.PgtaskSpec{}
	spec.Name = "backrest-restore-" + request.FromCluster + "-to-" + request.ToCluster
	spec.TaskType = crv1.PgtaskBackrestRestore
	spec.Parameters = make(map[string]string)
	spec.Parameters[util.LABEL_BACKREST_RESTORE_FROM_CLUSTER] = request.FromCluster
	spec.Parameters[util.LABEL_BACKREST_RESTORE_TO_CLUSTER] = request.ToCluster
	spec.Parameters[util.LABEL_BACKREST_RESTORE_TYPE] = request.RestoreType
	spec.Parameters[util.LABEL_BACKREST_RESTORE_PITR_TARGET] = request.PITRTarget

	newInstance = &crv1.Pgtask{
		ObjectMeta: meta_v1.ObjectMeta{
			Name: spec.Name,
		},
		Spec: spec,
	}
	return newInstance
}
<|endoftext|>


<commit_before>package integration_test

import (
	"github.com/APTrust/exchange/constants"
	"github.com/APTrust/exchange/models"
	"github.com/APTrust/exchange/util"
	"github.com/APTrust/exchange/util/storage"
	"github.com/APTrust/exchange/util/testutil"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"path/filepath"
	"testing"
)

/*
These tests check the results of the integration tests for
the app apt_record. See the ingest_test.sh script in
the scripts folder, which sets up an integration context, runs
the apt_record.
*/

func TestRecordResults(t *testing.T) {
	if !testutil.ShouldRunIntegrationTests() {
		t.Skip("Skipping integration test. Set ENV var RUN_EXCHANGE_INTEGRATION=true if you want to run them.")
	}
	// Load config
	configFile := filepath.Join("config", "integration.json")
	config, err := models.LoadConfigFile(configFile)
	require.Nil(t, err)
	config.ExpandFilePaths()

	// Find the log file that apt_record created when it was running
	// with the "config/integration.json" config options. We'll read
	// that file.
	pathToJsonLog := filepath.Join(config.LogDirectory, "apt_record.json")
	for _, bagName := range testutil.INTEGRATION_GOOD_BAGS {
		ingestManifest, err := testutil.FindIngestManifestInLog(pathToJsonLog, bagName)
		assert.Nil(t, err)
		if err != nil {
			continue
		}
		// TODO: Test WorkItem (stage, status, etc.) below.
		recordTestCommon(t, bagName, ingestManifest)
	}
}

func recordTestCommon(t *testing.T, bagName string, ingestManifest *models.IngestManifest) {
	// Test some basic object properties
	assert.NotEmpty(t, ingestManifest.WorkItemId, "WorkItemId should not be empty for %s", bagName)
	assert.NotEmpty(t, ingestManifest.S3Bucket, "S3Bucket should not be empty for %s", bagName)
	assert.NotEmpty(t, ingestManifest.S3Key, "S3Key should not be empty for %s", bagName)
	assert.NotEmpty(t, ingestManifest.ETag, "ETag should not be empty for %s", bagName)

	// Make sure the result has some basic info in RecordResult
	assert.True(t, ingestManifest.RecordResult.Attempted,
		"RecordResult.Attempted should be true for %s", bagName)
	assert.True(t, ingestManifest.RecordResult.AttemptNumber > 0,
		"RecordResult.AttemptNumber should be > 0 %s", bagName)
	assert.NotEmpty(t, ingestManifest.RecordResult.StartedAt,
		"RecordResult.StartedAt should not be empty for %s", bagName)
	assert.NotEmpty(t, ingestManifest.RecordResult.FinishedAt,
		"RecordResult.FinishedAt should not be empty for %s", bagName)
	assert.Empty(t, ingestManifest.RecordResult.Errors,
		"RecordResult.Errors should be empty for %s", bagName)
	assert.True(t, ingestManifest.RecordResult.Retry,
		"RecordResult.Retry should be true for %s", bagName)

	// Make sure the result has some basic info in CleanupResult
	assert.True(t, ingestManifest.CleanupResult.Attempted,
		"CleanupResult.Attempted should be true for %s", bagName)
	assert.True(t, ingestManifest.CleanupResult.AttemptNumber > 0,
		"CleanupResult.AttemptNumber should be > 0 %s", bagName)
	assert.NotEmpty(t, ingestManifest.CleanupResult.StartedAt,
		"CleanupResult.StartedAt should not be empty for %s", bagName)
	assert.NotEmpty(t, ingestManifest.CleanupResult.FinishedAt,
		"CleanupResult.FinishedAt should not be empty for %s", bagName)
	assert.Empty(t, ingestManifest.CleanupResult.Errors,
		"CleanupResult.Errors should be empty for %s", bagName)
	assert.True(t, ingestManifest.CleanupResult.Retry,
		"CleanupResult.Retry should be true for %s", bagName)

	// Make sure our IntellectualObject got all of its PremisEvents
	//obj := ingestManifest.Object
	db, err := storage.NewBoltDB(ingestManifest.DBPath)
	require.Nil(t, err)
	obj, err := db.GetIntellectualObject(db.ObjectIdentifier())
	require.Nil(t, err)
	require.Equal(t, 4, len(obj.PremisEvents))

	// Make sure this item was deleted from the receiving bucket
	// after ingest completed.
	assert.False(t, obj.IngestDeletedFromReceivingAt.IsZero(),
		"Object %s was not deleted from receiving bucket", bagName)
	assert.Empty(t, obj.IngestErrorMessage)

	// Check the object-level events
	creationEvents := obj.FindEventsByType(constants.EventCreation)
	idEvents := obj.FindEventsByType(constants.EventIdentifierAssignment)
	ingestEvents := obj.FindEventsByType(constants.EventIngestion)
	accessEvents := obj.FindEventsByType(constants.EventAccessAssignment)
	assert.Equal(t, 1, len(accessEvents), "Missing access event for %s", bagName)
	assert.Equal(t, 1, len(creationEvents), "Missing creation event for %s", bagName)
	assert.Equal(t, 1, len(idEvents), "Missing identifier assignment event for %s", bagName)
	assert.Equal(t, 1, len(ingestEvents), "Missing ingest event for %s", bagName)

	for _, event := range obj.PremisEvents {
		assert.True(t, event.Id > 0, "Event %s was not saved for %s", event.EventType, obj.Identifier)
		assert.True(t, event.IntellectualObjectId > 0,
			"event.IntellectualObjectId not set for %s %s", event.EventType, obj.Identifier)
		assert.False(t, event.DateTime.IsZero(),
			"event.DateTime was not set for %s %s", event.EventType, obj.Identifier)
		assert.False(t, event.CreatedAt.IsZero(),
			"event.CreatedAt was not set for %s %s", event.EventType, obj.Identifier)
		assert.False(t, event.UpdatedAt.IsZero(),
			"event.UpdatedAt was not set for %s %s", event.EventType, obj.Identifier)

		assert.True(t, util.LooksLikeUUID(event.Identifier),
			"Identifier for %s %s doesn't look like a UUID", event.EventType, obj.Identifier)
		assert.NotEmpty(t, event.EventType, "EventType missing for %s %s", obj.Identifier, event.Identifier)
		assert.NotEmpty(t, event.Detail, "Detail is empty for %s %s", event.EventType, obj.Identifier)
		assert.NotEmpty(t, event.Outcome, "Outcome is empty for %s %s", event.EventType, obj.Identifier)
		assert.NotEmpty(t, event.OutcomeDetail,
			"OutcomeDetail is empty for %s %s", event.EventType, obj.Identifier)
		assert.NotEmpty(t, event.Object, "Object is empty for %s %s", event.EventType, obj.Identifier)
		assert.NotEmpty(t, event.Agent, "Agent is empty for %s %s", event.EventType, obj.Identifier)
		assert.NotEmpty(t, event.OutcomeInformation,
			"OutcomeInformation is empty for %s %s", event.EventType, obj.Identifier)
		assert.Equal(t, obj.Identifier, event.IntellectualObjectIdentifier,
			"IntellectualObjectIdentifier is wrong for %s %s", event.EventType, obj.Identifier)
	}

	for _, gfIdentifier := range db.FileIdentifiers() {
		gf, err := db.GetGenericFile(gfIdentifier)
		require.Nil(t, err, gfIdentifier)

		// Skip these checks for files that didn't need to be saved.
		// Reasons for not needing to be saved:
		// 1. File has a non-savable name, according to util.HasSavableName
		// 2. File has not changed since last time we ingested this bag.
		if !gf.IngestNeedsSave {
			continue
		}

		// Make sure checksums are present
		require.Equal(t, 2, len(gf.Checksums),
			"Checksums should be %s, found %d for %s", 2, len(gf.Checksums), gf.Identifier)
		md5 := gf.GetChecksumByAlgorithm(constants.AlgMd5)
		sha256 := gf.GetChecksumByAlgorithm(constants.AlgSha256)
		require.NotNil(t, md5, "Missing md5 digest for for %s", gf.Identifier)
		require.NotNil(t, sha256, "Missing sha256 digest for for %s", gf.Identifier)

		// Make sure that these checksums were saved
		assert.True(t, md5.Id > 0, "md5 was not saved for %s", gf.Identifier)
		assert.True(t, md5.GenericFileId > 0, "md5.GenericFileId not set for %s", gf.Identifier)
		assert.False(t, md5.CreatedAt.IsZero(), "md5.CreatedAt was not set for %s", gf.Identifier)
		assert.False(t, md5.UpdatedAt.IsZero(), "md5.UpdatedAt was not set for %s", gf.Identifier)

		assert.True(t, sha256.Id > 0, "sha256 was not saved for %s", gf.Identifier)
		assert.True(t, sha256.GenericFileId > 0, "sha256.GenericFileId not set for %s", gf.Identifier)
		assert.False(t, sha256.CreatedAt.IsZero(), "sha256.CreatedAt was not set for %s", gf.Identifier)
		assert.False(t, sha256.UpdatedAt.IsZero(), "sha256.UpdatedAt was not set for %s", gf.Identifier)

		// Make sure PremisEvents are present
		require.Equal(t, 6, len(gf.PremisEvents),
			"PremisEvents count should be %s, found %d for %s", 6, len(gf.PremisEvents), gf.Identifier)
		assert.Equal(t, 1, len(gf.FindEventsByType(constants.EventFixityCheck)),
			"Missing fixity check event for %s", gf.Identifier)
		assert.Equal(t, 1, len(gf.FindEventsByType(constants.EventDigestCalculation)),
			"Missing digest calculation event for %s", gf.Identifier)
		assert.Equal(t, 2, len(gf.FindEventsByType(constants.EventIdentifierAssignment)),
			"Missing identifier assignment event(s) for %s", gf.Identifier)
		assert.Equal(t, 1, len(gf.FindEventsByType(constants.EventReplication)),
			"Missing replication event for %s", gf.Identifier)
		assert.Equal(t, 1, len(gf.FindEventsByType(constants.EventIngestion)),
			"Missing ingestion event for %s", gf.Identifier)

		for _, event := range gf.PremisEvents {
			assert.True(t, event.Id > 0, "Event %s was not saved for %s", event.EventType, gf.Identifier)
			assert.True(t, event.IntellectualObjectId > 0,
				"event.IntellectualObjectId not set for %s %s", event.EventType, gf.Identifier)
			assert.True(t, event.GenericFileId > 0,
				"event.GenericFileId not set for %s %s", event.EventType, gf.Identifier)
			assert.False(t, event.DateTime.IsZero(),
				"event.DateTime was not set for %s %s", event.EventType, gf.Identifier)
			assert.False(t, event.CreatedAt.IsZero(),
				"event.CreatedAt was not set for %s %s", event.EventType, gf.Identifier)
			assert.False(t, event.UpdatedAt.IsZero(),
				"event.UpdatedAt was not set for %s %s", event.EventType, gf.Identifier)

			assert.True(t, util.LooksLikeUUID(event.Identifier),
				"Identifier for %s %s doesn't look like a UUID", event.EventType, gf.Identifier)
			assert.NotEmpty(t, event.EventType, "EventType missing for %s %s", gf.Identifier, event.Identifier)
			assert.NotEmpty(t, event.Detail, "Detail is empty for %s %s", event.EventType, gf.Identifier)
			assert.NotEmpty(t, event.Outcome, "Outcome is empty for %s %s", event.EventType, gf.Identifier)
			assert.NotEmpty(t, event.OutcomeDetail,
				"OutcomeDetail is empty for %s %s", event.EventType, gf.Identifier)
			assert.NotEmpty(t, event.Object, "Object is empty for %s %s", event.EventType, gf.Identifier)
			assert.NotEmpty(t, event.Agent, "Agent is empty for %s %s", event.EventType, gf.Identifier)
			assert.NotEmpty(t, event.OutcomeInformation,
				"OutcomeInformation is empty for %s %s", event.EventType, gf.Identifier)
			assert.Equal(t, obj.Identifier, event.IntellectualObjectIdentifier,
				"IntellectualObjectIdentifier is wrong for %s %s", event.EventType, gf.Identifier)
			assert.Equal(t, gf.Identifier, event.GenericFileIdentifier,
				"GenericFileIdentifier is wrong for %s %s", event.EventType, gf.Identifier)
		}
	}
}
<commit_msg>Added glacier-only bags to apt_record_post_test<commit_after>package integration_test

import (
	"github.com/APTrust/exchange/constants"
	"github.com/APTrust/exchange/models"
	"github.com/APTrust/exchange/util"
	"github.com/APTrust/exchange/util/storage"
	"github.com/APTrust/exchange/util/testutil"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"path/filepath"
	"testing"
)

/*
These tests check the results of the integration tests for
the app apt_record. See the ingest_test.sh script in
the scripts folder, which sets up an integration context, runs
the apt_record.
*/

func TestRecordResults(t *testing.T) {
	if !testutil.ShouldRunIntegrationTests() {
		t.Skip("Skipping integration test. Set ENV var RUN_EXCHANGE_INTEGRATION=true if you want to run them.")
	}
	// Load config
	configFile := filepath.Join("config", "integration.json")
	config, err := models.LoadConfigFile(configFile)
	require.Nil(t, err)
	config.ExpandFilePaths()

	// Find the log file that apt_record created when it was running
	// with the "config/integration.json" config options. We'll read
	// that file.
	pathToJsonLog := filepath.Join(config.LogDirectory, "apt_record.json")
	bagNames := append(testutil.INTEGRATION_GOOD_BAGS, testutil.INTEGRATION_GLACIER_BAGS...)
	for _, bagName := range bagNames {
		ingestManifest, err := testutil.FindIngestManifestInLog(pathToJsonLog, bagName)
		assert.Nil(t, err)
		if err != nil {
			continue
		}
		// TODO: Test WorkItem (stage, status, etc.) below.
		recordTestCommon(t, bagName, ingestManifest)
	}
}

func recordTestCommon(t *testing.T, bagName string, ingestManifest *models.IngestManifest) {
	// Test some basic object properties
	assert.NotEmpty(t, ingestManifest.WorkItemId, "WorkItemId should not be empty for %s", bagName)
	assert.NotEmpty(t, ingestManifest.S3Bucket, "S3Bucket should not be empty for %s", bagName)
	assert.NotEmpty(t, ingestManifest.S3Key, "S3Key should not be empty for %s", bagName)
	assert.NotEmpty(t, ingestManifest.ETag, "ETag should not be empty for %s", bagName)

	// Make sure the result has some basic info in RecordResult
	assert.True(t, ingestManifest.RecordResult.Attempted,
		"RecordResult.Attempted should be true for %s", bagName)
	assert.True(t, ingestManifest.RecordResult.AttemptNumber > 0,
		"RecordResult.AttemptNumber should be > 0 %s", bagName)
	assert.NotEmpty(t, ingestManifest.RecordResult.StartedAt,
		"RecordResult.StartedAt should not be empty for %s", bagName)
	assert.NotEmpty(t, ingestManifest.RecordResult.FinishedAt,
		"RecordResult.FinishedAt should not be empty for %s", bagName)
	assert.Empty(t, ingestManifest.RecordResult.Errors,
		"RecordResult.Errors should be empty for %s", bagName)
	assert.True(t, ingestManifest.RecordResult.Retry,
		"RecordResult.Retry should be true for %s", bagName)

	// Make sure the result has some basic info in CleanupResult
	assert.True(t, ingestManifest.CleanupResult.Attempted,
		"CleanupResult.Attempted should be true for %s", bagName)
	assert.True(t, ingestManifest.CleanupResult.AttemptNumber > 0,
		"CleanupResult.AttemptNumber should be > 0 %s", bagName)
	assert.NotEmpty(t, ingestManifest.CleanupResult.StartedAt,
		"CleanupResult.StartedAt should not be empty for %s", bagName)
	assert.NotEmpty(t, ingestManifest.CleanupResult.FinishedAt,
		"CleanupResult.FinishedAt should not be empty for %s", bagName)
	assert.Empty(t, ingestManifest.CleanupResult.Errors,
		"CleanupResult.Errors should be empty for %s", bagName)
	assert.True(t, ingestManifest.CleanupResult.Retry,
		"CleanupResult.Retry should be true for %s", bagName)

	// Make sure our IntellectualObject got all of its PremisEvents
	//obj := ingestManifest.Object
	db, err := storage.NewBoltDB(ingestManifest.DBPath)
	require.Nil(t, err)
	obj, err := db.GetIntellectualObject(db.ObjectIdentifier())
	require.Nil(t, err)
	require.Equal(t, 4, len(obj.PremisEvents))

	// Make sure this item was deleted from the receiving bucket
	// after ingest completed.
	assert.False(t, obj.IngestDeletedFromReceivingAt.IsZero(),
		"Object %s was not deleted from receiving bucket", bagName)
	assert.Empty(t, obj.IngestErrorMessage)

	// Check the object-level events
	creationEvents := obj.FindEventsByType(constants.EventCreation)
	idEvents := obj.FindEventsByType(constants.EventIdentifierAssignment)
	ingestEvents := obj.FindEventsByType(constants.EventIngestion)
	accessEvents := obj.FindEventsByType(constants.EventAccessAssignment)
	assert.Equal(t, 1, len(accessEvents), "Missing access event for %s", bagName)
	assert.Equal(t, 1, len(creationEvents), "Missing creation event for %s", bagName)
	assert.Equal(t, 1, len(idEvents), "Missing identifier assignment event for %s", bagName)
	assert.Equal(t, 1, len(ingestEvents), "Missing ingest event for %s", bagName)

	for _, event := range obj.PremisEvents {
		assert.True(t, event.Id > 0, "Event %s was not saved for %s", event.EventType, obj.Identifier)
		assert.True(t, event.IntellectualObjectId > 0,
			"event.IntellectualObjectId not set for %s %s", event.EventType, obj.Identifier)
		assert.False(t, event.DateTime.IsZero(),
			"event.DateTime was not set for %s %s", event.EventType, obj.Identifier)
		assert.False(t, event.CreatedAt.IsZero(),
			"event.CreatedAt was not set for %s %s", event.EventType, obj.Identifier)
		assert.False(t, event.UpdatedAt.IsZero(),
			"event.UpdatedAt was not set for %s %s", event.EventType, obj.Identifier)

		assert.True(t, util.LooksLikeUUID(event.Identifier),
			"Identifier for %s %s doesn't look like a UUID", event.EventType, obj.Identifier)
		assert.NotEmpty(t, event.EventType, "EventType missing for %s %s", obj.Identifier, event.Identifier)
		assert.NotEmpty(t, event.Detail, "Detail is empty for %s %s", event.EventType, obj.Identifier)
		assert.NotEmpty(t, event.Outcome, "Outcome is empty for %s %s", event.EventType, obj.Identifier)
		assert.NotEmpty(t, event.OutcomeDetail,
			"OutcomeDetail is empty for %s %s", event.EventType, obj.Identifier)
		assert.NotEmpty(t, event.Object, "Object is empty for %s %s", event.EventType, obj.Identifier)
		assert.NotEmpty(t, event.Agent, "Agent is empty for %s %s", event.EventType, obj.Identifier)
		assert.NotEmpty(t, event.OutcomeInformation,
			"OutcomeInformation is empty for %s %s", event.EventType, obj.Identifier)
		assert.Equal(t, obj.Identifier, event.IntellectualObjectIdentifier,
			"IntellectualObjectIdentifier is wrong for %s %s", event.EventType, obj.Identifier)
	}

	for _, gfIdentifier := range db.FileIdentifiers() {
		gf, err := db.GetGenericFile(gfIdentifier)
		require.Nil(t, err, gfIdentifier)

		// Skip these checks for files that didn't need to be saved.
		// Reasons for not needing to be saved:
		// 1. File has a non-savable name, according to util.HasSavableName
		// 2. File has not changed since last time we ingested this bag.
		if !gf.IngestNeedsSave {
			continue
		}

		// Make sure checksums are present
		require.Equal(t, 2, len(gf.Checksums),
			"Checksums should be %s, found %d for %s", 2, len(gf.Checksums), gf.Identifier)
		md5 := gf.GetChecksumByAlgorithm(constants.AlgMd5)
		sha256 := gf.GetChecksumByAlgorithm(constants.AlgSha256)
		require.NotNil(t, md5, "Missing md5 digest for for %s", gf.Identifier)
		require.NotNil(t, sha256, "Missing sha256 digest for for %s", gf.Identifier)

		// Make sure that these checksums were saved
		assert.True(t, md5.Id > 0, "md5 was not saved for %s", gf.Identifier)
		assert.True(t, md5.GenericFileId > 0, "md5.GenericFileId not set for %s", gf.Identifier)
		assert.False(t, md5.CreatedAt.IsZero(), "md5.CreatedAt was not set for %s", gf.Identifier)
		assert.False(t, md5.UpdatedAt.IsZero(), "md5.UpdatedAt was not set for %s", gf.Identifier)

		assert.True(t, sha256.Id > 0, "sha256 was not saved for %s", gf.Identifier)
		assert.True(t, sha256.GenericFileId > 0, "sha256.GenericFileId not set for %s", gf.Identifier)
		assert.False(t, sha256.CreatedAt.IsZero(), "sha256.CreatedAt was not set for %s", gf.Identifier)
		assert.False(t, sha256.UpdatedAt.IsZero(), "sha256.UpdatedAt was not set for %s", gf.Identifier)

		// Make sure PremisEvents are present
		expectedEventCount := 6
		if gf.StorageOption != constants.StorageStandard {
			expectedEventCount = 5 // no replication event for Glacier-only files
		}
		require.Equal(t, expectedEventCount, len(gf.PremisEvents),
			"PremisEvents count should be %s, found %d for %s", 6, len(gf.PremisEvents), gf.Identifier)
		assert.Equal(t, 1, len(gf.FindEventsByType(constants.EventFixityCheck)),
			"Missing fixity check event for %s", gf.Identifier)
		assert.Equal(t, 1, len(gf.FindEventsByType(constants.EventDigestCalculation)),
			"Missing digest calculation event for %s", gf.Identifier)
		assert.Equal(t, 2, len(gf.FindEventsByType(constants.EventIdentifierAssignment)),
			"Missing identifier assignment event(s) for %s", gf.Identifier)
		if gf.StorageOption == constants.StorageStandard {
			assert.Equal(t, 1, len(gf.FindEventsByType(constants.EventReplication)),
				"Missing replication event for %s", gf.Identifier)
		}
		assert.Equal(t, 1, len(gf.FindEventsByType(constants.EventIngestion)),
			"Missing ingestion event for %s", gf.Identifier)

		for _, event := range gf.PremisEvents {
			assert.True(t, event.Id > 0, "Event %s was not saved for %s", event.EventType, gf.Identifier)
			assert.True(t, event.IntellectualObjectId > 0,
				"event.IntellectualObjectId not set for %s %s", event.EventType, gf.Identifier)
			assert.True(t, event.GenericFileId > 0,
				"event.GenericFileId not set for %s %s", event.EventType, gf.Identifier)
			assert.False(t, event.DateTime.IsZero(),
				"event.DateTime was not set for %s %s", event.EventType, gf.Identifier)
			assert.False(t, event.CreatedAt.IsZero(),
				"event.CreatedAt was not set for %s %s", event.EventType, gf.Identifier)
			assert.False(t, event.UpdatedAt.IsZero(),
				"event.UpdatedAt was not set for %s %s", event.EventType, gf.Identifier)

			assert.True(t, util.LooksLikeUUID(event.Identifier),
				"Identifier for %s %s doesn't look like a UUID", event.EventType, gf.Identifier)
			assert.NotEmpty(t, event.EventType, "EventType missing for %s %s", gf.Identifier, event.Identifier)
			assert.NotEmpty(t, event.Detail, "Detail is empty for %s %s", event.EventType, gf.Identifier)
			assert.NotEmpty(t, event.Outcome, "Outcome is empty for %s %s", event.EventType, gf.Identifier)
			assert.NotEmpty(t, event.OutcomeDetail,
				"OutcomeDetail is empty for %s %s", event.EventType, gf.Identifier)
			assert.NotEmpty(t, event.Object, "Object is empty for %s %s", event.EventType, gf.Identifier)
			assert.NotEmpty(t, event.Agent, "Agent is empty for %s %s", event.EventType, gf.Identifier)
			assert.NotEmpty(t, event.OutcomeInformation,
				"OutcomeInformation is empty for %s %s", event.EventType, gf.Identifier)
			assert.Equal(t, obj.Identifier, event.IntellectualObjectIdentifier,
				"IntellectualObjectIdentifier is wrong for %s %s", event.EventType, gf.Identifier)
			assert.Equal(t, gf.Identifier, event.GenericFileIdentifier,
				"GenericFileIdentifier is wrong for %s %s", event.EventType, gf.Identifier)
		}
	}
}
<|endoftext|>


<commit_before>package cmdtest

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"testing"
)

type Cmd struct {
	Path   string
	Env    []string
	Args   []string
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer
	Silent bool
}

func (cmd Cmd) WithEnv(env ...string) Cmd {
	newEnv := make([]string, len(cmd.Env))
	copy(newEnv, cmd.Env)
	newEnv = append(newEnv, env...)

	cmd.Env = newEnv
	return cmd
}

func (cmd Cmd) WithArgs(args ...string) Cmd {
	newArgs := make([]string, len(cmd.Args))
	copy(newArgs, cmd.Args)
	newArgs = append(newArgs, args...)

	cmd.Args = newArgs
	return cmd
}

func (cmd Cmd) WithTempHome(t *testing.T) (Cmd, string) {
	home := t.TempDir()
	return cmd.WithEnv("HOME=" + home), home
}

func (cmd Cmd) Silence() Cmd {
	cmd.Silent = true
	return cmd
}

func (cmd Cmd) WithInput(in io.Reader) Cmd {
	cmd.Stdin = in
	return cmd
}

func (cmd Cmd) OutputTo(out io.Writer) Cmd {
	cmd.Stdout = out
	return cmd
}

func (cmd Cmd) Run(args ...string) error {
	env := []string{
		// only inherit $PATH; we don't want to pass *everything* along because
		// then it's unclear what's necessary for the tests, but $PATH seems
		// necessary for basic functionality
		"PATH=" + os.Getenv("PATH"),
	}

	execCmd := exec.Command(cmd.Path, append(cmd.Args, args...)...)
	execCmd.Env = append(env, cmd.Env...)
	execCmd.Stdin = cmd.Stdin
	execCmd.Stdout = cmd.Stdout
	execCmd.Stderr = cmd.Stderr

	var verbose io.Writer = ioutil.Discard
	if testing.Verbose() {
		verbose = os.Stderr
	}

	if !cmd.Silent {
		if execCmd.Stdout != nil {
			execCmd.Stdout = io.MultiWriter(execCmd.Stdout, verbose)
		} else {
			execCmd.Stdout = verbose
		}

		if execCmd.Stderr != nil {
			execCmd.Stderr = io.MultiWriter(execCmd.Stderr, verbose)
		} else {
			execCmd.Stderr = verbose
		}
	}

	cmdStr := strings.Join(execCmd.Args, " ")

	fmt.Fprintf(verbose, "\x1b[33m==== EXEC %s\x1b[0m\n", cmdStr)

	err := execCmd.Run()
	if err != nil {
		fmt.Fprintf(verbose, "\x1b[33m==== %s\x1b[0m\n", err)
		return fmt.Errorf("run %s: %w", cmdStr, err)
	}

	return nil
}

func (cmd Cmd) Output(args ...string) (string, error) {
	buf := new(bytes.Buffer)

	cmd.Stdout = buf

	err := cmd.Run(args...)
	if err != nil {
		return "", err
	}

	return buf.String(), nil
}

func (cmd Cmd) OutputJSON(dest interface{}, args ...string) error {
	buf := new(bytes.Buffer)

	cmd.Stdout = buf

	err := cmd.Run(args...)
	if err != nil {
		return err
	}

	err = json.Unmarshal(buf.Bytes(), dest)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>integration: always write to os.Stderr<commit_after>package cmdtest

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
	"testing"
)

type Cmd struct {
	Path   string
	Env    []string
	Args   []string
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer
	Silent bool
}

func (cmd Cmd) WithEnv(env ...string) Cmd {
	newEnv := make([]string, len(cmd.Env))
	copy(newEnv, cmd.Env)
	newEnv = append(newEnv, env...)

	cmd.Env = newEnv
	return cmd
}

func (cmd Cmd) WithArgs(args ...string) Cmd {
	newArgs := make([]string, len(cmd.Args))
	copy(newArgs, cmd.Args)
	newArgs = append(newArgs, args...)

	cmd.Args = newArgs
	return cmd
}

func (cmd Cmd) WithTempHome(t *testing.T) (Cmd, string) {
	home := t.TempDir()
	return cmd.WithEnv("HOME=" + home), home
}

func (cmd Cmd) Silence() Cmd {
	cmd.Silent = true
	return cmd
}

func (cmd Cmd) WithInput(in io.Reader) Cmd {
	cmd.Stdin = in
	return cmd
}

func (cmd Cmd) OutputTo(out io.Writer) Cmd {
	cmd.Stdout = out
	return cmd
}

func (cmd Cmd) Run(args ...string) error {
	env := []string{
		// only inherit $PATH; we don't want to pass *everything* along because
		// then it's unclear what's necessary for the tests, but $PATH seems
		// necessary for basic functionality
		"PATH=" + os.Getenv("PATH"),
	}

	execCmd := exec.Command(cmd.Path, append(cmd.Args, args...)...)
	execCmd.Env = append(env, cmd.Env...)
	execCmd.Stdin = cmd.Stdin
	execCmd.Stdout = cmd.Stdout
	execCmd.Stderr = cmd.Stderr

	// always write to os.Stderr
	//
	// assuming these tests are run like go test ./..., we want output to be
	// shown when the test fails, and 'go test' already does that at the package
	// level.
	//
	// we could try to set this only if testing.Verbose(), but that would mean we
	// have to pass -v, which would result in showing all output even if the
	// tests pass, which is probably too noisy.
	verbose := os.Stderr

	if !cmd.Silent {
		if execCmd.Stdout != nil {
			execCmd.Stdout = io.MultiWriter(execCmd.Stdout, verbose)
		} else {
			execCmd.Stdout = verbose
		}

		if execCmd.Stderr != nil {
			execCmd.Stderr = io.MultiWriter(execCmd.Stderr, verbose)
		} else {
			execCmd.Stderr = verbose
		}
	}

	cmdStr := strings.Join(execCmd.Args, " ")

	fmt.Fprintf(verbose, "\x1b[33m==== EXEC %s\x1b[0m\n", cmdStr)

	err := execCmd.Run()
	if err != nil {
		fmt.Fprintf(verbose, "\x1b[33m==== %s\x1b[0m\n", err)
		return fmt.Errorf("run %s: %w", cmdStr, err)
	}

	return nil
}

func (cmd Cmd) Output(args ...string) (string, error) {
	buf := new(bytes.Buffer)

	cmd.Stdout = buf

	err := cmd.Run(args...)
	if err != nil {
		return "", err
	}

	return buf.String(), nil
}

func (cmd Cmd) OutputJSON(dest interface{}, args ...string) error {
	buf := new(bytes.Buffer)

	cmd.Stdout = buf

	err := cmd.Run(args...)
	if err != nil {
		return err
	}

	err = json.Unmarshal(buf.Bytes(), dest)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>