<commit_before>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cmd

import (
	"fmt"
	"github.com/blang/semver"
	"github.com/spf13/cobra"
	"io"
	"k8s.io/kops/channels/pkg/channels"
	"k8s.io/kops/util/pkg/tables"
	"net/url"
	"os"
	"strings"
)

type ApplyChannelOptions struct {
	Yes   bool
	Files []string
}

func NewCmdApplyChannel(f Factory, out io.Writer) *cobra.Command {
	var options ApplyChannelOptions

	cmd := &cobra.Command{
		Use:   "channel",
		Short: "Apply channel",
		RunE: func(cmd *cobra.Command, args []string) error {
			return RunApplyChannel(f, out, &options, args)
		},
	}

	cmd.Flags().BoolVar(&options.Yes, "yes", false, "Apply update")
	cmd.Flags().StringSliceVar(&options.Files, "f", []string{}, "Apply from a local file")

	return cmd
}

func RunApplyChannel(f Factory, out io.Writer, options *ApplyChannelOptions, args []string) error {
	k8sClient, err := f.KubernetesClient()
	if err != nil {
		return err
	}

	kubernetesVersionInfo, err := k8sClient.Discovery().ServerVersion()
	if err != nil {
		return fmt.Errorf("error querying kubernetes version: %v", err)
	}

	//kubernetesVersion, err := semver.Parse(kubernetesVersionInfo.Major + "." + kubernetesVersionInfo.Minor + ".0")
	//if err != nil {
	//	return fmt.Errorf("cannot parse kubernetes version %q", kubernetesVersionInfo.Major+"."+kubernetesVersionInfo.Minor + ".0")
	//}

	kubernetesVersion, err := semver.ParseTolerant(kubernetesVersionInfo.GitVersion)
	if err != nil {
		return fmt.Errorf("cannot parse kubernetes version %q", kubernetesVersionInfo.GitVersion)
	}

	// Remove Pre and Patch, as they make semver comparisons impractical
	kubernetesVersion.Pre = nil

	menu := channels.NewAddonMenu()

	for _, name := range args {
		location, err := url.Parse(name)
		if err != nil {
			return fmt.Errorf("unable to parse argument %q as url", name)
		}
		if !location.IsAbs() {
			// We recognize the following "well-known" format:
			// <name> with no slashes ->
			if strings.Contains(name, "/") {
				return fmt.Errorf("Channel format not recognized (did you mean to use `-f` to specify a local file?): %q", name)
			}
			expanded := "https://raw.githubusercontent.com/kubernetes/kops/master/addons/" + name + "/addon.yaml"
			location, err = url.Parse(expanded)
			if err != nil {
				return fmt.Errorf("unable to parse expanded argument %q as url", expanded)
			}
		}
		o, err := channels.LoadAddons(name, location)
		if err != nil {
			return fmt.Errorf("error loading channel %q: %v", location, err)
		}

		current, err := o.GetCurrent(kubernetesVersion)
		if err != nil {
			return fmt.Errorf("error processing latest versions in %q: %v", location, err)
		}
		menu.MergeAddons(current)
	}

	for _, f := range options.Files {
		location, err := url.Parse(f)
		if err != nil {
			return fmt.Errorf("unable to parse argument %q as url", f)
		}
		if !location.IsAbs() {
			cwd, err := os.Getwd()
			if err != nil {
				return fmt.Errorf("error getting current directory: %v", err)
			}
			baseURL, err := url.Parse(cwd + string(os.PathSeparator))
			if err != nil {
				return fmt.Errorf("error building url for current directory %q: %v", cwd, err)
			}
			location = baseURL.ResolveReference(location)
		}
		o, err := channels.LoadAddons(f, location)
		if err != nil {
			return fmt.Errorf("error loading file %q: %v", f, err)
		}

		current, err := o.GetCurrent(kubernetesVersion)
		if err != nil {
			return fmt.Errorf("error processing latest versions in %q: %v", f, err)
		}
		menu.MergeAddons(current)
	}

	var updates []*channels.AddonUpdate
	var needUpdates []*channels.Addon
	for _, addon := range menu.Addons {
		// TODO: Cache lookups to prevent repeated lookups?
		update, err := addon.GetRequiredUpdates(k8sClient)
		if err != nil {
			return fmt.Errorf("error checking for required update: %v", err)
		}
		if update != nil {
			updates = append(updates, update)
			needUpdates = append(needUpdates, addon)
		}
	}

	if len(updates) == 0 {
		fmt.Printf("No update required\n")
		return nil
	}

	{
		t := &tables.Table{}
		t.AddColumn("NAME", func(r *channels.AddonUpdate) string {
			return r.Name
		})
		t.AddColumn("CURRENT", func(r *channels.AddonUpdate) string {
			if r.ExistingVersion == nil {
				return "-"
			}
			if r.ExistingVersion.Version != nil {
				return *r.ExistingVersion.Version
			}
			return "?"
		})
		t.AddColumn("UPDATE", func(r *channels.AddonUpdate) string {
			if r.NewVersion == nil {
				return "-"
			}
			if r.NewVersion.Version != nil {
				return *r.NewVersion.Version
			}
			return "?"
		})

		columns := []string{"NAME", "CURRENT", "UPDATE"}
		err := t.Render(updates, os.Stdout, columns...)
		if err != nil {
			return err
		}
	}

	if !options.Yes {
		fmt.Printf("\nMust specify --yes to update\n")
		return nil
	}

	for _, needUpdate := range needUpdates {
		update, err := needUpdate.EnsureUpdated(k8sClient)
		if err != nil {
			return fmt.Errorf("error updating %q: %v", needUpdate.Name, err)
		}
		// Could have been a concurrent request
		if update != nil {
			if update.NewVersion.Version != nil {
				fmt.Printf("Updated %q to %s\n", update.Name, *update.NewVersion.Version)
			} else {
				fmt.Printf("Updated %q\n", update.Name)
			}
		}
	}

	fmt.Printf("\n")

	return nil
}
<commit_msg>channels: accept -f and --files<commit_after>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cmd

import (
	"fmt"
	"github.com/blang/semver"
	"github.com/spf13/cobra"
	"io"
	"k8s.io/kops/channels/pkg/channels"
	"k8s.io/kops/util/pkg/tables"
	"net/url"
	"os"
	"strings"
)

type ApplyChannelOptions struct {
	Yes   bool
	Files []string
}

func NewCmdApplyChannel(f Factory, out io.Writer) *cobra.Command {
	var options ApplyChannelOptions

	cmd := &cobra.Command{
		Use:   "channel",
		Short: "Apply channel",
		RunE: func(cmd *cobra.Command, args []string) error {
			return RunApplyChannel(f, out, &options, args)
		},
	}

	cmd.Flags().BoolVar(&options.Yes, "yes", false, "Apply update")
	cmd.Flags().StringSliceVarP(&options.Files, "filename", "f", []string{}, "Apply from a local file")

	return cmd
}

func RunApplyChannel(f Factory, out io.Writer, options *ApplyChannelOptions, args []string) error {
	k8sClient, err := f.KubernetesClient()
	if err != nil {
		return err
	}

	kubernetesVersionInfo, err := k8sClient.Discovery().ServerVersion()
	if err != nil {
		return fmt.Errorf("error querying kubernetes version: %v", err)
	}

	//kubernetesVersion, err := semver.Parse(kubernetesVersionInfo.Major + "." + kubernetesVersionInfo.Minor + ".0")
	//if err != nil {
	//	return fmt.Errorf("cannot parse kubernetes version %q", kubernetesVersionInfo.Major+"."+kubernetesVersionInfo.Minor + ".0")
	//}

	kubernetesVersion, err := semver.ParseTolerant(kubernetesVersionInfo.GitVersion)
	if err != nil {
		return fmt.Errorf("cannot parse kubernetes version %q", kubernetesVersionInfo.GitVersion)
	}

	// Remove Pre and Patch, as they make semver comparisons impractical
	kubernetesVersion.Pre = nil

	menu := channels.NewAddonMenu()

	for _, name := range args {
		location, err := url.Parse(name)
		if err != nil {
			return fmt.Errorf("unable to parse argument %q as url", name)
		}
		if !location.IsAbs() {
			// We recognize the following "well-known" format:
			// <name> with no slashes ->
			if strings.Contains(name, "/") {
				return fmt.Errorf("Channel format not recognized (did you mean to use `-f` to specify a local file?): %q", name)
			}
			expanded := "https://raw.githubusercontent.com/kubernetes/kops/master/addons/" + name + "/addon.yaml"
			location, err = url.Parse(expanded)
			if err != nil {
				return fmt.Errorf("unable to parse expanded argument %q as url", expanded)
			}
		}
		o, err := channels.LoadAddons(name, location)
		if err != nil {
			return fmt.Errorf("error loading channel %q: %v", location, err)
		}

		current, err := o.GetCurrent(kubernetesVersion)
		if err != nil {
			return fmt.Errorf("error processing latest versions in %q: %v", location, err)
		}
		menu.MergeAddons(current)
	}

	for _, f := range options.Files {
		location, err := url.Parse(f)
		if err != nil {
			return fmt.Errorf("unable to parse argument %q as url", f)
		}
		if !location.IsAbs() {
			cwd, err := os.Getwd()
			if err != nil {
				return fmt.Errorf("error getting current directory: %v", err)
			}
			baseURL, err := url.Parse(cwd + string(os.PathSeparator))
			if err != nil {
				return fmt.Errorf("error building url for current directory %q: %v", cwd, err)
			}
			location = baseURL.ResolveReference(location)
		}
		o, err := channels.LoadAddons(f, location)
		if err != nil {
			return fmt.Errorf("error loading file %q: %v", f, err)
		}

		current, err := o.GetCurrent(kubernetesVersion)
		if err != nil {
			return fmt.Errorf("error processing latest versions in %q: %v", f, err)
		}
		menu.MergeAddons(current)
	}

	var updates []*channels.AddonUpdate
	var needUpdates []*channels.Addon
	for _, addon := range menu.Addons {
		// TODO: Cache lookups to prevent repeated lookups?
		update, err := addon.GetRequiredUpdates(k8sClient)
		if err != nil {
			return fmt.Errorf("error checking for required update: %v", err)
		}
		if update != nil {
			updates = append(updates, update)
			needUpdates = append(needUpdates, addon)
		}
	}

	if len(updates) == 0 {
		fmt.Printf("No update required\n")
		return nil
	}

	{
		t := &tables.Table{}
		t.AddColumn("NAME", func(r *channels.AddonUpdate) string {
			return r.Name
		})
		t.AddColumn("CURRENT", func(r *channels.AddonUpdate) string {
			if r.ExistingVersion == nil {
				return "-"
			}
			if r.ExistingVersion.Version != nil {
				return *r.ExistingVersion.Version
			}
			return "?"
		})
		t.AddColumn("UPDATE", func(r *channels.AddonUpdate) string {
			if r.NewVersion == nil {
				return "-"
			}
			if r.NewVersion.Version != nil {
				return *r.NewVersion.Version
			}
			return "?"
		})

		columns := []string{"NAME", "CURRENT", "UPDATE"}
		err := t.Render(updates, os.Stdout, columns...)
		if err != nil {
			return err
		}
	}

	if !options.Yes {
		fmt.Printf("\nMust specify --yes to update\n")
		return nil
	}

	for _, needUpdate := range needUpdates {
		update, err := needUpdate.EnsureUpdated(k8sClient)
		if err != nil {
			return fmt.Errorf("error updating %q: %v", needUpdate.Name, err)
		}
		// Could have been a concurrent request
		if update != nil {
			if update.NewVersion.Version != nil {
				fmt.Printf("Updated %q to %s\n", update.Name, *update.NewVersion.Version)
			} else {
				fmt.Printf("Updated %q\n", update.Name)
			}
		}
	}

	fmt.Printf("\n")

	return nil
}
<|endoftext|>


<commit_before>package dynago

import (
	"encoding/json"
	"time"
)

const (
	SELECT_ALL        = "ALL_ATTRIBUTES"
	SELECT_PROJECTED  = "ALL_PROJECTED_ATTRIBUTES"
	SELECT_ATTRIBUTES = "SPECIFIC_ATTRIBUTES"
	SELECT_COUNT      = "COUNT"
)

var (
	RETURN_CONSUMED = map[bool]string{true: "TOTAL", false: "NONE"}

	RETURN_TOTAL_CONSUMED = "TOTAL"
	RETURN_INDEX_CONSUMED = "INDEXED"

	RETURN_METRICS = map[bool]string{true: "SIZE", false: "NONE"}

	RETURN_NONE        = "NONE"
	RETURN_ALL_OLD     = "ALL_OLD"
	RETURN_ALL_NEW     = "ALL_NEW"
	RETURN_UPDATED_OLD = "UPDATED_OLD"
	RETURN_UPDATED_NEW = "UPDATED_NEW"
)

type ConsumedCapacityDescription struct {
	CapacityUnits float32
	TableName     string
}

type KeyValue struct {
	Key   AttributeDefinition
	Value interface{}
}

// Items are maps of name/value pairs
type Item map[string]interface{}

func (pi *Item) UnmarshalJSON(data []byte) error {
	var dbitem AttributeNameValue

	if err := json.Unmarshal(data, &dbitem); err != nil {
		return err
	}

	item := make(Item)

	for k, v := range dbitem {
		item[k] = DecodeValue(v)
	}

	*pi = item
	return nil
}

func (pi *Item) MarshalJSON() ([]byte, error) {
	dbitem := AttributeNameValue{}

	for k, v := range *pi {
		dbitem[k] = EncodeValue(v)
	}

	return json.Marshal(dbitem)
}

//////////////////////////////////////////////////////////////////////////////
//
// Put/Update/Delete Item Result
//

type ItemResult struct {
	Attributes            Item
	ConsumedCapacity      ConsumedCapacityDescription
	ItemCollectionMetrics ItemCollectionMetrics
}

type ItemCollectionMetrics struct {
	ItemCollectionKey   AttributeNameValue
	SizeEstimateRangeGB []float64
}

//////////////////////////////////////////////////////////////////////////////
//
// Put/Update/Delete Item Request
//

type ItemRequest struct {
	TableName string

	Item             *Item              `json:",omitempty"` // PutItem
	Key              AttributeNameValue `json:",omitempty"` // UpdateItem/DeleteItem
	UpdateExpression string             `json:",omitempty"` // UpdateItem

	ConditionExpression       string             `json:",omitempty"`
	ExpressionAttributeNames  map[string]string  `json:",omitempty"`
	ExpressionAttributeValues AttributeNameValue `json:",omitempty"`

	ReturnConsumedCapacity      string `json:",omitempty"` // INDEXED | TOTAL | NONE
	ReturnItemCollectionMetrics string `json:",omitempty"` // SIZE | NONE
	ReturnValues                string `json:",omitempty"` // NONE | ALL_OLD | UPDATED_OLD | ALL_NEW | UPDATED_NEW
}

type ItemOption func(*ItemRequest)

func ConditionExpression(expr string) ItemOption {
	return func(req *ItemRequest) {
		req.ConditionExpression = expr
	}
}

func ExpressionAttributeNames(names map[string]string) ItemOption {
	return func(req *ItemRequest) {
		req.ExpressionAttributeNames = names
	}
}

func ExpressionAttributeValues(values map[string]interface{}) ItemOption {
	return func(req *ItemRequest) {
		req.ExpressionAttributeValues = EncodeItem(values)
	}
}

func ReturnConsumed(target string) ItemOption {
	return func(req *ItemRequest) {
		req.ReturnConsumedCapacity = target
	}
}

func ReturnMetrics(ret bool) ItemOption {
	return func(req *ItemRequest) {
		req.ReturnItemCollectionMetrics = RETURN_METRICS[ret]
	}
}

func ReturnValues(target string) ItemOption {
	return func(req *ItemRequest) {
		req.ReturnValues = target
	}
}

//////////////////////////////////////////////////////////////////////////////
//
// PutItem
//

func (db *DBClient) PutItem(tableName string, item Item, options ...ItemOption) (*Item, float32, error) {
	var req = ItemRequest{TableName: tableName, Item: &item}
	var res ItemResult

	for _, option := range options {
		option(&req)
	}

	if err := db.Query("PutItem", &req).Decode(&res); err != nil {
		return nil, 0.0, err
	} else {
		return &res.Attributes, res.ConsumedCapacity.CapacityUnits, err
	}
}

//////////////////////////////////////////////////////////////////////////////
//
// UpdateItem
//

func (db *DBClient) UpdateItem(tableName string, hashKey *KeyValue, rangeKey *KeyValue, updates string, options ...ItemOption) (*Item, float32, error) {
	var req = ItemRequest{TableName: tableName, UpdateExpression: updates}
	var res ItemResult

	req.Key = EncodeAttribute(hashKey.Key, hashKey.Value)
	if rangeKey != nil {
		req.Key[rangeKey.Key.AttributeName] = EncodeAttributeValue(rangeKey.Key, rangeKey.Value)
	}

	if rangeKey != nil {
		req.Key[rangeKey.Key.AttributeName] = EncodeAttributeValue(rangeKey.Key, rangeKey.Value)
	}

	for _, option := range options {
		option(&req)
	}

	if err := db.Query("UpdateItem", &req).Decode(&res); err != nil {
		return nil, 0.0, err
	} else {
		return &res.Attributes, res.ConsumedCapacity.CapacityUnits, err
	}
}

//////////////////////////////////////////////////////////////////////////////
//
// DeleteItem
//

func (db *DBClient) DeleteItem(tableName string, hashKey *KeyValue, rangeKey *KeyValue, options ...ItemOption) (*Item, float32, error) {
	var req = ItemRequest{TableName: tableName}
	var res ItemResult

	req.Key = EncodeAttribute(hashKey.Key, hashKey.Value)
	if rangeKey != nil {
		req.Key[rangeKey.Key.AttributeName] = EncodeAttributeValue(rangeKey.Key, rangeKey.Value)
	}

	for _, option := range options {
		option(&req)
	}

	if err := db.Query("DeleteItem", &req).Decode(&res); err != nil {
		return nil, 0.0, err
	} else {
		return &res.Attributes, res.ConsumedCapacity.CapacityUnits, err
	}
}

//////////////////////////////////////////////////////////////////////////////
//
// GetItem
//

type GetItemRequest struct {
	TableName              string
	Key                    AttributeNameValue
	AttributesToGet        []string
	ConsistentRead         bool
	ReturnConsumedCapacity string
}

type GetItemResult struct {
	ConsumedCapacity ConsumedCapacityDescription

	Item Item
}

func (db *DBClient) GetItem(tableName string, hashKey *KeyValue, rangeKey *KeyValue, attributes []string, consistent bool, consumed bool) (map[string]interface{}, float32, error) {

	req := GetItemRequest{TableName: tableName, AttributesToGet: attributes, ConsistentRead: consistent, ReturnConsumedCapacity: RETURN_CONSUMED[consumed]}
	req.Key = EncodeAttribute(hashKey.Key, hashKey.Value)
	if rangeKey != nil {
		req.Key[rangeKey.Key.AttributeName] = EncodeAttributeValue(rangeKey.Key, rangeKey.Value)
	}

	var res GetItemResult

	if err := db.Query("GetItem", req).Decode(&res); err != nil {
		return nil, 0.0, err
	}

	if len(res.Item) == 0 {
		return nil, res.ConsumedCapacity.CapacityUnits, nil
	}

	return res.Item, res.ConsumedCapacity.CapacityUnits, nil
}

//////////////////////////////////////////////////////////////////////////////
//
// Query
//

type QueryRequest struct {
	TableName        string
	AttributesToGet  []string `json:",omitempty"`
	ScanIndexForward bool

	ExclusiveStartKey AttributeNameValue   `json:",omitempty"`
	KeyConditions     map[string]Condition `json:",omitempty"`
	IndexName         string               `json:",omitempty"`

	FilterExpression          string             `json:",omitempty"`
	ProjectionExpression      string             `json:",omitempty"`
	ExpressionAttributeNames  map[string]string  `json:",omitempty"`
	ExpressionAttributeValues AttributeNameValue `json:",omitempty"`

	Limit                  *int   `json:",omitempty"`
	Select                 string `json:",omitempty"`
	ReturnConsumedCapacity string `json:",omitempty"`

	table *TableInstance
}

type QueryResult struct {
	Items            []Item
	ConsumedCapacity ConsumedCapacityDescription
	LastEvaluatedKey AttributeNameValue
	Count            int
	ScannedCount     int
}

func QueryTable(table *TableInstance) *QueryRequest {
	return &QueryRequest{TableName: table.Name, ScanIndexForward: true, KeyConditions: make(map[string]Condition), table: table}
}

func Query(tableName string) *QueryRequest {
	return &QueryRequest{TableName: tableName, ScanIndexForward: true, KeyConditions: make(map[string]Condition)}
}

func (req *QueryRequest) SetAttributes(attributes []string) *QueryRequest {
	req.AttributesToGet = attributes
	return req
}

func (req *QueryRequest) SetStartKey(startKey AttributeNameValue) *QueryRequest {
	req.ExclusiveStartKey = startKey
	return req
}

func (req *QueryRequest) SetIndex(indexName string) *QueryRequest {
	req.IndexName = indexName
	return req
}

func (req *QueryRequest) SetCondition(attrName string, condition Condition) *QueryRequest {
	req.KeyConditions[attrName] = condition
	return req
}

func (req *QueryRequest) SetAttrCondition(cond AttrCondition) *QueryRequest {
	for k, v := range cond {
		req.KeyConditions[k] = v
	}

	return req
}

func (req *QueryRequest) SetFilterExpression(filter string) *QueryRequest {
	req.FilterExpression = filter
	return req
}

func (req *QueryRequest) SetProjectionExpression(proj string) *QueryRequest {
	req.ProjectionExpression = proj
	return req
}

func (req *QueryRequest) SetLimit(limit int) *QueryRequest {
	req.Limit = &limit
	return req
}

func (req *QueryRequest) SetSelect(selectValue string) *QueryRequest {
	req.Select = selectValue
	return req
}

func (req *QueryRequest) SetConsumed(consumed bool) *QueryRequest {
	req.ReturnConsumedCapacity = RETURN_CONSUMED[consumed]
	return req
}

func (req *QueryRequest) Exec(db *DBClient) ([]Item, AttributeNameValue, float32, error) {
	if db == nil && req.table != nil {
		db = req.table.DB
	}

	var res QueryResult

	if err := db.Query("Query", req).Decode(&res); err != nil {
		return nil, nil, 0.0, err
	}

	return res.Items, res.LastEvaluatedKey, res.ConsumedCapacity.CapacityUnits, nil
}

//////////////////////////////////////////////////////////////////////////////
//
// Scan
//

type ScanRequest struct {
	TableName         string
	AttributesToGet   []string
	ExclusiveStartKey AttributeNameValue

	FilterExpression          string             `json:",omitempty"`
	ProjectionExpression      string             `json:",omitempty"`
	ExpressionAttributeNames  map[string]string  `json:",omitempty"`
	ExpressionAttributeValues AttributeNameValue `json:",omitempty"`

	Limit                  *int   `json:",omitempty"`
	Segment                *int   `json:",omitempty"`
	TotalSegments          *int   `json:",omitempty"`
	Select                 string `json:",omitempty"`
	ReturnConsumedCapacity string `json:",omitempty"`

	table *TableInstance
}

func ScanTable(table *TableInstance) *ScanRequest {
	return &ScanRequest{TableName: table.Name, table: table}
}

func Scan(tableName string) *ScanRequest {
	return &ScanRequest{TableName: tableName}
}

func (req *ScanRequest) SetAttributes(attributes []string) *ScanRequest {
	req.AttributesToGet = attributes
	return req
}

func (req *ScanRequest) SetStartKey(startKey AttributeNameValue) *ScanRequest {
	req.ExclusiveStartKey = startKey
	return req
}

func (req *ScanRequest) SetFilterExpression(filter string) *ScanRequest {
	req.FilterExpression = filter
	return req
}

func (req *ScanRequest) SetProjectionExpression(proj string) *ScanRequest {
	req.ProjectionExpression = proj
	return req
}

func (req *ScanRequest) SetLimit(limit int) *ScanRequest {
	req.Limit = &limit
	return req
}

func (req *ScanRequest) SetSegment(segment, totalSegments int) *ScanRequest {
	req.Segment = &segment
	req.TotalSegments = &totalSegments
	return req
}

func (req *ScanRequest) SetSelect(selectValue string) *ScanRequest {
	req.Select = selectValue
	return req
}

func (req *ScanRequest) SetConsumed(consumed bool) *ScanRequest {
	req.ReturnConsumedCapacity = RETURN_CONSUMED[consumed]
	return req
}

func (req *ScanRequest) Exec(db *DBClient) ([]Item, AttributeNameValue, float32, error) {
	var res QueryResult

	if err := db.Query("Scan", req).Decode(&res); err != nil {
		return nil, nil, 0.0, err
	}

	return res.Items, res.LastEvaluatedKey, res.ConsumedCapacity.CapacityUnits, nil
}

func (req *ScanRequest) Count(db *DBClient) (count int, scount int, consumed float32, err error) {
	return req.CountWithDelay(db, 0)
}

func (req *ScanRequest) CountWithDelay(db *DBClient, delay time.Duration) (count int, scount int, consumed float32, err error) {
	var res QueryResult

	creq := *req
	creq.Select = SELECT_COUNT

	for {
		res.LastEvaluatedKey = nil

		if err = db.Query("Scan", &creq).Decode(&res); err != nil {
			break
		}

		count += res.Count
		scount += res.ScannedCount
		consumed += res.ConsumedCapacity.CapacityUnits

		if res.LastEvaluatedKey == nil {
			break
		}

		creq.ExclusiveStartKey = make(AttributeNameValue)
		for k, v := range res.LastEvaluatedKey {
			creq.ExclusiveStartKey[k] = v
		}

		if delay > 0 {
			time.Sleep(delay)
		}
	}

	return
}
<commit_msg>make some parameters optionals (no need to pass nil)<commit_after>package dynago

import (
	"encoding/json"
	"time"
)

const (
	SELECT_ALL        = "ALL_ATTRIBUTES"
	SELECT_PROJECTED  = "ALL_PROJECTED_ATTRIBUTES"
	SELECT_ATTRIBUTES = "SPECIFIC_ATTRIBUTES"
	SELECT_COUNT      = "COUNT"
)

var (
	RETURN_CONSUMED = map[bool]string{true: "TOTAL", false: "NONE"}

	RETURN_TOTAL_CONSUMED = "TOTAL"
	RETURN_INDEX_CONSUMED = "INDEXED"

	RETURN_METRICS = map[bool]string{true: "SIZE", false: "NONE"}

	RETURN_NONE        = "NONE"
	RETURN_ALL_OLD     = "ALL_OLD"
	RETURN_ALL_NEW     = "ALL_NEW"
	RETURN_UPDATED_OLD = "UPDATED_OLD"
	RETURN_UPDATED_NEW = "UPDATED_NEW"
)

type ConsumedCapacityDescription struct {
	CapacityUnits float32
	TableName     string
}

type KeyValue struct {
	Key   AttributeDefinition
	Value interface{}
}

// Items are maps of name/value pairs
type Item map[string]interface{}

func (pi *Item) UnmarshalJSON(data []byte) error {
	var dbitem AttributeNameValue

	if err := json.Unmarshal(data, &dbitem); err != nil {
		return err
	}

	item := make(Item)

	for k, v := range dbitem {
		item[k] = DecodeValue(v)
	}

	*pi = item
	return nil
}

func (pi *Item) MarshalJSON() ([]byte, error) {
	dbitem := AttributeNameValue{}

	for k, v := range *pi {
		dbitem[k] = EncodeValue(v)
	}

	return json.Marshal(dbitem)
}

//////////////////////////////////////////////////////////////////////////////
//
// Put/Update/Delete Item Result
//

type ItemResult struct {
	Attributes            Item
	ConsumedCapacity      ConsumedCapacityDescription
	ItemCollectionMetrics ItemCollectionMetrics
}

type ItemCollectionMetrics struct {
	ItemCollectionKey   AttributeNameValue
	SizeEstimateRangeGB []float64
}

//////////////////////////////////////////////////////////////////////////////
//
// Put/Update/Delete Item Request
//

type ItemRequest struct {
	TableName string

	Item             *Item              `json:",omitempty"` // PutItem
	Key              AttributeNameValue `json:",omitempty"` // UpdateItem/DeleteItem
	UpdateExpression string             `json:",omitempty"` // UpdateItem

	ConditionExpression       string             `json:",omitempty"`
	ExpressionAttributeNames  map[string]string  `json:",omitempty"`
	ExpressionAttributeValues AttributeNameValue `json:",omitempty"`

	ReturnConsumedCapacity      string `json:",omitempty"` // INDEXED | TOTAL | NONE
	ReturnItemCollectionMetrics string `json:",omitempty"` // SIZE | NONE
	ReturnValues                string `json:",omitempty"` // NONE | ALL_OLD | UPDATED_OLD | ALL_NEW | UPDATED_NEW
}

type ItemOption func(*ItemRequest)

func ConditionExpression(expr string) ItemOption {
	return func(req *ItemRequest) {
		req.ConditionExpression = expr
	}
}

func ExpressionAttributeNames(names map[string]string) ItemOption {
	return func(req *ItemRequest) {
		req.ExpressionAttributeNames = names
	}
}

func ExpressionAttributeValues(values map[string]interface{}) ItemOption {
	return func(req *ItemRequest) {
		req.ExpressionAttributeValues = EncodeItem(values)
	}
}

func ReturnConsumed(target string) ItemOption {
	return func(req *ItemRequest) {
		req.ReturnConsumedCapacity = target
	}
}

func ReturnMetrics(ret bool) ItemOption {
	return func(req *ItemRequest) {
		req.ReturnItemCollectionMetrics = RETURN_METRICS[ret]
	}
}

func ReturnValues(target string) ItemOption {
	return func(req *ItemRequest) {
		req.ReturnValues = target
	}
}

//////////////////////////////////////////////////////////////////////////////
//
// PutItem
//

func (db *DBClient) PutItem(tableName string, item Item, options ...ItemOption) (*Item, float32, error) {
	var req = ItemRequest{TableName: tableName, Item: &item}
	var res ItemResult

	for _, option := range options {
		option(&req)
	}

	if err := db.Query("PutItem", &req).Decode(&res); err != nil {
		return nil, 0.0, err
	} else {
		return &res.Attributes, res.ConsumedCapacity.CapacityUnits, err
	}
}

//////////////////////////////////////////////////////////////////////////////
//
// UpdateItem
//

func (db *DBClient) UpdateItem(tableName string, hashKey *KeyValue, rangeKey *KeyValue, updates string, options ...ItemOption) (*Item, float32, error) {
	var req = ItemRequest{TableName: tableName, UpdateExpression: updates}
	var res ItemResult

	req.Key = EncodeAttribute(hashKey.Key, hashKey.Value)
	if rangeKey != nil {
		req.Key[rangeKey.Key.AttributeName] = EncodeAttributeValue(rangeKey.Key, rangeKey.Value)
	}

	if rangeKey != nil {
		req.Key[rangeKey.Key.AttributeName] = EncodeAttributeValue(rangeKey.Key, rangeKey.Value)
	}

	for _, option := range options {
		option(&req)
	}

	if err := db.Query("UpdateItem", &req).Decode(&res); err != nil {
		return nil, 0.0, err
	} else {
		return &res.Attributes, res.ConsumedCapacity.CapacityUnits, err
	}
}

//////////////////////////////////////////////////////////////////////////////
//
// DeleteItem
//

func (db *DBClient) DeleteItem(tableName string, hashKey *KeyValue, rangeKey *KeyValue, options ...ItemOption) (*Item, float32, error) {
	var req = ItemRequest{TableName: tableName}
	var res ItemResult

	req.Key = EncodeAttribute(hashKey.Key, hashKey.Value)
	if rangeKey != nil {
		req.Key[rangeKey.Key.AttributeName] = EncodeAttributeValue(rangeKey.Key, rangeKey.Value)
	}

	for _, option := range options {
		option(&req)
	}

	if err := db.Query("DeleteItem", &req).Decode(&res); err != nil {
		return nil, 0.0, err
	} else {
		return &res.Attributes, res.ConsumedCapacity.CapacityUnits, err
	}
}

//////////////////////////////////////////////////////////////////////////////
//
// GetItem
//

type GetItemRequest struct {
	TableName              string
	Key                    AttributeNameValue
	AttributesToGet        []string `json:",omitempty"`
	ConsistentRead         bool
	ReturnConsumedCapacity string `json:",omitempty"`
}

type GetItemResult struct {
	ConsumedCapacity ConsumedCapacityDescription

	Item Item
}

func (db *DBClient) GetItem(tableName string, hashKey *KeyValue, rangeKey *KeyValue, attributes []string, consistent bool, consumed bool) (map[string]interface{}, float32, error) {

	req := GetItemRequest{TableName: tableName, AttributesToGet: attributes, ConsistentRead: consistent, ReturnConsumedCapacity: RETURN_CONSUMED[consumed]}
	req.Key = EncodeAttribute(hashKey.Key, hashKey.Value)
	if rangeKey != nil {
		req.Key[rangeKey.Key.AttributeName] = EncodeAttributeValue(rangeKey.Key, rangeKey.Value)
	}

	var res GetItemResult

	if err := db.Query("GetItem", req).Decode(&res); err != nil {
		return nil, 0.0, err
	}

	if len(res.Item) == 0 {
		return nil, res.ConsumedCapacity.CapacityUnits, nil
	}

	return res.Item, res.ConsumedCapacity.CapacityUnits, nil
}

//////////////////////////////////////////////////////////////////////////////
//
// Query
//

type QueryRequest struct {
	TableName        string
	AttributesToGet  []string `json:",omitempty"`
	ScanIndexForward bool

	ExclusiveStartKey AttributeNameValue   `json:",omitempty"`
	KeyConditions     map[string]Condition `json:",omitempty"`
	IndexName         string               `json:",omitempty"`

	FilterExpression          string             `json:",omitempty"`
	ProjectionExpression      string             `json:",omitempty"`
	ExpressionAttributeNames  map[string]string  `json:",omitempty"`
	ExpressionAttributeValues AttributeNameValue `json:",omitempty"`

	Limit                  *int   `json:",omitempty"`
	Select                 string `json:",omitempty"`
	ReturnConsumedCapacity string `json:",omitempty"`

	table *TableInstance
}

type QueryResult struct {
	Items            []Item
	ConsumedCapacity ConsumedCapacityDescription
	LastEvaluatedKey AttributeNameValue
	Count            int
	ScannedCount     int
}

func QueryTable(table *TableInstance) *QueryRequest {
	return &QueryRequest{TableName: table.Name, ScanIndexForward: true, KeyConditions: make(map[string]Condition), table: table}
}

func Query(tableName string) *QueryRequest {
	return &QueryRequest{TableName: tableName, ScanIndexForward: true, KeyConditions: make(map[string]Condition)}
}

func (req *QueryRequest) SetAttributes(attributes []string) *QueryRequest {
	req.AttributesToGet = attributes
	return req
}

func (req *QueryRequest) SetStartKey(startKey AttributeNameValue) *QueryRequest {
	req.ExclusiveStartKey = startKey
	return req
}

func (req *QueryRequest) SetIndex(indexName string) *QueryRequest {
	req.IndexName = indexName
	return req
}

func (req *QueryRequest) SetCondition(attrName string, condition Condition) *QueryRequest {
	req.KeyConditions[attrName] = condition
	return req
}

func (req *QueryRequest) SetAttrCondition(cond AttrCondition) *QueryRequest {
	for k, v := range cond {
		req.KeyConditions[k] = v
	}

	return req
}

func (req *QueryRequest) SetFilterExpression(filter string) *QueryRequest {
	req.FilterExpression = filter
	return req
}

func (req *QueryRequest) SetProjectionExpression(proj string) *QueryRequest {
	req.ProjectionExpression = proj
	return req
}

func (req *QueryRequest) SetLimit(limit int) *QueryRequest {
	req.Limit = &limit
	return req
}

func (req *QueryRequest) SetSelect(selectValue string) *QueryRequest {
	req.Select = selectValue
	return req
}

func (req *QueryRequest) SetConsumed(consumed bool) *QueryRequest {
	req.ReturnConsumedCapacity = RETURN_CONSUMED[consumed]
	return req
}

func (req *QueryRequest) Exec(db *DBClient) ([]Item, AttributeNameValue, float32, error) {
	if db == nil && req.table != nil {
		db = req.table.DB
	}

	var res QueryResult

	if err := db.Query("Query", req).Decode(&res); err != nil {
		return nil, nil, 0.0, err
	}

	return res.Items, res.LastEvaluatedKey, res.ConsumedCapacity.CapacityUnits, nil
}

//////////////////////////////////////////////////////////////////////////////
//
// Scan
//

type ScanRequest struct {
	TableName         string
	AttributesToGet   []string
	ExclusiveStartKey AttributeNameValue

	FilterExpression          string             `json:",omitempty"`
	ProjectionExpression      string             `json:",omitempty"`
	ExpressionAttributeNames  map[string]string  `json:",omitempty"`
	ExpressionAttributeValues AttributeNameValue `json:",omitempty"`

	Limit                  *int   `json:",omitempty"`
	Segment                *int   `json:",omitempty"`
	TotalSegments          *int   `json:",omitempty"`
	Select                 string `json:",omitempty"`
	ReturnConsumedCapacity string `json:",omitempty"`

	table *TableInstance
}

func ScanTable(table *TableInstance) *ScanRequest {
	return &ScanRequest{TableName: table.Name, table: table}
}

func Scan(tableName string) *ScanRequest {
	return &ScanRequest{TableName: tableName}
}

func (req *ScanRequest) SetAttributes(attributes []string) *ScanRequest {
	req.AttributesToGet = attributes
	return req
}

func (req *ScanRequest) SetStartKey(startKey AttributeNameValue) *ScanRequest {
	req.ExclusiveStartKey = startKey
	return req
}

func (req *ScanRequest) SetFilterExpression(filter string) *ScanRequest {
	req.FilterExpression = filter
	return req
}

func (req *ScanRequest) SetProjectionExpression(proj string) *ScanRequest {
	req.ProjectionExpression = proj
	return req
}

func (req *ScanRequest) SetLimit(limit int) *ScanRequest {
	req.Limit = &limit
	return req
}

func (req *ScanRequest) SetSegment(segment, totalSegments int) *ScanRequest {
	req.Segment = &segment
	req.TotalSegments = &totalSegments
	return req
}

func (req *ScanRequest) SetSelect(selectValue string) *ScanRequest {
	req.Select = selectValue
	return req
}

func (req *ScanRequest) SetConsumed(consumed bool) *ScanRequest {
	req.ReturnConsumedCapacity = RETURN_CONSUMED[consumed]
	return req
}

func (req *ScanRequest) Exec(db *DBClient) ([]Item, AttributeNameValue, float32, error) {
	var res QueryResult

	if err := db.Query("Scan", req).Decode(&res); err != nil {
		return nil, nil, 0.0, err
	}

	return res.Items, res.LastEvaluatedKey, res.ConsumedCapacity.CapacityUnits, nil
}

func (req *ScanRequest) Count(db *DBClient) (count int, scount int, consumed float32, err error) {
	return req.CountWithDelay(db, 0)
}

func (req *ScanRequest) CountWithDelay(db *DBClient, delay time.Duration) (count int, scount int, consumed float32, err error) {
	var res QueryResult

	creq := *req
	creq.Select = SELECT_COUNT

	for {
		res.LastEvaluatedKey = nil

		if err = db.Query("Scan", &creq).Decode(&res); err != nil {
			break
		}

		count += res.Count
		scount += res.ScannedCount
		consumed += res.ConsumedCapacity.CapacityUnits

		if res.LastEvaluatedKey == nil {
			break
		}

		creq.ExclusiveStartKey = make(AttributeNameValue)
		for k, v := range res.LastEvaluatedKey {
			creq.ExclusiveStartKey[k] = v
		}

		if delay > 0 {
			time.Sleep(delay)
		}
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/prometheus/client_golang/api/prometheus"
	"github.com/prometheus/common/expfmt"
	"github.com/prometheus/common/model"
)

type Tag struct {
	Name  model.LabelName
	Value model.LabelValue
}

type Metric struct {
	Tags  []Tag
	Value float64
}

func CreateJSONMetrics(samples model.Vector) string {
	metrics := []Metric{}

	for _, sample := range samples {
		metric := Metric{}

		for name, value := range sample.Metric {
			tag := Tag{
				Name:  name,
				Value: value,
			}

			metric.Tags = append(metric.Tags, tag)
		}

		metric.Value = float64(sample.Value)

		metrics = append(metrics, metric)
	}

	jsonMetrics, _ := json.Marshal(metrics)

	return string(jsonMetrics)
}

func CreateGraphiteMetrics(samples model.Vector) string {
	metrics := ""

	for _, sample := range samples {
		name := sample.Metric["__name__"]

		value := sample.Value

		now := time.Now()
		timestamp := now.Unix()

		metric := fmt.Sprintf("%s %f %v\n", name, value, timestamp)

		metrics += metric
	}

	return metrics
}

func CreateInfluxMetrics(samples model.Vector) string {
	metrics := ""

	for _, sample := range samples {
		metric := string(sample.Metric["__name__"])

		for name, value := range sample.Metric {
			if name != "__name__" {
				metric += fmt.Sprintf(",%s=%s", name, value)
			}
		}

		value := sample.Value

		now := time.Now()
		timestamp := now.Unix()

		metric += fmt.Sprintf(" value=%f %v\n", value, timestamp)

		metrics += metric
	}

	return metrics
}

func OutputMetrics(samples model.Vector, outputFormat string) error {
	output := ""

	switch outputFormat {
	case "influx":
		output = CreateInfluxMetrics(samples)
	case "graphite":
		output = CreateGraphiteMetrics(samples)
	case "json":
		output = CreateJSONMetrics(samples)
	}

	fmt.Println(output)

	return nil
}

func QueryPrometheus(promURL string, queryString string) (model.Vector, error) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	promConfig := prometheus.Config{Address: promURL}
	promClient, err := prometheus.New(promConfig)

	if err != nil {
		return nil, err
	}

	promQueryClient := prometheus.NewQueryAPI(promClient)

	promResponse, err := promQueryClient.Query(ctx, queryString, time.Now())

	if err != nil {
		return nil, err
	}

	if promResponse.Type() == model.ValVector {
		return promResponse.(model.Vector), nil
	}

	return nil, errors.New("unexpected response type")
}

func QueryExporter(exporterURL string) (model.Vector, error) {
	expResponse, err := http.Get(exporterURL)
	defer expResponse.Body.Close()

	if err != nil {
		return nil, err
	}

	if expResponse.StatusCode != http.StatusOK {
		return nil, errors.New("exporter returned non OK HTTP response status")
	}

	var parser expfmt.TextParser

	metricFamilies, err := parser.TextToMetricFamilies(expResponse.Body)

	if err != nil {
		return nil, err
	}

	samples := model.Vector{}

	decodeOptions := &expfmt.DecodeOptions{
		Timestamp: model.Time(time.Now().Unix()),
	}

	for _, family := range metricFamilies {
		familySamples, _ := expfmt.ExtractSamples(decodeOptions, family)
		samples = append(samples, familySamples...)
	}

	return samples, nil
}

func main() {
	exporterURL := flag.String("exporter-url", "", "Prometheus exporter URL to pull metrics from")
	promURL := flag.String("prom-url", "http://localhost:9090", "Prometheus API URL")
	queryString := flag.String("prom-query", "up", "Prometheus API query string")
	outputFormat := flag.String("output-format", "influx", "The check output format to use for metrics {influx|graphite|json}")
	flag.Parse()

	samples := model.Vector{}
	var err error

	if *exporterURL != "" {
		samples, err = QueryExporter(*exporterURL)
	} else {
		samples, err = QueryPrometheus(*promURL, *queryString)
	}

	if err != nil {
		fmt.Errorf("%v", err)
		os.Exit(2)
	}

	err = OutputMetrics(samples, *outputFormat)

	if err != nil {
		fmt.Errorf("%v", err)
		os.Exit(2)
	}
}
<commit_msg>added metric prefix cli argument, an optional a metric name prefix for line protocols<commit_after>package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/prometheus/client_golang/api/prometheus"
	"github.com/prometheus/common/expfmt"
	"github.com/prometheus/common/model"
)

type Tag struct {
	Name  model.LabelName
	Value model.LabelValue
}

type Metric struct {
	Tags  []Tag
	Value float64
}

func CreateJSONMetrics(samples model.Vector) string {
	metrics := []Metric{}

	for _, sample := range samples {
		metric := Metric{}

		for name, value := range sample.Metric {
			tag := Tag{
				Name:  name,
				Value: value,
			}

			metric.Tags = append(metric.Tags, tag)
		}

		metric.Value = float64(sample.Value)

		metrics = append(metrics, metric)
	}

	jsonMetrics, _ := json.Marshal(metrics)

	return string(jsonMetrics)
}

func CreateGraphiteMetrics(samples model.Vector, metricPrefix string) string {
	metrics := ""

	for _, sample := range samples {
		name := fmt.Sprintf("%s%s", metricPrefix, sample.Metric["__name__"])

		value := sample.Value

		now := time.Now()
		timestamp := now.Unix()

		metric := fmt.Sprintf("%s %f %v\n", name, value, timestamp)

		metrics += metric
	}

	return metrics
}

func CreateInfluxMetrics(samples model.Vector, metricPrefix string) string {
	metrics := ""

	for _, sample := range samples {
		metric := fmt.Sprintf("%s%s", metricPrefix, sample.Metric["__name__"])

		for name, value := range sample.Metric {
			if name != "__name__" {
				metric += fmt.Sprintf(",%s=%s", name, value)
			}
		}

		value := sample.Value

		now := time.Now()
		timestamp := now.Unix()

		metric += fmt.Sprintf(" value=%f %v\n", value, timestamp)

		metrics += metric
	}

	return metrics
}

func OutputMetrics(samples model.Vector, outputFormat string, metricPrefix string) error {
	output := ""

	switch outputFormat {
	case "influx":
		output = CreateInfluxMetrics(samples, metricPrefix)
	case "graphite":
		output = CreateGraphiteMetrics(samples, metricPrefix)
	case "json":
		output = CreateJSONMetrics(samples)
	}

	fmt.Println(output)

	return nil
}

func QueryPrometheus(promURL string, queryString string) (model.Vector, error) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	promConfig := prometheus.Config{Address: promURL}
	promClient, err := prometheus.New(promConfig)

	if err != nil {
		return nil, err
	}

	promQueryClient := prometheus.NewQueryAPI(promClient)

	promResponse, err := promQueryClient.Query(ctx, queryString, time.Now())

	if err != nil {
		return nil, err
	}

	if promResponse.Type() == model.ValVector {
		return promResponse.(model.Vector), nil
	}

	return nil, errors.New("unexpected response type")
}

func QueryExporter(exporterURL string) (model.Vector, error) {
	expResponse, err := http.Get(exporterURL)
	defer expResponse.Body.Close()

	if err != nil {
		return nil, err
	}

	if expResponse.StatusCode != http.StatusOK {
		return nil, errors.New("exporter returned non OK HTTP response status")
	}

	var parser expfmt.TextParser

	metricFamilies, err := parser.TextToMetricFamilies(expResponse.Body)

	if err != nil {
		return nil, err
	}

	samples := model.Vector{}

	decodeOptions := &expfmt.DecodeOptions{
		Timestamp: model.Time(time.Now().Unix()),
	}

	for _, family := range metricFamilies {
		familySamples, _ := expfmt.ExtractSamples(decodeOptions, family)
		samples = append(samples, familySamples...)
	}

	return samples, nil
}

func main() {
	exporterURL := flag.String("exporter-url", "", "Prometheus exporter URL to pull metrics from.")
	promURL := flag.String("prom-url", "http://localhost:9090", "Prometheus API URL.")
	queryString := flag.String("prom-query", "up", "Prometheus API query string.")
	outputFormat := flag.String("output-format", "influx", "The check output format to use for metrics {influx|graphite|json}.")
	metricPrefix := flag.String("metric-prefix", "", "Metric name prefix, only supported by line protocol output formats.")
	flag.Parse()

	samples := model.Vector{}
	var err error

	if *exporterURL != "" {
		samples, err = QueryExporter(*exporterURL)
	} else {
		samples, err = QueryPrometheus(*promURL, *queryString)
	}

	if err != nil {
		fmt.Errorf("%v", err)
		os.Exit(2)
	}

	err = OutputMetrics(samples, *outputFormat, *metricPrefix)

	if err != nil {
		fmt.Errorf("%v", err)
		os.Exit(2)
	}
}
<|endoftext|>


<commit_before>package crypto

import (
	"crypto/rand"
	"math/big"
)

// RandBytes returns n bytes of random data.
func RandBytes(n int) ([]byte, error) {
	b := make([]byte, n)
	_, err := rand.Read(b)
	return b, err
}

// RandIntn returns a non-negative random integer in the range [0,n). It panics
// if n <= 0.
func RandIntn(n int) (int, error) {
	r, err := rand.Int(rand.Reader, big.NewInt(int64(n)))
	return int(r.Int64()), err
}

// Perm returns, as a slice of n ints, a random permutation of the integers
// [0,n).
func Perm(n int) ([]int, error) {
	m := make([]int, n)
	for i := 0; i < n; i++ {
		j, err := RandIntn(i + 1)
		if err != nil {
			return nil, err
		}

		m[i] = m[j]
		m[j] = i
	}
	return m, nil
}
<commit_msg>change crypto/rand.go to not be able to return errors<commit_after>package crypto

import (
	"crypto/rand"
	"math/big"
	"sync/atomic"
)

var (
	// Two counters is used in case the first overflows. The first is not at
	// risk of overflowing, but we're paranoid.
	counter     uint64
	counter2    uint64
	entropyBase Hash
)

// init will create an entropy pool that the RNG will draw from. On most
// systems, hashing out of a preset RNG pool will be substantially faster than
// using crypto/rand, which relies on syscalls.
func init() {
	// Use 64 bytes in case the first 32 aren't completely random.
	base := make([]byte, 64)
	_, err := rand.Read(base)
	if err != nil {
		panic("unable to set entropy for the crypto package")
	}
	entropyBase = HashObject(base)
}

// RandBytes returns n bytes of random data.
func RandBytes(n int) ([]byte, error) {
	b := make([]byte, n)
	for i := 0; i < n; i += HashSize {
		// Fetch and update the counter values before executing the hash.
		var newCounter, newCounter2 uint64
		newCounter = atomic.AddUint64(&counter, 1)
		if newCounter == 0 {
			newCounter2 = atomic.AddUint64(&counter2, 1)
		} else {
			newCounter2 = atomic.LoadUint64(&counter2)
		}
		// Grab some entropy using the unique counter set.
		entropy := HashAll(newCounter, newCounter2, entropyBase)

		// Fill out 'b'.
		copy(b[i:], entropy[:])
	}
	return b, nil
}

// RandIntn returns a non-negative random integer in the range [0,n). It panics
// if n <= 0.
func RandIntn(n int) (int, error) {
	if n <= 0 {
		panic("RandIntn must be called with a positive, nonzero number")
	}

	// Fetch and update the counter values before executing the hash.
	var newCounter, newCounter2 uint64
	newCounter = atomic.AddUint64(&counter, 1)
	if newCounter == 0 {
		newCounter2 = atomic.AddUint64(&counter2, 1)
	} else {
		newCounter2 = atomic.LoadUint64(&counter2)
	}
	// Grab some entropy using the unique counter set.
	entropy := HashAll(newCounter, newCounter2, entropyBase)

	// Convert the first 24 bytes into a big.Int, then grab the modulus of n.
	// 24 bytes means that there is at least 16 bytes of overflow, which means
	// early numbers are favored by less than 1-in-2^128 - a cryptographically
	// safe preference.
	b := new(big.Int)
	b.SetBytes(entropy[:24])

	// Take the modules of 'b'  and 'n' and return the result as an int.
	return int(b.Mod(b, big.NewInt(int64(n))).Int64()), nil
}

// Read will fill 'b' with completely random data.
func Read(b []byte) {
	n := len(b)
	for i := 0; i < n; i += HashSize {
		// Fetch and update the counter values before executing the hash.
		var newCounter, newCounter2 uint64
		newCounter = atomic.AddUint64(&counter, 1)
		if newCounter == 0 {
			newCounter2 = atomic.AddUint64(&counter2, 1)
		} else {
			newCounter2 = atomic.LoadUint64(&counter2)
		}
		// Grab some entropy using the unique counter set.
		entropy := HashAll(newCounter, newCounter2, entropyBase)

		// Fill out 'b'.
		copy(b[i:], entropy[:])
	}
}

// Perm returns, as a slice of n ints, a random permutation of the integers
// [0,n).
func Perm(n int) ([]int, error) {
	m := make([]int, n)
	for i := 0; i < n; i++ {
		j, _ := RandIntn(i + 1)
		m[i] = m[j]
		m[j] = i
	}
	return m, nil
}
<|endoftext|>


<commit_before>package packetdeserializers

import (
	"errors"
	"fmt"

	"github.com/piot/hasty-protocol/asciistring"
	"github.com/piot/hasty-protocol/commands"
	"github.com/piot/hasty-protocol/deserialize"
	"github.com/piot/hasty-protocol/packet"
	"github.com/piot/hasty-protocol/realmname"
	log "github.com/sirupsen/logrus"
)

// ToConnect : Channel to subscribe to
func ToConnect(in packet.Packet) (commands.Connect, error) {
	if in.Type() != packet.Connect {
		return commands.Connect{}, errors.New("Illegal packet type")
	}

	payload := in.Payload()
	pos := 0
	versionObject, versionOctetCount, versionErr := deserialize.ToVersion(payload[0:3])
	if versionErr != nil {
		return commands.Connect{}, errors.New("Illegal version")
	}
	pos += versionOctetCount
	realmString, _, realmStringErr := asciistring.FromOctets(payload[pos:])
	if realmStringErr != nil {
		return commands.Connect{}, fmt.Errorf("Illegal realm %s", realmStringErr)
	}
	realm, realmErr := realmname.NewName(realmString)
	if realmErr != nil {
		return commands.Connect{}, realmErr
	}
	connect := commands.NewConnect(realm, versionObject)
	log.Infof("connect %s\n", connect)
	return connect, nil
}
<commit_msg>minor logging<commit_after>package packetdeserializers

import (
	"errors"
	"fmt"

	"github.com/piot/hasty-protocol/asciistring"
	"github.com/piot/hasty-protocol/commands"
	"github.com/piot/hasty-protocol/deserialize"
	"github.com/piot/hasty-protocol/packet"
	"github.com/piot/hasty-protocol/realmname"
	log "github.com/sirupsen/logrus"
)

// ToConnect : Channel to subscribe to
func ToConnect(in packet.Packet) (commands.Connect, error) {
	if in.Type() != packet.Connect {
		return commands.Connect{}, errors.New("Illegal packet type")
	}

	payload := in.Payload()
	pos := 0
	versionObject, versionOctetCount, versionErr := deserialize.ToVersion(payload[0:3])
	if versionErr != nil {
		return commands.Connect{}, errors.New("Illegal version")
	}
	pos += versionOctetCount
	realmString, _, realmStringErr := asciistring.FromOctets(payload[pos:])
	if realmStringErr != nil {
		return commands.Connect{}, fmt.Errorf("Illegal realm %s", realmStringErr)
	}
	realm, realmErr := realmname.NewName(realmString)
	if realmErr != nil {
		return commands.Connect{}, realmErr
	}
	connect := commands.NewConnect(realm, versionObject)
	log.Infof("%s", connect)
	return connect, nil
}
<|endoftext|>


<commit_before>package input

import (
	"fmt"
	"io"
	"net/http"
	"time"

	"github.com/venicegeo/pzsvc-exec/worker/config"
)

var httpClient = http.Client{
	Timeout: 30 * time.Second,
}

type asyncDownloader interface {
	DownloadInputAsync(source config.InputSource) chan error
}

type defaultAsyncDownloader struct{}

var asyncDownloaderInstance asyncDownloader = defaultAsyncDownloader{}

func (dl defaultAsyncDownloader) DownloadInputAsync(source config.InputSource) chan error {
	errChan := make(chan error)

	go func() {
		var err error
		defer close(errChan)

		targetFile, err := fileCheckerInstance.CheckAndOpen(source.FileName, 0777)
		if err != nil {
			errChan <- err
			return
		}
		defer targetFile.Close()

		resp, err := httpClient.Get(source.URL)
		if err == nil && resp.StatusCode != http.StatusOK {
			err = fmt.Errorf("Unexpected HTTP status: %v", resp.StatusCode)
		}
		if err != nil {
			errChan <- err
			return
		}
		defer resp.Body.Close()

		_, err = io.Copy(targetFile, resp.Body)

		if err != nil {
			errChan <- err
			return
		}
	}()

	return errChan
}
<commit_msg>Configurable Client Timeout<commit_after>package input

import (
	"fmt"
	"io"
	"net/http"
	"time"

	"github.com/venicegeo/pzsvc-exec/worker/config"
)

func getClientTimeout() time.Duration {
	defaultTimeout := 30
	if envTimeout := os.Getenv("HTTP_TIMEOUT"); envTimeout != "" {
		defaultTimeout , _ = strconv.Atoi(envTimeout)
	}
	return (time.Duration(defaultTimeout) * time.Second)
}

var httpClient = http.Client{
	Timeout: getClientTimeout(),
}

type asyncDownloader interface {
	DownloadInputAsync(source config.InputSource) chan error
}

type defaultAsyncDownloader struct{}

var asyncDownloaderInstance asyncDownloader = defaultAsyncDownloader{}

func (dl defaultAsyncDownloader) DownloadInputAsync(source config.InputSource) chan error {
	errChan := make(chan error)

	go func() {
		var err error
		defer close(errChan)

		targetFile, err := fileCheckerInstance.CheckAndOpen(source.FileName, 0777)
		if err != nil {
			errChan <- err
			return
		}
		defer targetFile.Close()

		resp, err := httpClient.Get(source.URL)
		if err == nil && resp.StatusCode != http.StatusOK {
			err = fmt.Errorf("Unexpected HTTP status: %v", resp.StatusCode)
		}
		if err != nil {
			errChan <- err
			return
		}
		defer resp.Body.Close()

		_, err = io.Copy(targetFile, resp.Body)

		if err != nil {
			errChan <- err
			return
		}
	}()

	return errChan
}
<|endoftext|>


<commit_before>/*
   Copyright 2010-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.

   This file is licensed under the Apache License, Version 2.0 (the "License").
   You may not use this file except in compliance with the License. A copy of
   the License is located at

    http://aws.amazon.com/apache2.0/

   This file 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 (
    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/s3"

    "fmt"
)

// Tag S3 bucket MyBucket with cost center tag "123456" and stack tag "MyTestStack".
//
// See:
//    http://docs.aws.amazon.com/awsaccountbilling/latest/aboutv2/cost-alloc-tags.html
func main() {
    // Pre-defined values
    bucket := "MyBucket"
    tagName1 := "Cost Center"
    tagValue1 := "123456"
    tagName2 := "Stack"
    tagValue2 := "MyTestStack"

    // Initialize a session in us-west-2 that the SDK will use to load credentials
    // from the shared credentials file. (~/.aws/ccredentials).
    sess, err := session.NewSession(&aws.Config{
        Region: aws.String("us-west-2")},
    )
    if err != nil {
        fmt.Println(err.Error())
        return
    }

    // Create S3 service client
    svc := s3.New(sess)

    // Create input for PutBucket method
    input := &s3.PutBucketTaggingInput{
        Bucket: aws.String(bucket),
        Tagging: &s3.Tagging{
            TagSet: []*s3.Tag{
                {
                    Key:   aws.String(tagName1),
                    Value: aws.String(tagValue),
                },
                {
                    Key:   aws.String(tagName2),
                    Value: aws.String(tagValue2),
              },
            },
        },
    }

    _, err = svc.PutBucketTagging(input)
    if err != nil {
        fmt.Println(err.Error())
        return
    }

    // Now show the tags
    // Create input for GetBucket method
    input := &s3.GetBucketTaggingInput{
        Bucket: aws.String(bucket),
    }

    result, err := svc.GetBucketTagging(input)
    if err != nil {
        fmt.Println(err.Error())
        return
    }

    numTags := len(result.TagSet)

    if numTags > 0 {
        fmt.Println("Found", numTags, "Tag(s):")
        fmt.Println("")

        for _, t := range result.TagSet {
            fmt.Println("  Key:  ", *t.Key)
            fmt.Println("  Value:", *t.Value)
            fmt.Println("")
        }
    } else {
        fmt.Println("Did not find any tags")
    }
}
<commit_msg>Fixed typos in 'Using the AWS SDK for Go with AWS Services' topic<commit_after>/*
   Copyright 2010-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.

   This file is licensed under the Apache License, Version 2.0 (the "License").
   You may not use this file except in compliance with the License. A copy of
   the License is located at

    http://aws.amazon.com/apache2.0/

   This file 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 (
    "github.com/aws/aws-sdk-go/aws"
    "github.com/aws/aws-sdk-go/aws/session"
    "github.com/aws/aws-sdk-go/service/s3"

    "fmt"
)

// Tag S3 bucket MyBucket with cost center tag "123456" and stack tag "MyTestStack".
//
// See:
//    http://docs.aws.amazon.com/awsaccountbilling/latest/aboutv2/cost-alloc-tags.html
func main() {
    // Pre-defined values
    bucket := "MyBucket"
    tagName1 := "Cost Center"
    tagValue1 := "123456"
    tagName2 := "Stack"
    tagValue2 := "MyTestStack"

    // Initialize a session in us-west-2 that the SDK will use to load credentials
    // from the shared credentials file. (~/.aws/credentials).
    sess, err := session.NewSession(&aws.Config{
        Region: aws.String("us-west-2")},
    )
    if err != nil {
        fmt.Println(err.Error())
        return
    }

    // Create S3 service client
    svc := s3.New(sess)

    // Create input for PutBucket method
    input := &s3.PutBucketTaggingInput{
        Bucket: aws.String(bucket),
        Tagging: &s3.Tagging{
            TagSet: []*s3.Tag{
                {
                    Key:   aws.String(tagName1),
                    Value: aws.String(tagValue1),
                },
                {
                    Key:   aws.String(tagName2),
                    Value: aws.String(tagValue2),
              },
            },
        },
    }

    _, err = svc.PutBucketTagging(input)
    if err != nil {
        fmt.Println(err.Error())
        return
    }

    // Now show the tags
    // Create input for GetBucket method
    input := &s3.GetBucketTaggingInput{
        Bucket: aws.String(bucket),
    }

    result, err := svc.GetBucketTagging(input)
    if err != nil {
        fmt.Println(err.Error())
        return
    }

    numTags := len(result.TagSet)

    if numTags > 0 {
        fmt.Println("Found", numTags, "Tag(s):")
        fmt.Println("")

        for _, t := range result.TagSet {
            fmt.Println("  Key:  ", *t.Key)
            fmt.Println("  Value:", *t.Value)
            fmt.Println("")
        }
    } else {
        fmt.Println("Did not find any tags")
    }
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package cephbackupstorage implements the BackupStorage interface
// for Ceph Cloud Storage.
package cephbackupstorage

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
	"sync"

	"errors"

	minio "github.com/minio/minio-go"
	"golang.org/x/net/context"

	"vitess.io/vitess/go/vt/concurrency"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/mysqlctl/backupstorage"
)

var (
	// configFilePath is where the configs/credentials for backups will be stored.
	configFilePath = flag.String("ceph_backup_storage_config", "ceph_backup_config.json",
		"Path to JSON config file for ceph backup storage")
)

var storageConfig struct {
	AccessKey string `json:"accessKey"`
	SecretKey string `json:"secretKey"`
	EndPoint  string `json:"endPoint"`
	UseSSL    bool   `json:"useSSL"`
}

// CephBackupHandle implements BackupHandle for Ceph Cloud Storage.
type CephBackupHandle struct {
	client    *minio.Client
	bs        *CephBackupStorage
	dir       string
	name      string
	readOnly  bool
	errors    concurrency.AllErrorRecorder
	waitGroup sync.WaitGroup
}

// Directory implements BackupHandle.
func (bh *CephBackupHandle) Directory() string {
	return bh.dir
}

// Name implements BackupHandle.
func (bh *CephBackupHandle) Name() string {
	return bh.name
}

// AddFile implements BackupHandle.
func (bh *CephBackupHandle) AddFile(ctx context.Context, filename string, filesize int64) (io.WriteCloser, error) {
	if bh.readOnly {
		return nil, fmt.Errorf("AddFile cannot be called on read-only backup")
	}
	reader, writer := io.Pipe()
	bh.waitGroup.Add(1)
	go func() {
		defer bh.waitGroup.Done()

		// ceph bucket name is where the backups will go
		//backup handle dir field contains keyspace/shard value
		bucket := alterBucketName(bh.dir)

		// Give PutObject() the read end of the pipe.
		object := objName(bh.dir, bh.name, filename)
		_, err := bh.client.PutObjectWithContext(ctx, bucket, object, reader, -1, minio.PutObjectOptions{ContentType: "application/octet-stream"})
		if err != nil {
			// Signal the writer that an error occurred, in case it's not done writing yet.
			reader.CloseWithError(err)
			// In case the error happened after the writer finished, we need to remember it.
			bh.errors.RecordError(err)
		}
	}()
	// Give our caller the write end of the pipe.
	return writer, nil
}

// EndBackup implements BackupHandle.
func (bh *CephBackupHandle) EndBackup(ctx context.Context) error {
	if bh.readOnly {
		return fmt.Errorf("EndBackup cannot be called on read-only backup")
	}
	bh.waitGroup.Wait()
	// Return the saved PutObject() errors, if any.
	return bh.errors.Error()
}

// AbortBackup implements BackupHandle.
func (bh *CephBackupHandle) AbortBackup(ctx context.Context) error {
	if bh.readOnly {
		return fmt.Errorf("AbortBackup cannot be called on read-only backup")
	}
	return bh.bs.RemoveBackup(ctx, bh.dir, bh.name)
}

// ReadFile implements BackupHandle.
func (bh *CephBackupHandle) ReadFile(ctx context.Context, filename string) (io.ReadCloser, error) {
	if !bh.readOnly {
		return nil, fmt.Errorf("ReadFile cannot be called on read-write backup")
	}
	// ceph bucket name
	bucket := alterBucketName(bh.dir)
	object := objName(bh.dir, bh.name, filename)
	return bh.client.GetObjectWithContext(ctx, bucket, object, minio.GetObjectOptions{})
}

// CephBackupStorage implements BackupStorage for Ceph Cloud Storage.
type CephBackupStorage struct {
	// client is the instance of the Ceph Cloud Storage Go client.
	// Once this field is set, it must not be written again/unset to nil.
	_client *minio.Client
	// mu guards all fields.
	mu sync.Mutex
}

// ListBackups implements BackupStorage.
func (bs *CephBackupStorage) ListBackups(ctx context.Context, dir string) ([]backupstorage.BackupHandle, error) {
	c, err := bs.client()
	if err != nil {
		return nil, err
	}
	// ceph bucket name
	bucket := alterBucketName(dir)

	// List prefixes that begin with dir (i.e. list subdirs).
	var subdirs []string
	searchPrefix := objName(dir, "")

	doneCh := make(chan struct{})
	for object := range c.ListObjects(bucket, searchPrefix, false, doneCh) {
		if object.Err != nil {
			_, err := c.BucketExists(bucket)
			if err != nil {
				return nil, nil
			}
			return nil, object.Err
		}
		subdir := strings.TrimPrefix(object.Key, searchPrefix)
		subdir = strings.TrimSuffix(subdir, "/")
		subdirs = append(subdirs, subdir)
	}

	// Backups must be returned in order, oldest first.
	sort.Strings(subdirs)

	result := make([]backupstorage.BackupHandle, 0, len(subdirs))
	for _, subdir := range subdirs {
		result = append(result, &CephBackupHandle{
			client:   c,
			bs:       bs,
			dir:      dir,
			name:     subdir,
			readOnly: true,
		})
	}
	return result, nil
}

// StartBackup implements BackupStorage.
func (bs *CephBackupStorage) StartBackup(ctx context.Context, dir, name string) (backupstorage.BackupHandle, error) {
	c, err := bs.client()
	if err != nil {
		return nil, err
	}
	// ceph bucket name
	bucket := alterBucketName(dir)

	found, err := c.BucketExists(bucket)

	if err != nil {
		log.Info("Error from BucketExists: %v, quitting", bucket)
		return nil, errors.New("Error checking whether bucket exists: " + bucket)
	}
	if !found {
		log.Info("Bucket: %v doesn't exist, creating new bucket with the required name", bucket)
		err = c.MakeBucket(bucket, "")
		if err != nil {
			log.Info("Error creating Bucket: %v, quitting", bucket)
			return nil, errors.New("Error creating new bucket: " + bucket)
		}
	}

	return &CephBackupHandle{
		client:   c,
		bs:       bs,
		dir:      dir,
		name:     name,
		readOnly: false,
	}, nil
}

// RemoveBackup implements BackupStorage.
func (bs *CephBackupStorage) RemoveBackup(ctx context.Context, dir, name string) error {
	c, err := bs.client()
	if err != nil {
		return err
	}
	// ceph bucket name
	bucket := alterBucketName(dir)

	fullName := objName(dir, name, "")
	var arr []string
	doneCh := make(chan struct{})
	defer close(doneCh)
	for object := range c.ListObjects(bucket, fullName, true, doneCh) {
		if object.Err != nil {
			return object.Err
		}
		arr = append(arr, object.Key)
	}
	for _, obj := range arr {
		err = c.RemoveObject(bucket, obj)
		if err != nil {
			return err
		}
	}
	return nil
}

// Close implements BackupStorage.
func (bs *CephBackupStorage) Close() error {
	bs.mu.Lock()
	defer bs.mu.Unlock()

	if bs._client != nil {
		// a new client the next time one is needed.
		bs._client = nil
	}
	return nil
}

// client returns the Ceph Storage client instance.
// If there isn't one yet, it tries to create one.
func (bs *CephBackupStorage) client() (*minio.Client, error) {
	bs.mu.Lock()
	defer bs.mu.Unlock()

	if bs._client == nil {
		configFile, err := os.Open(*configFilePath)
		if err != nil {
			return nil, fmt.Errorf("file not present : %v", err)
		}
		defer configFile.Close()
		jsonParser := json.NewDecoder(configFile)
		if err = jsonParser.Decode(&storageConfig); err != nil {
			return nil, fmt.Errorf("error parsing the json file : %v", err)
		}

		accessKey := storageConfig.AccessKey
		secretKey := storageConfig.SecretKey
		url := storageConfig.EndPoint
		useSSL := storageConfig.UseSSL

		client, err := minio.NewV2(url, accessKey, secretKey, useSSL)
		if err != nil {
			return nil, err
		}
		bs._client = client
	}
	return bs._client, nil
}

func init() {
	backupstorage.BackupStorageMap["ceph"] = &CephBackupStorage{}
}

// objName joins path parts into an object name.
// Unlike path.Join, it doesn't collapse ".." or strip trailing slashes.
// It also adds the value of the -gcs_backup_storage_root flag if set.
func objName(parts ...string) string {
	return strings.Join(parts, "/")
}

// keeping in view the bucket naming conventions for ceph
// only keyspace informations is extracted and used for bucket name
func alterBucketName(dir string) string {
	bucket := strings.ToLower(dir)
	bucket = strings.Split(bucket, "/")[0]
	bucket = strings.Replace(bucket, "_", "-", -1)
	return bucket
}
<commit_msg>pass in correct filesize in ceph backup<commit_after>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package cephbackupstorage implements the BackupStorage interface
// for Ceph Cloud Storage.
package cephbackupstorage

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
	"sync"

	"errors"

	minio "github.com/minio/minio-go"
	"golang.org/x/net/context"

	"vitess.io/vitess/go/vt/concurrency"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/mysqlctl/backupstorage"
)

var (
	// configFilePath is where the configs/credentials for backups will be stored.
	configFilePath = flag.String("ceph_backup_storage_config", "ceph_backup_config.json",
		"Path to JSON config file for ceph backup storage")
)

var storageConfig struct {
	AccessKey string `json:"accessKey"`
	SecretKey string `json:"secretKey"`
	EndPoint  string `json:"endPoint"`
	UseSSL    bool   `json:"useSSL"`
}

// CephBackupHandle implements BackupHandle for Ceph Cloud Storage.
type CephBackupHandle struct {
	client    *minio.Client
	bs        *CephBackupStorage
	dir       string
	name      string
	readOnly  bool
	errors    concurrency.AllErrorRecorder
	waitGroup sync.WaitGroup
}

// Directory implements BackupHandle.
func (bh *CephBackupHandle) Directory() string {
	return bh.dir
}

// Name implements BackupHandle.
func (bh *CephBackupHandle) Name() string {
	return bh.name
}

// AddFile implements BackupHandle.
func (bh *CephBackupHandle) AddFile(ctx context.Context, filename string, filesize int64) (io.WriteCloser, error) {
	if bh.readOnly {
		return nil, fmt.Errorf("AddFile cannot be called on read-only backup")
	}
	reader, writer := io.Pipe()
	bh.waitGroup.Add(1)
	go func() {
		defer bh.waitGroup.Done()

		// ceph bucket name is where the backups will go
		//backup handle dir field contains keyspace/shard value
		bucket := alterBucketName(bh.dir)

		// Give PutObject() the read end of the pipe.
		object := objName(bh.dir, bh.name, filename)
		// if filesize is given as 0, pass it as -1 = UNKNOWN
		if filesize == 0 {
			filesize = -1
		}
		_, err := bh.client.PutObjectWithContext(ctx, bucket, object, reader, filesize, minio.PutObjectOptions{ContentType: "application/octet-stream"})
		if err != nil {
			// Signal the writer that an error occurred, in case it's not done writing yet.
			reader.CloseWithError(err)
			// In case the error happened after the writer finished, we need to remember it.
			bh.errors.RecordError(err)
		}
	}()
	// Give our caller the write end of the pipe.
	return writer, nil
}

// EndBackup implements BackupHandle.
func (bh *CephBackupHandle) EndBackup(ctx context.Context) error {
	if bh.readOnly {
		return fmt.Errorf("EndBackup cannot be called on read-only backup")
	}
	bh.waitGroup.Wait()
	// Return the saved PutObject() errors, if any.
	return bh.errors.Error()
}

// AbortBackup implements BackupHandle.
func (bh *CephBackupHandle) AbortBackup(ctx context.Context) error {
	if bh.readOnly {
		return fmt.Errorf("AbortBackup cannot be called on read-only backup")
	}
	return bh.bs.RemoveBackup(ctx, bh.dir, bh.name)
}

// ReadFile implements BackupHandle.
func (bh *CephBackupHandle) ReadFile(ctx context.Context, filename string) (io.ReadCloser, error) {
	if !bh.readOnly {
		return nil, fmt.Errorf("ReadFile cannot be called on read-write backup")
	}
	// ceph bucket name
	bucket := alterBucketName(bh.dir)
	object := objName(bh.dir, bh.name, filename)
	return bh.client.GetObjectWithContext(ctx, bucket, object, minio.GetObjectOptions{})
}

// CephBackupStorage implements BackupStorage for Ceph Cloud Storage.
type CephBackupStorage struct {
	// client is the instance of the Ceph Cloud Storage Go client.
	// Once this field is set, it must not be written again/unset to nil.
	_client *minio.Client
	// mu guards all fields.
	mu sync.Mutex
}

// ListBackups implements BackupStorage.
func (bs *CephBackupStorage) ListBackups(ctx context.Context, dir string) ([]backupstorage.BackupHandle, error) {
	c, err := bs.client()
	if err != nil {
		return nil, err
	}
	// ceph bucket name
	bucket := alterBucketName(dir)

	// List prefixes that begin with dir (i.e. list subdirs).
	var subdirs []string
	searchPrefix := objName(dir, "")

	doneCh := make(chan struct{})
	for object := range c.ListObjects(bucket, searchPrefix, false, doneCh) {
		if object.Err != nil {
			_, err := c.BucketExists(bucket)
			if err != nil {
				return nil, nil
			}
			return nil, object.Err
		}
		subdir := strings.TrimPrefix(object.Key, searchPrefix)
		subdir = strings.TrimSuffix(subdir, "/")
		subdirs = append(subdirs, subdir)
	}

	// Backups must be returned in order, oldest first.
	sort.Strings(subdirs)

	result := make([]backupstorage.BackupHandle, 0, len(subdirs))
	for _, subdir := range subdirs {
		result = append(result, &CephBackupHandle{
			client:   c,
			bs:       bs,
			dir:      dir,
			name:     subdir,
			readOnly: true,
		})
	}
	return result, nil
}

// StartBackup implements BackupStorage.
func (bs *CephBackupStorage) StartBackup(ctx context.Context, dir, name string) (backupstorage.BackupHandle, error) {
	c, err := bs.client()
	if err != nil {
		return nil, err
	}
	// ceph bucket name
	bucket := alterBucketName(dir)

	found, err := c.BucketExists(bucket)

	if err != nil {
		log.Info("Error from BucketExists: %v, quitting", bucket)
		return nil, errors.New("Error checking whether bucket exists: " + bucket)
	}
	if !found {
		log.Info("Bucket: %v doesn't exist, creating new bucket with the required name", bucket)
		err = c.MakeBucket(bucket, "")
		if err != nil {
			log.Info("Error creating Bucket: %v, quitting", bucket)
			return nil, errors.New("Error creating new bucket: " + bucket)
		}
	}

	return &CephBackupHandle{
		client:   c,
		bs:       bs,
		dir:      dir,
		name:     name,
		readOnly: false,
	}, nil
}

// RemoveBackup implements BackupStorage.
func (bs *CephBackupStorage) RemoveBackup(ctx context.Context, dir, name string) error {
	c, err := bs.client()
	if err != nil {
		return err
	}
	// ceph bucket name
	bucket := alterBucketName(dir)

	fullName := objName(dir, name, "")
	var arr []string
	doneCh := make(chan struct{})
	defer close(doneCh)
	for object := range c.ListObjects(bucket, fullName, true, doneCh) {
		if object.Err != nil {
			return object.Err
		}
		arr = append(arr, object.Key)
	}
	for _, obj := range arr {
		err = c.RemoveObject(bucket, obj)
		if err != nil {
			return err
		}
	}
	return nil
}

// Close implements BackupStorage.
func (bs *CephBackupStorage) Close() error {
	bs.mu.Lock()
	defer bs.mu.Unlock()

	if bs._client != nil {
		// a new client the next time one is needed.
		bs._client = nil
	}
	return nil
}

// client returns the Ceph Storage client instance.
// If there isn't one yet, it tries to create one.
func (bs *CephBackupStorage) client() (*minio.Client, error) {
	bs.mu.Lock()
	defer bs.mu.Unlock()

	if bs._client == nil {
		configFile, err := os.Open(*configFilePath)
		if err != nil {
			return nil, fmt.Errorf("file not present : %v", err)
		}
		defer configFile.Close()
		jsonParser := json.NewDecoder(configFile)
		if err = jsonParser.Decode(&storageConfig); err != nil {
			return nil, fmt.Errorf("error parsing the json file : %v", err)
		}

		accessKey := storageConfig.AccessKey
		secretKey := storageConfig.SecretKey
		url := storageConfig.EndPoint
		useSSL := storageConfig.UseSSL

		client, err := minio.NewV2(url, accessKey, secretKey, useSSL)
		if err != nil {
			return nil, err
		}
		bs._client = client
	}
	return bs._client, nil
}

func init() {
	backupstorage.BackupStorageMap["ceph"] = &CephBackupStorage{}
}

// objName joins path parts into an object name.
// Unlike path.Join, it doesn't collapse ".." or strip trailing slashes.
// It also adds the value of the -gcs_backup_storage_root flag if set.
func objName(parts ...string) string {
	return strings.Join(parts, "/")
}

// keeping in view the bucket naming conventions for ceph
// only keyspace informations is extracted and used for bucket name
func alterBucketName(dir string) string {
	bucket := strings.ToLower(dir)
	bucket = strings.Split(bucket, "/")[0]
	bucket = strings.Replace(bucket, "_", "-", -1)
	return bucket
}
<|endoftext|>


<commit_before>package dormantdatabase

import (
	"sync"

	hookapi "github.com/appscode/kubernetes-webhook-util/admission/v1beta1"
	dynamic_util "github.com/appscode/kutil/dynamic"
	meta_util "github.com/appscode/kutil/meta"
	api "github.com/kubedb/apimachinery/apis/kubedb/v1alpha1"
	cs "github.com/kubedb/apimachinery/client/clientset/versioned"
	plugin "github.com/kubedb/apimachinery/pkg/admission"
	admission "k8s.io/api/admission/v1beta1"
	core "k8s.io/api/core/v1"
	kerr "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/kubernetes"
	clientsetscheme "k8s.io/client-go/kubernetes/scheme"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/reference"
)

type DormantDatabaseValidator struct {
	client      kubernetes.Interface
	dc          dynamic.Interface
	extClient   cs.Interface
	lock        sync.RWMutex
	initialized bool
}

var _ hookapi.AdmissionHook = &DormantDatabaseValidator{}

func (a *DormantDatabaseValidator) Resource() (plural schema.GroupVersionResource, singular string) {
	return schema.GroupVersionResource{
			Group:    "validators.kubedb.com",
			Version:  "v1alpha1",
			Resource: "dormantdatabases",
		},
		"dormantdatabase"
}

func (a *DormantDatabaseValidator) Initialize(config *rest.Config, stopCh <-chan struct{}) error {
	a.lock.Lock()
	defer a.lock.Unlock()

	a.initialized = true

	var err error
	if a.client, err = kubernetes.NewForConfig(config); err != nil {
		return err
	}
	if a.dc, err = dynamic.NewForConfig(config); err != nil {
		return err
	}
	if a.extClient, err = cs.NewForConfig(config); err != nil {
		return err
	}
	return err
}

func (a *DormantDatabaseValidator) Admit(req *admission.AdmissionRequest) *admission.AdmissionResponse {
	status := &admission.AdmissionResponse{}

	// No validation on CREATE
	if (req.Operation != admission.Update && req.Operation != admission.Delete) ||
		len(req.SubResource) != 0 ||
		req.Kind.Group != api.SchemeGroupVersion.Group ||
		req.Kind.Kind != api.ResourceKindDormantDatabase {
		status.Allowed = true
		return status
	}

	a.lock.RLock()
	defer a.lock.RUnlock()
	if !a.initialized {
		return hookapi.StatusUninitialized()
	}

	switch req.Operation {
	case admission.Delete:
		if req.Name != "" {
			// req.Object.Raw = nil, so read from kubernetes
			obj, err := a.extClient.KubedbV1alpha1().DormantDatabases(req.Namespace).Get(req.Name, metav1.GetOptions{})
			if err != nil && !kerr.IsNotFound(err) {
				return hookapi.StatusInternalServerError(err)
			} else if kerr.IsNotFound(err) {
				break
			}
			if err := a.handleOwnerReferences(obj); err != nil {
				return hookapi.StatusInternalServerError(err)
			}
		}
	case admission.Update:
		// validate the operation made by User
		obj, err := meta_util.UnmarshalFromJSON(req.Object.Raw, api.SchemeGroupVersion)
		if err != nil {
			return hookapi.StatusBadRequest(err)
		}
		OldObj, err := meta_util.UnmarshalFromJSON(req.OldObject.Raw, api.SchemeGroupVersion)
		if err != nil {
			return hookapi.StatusBadRequest(err)
		}
		if err := plugin.ValidateUpdate(obj, OldObj, req.Kind.Kind); err != nil {
			return hookapi.StatusBadRequest(err)
		}
	}

	status.Allowed = true
	return status
}

func (a *DormantDatabaseValidator) handleOwnerReferences(dormantDatabase *api.DormantDatabase) error {
	if dormantDatabase.Spec.WipeOut {
		if err := a.setOwnerReferenceToObjects(dormantDatabase); err != nil {
			return err
		}
	} else {
		if err := a.removeOwnerReferenceFromObjects(dormantDatabase); err != nil {
			return err
		}
	}
	return nil
}

func (a *DormantDatabaseValidator) setOwnerReferenceToObjects(dormantDatabase *api.DormantDatabase) error {
	// Get LabelSelector for Other Components first
	dbKind, err := meta_util.GetStringValue(dormantDatabase.ObjectMeta.Labels, api.LabelDatabaseKind)
	if err != nil {
		return err
	}
	labelMap := map[string]string{
		api.LabelDatabaseName: dormantDatabase.Name,
		api.LabelDatabaseKind: dbKind,
	}
	selector := labels.SelectorFromSet(labelMap)

	// Get object reference of dormant database
	ref, rerr := reference.GetReference(clientsetscheme.Scheme, dormantDatabase)
	if rerr != nil {
		return rerr
	}
	if err := dynamic_util.EnsureOwnerReferenceForSelector(
		a.dc,
		api.SchemeGroupVersion.WithResource(api.ResourcePluralSnapshot),
		dormantDatabase.Namespace,
		selector,
		ref); err != nil {
		return nil
	}
	if err := dynamic_util.EnsureOwnerReferenceForSelector(
		a.dc,
		core.SchemeGroupVersion.WithResource("persistentvolumeclaims"),
		"", // non-namespaced
		selector,
		ref); err != nil {
		return nil
	}
	if err := dynamic_util.EnsureOwnerReferenceForItems(
		a.dc,
		core.SchemeGroupVersion.WithResource("secrets"),
		dormantDatabase.Namespace,
		dormantDatabase.GetDatabaseSecrets(),
		ref); err != nil {
		return nil
	}
	return nil
}

func (a *DormantDatabaseValidator) removeOwnerReferenceFromObjects(dormantDatabase *api.DormantDatabase) error {
	// First, Get LabelSelector for Other Components
	dbKind, err := meta_util.GetStringValue(dormantDatabase.ObjectMeta.Labels, api.LabelDatabaseKind)
	if err != nil {
		return err
	}
	labelMap := map[string]string{
		api.LabelDatabaseName: dormantDatabase.Name,
		api.LabelDatabaseKind: dbKind,
	}
	selector := labels.SelectorFromSet(labelMap)

	// Get object reference of dormant database
	ref, rerr := reference.GetReference(clientsetscheme.Scheme, dormantDatabase)
	if rerr != nil {
		return rerr
	}
	if err := dynamic_util.RemoveOwnerReferenceForSelector(
		a.dc,
		api.SchemeGroupVersion.WithResource(api.ResourcePluralSnapshot),
		dormantDatabase.Namespace,
		selector,
		ref); err != nil {
		return nil
	}
	if err := dynamic_util.RemoveOwnerReferenceForSelector(
		a.dc,
		core.SchemeGroupVersion.WithResource("persistentvolumeclaims"),
		"", // non-namespaced
		selector,
		ref); err != nil {
		return nil
	}
	if err := dynamic_util.RemoveOwnerReferenceForItems(
		a.dc,
		core.SchemeGroupVersion.WithResource("secrets"),
		dormantDatabase.Namespace,
		dormantDatabase.GetDatabaseSecrets(),
		ref); err != nil {
		return nil
	}
	return nil
}
<commit_msg>Provide namespace to DynamicClient for PersistentVolumeClaims resource (#308)<commit_after>package dormantdatabase

import (
	"sync"

	hookapi "github.com/appscode/kubernetes-webhook-util/admission/v1beta1"
	dynamic_util "github.com/appscode/kutil/dynamic"
	meta_util "github.com/appscode/kutil/meta"
	api "github.com/kubedb/apimachinery/apis/kubedb/v1alpha1"
	cs "github.com/kubedb/apimachinery/client/clientset/versioned"
	plugin "github.com/kubedb/apimachinery/pkg/admission"
	admission "k8s.io/api/admission/v1beta1"
	core "k8s.io/api/core/v1"
	kerr "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/kubernetes"
	clientsetscheme "k8s.io/client-go/kubernetes/scheme"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/reference"
)

type DormantDatabaseValidator struct {
	client      kubernetes.Interface
	dc          dynamic.Interface
	extClient   cs.Interface
	lock        sync.RWMutex
	initialized bool
}

var _ hookapi.AdmissionHook = &DormantDatabaseValidator{}

func (a *DormantDatabaseValidator) Resource() (plural schema.GroupVersionResource, singular string) {
	return schema.GroupVersionResource{
			Group:    "validators.kubedb.com",
			Version:  "v1alpha1",
			Resource: "dormantdatabases",
		},
		"dormantdatabase"
}

func (a *DormantDatabaseValidator) Initialize(config *rest.Config, stopCh <-chan struct{}) error {
	a.lock.Lock()
	defer a.lock.Unlock()

	a.initialized = true

	var err error
	if a.client, err = kubernetes.NewForConfig(config); err != nil {
		return err
	}
	if a.dc, err = dynamic.NewForConfig(config); err != nil {
		return err
	}
	if a.extClient, err = cs.NewForConfig(config); err != nil {
		return err
	}
	return err
}

func (a *DormantDatabaseValidator) Admit(req *admission.AdmissionRequest) *admission.AdmissionResponse {
	status := &admission.AdmissionResponse{}

	// No validation on CREATE
	if (req.Operation != admission.Update && req.Operation != admission.Delete) ||
		len(req.SubResource) != 0 ||
		req.Kind.Group != api.SchemeGroupVersion.Group ||
		req.Kind.Kind != api.ResourceKindDormantDatabase {
		status.Allowed = true
		return status
	}

	a.lock.RLock()
	defer a.lock.RUnlock()
	if !a.initialized {
		return hookapi.StatusUninitialized()
	}

	switch req.Operation {
	case admission.Delete:
		if req.Name != "" {
			// req.Object.Raw = nil, so read from kubernetes
			obj, err := a.extClient.KubedbV1alpha1().DormantDatabases(req.Namespace).Get(req.Name, metav1.GetOptions{})
			if err != nil && !kerr.IsNotFound(err) {
				return hookapi.StatusInternalServerError(err)
			} else if kerr.IsNotFound(err) {
				break
			}
			if err := a.handleOwnerReferences(obj); err != nil {
				return hookapi.StatusInternalServerError(err)
			}
		}
	case admission.Update:
		// validate the operation made by User
		obj, err := meta_util.UnmarshalFromJSON(req.Object.Raw, api.SchemeGroupVersion)
		if err != nil {
			return hookapi.StatusBadRequest(err)
		}
		OldObj, err := meta_util.UnmarshalFromJSON(req.OldObject.Raw, api.SchemeGroupVersion)
		if err != nil {
			return hookapi.StatusBadRequest(err)
		}
		if err := plugin.ValidateUpdate(obj, OldObj, req.Kind.Kind); err != nil {
			return hookapi.StatusBadRequest(err)
		}
	}

	status.Allowed = true
	return status
}

func (a *DormantDatabaseValidator) handleOwnerReferences(dormantDatabase *api.DormantDatabase) error {
	if dormantDatabase.Spec.WipeOut {
		if err := a.setOwnerReferenceToObjects(dormantDatabase); err != nil {
			return err
		}
	} else {
		if err := a.removeOwnerReferenceFromObjects(dormantDatabase); err != nil {
			return err
		}
	}
	return nil
}

func (a *DormantDatabaseValidator) setOwnerReferenceToObjects(dormantDatabase *api.DormantDatabase) error {
	// Get LabelSelector for Other Components first
	dbKind, err := meta_util.GetStringValue(dormantDatabase.ObjectMeta.Labels, api.LabelDatabaseKind)
	if err != nil {
		return err
	}
	labelMap := map[string]string{
		api.LabelDatabaseName: dormantDatabase.Name,
		api.LabelDatabaseKind: dbKind,
	}
	selector := labels.SelectorFromSet(labelMap)

	// Get object reference of dormant database
	ref, rerr := reference.GetReference(clientsetscheme.Scheme, dormantDatabase)
	if rerr != nil {
		return rerr
	}
	if err := dynamic_util.EnsureOwnerReferenceForSelector(
		a.dc,
		api.SchemeGroupVersion.WithResource(api.ResourcePluralSnapshot),
		dormantDatabase.Namespace,
		selector,
		ref); err != nil {
		return nil
	}
	if err := dynamic_util.EnsureOwnerReferenceForSelector(
		a.dc,
		core.SchemeGroupVersion.WithResource("persistentvolumeclaims"),
		dormantDatabase.Namespace,
		selector,
		ref); err != nil {
		return nil
	}
	if err := dynamic_util.EnsureOwnerReferenceForItems(
		a.dc,
		core.SchemeGroupVersion.WithResource("secrets"),
		dormantDatabase.Namespace,
		dormantDatabase.GetDatabaseSecrets(),
		ref); err != nil {
		return nil
	}
	return nil
}

func (a *DormantDatabaseValidator) removeOwnerReferenceFromObjects(dormantDatabase *api.DormantDatabase) error {
	// First, Get LabelSelector for Other Components
	dbKind, err := meta_util.GetStringValue(dormantDatabase.ObjectMeta.Labels, api.LabelDatabaseKind)
	if err != nil {
		return err
	}
	labelMap := map[string]string{
		api.LabelDatabaseName: dormantDatabase.Name,
		api.LabelDatabaseKind: dbKind,
	}
	selector := labels.SelectorFromSet(labelMap)

	// Get object reference of dormant database
	ref, rerr := reference.GetReference(clientsetscheme.Scheme, dormantDatabase)
	if rerr != nil {
		return rerr
	}
	if err := dynamic_util.RemoveOwnerReferenceForSelector(
		a.dc,
		api.SchemeGroupVersion.WithResource(api.ResourcePluralSnapshot),
		dormantDatabase.Namespace,
		selector,
		ref); err != nil {
		return nil
	}
	if err := dynamic_util.RemoveOwnerReferenceForSelector(
		a.dc,
		core.SchemeGroupVersion.WithResource("persistentvolumeclaims"),
		dormantDatabase.Namespace,
		selector,
		ref); err != nil {
		return nil
	}
	if err := dynamic_util.RemoveOwnerReferenceForItems(
		a.dc,
		core.SchemeGroupVersion.WithResource("secrets"),
		dormantDatabase.Namespace,
		dormantDatabase.GetDatabaseSecrets(),
		ref); err != nil {
		return nil
	}
	return nil
}
<|endoftext|>


<commit_before>package derivatives_test

import (
	"testing"

	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/derivatives"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestNewDerivativeStatusFromWsRaw(t *testing.T) {
	t.Run("insufficient arguments", func(t *testing.T) {
		payload := []interface{}{float64(1591614631576)}

		d, err := derivatives.NewDerivativeStatusFromWsRaw("tBTCF0:USTF0", payload)
		require.NotNil(t, err)
		require.Nil(t, d)
	})

	t.Run("valid arguments", func(t *testing.T) {
		payload := []interface{}{
			float64(1591614631576),
			nil,
			9271.1234567,
			9275.3,
			nil,
			1391472.27686063,
			nil,
			1594656000000,
			-0.00011968,
			3144,
			nil,
			0,
			nil,
			nil,
			9276.06,
			nil,
			nil,
			3813.72957182,
		}

		d, err := derivatives.NewDerivativeStatusFromWsRaw("tBTCF0:USTF0", payload)
		require.Nil(t, err)

		expected := &derivatives.DerivativeStatus{
			Symbol:               "tBTCF0:USTF0",
			MTS:                  1591614631576,
			Price:                9271.1234567,
			SpotPrice:            9275.3,
			InsuranceFundBalance: 1.39147227686063e+06,
			FundingAccrued:       -0.00011968,
			MarkPrice:            9276.06,
			OpenInterest:         3813.72957182,
		}
		assert.Equal(t, expected, d)
	})
}

func TestNewDerivativeStatusFromRaw(t *testing.T) {
	t.Run("insufficient arguments", func(t *testing.T) {
		payload := []interface{}{"tBTCF0:USTF0"}

		d, err := derivatives.NewDerivativeStatusFromRaw(payload)
		require.NotNil(t, err)
		require.Nil(t, d)
	})

	t.Run("valid arguments", func(t *testing.T) {
		payload := []interface{}{
			"tBTCF0:USTF0",
			float64(1591614631576),
			nil,
			9271.1234567,
			9275.3,
			nil,
			1391472.27686063,
			nil,
			1594656000000,
			-0.00011968,
			3144,
			nil,
			0,
			nil,
			nil,
			9276.06,
			nil,
			nil,
			3813.72957182,
		}

		d, err := derivatives.NewDerivativeStatusFromRaw(payload)
		require.Nil(t, err)

		expected := &derivatives.DerivativeStatus{
			Symbol:               "tBTCF0:USTF0",
			MTS:                  1591614631576,
			Price:                9271.1234567,
			SpotPrice:            9275.3,
			InsuranceFundBalance: 1.39147227686063e+06,
			FundingAccrued:       -0.00011968,
			MarkPrice:            9276.06,
			OpenInterest:         3813.72957182,
		}
		assert.Equal(t, expected, d)
	})
}
<commit_msg>adding more tests to cover snapshot functionaity for derivatives status<commit_after>package derivatives_test

import (
	"testing"

	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/derivatives"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestNewDerivativeStatusFromWsRaw(t *testing.T) {
	t.Run("insufficient arguments", func(t *testing.T) {
		payload := []interface{}{float64(1591614631576)}

		d, err := derivatives.FromWsRaw("tBTCF0:USTF0", payload)
		require.NotNil(t, err)
		require.Nil(t, d)
	})

	t.Run("valid arguments", func(t *testing.T) {
		payload := []interface{}{
			float64(1591614631576),
			nil,
			9271.1234567,
			9275.3,
			nil,
			1391472.27686063,
			nil,
			1594656000000,
			-0.00011968,
			3144,
			nil,
			0,
			nil,
			nil,
			9276.06,
			nil,
			nil,
			3813.72957182,
		}

		d, err := derivatives.FromWsRaw("tBTCF0:USTF0", payload)
		require.Nil(t, err)

		expected := &derivatives.DerivativeStatus{
			Symbol:               "tBTCF0:USTF0",
			MTS:                  1591614631576,
			Price:                9271.1234567,
			SpotPrice:            9275.3,
			InsuranceFundBalance: 1.39147227686063e+06,
			FundingAccrued:       -0.00011968,
			FundingStep:          3144,
			MarkPrice:            9276.06,
			OpenInterest:         3813.72957182,
		}
		assert.Equal(t, expected, d)
	})
}

func TestNewDerivativeStatusFromRaw(t *testing.T) {
	t.Run("insufficient arguments", func(t *testing.T) {
		payload := []interface{}{"tBTCF0:USTF0"}

		d, err := derivatives.FromRaw(payload)
		require.NotNil(t, err)
		require.Nil(t, d)
	})

	t.Run("valid arguments", func(t *testing.T) {
		payload := []interface{}{
			"tBTCF0:USTF0",
			float64(1591614631576),
			nil,
			9271.1234567,
			9275.3,
			nil,
			1391472.27686063,
			nil,
			1594656000000,
			-0.00011968,
			3144,
			nil,
			0,
			nil,
			nil,
			9276.06,
			nil,
			nil,
			3813.72957182,
		}

		d, err := derivatives.FromRaw(payload)
		require.Nil(t, err)

		expected := &derivatives.DerivativeStatus{
			Symbol:               "tBTCF0:USTF0",
			MTS:                  1591614631576,
			Price:                9271.1234567,
			SpotPrice:            9275.3,
			InsuranceFundBalance: 1.39147227686063e+06,
			FundingAccrued:       -0.00011968,
			FundingStep:          3144,
			MarkPrice:            9276.06,
			OpenInterest:         3813.72957182,
		}
		assert.Equal(t, expected, d)
	})
}

func TestSnapshotFromRaw(t *testing.T) {
	t.Run("invalid arguments", func(t *testing.T) {
		payload := [][]interface{}{{"tBTCF0:USTF0"}}
		ss, err := derivatives.SnapshotFromRaw(payload)
		require.NotNil(t, err)
		require.Nil(t, ss)
	})

	t.Run("valid arguments", func(t *testing.T) {
		payload := [][]interface{}{
			{
				"tBTCF0:USTF0",
				float64(1591614631576),
				nil,
				9271.1234567,
				9275.3,
				nil,
				1391472.27686063,
				nil,
				1594656000000,
				-0.00011968,
				3144,
				nil,
				0,
				nil,
				nil,
				9276.06,
				nil,
				nil,
				3813.72957182,
			},
			{
				"tBTCF0:USTF0",
				float64(1591614631576),
				nil,
				9271.1234567,
				9275.3,
				nil,
				1391472.27686063,
				nil,
				1594656000000,
				-0.00011968,
				3200,
				nil,
				0,
				nil,
				nil,
				9276.06,
				nil,
				nil,
				3813.72957182,
			},
		}
		ss, err := derivatives.SnapshotFromRaw(payload)
		require.Nil(t, err)

		expected := &derivatives.DerivativeStatusSnapshot{
			Snapshot: []*derivatives.DerivativeStatus{
				{
					Symbol:               "tBTCF0:USTF0",
					MTS:                  1591614631576,
					Price:                9271.1234567,
					SpotPrice:            9275.3,
					InsuranceFundBalance: 1.39147227686063e+06,
					FundingAccrued:       -0.00011968,
					FundingStep:          3144,
					MarkPrice:            9276.06,
					OpenInterest:         3813.72957182,
				},
				{
					Symbol:               "tBTCF0:USTF0",
					MTS:                  1591614631576,
					Price:                9271.1234567,
					SpotPrice:            9275.3,
					InsuranceFundBalance: 1.39147227686063e+06,
					FundingAccrued:       -0.00011968,
					FundingStep:          3200,
					MarkPrice:            9276.06,
					OpenInterest:         3813.72957182,
				},
			},
		}

		assert.Equal(t, expected, ss)
	})
}
<|endoftext|>