<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 planbuilder

import (
	"fmt"
	"strconv"

	"vitess.io/vitess/go/mysql/collations"

	"vitess.io/vitess/go/sqltypes"

	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/vterrors"

	"vitess.io/vitess/go/vt/vtgate/semantics"

	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vtgate/engine"
)

var _ logicalPlan = (*orderedAggregate)(nil)

// orderedAggregate is the logicalPlan for engine.OrderedAggregate.
// This gets built if there are aggregations on a SelectScatter
// route. The primitive requests the underlying route to order
// the results by the grouping columns. This will allow the
// engine code to aggregate the results as they come.
// For example: 'select col1, col2, count(*) from t group by col1, col2'
// will be sent to the scatter route as:
// 'select col1, col2, count(*) from t group by col1, col2 order by col1, col2`
// The orderAggregate primitive built for this will be:
//    &engine.OrderedAggregate {
//      // Aggregates has one column. It computes the count
//      // using column 2 of the underlying route.
//      Aggregates: []AggregateParams{{
//        Opcode: AggregateCount,
//        Col: 2,
//      }},
//
//      // Keys has the two group by values for col1 and col2.
//      // The column numbers are from the underlying route.
//      // These values will be used to perform the grouping
//      // of the ordered results as they come from the underlying
//      // route.
//      Keys: []int{0, 1},
//      Input: (Scatter Route with the order by request),
//    }
type orderedAggregate struct {
	resultsBuilder
	extraDistinct *sqlparser.ColName

	// preProcess is true if one of the aggregates needs preprocessing.
	preProcess bool

	// aggregates specifies the aggregation parameters for each
	// aggregation function: function opcode and input column number.
	aggregates []*engine.AggregateParams

	// groupByKeys specifies the input values that must be used for
	// the aggregation key.
	groupByKeys []*engine.GroupByParams

	truncateColumnCount int
}

// checkAggregates analyzes the select expression for aggregates. If it determines
// that a primitive is needed to handle the aggregation, it builds an orderedAggregate
// primitive and returns it. It returns a groupByHandler if there is aggregation it
// can handle.
func (pb *primitiveBuilder) checkAggregates(sel *sqlparser.Select) error {
	rb, isRoute := pb.plan.(*route)
	if isRoute && rb.isSingleShard() {
		// since we can push down all of the aggregation to the route,
		// we don't need to do anything else here
		return nil
	}

	// Check if we can allow aggregates.
	hasAggregates := sqlparser.ContainsAggregation(sel.SelectExprs) || len(sel.GroupBy) > 0
	if !hasAggregates && !sel.Distinct {
		return nil
	}

	// The query has aggregates. We can proceed only
	// if the underlying primitive is a route because
	// we need the ability to push down group by and
	// order by clauses.
	if !isRoute {
		if hasAggregates {
			return vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: cross-shard query with aggregates")
		}
		pb.plan = newDistinct(pb.plan, nil)
		return nil
	}

	// If there is a distinct clause, we can check the select list
	// to see if it has a unique vindex reference. For example,
	// if the query was 'select distinct id, col from t' (with id
	// as a unique vindex), then the distinct operation can be
	// safely pushed down because the unique vindex guarantees
	// that each id can only be in a single shard. Without the
	// unique vindex property, the id could come from multiple
	// shards, which will require us to perform the grouping
	// at the vtgate level.
	if sel.Distinct {
		for _, selectExpr := range sel.SelectExprs {
			switch selectExpr := selectExpr.(type) {
			case *sqlparser.AliasedExpr:
				vindex := pb.st.Vindex(selectExpr.Expr, rb)
				if vindex != nil && vindex.IsUnique() {
					return nil
				}
			}
		}
	}

	// The group by clause could also reference a unique vindex. The above
	// example could itself have been written as
	// 'select id, col from t group by id, col', or a query could be like
	// 'select id, count(*) from t group by id'. In the above cases,
	// the grouping can be done at the shard level, which allows the entire query
	// to be pushed down. In order to perform this analysis, we're going to look
	// ahead at the group by clause to see if it references a unique vindex.
	if pb.groupByHasUniqueVindex(sel, rb) {
		return nil
	}

	// We need an aggregator primitive.
	oa := &orderedAggregate{}
	oa.resultsBuilder = newResultsBuilder(rb, oa)
	pb.plan = oa
	pb.plan.Reorder(0)
	return nil
}

// groupbyHasUniqueVindex looks ahead at the group by expression to see if
// it references a unique vindex.
//
// The vitess group by rules are different from MySQL because it's not possible
// to match the MySQL behavior without knowing the schema. For example:
// 'select id as val from t group by val' will have different interpretations
// under MySQL depending on whether t has a val column or not.
// In vitess, we always assume that 'val' references 'id'. This is achieved
// by the symbol table resolving against the select list before searching
// the tables.
//
// In order to look ahead, we have to overcome the chicken-and-egg problem:
// group by needs the select aliases to be built. Select aliases are built
// on push-down. But push-down decision depends on whether group by expressions
// reference a vindex.
// To overcome this, the look-ahead has to perform a search that matches
// the group by analyzer. The flow is similar to oa.PushGroupBy, except that
// we don't search the ResultColumns because they're not created yet. Also,
// error conditions are treated as no match for simplicity; They will be
// subsequently caught downstream.
func (pb *primitiveBuilder) groupByHasUniqueVindex(sel *sqlparser.Select, rb *route) bool {
	for _, expr := range sel.GroupBy {
		var matchedExpr sqlparser.Expr
		switch node := expr.(type) {
		case *sqlparser.ColName:
			if expr := findAlias(node, sel.SelectExprs); expr != nil {
				matchedExpr = expr
			} else {
				matchedExpr = node
			}
		case *sqlparser.Literal:
			if node.Type != sqlparser.IntVal {
				continue
			}
			num, err := strconv.ParseInt(string(node.Val), 0, 64)
			if err != nil {
				continue
			}
			if num < 1 || num > int64(len(sel.SelectExprs)) {
				continue
			}
			expr, ok := sel.SelectExprs[num-1].(*sqlparser.AliasedExpr)
			if !ok {
				continue
			}
			matchedExpr = expr.Expr
		default:
			continue
		}
		vindex := pb.st.Vindex(matchedExpr, rb)
		if vindex != nil && vindex.IsUnique() {
			return true
		}
	}
	return false
}

func findAlias(colname *sqlparser.ColName, selects sqlparser.SelectExprs) sqlparser.Expr {
	// Qualified column names cannot match an (unqualified) alias.
	if !colname.Qualifier.IsEmpty() {
		return nil
	}
	// See if this references an alias.
	for _, selectExpr := range selects {
		selectExpr, ok := selectExpr.(*sqlparser.AliasedExpr)
		if !ok {
			continue
		}
		if colname.Name.Equal(selectExpr.As) {
			return selectExpr.Expr
		}
	}
	return nil
}

// Primitive implements the logicalPlan interface
func (oa *orderedAggregate) Primitive() engine.Primitive {
	colls := map[int]collations.ID{}
	for _, key := range oa.aggregates {
		if key.CollationID != collations.Unknown {
			colls[key.KeyCol] = key.CollationID
		}
	}
	for _, key := range oa.groupByKeys {
		if key.CollationID != collations.Unknown {
			colls[key.KeyCol] = key.CollationID
		}
	}

	input := oa.input.Primitive()
	if len(oa.groupByKeys) == 0 {
		return &engine.ScalarAggregate{
			PreProcess:          oa.preProcess,
			Aggregates:          oa.aggregates,
			TruncateColumnCount: oa.truncateColumnCount,
			Collations:          colls,
			Input:               input,
		}
	}

	return &engine.OrderedAggregate{
		PreProcess:          oa.preProcess,
		Aggregates:          oa.aggregates,
		GroupByKeys:         oa.groupByKeys,
		TruncateColumnCount: oa.truncateColumnCount,
		Input:               input,
	}
}

func (oa *orderedAggregate) pushAggr(pb *primitiveBuilder, expr *sqlparser.AliasedExpr, origin logicalPlan) (rc *resultColumn, colNumber int, err error) {
	funcExpr := expr.Expr.(*sqlparser.FuncExpr)
	opcode := engine.SupportedAggregates[funcExpr.Name.Lowered()]
	if len(funcExpr.Exprs) != 1 {
		return nil, 0, fmt.Errorf("unsupported: only one expression allowed inside aggregates: %s", sqlparser.String(funcExpr))
	}
	handleDistinct, innerAliased, err := oa.needDistinctHandling(pb, funcExpr, opcode)
	if err != nil {
		return nil, 0, err
	}
	if handleDistinct {
		if oa.extraDistinct != nil {
			return nil, 0, fmt.Errorf("unsupported: only one distinct aggregation allowed in a select: %s", sqlparser.String(funcExpr))
		}
		// Push the expression that's inside the aggregate.
		// The column will eventually get added to the group by and order by clauses.
		newBuilder, _, innerCol, err := planProjection(pb, oa.input, innerAliased, origin)
		if err != nil {
			return nil, 0, err
		}
		pb.plan = newBuilder
		col, err := BuildColName(oa.input.ResultColumns(), innerCol)
		if err != nil {
			return nil, 0, err
		}
		oa.extraDistinct = col
		oa.preProcess = true
		var alias string
		if expr.As.IsEmpty() {
			alias = sqlparser.String(expr.Expr)
		} else {
			alias = expr.As.String()
		}
		switch opcode {
		case engine.AggregateCount:
			opcode = engine.AggregateCountDistinct
		case engine.AggregateSum:
			opcode = engine.AggregateSumDistinct
		}
		oa.aggregates = append(oa.aggregates, &engine.AggregateParams{
			Opcode: opcode,
			Col:    innerCol,
			Alias:  alias,
		})
	} else {
		newBuilder, _, innerCol, err := planProjection(pb, oa.input, expr, origin)
		if err != nil {
			return nil, 0, err
		}
		pb.plan = newBuilder
		oa.aggregates = append(oa.aggregates, &engine.AggregateParams{
			Opcode: opcode,
			Col:    innerCol,
		})
	}

	// Build a new rc with oa as origin because it's semantically different
	// from the expression we pushed down.
	rc = newResultColumn(expr, oa)
	oa.resultColumns = append(oa.resultColumns, rc)
	return rc, len(oa.resultColumns) - 1, nil
}

// needDistinctHandling returns true if oa needs to handle the distinct clause.
// If true, it will also return the aliased expression that needs to be pushed
// down into the underlying route.
func (oa *orderedAggregate) needDistinctHandling(pb *primitiveBuilder, funcExpr *sqlparser.FuncExpr, opcode engine.AggregateOpcode) (bool, *sqlparser.AliasedExpr, error) {
	if !funcExpr.Distinct {
		return false, nil, nil
	}
	if opcode != engine.AggregateCount && opcode != engine.AggregateSum {
		return false, nil, nil
	}
	innerAliased, ok := funcExpr.Exprs[0].(*sqlparser.AliasedExpr)
	if !ok {
		return false, nil, fmt.Errorf("syntax error: %s", sqlparser.String(funcExpr))
	}
	rb, ok := oa.input.(*route)
	if !ok {
		// Unreachable
		return true, innerAliased, nil
	}
	vindex := pb.st.Vindex(innerAliased.Expr, rb)
	if vindex != nil && vindex.IsUnique() {
		return false, nil, nil
	}
	return true, innerAliased, nil
}

// Wireup implements the logicalPlan interface
// If text columns are detected in the keys, then the function modifies
// the primitive to pull a corresponding weight_string from mysql and
// compare those instead. This is because we currently don't have the
// ability to mimic mysql's collation behavior.
func (oa *orderedAggregate) Wireup(plan logicalPlan, jt *jointab) error {
	for i, gbk := range oa.groupByKeys {
		rc := oa.resultColumns[gbk.KeyCol]
		if sqltypes.IsText(rc.column.typ) {
			weightcolNumber, err := oa.input.SupplyWeightString(gbk.KeyCol, gbk.FromGroupBy)
			if err != nil {
				_, isUnsupportedErr := err.(UnsupportedSupplyWeightString)
				if isUnsupportedErr {
					continue
				}
				return err
			}
			oa.weightStrings[rc] = weightcolNumber
			oa.groupByKeys[i].WeightStringCol = weightcolNumber
			oa.groupByKeys[i].KeyCol = weightcolNumber
			oa.truncateColumnCount = len(oa.resultColumns)
		}
	}
	return oa.input.Wireup(plan, jt)
}

func (oa *orderedAggregate) WireupGen4(semTable *semantics.SemTable) error {
	return oa.input.WireupGen4(semTable)
}

// OutputColumns implements the logicalPlan interface
func (oa *orderedAggregate) OutputColumns() []sqlparser.SelectExpr {
	outputCols := sqlparser.CloneSelectExprs(oa.input.OutputColumns())
	for _, aggr := range oa.aggregates {
		outputCols[aggr.Col] = &sqlparser.AliasedExpr{Expr: aggr.Expr, As: sqlparser.NewColIdent(aggr.Alias)}
	}
	if oa.truncateColumnCount > 0 {
		return outputCols[:oa.truncateColumnCount]
	}
	return outputCols
}

// SetTruncateColumnCount sets the truncate column count.
func (oa *orderedAggregate) SetTruncateColumnCount(count int) {
	oa.truncateColumnCount = count
}
<commit_msg>fix: add collations to non-scalar aggregates<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 planbuilder

import (
	"fmt"
	"strconv"

	"vitess.io/vitess/go/mysql/collations"

	"vitess.io/vitess/go/sqltypes"

	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/vterrors"

	"vitess.io/vitess/go/vt/vtgate/semantics"

	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vtgate/engine"
)

var _ logicalPlan = (*orderedAggregate)(nil)

// orderedAggregate is the logicalPlan for engine.OrderedAggregate.
// This gets built if there are aggregations on a SelectScatter
// route. The primitive requests the underlying route to order
// the results by the grouping columns. This will allow the
// engine code to aggregate the results as they come.
// For example: 'select col1, col2, count(*) from t group by col1, col2'
// will be sent to the scatter route as:
// 'select col1, col2, count(*) from t group by col1, col2 order by col1, col2`
// The orderAggregate primitive built for this will be:
//    &engine.OrderedAggregate {
//      // Aggregates has one column. It computes the count
//      // using column 2 of the underlying route.
//      Aggregates: []AggregateParams{{
//        Opcode: AggregateCount,
//        Col: 2,
//      }},
//
//      // Keys has the two group by values for col1 and col2.
//      // The column numbers are from the underlying route.
//      // These values will be used to perform the grouping
//      // of the ordered results as they come from the underlying
//      // route.
//      Keys: []int{0, 1},
//      Input: (Scatter Route with the order by request),
//    }
type orderedAggregate struct {
	resultsBuilder
	extraDistinct *sqlparser.ColName

	// preProcess is true if one of the aggregates needs preprocessing.
	preProcess bool

	// aggregates specifies the aggregation parameters for each
	// aggregation function: function opcode and input column number.
	aggregates []*engine.AggregateParams

	// groupByKeys specifies the input values that must be used for
	// the aggregation key.
	groupByKeys []*engine.GroupByParams

	truncateColumnCount int
}

// checkAggregates analyzes the select expression for aggregates. If it determines
// that a primitive is needed to handle the aggregation, it builds an orderedAggregate
// primitive and returns it. It returns a groupByHandler if there is aggregation it
// can handle.
func (pb *primitiveBuilder) checkAggregates(sel *sqlparser.Select) error {
	rb, isRoute := pb.plan.(*route)
	if isRoute && rb.isSingleShard() {
		// since we can push down all of the aggregation to the route,
		// we don't need to do anything else here
		return nil
	}

	// Check if we can allow aggregates.
	hasAggregates := sqlparser.ContainsAggregation(sel.SelectExprs) || len(sel.GroupBy) > 0
	if !hasAggregates && !sel.Distinct {
		return nil
	}

	// The query has aggregates. We can proceed only
	// if the underlying primitive is a route because
	// we need the ability to push down group by and
	// order by clauses.
	if !isRoute {
		if hasAggregates {
			return vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: cross-shard query with aggregates")
		}
		pb.plan = newDistinct(pb.plan, nil)
		return nil
	}

	// If there is a distinct clause, we can check the select list
	// to see if it has a unique vindex reference. For example,
	// if the query was 'select distinct id, col from t' (with id
	// as a unique vindex), then the distinct operation can be
	// safely pushed down because the unique vindex guarantees
	// that each id can only be in a single shard. Without the
	// unique vindex property, the id could come from multiple
	// shards, which will require us to perform the grouping
	// at the vtgate level.
	if sel.Distinct {
		for _, selectExpr := range sel.SelectExprs {
			switch selectExpr := selectExpr.(type) {
			case *sqlparser.AliasedExpr:
				vindex := pb.st.Vindex(selectExpr.Expr, rb)
				if vindex != nil && vindex.IsUnique() {
					return nil
				}
			}
		}
	}

	// The group by clause could also reference a unique vindex. The above
	// example could itself have been written as
	// 'select id, col from t group by id, col', or a query could be like
	// 'select id, count(*) from t group by id'. In the above cases,
	// the grouping can be done at the shard level, which allows the entire query
	// to be pushed down. In order to perform this analysis, we're going to look
	// ahead at the group by clause to see if it references a unique vindex.
	if pb.groupByHasUniqueVindex(sel, rb) {
		return nil
	}

	// We need an aggregator primitive.
	oa := &orderedAggregate{}
	oa.resultsBuilder = newResultsBuilder(rb, oa)
	pb.plan = oa
	pb.plan.Reorder(0)
	return nil
}

// groupbyHasUniqueVindex looks ahead at the group by expression to see if
// it references a unique vindex.
//
// The vitess group by rules are different from MySQL because it's not possible
// to match the MySQL behavior without knowing the schema. For example:
// 'select id as val from t group by val' will have different interpretations
// under MySQL depending on whether t has a val column or not.
// In vitess, we always assume that 'val' references 'id'. This is achieved
// by the symbol table resolving against the select list before searching
// the tables.
//
// In order to look ahead, we have to overcome the chicken-and-egg problem:
// group by needs the select aliases to be built. Select aliases are built
// on push-down. But push-down decision depends on whether group by expressions
// reference a vindex.
// To overcome this, the look-ahead has to perform a search that matches
// the group by analyzer. The flow is similar to oa.PushGroupBy, except that
// we don't search the ResultColumns because they're not created yet. Also,
// error conditions are treated as no match for simplicity; They will be
// subsequently caught downstream.
func (pb *primitiveBuilder) groupByHasUniqueVindex(sel *sqlparser.Select, rb *route) bool {
	for _, expr := range sel.GroupBy {
		var matchedExpr sqlparser.Expr
		switch node := expr.(type) {
		case *sqlparser.ColName:
			if expr := findAlias(node, sel.SelectExprs); expr != nil {
				matchedExpr = expr
			} else {
				matchedExpr = node
			}
		case *sqlparser.Literal:
			if node.Type != sqlparser.IntVal {
				continue
			}
			num, err := strconv.ParseInt(string(node.Val), 0, 64)
			if err != nil {
				continue
			}
			if num < 1 || num > int64(len(sel.SelectExprs)) {
				continue
			}
			expr, ok := sel.SelectExprs[num-1].(*sqlparser.AliasedExpr)
			if !ok {
				continue
			}
			matchedExpr = expr.Expr
		default:
			continue
		}
		vindex := pb.st.Vindex(matchedExpr, rb)
		if vindex != nil && vindex.IsUnique() {
			return true
		}
	}
	return false
}

func findAlias(colname *sqlparser.ColName, selects sqlparser.SelectExprs) sqlparser.Expr {
	// Qualified column names cannot match an (unqualified) alias.
	if !colname.Qualifier.IsEmpty() {
		return nil
	}
	// See if this references an alias.
	for _, selectExpr := range selects {
		selectExpr, ok := selectExpr.(*sqlparser.AliasedExpr)
		if !ok {
			continue
		}
		if colname.Name.Equal(selectExpr.As) {
			return selectExpr.Expr
		}
	}
	return nil
}

// Primitive implements the logicalPlan interface
func (oa *orderedAggregate) Primitive() engine.Primitive {
	colls := map[int]collations.ID{}
	for _, key := range oa.aggregates {
		if key.CollationID != collations.Unknown {
			colls[key.KeyCol] = key.CollationID
		}
	}
	for _, key := range oa.groupByKeys {
		if key.CollationID != collations.Unknown {
			colls[key.KeyCol] = key.CollationID
		}
	}

	input := oa.input.Primitive()
	if len(oa.groupByKeys) == 0 {
		return &engine.ScalarAggregate{
			PreProcess:          oa.preProcess,
			Aggregates:          oa.aggregates,
			TruncateColumnCount: oa.truncateColumnCount,
			Collations:          colls,
			Input:               input,
		}
	}

	return &engine.OrderedAggregate{
		PreProcess:          oa.preProcess,
		Aggregates:          oa.aggregates,
		GroupByKeys:         oa.groupByKeys,
		TruncateColumnCount: oa.truncateColumnCount,
		Collations:          colls,
		Input:               input,
	}
}

func (oa *orderedAggregate) pushAggr(pb *primitiveBuilder, expr *sqlparser.AliasedExpr, origin logicalPlan) (rc *resultColumn, colNumber int, err error) {
	funcExpr := expr.Expr.(*sqlparser.FuncExpr)
	opcode := engine.SupportedAggregates[funcExpr.Name.Lowered()]
	if len(funcExpr.Exprs) != 1 {
		return nil, 0, fmt.Errorf("unsupported: only one expression allowed inside aggregates: %s", sqlparser.String(funcExpr))
	}
	handleDistinct, innerAliased, err := oa.needDistinctHandling(pb, funcExpr, opcode)
	if err != nil {
		return nil, 0, err
	}
	if handleDistinct {
		if oa.extraDistinct != nil {
			return nil, 0, fmt.Errorf("unsupported: only one distinct aggregation allowed in a select: %s", sqlparser.String(funcExpr))
		}
		// Push the expression that's inside the aggregate.
		// The column will eventually get added to the group by and order by clauses.
		newBuilder, _, innerCol, err := planProjection(pb, oa.input, innerAliased, origin)
		if err != nil {
			return nil, 0, err
		}
		pb.plan = newBuilder
		col, err := BuildColName(oa.input.ResultColumns(), innerCol)
		if err != nil {
			return nil, 0, err
		}
		oa.extraDistinct = col
		oa.preProcess = true
		var alias string
		if expr.As.IsEmpty() {
			alias = sqlparser.String(expr.Expr)
		} else {
			alias = expr.As.String()
		}
		switch opcode {
		case engine.AggregateCount:
			opcode = engine.AggregateCountDistinct
		case engine.AggregateSum:
			opcode = engine.AggregateSumDistinct
		}
		oa.aggregates = append(oa.aggregates, &engine.AggregateParams{
			Opcode: opcode,
			Col:    innerCol,
			Alias:  alias,
		})
	} else {
		newBuilder, _, innerCol, err := planProjection(pb, oa.input, expr, origin)
		if err != nil {
			return nil, 0, err
		}
		pb.plan = newBuilder
		oa.aggregates = append(oa.aggregates, &engine.AggregateParams{
			Opcode: opcode,
			Col:    innerCol,
		})
	}

	// Build a new rc with oa as origin because it's semantically different
	// from the expression we pushed down.
	rc = newResultColumn(expr, oa)
	oa.resultColumns = append(oa.resultColumns, rc)
	return rc, len(oa.resultColumns) - 1, nil
}

// needDistinctHandling returns true if oa needs to handle the distinct clause.
// If true, it will also return the aliased expression that needs to be pushed
// down into the underlying route.
func (oa *orderedAggregate) needDistinctHandling(pb *primitiveBuilder, funcExpr *sqlparser.FuncExpr, opcode engine.AggregateOpcode) (bool, *sqlparser.AliasedExpr, error) {
	if !funcExpr.Distinct {
		return false, nil, nil
	}
	if opcode != engine.AggregateCount && opcode != engine.AggregateSum {
		return false, nil, nil
	}
	innerAliased, ok := funcExpr.Exprs[0].(*sqlparser.AliasedExpr)
	if !ok {
		return false, nil, fmt.Errorf("syntax error: %s", sqlparser.String(funcExpr))
	}
	rb, ok := oa.input.(*route)
	if !ok {
		// Unreachable
		return true, innerAliased, nil
	}
	vindex := pb.st.Vindex(innerAliased.Expr, rb)
	if vindex != nil && vindex.IsUnique() {
		return false, nil, nil
	}
	return true, innerAliased, nil
}

// Wireup implements the logicalPlan interface
// If text columns are detected in the keys, then the function modifies
// the primitive to pull a corresponding weight_string from mysql and
// compare those instead. This is because we currently don't have the
// ability to mimic mysql's collation behavior.
func (oa *orderedAggregate) Wireup(plan logicalPlan, jt *jointab) error {
	for i, gbk := range oa.groupByKeys {
		rc := oa.resultColumns[gbk.KeyCol]
		if sqltypes.IsText(rc.column.typ) {
			weightcolNumber, err := oa.input.SupplyWeightString(gbk.KeyCol, gbk.FromGroupBy)
			if err != nil {
				_, isUnsupportedErr := err.(UnsupportedSupplyWeightString)
				if isUnsupportedErr {
					continue
				}
				return err
			}
			oa.weightStrings[rc] = weightcolNumber
			oa.groupByKeys[i].WeightStringCol = weightcolNumber
			oa.groupByKeys[i].KeyCol = weightcolNumber
			oa.truncateColumnCount = len(oa.resultColumns)
		}
	}
	return oa.input.Wireup(plan, jt)
}

func (oa *orderedAggregate) WireupGen4(semTable *semantics.SemTable) error {
	return oa.input.WireupGen4(semTable)
}

// OutputColumns implements the logicalPlan interface
func (oa *orderedAggregate) OutputColumns() []sqlparser.SelectExpr {
	outputCols := sqlparser.CloneSelectExprs(oa.input.OutputColumns())
	for _, aggr := range oa.aggregates {
		outputCols[aggr.Col] = &sqlparser.AliasedExpr{Expr: aggr.Expr, As: sqlparser.NewColIdent(aggr.Alias)}
	}
	if oa.truncateColumnCount > 0 {
		return outputCols[:oa.truncateColumnCount]
	}
	return outputCols
}

// SetTruncateColumnCount sets the truncate column count.
func (oa *orderedAggregate) SetTruncateColumnCount(count int) {
	oa.truncateColumnCount = count
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/kms"
)

func resourceAwsKmsKey() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsKmsKeyCreate,
		Read:   resourceAwsKmsKeyRead,
		Update: resourceAwsKmsKeyUpdate,
		Delete: resourceAwsKmsKeyDelete,

		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"key_id": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"key_usage": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
				ValidateFunc: func(v interface{}, k string) (ws []string, es []error) {
					value := v.(string)
					if !(value == "ENCRYPT_DECRYPT" || value == "") {
						es = append(es, fmt.Errorf(
							"%q must be ENCRYPT_DECRYPT or not specified", k))
					}
					return
				},
			},
			"policy": &schema.Schema{
				Type:             schema.TypeString,
				Optional:         true,
				Computed:         true,
				ValidateFunc:     validateJsonString,
				DiffSuppressFunc: suppressEquivalentAwsPolicyDiffs,
			},
			"is_enabled": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},
			"enable_key_rotation": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  false,
			},
			"deletion_window_in_days": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				ValidateFunc: func(v interface{}, k string) (ws []string, es []error) {
					value := v.(int)
					if value > 30 || value < 7 {
						es = append(es, fmt.Errorf(
							"%q must be between 7 and 30 days inclusive", k))
					}
					return
				},
			},
		},
	}
}

func resourceAwsKmsKeyCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kmsconn

	// Allow aws to chose default values if we don't pass them
	var req kms.CreateKeyInput
	if v, exists := d.GetOk("description"); exists {
		req.Description = aws.String(v.(string))
	}
	if v, exists := d.GetOk("key_usage"); exists {
		req.KeyUsage = aws.String(v.(string))
	}
	if v, exists := d.GetOk("policy"); exists {
		req.Policy = aws.String(v.(string))
	}

	resp, err := conn.CreateKey(&req)
	if err != nil {
		return err
	}

	d.SetId(*resp.KeyMetadata.KeyId)
	d.Set("key_id", resp.KeyMetadata.KeyId)

	return _resourceAwsKmsKeyUpdate(d, meta, true)
}

func resourceAwsKmsKeyRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kmsconn

	req := &kms.DescribeKeyInput{
		KeyId: aws.String(d.Id()),
	}
	resp, err := conn.DescribeKey(req)
	if err != nil {
		return err
	}
	metadata := resp.KeyMetadata

	if *metadata.KeyState == "PendingDeletion" {
		log.Printf("[WARN] Removing KMS key %s because it's already gone", d.Id())
		d.SetId("")
		return nil
	}

	d.SetId(*metadata.KeyId)

	d.Set("arn", metadata.Arn)
	d.Set("key_id", metadata.KeyId)
	d.Set("description", metadata.Description)
	d.Set("key_usage", metadata.KeyUsage)
	d.Set("is_enabled", metadata.Enabled)

	p, err := conn.GetKeyPolicy(&kms.GetKeyPolicyInput{
		KeyId:      metadata.KeyId,
		PolicyName: aws.String("default"),
	})
	if err != nil {
		return err
	}

	policy, _ := normalizeJsonString(*p.Policy)
	d.Set("policy", policy)

	krs, err := conn.GetKeyRotationStatus(&kms.GetKeyRotationStatusInput{
		KeyId: metadata.KeyId,
	})
	if err != nil {
		return err
	}
	d.Set("enable_key_rotation", krs.KeyRotationEnabled)

	return nil
}

func resourceAwsKmsKeyUpdate(d *schema.ResourceData, meta interface{}) error {
	return _resourceAwsKmsKeyUpdate(d, meta, false)
}

// We expect new keys to be enabled already
// but there is no easy way to differentiate between Update()
// called from Create() and regular update, so we have this wrapper
func _resourceAwsKmsKeyUpdate(d *schema.ResourceData, meta interface{}, isFresh bool) error {
	conn := meta.(*AWSClient).kmsconn

	if d.HasChange("is_enabled") && d.Get("is_enabled").(bool) && !isFresh {
		// Enable before any attributes will be modified
		if err := updateKmsKeyStatus(conn, d.Id(), d.Get("is_enabled").(bool)); err != nil {
			return err
		}
	}

	if d.HasChange("enable_key_rotation") {
		if err := updateKmsKeyRotationStatus(conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("description") {
		if err := resourceAwsKmsKeyDescriptionUpdate(conn, d); err != nil {
			return err
		}
	}
	if d.HasChange("policy") {
		if err := resourceAwsKmsKeyPolicyUpdate(conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("is_enabled") && !d.Get("is_enabled").(bool) {
		// Only disable when all attributes are modified
		// because we cannot modify disabled keys
		if err := updateKmsKeyStatus(conn, d.Id(), d.Get("is_enabled").(bool)); err != nil {
			return err
		}
	}

	return resourceAwsKmsKeyRead(d, meta)
}

func resourceAwsKmsKeyDescriptionUpdate(conn *kms.KMS, d *schema.ResourceData) error {
	description := d.Get("description").(string)
	keyId := d.Get("key_id").(string)

	log.Printf("[DEBUG] KMS key: %s, update description: %s", keyId, description)

	req := &kms.UpdateKeyDescriptionInput{
		Description: aws.String(description),
		KeyId:       aws.String(keyId),
	}
	_, err := conn.UpdateKeyDescription(req)
	return err
}

func resourceAwsKmsKeyPolicyUpdate(conn *kms.KMS, d *schema.ResourceData) error {
	policy, _ := normalizeJsonString(d.Get("policy").(string))
	keyId := d.Get("key_id").(string)

	log.Printf("[DEBUG] KMS key: %s, update policy: %s", keyId, policy)

	req := &kms.PutKeyPolicyInput{
		KeyId:      aws.String(keyId),
		Policy:     aws.String(policy),
		PolicyName: aws.String("default"),
	}
	_, err := conn.PutKeyPolicy(req)
	return err
}

func updateKmsKeyStatus(conn *kms.KMS, id string, shouldBeEnabled bool) error {
	var err error

	if shouldBeEnabled {
		log.Printf("[DEBUG] Enabling KMS key %q", id)
		_, err = conn.EnableKey(&kms.EnableKeyInput{
			KeyId: aws.String(id),
		})
	} else {
		log.Printf("[DEBUG] Disabling KMS key %q", id)
		_, err = conn.DisableKey(&kms.DisableKeyInput{
			KeyId: aws.String(id),
		})
	}

	if err != nil {
		return fmt.Errorf("Failed to set KMS key %q status to %t: %q",
			id, shouldBeEnabled, err.Error())
	}

	// Wait for propagation since KMS is eventually consistent
	wait := resource.StateChangeConf{
		Pending:                   []string{fmt.Sprintf("%t", !shouldBeEnabled)},
		Target:                    []string{fmt.Sprintf("%t", shouldBeEnabled)},
		Timeout:                   20 * time.Minute,
		MinTimeout:                2 * time.Second,
		ContinuousTargetOccurence: 10,
		Refresh: func() (interface{}, string, error) {
			log.Printf("[DEBUG] Checking if KMS key %s enabled status is %t",
				id, shouldBeEnabled)
			resp, err := conn.DescribeKey(&kms.DescribeKeyInput{
				KeyId: aws.String(id),
			})
			if err != nil {
				return resp, "FAILED", err
			}
			status := fmt.Sprintf("%t", *resp.KeyMetadata.Enabled)
			log.Printf("[DEBUG] KMS key %s status received: %s, retrying", id, status)

			return resp, status, nil
		},
	}

	_, err = wait.WaitForState()
	if err != nil {
		return fmt.Errorf("Failed setting KMS key status to %t: %s", shouldBeEnabled, err)
	}

	return nil
}

func updateKmsKeyRotationStatus(conn *kms.KMS, d *schema.ResourceData) error {
	var err error
	shouldEnableRotation := d.Get("enable_key_rotation").(bool)
	if shouldEnableRotation {
		log.Printf("[DEBUG] Enabling key rotation for KMS key %q", d.Id())
		_, err = conn.EnableKeyRotation(&kms.EnableKeyRotationInput{
			KeyId: aws.String(d.Id()),
		})
	} else {
		log.Printf("[DEBUG] Disabling key rotation for KMS key %q", d.Id())
		_, err = conn.DisableKeyRotation(&kms.DisableKeyRotationInput{
			KeyId: aws.String(d.Id()),
		})
	}

	if err != nil {
		return fmt.Errorf("Failed to set key rotation for %q to %t: %q",
			d.Id(), shouldEnableRotation, err.Error())
	}

	// Wait for propagation since KMS is eventually consistent
	wait := resource.StateChangeConf{
		Pending:                   []string{fmt.Sprintf("%t", !shouldEnableRotation)},
		Target:                    []string{fmt.Sprintf("%t", shouldEnableRotation)},
		Timeout:                   5 * time.Minute,
		MinTimeout:                1 * time.Second,
		ContinuousTargetOccurence: 5,
		Refresh: func() (interface{}, string, error) {
			log.Printf("[DEBUG] Checking if KMS key %s rotation status is %t",
				d.Id(), shouldEnableRotation)
			resp, err := conn.GetKeyRotationStatus(&kms.GetKeyRotationStatusInput{
				KeyId: aws.String(d.Id()),
			})
			if err != nil {
				return resp, "FAILED", err
			}
			status := fmt.Sprintf("%t", *resp.KeyRotationEnabled)
			log.Printf("[DEBUG] KMS key %s rotation status received: %s, retrying", d.Id(), status)

			return resp, status, nil
		},
	}

	_, err = wait.WaitForState()
	if err != nil {
		return fmt.Errorf("Failed setting KMS key rotation status to %t: %s", shouldEnableRotation, err)
	}

	return nil
}

func resourceAwsKmsKeyDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kmsconn
	keyId := d.Get("key_id").(string)

	req := &kms.ScheduleKeyDeletionInput{
		KeyId: aws.String(keyId),
	}
	if v, exists := d.GetOk("deletion_window_in_days"); exists {
		req.PendingWindowInDays = aws.Int64(int64(v.(int)))
	}
	_, err := conn.ScheduleKeyDeletion(req)
	if err != nil {
		return err
	}

	// Wait for propagation since KMS is eventually consistent
	wait := resource.StateChangeConf{
		Pending:                   []string{"Enabled", "Disabled"},
		Target:                    []string{"PendingDeletion"},
		Timeout:                   20 * time.Minute,
		MinTimeout:                2 * time.Second,
		ContinuousTargetOccurence: 10,
		Refresh: func() (interface{}, string, error) {
			log.Printf("[DEBUG] Checking if KMS key %s state is PendingDeletion", keyId)
			resp, err := conn.DescribeKey(&kms.DescribeKeyInput{
				KeyId: aws.String(keyId),
			})
			if err != nil {
				return resp, "Failed", err
			}

			metadata := *resp.KeyMetadata
			log.Printf("[DEBUG] KMS key %s state is %s, retrying", keyId, *metadata.KeyState)

			return resp, *metadata.KeyState, nil
		},
	}

	_, err = wait.WaitForState()
	if err != nil {
		return fmt.Errorf("Failed deactivating KMS key %s: %s", keyId, err)
	}

	log.Printf("[DEBUG] KMS Key %s deactivated.", keyId)
	d.SetId("")
	return nil
}
<commit_msg>Handle JSON parsing error in the ReadFunc for the policy document.<commit_after>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/kms"
	"github.com/hashicorp/errwrap"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

func resourceAwsKmsKey() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsKmsKeyCreate,
		Read:   resourceAwsKmsKeyRead,
		Update: resourceAwsKmsKeyUpdate,
		Delete: resourceAwsKmsKeyDelete,

		Importer: &schema.ResourceImporter{
			State: schema.ImportStatePassthrough,
		},

		Schema: map[string]*schema.Schema{
			"arn": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"key_id": &schema.Schema{
				Type:     schema.TypeString,
				Computed: true,
			},
			"description": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
			},
			"key_usage": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
				ValidateFunc: func(v interface{}, k string) (ws []string, es []error) {
					value := v.(string)
					if !(value == "ENCRYPT_DECRYPT" || value == "") {
						es = append(es, fmt.Errorf(
							"%q must be ENCRYPT_DECRYPT or not specified", k))
					}
					return
				},
			},
			"policy": &schema.Schema{
				Type:             schema.TypeString,
				Optional:         true,
				Computed:         true,
				ValidateFunc:     validateJsonString,
				DiffSuppressFunc: suppressEquivalentAwsPolicyDiffs,
			},
			"is_enabled": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},
			"enable_key_rotation": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  false,
			},
			"deletion_window_in_days": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				ValidateFunc: func(v interface{}, k string) (ws []string, es []error) {
					value := v.(int)
					if value > 30 || value < 7 {
						es = append(es, fmt.Errorf(
							"%q must be between 7 and 30 days inclusive", k))
					}
					return
				},
			},
		},
	}
}

func resourceAwsKmsKeyCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kmsconn

	// Allow aws to chose default values if we don't pass them
	var req kms.CreateKeyInput
	if v, exists := d.GetOk("description"); exists {
		req.Description = aws.String(v.(string))
	}
	if v, exists := d.GetOk("key_usage"); exists {
		req.KeyUsage = aws.String(v.(string))
	}
	if v, exists := d.GetOk("policy"); exists {
		req.Policy = aws.String(v.(string))
	}

	resp, err := conn.CreateKey(&req)
	if err != nil {
		return err
	}

	d.SetId(*resp.KeyMetadata.KeyId)
	d.Set("key_id", resp.KeyMetadata.KeyId)

	return _resourceAwsKmsKeyUpdate(d, meta, true)
}

func resourceAwsKmsKeyRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kmsconn

	req := &kms.DescribeKeyInput{
		KeyId: aws.String(d.Id()),
	}
	resp, err := conn.DescribeKey(req)
	if err != nil {
		return err
	}
	metadata := resp.KeyMetadata

	if *metadata.KeyState == "PendingDeletion" {
		log.Printf("[WARN] Removing KMS key %s because it's already gone", d.Id())
		d.SetId("")
		return nil
	}

	d.SetId(*metadata.KeyId)

	d.Set("arn", metadata.Arn)
	d.Set("key_id", metadata.KeyId)
	d.Set("description", metadata.Description)
	d.Set("key_usage", metadata.KeyUsage)
	d.Set("is_enabled", metadata.Enabled)

	p, err := conn.GetKeyPolicy(&kms.GetKeyPolicyInput{
		KeyId:      metadata.KeyId,
		PolicyName: aws.String("default"),
	})
	if err != nil {
		return err
	}

	policy, err := normalizeJsonString(*p.Policy)
	if err != nil {
		return errwrap.Wrapf("policy contains an invalid JSON: {{err}}", err)
	}
	d.Set("policy", policy)

	krs, err := conn.GetKeyRotationStatus(&kms.GetKeyRotationStatusInput{
		KeyId: metadata.KeyId,
	})
	if err != nil {
		return err
	}
	d.Set("enable_key_rotation", krs.KeyRotationEnabled)

	return nil
}

func resourceAwsKmsKeyUpdate(d *schema.ResourceData, meta interface{}) error {
	return _resourceAwsKmsKeyUpdate(d, meta, false)
}

// We expect new keys to be enabled already
// but there is no easy way to differentiate between Update()
// called from Create() and regular update, so we have this wrapper
func _resourceAwsKmsKeyUpdate(d *schema.ResourceData, meta interface{}, isFresh bool) error {
	conn := meta.(*AWSClient).kmsconn

	if d.HasChange("is_enabled") && d.Get("is_enabled").(bool) && !isFresh {
		// Enable before any attributes will be modified
		if err := updateKmsKeyStatus(conn, d.Id(), d.Get("is_enabled").(bool)); err != nil {
			return err
		}
	}

	if d.HasChange("enable_key_rotation") {
		if err := updateKmsKeyRotationStatus(conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("description") {
		if err := resourceAwsKmsKeyDescriptionUpdate(conn, d); err != nil {
			return err
		}
	}
	if d.HasChange("policy") {
		if err := resourceAwsKmsKeyPolicyUpdate(conn, d); err != nil {
			return err
		}
	}

	if d.HasChange("is_enabled") && !d.Get("is_enabled").(bool) {
		// Only disable when all attributes are modified
		// because we cannot modify disabled keys
		if err := updateKmsKeyStatus(conn, d.Id(), d.Get("is_enabled").(bool)); err != nil {
			return err
		}
	}

	return resourceAwsKmsKeyRead(d, meta)
}

func resourceAwsKmsKeyDescriptionUpdate(conn *kms.KMS, d *schema.ResourceData) error {
	description := d.Get("description").(string)
	keyId := d.Get("key_id").(string)

	log.Printf("[DEBUG] KMS key: %s, update description: %s", keyId, description)

	req := &kms.UpdateKeyDescriptionInput{
		Description: aws.String(description),
		KeyId:       aws.String(keyId),
	}
	_, err := conn.UpdateKeyDescription(req)
	return err
}

func resourceAwsKmsKeyPolicyUpdate(conn *kms.KMS, d *schema.ResourceData) error {
	policy, err := normalizeJsonString(d.Get("policy").(string))
	if err != nil {
		return errwrap.Wrapf("policy contains an invalid JSON: {{err}}", err)
	}
	keyId := d.Get("key_id").(string)

	log.Printf("[DEBUG] KMS key: %s, update policy: %s", keyId, policy)

	req := &kms.PutKeyPolicyInput{
		KeyId:      aws.String(keyId),
		Policy:     aws.String(policy),
		PolicyName: aws.String("default"),
	}
	_, err = conn.PutKeyPolicy(req)
	return err
}

func updateKmsKeyStatus(conn *kms.KMS, id string, shouldBeEnabled bool) error {
	var err error

	if shouldBeEnabled {
		log.Printf("[DEBUG] Enabling KMS key %q", id)
		_, err = conn.EnableKey(&kms.EnableKeyInput{
			KeyId: aws.String(id),
		})
	} else {
		log.Printf("[DEBUG] Disabling KMS key %q", id)
		_, err = conn.DisableKey(&kms.DisableKeyInput{
			KeyId: aws.String(id),
		})
	}

	if err != nil {
		return fmt.Errorf("Failed to set KMS key %q status to %t: %q",
			id, shouldBeEnabled, err.Error())
	}

	// Wait for propagation since KMS is eventually consistent
	wait := resource.StateChangeConf{
		Pending:                   []string{fmt.Sprintf("%t", !shouldBeEnabled)},
		Target:                    []string{fmt.Sprintf("%t", shouldBeEnabled)},
		Timeout:                   20 * time.Minute,
		MinTimeout:                2 * time.Second,
		ContinuousTargetOccurence: 10,
		Refresh: func() (interface{}, string, error) {
			log.Printf("[DEBUG] Checking if KMS key %s enabled status is %t",
				id, shouldBeEnabled)
			resp, err := conn.DescribeKey(&kms.DescribeKeyInput{
				KeyId: aws.String(id),
			})
			if err != nil {
				return resp, "FAILED", err
			}
			status := fmt.Sprintf("%t", *resp.KeyMetadata.Enabled)
			log.Printf("[DEBUG] KMS key %s status received: %s, retrying", id, status)

			return resp, status, nil
		},
	}

	_, err = wait.WaitForState()
	if err != nil {
		return fmt.Errorf("Failed setting KMS key status to %t: %s", shouldBeEnabled, err)
	}

	return nil
}

func updateKmsKeyRotationStatus(conn *kms.KMS, d *schema.ResourceData) error {
	var err error
	shouldEnableRotation := d.Get("enable_key_rotation").(bool)
	if shouldEnableRotation {
		log.Printf("[DEBUG] Enabling key rotation for KMS key %q", d.Id())
		_, err = conn.EnableKeyRotation(&kms.EnableKeyRotationInput{
			KeyId: aws.String(d.Id()),
		})
	} else {
		log.Printf("[DEBUG] Disabling key rotation for KMS key %q", d.Id())
		_, err = conn.DisableKeyRotation(&kms.DisableKeyRotationInput{
			KeyId: aws.String(d.Id()),
		})
	}

	if err != nil {
		return fmt.Errorf("Failed to set key rotation for %q to %t: %q",
			d.Id(), shouldEnableRotation, err.Error())
	}

	// Wait for propagation since KMS is eventually consistent
	wait := resource.StateChangeConf{
		Pending:                   []string{fmt.Sprintf("%t", !shouldEnableRotation)},
		Target:                    []string{fmt.Sprintf("%t", shouldEnableRotation)},
		Timeout:                   5 * time.Minute,
		MinTimeout:                1 * time.Second,
		ContinuousTargetOccurence: 5,
		Refresh: func() (interface{}, string, error) {
			log.Printf("[DEBUG] Checking if KMS key %s rotation status is %t",
				d.Id(), shouldEnableRotation)
			resp, err := conn.GetKeyRotationStatus(&kms.GetKeyRotationStatusInput{
				KeyId: aws.String(d.Id()),
			})
			if err != nil {
				return resp, "FAILED", err
			}
			status := fmt.Sprintf("%t", *resp.KeyRotationEnabled)
			log.Printf("[DEBUG] KMS key %s rotation status received: %s, retrying", d.Id(), status)

			return resp, status, nil
		},
	}

	_, err = wait.WaitForState()
	if err != nil {
		return fmt.Errorf("Failed setting KMS key rotation status to %t: %s", shouldEnableRotation, err)
	}

	return nil
}

func resourceAwsKmsKeyDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).kmsconn
	keyId := d.Get("key_id").(string)

	req := &kms.ScheduleKeyDeletionInput{
		KeyId: aws.String(keyId),
	}
	if v, exists := d.GetOk("deletion_window_in_days"); exists {
		req.PendingWindowInDays = aws.Int64(int64(v.(int)))
	}
	_, err := conn.ScheduleKeyDeletion(req)
	if err != nil {
		return err
	}

	// Wait for propagation since KMS is eventually consistent
	wait := resource.StateChangeConf{
		Pending:                   []string{"Enabled", "Disabled"},
		Target:                    []string{"PendingDeletion"},
		Timeout:                   20 * time.Minute,
		MinTimeout:                2 * time.Second,
		ContinuousTargetOccurence: 10,
		Refresh: func() (interface{}, string, error) {
			log.Printf("[DEBUG] Checking if KMS key %s state is PendingDeletion", keyId)
			resp, err := conn.DescribeKey(&kms.DescribeKeyInput{
				KeyId: aws.String(keyId),
			})
			if err != nil {
				return resp, "Failed", err
			}

			metadata := *resp.KeyMetadata
			log.Printf("[DEBUG] KMS key %s state is %s, retrying", keyId, *metadata.KeyState)

			return resp, *metadata.KeyState, nil
		},
	}

	_, err = wait.WaitForState()
	if err != nil {
		return fmt.Errorf("Failed deactivating KMS key %s: %s", keyId, err)
	}

	log.Printf("[DEBUG] KMS Key %s deactivated.", keyId)
	d.SetId("")
	return nil
}
<|endoftext|>


<commit_before>// +build !linux
package placer

import (
	"go.polydawn.net/repeatr/rio"
)

var _ rio.Placer = BindPlacer

func BindPlacer(srcPath, destPath string, writable bool, _ bool) rio.Emplacement {
	panic("BindPlacer unsupported on darwin")
}

func NewAufsPlacer(workPath string) rio.Placer {
	panic("AufsPlacer unsupported on darwin")
}
<commit_msg>less specific panics for unsupported functions<commit_after>// +build !linux
package placer

import (
	"go.polydawn.net/repeatr/rio"
)

var _ rio.Placer = BindPlacer

func BindPlacer(srcPath, destPath string, writable bool, _ bool) rio.Emplacement {
	panic("BindPlacer unsupported on this platform")
}

func NewAufsPlacer(workPath string) rio.Placer {
	panic("AufsPlacer unsupported on this platform")
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"os"
	"path/filepath"
	"regexp"

	"github.com/coreos/go-etcd/etcd"
	"github.com/prometheus/log"
)

const servicesPrefix = "/services"

var pathPat = regexp.MustCompile(`/services/([^/]+)(?:/(\d+))?`)

// TargetGroup is the target group read by Prometheus.
type TargetGroup struct {
	Targets []string          `json:"targets,omitempty"`
	Labels  map[string]string `json:"labels,omitempty"`
}

type Instances map[string]string

// services are the services stored in etcd.
type services struct {
	m   map[string]Instances // The current services.
	del []string             // Services deleted in the last update.
}

var (
	etcdServer = flag.String("server", "http://127.0.0.1:4001", "etcd server to connect to")
	targetDir  = flag.String("target-dir", "tgroups", "directory to store the target group files")
)

func main() {
	flag.Parse()

	client := etcd.NewClient([]string{*etcdServer})

	srvs := &services{
		m: map[string]Instances{},
	}
	updates := make(chan *etcd.Response)

	// Perform an initial read of all services.
	res, err := client.Get(servicesPrefix, false, true)
	if err != nil {
		log.Fatalf("Error on initial retrieval: %s", err)
	}
	srvs.handle(res.Node, srvs.update)
	srvs.persist()

	// Start watching for updates.
	go func() {
		_, err := client.Watch(servicesPrefix, 0, true, updates, nil)
		if err != nil {
			log.Errorln(err)
		}
	}()

	// Apply updates sent on the channel.
	for res := range updates {
		h := srvs.update
		if res.Action == "delete" {
			h = srvs.delete
			log.Println("delete", res.Node)
		} else {
			log.Println("update", res.Node)
		}
		srvs.handle(res.Node, h)
		srvs.persist()
	}
}

func (srvs *services) handle(node *etcd.Node, h func(*etcd.Node)) {
	if node.Dir {
		for _, n := range node.Nodes {
			srvs.handle(n, h)
		}
	}
	if !pathPat.MatchString(node.Key) {
		log.Warnf("unhandled key %q", node.Key)
		return
	}
	h(node)
}

// delete services or instances based on the given node.
func (srvs *services) delete(node *etcd.Node) {
	log.Println("delete", node)
	match := pathPat.FindStringSubmatch(node.Key)
	srv := match[1]
	// Deletion of an entire service.
	if match[2] == "" {
		srvs.del = append(srvs.del, srv)
		delete(srvs.m, srv)
		return
	}

	instances, ok := srvs.m[srv]
	if !ok {
		log.Errorf("Received delete for unknown service %s", srv)
		return
	}
	delete(instances, match[2])
}

// update the services based on the given node.
func (srvs *services) update(node *etcd.Node) {
	log.Println("update", node)
	match := pathPat.FindStringSubmatch(node.Key)
	// Creating a new job dir does not require an action.
	if match[2] == "" {
		return
	}
	srv := match[1]

	instances, ok := srvs.m[srv]
	if !ok {
		instances = Instances{}
	}
	instances[match[2]] = node.Value
	srvs.m[srv] = instances
}

// persist writes the current services to disc.
func (srvs *services) persist() {
	// Write files for current services.
	for job, instances := range srvs.m {
		var targets []string
		for _, addr := range instances {
			targets = append(targets, addr)
		}
		content, err := json.Marshal([]*TargetGroup{
			{
				Targets: targets,
				Labels:  map[string]string{"job": job},
			},
		})
		if err != nil {
			log.Errorln(err)
			continue
		}

		f, err := create(filepath.Join(*targetDir, job+".json"))
		if err != nil {
			log.Errorln(err)
			continue
		}
		if _, err := f.Write(content); err != nil {
			log.Errorln(err)
		}
		f.Close()
	}
	// Remove files for disappeared services.
	for _, job := range srvs.del {
		if err := os.Remove(filepath.Join(*targetDir, job+".json")); err != nil {
			log.Errorln(err)
		}
	}
	srvs.del = nil
}
<commit_msg>Improve info logging<commit_after>package main

import (
	"encoding/json"
	"flag"
	"os"
	"path/filepath"
	"regexp"

	"github.com/coreos/go-etcd/etcd"
	"github.com/prometheus/log"
)

const servicesPrefix = "/services"

var pathPat = regexp.MustCompile(`/services/([^/]+)(?:/(\d+))?`)

// TargetGroup is the target group read by Prometheus.
type TargetGroup struct {
	Targets []string          `json:"targets,omitempty"`
	Labels  map[string]string `json:"labels,omitempty"`
}

type Instances map[string]string

// services are the services stored in etcd.
type services struct {
	m   map[string]Instances // The current services.
	del []string             // Services deleted in the last update.
}

var (
	etcdServer = flag.String("server", "http://127.0.0.1:4001", "etcd server to connect to")
	targetDir  = flag.String("target-dir", "tgroups", "directory to store the target group files")
)

func main() {
	flag.Parse()

	client := etcd.NewClient([]string{*etcdServer})

	srvs := &services{
		m: map[string]Instances{},
	}
	updates := make(chan *etcd.Response)

	// Perform an initial read of all services.
	res, err := client.Get(servicesPrefix, false, true)
	if err != nil {
		log.Fatalf("Error on initial retrieval: %s", err)
	}
	srvs.handle(res.Node, srvs.update)
	srvs.persist()

	// Start watching for updates.
	go func() {
		_, err := client.Watch(servicesPrefix, 0, true, updates, nil)
		if err != nil {
			log.Errorln(err)
		}
	}()

	// Apply updates sent on the channel.
	for res := range updates {
		h := srvs.update
		if res.Action == "delete" {
			h = srvs.delete
			log.Infoln("delete", res.Node.Key, res.Node.Value)
		} else {
			log.Infoln("update", res.Node.Key, res.Node.Value)
		}
		srvs.handle(res.Node, h)
		srvs.persist()
	}
}

func (srvs *services) handle(node *etcd.Node, h func(*etcd.Node)) {
	if node.Dir {
		for _, n := range node.Nodes {
			srvs.handle(n, h)
		}
	}
	if !pathPat.MatchString(node.Key) {
		log.Warnf("unhandled key %q", node.Key)
		return
	}
	h(node)
}

// delete services or instances based on the given node.
func (srvs *services) delete(node *etcd.Node) {
	match := pathPat.FindStringSubmatch(node.Key)
	srv := match[1]
	// Deletion of an entire service.
	if match[2] == "" {
		srvs.del = append(srvs.del, srv)
		delete(srvs.m, srv)
		return
	}

	instances, ok := srvs.m[srv]
	if !ok {
		log.Errorf("Received delete for unknown service %s", srv)
		return
	}
	delete(instances, match[2])
}

// update the services based on the given node.
func (srvs *services) update(node *etcd.Node) {
	match := pathPat.FindStringSubmatch(node.Key)
	// Creating a new job dir does not require an action.
	if match[2] == "" {
		return
	}
	srv := match[1]

	instances, ok := srvs.m[srv]
	if !ok {
		instances = Instances{}
	}
	instances[match[2]] = node.Value
	srvs.m[srv] = instances
}

// persist writes the current services to disc.
func (srvs *services) persist() {
	// Write files for current services.
	for job, instances := range srvs.m {
		var targets []string
		for _, addr := range instances {
			targets = append(targets, addr)
		}
		content, err := json.Marshal([]*TargetGroup{
			{
				Targets: targets,
				Labels:  map[string]string{"job": job},
			},
		})
		if err != nil {
			log.Errorln(err)
			continue
		}

		f, err := create(filepath.Join(*targetDir, job+".json"))
		if err != nil {
			log.Errorln(err)
			continue
		}
		if _, err := f.Write(content); err != nil {
			log.Errorln(err)
		}
		f.Close()
	}
	// Remove files for disappeared services.
	for _, job := range srvs.del {
		if err := os.Remove(filepath.Join(*targetDir, job+".json")); err != nil {
			log.Errorln(err)
		}
	}
	srvs.del = nil
}
<|endoftext|>


<commit_before>package workq

import "sync"

type Queue struct {
	items []*Item
	mutex sync.Mutex
}

func (q *Queue) Push(item *Item) {
	q.mutex.Lock()
	defer q.mutex.Unlock()

	q.items = append(q.items, item)
	go item.Translate()
}

func (q *Queue) Pop() *Item {
	q.mutex.Lock()
	defer q.mutex.Unlock()

	if !q.IsEmpty() {
		item := (q.items)[0]
		<-item.Done
		q.items = q.items[1:len(q.items)]
		return item
	}
	return nil
}

func (q *Queue) Len() int {
	return len(q.items)
}

func (q *Queue) IsEmpty() bool {
	return q.Len() == 0
}
<commit_msg>Use a channel inside the Queue<commit_after>package workq

type Queue struct {
	ch chan *Item
}

func NewQueue() *Queue {
	q := new(Queue)
	q.ch = make(chan *Item, 10)
	return q
}

func (q *Queue) Push(item *Item) {
	q.ch <- item
	go item.Translate()
}

func (q *Queue) Pop() <-chan *Item {
	ch := make(chan *Item)
	go func() {
		for item := range q.ch {
			<-item.Done
			ch <- item
		}
	}()
	return ch
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"math"
	"math/rand"
	"os"
	"strconv"
	"time"

	"github.com/piotrkowalczuk/pqt"
	"github.com/piotrkowalczuk/pqt/example/app/internal/model"
	"github.com/piotrkowalczuk/sklog"
)

var (
	address string
	dbg     bool
)

func init() {
	flag.StringVar(&address, "addr", "postgres://localhost:5432/test?sslmode=disable", "postgres connection string")
	flag.BoolVar(&dbg, "dbg", true, "debug mode")
}

func main() {
	flag.Parse()
	log := sklog.NewHumaneLogger(os.Stdout, sklog.DefaultHTTPFormatter)
	db, err := sql.Open("postgres", address)
	if err != nil {
		sklog.Fatal(log, err)
	}

	_, err = db.Exec(model.SQL)
	if err != nil {
		sklog.Fatal(log, err)
	}

	repo := struct {
		news     model.NewsRepositoryBase
		comment  model.CommentRepositoryBase
		category model.CategoryRepositoryBase
	}{
		news: model.NewsRepositoryBase{
			DB:      db,
			Table:   model.TableNews,
			Columns: model.TableNewsColumns,
			Debug:   true,
			Log:     log,
		},
		comment: model.CommentRepositoryBase{
			DB:      db,
			Table:   model.TableComment,
			Columns: model.TableCommentColumns,
			Debug:   true,
			Log:     log,
		},
		category: model.CategoryRepositoryBase{
			DB:      db,
			Table:   model.TableCategory,
			Columns: model.TableCategoryColumns,
			Debug:   true,
			Log:     log,
		},
	}

	ctx := context.Background()

	count, err := repo.news.Count(ctx, &model.NewsCriteria{})
	if err != nil {
		sklog.Fatal(log, err)
	}
	sklog.Debug(log, "number of news fetched", "count", count)

	news, err := repo.news.Insert(ctx, &model.NewsEntity{
		Title:   fmt.Sprintf("Lorem Ipsum - %d - %d", time.Now().Unix(), rand.Int63()),
		Content: "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Nullam a felis vel erat gravida luctus at id nisi. Cum sociis natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Vivamus a nibh massa. Cum sociis natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Fusce viverra quam id dolor facilisis ultrices. Donec blandit, justo sit amet consequat gravida, nisi velit efficitur neque, ac ullamcorper leo dui vitae lorem. Pellentesque vitae ligula id massa fringilla facilisis eu sit amet neque. Ut ac fringilla mi. Maecenas id fermentum massa. Duis at tristique felis, nec aliquet nisi. Suspendisse potenti. In sed dolor maximus, dapibus arcu vitae, vehicula ligula. Nunc imperdiet eu ipsum sed pretium. Nullam iaculis nunc id dictum auctor.",
		Lead:    sql.NullString{String: "Neque porro quisquam est qui dolorem ipsum quia dolor sit amet, consectetur, adipisci velit...", Valid: true},
	})
	if err != nil {
		switch pqt.ErrorConstraint(err) {
		case model.TableNewsConstraintTitleUnique:
			sklog.Fatal(log, errors.New("news with such title already exists"))
		default:
			sklog.Fatal(log, err)
		}
	}

	nb := 20
	for i := 0; i < nb; i++ {
		_, err = repo.comment.Insert(ctx, &model.CommentEntity{
			NewsID:    news.ID,
			NewsTitle: news.Title,
			Content:   "Etiam eget nunc vel tellus placerat accumsan. Quisque dictum commodo orci, a eleifend nulla viverra malesuada. Etiam dui purus, dapibus a risus sed, porta scelerisque lorem. Sed vehicula mauris tellus, at dapibus risus facilisis vitae. Sed at lacus mollis, cursus sapien eu, egestas ligula. Cras blandit, arcu quis aliquam dictum, nibh purus pulvinar turpis, in dapibus est nibh et enim. Donec ex arcu, iaculis eget euismod id, lobortis nec enim. Quisque sed massa vel dui convallis ultrices. Nulla rutrum sed lacus vel ornare. Aliquam vulputate condimentum elit at pellentesque. Curabitur vitae sem tincidunt, volutpat urna ut, consequat turpis. Pellentesque varius justo libero, a volutpat lacus vulputate at. Integer tristique pharetra urna vel pharetra. In porttitor tincidunt eros, vel eleifend quam elementum a.",
		})
		if err != nil {
			sklog.Fatal(log, err)
		}
	}

	iter, err := repo.comment.FindIter(ctx, &model.CommentFindExpr{
		OrderBy: map[string]bool{
			"id": false,
			"non_existing_column": true,
		},
		Where: &model.CommentCriteria{
			NewsID: sql.NullInt64{Int64: news.ID, Valid: true},
		},
	})
	if err != nil {
		sklog.Fatal(log, err)
	}
	got := 0
	for iter.Next() {
		com, err := iter.Comment()
		if err != nil {
			sklog.Fatal(log, err)
		}
		sklog.Debug(log, "comment fetched", "comment_id", com.ID)
		got++
	}
	if err = iter.Err(); err != nil {
		sklog.Fatal(log, err)
	}
	if nb != got {
		sklog.Fatal(log, fmt.Errorf("wrong number of comments, expected %d but got %d", nb, got))
	} else {
		sklog.Info(log, "proper number of comments")
	}

	category, err := repo.category.Insert(ctx, &model.CategoryEntity{
		Name: "parent",
	})
	if err != nil {
		sklog.Fatal(log, err)
	}

	for i := 0; i < nb; i++ {
		_, err := repo.category.Insert(ctx, &model.CategoryEntity{
			ParentID: sql.NullInt64{Int64: category.ID, Valid: true},
			Name:     "child_category" + strconv.Itoa(i),
		})
		if err != nil {
			sklog.Fatal(log, err)
		}
	}

	count, err = repo.category.Count(ctx, &model.CategoryCriteria{
		ParentID: sql.NullInt64{Int64: category.ID, Valid: true},
	})
	if err != nil {
		sklog.Fatal(log, err)
	}
	if count != int64(nb) {
		sklog.Fatal(log, fmt.Errorf("wrong number of categories, expected %d but got %d", nb, count))
	} else {
		sklog.Info(log, "proper number of categories")
	}

	_, err = repo.category.Insert(ctx, &model.CategoryEntity{
		ParentID: sql.NullInt64{Int64: int64(math.MaxInt64 - 1), Valid: true},
		Name:     "does not work",
	})
	if err != nil {
		switch pqt.ErrorConstraint(err) {
		case model.TableCategoryConstraintParentIDForeignKey:
			sklog.Info(log, "category parent id constraint properly catched, category with such id does not exists")
		default:
			sklog.Fatal(log, fmt.Errorf("category constraint not catched properly, expected %s but got %s", model.TableCategoryConstraintParentIDForeignKey, pqt.ErrorConstraint(err)))
		}
	}

	ctx, cancel := context.WithTimeout(ctx, 1*time.Nanosecond)
	defer cancel()

	count, err = repo.news.Count(ctx, &model.NewsCriteria{})
	if err != nil {
		if err != context.DeadlineExceeded {
			sklog.Fatal(log, err)
		}
		sklog.Debug(log, "as expected, news count failed due to deadline")
	}
}
<commit_msg>joins - example app fixes<commit_after>package main

import (
	"context"
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"math"
	"math/rand"
	"os"
	"strconv"
	"time"

	"github.com/piotrkowalczuk/pqt"
	"github.com/piotrkowalczuk/pqt/example/app/internal/model"
	"github.com/piotrkowalczuk/sklog"
)

var (
	address string
	dbg     bool
)

func init() {
	flag.StringVar(&address, "addr", "postgres://localhost:5432/test?sslmode=disable", "postgres connection string")
	flag.BoolVar(&dbg, "dbg", true, "debug mode")
}

func main() {
	flag.Parse()
	log := sklog.NewHumaneLogger(os.Stdout, sklog.DefaultHTTPFormatter)
	db, err := sql.Open("postgres", address)
	if err != nil {
		sklog.Fatal(log, err)
	}

	_, err = db.Exec(model.SQL)
	if err != nil {
		sklog.Fatal(log, err)
	}

	repo := struct {
		news     model.NewsRepositoryBase
		comment  model.CommentRepositoryBase
		category model.CategoryRepositoryBase
	}{
		news: model.NewsRepositoryBase{
			DB:      db,
			Table:   model.TableNews,
			Columns: model.TableNewsColumns,
			Debug:   true,
			Log:     log,
		},
		comment: model.CommentRepositoryBase{
			DB:      db,
			Table:   model.TableComment,
			Columns: model.TableCommentColumns,
			Debug:   true,
			Log:     log,
		},
		category: model.CategoryRepositoryBase{
			DB:      db,
			Table:   model.TableCategory,
			Columns: model.TableCategoryColumns,
			Debug:   true,
			Log:     log,
		},
	}

	ctx := context.Background()

	count, err := repo.news.Count(ctx, &model.NewsCountExpr{})
	if err != nil {
		sklog.Fatal(log, err)
	}
	sklog.Debug(log, "number of news fetched", "count", count)

	news, err := repo.news.Insert(ctx, &model.NewsEntity{
		Title:   fmt.Sprintf("Lorem Ipsum - %d - %d", time.Now().Unix(), rand.Int63()),
		Content: "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Nullam a felis vel erat gravida luctus at id nisi. Cum sociis natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Vivamus a nibh massa. Cum sociis natoque penatibus et magnis dis parturient montes, nascetur ridiculus mus. Fusce viverra quam id dolor facilisis ultrices. Donec blandit, justo sit amet consequat gravida, nisi velit efficitur neque, ac ullamcorper leo dui vitae lorem. Pellentesque vitae ligula id massa fringilla facilisis eu sit amet neque. Ut ac fringilla mi. Maecenas id fermentum massa. Duis at tristique felis, nec aliquet nisi. Suspendisse potenti. In sed dolor maximus, dapibus arcu vitae, vehicula ligula. Nunc imperdiet eu ipsum sed pretium. Nullam iaculis nunc id dictum auctor.",
		Lead:    sql.NullString{String: "Neque porro quisquam est qui dolorem ipsum quia dolor sit amet, consectetur, adipisci velit...", Valid: true},
	})
	if err != nil {
		switch pqt.ErrorConstraint(err) {
		case model.TableNewsConstraintTitleUnique:
			sklog.Fatal(log, errors.New("news with such title already exists"))
		default:
			sklog.Fatal(log, err)
		}
	}

	nb := 20
	for i := 0; i < nb; i++ {
		_, err = repo.comment.Insert(ctx, &model.CommentEntity{
			NewsID:    news.ID,
			NewsTitle: news.Title,
			Content:   "Etiam eget nunc vel tellus placerat accumsan. Quisque dictum commodo orci, a eleifend nulla viverra malesuada. Etiam dui purus, dapibus a risus sed, porta scelerisque lorem. Sed vehicula mauris tellus, at dapibus risus facilisis vitae. Sed at lacus mollis, cursus sapien eu, egestas ligula. Cras blandit, arcu quis aliquam dictum, nibh purus pulvinar turpis, in dapibus est nibh et enim. Donec ex arcu, iaculis eget euismod id, lobortis nec enim. Quisque sed massa vel dui convallis ultrices. Nulla rutrum sed lacus vel ornare. Aliquam vulputate condimentum elit at pellentesque. Curabitur vitae sem tincidunt, volutpat urna ut, consequat turpis. Pellentesque varius justo libero, a volutpat lacus vulputate at. Integer tristique pharetra urna vel pharetra. In porttitor tincidunt eros, vel eleifend quam elementum a.",
		})
		if err != nil {
			sklog.Fatal(log, err)
		}
	}

	iter, err := repo.comment.FindIter(ctx, &model.CommentFindExpr{
		OrderBy: map[string]bool{
			"id": false,
			"non_existing_column": true,
		},
		Where: &model.CommentCriteria{
			NewsID: sql.NullInt64{Int64: news.ID, Valid: true},
		},
	})
	if err != nil {
		sklog.Fatal(log, err)
	}
	got := 0
	for iter.Next() {
		com, err := iter.Comment()
		if err != nil {
			sklog.Fatal(log, err)
		}
		sklog.Debug(log, "comment fetched", "comment_id", com.ID)
		got++
	}
	if err = iter.Err(); err != nil {
		sklog.Fatal(log, err)
	}
	if nb != got {
		sklog.Fatal(log, fmt.Errorf("wrong number of comments, expected %d but got %d", nb, got))
	} else {
		sklog.Info(log, "proper number of comments")
	}

	category, err := repo.category.Insert(ctx, &model.CategoryEntity{
		Name: "parent",
	})
	if err != nil {
		sklog.Fatal(log, err)
	}

	for i := 0; i < nb; i++ {
		_, err := repo.category.Insert(ctx, &model.CategoryEntity{
			ParentID: sql.NullInt64{Int64: category.ID, Valid: true},
			Name:     "child_category" + strconv.Itoa(i),
		})
		if err != nil {
			sklog.Fatal(log, err)
		}
	}

	count, err = repo.category.Count(ctx, &model.CategoryCountExpr{
		Where: &model.CategoryCriteria{
			ParentID: sql.NullInt64{Int64: category.ID, Valid: true},
		},
	})
	if err != nil {
		sklog.Fatal(log, err)
	}
	if count != int64(nb) {
		sklog.Fatal(log, fmt.Errorf("wrong number of categories, expected %d but got %d", nb, count))
	} else {
		sklog.Info(log, "proper number of categories")
	}

	_, err = repo.category.Insert(ctx, &model.CategoryEntity{
		ParentID: sql.NullInt64{Int64: int64(math.MaxInt64 - 1), Valid: true},
		Name:     "does not work",
	})
	if err != nil {
		switch pqt.ErrorConstraint(err) {
		case model.TableCategoryConstraintParentIDForeignKey:
			sklog.Info(log, "category parent id constraint properly catched, category with such id does not exists")
		default:
			sklog.Fatal(log, fmt.Errorf("category constraint not catched properly, expected %s but got %s", model.TableCategoryConstraintParentIDForeignKey, pqt.ErrorConstraint(err)))
		}
	}

	ctx, cancel := context.WithTimeout(ctx, 1*time.Nanosecond)
	defer cancel()

	count, err = repo.news.Count(ctx, &model.NewsCountExpr{})
	if err != nil {
		if err != context.DeadlineExceeded {
			sklog.Fatal(log, err)
		}
		sklog.Debug(log, "as expected, news count failed due to deadline")
	}
}
<|endoftext|>


<commit_before>package main

import "fmt"

func main() {
	iter := iterator()

	iter(5)
	for iter() {
		fmt.Println("hello")
	}
}

func iterator() func(...int) bool {
	counter := 0
	running := false
	return func(i ...int) bool {
		if i == nil {
			if counter <= 0 {
				running = false
			}
			counter--
		} else {
			counter = i[0]
			running = true
		}
		return running
	}
}
<commit_msg>Update iterator.go<commit_after>package main

import "fmt"

func main() {
	iter := iterator()

	iter(5)
	for iter() {
		fmt.Println("hello")
	}
	// OUTPUT:
	// hello
	// hello
	// hello
	// hello
	// hello

}

func iterator() func(...int) bool {
	counter := 0
	running := false
	return func(i ...int) bool {
		if i == nil {
			if counter <= 0 {
				running = false
			}
			counter--
		} else {
			counter = i[0]
			running = true
		}
		return running
	}
}
<|endoftext|>


<commit_before><commit_msg>Removed value from Table.Equals() error message.<commit_after><|endoftext|>


<commit_before>package cloudatgost

import(
	"bytes"
	"net/http"
	"net/url"
)

// A PowerOp represents a successful power operation job response.
// It contains miscellaneous informations about the API and the job,
// like the task id and the result.
type PowerOp struct {
	Status string `json:"status"`
	Time int `json:"time"`
	API string `json:"api"`
	Serverid string `json:"serverid"`
	Action string `json:"action"`
	Taskid int64 `json:"taskid"`
	Result string `json:"result"`
}

// Action is a function that behaves as the actual component of the power management.
// It accepts a serverID and an operation string as its parameters, then forms an
// HTTP POST request to the endpoint. It can be used as a standalone support for
// power operations, but it serves well as a base for shorthands.
func (c *Client) Action(serverID string, operation string) (*PowerOp) {
	v := &PowerOp{}
	URL, err := url.Parse(c.BaseURL)
	if err != nil {
		panic("boom! Busted :F")
	}
	URL.Path += "powerop.php"
	parameters := url.Values{}
	parameters.Add("key", c.Token)
	parameters.Add("login", c.Login)
	parameters.Add("sid", serverID)
	parameters.Add("action", operation)

	request, err := http.NewRequest("POST", URL.String(), bytes.NewBufferString(parameters.Encode()))
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	if err != nil {
		return nil
	}

	c.Do(request, &v)
	return v
}

// PowerOn is a shorthand function that performs a "poweron" request through
// the Action function. It accepts a serverID as its unique parameter.
func (c *Client) PowerOn(serverID string) (*PowerOp) {
	return Action(serverID, "poweron")
}

// PowerOff is a shorthand function that performs a "poweroff" request through
// the Action function. It accepts a serverID as its unique parameter.
func (c *Client) PowerOff(serverID string) (*PowerOp) {
	return Action(serverID, "poweroff")
}

// Reboot is a shorthand function that performs a "reset" request through
// the Action function. It accepts a serverID as its unique parameter.
func (c *Client) Reboot(serverID string) (*PowerOp) {
	return Action(serverID, "reset")
}
<commit_msg>fixed a bug<commit_after>package cloudatgost

import(
	"bytes"
	"net/http"
	"net/url"
)

// A PowerOp represents a successful power operation job response.
// It contains miscellaneous informations about the API and the job,
// like the task id and the result.
type PowerOp struct {
	Status string `json:"status"`
	Time int `json:"time"`
	API string `json:"api"`
	Serverid string `json:"serverid"`
	Action string `json:"action"`
	Taskid int64 `json:"taskid"`
	Result string `json:"result"`
}

// Action is a function that behaves as the actual component of the power management.
// It accepts a serverID and an operation string as its parameters, then forms an
// HTTP POST request to the endpoint. It can be used as a standalone support for
// power operations, but it serves well as a base for shorthands.
func (c *Client) Action(serverID string, operation string) (*PowerOp) {
	v := &PowerOp{}
	URL, err := url.Parse(c.BaseURL)
	if err != nil {
		panic("boom! Busted :F")
	}
	URL.Path += "powerop.php"
	parameters := url.Values{}
	parameters.Add("key", c.Token)
	parameters.Add("login", c.Login)
	parameters.Add("sid", serverID)
	parameters.Add("action", operation)

	request, err := http.NewRequest("POST", URL.String(), bytes.NewBufferString(parameters.Encode()))
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	if err != nil {
		return nil
	}

	c.Do(request, &v)
	return v
}

// PowerOn is a shorthand function that performs a "poweron" request through
// the Action function. It accepts a serverID as its unique parameter.
func (c *Client) PowerOn(serverID string) (*PowerOp) {
	return c.Action(serverID, "poweron")
}

// PowerOff is a shorthand function that performs a "poweroff" request through
// the Action function. It accepts a serverID as its unique parameter.
func (c *Client) PowerOff(serverID string) (*PowerOp) {
	return c.Action(serverID, "poweroff")
}

// Reboot is a shorthand function that performs a "reset" request through
// the Action function. It accepts a serverID as its unique parameter.
func (c *Client) Reboot(serverID string) (*PowerOp) {
	return c.Action(serverID, "reset")
}
<|endoftext|>


<commit_before>package tpls

import (
	"os"
	"path/filepath"
	"text/template"

	"github.com/mattn/go-zglob"
	"github.com/mh-cbon/go-msi/manifest"
)

// find all wxs fies in given directory
func Find(srcDir string, pattern string) ([]string, error) {
	glob := filepath.Join(srcDir, pattern)
	return zglob.Glob(glob)
}

// Generate given src file to out file using given manifest
func GenerateTemplate(wixFile *manifest.WixManifest, src string, out string) error {
	tpl, err := template.ParseFiles(src)
	if err != nil {
		return err
	}

	fileWriter, err := os.Create(out)
	if err != nil {
		return err
	}
	defer fileWriter.Close()
	err = tpl.Execute(fileWriter, wixFile)
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>doc and lint<commit_after>package tpls

import (
	"os"
	"path/filepath"
	"text/template"

	"github.com/mattn/go-zglob"
	"github.com/mh-cbon/go-msi/manifest"
)

// Find all wxs fies in given directory
func Find(srcDir string, pattern string) ([]string, error) {
	glob := filepath.Join(srcDir, pattern)
	return zglob.Glob(glob)
}

// GenerateTemplate generates given src template to out file using given manifest
func GenerateTemplate(wixFile *manifest.WixManifest, src string, out string) error {
	tpl, err := template.ParseFiles(src)
	if err != nil {
		return err
	}

	fileWriter, err := os.Create(out)
	if err != nil {
		return err
	}
	defer fileWriter.Close()
	err = tpl.Execute(fileWriter, wixFile)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Gravitational, 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 trace

import (
	"bytes"
	"fmt"
	"net/http"
	"os"
	"strings"
	"testing"

	log "github.com/Sirupsen/logrus"

	"golang.org/x/net/context"
	. "gopkg.in/check.v1"
)

func TestTrace(t *testing.T) { TestingT(t) }

type TraceSuite struct {
}

var _ = Suite(&TraceSuite{})

func (s *TraceSuite) TestEmpty(c *C) {
	c.Assert(DebugReport(nil), Equals, "")
	c.Assert(UserMessage(nil), Equals, "")
}

func (s *TraceSuite) TestWrap(c *C) {
	testErr := &TestError{Param: "param"}
	err := Wrap(Wrap(testErr))

	c.Assert(line(DebugReport(err)), Matches, ".*trace_test.go.*")
	c.Assert(line(UserMessage(err)), Not(Matches), ".*trace_test.go.*")
}

func (s *TraceSuite) TestOrigError(c *C) {
	testErr := fmt.Errorf("some error")
	err := Wrap(Wrap(testErr))
	c.Assert(err.OrigError(), Equals, testErr)
}

func (s *TraceSuite) TestWrapMessage(c *C) {
	testErr := fmt.Errorf("description")

	err := Wrap(testErr)

	SetDebug(true)
	c.Assert(line(err.Error()), Matches, ".*trace_test.go.*")
	c.Assert(line(err.Error()), Matches, ".*description.*")

	SetDebug(false)
	c.Assert(line(err.Error()), Not(Matches), ".*trace_test.go.*")
	c.Assert(line(err.Error()), Matches, ".*description.*")
}

func (s *TraceSuite) TestWrapUserMessage(c *C) {
	testErr := fmt.Errorf("description")

	err := Wrap(testErr, "user message")
	c.Assert(line(UserMessage(err)), Equals, "user message")

	err = Wrap(err, "user message 2")
	c.Assert(line(UserMessage(err)), Equals, "user message, user message 2")
}

func (s *TraceSuite) TestWrapNil(c *C) {
	err1 := Wrap(nil)
	c.Assert(err1, IsNil)

	var err2 error
	err2 = nil

	err3 := Wrap(err2)
	c.Assert(err3, IsNil)

	err4 := Wrap(err3)
	c.Assert(err4, IsNil)
}

func (s *TraceSuite) TestWrapStdlibErrors(c *C) {
	c.Assert(IsNotFound(os.ErrNotExist), Equals, true)
}

func (s *TraceSuite) TestLogFormatter(c *C) {

	for _, f := range []log.Formatter{&TextFormatter{}, &JSONFormatter{}} {
		log.SetFormatter(f)

		// check case with global Infof
		buf := &bytes.Buffer{}
		log.SetOutput(buf)
		log.Infof("hello")
		c.Assert(line(buf.String()), Matches, ".*trace_test.go.*")

		// check case with embedded Infof
		buf = &bytes.Buffer{}
		log.SetOutput(buf)
		log.WithFields(log.Fields{"a": "b"}).Infof("hello")
		c.Assert(line(buf.String()), Matches, ".*trace_test.go.*")
	}
}

func (s *TraceSuite) TestGenericErrors(c *C) {
	testCases := []struct {
		Err        error
		Predicate  func(error) bool
		StatusCode int
	}{
		{
			Err:        NotFound("not found"),
			Predicate:  IsNotFound,
			StatusCode: http.StatusNotFound,
		},
		{
			Err:        AlreadyExists("already exists"),
			Predicate:  IsAlreadyExists,
			StatusCode: http.StatusConflict,
		},
		{
			Err:        BadParameter("is bad"),
			Predicate:  IsBadParameter,
			StatusCode: http.StatusBadRequest,
		},
		{
			Err:        CompareFailed("is bad"),
			Predicate:  IsCompareFailed,
			StatusCode: http.StatusPreconditionFailed,
		},
		{
			Err:        AccessDenied("denied"),
			Predicate:  IsAccessDenied,
			StatusCode: http.StatusForbidden,
		},
		{
			Err:        ConnectionProblem(nil, "prob"),
			Predicate:  IsConnectionProblem,
			StatusCode: http.StatusRequestTimeout,
		},
		{
			Err:        LimitExceeded("limit exceeded"),
			Predicate:  IsLimitExceeded,
			StatusCode: statusTooManyRequests,
		},
	}

	for i, testCase := range testCases {
		comment := Commentf("test case #%v", i+1)
		SetDebug(true)
		err := testCase.Err

		t := err.(*TraceErr)
		c.Assert(len(t.Traces), Not(Equals), 0, comment)
		c.Assert(line(err.Error()), Matches, "*.trace_test.go.*", comment)
		c.Assert(testCase.Predicate(err), Equals, true, comment)

		w := newTestWriter()
		WriteError(w, err)
		outerr := ReadError(w.StatusCode, w.Body)
		c.Assert(testCase.Predicate(outerr), Equals, true, comment)
		t = outerr.(*TraceErr)
		c.Assert(len(t.Traces), Not(Equals), 0, comment)

		SetDebug(false)
		w = newTestWriter()
		WriteError(w, err)
		outerr = ReadError(w.StatusCode, w.Body)
		c.Assert(testCase.Predicate(outerr), Equals, true, comment)
	}
}

// Make sure we write some output produced by standard errors
func (s *TraceSuite) TestWriteExternalErrors(c *C) {
	err := fmt.Errorf("snap!")

	SetDebug(true)
	w := newTestWriter()
	WriteError(w, err)
	c.Assert(w.StatusCode, Equals, http.StatusInternalServerError)
	c.Assert(strings.Replace(string(w.Body), "\n", "", -1), Matches, "*.snap.*")

	SetDebug(false)
	w = newTestWriter()
	WriteError(w, err)
	c.Assert(w.StatusCode, Equals, http.StatusInternalServerError)
	c.Assert(strings.Replace(string(w.Body), "\n", "", -1), Matches, "*.snap.*")
}

func (s *TraceSuite) TestAggregates(c *C) {
	err1 := Errorf("failed one")
	err2 := Errorf("failed two")
	err := NewAggregate(err1, err2)
	c.Assert(IsAggregate(err), Equals, true)
	agg := Unwrap(err).(Aggregate)
	c.Assert(agg.Errors(), DeepEquals, []error{err1, err2})
	c.Assert(err.Error(), DeepEquals, "failed one, failed two")
}

func (s *TraceSuite) TestErrorf(c *C) {
	err := Errorf("error")
	c.Assert(line(DebugReport(err)), Matches, "*.trace_test.go.*")
	c.Assert(line(UserMessage(err)), Equals, "error")
}

func (s *TraceSuite) TestAggregateConvertsToCommonErrors(c *C) {
	testCases := []struct {
		Err                error
		Predicate          func(error) bool
		RoundtripPredicate func(error) bool
		StatusCode         int
	}{
		{
			// Aggregate unwraps to first aggregated error
			Err: NewAggregate(BadParameter("invalid value of foo"),
				LimitExceeded("limit exceeded")),
			Predicate:          IsAggregate,
			RoundtripPredicate: IsBadParameter,
			StatusCode:         http.StatusBadRequest,
		},
		{
			// Nested aggregate unwraps recursively
			Err: NewAggregate(NewAggregate(BadParameter("invalid value of foo"),
				LimitExceeded("limit exceeded"))),
			Predicate:          IsAggregate,
			RoundtripPredicate: IsBadParameter,
			StatusCode:         http.StatusBadRequest,
		},
	}
	for i, testCase := range testCases {
		comment := Commentf("test case #%v", i+1)
		SetDebug(true)
		err := testCase.Err

		c.Assert(line(err.Error()), Matches, "*.trace_test.go.*", comment)
		c.Assert(testCase.Predicate(err), Equals, true, comment)

		w := newTestWriter()
		WriteError(w, err)
		outerr := ReadError(w.StatusCode, w.Body)
		c.Assert(testCase.RoundtripPredicate(outerr), Equals, true, comment)

		t := outerr.(*TraceErr)
		c.Assert(len(t.Traces), Not(Equals), 0, comment)

		SetDebug(false)
		w = newTestWriter()
		WriteError(w, err)
		outerr = ReadError(w.StatusCode, w.Body)
		c.Assert(testCase.RoundtripPredicate(outerr), Equals, true, comment)
	}
}

func (s *TraceSuite) TestAggregateThrowAwayNils(c *C) {
	err := NewAggregate(fmt.Errorf("error1"), nil, fmt.Errorf("error2"))
	c.Assert(err.Error(), Not(Matches), ".*nil.*")
}

func (s *TraceSuite) TestAggregateAllNils(c *C) {
	c.Assert(NewAggregate(nil, nil, nil), IsNil)
}

func (s *TraceSuite) TestAggregateFromChannel(c *C) {
	errCh := make(chan error, 3)
	errCh <- fmt.Errorf("Snap!")
	errCh <- fmt.Errorf("BAM")
	errCh <- fmt.Errorf("omg")
	close(errCh)
	err := NewAggregateFromChannel(errCh, context.Background())
	c.Assert(err.Error(), Matches, ".*Snap!.*")
	c.Assert(err.Error(), Matches, ".*BAM.*")
	c.Assert(err.Error(), Matches, ".*omg.*")
}

func (s *TraceSuite) TestAggregateFromChannelCancel(c *C) {
	errCh := make(chan error, 3)
	errCh <- fmt.Errorf("Snap!")
	errCh <- fmt.Errorf("BAM")
	errCh <- fmt.Errorf("omg")
	ctx, cancel := context.WithCancel(context.Background())
	// we never closed the channel so we just need to make sure
	// the function exits when we cancel it
	cancel()
	NewAggregateFromChannel(errCh, ctx)
}

type TestError struct {
	Traces
	Param string
}

func (n *TestError) Error() string {
	return fmt.Sprintf("TestError(param=%v,trace=%v)", n.Param, n.Traces)
}

func (n *TestError) OrigError() error {
	return n
}

func newTestWriter() *testWriter {
	return &testWriter{
		H: make(http.Header),
	}
}

type testWriter struct {
	H          http.Header
	Body       []byte
	StatusCode int
}

func (tw *testWriter) Header() http.Header {
	return tw.H
}

func (tw *testWriter) Write(body []byte) (int, error) {
	tw.Body = body
	return len(tw.Body), nil
}

func (tw *testWriter) WriteHeader(code int) {
	tw.StatusCode = code
}

func line(s string) string {
	return strings.Replace(s, "\n", "", -1)
}
<commit_msg>Add a test<commit_after>/*
Copyright 2015 Gravitational, 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 trace

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"testing"

	log "github.com/Sirupsen/logrus"

	"golang.org/x/net/context"
	. "gopkg.in/check.v1"
)

func TestTrace(t *testing.T) { TestingT(t) }

type TraceSuite struct {
}

var _ = Suite(&TraceSuite{})

func (s *TraceSuite) TestEmpty(c *C) {
	c.Assert(DebugReport(nil), Equals, "")
	c.Assert(UserMessage(nil), Equals, "")
}

func (s *TraceSuite) TestWrap(c *C) {
	testErr := &TestError{Param: "param"}
	err := Wrap(Wrap(testErr))

	c.Assert(line(DebugReport(err)), Matches, ".*trace_test.go.*")
	c.Assert(line(UserMessage(err)), Not(Matches), ".*trace_test.go.*")
}

func (s *TraceSuite) TestOrigError(c *C) {
	testErr := fmt.Errorf("some error")
	err := Wrap(Wrap(testErr))
	c.Assert(err.OrigError(), Equals, testErr)
}

func (s *TraceSuite) TestIsEOF(c *C) {
	c.Assert(IsEOF(io.EOF), Equals, true)
	c.Assert(IsEOF(Wrap(io.EOF)), Equals, true)
}

func (s *TraceSuite) TestWrapMessage(c *C) {
	testErr := fmt.Errorf("description")

	err := Wrap(testErr)

	SetDebug(true)
	c.Assert(line(err.Error()), Matches, ".*trace_test.go.*")
	c.Assert(line(err.Error()), Matches, ".*description.*")

	SetDebug(false)
	c.Assert(line(err.Error()), Not(Matches), ".*trace_test.go.*")
	c.Assert(line(err.Error()), Matches, ".*description.*")
}

func (s *TraceSuite) TestWrapUserMessage(c *C) {
	testErr := fmt.Errorf("description")

	err := Wrap(testErr, "user message")
	c.Assert(line(UserMessage(err)), Equals, "user message")

	err = Wrap(err, "user message 2")
	c.Assert(line(UserMessage(err)), Equals, "user message, user message 2")
}

func (s *TraceSuite) TestWrapNil(c *C) {
	err1 := Wrap(nil)
	c.Assert(err1, IsNil)

	var err2 error
	err2 = nil

	err3 := Wrap(err2)
	c.Assert(err3, IsNil)

	err4 := Wrap(err3)
	c.Assert(err4, IsNil)
}

func (s *TraceSuite) TestWrapStdlibErrors(c *C) {
	c.Assert(IsNotFound(os.ErrNotExist), Equals, true)
}

func (s *TraceSuite) TestLogFormatter(c *C) {

	for _, f := range []log.Formatter{&TextFormatter{}, &JSONFormatter{}} {
		log.SetFormatter(f)

		// check case with global Infof
		buf := &bytes.Buffer{}
		log.SetOutput(buf)
		log.Infof("hello")
		c.Assert(line(buf.String()), Matches, ".*trace_test.go.*")

		// check case with embedded Infof
		buf = &bytes.Buffer{}
		log.SetOutput(buf)
		log.WithFields(log.Fields{"a": "b"}).Infof("hello")
		c.Assert(line(buf.String()), Matches, ".*trace_test.go.*")
	}
}

func (s *TraceSuite) TestGenericErrors(c *C) {
	testCases := []struct {
		Err        error
		Predicate  func(error) bool
		StatusCode int
	}{
		{
			Err:        NotFound("not found"),
			Predicate:  IsNotFound,
			StatusCode: http.StatusNotFound,
		},
		{
			Err:        AlreadyExists("already exists"),
			Predicate:  IsAlreadyExists,
			StatusCode: http.StatusConflict,
		},
		{
			Err:        BadParameter("is bad"),
			Predicate:  IsBadParameter,
			StatusCode: http.StatusBadRequest,
		},
		{
			Err:        CompareFailed("is bad"),
			Predicate:  IsCompareFailed,
			StatusCode: http.StatusPreconditionFailed,
		},
		{
			Err:        AccessDenied("denied"),
			Predicate:  IsAccessDenied,
			StatusCode: http.StatusForbidden,
		},
		{
			Err:        ConnectionProblem(nil, "prob"),
			Predicate:  IsConnectionProblem,
			StatusCode: http.StatusRequestTimeout,
		},
		{
			Err:        LimitExceeded("limit exceeded"),
			Predicate:  IsLimitExceeded,
			StatusCode: statusTooManyRequests,
		},
	}

	for i, testCase := range testCases {
		comment := Commentf("test case #%v", i+1)
		SetDebug(true)
		err := testCase.Err

		t := err.(*TraceErr)
		c.Assert(len(t.Traces), Not(Equals), 0, comment)
		c.Assert(line(err.Error()), Matches, "*.trace_test.go.*", comment)
		c.Assert(testCase.Predicate(err), Equals, true, comment)

		w := newTestWriter()
		WriteError(w, err)
		outerr := ReadError(w.StatusCode, w.Body)
		c.Assert(testCase.Predicate(outerr), Equals, true, comment)
		t = outerr.(*TraceErr)
		c.Assert(len(t.Traces), Not(Equals), 0, comment)

		SetDebug(false)
		w = newTestWriter()
		WriteError(w, err)
		outerr = ReadError(w.StatusCode, w.Body)
		c.Assert(testCase.Predicate(outerr), Equals, true, comment)
	}
}

// Make sure we write some output produced by standard errors
func (s *TraceSuite) TestWriteExternalErrors(c *C) {
	err := fmt.Errorf("snap!")

	SetDebug(true)
	w := newTestWriter()
	WriteError(w, err)
	c.Assert(w.StatusCode, Equals, http.StatusInternalServerError)
	c.Assert(strings.Replace(string(w.Body), "\n", "", -1), Matches, "*.snap.*")

	SetDebug(false)
	w = newTestWriter()
	WriteError(w, err)
	c.Assert(w.StatusCode, Equals, http.StatusInternalServerError)
	c.Assert(strings.Replace(string(w.Body), "\n", "", -1), Matches, "*.snap.*")
}

func (s *TraceSuite) TestAggregates(c *C) {
	err1 := Errorf("failed one")
	err2 := Errorf("failed two")
	err := NewAggregate(err1, err2)
	c.Assert(IsAggregate(err), Equals, true)
	agg := Unwrap(err).(Aggregate)
	c.Assert(agg.Errors(), DeepEquals, []error{err1, err2})
	c.Assert(err.Error(), DeepEquals, "failed one, failed two")
}

func (s *TraceSuite) TestErrorf(c *C) {
	err := Errorf("error")
	c.Assert(line(DebugReport(err)), Matches, "*.trace_test.go.*")
	c.Assert(line(UserMessage(err)), Equals, "error")
}

func (s *TraceSuite) TestAggregateConvertsToCommonErrors(c *C) {
	testCases := []struct {
		Err                error
		Predicate          func(error) bool
		RoundtripPredicate func(error) bool
		StatusCode         int
	}{
		{
			// Aggregate unwraps to first aggregated error
			Err: NewAggregate(BadParameter("invalid value of foo"),
				LimitExceeded("limit exceeded")),
			Predicate:          IsAggregate,
			RoundtripPredicate: IsBadParameter,
			StatusCode:         http.StatusBadRequest,
		},
		{
			// Nested aggregate unwraps recursively
			Err: NewAggregate(NewAggregate(BadParameter("invalid value of foo"),
				LimitExceeded("limit exceeded"))),
			Predicate:          IsAggregate,
			RoundtripPredicate: IsBadParameter,
			StatusCode:         http.StatusBadRequest,
		},
	}
	for i, testCase := range testCases {
		comment := Commentf("test case #%v", i+1)
		SetDebug(true)
		err := testCase.Err

		c.Assert(line(err.Error()), Matches, "*.trace_test.go.*", comment)
		c.Assert(testCase.Predicate(err), Equals, true, comment)

		w := newTestWriter()
		WriteError(w, err)
		outerr := ReadError(w.StatusCode, w.Body)
		c.Assert(testCase.RoundtripPredicate(outerr), Equals, true, comment)

		t := outerr.(*TraceErr)
		c.Assert(len(t.Traces), Not(Equals), 0, comment)

		SetDebug(false)
		w = newTestWriter()
		WriteError(w, err)
		outerr = ReadError(w.StatusCode, w.Body)
		c.Assert(testCase.RoundtripPredicate(outerr), Equals, true, comment)
	}
}

func (s *TraceSuite) TestAggregateThrowAwayNils(c *C) {
	err := NewAggregate(fmt.Errorf("error1"), nil, fmt.Errorf("error2"))
	c.Assert(err.Error(), Not(Matches), ".*nil.*")
}

func (s *TraceSuite) TestAggregateAllNils(c *C) {
	c.Assert(NewAggregate(nil, nil, nil), IsNil)
}

func (s *TraceSuite) TestAggregateFromChannel(c *C) {
	errCh := make(chan error, 3)
	errCh <- fmt.Errorf("Snap!")
	errCh <- fmt.Errorf("BAM")
	errCh <- fmt.Errorf("omg")
	close(errCh)
	err := NewAggregateFromChannel(errCh, context.Background())
	c.Assert(err.Error(), Matches, ".*Snap!.*")
	c.Assert(err.Error(), Matches, ".*BAM.*")
	c.Assert(err.Error(), Matches, ".*omg.*")
}

func (s *TraceSuite) TestAggregateFromChannelCancel(c *C) {
	errCh := make(chan error, 3)
	errCh <- fmt.Errorf("Snap!")
	errCh <- fmt.Errorf("BAM")
	errCh <- fmt.Errorf("omg")
	ctx, cancel := context.WithCancel(context.Background())
	// we never closed the channel so we just need to make sure
	// the function exits when we cancel it
	cancel()
	NewAggregateFromChannel(errCh, ctx)
}

type TestError struct {
	Traces
	Param string
}

func (n *TestError) Error() string {
	return fmt.Sprintf("TestError(param=%v,trace=%v)", n.Param, n.Traces)
}

func (n *TestError) OrigError() error {
	return n
}

func newTestWriter() *testWriter {
	return &testWriter{
		H: make(http.Header),
	}
}

type testWriter struct {
	H          http.Header
	Body       []byte
	StatusCode int
}

func (tw *testWriter) Header() http.Header {
	return tw.H
}

func (tw *testWriter) Write(body []byte) (int, error) {
	tw.Body = body
	return len(tw.Body), nil
}

func (tw *testWriter) WriteHeader(code int) {
	tw.StatusCode = code
}

func line(s string) string {
	return strings.Replace(s, "\n", "", -1)
}
<|endoftext|>