<commit_before>package mackerel

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
)

// MetricValue metric value
type MetricValue struct {
	Name  string      `json:"name,omitempty"`
	Time  int64       `json:"time,omitempty"`
	Value interface{} `json:"value,omitempty"`
}

// HostMetricValue host metric value
type HostMetricValue struct {
	HostID string `json:"hostId,omitempty"`
	*MetricValue
}

// LatestMetricValues latest metric value
type LatestMetricValues map[string]map[string]*MetricValue

// PostHostMetricValues post host metrics
func (c *Client) PostHostMetricValues(metricValues [](*HostMetricValue)) error {
	resp, err := c.PostJSON("/api/v0/tsdb", metricValues)
	defer closeResponse(resp)
	return err
}

// PostHostMetricValuesByHostID post host metrics
func (c *Client) PostHostMetricValuesByHostID(hostID string, metricValues [](*MetricValue)) error {
	var hostMetricValues []*HostMetricValue
	for _, metricValue := range metricValues {
		hostMetricValues = append(hostMetricValues, &HostMetricValue{
			HostID:      hostID,
			MetricValue: metricValue,
		})
	}
	return c.PostHostMetricValues(hostMetricValues)
}

// PostServiceMetricValues post service metrics
func (c *Client) PostServiceMetricValues(serviceName string, metricValues [](*MetricValue)) error {
	resp, err := c.PostJSON(fmt.Sprintf("/api/v0/services/%s/tsdb", serviceName), metricValues)
	defer closeResponse(resp)
	return err
}

// FetchLatestMetricValues fetch latest metrics
func (c *Client) FetchLatestMetricValues(hostIDs []string, metricNames []string) (LatestMetricValues, error) {
	v := url.Values{}
	for _, hostID := range hostIDs {
		v.Add("hostId", hostID)
	}
	for _, metricName := range metricNames {
		v.Add("name", metricName)
	}

	req, err := http.NewRequest("GET", fmt.Sprintf("%s?%s", c.urlFor("/api/v0/tsdb/latest").String(), v.Encode()), nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.Request(req)
	defer closeResponse(resp)
	if err != nil {
		return nil, err
	}

	var data struct {
		LatestMetricValues LatestMetricValues `json:"tsdbLatest"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return nil, err
	}

	return data.LatestMetricValues, err
}

// FetchHostMetricValues retrieves the metric values for a Host
func (c *Client) FetchHostMetricValues(hostID string, metricName string, from int64, to int64) ([]MetricValue, error) {
	return c.fetchMetricValues(&hostID, nil, metricName, from, to)
}

// FetchServiceMetricValues retrieves the metric values for a Service
func (c *Client) FetchServiceMetricValues(serviceName string, metricName string, from int64, to int64) ([]MetricValue, error) {
	return c.fetchMetricValues(nil, &serviceName, metricName, from, to)
}

func (c *Client) fetchMetricValues(hostID *string, serviceName *string, metricName string, from int64, to int64) ([]MetricValue, error) {
	v := url.Values{}
	v.Add("name", metricName)
	v.Add("from", strconv.FormatInt(from, 10))
	v.Add("to", strconv.FormatInt(to, 10))

	url := ""
	if hostID != nil {
		url = "/api/v0/hosts/" + *hostID + "/metrics"
	} else if serviceName != nil {
		url = "/api/v0/services/" + *serviceName + "/metrics"
	} else {
		return nil, errors.New("specify either host or service")
	}
	req, err := http.NewRequest("GET", fmt.Sprintf("%s?%s", c.urlFor(url).String(), v.Encode()), nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.Request(req)
	defer closeResponse(resp)
	if err != nil {
		return nil, err
	}

	var data struct {
		MetricValues []MetricValue `json:"metrics"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return nil, err
	}
	return data.MetricValues, err
}
<commit_msg>avoid string pointers<commit_after>package mackerel

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
)

// MetricValue metric value
type MetricValue struct {
	Name  string      `json:"name,omitempty"`
	Time  int64       `json:"time,omitempty"`
	Value interface{} `json:"value,omitempty"`
}

// HostMetricValue host metric value
type HostMetricValue struct {
	HostID string `json:"hostId,omitempty"`
	*MetricValue
}

// LatestMetricValues latest metric value
type LatestMetricValues map[string]map[string]*MetricValue

// PostHostMetricValues post host metrics
func (c *Client) PostHostMetricValues(metricValues [](*HostMetricValue)) error {
	resp, err := c.PostJSON("/api/v0/tsdb", metricValues)
	defer closeResponse(resp)
	return err
}

// PostHostMetricValuesByHostID post host metrics
func (c *Client) PostHostMetricValuesByHostID(hostID string, metricValues [](*MetricValue)) error {
	var hostMetricValues []*HostMetricValue
	for _, metricValue := range metricValues {
		hostMetricValues = append(hostMetricValues, &HostMetricValue{
			HostID:      hostID,
			MetricValue: metricValue,
		})
	}
	return c.PostHostMetricValues(hostMetricValues)
}

// PostServiceMetricValues post service metrics
func (c *Client) PostServiceMetricValues(serviceName string, metricValues [](*MetricValue)) error {
	resp, err := c.PostJSON(fmt.Sprintf("/api/v0/services/%s/tsdb", serviceName), metricValues)
	defer closeResponse(resp)
	return err
}

// FetchLatestMetricValues fetch latest metrics
func (c *Client) FetchLatestMetricValues(hostIDs []string, metricNames []string) (LatestMetricValues, error) {
	v := url.Values{}
	for _, hostID := range hostIDs {
		v.Add("hostId", hostID)
	}
	for _, metricName := range metricNames {
		v.Add("name", metricName)
	}

	req, err := http.NewRequest("GET", fmt.Sprintf("%s?%s", c.urlFor("/api/v0/tsdb/latest").String(), v.Encode()), nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.Request(req)
	defer closeResponse(resp)
	if err != nil {
		return nil, err
	}

	var data struct {
		LatestMetricValues LatestMetricValues `json:"tsdbLatest"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return nil, err
	}

	return data.LatestMetricValues, err
}

// FetchHostMetricValues retrieves the metric values for a Host
func (c *Client) FetchHostMetricValues(hostID string, metricName string, from int64, to int64) ([]MetricValue, error) {
	return c.fetchMetricValues(hostID, "", metricName, from, to)
}

// FetchServiceMetricValues retrieves the metric values for a Service
func (c *Client) FetchServiceMetricValues(serviceName string, metricName string, from int64, to int64) ([]MetricValue, error) {
	return c.fetchMetricValues("", serviceName, metricName, from, to)
}

func (c *Client) fetchMetricValues(hostID string, serviceName string, metricName string, from int64, to int64) ([]MetricValue, error) {
	v := url.Values{}
	v.Add("name", metricName)
	v.Add("from", strconv.FormatInt(from, 10))
	v.Add("to", strconv.FormatInt(to, 10))

	url := ""
	if hostID != "" {
		url = "/api/v0/hosts/" + hostID + "/metrics"
	} else if serviceName != "" {
		url = "/api/v0/services/" + serviceName + "/metrics"
	} else {
		return nil, errors.New("specify either host or service")
	}
	req, err := http.NewRequest("GET", fmt.Sprintf("%s?%s", c.urlFor(url).String(), v.Encode()), nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.Request(req)
	defer closeResponse(resp)
	if err != nil {
		return nil, err
	}

	var data struct {
		MetricValues []MetricValue `json:"metrics"`
	}
	err = json.NewDecoder(resp.Body).Decode(&data)
	if err != nil {
		return nil, err
	}
	return data.MetricValues, err
}
<|endoftext|>


<commit_before>package aggro

import (
	"fmt"
	"math"
	"sort"

	"github.com/shopspring/decimal"
)

type Metric struct {
	Type  string
	Field string
}

func (m *Metric) measurer() (measurer, error) {
	switch m.Type {
	case "mean":
		return &mean{}, nil
	case "median":
		return &median{}, nil
	case "mode":
		return &mode{}, nil
	case "min":
		return &min{}, nil
	case "max":
		return &max{}, nil
	case "cardinality":
		return &cardinality{}, nil
	case "sum":
		return &sum{}, nil
	case "stdev":
		return &stdev{}, nil
	default:
		return nil, fmt.Errorf("Unknown metric: %s", m.Type)
	}
}

type measurer interface {
	AddDatum(interface{})
	Result() interface{}
}

// Mean
type mean struct {
	count int
	sum   decimal.Decimal
}

func (a *mean) AddDatum(datum interface{}) {
	amount := datum.(*decimal.Decimal)
	a.count++
	a.sum = a.sum.Add(*amount)
}

func (a *mean) Result() interface{} {
	if a.count == 0 {
		return nil
	}
	result, _ := a.sum.Div(decimal.NewFromFloat(float64(a.count))).Float64()
	return result
}

// Median
type median struct {
	list []decimal.Decimal
}

func (a *median) AddDatum(datum interface{}) {
	amount := datum.(*decimal.Decimal)
	a.list = append(a.list, *amount)
}

func (a *median) Result() interface{} {
	if len(a.list) == 0 {
		return nil
	}

	// Sort our list in numerical order.
	sort.Sort(decimalSortNumerical(a.list))

	// Determine median value.
	middle := len(a.list) / 2
	median, _ := a.list[middle].Float64()

	// Even slice? Get the mean of the middle values.
	if len(a.list)%2 == 0 {
		prev, _ := a.list[middle-1].Float64()
		median = (median + prev) / 2
	}

	return median
}

// Mode
type mode struct {
	count int
	list  []decimal.Decimal
	value decimal.Decimal
}

func (a *mode) AddDatum(datum interface{}) {
	amount := datum.(*decimal.Decimal)
	a.list = append(a.list, *amount)
	a.count++
}

func (a *mode) Result() interface{} {
	if len(a.list) == 0 {
		return nil
	}

	numbers := []float64{}
	for _, number := range a.list {
		value, _ := number.Float64()
		numbers = append(numbers, value)
	}

	modes := []float64{}
	freq := make(map[float64]int, len(numbers))
	tip := 0

	// Range our values determining the highest frequency (tip).
	for _, x := range numbers {
		freq[x]++
		if freq[x] > tip {
			tip = freq[x]
		}
	}
	for x, f := range freq {
		if f == tip {
			modes = append(modes, x)
		}
	}
	if tip == 1 || len(modes) == len(numbers) {
		modes = []float64{}
	}
	return modes
}

// Min
type min struct {
	amount *decimal.Decimal
}

func (a *min) AddDatum(datum interface{}) {
	amount := datum.(*decimal.Decimal)
	if a.amount == nil {
		a.amount = amount
	}
	if (a.amount).Cmp(*amount) > -1 {
		a.amount = amount
	}
}

func (a *min) Result() interface{} {
	if a.amount == nil {
		return nil
	}
	result, _ := a.amount.Float64()
	return result
}

// Max
type max struct {
	amount *decimal.Decimal
}

func (a *max) AddDatum(datum interface{}) {
	amount := datum.(*decimal.Decimal)
	if a.amount == nil {
		a.amount = amount
	}
	if (a.amount).Cmp(*amount) < 0 {
		a.amount = amount
	}
}

func (a *max) Result() interface{} {
	if a.amount == nil {
		return nil
	}
	result, _ := a.amount.Float64()
	return result
}

// Cadinality
type cardinality struct {
	size int
}

func (a *cardinality) AddDatum(datum interface{}) {
	a.size++
}

func (a *cardinality) Result() interface{} {
	return a.size
}

// Sum
type sum struct {
	sum decimal.Decimal
}

func (a *sum) AddDatum(datum interface{}) {
	amount := datum.(*decimal.Decimal)
	a.sum = a.sum.Add(*amount)
}

func (a *sum) Result() interface{} {
	result, _ := a.sum.Float64()
	return result
}

// Standard deviation
type stdev struct {
	count int
	sum   decimal.Decimal
	list  []decimal.Decimal
}

func (a *stdev) AddDatum(datum interface{}) {
	amount := datum.(*decimal.Decimal)
	a.count++
	a.sum = a.sum.Add(*amount)
	a.list = append(a.list, *amount)
}

func (a *stdev) Result() interface{} {
	if a.count == 0 {
		return nil
	}

	// Find the mean.
	mean, _ := a.sum.Div(decimal.NewFromFloat(float64(a.count))).Float64()

	// Find and add each values distance to the mean.
	total := 0.0
	for _, number := range a.list {
		val, _ := number.Float64()
		total += math.Pow(val-mean, 2)
	}

	// Divide by result set length.
	variance := total / float64(len(a.list)-1)
	return math.Sqrt(variance)
}

// Sorters
type decimalSortNumerical []decimal.Decimal

func (s decimalSortNumerical) Len() int      { return len(s) }
func (s decimalSortNumerical) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s decimalSortNumerical) Less(i, j int) bool {
	a, _ := s[i].Float64()
	b, _ := s[j].Float64()
	return a < b
}
<commit_msg>Comments. Implemented cardinality correctly. Added valueCount metric<commit_after>package aggro

import (
	"fmt"
	"math"
	"sort"

	"github.com/shopspring/decimal"
)

type Metric struct {
	Type  string
	Field string
}

func (m *Metric) measurer() (measurer, error) {
	switch m.Type {
	case "mean":
		return &mean{}, nil
	case "median":
		return &median{}, nil
	case "mode":
		return &mode{}, nil
	case "min":
		return &min{}, nil
	case "max":
		return &max{}, nil
	case "cardinality":
		return &cardinality{}, nil
	case "sum":
		return &sum{}, nil
	case "stdev":
		return &stdev{}, nil
	case "count":
		return &valueCount{}, nil
	default:
		return nil, fmt.Errorf("Unknown metric: %s", m.Type)
	}
}

type measurer interface {
	AddDatum(interface{})
	Result() interface{}
}

// Mean
// Your standard average. Sum all values and divide by the number of values.
type mean struct {
	count int
	sum   decimal.Decimal
}

func (a *mean) AddDatum(datum interface{}) {
	// Cast to *decimal.Decimal.
	amount := datum.(*decimal.Decimal)

	// Increase our count.
	a.count++

	// Add our value to our existing sum.
	a.sum = a.sum.Add(*amount)
}

func (a *mean) Result() interface{} {
	if a.count == 0 {
		return nil
	}

	// Divide our sum by the count.
	result, _ := a.sum.Div(decimal.NewFromFloat(float64(a.count))).Float64()
	return result
}

// Median
// Dataset should be in numerical order. Median is the middle value of our dataset.
// If there is no middle value (due to dataset having even number of values) then
// the median is the mean (average) of the middle two values.
type median struct {
	list []decimal.Decimal
}

func (a *median) AddDatum(datum interface{}) {
	// Cast to *decimal.Decimal.
	amount := datum.(*decimal.Decimal)

	// Append value to median slice.
	a.list = append(a.list, *amount)
}

func (a *median) Result() interface{} {
	if len(a.list) == 0 {
		return nil
	}

	// Sort our list in numerical order.
	sort.Sort(decimalSortNumerical(a.list))

	// Find the middle of our dataset.
	middle := len(a.list) / 2
	median, _ := a.list[middle].Float64()

	// Is our dataset length even? If so, we don't have a correct middle value.
	// In this case, take the middle two values of our dataset and determine the
	// mean of them.
	if len(a.list)%2 == 0 {
		prev, _ := a.list[middle-1].Float64()
		median = (median + prev) / 2
	}

	return median
}

// Mode
// Mode is the value(s) that occur most often within the dataset. If no values
// are repeated (or all values are repeated), then the dataset has no mode.
//
// 1) Range each value in our dataset.
type mode struct {
	count int
	list  []decimal.Decimal
	value decimal.Decimal
}

func (a *mode) AddDatum(datum interface{}) {
	// Cast to *decimal.Decimal.
	amount := datum.(*decimal.Decimal)

	// Append value to mode slice.
	a.list = append(a.list, *amount)

	// Increase our count.
	a.count++
}

func (a *mode) Result() interface{} {
	if len(a.list) == 0 {
		return nil
	}

	// Results slice.
	modes := []float64{}

	// 'tip' represents our highest frequency count across our entire dataset. A
	// dataset with a tip of '1' means no repeated values were found.
	tip := 0

	// Range our values building a frequency map. This represents each value and
	// the number of times it appears in our dataset.
	freq := make(map[float64]int, len(a.list))
	for _, val := range a.list {
		value, _ := val.Float64()
		freq[value]++
		if freq[value] > tip {
			tip = freq[value]
		}
	}

	// Range our frequency map, checking if our values count matches our tip. If so
	// we have a mode!
	for val, c := range freq {
		if c == tip {
			modes = append(modes, val)
		}
	}

	// If tip is 1 (no repeating values found), or length of resulting modes slice
	// matches our dataset, then return no mode (empty).
	if tip == 1 || len(modes) == len(a.list) {
		modes = []float64{}
	}
	return modes
}

// Min
// Min is the smallest value within the dataset.
type min struct {
	amount *decimal.Decimal
}

func (a *min) AddDatum(datum interface{}) {
	// Cast to *decimal.Decimal.
	amount := datum.(*decimal.Decimal)
	if a.amount == nil {
		a.amount = amount
	}

	// If value is < existing min.amount, assign as min.amount.
	if (a.amount).Cmp(*amount) > -1 {
		a.amount = amount
	}
}

func (a *min) Result() interface{} {
	if a.amount == nil {
		return nil
	}
	result, _ := a.amount.Float64()
	return result
}

// Max
// Max is the largest value within the dataset.
type max struct {
	amount *decimal.Decimal
}

func (a *max) AddDatum(datum interface{}) {
	// Cast to *decimal.Decimal.
	amount := datum.(*decimal.Decimal)
	if a.amount == nil {
		a.amount = amount
	}

	// If value is > existing max.amount, assign as max.amount.
	if (a.amount).Cmp(*amount) < 0 {
		a.amount = amount
	}
}

func (a *max) Result() interface{} {
	if a.amount == nil {
		return nil
	}
	result, _ := a.amount.Float64()
	return result
}

// Cardinality
// Cardinality is a count of unique values in our dataset.
type cardinality struct {
	values map[interface{}]int
}

func (a *cardinality) AddDatum(datum interface{}) {
	if a.values == nil {
		a.values = map[interface{}]int{}
	}

	// Track frequency of our values within the dataset. Handle
	switch t := datum.(type) {
	case *decimal.Decimal:
		floatVal, _ := t.Float64()
		a.values[floatVal]++
	case string:
		a.values[t]++
	}
}

func (a *cardinality) Result() interface{} {
	return len(a.values)
}

// Value Count
// valueCount is the total number of values in the dataset.
type valueCount struct {
	size int
}

func (a *valueCount) AddDatum(datum interface{}) {
	a.size++
}

func (a *valueCount) Result() interface{} {
	return a.size
}

// Sum
// Sum is all dataset values added together.
type sum struct {
	sum decimal.Decimal
}

func (a *sum) AddDatum(datum interface{}) {
	amount := datum.(*decimal.Decimal)

	// Add our value to existing sum.
	a.sum = a.sum.Add(*amount)
}

func (a *sum) Result() interface{} {
	result, _ := a.sum.Float64()
	return result
}

// Standard deviation
// Standard deviation is a representation of how spread out values in the dataset are.
// It's calculated as square root of 'variance'. Variance is the average of the
// squared differences from the mean.
//
// 1) Determine mean.
// 2) Then for value in the dataset, subtract the mean and square the result.
// 3) Calculate the mean of those squared differences (variance).
// 4) Square root the result.
type stdev struct {
	count int
	sum   decimal.Decimal
	list  []decimal.Decimal
}

func (a *stdev) AddDatum(datum interface{}) {
	// Cast to *decimal.Decimal.
	amount := datum.(*decimal.Decimal)

	// Increase our count.
	a.count++

	// Add our value to existing sum.
	a.sum = a.sum.Add(*amount)

	// Append value to stdev slice.
	a.list = append(a.list, *amount)
}

func (a *stdev) Result() interface{} {
	// stdev requires two or more rows to work with.
	if a.count < 2 {
		return nil
	}

	// 1) Determine the mean (avg).
	mean, _ := a.sum.Div(decimal.NewFromFloat(float64(a.count))).Float64()

	// 2) Ranging our dataset, subtract mean and square the result from each value.
	total := 0.0
	for _, number := range a.list {
		val, _ := number.Float64()
		total += math.Pow(val-mean, 2)
	}

	// 3) Calculate the variance (mean of our squared results).
	variance := total / float64(len(a.list)-1)

	// 4) Square the result.
	return math.Sqrt(variance)
}

// Sorters
type decimalSortNumerical []decimal.Decimal

func (s decimalSortNumerical) Len() int      { return len(s) }
func (s decimalSortNumerical) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s decimalSortNumerical) Less(i, j int) bool {
	a, _ := s[i].Float64()
	b, _ := s[j].Float64()
	return a < b
}
<|endoftext|>


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

package sqmetrics

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"runtime"
	"time"

	"github.com/rcrowley/go-metrics"
)

// SquareMetrics posts metrics to an HTTP/JSON bridge endpoint
type SquareMetrics struct {
	registry metrics.Registry
	url      string
	prefix   string
	hostname string
}

// NewMetrics is the entry point for this code
func NewMetrics(metricsURL, metricsPrefix string, registry metrics.Registry) *SquareMetrics {
	hostname, err := os.Hostname()
	if err != nil {
		panic(err)
	}

	metrics := &SquareMetrics{
		registry: registry,
		url:      metricsURL,
		prefix:   metricsPrefix,
		hostname: hostname,
	}

	if metricsURL != "" {
		go metrics.publishMetrics()
	}

	go metrics.collectSystemMetrics()
	return metrics
}

func (mb *SquareMetrics) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	metrics := mb.SerializeMetrics()
	raw, err := json.Marshal(metrics)
	if err != nil {
		panic(err)
	}
	w.Write(raw)
}

// Publish metrics to bridge
func (mb *SquareMetrics) publishMetrics() {
	for range time.Tick(1 * time.Second) {
		mb.postMetrics()
	}
}

// Collect memory usage metrics
func (mb *SquareMetrics) collectSystemMetrics() {
	var mem runtime.MemStats

	alloc := metrics.GetOrRegisterGauge("runtime.mem.alloc", mb.registry)
	totalAlloc := metrics.GetOrRegisterGauge("runtime.mem.total-alloc", mb.registry)
	sys := metrics.GetOrRegisterGauge("runtime.mem.sys", mb.registry)
	heapAlloc := metrics.GetOrRegisterGauge("runtime.mem.heap.alloc", mb.registry)
	heapSys := metrics.GetOrRegisterGauge("runtime.mem.heap.sys", mb.registry)
	heapInUse := metrics.GetOrRegisterGauge("runtime.mem.heap.in-use", mb.registry)
	stackSys := metrics.GetOrRegisterGauge("runtime.mem.stack.sys", mb.registry)
	stackInUse := metrics.GetOrRegisterGauge("runtime.mem.stack.in-use", mb.registry)
	gcPauseTotal := metrics.GetOrRegisterGauge("runtime.mem.gc.pause-total", mb.registry)
	gcCPUFraction := metrics.GetOrRegisterGaugeFloat64("runtime.mem.gc.cpu-fraction", mb.registry)
	numGoRoutines := metrics.GetOrRegisterGauge("runtime.goroutines", mb.registry)
	numCgoCalls := metrics.GetOrRegisterGauge("runtime.cgo-calls", mb.registry)

	for range time.Tick(1 * time.Second) {
		runtime.ReadMemStats(&mem)

		alloc.Update(int64(mem.Alloc))
		totalAlloc.Update(int64(mem.TotalAlloc))
		sys.Update(int64(mem.Sys))

		heapAlloc.Update(int64(mem.HeapAlloc))
		heapSys.Update(int64(mem.HeapSys))
		heapInUse.Update(int64(mem.HeapInuse))

		stackSys.Update(int64(mem.StackSys))
		stackInUse.Update(int64(mem.StackInuse))

		gcPauseTotal.Update(int64(mem.PauseTotalNs))
		gcCPUFraction.Update(mem.GCCPUFraction)

		numGoRoutines.Update(int64(runtime.NumGoroutine()))
		numCgoCalls.Update(int64(runtime.NumCgoCall()))
	}
}

func (mb *SquareMetrics) postMetrics() {
	metrics := mb.SerializeMetrics()
	raw, err := json.Marshal(metrics)
	if err != nil {
		panic(err)
	}
	resp, err := http.Post(mb.url, "application/json", bytes.NewReader(raw))
	if err == nil {
		resp.Body.Close()
	}
}

func (mb *SquareMetrics) serializeMetric(now int64, metric tuple) map[string]interface{} {
	return map[string]interface{}{
		"timestamp": now,
		"metric":    fmt.Sprintf("%s.%s", mb.prefix, metric.name),
		"value":     metric.value,
		"hostname":  mb.hostname,
	}
}

type tuple struct {
	name  string
	value interface{}
}

// SerializeMetrics returns a map of the collected metrics, suitable for JSON marshalling
func (mb *SquareMetrics) SerializeMetrics() []map[string]interface{} {
	nvs := []tuple{}

	mb.registry.Each(func(name string, i interface{}) {
		switch metric := i.(type) {
		case metrics.Counter:
			nvs = append(nvs, tuple{name, metric.Count()})
		case metrics.Gauge:
			nvs = append(nvs, tuple{name, metric.Value()})
		case metrics.GaugeFloat64:
			nvs = append(nvs, tuple{name, metric.Value()})
		case metrics.Timer:
			timer := metric.Snapshot()
			nvs = append(nvs, []tuple{
				{fmt.Sprintf("%s.count", name), timer.Count()},
				{fmt.Sprintf("%s.min", name), timer.Min()},
				{fmt.Sprintf("%s.max", name), timer.Max()},
				{fmt.Sprintf("%s.mean", name), timer.Mean()},
				{fmt.Sprintf("%s.std-dev", name), timer.StdDev()},
				{fmt.Sprintf("%s.one-minute", name), timer.Rate1()},
				{fmt.Sprintf("%s.five-minute", name), timer.Rate5()},
				{fmt.Sprintf("%s.fifteen-minute", name), timer.Rate15()},
				{fmt.Sprintf("%s.mean-rate", name), timer.RateMean()},
				{fmt.Sprintf("%s.50-percentile", name), timer.Percentile(0.5)},
				{fmt.Sprintf("%s.75-percentile", name), timer.Percentile(0.75)},
				{fmt.Sprintf("%s.95-percentile", name), timer.Percentile(0.95)},
				{fmt.Sprintf("%s.99-percentile", name), timer.Percentile(0.99)},
				{fmt.Sprintf("%s.999-percentile", name), timer.Percentile(0.999)},
			}...)
		}
	})

	now := time.Now().Unix()
	out := []map[string]interface{}{}
	for _, nv := range nvs {
		out = append(out, mb.serializeMetric(now, nv))
	}

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

package sqmetrics

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"runtime"
	"time"

	"github.com/rcrowley/go-metrics"
)

// SquareMetrics posts metrics to an HTTP/JSON bridge endpoint
type SquareMetrics struct {
	registry metrics.Registry
	url      string
	prefix   string
	hostname string
}

// NewMetrics is the entry point for this code
func NewMetrics(metricsURL, metricsPrefix string, registry metrics.Registry) *SquareMetrics {
	hostname, err := os.Hostname()
	if err != nil {
		panic(err)
	}

	metrics := &SquareMetrics{
		registry: registry,
		url:      metricsURL,
		prefix:   metricsPrefix,
		hostname: hostname,
	}

	if metricsURL != "" {
		go metrics.publishMetrics()
	}

	go metrics.collectSystemMetrics()
	return metrics
}

func (mb *SquareMetrics) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	metrics := mb.SerializeMetrics()
	raw, err := json.Marshal(metrics)
	if err != nil {
		panic(err)
	}
	w.Write(raw)
}

// Publish metrics to bridge
func (mb *SquareMetrics) publishMetrics() {
	for range time.Tick(1 * time.Second) {
		mb.postMetrics()
	}
}

// Collect memory usage metrics
func (mb *SquareMetrics) collectSystemMetrics() {
	var mem runtime.MemStats

	update := func(name string, value uint64) {
		metrics.GetOrRegisterGauge(name, mb.registry).Update(int64(value))
	}

	updateFloat := func(name string, value float64) {
		metrics.GetOrRegisterGaugeFloat64(name, mb.registry).Update(value)
	}

	sample := metrics.NewExpDecaySample(1028, 0.015)
	gcHistogram := metrics.GetOrRegisterHistogram("runtime.mem.gc.duration", mb.registry, sample)

	var observedPauses uint32 = 0
	for range time.Tick(1 * time.Second) {
		runtime.ReadMemStats(&mem)

		update("runtime.mem.alloc", mem.Alloc)
		update("runtime.mem.total-alloc", mem.TotalAlloc)
		update("runtime.mem.sys", mem.Sys)
		update("runtime.mem.lookups", mem.Lookups)
		update("runtime.mem.mallocs", mem.Mallocs)
		update("runtime.mem.frees", mem.Frees)

		update("runtime.mem.heap.alloc", mem.HeapAlloc)
		update("runtime.mem.heap.sys", mem.HeapSys)
		update("runtime.mem.heap.idle", mem.HeapIdle)
		update("runtime.mem.heap.inuse", mem.HeapInuse)
		update("runtime.mem.heap.released", mem.HeapReleased)
		update("runtime.mem.heap.objects", mem.HeapObjects)

		update("runtime.mem.stack.inuse", mem.StackInuse)
		update("runtime.mem.stack.sys", mem.StackSys)
		update("runtime.mem.stack.sys", mem.StackSys)

		update("runtime.goroutines", uint64(runtime.NumGoroutine()))
		update("runtime.cgo-calls", uint64(runtime.NumCgoCall()))

		update("runtime.mem.gc.num-gc", uint64(mem.NumGC))
		updateFloat("runtime.mem.gc.cpu-fraction", mem.GCCPUFraction)

		// Update histogram of GC pauses
		for ; observedPauses < mem.NumGC; observedPauses++ {
			gcHistogram.Update(int64(mem.PauseNs[(observedPauses+1)%256]))
		}
	}
}

func (mb *SquareMetrics) postMetrics() {
	metrics := mb.SerializeMetrics()
	raw, err := json.Marshal(metrics)
	if err != nil {
		panic(err)
	}
	resp, err := http.Post(mb.url, "application/json", bytes.NewReader(raw))
	if err == nil {
		resp.Body.Close()
	}
}

func (mb *SquareMetrics) serializeMetric(now int64, metric tuple) map[string]interface{} {
	return map[string]interface{}{
		"timestamp": now,
		"metric":    fmt.Sprintf("%s.%s", mb.prefix, metric.name),
		"value":     metric.value,
		"hostname":  mb.hostname,
	}
}

type tuple struct {
	name  string
	value interface{}
}

// SerializeMetrics returns a map of the collected metrics, suitable for JSON marshalling
func (mb *SquareMetrics) SerializeMetrics() []map[string]interface{} {
	nvs := []tuple{}

	mb.registry.Each(func(name string, i interface{}) {
		switch metric := i.(type) {
		case metrics.Counter:
			nvs = append(nvs, tuple{name, metric.Count()})
		case metrics.Gauge:
			nvs = append(nvs, tuple{name, metric.Value()})
		case metrics.GaugeFloat64:
			nvs = append(nvs, tuple{name, metric.Value()})
		case metrics.Timer:
			timer := metric.Snapshot()
			nvs = append(nvs, []tuple{
				{fmt.Sprintf("%s.count", name), timer.Count()},
				{fmt.Sprintf("%s.min", name), timer.Min()},
				{fmt.Sprintf("%s.max", name), timer.Max()},
				{fmt.Sprintf("%s.mean", name), timer.Mean()},
				{fmt.Sprintf("%s.std-dev", name), timer.StdDev()},
				{fmt.Sprintf("%s.one-minute", name), timer.Rate1()},
				{fmt.Sprintf("%s.five-minute", name), timer.Rate5()},
				{fmt.Sprintf("%s.fifteen-minute", name), timer.Rate15()},
				{fmt.Sprintf("%s.mean-rate", name), timer.RateMean()},
				{fmt.Sprintf("%s.50-percentile", name), timer.Percentile(0.5)},
				{fmt.Sprintf("%s.75-percentile", name), timer.Percentile(0.75)},
				{fmt.Sprintf("%s.95-percentile", name), timer.Percentile(0.95)},
				{fmt.Sprintf("%s.99-percentile", name), timer.Percentile(0.99)},
				{fmt.Sprintf("%s.999-percentile", name), timer.Percentile(0.999)},
			}...)
		}
	})

	now := time.Now().Unix()
	out := []map[string]interface{}{}
	for _, nv := range nvs {
		out = append(out, mb.serializeMetric(now, nv))
	}

	return out
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2017 Jon Carlson.  All rights reserved.
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.
//

package main

import (
	"bytes"
	"database/sql"
	"fmt"
	"github.com/joncrlsn/misc"
	"github.com/joncrlsn/pgutil"
	"sort"
	"strconv"
	"strings"
	"text/template"
)

var (
	columnSqlTemplate = initColumnSqlTemplate()
)

// Initializes the Sql template
func initColumnSqlTemplate() *template.Template {
	sql := `
SELECT table_schema
    ,  {{if eq $.DbSchema "*" }}table_schema || '.' || {{end}}table_name || '.' ||lpad(cast (ordinal_position as varchar), 5, '0')|| column_name AS compare_name
	, table_name
    , column_name
    , data_type
    , is_nullable
    , column_default
    , character_maximum_length
    , is_identity
    , identity_generation
    , substring(udt_name from 2) AS array_type
FROM information_schema.columns
WHERE is_updatable = 'YES'
{{if eq $.DbSchema "*" }}
AND table_schema NOT LIKE 'pg_%'
AND table_schema <> 'information_schema'
{{else}}
AND table_schema = '{{$.DbSchema}}'
{{end}}
ORDER BY compare_name ASC;
`
	t := template.New("ColumnSqlTmpl")
	template.Must(t.Parse(sql))
	return t
}

var (
	tableColumnSqlTemplate = initTableColumnSqlTemplate()
)

// Initializes the Sql template
func initTableColumnSqlTemplate() *template.Template {
	sql := `
SELECT a.table_schema
    , {{if eq $.DbSchema "*" }}a.table_schema || '.' || {{end}}a.table_name || '.' || column_name  AS compare_name
	, a.table_name
    , column_name
    , data_type
    , is_nullable
    , column_default
    , character_maximum_length
FROM information_schema.columns a
INNER JOIN information_schema.tables b
    ON a.table_schema = b.table_schema AND
       a.table_name = b.table_name AND
       b.table_type = 'BASE TABLE'
WHERE is_updatable = 'YES'
{{if eq $.DbSchema "*" }}
AND a.table_schema NOT LIKE 'pg_%'
AND a.table_schema <> 'information_schema'
{{else}}
AND a.table_schema = '{{$.DbSchema}}'
{{end}}
{{ if $.TableType }}
AND b.table_type = '{{ $.TableType }}'
{{ end }}
ORDER BY compare_name ASC;
`
	t := template.New("ColumnSqlTmpl")
	template.Must(t.Parse(sql))
	return t
}

// ==================================
// Column Rows definition
// ==================================

// ColumnRows is a sortable slice of string maps
type ColumnRows []map[string]string

func (slice ColumnRows) Len() int {
	return len(slice)
}

func (slice ColumnRows) Less(i, j int) bool {
	return slice[i]["compare_name"] < slice[j]["compare_name"]
}

func (slice ColumnRows) Swap(i, j int) {
	slice[i], slice[j] = slice[j], slice[i]
}

// ==================================
// ColumnSchema definition
// (implements Schema -- defined in pgdiff.go)
// ==================================

// ColumnSchema holds a slice of rows from one of the databases as well as
// a reference to the current row of data we're viewing.
type ColumnSchema struct {
	rows   ColumnRows
	rowNum int
	done   bool
}

// get returns the value from the current row for the given key
func (c *ColumnSchema) get(key string) string {
	if c.rowNum >= len(c.rows) {
		return ""
	}
	return c.rows[c.rowNum][key]
}

// NextRow increments the rowNum and tells you whether or not there are more
func (c *ColumnSchema) NextRow() bool {
	if c.rowNum >= len(c.rows)-1 {
		c.done = true
	}
	c.rowNum = c.rowNum + 1
	return !c.done
}

// Compare tells you, in one pass, whether or not the first row matches, is less than, or greater than the second row
func (c *ColumnSchema) Compare(obj interface{}) int {
	c2, ok := obj.(*ColumnSchema)
	if !ok {
		fmt.Println("Error!!!, Compare needs a ColumnSchema instance", c2)
	}

	val := misc.CompareStrings(c.get("compare_name"), c2.get("compare_name"))
	return val
}

// Add prints SQL to add the column
func (c *ColumnSchema) Add() {

	schema := dbInfo2.DbSchema
	if schema == "*" {
		schema = c.get("table_schema")
	}

	// Knowing the version of db2 would eliminate the need for this warning
	if c.get("is_identity") == "YES" {
		fmt.Println("-- WARNING: identity columns are not supported in PostgreSQL versions < 10.")
		fmt.Println("-- Attempting to create identity columns in earlier versions will probably result in errors.")
	}

	if c.get("data_type") == "character varying" {
		maxLength, valid := getMaxLength(c.get("character_maximum_length"))
		if !valid {
			fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s character varying", schema, c.get("table_name"), c.get("column_name"))
		} else {
			fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s character varying(%s)", schema, c.get("table_name"), c.get("column_name"), maxLength)
		}
	} else {
		dataType := c.get("data_type")
		//if c.get("data_type") == "ARRAY" {
			//fmt.Println("-- Note that adding of array data types are not yet generated properly.")
		//}
		if dataType == "ARRAY" {
			dataType = c.get("array_type")+"[]"
		}
		//fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s %s", schema, c.get("table_name"), c.get("column_name"), c.get("data_type"))
		fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s %s", schema, c.get("table_name"), c.get("column_name"), dataType)
	}

	if c.get("is_nullable") == "NO" {
		fmt.Printf(" NOT NULL")
	}
	if c.get("column_default") != "null" {
		fmt.Printf(" DEFAULT %s", c.get("column_default"))
	}
	// NOTE: there are more identity column sequence options according to the PostgreSQL
	// CREATE TABLE docs, but these do not appear to be available as of version 10.1
	if c.get("is_identity") == "YES" {
		fmt.Printf(" GENERATED %s AS IDENTITY", c.get("identity_generation"))
	}
	fmt.Printf(";\n")
}

// Drop prints SQL to drop the column
func (c *ColumnSchema) Drop() {
	// if dropping column
	fmt.Printf("ALTER TABLE %s.%s DROP COLUMN IF EXISTS %s;\n", c.get("table_schema"), c.get("table_name"), c.get("column_name"))
}

// Change handles the case where the table and column match, but the details do not
func (c *ColumnSchema) Change(obj interface{}) {
	c2, ok := obj.(*ColumnSchema)
	if !ok {
		fmt.Println("Error!!!, ColumnSchema.Change(obj) needs a ColumnSchema instance", c2)
	}

	// Adjust data type for array columns
	dataType1 := c.get("data_type")
	if dataType1 == "ARRAY" {
		dataType1 = c.get("array_type")+"[]"
	}
	dataType2 := c2.get("data_type")
	if dataType2 == "ARRAY" {
		dataType2 = c2.get("array_type")+"[]"
	}

	// Detect column type change (mostly varchar length, or number size increase)
	// (integer to/from bigint is OK)
	if dataType1 == dataType2 {
		if dataType1 == "character varying" {
			max1, max1Valid := getMaxLength(c.get("character_maximum_length"))
			max2, max2Valid := getMaxLength(c2.get("character_maximum_length"))
			if !max1Valid && !max2Valid {
				// Leave them alone, they both have undefined max lengths
			} else if (max1Valid || !max2Valid) && (max1 != c2.get("character_maximum_length")) {
				//if !max1Valid {
				//    fmt.Println("-- WARNING: varchar column has no maximum length.  Setting to 1024, which may result in data loss.")
				//}
				max1Int, err1 := strconv.Atoi(max1)
				check("converting string to int", err1)
				max2Int, err2 := strconv.Atoi(max2)
				check("converting string to int", err2)
				if max1Int < max2Int {
					fmt.Println("-- WARNING: The next statement will shorten a character varying column, which may result in data loss.")
				}
				fmt.Printf("-- max1Valid: %v  max2Valid: %v \n", max1Valid, max2Valid)
				fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s TYPE character varying(%s);\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), max1)
			}
		}
	}

	// Code and test a column change from integer to bigint
	if dataType1 != dataType2 {
		fmt.Printf("-- WARNING: This type change may not work well: (%s to %s).\n", dataType2, dataType1)
		if strings.HasPrefix(dataType1, "character") {
			max1, max1Valid := getMaxLength(c.get("character_maximum_length"))
			if !max1Valid {
				fmt.Println("-- WARNING: varchar column has no maximum length.  Setting to 1024")
			}
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s TYPE %s(%s);\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), dataType1, max1)
		} else {
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s TYPE %s;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), dataType1)
		}
	}

	// Detect column default change (or added, dropped)
	if c.get("column_default") == "null" {
		if c.get("column_default") != "null" {
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s DROP DEFAULT;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
		}
	} else if c.get("column_default") != c2.get("column_default") {
		fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s SET DEFAULT %s;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), c.get("column_default"))
	}

	// Detect identity column change
	// Save result to variable instead of printing because order for adding/removing
	// is_nullable affects identity columns
	var identitySql string
	if c.get("is_identity") != c2.get("is_identity") {
		// Knowing the version of db2 would eliminate the need for this warning
		fmt.Println("-- WARNING: identity columns are not supported in PostgreSQL versions < 10.")
		fmt.Println("-- Attempting to create identity columns in earlier versions will probably result in errors.")
		if c.get("is_identity") == "YES" {
			identitySql = fmt.Sprintf("ALTER TABLE \"%s\".\"%s\" ALTER COLUMN \"%s\" ADD GENERATED %s AS IDENTITY;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), c.get("identity_generation"))
		} else {
			identitySql = fmt.Sprintf("ALTER TABLE \"%s\".\"%s\" ALTER COLUMN \"%s\" DROP IDENTITY;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
		}
	}

	// Detect not-null and nullable change
	if c.get("is_nullable") != c2.get("is_nullable") {
		if c.get("is_nullable") == "YES" {
			if identitySql != "" {
				fmt.Printf(identitySql)
			}
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s DROP NOT NULL;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
		} else {
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s SET NOT NULL;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
			if identitySql != "" {
				fmt.Printf(identitySql)
			}
		}
	} else {
		if identitySql != "" {
			fmt.Printf(identitySql)
		}
	}
}

// ==================================
// Standalone Functions
// ==================================

// compare outputs SQL to make the columns match between two databases or schemas
func compare(conn1 *sql.DB, conn2 *sql.DB, tpl *template.Template) {
	buf1 := new(bytes.Buffer)
	tpl.Execute(buf1, dbInfo1)

	buf2 := new(bytes.Buffer)
	tpl.Execute(buf2, dbInfo2)

	rowChan1, _ := pgutil.QueryStrings(conn1, buf1.String())
	rowChan2, _ := pgutil.QueryStrings(conn2, buf2.String())

	//rows1 := make([]map[string]string, 500)
	rows1 := make(ColumnRows, 0)
	for row := range rowChan1 {
		rows1 = append(rows1, row)
	}
	sort.Sort(rows1)

	//rows2 := make([]map[string]string, 500)
	rows2 := make(ColumnRows, 0)
	for row := range rowChan2 {
		rows2 = append(rows2, row)
	}
	sort.Sort(&rows2)

	// We have to explicitly type this as Schema here for some unknown reason
	var schema1 Schema = &ColumnSchema{rows: rows1, rowNum: -1}
	var schema2 Schema = &ColumnSchema{rows: rows2, rowNum: -1}

	// Compare the columns
	doDiff(schema1, schema2)

}

// compareColumns outputs SQL to make the columns match between two databases or schemas
func compareColumns(conn1 *sql.DB, conn2 *sql.DB) {

    compare(conn1, conn2, columnSqlTemplate)

}

// compareColumns outputs SQL to make the tables columns (without views columns) match between two databases or schemas
func compareTableColumns(conn1 *sql.DB, conn2 *sql.DB) {

    compare(conn1, conn2, tableColumnSqlTemplate)

}

// getMaxLength returns the maximum length and whether or not it is valid
func getMaxLength(maxLength string) (string, bool) {

	if maxLength == "null" {
		// default to 1024
		return "1024", false
	}
	return maxLength, true
}

<commit_msg>fix:  comparison for column default value<commit_after>//
// Copyright (c) 2017 Jon Carlson.  All rights reserved.
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.
//

package main

import (
	"bytes"
	"database/sql"
	"fmt"
	"github.com/joncrlsn/misc"
	"github.com/joncrlsn/pgutil"
	"sort"
	"strconv"
	"strings"
	"text/template"
)

var (
	columnSqlTemplate = initColumnSqlTemplate()
)

// Initializes the Sql template
func initColumnSqlTemplate() *template.Template {
	sql := `
SELECT table_schema
    ,  {{if eq $.DbSchema "*" }}table_schema || '.' || {{end}}table_name || '.' ||lpad(cast (ordinal_position as varchar), 5, '0')|| column_name AS compare_name
	, table_name
    , column_name
    , data_type
    , is_nullable
    , column_default
    , character_maximum_length
    , is_identity
    , identity_generation
    , substring(udt_name from 2) AS array_type
FROM information_schema.columns
WHERE is_updatable = 'YES'
{{if eq $.DbSchema "*" }}
AND table_schema NOT LIKE 'pg_%'
AND table_schema <> 'information_schema'
{{else}}
AND table_schema = '{{$.DbSchema}}'
{{end}}
ORDER BY compare_name ASC;
`
	t := template.New("ColumnSqlTmpl")
	template.Must(t.Parse(sql))
	return t
}

var (
	tableColumnSqlTemplate = initTableColumnSqlTemplate()
)

// Initializes the Sql template
func initTableColumnSqlTemplate() *template.Template {
	sql := `
SELECT a.table_schema
    , {{if eq $.DbSchema "*" }}a.table_schema || '.' || {{end}}a.table_name || '.' || column_name  AS compare_name
	, a.table_name
    , column_name
    , data_type
    , is_nullable
    , column_default
    , character_maximum_length
FROM information_schema.columns a
INNER JOIN information_schema.tables b
    ON a.table_schema = b.table_schema AND
       a.table_name = b.table_name AND
       b.table_type = 'BASE TABLE'
WHERE is_updatable = 'YES'
{{if eq $.DbSchema "*" }}
AND a.table_schema NOT LIKE 'pg_%'
AND a.table_schema <> 'information_schema'
{{else}}
AND a.table_schema = '{{$.DbSchema}}'
{{end}}
{{ if $.TableType }}
AND b.table_type = '{{ $.TableType }}'
{{ end }}
ORDER BY compare_name ASC;
`
	t := template.New("ColumnSqlTmpl")
	template.Must(t.Parse(sql))
	return t
}

// ==================================
// Column Rows definition
// ==================================

// ColumnRows is a sortable slice of string maps
type ColumnRows []map[string]string

func (slice ColumnRows) Len() int {
	return len(slice)
}

func (slice ColumnRows) Less(i, j int) bool {
	return slice[i]["compare_name"] < slice[j]["compare_name"]
}

func (slice ColumnRows) Swap(i, j int) {
	slice[i], slice[j] = slice[j], slice[i]
}

// ==================================
// ColumnSchema definition
// (implements Schema -- defined in pgdiff.go)
// ==================================

// ColumnSchema holds a slice of rows from one of the databases as well as
// a reference to the current row of data we're viewing.
type ColumnSchema struct {
	rows   ColumnRows
	rowNum int
	done   bool
}

// get returns the value from the current row for the given key
func (c *ColumnSchema) get(key string) string {
	if c.rowNum >= len(c.rows) {
		return ""
	}
	return c.rows[c.rowNum][key]
}

// NextRow increments the rowNum and tells you whether or not there are more
func (c *ColumnSchema) NextRow() bool {
	if c.rowNum >= len(c.rows)-1 {
		c.done = true
	}
	c.rowNum = c.rowNum + 1
	return !c.done
}

// Compare tells you, in one pass, whether or not the first row matches, is less than, or greater than the second row
func (c *ColumnSchema) Compare(obj interface{}) int {
	c2, ok := obj.(*ColumnSchema)
	if !ok {
		fmt.Println("Error!!!, Compare needs a ColumnSchema instance", c2)
	}

	val := misc.CompareStrings(c.get("compare_name"), c2.get("compare_name"))
	return val
}

// Add prints SQL to add the column
func (c *ColumnSchema) Add() {

	schema := dbInfo2.DbSchema
	if schema == "*" {
		schema = c.get("table_schema")
	}

	// Knowing the version of db2 would eliminate the need for this warning
	if c.get("is_identity") == "YES" {
		fmt.Println("-- WARNING: identity columns are not supported in PostgreSQL versions < 10.")
		fmt.Println("-- Attempting to create identity columns in earlier versions will probably result in errors.")
	}

	if c.get("data_type") == "character varying" {
		maxLength, valid := getMaxLength(c.get("character_maximum_length"))
		if !valid {
			fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s character varying", schema, c.get("table_name"), c.get("column_name"))
		} else {
			fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s character varying(%s)", schema, c.get("table_name"), c.get("column_name"), maxLength)
		}
	} else {
		dataType := c.get("data_type")
		//if c.get("data_type") == "ARRAY" {
			//fmt.Println("-- Note that adding of array data types are not yet generated properly.")
		//}
		if dataType == "ARRAY" {
			dataType = c.get("array_type")+"[]"
		}
		//fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s %s", schema, c.get("table_name"), c.get("column_name"), c.get("data_type"))
		fmt.Printf("ALTER TABLE %s.%s ADD COLUMN %s %s", schema, c.get("table_name"), c.get("column_name"), dataType)
	}

	if c.get("is_nullable") == "NO" {
		fmt.Printf(" NOT NULL")
	}
	if c.get("column_default") != "null" {
		fmt.Printf(" DEFAULT %s", c.get("column_default"))
	}
	// NOTE: there are more identity column sequence options according to the PostgreSQL
	// CREATE TABLE docs, but these do not appear to be available as of version 10.1
	if c.get("is_identity") == "YES" {
		fmt.Printf(" GENERATED %s AS IDENTITY", c.get("identity_generation"))
	}
	fmt.Printf(";\n")
}

// Drop prints SQL to drop the column
func (c *ColumnSchema) Drop() {
	// if dropping column
	fmt.Printf("ALTER TABLE %s.%s DROP COLUMN IF EXISTS %s;\n", c.get("table_schema"), c.get("table_name"), c.get("column_name"))
}

// Change handles the case where the table and column match, but the details do not
func (c *ColumnSchema) Change(obj interface{}) {
	c2, ok := obj.(*ColumnSchema)
	if !ok {
		fmt.Println("Error!!!, ColumnSchema.Change(obj) needs a ColumnSchema instance", c2)
	}

	// Adjust data type for array columns
	dataType1 := c.get("data_type")
	if dataType1 == "ARRAY" {
		dataType1 = c.get("array_type")+"[]"
	}
	dataType2 := c2.get("data_type")
	if dataType2 == "ARRAY" {
		dataType2 = c2.get("array_type")+"[]"
	}

	// Detect column type change (mostly varchar length, or number size increase)
	// (integer to/from bigint is OK)
	if dataType1 == dataType2 {
		if dataType1 == "character varying" {
			max1, max1Valid := getMaxLength(c.get("character_maximum_length"))
			max2, max2Valid := getMaxLength(c2.get("character_maximum_length"))
			if !max1Valid && !max2Valid {
				// Leave them alone, they both have undefined max lengths
			} else if (max1Valid || !max2Valid) && (max1 != c2.get("character_maximum_length")) {
				//if !max1Valid {
				//    fmt.Println("-- WARNING: varchar column has no maximum length.  Setting to 1024, which may result in data loss.")
				//}
				max1Int, err1 := strconv.Atoi(max1)
				check("converting string to int", err1)
				max2Int, err2 := strconv.Atoi(max2)
				check("converting string to int", err2)
				if max1Int < max2Int {
					fmt.Println("-- WARNING: The next statement will shorten a character varying column, which may result in data loss.")
				}
				fmt.Printf("-- max1Valid: %v  max2Valid: %v \n", max1Valid, max2Valid)
				fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s TYPE character varying(%s);\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), max1)
			}
		}
	}

	// Code and test a column change from integer to bigint
	if dataType1 != dataType2 {
		fmt.Printf("-- WARNING: This type change may not work well: (%s to %s).\n", dataType2, dataType1)
		if strings.HasPrefix(dataType1, "character") {
			max1, max1Valid := getMaxLength(c.get("character_maximum_length"))
			if !max1Valid {
				fmt.Println("-- WARNING: varchar column has no maximum length.  Setting to 1024")
			}
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s TYPE %s(%s);\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), dataType1, max1)
		} else {
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s TYPE %s;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), dataType1)
		}
	}

	// Detect column default change (or added, dropped)
	if c.get("column_default") == "null" {
		if c2.get("column_default") != "null" {
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s DROP DEFAULT;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
		}
	} else if c.get("column_default") != c2.get("column_default") {
		fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s SET DEFAULT %s;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), c.get("column_default"))
	}

	// Detect identity column change
	// Save result to variable instead of printing because order for adding/removing
	// is_nullable affects identity columns
	var identitySql string
	if c.get("is_identity") != c2.get("is_identity") {
		// Knowing the version of db2 would eliminate the need for this warning
		fmt.Println("-- WARNING: identity columns are not supported in PostgreSQL versions < 10.")
		fmt.Println("-- Attempting to create identity columns in earlier versions will probably result in errors.")
		if c.get("is_identity") == "YES" {
			identitySql = fmt.Sprintf("ALTER TABLE \"%s\".\"%s\" ALTER COLUMN \"%s\" ADD GENERATED %s AS IDENTITY;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"), c.get("identity_generation"))
		} else {
			identitySql = fmt.Sprintf("ALTER TABLE \"%s\".\"%s\" ALTER COLUMN \"%s\" DROP IDENTITY;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
		}
	}

	// Detect not-null and nullable change
	if c.get("is_nullable") != c2.get("is_nullable") {
		if c.get("is_nullable") == "YES" {
			if identitySql != "" {
				fmt.Printf(identitySql)
			}
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s DROP NOT NULL;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
		} else {
			fmt.Printf("ALTER TABLE %s.%s ALTER COLUMN %s SET NOT NULL;\n", c2.get("table_schema"), c.get("table_name"), c.get("column_name"))
			if identitySql != "" {
				fmt.Printf(identitySql)
			}
		}
	} else {
		if identitySql != "" {
			fmt.Printf(identitySql)
		}
	}
}

// ==================================
// Standalone Functions
// ==================================

// compare outputs SQL to make the columns match between two databases or schemas
func compare(conn1 *sql.DB, conn2 *sql.DB, tpl *template.Template) {
	buf1 := new(bytes.Buffer)
	tpl.Execute(buf1, dbInfo1)

	buf2 := new(bytes.Buffer)
	tpl.Execute(buf2, dbInfo2)

	rowChan1, _ := pgutil.QueryStrings(conn1, buf1.String())
	rowChan2, _ := pgutil.QueryStrings(conn2, buf2.String())

	//rows1 := make([]map[string]string, 500)
	rows1 := make(ColumnRows, 0)
	for row := range rowChan1 {
		rows1 = append(rows1, row)
	}
	sort.Sort(rows1)

	//rows2 := make([]map[string]string, 500)
	rows2 := make(ColumnRows, 0)
	for row := range rowChan2 {
		rows2 = append(rows2, row)
	}
	sort.Sort(&rows2)

	// We have to explicitly type this as Schema here for some unknown reason
	var schema1 Schema = &ColumnSchema{rows: rows1, rowNum: -1}
	var schema2 Schema = &ColumnSchema{rows: rows2, rowNum: -1}

	// Compare the columns
	doDiff(schema1, schema2)

}

// compareColumns outputs SQL to make the columns match between two databases or schemas
func compareColumns(conn1 *sql.DB, conn2 *sql.DB) {

    compare(conn1, conn2, columnSqlTemplate)

}

// compareColumns outputs SQL to make the tables columns (without views columns) match between two databases or schemas
func compareTableColumns(conn1 *sql.DB, conn2 *sql.DB) {

    compare(conn1, conn2, tableColumnSqlTemplate)

}

// getMaxLength returns the maximum length and whether or not it is valid
func getMaxLength(maxLength string) (string, bool) {

	if maxLength == "null" {
		// default to 1024
		return "1024", false
	}
	return maxLength, true
}

<|endoftext|>


<commit_before>package metrics

import (
	"errors"
	"fmt"
)

var Backends = map[string]Interface{}
var current Interface = nil

type Event struct {
	State      string
	Host       string
	Service    string
	Metric     int64
	Ttl        float32
	Tags       []string
	Attributes map[string]interface{}
}

func (e *Event) Set(k string, v interface{}) {
	if e.Attributes == nil {
		e.Attributes = map[string]interface{}{}
	}
	e.Attributes[k] = v
}

type Interface interface {
	Publish(e ...*Event) error
}

func Register(name string, m Interface) {
	Backends[name] = m
}

func Use(name string) error {
	if _, ok := Backends[name]; !ok {
		return errors.New("backend not found")
	}
	current = Backends[name]
	return nil
}

func Publish(e *Event) error {
	if current != nil {
		err := current.Publish(e)
		if err != nil {
			return fmt.Errorf("error publishing metric '%s': %s", e.Service, err)
		}
	}
	return nil
}
<commit_msg>metrics.SetPrefix<commit_after>package metrics

import (
	"errors"
	"fmt"
)

var Backends = map[string]Interface{}
var current Interface = nil
var prefix = ""

type Event struct {
	State      string
	Host       string
	Service    string
	Metric     int64
	Ttl        float32
	Tags       []string
	Attributes map[string]interface{}
}

func (e *Event) Set(k string, v interface{}) {
	if e.Attributes == nil {
		e.Attributes = map[string]interface{}{}
	}
	e.Attributes[k] = v
}

type Interface interface {
	Publish(e ...*Event) error
}

func Register(name string, m Interface) {
	Backends[name] = m
}

func Use(name string) error {
	if _, ok := Backends[name]; !ok {
		return errors.New("backend not found")
	}
	current = Backends[name]
	return nil
}

func Publish(e *Event) error {
	if current != nil {
		// TODO Make copy of event
		e.Service = prefix + e.Service

		err := current.Publish(e)
		if err != nil {
			return fmt.Errorf("error publishing metric '%s': %s", e.Service, err)
		}
	}
	return nil
}

func SetPrefix(pre string) {
	prefix = pre
}
<|endoftext|>


<commit_before>package mfs

import (
	"errors"
	"fmt"
	"os"
	"path"
	"sort"
	"sync"
	"time"

	context "gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"

	dag "github.com/ipfs/go-ipfs/merkledag"
	ft "github.com/ipfs/go-ipfs/unixfs"
	ufspb "github.com/ipfs/go-ipfs/unixfs/pb"
)

var ErrNotYetImplemented = errors.New("not yet implemented")
var ErrInvalidChild = errors.New("invalid child node")
var ErrDirExists = errors.New("directory already has entry by that name")

type Directory struct {
	dserv  dag.DAGService
	parent childCloser

	childDirs map[string]*Directory
	files     map[string]*File

	lock sync.Mutex
	node *dag.Node
	ctx  context.Context

	modTime time.Time

	name string
}

func NewDirectory(ctx context.Context, name string, node *dag.Node, parent childCloser, dserv dag.DAGService) *Directory {
	return &Directory{
		dserv:     dserv,
		ctx:       ctx,
		name:      name,
		node:      node,
		parent:    parent,
		childDirs: make(map[string]*Directory),
		files:     make(map[string]*File),
		modTime:   time.Now(),
	}
}

// closeChild updates the child by the given name to the dag node 'nd'
// and changes its own dag node
func (d *Directory) closeChild(name string, nd *dag.Node, sync bool) error {
	mynd, err := d.closeChildUpdate(name, nd, sync)
	if err != nil {
		return err
	}

	if sync {
		return d.parent.closeChild(d.name, mynd, true)
	}
	return nil
}

// closeChildUpdate is the portion of closeChild that needs to be locked around
func (d *Directory) closeChildUpdate(name string, nd *dag.Node, sync bool) (*dag.Node, error) {
	d.lock.Lock()
	defer d.lock.Unlock()

	err := d.updateChild(name, nd)
	if err != nil {
		return nil, err
	}

	if sync {
		return d.flushCurrentNode()
	}
	return nil, nil
}

func (d *Directory) flushCurrentNode() (*dag.Node, error) {
	_, err := d.dserv.Add(d.node)
	if err != nil {
		return nil, err
	}

	return d.node.Copy(), nil
}

func (d *Directory) updateChild(name string, nd *dag.Node) error {
	err := d.node.RemoveNodeLink(name)
	if err != nil && err != dag.ErrNotFound {
		return err
	}

	err = d.node.AddNodeLinkClean(name, nd)
	if err != nil {
		return err
	}

	d.modTime = time.Now()

	return nil
}

func (d *Directory) Type() NodeType {
	return TDir
}

// childNode returns a FSNode under this directory by the given name if it exists.
// it does *not* check the cached dirs and files
func (d *Directory) childNode(name string) (FSNode, error) {
	nd, err := d.childFromDag(name)
	if err != nil {
		return nil, err
	}

	i, err := ft.FromBytes(nd.Data)
	if err != nil {
		return nil, err
	}

	switch i.GetType() {
	case ufspb.Data_Directory:
		ndir := NewDirectory(d.ctx, name, nd, d, d.dserv)
		d.childDirs[name] = ndir
		return ndir, nil
	case ufspb.Data_File, ufspb.Data_Raw:
		nfi, err := NewFile(name, nd, d, d.dserv)
		if err != nil {
			return nil, err
		}
		d.files[name] = nfi
		return nfi, nil
	case ufspb.Data_Metadata:
		return nil, ErrNotYetImplemented
	default:
		return nil, ErrInvalidChild
	}
}

// Child returns the child of this directory by the given name
func (d *Directory) Child(name string) (FSNode, error) {
	d.lock.Lock()
	defer d.lock.Unlock()
	return d.childUnsync(name)
}

// childFromDag searches through this directories dag node for a child link
// with the given name
func (d *Directory) childFromDag(name string) (*dag.Node, error) {
	for _, lnk := range d.node.Links {
		if lnk.Name == name {
			return lnk.GetNode(d.ctx, d.dserv)
		}
	}

	return nil, os.ErrNotExist
}

// childUnsync returns the child under this directory by the given name
// without locking, useful for operations which already hold a lock
func (d *Directory) childUnsync(name string) (FSNode, error) {
	cdir, ok := d.childDirs[name]
	if ok {
		return cdir, nil
	}

	cfile, ok := d.files[name]
	if ok {
		return cfile, nil
	}

	return d.childNode(name)
}

type NodeListing struct {
	Name string
	Type int
	Size int64
	Hash string
}

func (d *Directory) ListNames() []string {
	d.lock.Lock()
	defer d.lock.Unlock()

	names := make(map[string]struct{})
	for n, _ := range d.childDirs {
		names[n] = struct{}{}
	}
	for n, _ := range d.files {
		names[n] = struct{}{}
	}

	for _, l := range d.node.Links {
		names[l.Name] = struct{}{}
	}

	var out []string
	for n, _ := range names {
		out = append(out, n)
	}
	sort.Strings(out)

	return out
}

func (d *Directory) List() ([]NodeListing, error) {
	d.lock.Lock()
	defer d.lock.Unlock()

	var out []NodeListing
	for _, l := range d.node.Links {
		child := NodeListing{}
		child.Name = l.Name

		c, err := d.childUnsync(l.Name)
		if err != nil {
			return nil, err
		}

		child.Type = int(c.Type())
		if c, ok := c.(*File); ok {
			size, err := c.Size()
			if err != nil {
				return nil, err
			}
			child.Size = size
		}
		nd, err := c.GetNode()
		if err != nil {
			return nil, err
		}

		k, err := nd.Key()
		if err != nil {
			return nil, err
		}

		child.Hash = k.B58String()

		out = append(out, child)
	}

	return out, nil
}

func (d *Directory) Mkdir(name string) (*Directory, error) {
	d.lock.Lock()
	defer d.lock.Unlock()

	fsn, err := d.childUnsync(name)
	if err == nil {
		switch fsn := fsn.(type) {
		case *Directory:
			return fsn, os.ErrExist
		case *File:
			return nil, os.ErrExist
		default:
			return nil, fmt.Errorf("unrecognized type: %#v", fsn)
		}
	}

	ndir := &dag.Node{Data: ft.FolderPBData()}

	_, err = d.dserv.Add(ndir)
	if err != nil {
		return nil, err
	}

	err = d.node.AddNodeLinkClean(name, ndir)
	if err != nil {
		return nil, err
	}

	dirobj := NewDirectory(d.ctx, name, ndir, d, d.dserv)
	d.childDirs[name] = dirobj
	return dirobj, nil
}

func (d *Directory) Unlink(name string) error {
	d.lock.Lock()
	defer d.lock.Unlock()

	delete(d.childDirs, name)
	delete(d.files, name)

	err := d.node.RemoveNodeLink(name)
	if err != nil {
		return err
	}

	_, err = d.dserv.Add(d.node)
	if err != nil {
		return err
	}

	return nil
}

func (d *Directory) Flush() error {
	d.lock.Lock()
	nd, err := d.flushCurrentNode()
	if err != nil {
		d.lock.Unlock()
		return err
	}
	d.lock.Unlock()

	return d.parent.closeChild(d.name, nd, true)
}

// AddChild adds the node 'nd' under this directory giving it the name 'name'
func (d *Directory) AddChild(name string, nd *dag.Node) error {
	d.lock.Lock()
	defer d.lock.Unlock()

	_, err := d.childUnsync(name)
	if err == nil {
		return ErrDirExists
	}

	_, err = d.dserv.Add(nd)
	if err != nil {
		return err
	}

	err = d.node.AddNodeLinkClean(name, nd)
	if err != nil {
		return err
	}

	d.modTime = time.Now()

	return nil
}

func (d *Directory) sync() error {
	for name, dir := range d.childDirs {
		nd, err := dir.GetNode()
		if err != nil {
			return err
		}

		err = d.updateChild(name, nd)
		if err != nil {
			return err
		}
	}

	for name, file := range d.files {
		nd, err := file.GetNode()
		if err != nil {
			return err
		}

		err = d.updateChild(name, nd)
		if err != nil {
			return err
		}
	}

	return nil
}

func (d *Directory) Path() string {
	cur := d
	var out string
	for cur != nil {
		out = path.Join(cur.name, out)
		cur = cur.parent.(*Directory)
	}
	return out
}

func (d *Directory) GetNode() (*dag.Node, error) {
	d.lock.Lock()
	defer d.lock.Unlock()

	err := d.sync()
	if err != nil {
		return nil, err
	}

	_, err = d.dserv.Add(d.node)
	if err != nil {
		return nil, err
	}

	return d.node.Copy(), nil
}
<commit_msg>Cache files/dirs when added.<commit_after>package mfs

import (
	"errors"
	"fmt"
	"os"
	"path"
	"sort"
	"sync"
	"time"

	context "gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"

	dag "github.com/ipfs/go-ipfs/merkledag"
	ft "github.com/ipfs/go-ipfs/unixfs"
	ufspb "github.com/ipfs/go-ipfs/unixfs/pb"
)

var ErrNotYetImplemented = errors.New("not yet implemented")
var ErrInvalidChild = errors.New("invalid child node")
var ErrDirExists = errors.New("directory already has entry by that name")

type Directory struct {
	dserv  dag.DAGService
	parent childCloser

	childDirs map[string]*Directory
	files     map[string]*File

	lock sync.Mutex
	node *dag.Node
	ctx  context.Context

	modTime time.Time

	name string
}

func NewDirectory(ctx context.Context, name string, node *dag.Node, parent childCloser, dserv dag.DAGService) *Directory {
	return &Directory{
		dserv:     dserv,
		ctx:       ctx,
		name:      name,
		node:      node,
		parent:    parent,
		childDirs: make(map[string]*Directory),
		files:     make(map[string]*File),
		modTime:   time.Now(),
	}
}

// closeChild updates the child by the given name to the dag node 'nd'
// and changes its own dag node
func (d *Directory) closeChild(name string, nd *dag.Node, sync bool) error {
	mynd, err := d.closeChildUpdate(name, nd, sync)
	if err != nil {
		return err
	}

	if sync {
		return d.parent.closeChild(d.name, mynd, true)
	}
	return nil
}

// closeChildUpdate is the portion of closeChild that needs to be locked around
func (d *Directory) closeChildUpdate(name string, nd *dag.Node, sync bool) (*dag.Node, error) {
	d.lock.Lock()
	defer d.lock.Unlock()

	err := d.updateChild(name, nd)
	if err != nil {
		return nil, err
	}

	if sync {
		return d.flushCurrentNode()
	}
	return nil, nil
}

func (d *Directory) flushCurrentNode() (*dag.Node, error) {
	_, err := d.dserv.Add(d.node)
	if err != nil {
		return nil, err
	}

	return d.node.Copy(), nil
}

func (d *Directory) updateChild(name string, nd *dag.Node) error {
	err := d.node.RemoveNodeLink(name)
	if err != nil && err != dag.ErrNotFound {
		return err
	}

	err = d.node.AddNodeLinkClean(name, nd)
	if err != nil {
		return err
	}

	d.modTime = time.Now()

	return nil
}

func (d *Directory) Type() NodeType {
	return TDir
}

// childNode returns a FSNode under this directory by the given name if it exists.
// it does *not* check the cached dirs and files
func (d *Directory) childNode(name string) (FSNode, error) {
	nd, err := d.childFromDag(name)
	if err != nil {
		return nil, err
	}

	return d.cacheNode(name, nd)
}

// cacheNode caches a node into d.childDirs or d.files and returns the FSNode.
func (d *Directory) cacheNode(name string, nd *dag.Node) (FSNode, error) {
	i, err := ft.FromBytes(nd.Data)
	if err != nil {
		return nil, err
	}

	switch i.GetType() {
	case ufspb.Data_Directory:
		ndir := NewDirectory(d.ctx, name, nd, d, d.dserv)
		d.childDirs[name] = ndir
		return ndir, nil
	case ufspb.Data_File, ufspb.Data_Raw:
		nfi, err := NewFile(name, nd, d, d.dserv)
		if err != nil {
			return nil, err
		}
		d.files[name] = nfi
		return nfi, nil
	case ufspb.Data_Metadata:
		return nil, ErrNotYetImplemented
	default:
		return nil, ErrInvalidChild
	}
}

// Child returns the child of this directory by the given name
func (d *Directory) Child(name string) (FSNode, error) {
	d.lock.Lock()
	defer d.lock.Unlock()
	return d.childUnsync(name)
}

// childFromDag searches through this directories dag node for a child link
// with the given name
func (d *Directory) childFromDag(name string) (*dag.Node, error) {
	for _, lnk := range d.node.Links {
		if lnk.Name == name {
			return lnk.GetNode(d.ctx, d.dserv)
		}
	}

	return nil, os.ErrNotExist
}

// childUnsync returns the child under this directory by the given name
// without locking, useful for operations which already hold a lock
func (d *Directory) childUnsync(name string) (FSNode, error) {
	cdir, ok := d.childDirs[name]
	if ok {
		return cdir, nil
	}

	cfile, ok := d.files[name]
	if ok {
		return cfile, nil
	}

	return d.childNode(name)
}

type NodeListing struct {
	Name string
	Type int
	Size int64
	Hash string
}

func (d *Directory) ListNames() []string {
	d.lock.Lock()
	defer d.lock.Unlock()

	names := make(map[string]struct{})
	for n, _ := range d.childDirs {
		names[n] = struct{}{}
	}
	for n, _ := range d.files {
		names[n] = struct{}{}
	}

	for _, l := range d.node.Links {
		names[l.Name] = struct{}{}
	}

	var out []string
	for n, _ := range names {
		out = append(out, n)
	}
	sort.Strings(out)

	return out
}

func (d *Directory) List() ([]NodeListing, error) {
	d.lock.Lock()
	defer d.lock.Unlock()

	var out []NodeListing
	for _, l := range d.node.Links {
		child := NodeListing{}
		child.Name = l.Name

		c, err := d.childUnsync(l.Name)
		if err != nil {
			return nil, err
		}

		child.Type = int(c.Type())
		if c, ok := c.(*File); ok {
			size, err := c.Size()
			if err != nil {
				return nil, err
			}
			child.Size = size
		}
		nd, err := c.GetNode()
		if err != nil {
			return nil, err
		}

		k, err := nd.Key()
		if err != nil {
			return nil, err
		}

		child.Hash = k.B58String()

		out = append(out, child)
	}

	return out, nil
}

func (d *Directory) Mkdir(name string) (*Directory, error) {
	d.lock.Lock()
	defer d.lock.Unlock()

	fsn, err := d.childUnsync(name)
	if err == nil {
		switch fsn := fsn.(type) {
		case *Directory:
			return fsn, os.ErrExist
		case *File:
			return nil, os.ErrExist
		default:
			return nil, fmt.Errorf("unrecognized type: %#v", fsn)
		}
	}

	ndir := &dag.Node{Data: ft.FolderPBData()}

	_, err = d.dserv.Add(ndir)
	if err != nil {
		return nil, err
	}

	err = d.node.AddNodeLinkClean(name, ndir)
	if err != nil {
		return nil, err
	}

	dirobj := NewDirectory(d.ctx, name, ndir, d, d.dserv)
	d.childDirs[name] = dirobj
	return dirobj, nil
}

func (d *Directory) Unlink(name string) error {
	d.lock.Lock()
	defer d.lock.Unlock()

	delete(d.childDirs, name)
	delete(d.files, name)

	err := d.node.RemoveNodeLink(name)
	if err != nil {
		return err
	}

	_, err = d.dserv.Add(d.node)
	if err != nil {
		return err
	}

	return nil
}

func (d *Directory) Flush() error {
	d.lock.Lock()
	nd, err := d.flushCurrentNode()
	if err != nil {
		d.lock.Unlock()
		return err
	}
	d.lock.Unlock()

	return d.parent.closeChild(d.name, nd, true)
}

// AddChild adds the node 'nd' under this directory giving it the name 'name'
func (d *Directory) AddChild(name string, nd *dag.Node) error {
	d.lock.Lock()
	defer d.lock.Unlock()

	_, err := d.childUnsync(name)
	if err == nil {
		return ErrDirExists
	}

	_, err = d.dserv.Add(nd)
	if err != nil {
		return err
	}

	err = d.node.AddNodeLinkClean(name, nd)
	if err != nil {
		return err
	}

	d.modTime = time.Now()

	if len(nd.Links) == 0 {
		nfi, err := NewFile(name, nd, d, d.dserv)
		if err != nil {
			return err
		}
		d.files[name] = nfi
	} else {
		ndir := NewDirectory(d.ctx, name, nd, d, d.dserv)
		d.childDirs[name] = ndir
	}

	return nil
}

func (d *Directory) sync() error {
	for name, dir := range d.childDirs {
		nd, err := dir.GetNode()
		if err != nil {
			return err
		}

		err = d.updateChild(name, nd)
		if err != nil {
			return err
		}
	}

	for name, file := range d.files {
		nd, err := file.GetNode()
		if err != nil {
			return err
		}

		err = d.updateChild(name, nd)
		if err != nil {
			return err
		}
	}

	return nil
}

func (d *Directory) Path() string {
	cur := d
	var out string
	for cur != nil {
		out = path.Join(cur.name, out)
		cur = cur.parent.(*Directory)
	}
	return out
}

func (d *Directory) GetNode() (*dag.Node, error) {
	d.lock.Lock()
	defer d.lock.Unlock()

	err := d.sync()
	if err != nil {
		return nil, err
	}

	_, err = d.dserv.Add(d.node)
	if err != nil {
		return nil, err
	}

	return d.node.Copy(), nil
}
<|endoftext|>


<commit_before>package main

//
// Migrate data from one redis cluster to another.
//

import (
	"flag"
	"fmt"
	"gopkg.in/redis.v3" // http://godoc.org/gopkg.in/redis.v3
	"log"
	"strings"
)

// redis server connections
var sourceCluster *redis.ClusterClient
var destinationCluster *redis.ClusterClient
var sourceHost *redis.Client
var destinationHost *redis.Client

// redis server arrays for connecting
var sourceHostsArray []string
var destinationHostsArray []string

// cluster indication flags
var sourceIsCluster = false
var destinationIsCluster = false

// connected states
var sourceClusterConnected = false
var destinationClusterConnected = false

func main() {

	//
	// Command line argument handling
	//

	// parse command line arguments
	var sourceHosts = flag.String("sourceHosts", "", "A list of source cluster host:port servers seperated by commas. EX) 127.0.0.1:6379,127.0.0.1:6380")
	var destinationHosts = flag.String("destinationHosts", "", "A list of source cluster host:port servers seperated by spaces. EX) 127.0.0.1:6379,127.0.0.1:6380")
	var getKeys = flag.Bool("getKeys", false, "Fetches and prints keys from the source host.")
	var copyData = flag.Bool("copyData", false, "Copies all keys in a specified list to the destination cluster from the source cluster.")
	var keyFilePath = flag.String("keyFile", "", "The file path which contains the list of keys to migrate.")

	// parse the args we are looking for
	flag.Parse()

	//
	// Connect to redis servers or clusters
	//

	// connect to the host if servers found
	// break source hosts comma list into an array
	var sourceHostsString = *sourceHosts
	if len(sourceHostsString) > 0 {
		//log.Println("Source hosts detected: " + sourceHostsString)
		// break source hosts string at commas into a slice
		sourceHostsArray = strings.Split(sourceHostsString, ",")
		connectSourceCluster()
	}

	// connect to the host if servers found
	// break destination hosts comma list into an array
	var destinationHostsString = *destinationHosts
	if len(destinationHostsString) > 0 {
		//log.Println("Destination hosts detected: " + destinationHostsString)
		// break source hosts string at commas into a slice
		destinationHostsArray = strings.Split(destinationHostsString, ",")
		connectDestinationCluster()
	}

	//
	// Do the right thing depending on the operations passed from cli
	//

	// Get and display a key list
	if *getKeys == true {
		// ensure we are connected
		if sourceClusterConnected != true {
			log.Fatalln("Please specify a source cluster using -sourceCluster=127.0.0.1:6379.")
		}

		//log.Println("Getting full key list...")
		// iterate through each host in the destination cluster, connect, and
		// run KEYS *
		var allKeys = getSourceKeys()

		// loop through all keys and print them plainly one per line
		for i := 0; i < len(allKeys); i++ {
			fmt.Println(allKeys[i])
		}
	}

	// Copy all or some keys to the new server/cluster
	if *copyData == true {

		// ensure we are connected
		if sourceClusterConnected != true {
			log.Fatalln("Please specify a source cluster using -sourceCluster=127.0.0.1:6379.")
		}
		if destinationClusterConnected != true {
			log.Fatalln("Please specify a destination cluster using -destinationCluster=127.0.0.1:6379")
		}

		// check if a keyfile was specified
		var keyFile = *keyFilePath

		// if the key file path was set, open the file and read all the keys
		// into an array

		// loop through the array of key strings
		// read key from source cluster
		// write key to destination cluster

		// load the list of keys from the keyfile
		log.Println("Destination hosts: " + destinationHostsString + keyFile)
	}
}

// ping testing functions
func clusterPingTest(redisClient *redis.ClusterClient) {
	var pingTest = redisClient.Ping()
	var pingMessage, pingError = pingTest.Result()
	if pingError != nil {
		log.Fatalln("Error when pinging a Redis connection:" + pingMessage)
	}
}
func hostPingTest(redisClient *redis.Client) {
	var pingTest = redisClient.Ping()
	var pingMessage, pingError = pingTest.Result()
	if pingError != nil {
		log.Fatalln("Error when pinging a Redis connection:" + pingMessage)
	}
}

// Connects to the source host/cluster
func connectSourceCluster() {

	// connect to source cluster and ping it
	if len(sourceHostsArray) == 1 {
		sourceHost = redis.NewClient(&redis.Options{
			Addr: sourceHostsArray[0],
		})
		sourceIsCluster = false
		//log.Println("Source is a single host.")
		hostPingTest(sourceHost)
	} else {
		sourceCluster = redis.NewClusterClient(&redis.ClusterOptions{
			Addrs: sourceHostsArray,
		})
		sourceIsCluster = true
		//log.Println("Source is a cluster.")
		clusterPingTest(sourceCluster)
	}
	sourceClusterConnected = true
	//log.Println("Source connected")
}

// Connects to the destination host/cluster
func connectDestinationCluster() {

	// connect to destination cluster and ping it
	if len(destinationHostsArray) == 1 {
		destinationHost = redis.NewClient(&redis.Options{
			Addr: destinationHostsArray[0],
		})
		destinationIsCluster = false
		//log.Println("Destination is a single host.")
		hostPingTest(destinationHost)
	} else {
		destinationCluster = redis.NewClusterClient(&redis.ClusterOptions{
			Addrs: destinationHostsArray,
		})
		destinationIsCluster = true
		//log.Println("Destination is a cluster.")
		clusterPingTest(destinationCluster)
	}
	destinationClusterConnected = true
	//log.Println("Destination connected.")
}

// Gets all the keys from the source server/cluster
func getSourceKeys() []string {

	var allKeys *redis.StringSliceCmd
	if destinationIsCluster == true {
		allKeys = sourceCluster.Keys("*")
	} else {
		allKeys = sourceHost.Keys("*")
	}

	return allKeys.Val()
}
<commit_msg>fully working key migration with help<commit_after>package main

//
// Migrate data from one redis cluster to another.
//

import (
	"bufio"
	"flag"
	"fmt"
	"gopkg.in/redis.v3" // http://godoc.org/gopkg.in/redis.v3
	"log"
	"os"
	"strconv"
	"strings"
	"time"
)

// redis server connections
var sourceCluster *redis.ClusterClient
var destinationCluster *redis.ClusterClient
var sourceHost *redis.Client
var destinationHost *redis.Client

// redis server arrays for connecting
var sourceHostsArray []string
var destinationHostsArray []string

// cluster indication flags
var sourceIsCluster = false
var destinationIsCluster = false

// connected states
var sourceClusterConnected = false
var destinationClusterConnected = false

// counter for number of keys migrated
var keysMigrated int64 = 0

func main() {

	//
	// Command line argument handling
	//

	// parse command line arguments
	var sourceHosts = flag.String("sourceHosts", "", "A list of source cluster host:port servers seperated by commas. EX) 127.0.0.1:6379,127.0.0.1:6380")
	var destinationHosts = flag.String("destinationHosts", "", "A list of source cluster host:port servers seperated by spaces. EX) 127.0.0.1:6379,127.0.0.1:6380")
	var getKeys = flag.Bool("getKeys", false, "Fetches and prints keys from the source host.")
	var copyData = flag.Bool("copyData", false, "Copies all keys in a specified list to the destination cluster from the source cluster.")
	var keyFilePath = flag.String("keyFile", "", "The file path which contains the list of keys to migrate.")

	// parse the args we are looking for
	flag.Parse()

	// Ensure a valid operation was passed
	if *getKeys == false && *copyData == false {
		showHelp()
	}

	//
	// Connect to redis servers or clusters
	//

	// connect to the host if servers found
	// break source hosts comma list into an array
	var sourceHostsString = *sourceHosts
	if len(sourceHostsString) > 0 {
		//log.Println("Source hosts detected: " + sourceHostsString)
		// break source hosts string at commas into a slice
		sourceHostsArray = strings.Split(sourceHostsString, ",")
		connectSourceCluster()
	}

	// connect to the host if servers found
	// break destination hosts comma list into an array
	var destinationHostsString = *destinationHosts
	if len(destinationHostsString) > 0 {
		//log.Println("Destination hosts detected: " + destinationHostsString)
		// break source hosts string at commas into a slice
		destinationHostsArray = strings.Split(destinationHostsString, ",")
		connectDestinationCluster()
	}

	//
	// Do the right thing depending on the operations passed from cli
	//

	// Get and display a key list
	if *getKeys == true {

		// ensure we are connected
		if sourceClusterConnected != true {
			log.Fatalln("Please specify a source cluster using -sourceCluster=127.0.0.1:6379.")
		}

		//log.Println("Getting full key list...")
		// iterate through each host in the destination cluster, connect, and
		// run KEYS *
		var allKeys = getSourceKeys()

		// see how many keys we fetched
		if len(allKeys) > 0 {
			// loop through all keys and print them plainly one per line
			for i := 0; i < len(allKeys); i++ {
				fmt.Println(allKeys[i])
			}
		} else {
			fmt.Println("No keys found in source cluster.")
		}
	}

	// Copy all or some keys to the new server/cluster
	if *copyData == true {

		// ensure we are connected
		if sourceClusterConnected != true {
			log.Fatalln("Please specify a source cluster using -sourceCluster=127.0.0.1:6379.")
		}
		if destinationClusterConnected != true {
			log.Fatalln("Please specify a destination cluster using -destinationCluster=127.0.0.1:6379")
		}

		// if the key file path was set, open the file
		if len(*keyFilePath) > 0 {
			var keyFile, err = os.Open(*keyFilePath)
			if err != nil {
				log.Fatalln("Unable to open key file specified.")
			}
			// create a new scanner for parsing the io reader returned by the
			// os.Open call earlier
			var keyFileScanner = bufio.NewScanner(keyFile)

			// read the entire key file
			for keyFileScanner.Scan() {

				// fetch the text for this line
				var key = keyFileScanner.Text()

				// migrate the key from source to destination
				migrateKey(key)

			}
		} else {
			// This is what we do if no key file was specified

			var allKeys = getSourceKeys()

			// loop through all keys and print them plainly one per line
			for i := 0; i < len(allKeys); i++ {
				var key = allKeys[i]
				migrateKey(key)
			}

		}

	}

	// Finish up with some stats
	fmt.Println("Migrated " + strconv.FormatInt(keysMigrated, 10) + " keys.")
}

// ping testing functions
func clusterPingTest(redisClient *redis.ClusterClient) {
	var pingTest = redisClient.Ping()
	var pingMessage, pingError = pingTest.Result()
	if pingError != nil {
		log.Fatalln("Error when pinging a Redis connection:" + pingMessage)
	}
}
func hostPingTest(redisClient *redis.Client) {
	var pingTest = redisClient.Ping()
	var pingMessage, pingError = pingTest.Result()
	if pingError != nil {
		log.Fatalln("Error when pinging a Redis connection:" + pingMessage)
	}
}

// Connects to the source host/cluster
func connectSourceCluster() {

	// connect to source cluster and ping it
	if len(sourceHostsArray) == 1 {
		sourceHost = redis.NewClient(&redis.Options{
			Addr: sourceHostsArray[0],
		})
		sourceIsCluster = false
		//log.Println("Source is a single host.")
		hostPingTest(sourceHost)
	} else {
		sourceCluster = redis.NewClusterClient(&redis.ClusterOptions{
			Addrs: sourceHostsArray,
		})
		sourceIsCluster = true
		//log.Println("Source is a cluster.")
		clusterPingTest(sourceCluster)
	}
	sourceClusterConnected = true
	//log.Println("Source connected")
}

// Connects to the destination host/cluster
func connectDestinationCluster() {

	// connect to destination cluster and ping it
	if len(destinationHostsArray) == 1 {
		destinationHost = redis.NewClient(&redis.Options{
			Addr: destinationHostsArray[0],
		})
		destinationIsCluster = false
		//log.Println("Destination is a single host.")
		hostPingTest(destinationHost)
	} else {
		destinationCluster = redis.NewClusterClient(&redis.ClusterOptions{
			Addrs: destinationHostsArray,
		})
		destinationIsCluster = true
		//log.Println("Destination is a cluster.")
		clusterPingTest(destinationCluster)
	}
	destinationClusterConnected = true
	//log.Println("Destination connected.")
}

// Gets all the keys from the source server/cluster
func getSourceKeys() []string {

	var allKeys *redis.StringSliceCmd
	if destinationIsCluster == true {
		allKeys = sourceCluster.Keys("*")
	} else {
		allKeys = sourceHost.Keys("*")
	}

	return allKeys.Val()
}

// Migrates a key from the source cluster to the deestination one
func migrateKey(key string) {

	//log.Println("migrating key:" + key)

	keysMigrated = keysMigrated + 1

	// init a value to hold the key data
	var data string

	// init a value to hold the key ttl
	var ttl time.Duration

	// get the key from the source
	if sourceIsCluster == true {
		data = sourceCluster.Get(key).Val()
		ttl = sourceCluster.PTTL(key).Val()

	} else {
		data = sourceHost.Get(key).Val()
		ttl = sourceHost.PTTL(key).Val()
	}

	// put the key in the destination cluster and set the ttl
	if destinationIsCluster == true {
		destinationCluster.Set(key, data, ttl)
	} else {
		destinationHost.Set(key, data, ttl)
	}

	return
}

// Displays the help content
func showHelp() {
	fmt.Println(`
- Redis Key Migrator -
https://github.com/integrii/go-redis-migrator

Migrates all or some of the keys from a Redis host or cluster to a specified host or cluster.  Supports migrating TTLs.
go-redis-migrator can also be used to list the keys in a cluster.  Run with the -getKey=true and -sourceHosts= flags to do so.

You must run this program with an operation before it will do anything.

Flags:
  -getKeys=false: Fetches and prints keys from the source host.
  -copyData=false: Copies all keys in a list specified by -keyFile= to the destination cluster from the source cluster.
  -keyFile="": The file path which contains the list of keys to migrate.  One per line.
  -destinationHosts="": A list of source cluster host:port servers seperated by spaces. Use a single host without a ',' if there is no cluster. EX) 127.0.0.1:6379,127.0.0.1:6380
  -sourceHosts="": A list of source cluster host:port servers seperated by commas. Use a single host without a ',' if there is no cluster. EX) 127.0.0.1:6379,127.0.0.1:6380
	`)

	os.Exit(0)
}
<|endoftext|>