<commit_before>// Copyright 2016 PingCAP, 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package executor

import (
	"bytes"
	"fmt"
	"sort"
	"strings"

	"github.com/juju/errors"
	"github.com/pingcap/tidb/ast"
	"github.com/pingcap/tidb/column"
	"github.com/pingcap/tidb/context"
	"github.com/pingcap/tidb/infoschema"
	"github.com/pingcap/tidb/model"
	"github.com/pingcap/tidb/mysql"
	"github.com/pingcap/tidb/privilege"
	"github.com/pingcap/tidb/sessionctx/variable"
	"github.com/pingcap/tidb/table"
	"github.com/pingcap/tidb/util/charset"
	"github.com/pingcap/tidb/util/types"
)

// ShowExec represents a show executor.
type ShowExec struct {
	Tp     ast.ShowStmtType // Databases/Tables/Columns/....
	DBName model.CIStr
	Table  *ast.TableName  // Used for showing columns.
	Column *ast.ColumnName // Used for `desc table column`.
	Flag   int             // Some flag parsed from sql, such as FULL.
	Full   bool
	User   string // Used for show grants.

	// Used by show variables
	GlobalScope bool

	fields []*ast.ResultField
	ctx    context.Context
	is     infoschema.InfoSchema

	fetched bool
	rows    []*Row
	cursor  int
}

// Fields implements Executor Fields interface.
func (e *ShowExec) Fields() []*ast.ResultField {
	return e.fields
}

// Next implements Execution Next interface.
func (e *ShowExec) Next() (*Row, error) {
	if e.rows == nil {
		err := e.fetchAll()
		if err != nil {
			return nil, errors.Trace(err)
		}
	}
	if e.cursor >= len(e.rows) {
		return nil, nil
	}
	row := e.rows[e.cursor]
	for i, field := range e.fields {
		field.Expr.SetValue(row.Data[i])
	}
	e.cursor++
	return row, nil
}

func (e *ShowExec) fetchAll() error {
	switch e.Tp {
	case ast.ShowCharset:
		return e.fetchShowCharset()
	case ast.ShowCollation:
		return e.fetchShowCollation()
	case ast.ShowColumns:
		return e.fetchShowColumns()
	case ast.ShowCreateTable:
		return e.fetchShowCreateTable()
	case ast.ShowDatabases:
		return e.fetchShowDatabases()
	case ast.ShowEngines:
		return e.fetchShowEngines()
	case ast.ShowGrants:
		return e.fetchShowGrants()
	case ast.ShowIndex:
		return e.fetchShowIndex()
	case ast.ShowProcedureStatus:
		return e.fetchShowProcedureStatus()
	case ast.ShowStatus:
		return e.fetchShowStatus()
	case ast.ShowTables:
		return e.fetchShowTables()
	case ast.ShowTableStatus:
		return e.fetchShowTableStatus()
	case ast.ShowTriggers:
		return e.fetchShowTriggers()
	case ast.ShowVariables:
		return e.fetchShowVariables()
	case ast.ShowWarnings:
		// empty result
	}
	return nil
}

func (e *ShowExec) fetchShowEngines() error {
	row := &Row{
		Data: []interface{}{"InnoDB", "DEFAULT", "Supports transactions, row-level locking, and foreign keys", "YES", "YES", "YES"},
	}
	e.rows = append(e.rows, row)
	return nil
}

func (e *ShowExec) fetchShowDatabases() error {
	dbs := e.is.AllSchemaNames()
	// TODO: let information_schema be the first database
	sort.Strings(dbs)
	for _, d := range dbs {
		e.rows = append(e.rows, &Row{Data: []interface{}{d}})
	}
	return nil
}

func (e *ShowExec) fetchShowTables() error {
	if !e.is.SchemaExists(e.DBName) {
		return errors.Errorf("Can not find DB: %s", e.DBName)
	}
	// sort for tables
	var tableNames []string
	for _, v := range e.is.SchemaTables(e.DBName) {
		tableNames = append(tableNames, v.TableName().L)
	}
	sort.Strings(tableNames)
	for _, v := range tableNames {
		data := []interface{}{v}
		if e.Full {
			// TODO: support "VIEW" later if we have supported view feature.
			// now, just use "BASE TABLE".
			data = append(data, "BASE TABLE")
		}
		e.rows = append(e.rows, &Row{Data: data})
	}
	return nil
}

func (e *ShowExec) fetchShowTableStatus() error {
	if !e.is.SchemaExists(e.DBName) {
		return errors.Errorf("Can not find DB: %s", e.DBName)
	}

	// sort for tables
	var tableNames []string
	for _, v := range e.is.SchemaTables(e.DBName) {
		tableNames = append(tableNames, v.TableName().L)
	}
	sort.Strings(tableNames)

	for _, v := range tableNames {
		now := mysql.CurrentTime(mysql.TypeDatetime)
		data := []interface{}{
			v, "InnoDB", "10", "Compact", 100, 100,
			100, 100, 100, 100, 100,
			now, now, now, "utf8_general_ci", "",
			"", "",
		}
		e.rows = append(e.rows, &Row{Data: data})
	}
	return nil
}

func (e *ShowExec) fetchShowColumns() error {
	tb, err := e.getTable()
	if err != nil {
		return errors.Trace(err)
	}
	cols := tb.Cols()
	for _, col := range cols {
		if e.Column != nil && e.Column.Name.L != col.Name.L {
			continue
		}

		desc := column.NewColDesc(col)

		// The FULL keyword causes the output to include the column collation and comments,
		// as well as the privileges you have for each column.
		row := &Row{}
		if e.Full {
			row.Data = []interface{}{
				desc.Field,
				desc.Type,
				desc.Collation,
				desc.Null,
				desc.Key,
				desc.DefaultValue,
				desc.Extra,
				desc.Privileges,
				desc.Comment,
			}
		} else {
			row.Data = []interface{}{
				desc.Field,
				desc.Type,
				desc.Null,
				desc.Key,
				desc.DefaultValue,
				desc.Extra,
			}
		}
		e.rows = append(e.rows, row)
	}
	return nil
}

func (e *ShowExec) fetchShowIndex() error {
	tb, err := e.getTable()
	if err != nil {
		return errors.Trace(err)
	}
	for _, idx := range tb.Indices() {
		for i, col := range idx.Columns {
			nonUniq := 1
			if idx.Unique {
				nonUniq = 0
			}
			var subPart interface{}
			if col.Length != types.UnspecifiedLength {
				subPart = col.Length
			}
			data := []interface{}{
				tb.TableName().O, // Table
				nonUniq,          // Non_unique
				idx.Name.O,       // Key_name
				i + 1,            // Seq_in_index
				col.Name.O,       // Column_name
				"utf8_bin",       // Colation
				0,                // Cardinality
				subPart,          // Sub_part
				nil,              // Packed
				"YES",            // Null
				"BTREE",          // Index_type
				"",               // Comment
				"",               // Index_comment
			}
			e.rows = append(e.rows, &Row{Data: data})
		}
	}
	return nil
}

func (e *ShowExec) fetchShowCharset() error {
	// See: http://dev.mysql.com/doc/refman/5.7/en/show-character-set.html
	descs := charset.GetAllCharsets()
	for _, desc := range descs {
		row := &Row{
			Data: []interface{}{desc.Name, desc.Desc, desc.DefaultCollation, desc.Maxlen},
		}
		e.rows = append(e.rows, row)
	}
	return nil
}

func (e *ShowExec) fetchShowVariables() error {
	sessionVars := variable.GetSessionVars(e.ctx)
	globalVars := variable.GetGlobalVarAccessor(e.ctx)
	for _, v := range variable.SysVars {
		var err error
		var value string
		if !e.GlobalScope {
			// Try to get Session Scope variable value first.
			sv, ok := sessionVars.Systems[v.Name]
			if ok {
				value = sv
			} else {
				// If session scope variable is not set, get the global scope value.
				value, err = globalVars.GetGlobalSysVar(e.ctx, v.Name)
				if err != nil {
					return errors.Trace(err)
				}
			}
		} else {
			value, err = globalVars.GetGlobalSysVar(e.ctx, v.Name)
			if err != nil {
				return errors.Trace(err)
			}
		}
		row := &Row{Data: []interface{}{v.Name, value}}
		e.rows = append(e.rows, row)
	}
	return nil
}

func (e *ShowExec) fetchShowStatus() error {
	statusVars, err := variable.GetStatusVars()
	if err != nil {
		return errors.Trace(err)
	}
	for status, v := range statusVars {
		if e.GlobalScope && v.Scope == variable.ScopeSession {
			continue
		}
		value, err := types.ToString(v.Value)
		if err != nil {
			return errors.Trace(err)
		}
		row := &Row{Data: []interface{}{status, value}}
		e.rows = append(e.rows, row)
	}
	return nil
}

func (e *ShowExec) fetchShowCreateTable() error {
	tb, err := e.getTable()
	if err != nil {
		return errors.Trace(err)
	}

	// TODO: let the result more like MySQL.
	var buf bytes.Buffer
	buf.WriteString(fmt.Sprintf("CREATE TABLE `%s` (\n", tb.TableName().O))
	for i, col := range tb.Cols() {
		buf.WriteString(fmt.Sprintf("  `%s` %s", col.Name.O, col.GetTypeDesc()))
		if mysql.HasAutoIncrementFlag(col.Flag) {
			buf.WriteString(" NOT NULL AUTO_INCREMENT")
		} else {
			if mysql.HasNotNullFlag(col.Flag) {
				buf.WriteString(" NOT NULL")
			}
			switch col.DefaultValue {
			case nil:
				buf.WriteString(" DEFAULT NULL")
			case "CURRENT_TIMESTAMP":
				buf.WriteString(" DEFAULT CURRENT_TIMESTAMP")
			default:
				buf.WriteString(fmt.Sprintf(" DEFAULT '%v'", col.DefaultValue))
			}

			if mysql.HasOnUpdateNowFlag(col.Flag) {
				buf.WriteString(" ON UPDATE CURRENT_TIMESTAMP")
			}
		}
		if i != len(tb.Cols())-1 {
			buf.WriteString(",\n")
		}
	}

	if len(tb.Indices()) > 0 {
		buf.WriteString(",\n")
	}

	for i, idx := range tb.Indices() {
		if idx.Primary {
			buf.WriteString("  PRIMARY KEY ")
		} else if idx.Unique {
			buf.WriteString(fmt.Sprintf("  UNIQUE KEY `%s` ", idx.Name.O))
		} else {
			buf.WriteString(fmt.Sprintf("  KEY `%s` ", idx.Name.O))
		}

		cols := make([]string, 0, len(idx.Columns))
		for _, c := range idx.Columns {
			cols = append(cols, c.Name.O)
		}
		buf.WriteString(fmt.Sprintf("(`%s`)", strings.Join(cols, "`,`")))
		if i != len(tb.Indices())-1 {
			buf.WriteString(",\n")
		}
	}
	buf.WriteString("\n")

	buf.WriteString(") ENGINE=InnoDB")
	if s := tb.Meta().Charset; len(s) > 0 {
		buf.WriteString(fmt.Sprintf(" DEFAULT CHARSET=%s", s))
	} else {
		buf.WriteString(" DEFAULT CHARSET=latin1")
	}

	data := []interface{}{
		tb.TableName().O,
		buf.String(),
	}
	e.rows = append(e.rows, &Row{Data: data})
	return nil
}

func (e *ShowExec) fetchShowCollation() error {
	collations := charset.GetCollations()
	for _, v := range collations {
		isDefault := ""
		if v.IsDefault {
			isDefault = "Yes"
		}
		row := &Row{Data: []interface{}{v.Name, v.CharsetName, v.ID, isDefault, "Yes", 1}}
		e.rows = append(e.rows, row)
	}
	return nil
}

func (e *ShowExec) fetchShowGrants() error {
	// Get checker
	checker := privilege.GetPrivilegeChecker(e.ctx)
	if checker == nil {
		return errors.New("Miss privilege checker!")
	}
	gs, err := checker.ShowGrants(e.ctx, e.User)
	if err != nil {
		return errors.Trace(err)
	}
	for _, g := range gs {
		data := []interface{}{g}
		e.rows = append(e.rows, &Row{Data: data})
	}
	return nil
}

func (e *ShowExec) fetchShowTriggers() error {
	return nil
}

func (e *ShowExec) fetchShowProcedureStatus() error {
	return nil
}

func (e *ShowExec) getTable() (table.Table, error) {
	if e.Table == nil {
		return nil, errors.New("table not found")
	}
	tb, ok := e.is.TableByID(e.Table.TableInfo.ID)
	if !ok {
		return nil, errors.Errorf("table %s not found", e.Table.Name)
	}
	return tb, nil
}

// Close implements Executor Close interface.
func (e *ShowExec) Close() error {
	return nil
}
<commit_msg>*:fix build after Table interface change.<commit_after>// Copyright 2016 PingCAP, 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package executor

import (
	"bytes"
	"fmt"
	"sort"
	"strings"

	"github.com/juju/errors"
	"github.com/pingcap/tidb/ast"
	"github.com/pingcap/tidb/column"
	"github.com/pingcap/tidb/context"
	"github.com/pingcap/tidb/infoschema"
	"github.com/pingcap/tidb/model"
	"github.com/pingcap/tidb/mysql"
	"github.com/pingcap/tidb/privilege"
	"github.com/pingcap/tidb/sessionctx/variable"
	"github.com/pingcap/tidb/table"
	"github.com/pingcap/tidb/util/charset"
	"github.com/pingcap/tidb/util/types"
)

// ShowExec represents a show executor.
type ShowExec struct {
	Tp     ast.ShowStmtType // Databases/Tables/Columns/....
	DBName model.CIStr
	Table  *ast.TableName  // Used for showing columns.
	Column *ast.ColumnName // Used for `desc table column`.
	Flag   int             // Some flag parsed from sql, such as FULL.
	Full   bool
	User   string // Used for show grants.

	// Used by show variables
	GlobalScope bool

	fields []*ast.ResultField
	ctx    context.Context
	is     infoschema.InfoSchema

	fetched bool
	rows    []*Row
	cursor  int
}

// Fields implements Executor Fields interface.
func (e *ShowExec) Fields() []*ast.ResultField {
	return e.fields
}

// Next implements Execution Next interface.
func (e *ShowExec) Next() (*Row, error) {
	if e.rows == nil {
		err := e.fetchAll()
		if err != nil {
			return nil, errors.Trace(err)
		}
	}
	if e.cursor >= len(e.rows) {
		return nil, nil
	}
	row := e.rows[e.cursor]
	for i, field := range e.fields {
		field.Expr.SetValue(row.Data[i])
	}
	e.cursor++
	return row, nil
}

func (e *ShowExec) fetchAll() error {
	switch e.Tp {
	case ast.ShowCharset:
		return e.fetchShowCharset()
	case ast.ShowCollation:
		return e.fetchShowCollation()
	case ast.ShowColumns:
		return e.fetchShowColumns()
	case ast.ShowCreateTable:
		return e.fetchShowCreateTable()
	case ast.ShowDatabases:
		return e.fetchShowDatabases()
	case ast.ShowEngines:
		return e.fetchShowEngines()
	case ast.ShowGrants:
		return e.fetchShowGrants()
	case ast.ShowIndex:
		return e.fetchShowIndex()
	case ast.ShowProcedureStatus:
		return e.fetchShowProcedureStatus()
	case ast.ShowStatus:
		return e.fetchShowStatus()
	case ast.ShowTables:
		return e.fetchShowTables()
	case ast.ShowTableStatus:
		return e.fetchShowTableStatus()
	case ast.ShowTriggers:
		return e.fetchShowTriggers()
	case ast.ShowVariables:
		return e.fetchShowVariables()
	case ast.ShowWarnings:
		// empty result
	}
	return nil
}

func (e *ShowExec) fetchShowEngines() error {
	row := &Row{
		Data: []interface{}{"InnoDB", "DEFAULT", "Supports transactions, row-level locking, and foreign keys", "YES", "YES", "YES"},
	}
	e.rows = append(e.rows, row)
	return nil
}

func (e *ShowExec) fetchShowDatabases() error {
	dbs := e.is.AllSchemaNames()
	// TODO: let information_schema be the first database
	sort.Strings(dbs)
	for _, d := range dbs {
		e.rows = append(e.rows, &Row{Data: []interface{}{d}})
	}
	return nil
}

func (e *ShowExec) fetchShowTables() error {
	if !e.is.SchemaExists(e.DBName) {
		return errors.Errorf("Can not find DB: %s", e.DBName)
	}
	// sort for tables
	var tableNames []string
	for _, v := range e.is.SchemaTables(e.DBName) {
		tableNames = append(tableNames, v.Meta().Name.L)
	}
	sort.Strings(tableNames)
	for _, v := range tableNames {
		data := []interface{}{v}
		if e.Full {
			// TODO: support "VIEW" later if we have supported view feature.
			// now, just use "BASE TABLE".
			data = append(data, "BASE TABLE")
		}
		e.rows = append(e.rows, &Row{Data: data})
	}
	return nil
}

func (e *ShowExec) fetchShowTableStatus() error {
	if !e.is.SchemaExists(e.DBName) {
		return errors.Errorf("Can not find DB: %s", e.DBName)
	}

	// sort for tables
	var tableNames []string
	for _, v := range e.is.SchemaTables(e.DBName) {
		tableNames = append(tableNames, v.Meta().Name.L)
	}
	sort.Strings(tableNames)

	for _, v := range tableNames {
		now := mysql.CurrentTime(mysql.TypeDatetime)
		data := []interface{}{
			v, "InnoDB", "10", "Compact", 100, 100,
			100, 100, 100, 100, 100,
			now, now, now, "utf8_general_ci", "",
			"", "",
		}
		e.rows = append(e.rows, &Row{Data: data})
	}
	return nil
}

func (e *ShowExec) fetchShowColumns() error {
	tb, err := e.getTable()
	if err != nil {
		return errors.Trace(err)
	}
	cols := tb.Cols()
	for _, col := range cols {
		if e.Column != nil && e.Column.Name.L != col.Name.L {
			continue
		}

		desc := column.NewColDesc(col)

		// The FULL keyword causes the output to include the column collation and comments,
		// as well as the privileges you have for each column.
		row := &Row{}
		if e.Full {
			row.Data = []interface{}{
				desc.Field,
				desc.Type,
				desc.Collation,
				desc.Null,
				desc.Key,
				desc.DefaultValue,
				desc.Extra,
				desc.Privileges,
				desc.Comment,
			}
		} else {
			row.Data = []interface{}{
				desc.Field,
				desc.Type,
				desc.Null,
				desc.Key,
				desc.DefaultValue,
				desc.Extra,
			}
		}
		e.rows = append(e.rows, row)
	}
	return nil
}

func (e *ShowExec) fetchShowIndex() error {
	tb, err := e.getTable()
	if err != nil {
		return errors.Trace(err)
	}
	for _, idx := range tb.Indices() {
		for i, col := range idx.Columns {
			nonUniq := 1
			if idx.Unique {
				nonUniq = 0
			}
			var subPart interface{}
			if col.Length != types.UnspecifiedLength {
				subPart = col.Length
			}
			data := []interface{}{
				tb.Meta().Name.O, // Table
				nonUniq,          // Non_unique
				idx.Name.O,       // Key_name
				i + 1,            // Seq_in_index
				col.Name.O,       // Column_name
				"utf8_bin",       // Colation
				0,                // Cardinality
				subPart,          // Sub_part
				nil,              // Packed
				"YES",            // Null
				"BTREE",          // Index_type
				"",               // Comment
				"",               // Index_comment
			}
			e.rows = append(e.rows, &Row{Data: data})
		}
	}
	return nil
}

func (e *ShowExec) fetchShowCharset() error {
	// See: http://dev.mysql.com/doc/refman/5.7/en/show-character-set.html
	descs := charset.GetAllCharsets()
	for _, desc := range descs {
		row := &Row{
			Data: []interface{}{desc.Name, desc.Desc, desc.DefaultCollation, desc.Maxlen},
		}
		e.rows = append(e.rows, row)
	}
	return nil
}

func (e *ShowExec) fetchShowVariables() error {
	sessionVars := variable.GetSessionVars(e.ctx)
	globalVars := variable.GetGlobalVarAccessor(e.ctx)
	for _, v := range variable.SysVars {
		var err error
		var value string
		if !e.GlobalScope {
			// Try to get Session Scope variable value first.
			sv, ok := sessionVars.Systems[v.Name]
			if ok {
				value = sv
			} else {
				// If session scope variable is not set, get the global scope value.
				value, err = globalVars.GetGlobalSysVar(e.ctx, v.Name)
				if err != nil {
					return errors.Trace(err)
				}
			}
		} else {
			value, err = globalVars.GetGlobalSysVar(e.ctx, v.Name)
			if err != nil {
				return errors.Trace(err)
			}
		}
		row := &Row{Data: []interface{}{v.Name, value}}
		e.rows = append(e.rows, row)
	}
	return nil
}

func (e *ShowExec) fetchShowStatus() error {
	statusVars, err := variable.GetStatusVars()
	if err != nil {
		return errors.Trace(err)
	}
	for status, v := range statusVars {
		if e.GlobalScope && v.Scope == variable.ScopeSession {
			continue
		}
		value, err := types.ToString(v.Value)
		if err != nil {
			return errors.Trace(err)
		}
		row := &Row{Data: []interface{}{status, value}}
		e.rows = append(e.rows, row)
	}
	return nil
}

func (e *ShowExec) fetchShowCreateTable() error {
	tb, err := e.getTable()
	if err != nil {
		return errors.Trace(err)
	}

	// TODO: let the result more like MySQL.
	var buf bytes.Buffer
	buf.WriteString(fmt.Sprintf("CREATE TABLE `%s` (\n", tb.Meta().Name.O))
	for i, col := range tb.Cols() {
		buf.WriteString(fmt.Sprintf("  `%s` %s", col.Name.O, col.GetTypeDesc()))
		if mysql.HasAutoIncrementFlag(col.Flag) {
			buf.WriteString(" NOT NULL AUTO_INCREMENT")
		} else {
			if mysql.HasNotNullFlag(col.Flag) {
				buf.WriteString(" NOT NULL")
			}
			switch col.DefaultValue {
			case nil:
				buf.WriteString(" DEFAULT NULL")
			case "CURRENT_TIMESTAMP":
				buf.WriteString(" DEFAULT CURRENT_TIMESTAMP")
			default:
				buf.WriteString(fmt.Sprintf(" DEFAULT '%v'", col.DefaultValue))
			}

			if mysql.HasOnUpdateNowFlag(col.Flag) {
				buf.WriteString(" ON UPDATE CURRENT_TIMESTAMP")
			}
		}
		if i != len(tb.Cols())-1 {
			buf.WriteString(",\n")
		}
	}

	if len(tb.Indices()) > 0 {
		buf.WriteString(",\n")
	}

	for i, idx := range tb.Indices() {
		if idx.Primary {
			buf.WriteString("  PRIMARY KEY ")
		} else if idx.Unique {
			buf.WriteString(fmt.Sprintf("  UNIQUE KEY `%s` ", idx.Name.O))
		} else {
			buf.WriteString(fmt.Sprintf("  KEY `%s` ", idx.Name.O))
		}

		cols := make([]string, 0, len(idx.Columns))
		for _, c := range idx.Columns {
			cols = append(cols, c.Name.O)
		}
		buf.WriteString(fmt.Sprintf("(`%s`)", strings.Join(cols, "`,`")))
		if i != len(tb.Indices())-1 {
			buf.WriteString(",\n")
		}
	}
	buf.WriteString("\n")

	buf.WriteString(") ENGINE=InnoDB")
	if s := tb.Meta().Charset; len(s) > 0 {
		buf.WriteString(fmt.Sprintf(" DEFAULT CHARSET=%s", s))
	} else {
		buf.WriteString(" DEFAULT CHARSET=latin1")
	}

	data := []interface{}{
		tb.Meta().Name.O,
		buf.String(),
	}
	e.rows = append(e.rows, &Row{Data: data})
	return nil
}

func (e *ShowExec) fetchShowCollation() error {
	collations := charset.GetCollations()
	for _, v := range collations {
		isDefault := ""
		if v.IsDefault {
			isDefault = "Yes"
		}
		row := &Row{Data: []interface{}{v.Name, v.CharsetName, v.ID, isDefault, "Yes", 1}}
		e.rows = append(e.rows, row)
	}
	return nil
}

func (e *ShowExec) fetchShowGrants() error {
	// Get checker
	checker := privilege.GetPrivilegeChecker(e.ctx)
	if checker == nil {
		return errors.New("Miss privilege checker!")
	}
	gs, err := checker.ShowGrants(e.ctx, e.User)
	if err != nil {
		return errors.Trace(err)
	}
	for _, g := range gs {
		data := []interface{}{g}
		e.rows = append(e.rows, &Row{Data: data})
	}
	return nil
}

func (e *ShowExec) fetchShowTriggers() error {
	return nil
}

func (e *ShowExec) fetchShowProcedureStatus() error {
	return nil
}

func (e *ShowExec) getTable() (table.Table, error) {
	if e.Table == nil {
		return nil, errors.New("table not found")
	}
	tb, ok := e.is.TableByID(e.Table.TableInfo.ID)
	if !ok {
		return nil, errors.Errorf("table %s not found", e.Table.Name)
	}
	return tb, nil
}

// Close implements Executor Close interface.
func (e *ShowExec) Close() error {
	return nil
}
<|endoftext|>


<commit_before>package inigo_test

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	"github.com/cloudfoundry-incubator/inigo/executor_runner"
	"github.com/cloudfoundry-incubator/inigo/inigo_server"
	"github.com/cloudfoundry-incubator/inigo/loggredile"
	Bbs "github.com/cloudfoundry-incubator/runtime-schema/bbs"
	"github.com/cloudfoundry-incubator/runtime-schema/models"
	"github.com/cloudfoundry-incubator/runtime-schema/models/factories"
	"github.com/cloudfoundry/loggregatorlib/logmessage"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/vito/cmdtest/matchers"
)

var _ = Describe("Executor", func() {
	var bbs *Bbs.BBS

	BeforeEach(func() {
		bbs = Bbs.New(etcdRunner.Adapter())
	})

	Describe("starting without a snaphsot", func() {
		It("should come up, just fine", func() {
			executorRunner.Start()
			executorRunner.Stop()
		})
	})

	Describe("when starting with invalid memory/disk", func() {
		It("should exit with failure", func() {
			executorRunner.StartWithoutCheck(executor_runner.Config{MemoryMB: -1, DiskMB: -1, SnapshotFile: "/tmp/i_dont_exist"})
			Ω(executorRunner.Session).Should(SayWithTimeout("valid memory and disk capacity must be specified", time.Second))
			Ω(executorRunner.Session).Should(ExitWith(1))
		})
	})

	Describe("when restarted with running tasks", func() {
		var tmpdir string
		var registrySnapshotFile string
		var executorConfig executor_runner.Config

		BeforeEach(func() {
			tmpdir, err := ioutil.TempDir(os.TempDir(), "executor-registry")
			Ω(err).ShouldNot(HaveOccurred())

			registrySnapshotFile = filepath.Join(tmpdir, "snapshot.json")

			executorConfig = executor_runner.Config{
				MemoryMB:     1024,
				DiskMB:       1024,
				SnapshotFile: registrySnapshotFile,
			}

			executorRunner.Start(executorConfig)

			existingGuid := factories.GenerateGuid()

			existingRunOnce := factories.BuildRunOnceWithRunAction(
				1024,
				1024,
				inigoserver.CurlCommand(existingGuid)+"; sleep 60",
			)

			bbs.DesireRunOnce(existingRunOnce)

			Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(existingGuid))

			executorRunner.Stop()
		})

		AfterEach(func() {
			executorRunner.Stop()

			os.RemoveAll(tmpdir)
		})

		Context("when the snapshot is corrupted", func() {
			It("should exit with failure", func() {
				file, err := ioutil.TempFile(os.TempDir(), "executor-invalid-snapshot")
				Ω(err).ShouldNot(HaveOccurred())

				_, err = file.Write([]byte("ß"))
				Ω(err).ShouldNot(HaveOccurred())

				executorConfig.SnapshotFile = file.Name()

				executorRunner.StartWithoutCheck(executorConfig)

				Ω(executorRunner.Session).Should(SayWithTimeout("corrupt registry", time.Second))
				Ω(executorRunner.Session).Should(ExitWith(1))
			})
		})
	})

	Describe("Heartbeating", func() {
		It("should heartbeat its presence", func() {
			executorRunner.Start()

			Eventually(func() interface{} {
				executors, _ := bbs.GetAllExecutors()
				return executors
			}).Should(HaveLen(1))
		})
	})

	Describe("Resource limits", func() {
		BeforeEach(func() {
			executorRunner.Start()
		})

		It("should only pick up tasks if it has capacity", func() {
			firstGuyGuid := factories.GenerateGuid()
			secondGuyGuid := factories.GenerateGuid()
			firstGuyRunOnce := factories.BuildRunOnceWithRunAction(1024, 1024, inigoserver.CurlCommand(firstGuyGuid)+"; sleep 5")
			bbs.DesireRunOnce(firstGuyRunOnce)

			Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(firstGuyGuid))

			secondGuyRunOnce := factories.BuildRunOnceWithRunAction(1024, 1024, inigoserver.CurlCommand(secondGuyGuid))
			bbs.DesireRunOnce(secondGuyRunOnce)

			Consistently(inigoserver.ReportingGuids, 2.0).ShouldNot(ContainElement(secondGuyGuid))
		})
	})

	Describe("Stack", func() {
		BeforeEach(func() {
			executorRunner.Start(executor_runner.Config{Stack: "penguin"})
		})

		It("should only pick up tasks if the stacks match", func() {
			matchingGuid := factories.GenerateGuid()
			matchingRunOnce := factories.BuildRunOnceWithRunAction(100, 100, inigoserver.CurlCommand(matchingGuid)+"; sleep 10")
			matchingRunOnce.Stack = "penguin"

			nonMatchingGuid := factories.GenerateGuid()
			nonMatchingRunOnce := factories.BuildRunOnceWithRunAction(100, 100, inigoserver.CurlCommand(nonMatchingGuid)+"; sleep 10")
			nonMatchingRunOnce.Stack = "lion"

			bbs.DesireRunOnce(matchingRunOnce)
			bbs.DesireRunOnce(nonMatchingRunOnce)

			Consistently(inigoserver.ReportingGuids, 2.0).ShouldNot(ContainElement(nonMatchingGuid), "Did not expect to see this app running, as it has the wrong stack.")
			Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(matchingGuid))
		})
	})

	Describe("Running a command", func() {
		var guid string
		BeforeEach(func() {
			executorRunner.Start()
			guid = factories.GenerateGuid()
		})

		It("should run the command with the provided environment", func() {
			env := [][]string{
				{"FOO", "BAR"},
				{"BAZ", "WIBBLE"},
				{"FOO", "$FOO-$BAZ"},
			}
			runOnce := models.RunOnce{
				Guid:     factories.GenerateGuid(),
				MemoryMB: 1024,
				DiskMB:   1024,
				Actions: []models.ExecutorAction{
					{Action: models.RunAction{Script: `echo $FOO > out`, Env: env}},
					{Action: models.UploadAction{From: "out", To: inigoserver.UploadUrl("out")}},
					{Action: models.RunAction{Script: inigoserver.CurlCommand(guid)}},
				},
			}

			bbs.DesireRunOnce(runOnce)

			Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(guid))
			Ω(inigoserver.DownloadFileString("out")).Should(Equal("BAR-WIBBLE\n"))

			Eventually(bbs.GetAllCompletedRunOnces, 5.0).Should(HaveLen(1))
			runOnces, _ := bbs.GetAllCompletedRunOnces()
			Ω(runOnces[0].Failed).Should(BeFalse())
		})

		Context("when the command exceeds its memory limit", func() {
			var otherGuid string

			It("should fail the RunOnce", func() {
				otherGuid = factories.GenerateGuid()
				runOnce := models.RunOnce{
					Guid:     factories.GenerateGuid(),
					MemoryMB: 10,
					DiskMB:   1024,
					Actions: []models.ExecutorAction{
						{Action: models.RunAction{Script: inigoserver.CurlCommand(guid)}},
						{Action: models.RunAction{Script: `ruby -e 'arr = "m"*1024*1024*100'`}},
						{Action: models.RunAction{Script: inigoserver.CurlCommand(otherGuid)}},
					},
				}

				bbs.DesireRunOnce(runOnce)

				Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(guid))

				Eventually(bbs.GetAllCompletedRunOnces, 5.0).Should(HaveLen(1))
				runOnces, _ := bbs.GetAllCompletedRunOnces()
				Ω(runOnces[0].Failed).Should(BeTrue())
				Ω(runOnces[0].FailureReason).Should(ContainSubstring("137"))

				Ω(inigoserver.ReportingGuids()).ShouldNot(ContainElement(otherGuid))
			})
		})

		Context("when the command exceeds its file descriptor limit", func() {
			It("should fail the RunOnce", func() {
				runOnce := models.RunOnce{
					Guid:            factories.GenerateGuid(),
					MemoryMB:        10,
					DiskMB:          1024,
					FileDescriptors: 1,
					Actions: []models.ExecutorAction{
						{Action: models.RunAction{Script: `ruby -e '10.times.each { |x| File.open("#{x}","w") }'`}},
					},
				}

				bbs.DesireRunOnce(runOnce)

				Eventually(bbs.GetAllCompletedRunOnces, 5.0).Should(HaveLen(1))
				runOnces, _ := bbs.GetAllCompletedRunOnces()
				Ω(runOnces[0].Failed).Should(BeTrue())
				Ω(runOnces[0].FailureReason).Should(ContainSubstring("127"))
			})
		})

		Context("when the command times out", func() {
			It("should fail the RunOnce", func() {
				runOnce := models.RunOnce{
					Guid:     factories.GenerateGuid(),
					MemoryMB: 1024,
					DiskMB:   1024,
					Actions: []models.ExecutorAction{
						{Action: models.RunAction{Script: inigoserver.CurlCommand(guid)}},
						{Action: models.RunAction{Script: `sleep 0.8`, Timeout: 500 * time.Millisecond}},
					},
				}

				bbs.DesireRunOnce(runOnce)

				Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(guid))
				Eventually(bbs.GetAllCompletedRunOnces, 5.0).Should(HaveLen(1))
				runOnces, _ := bbs.GetAllCompletedRunOnces()
				Ω(runOnces[0].Failed).Should(BeTrue())
				Ω(runOnces[0].FailureReason).Should(ContainSubstring("timed out"))
			})
		})
	})

	Describe("Running a downloaded file", func() {
		var guid string
		BeforeEach(func() {
			executorRunner.Start()

			guid = factories.GenerateGuid()
			inigoserver.UploadFileString("curling.sh", inigoserver.CurlCommand(guid))
		})

		It("downloads the file", func() {
			runOnce := models.RunOnce{
				Guid:     factories.GenerateGuid(),
				MemoryMB: 1024,
				DiskMB:   1024,
				Actions: []models.ExecutorAction{
					{Action: models.DownloadAction{From: inigoserver.DownloadUrl("curling.sh"), To: "curling.sh", Extract: false}},
					{Action: models.RunAction{Script: "bash curling.sh"}},
				},
			}

			bbs.DesireRunOnce(runOnce)

			Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(guid))
		})
	})

	Describe("Uploading a file", func() {
		var guid string
		BeforeEach(func() {
			executorRunner.Start()

			guid = factories.GenerateGuid()
		})

		It("uploads the file", func() {
			runOnce := models.RunOnce{
				Guid:     factories.GenerateGuid(),
				MemoryMB: 1024,
				DiskMB:   1024,
				Actions: []models.ExecutorAction{
					{Action: models.RunAction{Script: `echo "tasty thingy" > thingy`}},
					{Action: models.UploadAction{From: "thingy", To: inigoserver.UploadUrl("thingy")}},
					{Action: models.RunAction{Script: inigoserver.CurlCommand(guid)}},
				},
			}

			bbs.DesireRunOnce(runOnce)

			Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(guid))
			Ω(inigoserver.DownloadFileString("thingy")).Should(Equal("tasty thingy\n"))
		})
	})

	Describe("Fetching results", func() {
		BeforeEach(func() {
			executorRunner.Start()
		})

		It("should fetch the contents of the requested file and provide the content in the completed RunOnce", func() {
			runOnce := models.RunOnce{
				Guid:     factories.GenerateGuid(),
				MemoryMB: 1024,
				DiskMB:   1024,
				Actions: []models.ExecutorAction{
					{Action: models.RunAction{Script: `echo "tasty thingy" > thingy`}},
					{Action: models.FetchResultAction{File: "thingy"}},
				},
			}

			bbs.DesireRunOnce(runOnce)

			Eventually(bbs.GetAllCompletedRunOnces, 5.0).Should(HaveLen(1))

			runOnces, _ := bbs.GetAllCompletedRunOnces()
			Ω(runOnces[0].Result).Should(Equal("tasty thingy\n"))
		})
	})

	Describe("A RunOnce with logging configured", func() {
		BeforeEach(func() {
			executorRunner.Start()
		})

		It("has its stdout and stderr emitted to Loggregator", func(done Done) {
			logGuid := factories.GenerateGuid()

			messages, stop := loggredile.StreamMessages(
				loggregatorRunner.Config.OutgoingPort,
				"/tail/?app="+logGuid,
			)

			runOnce := factories.BuildRunOnceWithRunAction(
				1024,
				1024,
				"echo out A; echo out B; echo out C; echo err A 1>&2; echo err B 1>&2; echo err C 1>&2",
			)
			runOnce.Log.Guid = logGuid
			runOnce.Log.SourceName = "APP"

			bbs.DesireRunOnce(runOnce)

			outStream := []string{}
			errStream := []string{}

			for i := 0; i < 6; i++ {
				message := <-messages
				switch message.GetMessageType() {
				case logmessage.LogMessage_OUT:
					outStream = append(outStream, string(message.GetMessage()))
				case logmessage.LogMessage_ERR:
					errStream = append(errStream, string(message.GetMessage()))
				}
			}

			Ω(outStream).Should(Equal([]string{"out A", "out B", "out C"}))
			Ω(errStream).Should(Equal([]string{"err A", "err B", "err C"}))

			close(stop)
			close(done)
		}, 10.0)
	})
})
<commit_msg>137 => out of memory<commit_after>package inigo_test

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	"github.com/cloudfoundry-incubator/inigo/executor_runner"
	"github.com/cloudfoundry-incubator/inigo/inigo_server"
	"github.com/cloudfoundry-incubator/inigo/loggredile"
	Bbs "github.com/cloudfoundry-incubator/runtime-schema/bbs"
	"github.com/cloudfoundry-incubator/runtime-schema/models"
	"github.com/cloudfoundry-incubator/runtime-schema/models/factories"
	"github.com/cloudfoundry/loggregatorlib/logmessage"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/vito/cmdtest/matchers"
)

var _ = Describe("Executor", func() {
	var bbs *Bbs.BBS

	BeforeEach(func() {
		bbs = Bbs.New(etcdRunner.Adapter())
	})

	Describe("starting without a snaphsot", func() {
		It("should come up, just fine", func() {
			executorRunner.Start()
			executorRunner.Stop()
		})
	})

	Describe("when starting with invalid memory/disk", func() {
		It("should exit with failure", func() {
			executorRunner.StartWithoutCheck(executor_runner.Config{MemoryMB: -1, DiskMB: -1, SnapshotFile: "/tmp/i_dont_exist"})
			Ω(executorRunner.Session).Should(SayWithTimeout("valid memory and disk capacity must be specified", time.Second))
			Ω(executorRunner.Session).Should(ExitWith(1))
		})
	})

	Describe("when restarted with running tasks", func() {
		var tmpdir string
		var registrySnapshotFile string
		var executorConfig executor_runner.Config

		BeforeEach(func() {
			tmpdir, err := ioutil.TempDir(os.TempDir(), "executor-registry")
			Ω(err).ShouldNot(HaveOccurred())

			registrySnapshotFile = filepath.Join(tmpdir, "snapshot.json")

			executorConfig = executor_runner.Config{
				MemoryMB:     1024,
				DiskMB:       1024,
				SnapshotFile: registrySnapshotFile,
			}

			executorRunner.Start(executorConfig)

			existingGuid := factories.GenerateGuid()

			existingRunOnce := factories.BuildRunOnceWithRunAction(
				1024,
				1024,
				inigoserver.CurlCommand(existingGuid)+"; sleep 60",
			)

			bbs.DesireRunOnce(existingRunOnce)

			Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(existingGuid))

			executorRunner.Stop()
		})

		AfterEach(func() {
			executorRunner.Stop()

			os.RemoveAll(tmpdir)
		})

		Context("when the snapshot is corrupted", func() {
			It("should exit with failure", func() {
				file, err := ioutil.TempFile(os.TempDir(), "executor-invalid-snapshot")
				Ω(err).ShouldNot(HaveOccurred())

				_, err = file.Write([]byte("ß"))
				Ω(err).ShouldNot(HaveOccurred())

				executorConfig.SnapshotFile = file.Name()

				executorRunner.StartWithoutCheck(executorConfig)

				Ω(executorRunner.Session).Should(SayWithTimeout("corrupt registry", time.Second))
				Ω(executorRunner.Session).Should(ExitWith(1))
			})
		})
	})

	Describe("Heartbeating", func() {
		It("should heartbeat its presence", func() {
			executorRunner.Start()

			Eventually(func() interface{} {
				executors, _ := bbs.GetAllExecutors()
				return executors
			}).Should(HaveLen(1))
		})
	})

	Describe("Resource limits", func() {
		BeforeEach(func() {
			executorRunner.Start()
		})

		It("should only pick up tasks if it has capacity", func() {
			firstGuyGuid := factories.GenerateGuid()
			secondGuyGuid := factories.GenerateGuid()
			firstGuyRunOnce := factories.BuildRunOnceWithRunAction(1024, 1024, inigoserver.CurlCommand(firstGuyGuid)+"; sleep 5")
			bbs.DesireRunOnce(firstGuyRunOnce)

			Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(firstGuyGuid))

			secondGuyRunOnce := factories.BuildRunOnceWithRunAction(1024, 1024, inigoserver.CurlCommand(secondGuyGuid))
			bbs.DesireRunOnce(secondGuyRunOnce)

			Consistently(inigoserver.ReportingGuids, 2.0).ShouldNot(ContainElement(secondGuyGuid))
		})
	})

	Describe("Stack", func() {
		BeforeEach(func() {
			executorRunner.Start(executor_runner.Config{Stack: "penguin"})
		})

		It("should only pick up tasks if the stacks match", func() {
			matchingGuid := factories.GenerateGuid()
			matchingRunOnce := factories.BuildRunOnceWithRunAction(100, 100, inigoserver.CurlCommand(matchingGuid)+"; sleep 10")
			matchingRunOnce.Stack = "penguin"

			nonMatchingGuid := factories.GenerateGuid()
			nonMatchingRunOnce := factories.BuildRunOnceWithRunAction(100, 100, inigoserver.CurlCommand(nonMatchingGuid)+"; sleep 10")
			nonMatchingRunOnce.Stack = "lion"

			bbs.DesireRunOnce(matchingRunOnce)
			bbs.DesireRunOnce(nonMatchingRunOnce)

			Consistently(inigoserver.ReportingGuids, 2.0).ShouldNot(ContainElement(nonMatchingGuid), "Did not expect to see this app running, as it has the wrong stack.")
			Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(matchingGuid))
		})
	})

	Describe("Running a command", func() {
		var guid string
		BeforeEach(func() {
			executorRunner.Start()
			guid = factories.GenerateGuid()
		})

		It("should run the command with the provided environment", func() {
			env := [][]string{
				{"FOO", "BAR"},
				{"BAZ", "WIBBLE"},
				{"FOO", "$FOO-$BAZ"},
			}
			runOnce := models.RunOnce{
				Guid:     factories.GenerateGuid(),
				MemoryMB: 1024,
				DiskMB:   1024,
				Actions: []models.ExecutorAction{
					{Action: models.RunAction{Script: `echo $FOO > out`, Env: env}},
					{Action: models.UploadAction{From: "out", To: inigoserver.UploadUrl("out")}},
					{Action: models.RunAction{Script: inigoserver.CurlCommand(guid)}},
				},
			}

			bbs.DesireRunOnce(runOnce)

			Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(guid))
			Ω(inigoserver.DownloadFileString("out")).Should(Equal("BAR-WIBBLE\n"))

			Eventually(bbs.GetAllCompletedRunOnces, 5.0).Should(HaveLen(1))
			runOnces, _ := bbs.GetAllCompletedRunOnces()
			Ω(runOnces[0].Failed).Should(BeFalse())
		})

		Context("when the command exceeds its memory limit", func() {
			var otherGuid string

			It("should fail the RunOnce", func() {
				otherGuid = factories.GenerateGuid()
				runOnce := models.RunOnce{
					Guid:     factories.GenerateGuid(),
					MemoryMB: 10,
					DiskMB:   1024,
					Actions: []models.ExecutorAction{
						{Action: models.RunAction{Script: inigoserver.CurlCommand(guid)}},
						{Action: models.RunAction{Script: `ruby -e 'arr = "m"*1024*1024*100'`}},
						{Action: models.RunAction{Script: inigoserver.CurlCommand(otherGuid)}},
					},
				}

				bbs.DesireRunOnce(runOnce)

				Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(guid))

				Eventually(bbs.GetAllCompletedRunOnces, 5.0).Should(HaveLen(1))
				runOnces, _ := bbs.GetAllCompletedRunOnces()
				Ω(runOnces[0].Failed).Should(BeTrue())
				Ω(runOnces[0].FailureReason).Should(ContainSubstring("out of memory"))

				Ω(inigoserver.ReportingGuids()).ShouldNot(ContainElement(otherGuid))
			})
		})

		Context("when the command exceeds its file descriptor limit", func() {
			It("should fail the RunOnce", func() {
				runOnce := models.RunOnce{
					Guid:            factories.GenerateGuid(),
					MemoryMB:        10,
					DiskMB:          1024,
					FileDescriptors: 1,
					Actions: []models.ExecutorAction{
						{Action: models.RunAction{Script: `ruby -e '10.times.each { |x| File.open("#{x}","w") }'`}},
					},
				}

				bbs.DesireRunOnce(runOnce)

				Eventually(bbs.GetAllCompletedRunOnces, 5.0).Should(HaveLen(1))
				runOnces, _ := bbs.GetAllCompletedRunOnces()
				Ω(runOnces[0].Failed).Should(BeTrue())
				Ω(runOnces[0].FailureReason).Should(ContainSubstring("127"))
			})
		})

		Context("when the command times out", func() {
			It("should fail the RunOnce", func() {
				runOnce := models.RunOnce{
					Guid:     factories.GenerateGuid(),
					MemoryMB: 1024,
					DiskMB:   1024,
					Actions: []models.ExecutorAction{
						{Action: models.RunAction{Script: inigoserver.CurlCommand(guid)}},
						{Action: models.RunAction{Script: `sleep 0.8`, Timeout: 500 * time.Millisecond}},
					},
				}

				bbs.DesireRunOnce(runOnce)

				Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(guid))
				Eventually(bbs.GetAllCompletedRunOnces, 5.0).Should(HaveLen(1))
				runOnces, _ := bbs.GetAllCompletedRunOnces()
				Ω(runOnces[0].Failed).Should(BeTrue())
				Ω(runOnces[0].FailureReason).Should(ContainSubstring("timed out"))
			})
		})
	})

	Describe("Running a downloaded file", func() {
		var guid string
		BeforeEach(func() {
			executorRunner.Start()

			guid = factories.GenerateGuid()
			inigoserver.UploadFileString("curling.sh", inigoserver.CurlCommand(guid))
		})

		It("downloads the file", func() {
			runOnce := models.RunOnce{
				Guid:     factories.GenerateGuid(),
				MemoryMB: 1024,
				DiskMB:   1024,
				Actions: []models.ExecutorAction{
					{Action: models.DownloadAction{From: inigoserver.DownloadUrl("curling.sh"), To: "curling.sh", Extract: false}},
					{Action: models.RunAction{Script: "bash curling.sh"}},
				},
			}

			bbs.DesireRunOnce(runOnce)

			Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(guid))
		})
	})

	Describe("Uploading a file", func() {
		var guid string
		BeforeEach(func() {
			executorRunner.Start()

			guid = factories.GenerateGuid()
		})

		It("uploads the file", func() {
			runOnce := models.RunOnce{
				Guid:     factories.GenerateGuid(),
				MemoryMB: 1024,
				DiskMB:   1024,
				Actions: []models.ExecutorAction{
					{Action: models.RunAction{Script: `echo "tasty thingy" > thingy`}},
					{Action: models.UploadAction{From: "thingy", To: inigoserver.UploadUrl("thingy")}},
					{Action: models.RunAction{Script: inigoserver.CurlCommand(guid)}},
				},
			}

			bbs.DesireRunOnce(runOnce)

			Eventually(inigoserver.ReportingGuids, 5.0).Should(ContainElement(guid))
			Ω(inigoserver.DownloadFileString("thingy")).Should(Equal("tasty thingy\n"))
		})
	})

	Describe("Fetching results", func() {
		BeforeEach(func() {
			executorRunner.Start()
		})

		It("should fetch the contents of the requested file and provide the content in the completed RunOnce", func() {
			runOnce := models.RunOnce{
				Guid:     factories.GenerateGuid(),
				MemoryMB: 1024,
				DiskMB:   1024,
				Actions: []models.ExecutorAction{
					{Action: models.RunAction{Script: `echo "tasty thingy" > thingy`}},
					{Action: models.FetchResultAction{File: "thingy"}},
				},
			}

			bbs.DesireRunOnce(runOnce)

			Eventually(bbs.GetAllCompletedRunOnces, 5.0).Should(HaveLen(1))

			runOnces, _ := bbs.GetAllCompletedRunOnces()
			Ω(runOnces[0].Result).Should(Equal("tasty thingy\n"))
		})
	})

	Describe("A RunOnce with logging configured", func() {
		BeforeEach(func() {
			executorRunner.Start()
		})

		It("has its stdout and stderr emitted to Loggregator", func(done Done) {
			logGuid := factories.GenerateGuid()

			messages, stop := loggredile.StreamMessages(
				loggregatorRunner.Config.OutgoingPort,
				"/tail/?app="+logGuid,
			)

			runOnce := factories.BuildRunOnceWithRunAction(
				1024,
				1024,
				"echo out A; echo out B; echo out C; echo err A 1>&2; echo err B 1>&2; echo err C 1>&2",
			)
			runOnce.Log.Guid = logGuid
			runOnce.Log.SourceName = "APP"

			bbs.DesireRunOnce(runOnce)

			outStream := []string{}
			errStream := []string{}

			for i := 0; i < 6; i++ {
				message := <-messages
				switch message.GetMessageType() {
				case logmessage.LogMessage_OUT:
					outStream = append(outStream, string(message.GetMessage()))
				case logmessage.LogMessage_ERR:
					errStream = append(errStream, string(message.GetMessage()))
				}
			}

			Ω(outStream).Should(Equal([]string{"out A", "out B", "out C"}))
			Ω(errStream).Should(Equal([]string{"err A", "err B", "err C"}))

			close(stop)
			close(done)
		}, 10.0)
	})
})
<|endoftext|>


<commit_before>package exporter

import (
	"fmt"
	"net/url"
	"regexp"
	"strconv"
	"strings"

	"github.com/gomodule/redigo/redis"
	"github.com/prometheus/client_golang/prometheus"
	log "github.com/sirupsen/logrus"
)

type dbKeyPair struct {
	db, key string
}

type keyInfo struct {
	size    float64
	keyType string
}

var errKeyTypeNotFound = fmt.Errorf("key not found")

// getKeyInfo takes a key and returns the type, and the size or length of the value stored at that key.
func getKeyInfo(c redis.Conn, key string) (info keyInfo, err error) {
	if info.keyType, err = redis.String(doRedisCmd(c, "TYPE", key)); err != nil {
		return info, err
	}

	switch info.keyType {
	case "none":
		return info, errKeyTypeNotFound
	case "string":
		if size, err := redis.Int64(doRedisCmd(c, "PFCOUNT", key)); err == nil {
			// hyperloglog
			info.size = float64(size)
		} else if size, err := redis.Int64(doRedisCmd(c, "STRLEN", key)); err == nil {
			info.size = float64(size)
		}
	case "list":
		if size, err := redis.Int64(doRedisCmd(c, "LLEN", key)); err == nil {
			info.size = float64(size)
		}
	case "set":
		if size, err := redis.Int64(doRedisCmd(c, "SCARD", key)); err == nil {
			info.size = float64(size)
		}
	case "zset":
		if size, err := redis.Int64(doRedisCmd(c, "ZCARD", key)); err == nil {
			info.size = float64(size)
		}
	case "hash":
		if size, err := redis.Int64(doRedisCmd(c, "HLEN", key)); err == nil {
			info.size = float64(size)
		}
	case "stream":
		if size, err := redis.Int64(doRedisCmd(c, "XLEN", key)); err == nil {
			info.size = float64(size)
		}
	default:
		err = fmt.Errorf("unknown type: %v for key: %v", info.keyType, key)
	}

	return info, err
}

func (e *Exporter) extractCheckKeyMetrics(ch chan<- prometheus.Metric, c redis.Conn) {
	keys, err := parseKeyArg(e.options.CheckKeys)
	if err != nil {
		log.Errorf("Couldn't parse check-keys: %#v", err)
		return
	}
	log.Debugf("keys: %#v", keys)

	singleKeys, err := parseKeyArg(e.options.CheckSingleKeys)
	if err != nil {
		log.Errorf("Couldn't parse check-single-keys: %#v", err)
		return
	}
	log.Debugf("e.singleKeys: %#v", singleKeys)

	allKeys := append([]dbKeyPair{}, singleKeys...)

	log.Debugf("e.keys: %#v", keys)
	scannedKeys, err := getKeysFromPatterns(c, keys, e.options.CheckKeysBatchSize)
	if err != nil {
		log.Errorf("Error expanding key patterns: %#v", err)
	} else {
		allKeys = append(allKeys, scannedKeys...)
	}

	log.Debugf("allKeys: %#v", allKeys)
	for _, k := range allKeys {
		if _, err := doRedisCmd(c, "SELECT", k.db); err != nil {
			log.Errorf("Couldn't select database %#v when getting key info.", k.db)
			continue
		}

		dbLabel := "db" + k.db
		info, err := getKeyInfo(c, k.key)
		switch err {
		case errKeyTypeNotFound:
			log.Debugf("Key '%s' not found when trying to get type and size: using default '0.0'", k.key)
			e.registerConstMetricGauge(ch, "key_size", 0.0, dbLabel, k.key)
		case nil:
			e.registerConstMetricGauge(ch, "key_size", info.size, dbLabel, k.key)

			// Only record value metric if value is float-y
			if val, err := redis.Float64(doRedisCmd(c, "GET", k.key)); err == nil {
				e.registerConstMetricGauge(ch, "key_value", val, dbLabel, k.key)
			}
		default:
			log.Error(err)
		}
	}
}

func (e *Exporter) extractCountKeysMetrics(ch chan<- prometheus.Metric, c redis.Conn) {
	cntKeys, err := parseKeyArg(e.options.CountKeys)
	if err != nil {
		log.Errorf("Couldn't parse given count keys: %s", err)
		return
	}

	for _, k := range cntKeys {
		if _, err := doRedisCmd(c, "SELECT", k.db); err != nil {
			log.Errorf("Couldn't select database '%s' when getting stream info", k.db)
			continue
		}
		cnt, err := getKeysCount(c, k.key, e.options.CheckKeysBatchSize)
		if err != nil {
			log.Errorf("couldn't get key count for '%s', err: %s", k.key, err)
			continue
		}
		dbLabel := "db" + k.db
		e.registerConstMetricGauge(ch, "keys_count", float64(cnt), dbLabel, k.key)
	}
}

func getKeysCount(c redis.Conn, pattern string, count int64) (int, error) {
	keysCount := 0

	keys, err := scanKeys(c, pattern, count)
	if err != nil {
		return keysCount, fmt.Errorf("error retrieving '%s' keys err: %s", pattern, err)
	}
	keysCount = len(keys)

	return keysCount, nil
}

// Regexp pattern to check if given key contains any
// glob-style pattern symbol.
//
// https://redis.io/commands/scan#the-match-option
var globPattern = regexp.MustCompile(`[\?\*\[\]\^]+`)

// getKeysFromPatterns does a SCAN for a key if the key contains pattern characters
func getKeysFromPatterns(c redis.Conn, keys []dbKeyPair, count int64) (expandedKeys []dbKeyPair, err error) {
	expandedKeys = []dbKeyPair{}
	for _, k := range keys {
		if globPattern.MatchString(k.key) {
			if _, err := doRedisCmd(c, "SELECT", k.db); err != nil {
				return expandedKeys, err
			}
			keyNames, err := redis.Strings(scanKeys(c, k.key, count))
			if err != nil {
				log.Errorf("error with SCAN for pattern: %#v err: %s", k.key, err)
				continue
			}
			for _, keyName := range keyNames {
				expandedKeys = append(expandedKeys, dbKeyPair{db: k.db, key: keyName})
			}
		} else {
			expandedKeys = append(expandedKeys, k)
		}
	}

	return expandedKeys, err
}

// parseKeyArgs splits a command-line supplied argument into a slice of dbKeyPairs.
func parseKeyArg(keysArgString string) (keys []dbKeyPair, err error) {
	if keysArgString == "" {
		log.Debugf("parseKeyArg(): Got empty key arguments, parsing skipped")
		return keys, err
	}
	for _, k := range strings.Split(keysArgString, ",") {
		var db string
		var key string
		if k == "" {
			continue
		}
		frags := strings.Split(k, "=")
		switch len(frags) {
		case 1:
			db = "0"
			key, err = url.QueryUnescape(strings.TrimSpace(frags[0]))
		case 2:
			db = strings.Replace(strings.TrimSpace(frags[0]), "db", "", -1)
			key, err = url.QueryUnescape(strings.TrimSpace(frags[1]))
		default:
			return keys, fmt.Errorf("invalid key list argument: %s", k)
		}
		if err != nil {
			return keys, fmt.Errorf("couldn't parse db/key string: %s", k)
		}

		// We want to guarantee at the top level that invalid values
		// will not fall into the final Redis call.
		if db == "" || key == "" {
			log.Errorf("parseKeyArg(): Empty value parsed in pair '%s=%s', skip", db, key)
			continue
		}

		number, err := strconv.Atoi(db)
		if err != nil || number < 0 {
			return keys, fmt.Errorf("Invalid database index for db \"%s\": %s", db, err)
		}

		keys = append(keys, dbKeyPair{db, key})
	}
	return keys, err
}

// scanForKeys returns a list of keys matching `pattern` by using `SCAN`, which is safer for production systems than using `KEYS`.
// This function was adapted from: https://github.com/reisinger/examples-redigo
func scanKeys(c redis.Conn, pattern string, count int64) (keys []interface{}, err error) {
	if pattern == "" {
		return keys, fmt.Errorf("Pattern shouldn't be empty")
	}

	iter := 0
	for {
		arr, err := redis.Values(doRedisCmd(c, "SCAN", iter, "MATCH", pattern, "COUNT", count))
		if err != nil {
			return keys, fmt.Errorf("error retrieving '%s' keys err: %s", pattern, err)
		}
		if len(arr) != 2 {
			return keys, fmt.Errorf("invalid response from SCAN for pattern: %s", pattern)
		}

		k, _ := redis.Values(arr[1], nil)
		keys = append(keys, k...)

		if iter, _ = redis.Int(arr[0], nil); iter == 0 {
			break
		}
	}

	return keys, nil
}
<commit_msg>Add condition to check if key is of type string before running GET to increase performance (#538)<commit_after>package exporter

import (
	"fmt"
	"net/url"
	"regexp"
	"strconv"
	"strings"

	"github.com/gomodule/redigo/redis"
	"github.com/prometheus/client_golang/prometheus"
	log "github.com/sirupsen/logrus"
)

type dbKeyPair struct {
	db, key string
}

type keyInfo struct {
	size    float64
	keyType string
}

var errKeyTypeNotFound = fmt.Errorf("key not found")

// getKeyInfo takes a key and returns the type, and the size or length of the value stored at that key.
func getKeyInfo(c redis.Conn, key string) (info keyInfo, err error) {
	if info.keyType, err = redis.String(doRedisCmd(c, "TYPE", key)); err != nil {
		return info, err
	}

	switch info.keyType {
	case "none":
		return info, errKeyTypeNotFound
	case "string":
        // Use PFCOUNT first because STRLEN on HyperLogLog strings returns the wrong length
        // while PFCOUNT only works on HLL strings and returns an error on regular strings.
		if size, err := redis.Int64(doRedisCmd(c, "PFCOUNT", key)); err == nil {
			// hyperloglog
			info.size = float64(size)
		} else if size, err := redis.Int64(doRedisCmd(c, "STRLEN", key)); err == nil {
			info.size = float64(size)
		}
	case "list":
		if size, err := redis.Int64(doRedisCmd(c, "LLEN", key)); err == nil {
			info.size = float64(size)
		}
	case "set":
		if size, err := redis.Int64(doRedisCmd(c, "SCARD", key)); err == nil {
			info.size = float64(size)
		}
	case "zset":
		if size, err := redis.Int64(doRedisCmd(c, "ZCARD", key)); err == nil {
			info.size = float64(size)
		}
	case "hash":
		if size, err := redis.Int64(doRedisCmd(c, "HLEN", key)); err == nil {
			info.size = float64(size)
		}
	case "stream":
		if size, err := redis.Int64(doRedisCmd(c, "XLEN", key)); err == nil {
			info.size = float64(size)
		}
	default:
		err = fmt.Errorf("unknown type: %v for key: %v", info.keyType, key)
	}

	return info, err
}

func (e *Exporter) extractCheckKeyMetrics(ch chan<- prometheus.Metric, c redis.Conn) {
	keys, err := parseKeyArg(e.options.CheckKeys)
	if err != nil {
		log.Errorf("Couldn't parse check-keys: %#v", err)
		return
	}
	log.Debugf("keys: %#v", keys)

	singleKeys, err := parseKeyArg(e.options.CheckSingleKeys)
	if err != nil {
		log.Errorf("Couldn't parse check-single-keys: %#v", err)
		return
	}
	log.Debugf("e.singleKeys: %#v", singleKeys)

	allKeys := append([]dbKeyPair{}, singleKeys...)

	log.Debugf("e.keys: %#v", keys)
	scannedKeys, err := getKeysFromPatterns(c, keys, e.options.CheckKeysBatchSize)
	if err != nil {
		log.Errorf("Error expanding key patterns: %#v", err)
	} else {
		allKeys = append(allKeys, scannedKeys...)
	}

	log.Debugf("allKeys: %#v", allKeys)
	for _, k := range allKeys {
		if _, err := doRedisCmd(c, "SELECT", k.db); err != nil {
			log.Errorf("Couldn't select database %#v when getting key info.", k.db)
			continue
		}

		dbLabel := "db" + k.db
		info, err := getKeyInfo(c, k.key)
		switch err {
		case errKeyTypeNotFound:
			log.Debugf("Key '%s' not found when trying to get type and size: using default '0.0'", k.key)
			e.registerConstMetricGauge(ch, "key_size", 0.0, dbLabel, k.key)
		case nil:
			e.registerConstMetricGauge(ch, "key_size", info.size, dbLabel, k.key)

            // Only run on single value strings
            if info.keyType == "string" {
			    // Only record value metric if value is float-y
			    if val, err := redis.Float64(doRedisCmd(c, "GET", k.key)); err == nil {
			    	e.registerConstMetricGauge(ch, "key_value", val, dbLabel, k.key)
			    }
            }
		default:
			log.Error(err)
		}
	}
}

func (e *Exporter) extractCountKeysMetrics(ch chan<- prometheus.Metric, c redis.Conn) {
	cntKeys, err := parseKeyArg(e.options.CountKeys)
	if err != nil {
		log.Errorf("Couldn't parse given count keys: %s", err)
		return
	}

	for _, k := range cntKeys {
		if _, err := doRedisCmd(c, "SELECT", k.db); err != nil {
			log.Errorf("Couldn't select database '%s' when getting stream info", k.db)
			continue
		}
		cnt, err := getKeysCount(c, k.key, e.options.CheckKeysBatchSize)
		if err != nil {
			log.Errorf("couldn't get key count for '%s', err: %s", k.key, err)
			continue
		}
		dbLabel := "db" + k.db
		e.registerConstMetricGauge(ch, "keys_count", float64(cnt), dbLabel, k.key)
	}
}

func getKeysCount(c redis.Conn, pattern string, count int64) (int, error) {
	keysCount := 0

	keys, err := scanKeys(c, pattern, count)
	if err != nil {
		return keysCount, fmt.Errorf("error retrieving '%s' keys err: %s", pattern, err)
	}
	keysCount = len(keys)

	return keysCount, nil
}

// Regexp pattern to check if given key contains any
// glob-style pattern symbol.
//
// https://redis.io/commands/scan#the-match-option
var globPattern = regexp.MustCompile(`[\?\*\[\]\^]+`)

// getKeysFromPatterns does a SCAN for a key if the key contains pattern characters
func getKeysFromPatterns(c redis.Conn, keys []dbKeyPair, count int64) (expandedKeys []dbKeyPair, err error) {
	expandedKeys = []dbKeyPair{}
	for _, k := range keys {
		if globPattern.MatchString(k.key) {
			if _, err := doRedisCmd(c, "SELECT", k.db); err != nil {
				return expandedKeys, err
			}
			keyNames, err := redis.Strings(scanKeys(c, k.key, count))
			if err != nil {
				log.Errorf("error with SCAN for pattern: %#v err: %s", k.key, err)
				continue
			}
			for _, keyName := range keyNames {
				expandedKeys = append(expandedKeys, dbKeyPair{db: k.db, key: keyName})
			}
		} else {
			expandedKeys = append(expandedKeys, k)
		}
	}

	return expandedKeys, err
}

// parseKeyArgs splits a command-line supplied argument into a slice of dbKeyPairs.
func parseKeyArg(keysArgString string) (keys []dbKeyPair, err error) {
	if keysArgString == "" {
		log.Debugf("parseKeyArg(): Got empty key arguments, parsing skipped")
		return keys, err
	}
	for _, k := range strings.Split(keysArgString, ",") {
		var db string
		var key string
		if k == "" {
			continue
		}
		frags := strings.Split(k, "=")
		switch len(frags) {
		case 1:
			db = "0"
			key, err = url.QueryUnescape(strings.TrimSpace(frags[0]))
		case 2:
			db = strings.Replace(strings.TrimSpace(frags[0]), "db", "", -1)
			key, err = url.QueryUnescape(strings.TrimSpace(frags[1]))
		default:
			return keys, fmt.Errorf("invalid key list argument: %s", k)
		}
		if err != nil {
			return keys, fmt.Errorf("couldn't parse db/key string: %s", k)
		}

		// We want to guarantee at the top level that invalid values
		// will not fall into the final Redis call.
		if db == "" || key == "" {
			log.Errorf("parseKeyArg(): Empty value parsed in pair '%s=%s', skip", db, key)
			continue
		}

		number, err := strconv.Atoi(db)
		if err != nil || number < 0 {
			return keys, fmt.Errorf("Invalid database index for db \"%s\": %s", db, err)
		}

		keys = append(keys, dbKeyPair{db, key})
	}
	return keys, err
}

// scanForKeys returns a list of keys matching `pattern` by using `SCAN`, which is safer for production systems than using `KEYS`.
// This function was adapted from: https://github.com/reisinger/examples-redigo
func scanKeys(c redis.Conn, pattern string, count int64) (keys []interface{}, err error) {
	if pattern == "" {
		return keys, fmt.Errorf("Pattern shouldn't be empty")
	}

	iter := 0
	for {
		arr, err := redis.Values(doRedisCmd(c, "SCAN", iter, "MATCH", pattern, "COUNT", count))
		if err != nil {
			return keys, fmt.Errorf("error retrieving '%s' keys err: %s", pattern, err)
		}
		if len(arr) != 2 {
			return keys, fmt.Errorf("invalid response from SCAN for pattern: %s", pattern)
		}

		k, _ := redis.Values(arr[1], nil)
		keys = append(keys, k...)

		if iter, _ = redis.Int(arr[0], nil); iter == 0 {
			break
		}
	}

	return keys, nil
}
<|endoftext|>


<commit_before>// Copyright 2010 Bill Burdick. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// dice is a program that uses seq to calculate some probabilities for a game:
// each player rolls 3 dice, each of which can have 4, 6, 8, or 10 sides
// then, they compare the maximum numbers.  The one with the highest number wins
// this program calculates probabilities of winning by a certain margin (which can be from 1 - 9)
package main

import "fmt"
import "io"
import "bytes"
import "container/vector"
import "sort"
import . "github.com/zot/bills-tools/seq"
//import "reflect"

func add(i int, s Seq) Seq {
	return s.Map(func(el El)El {
		return i + el.(int)
	})
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func main() {
	d4 := add(1, AUpto(4))
	d6 := add(1, AUpto(6))
	d8 := add(1, AUpto(8))
	d10 := add(1, AUpto(10))
	names := map[interface{}]string{d4:"d4", d6:"d6", d8:"d8", d10:"d10"}
	dice := From(d4, d6, d8, d10)
	rank := map[Seq]int{d4:0, d6:1, d8:2, d10:3}
	sets := map[string]int{}
	//attempts is [[label, [score, ...]]...]
	attempts := Map(Filter(Product(From(dice, dice, dice)), func(d El)bool{
		oldRank := -1
		result := true
		Do(d.(Seq), func(set El){
			newRank := rank[set.(Seq)]
			result = result && newRank >= oldRank
			oldRank = newRank
		})
		return result
	}), func(el El)El{
		buf := bytes.NewBuffer(make([]byte, 0, 10))
		io.WriteString(buf, "<")
		Pretty(el.(Seq), names, buf)
		io.WriteString(buf, ">")
		return From(buf.String(), Map(Product(el.(Seq)), func(el El)El{
			return Fold(el.(Seq), 0, func(acc, el El)El{return max(acc.(int), el.(int))})
		}))

	})
	println("#sets:", len(sets))
	fmt.Println("#Attempts:", Len(attempts))
	println("results...")
	Do(attempts, func(el El){
		label, rolls := First2(el.(Seq))
		fmt.Printf("%s: %d\n", label, Len(rolls.(Seq)))
	})
	Do(CFlatMap(attempts, func(el El) Seq {
		label, sc := First2(el.(Seq))
		return CMap(attempts, func(del El) El {
			rolls, wins := 0, 0
			margins := map[int]int{}
			dlabel, dsc := First2(del.(Seq))
			Do(Product(From(sc,dsc)), func(rel El){
				rolls++
				attack, defense := First2(rel.(Seq))
				margin := attack.(int) - defense.(int)
				if margin > 0 {
					wins++
					margins[margin]++
				}

			})
			return From(label, dlabel, rolls, wins, margins)
		})
	}), func(el El){
		l, d, r, w, m := First5(el.(Seq))
		printResult(l.(string), d.(string), r.(int), w.(int), m.(map[int]int))
	})
}

func printResult(label string, dlabel string, rolls int, wins int, margins map[int]int) {
	fmt.Printf("%s vs %s rolls: %d wins: %d margins:", label, dlabel, rolls, wins)
	for i := 1; i <= 9; i++ {
		v := margins[i]
		if v > 0 {
			fmt.Printf(" %d %.2f", v, float(v)*100/float(wins))
		}
	}
	println()
	dumpMargin(wins, margins)
}

func round(value float) int {
	floor := int(value)
	if value - float(floor) > 0.5 {
		if value > 0 {
			return floor + 1
		}
		return floor - 1
	}
	return floor
}

func dumpMargin(totMargin int, margins map[int]int) {
	for k := 1; k <= 9; k++ {
		v := margins[k]
		if v > 0 {
			percent := float(v)*100/float(totMargin)
			fmt.Printf("%d: %10d (%6.2f) ", k, v, percent)
			for i := 0; i < int(round(percent)); i++ {
				print(".")
			}
			println()
		}
	}
}
func dumpResults(totMargin int, margins map[string]map[int]int) {
	vec := vector.StringVector(make([]string, 0, 32))
	for k := range margins {
		vec.Push(k)
	}
	sort.StringArray(vec).Sort()
	for _, dice := range vec {
		println("Margins for", dice)
		dumpMargin(totMargin, margins[dice])
	}
}
<commit_msg>fixed dice.go<commit_after>// Copyright 2010 Bill Burdick. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// dice is a program that uses seq to calculate some probabilities for a game:
// each player rolls 3 dice, each of which can have 4, 6, 8, or 10 sides
// then, they compare the maximum numbers.  The one with the highest number wins
// this program calculates probabilities of winning by a certain margin (which can be from 1 - 9)
package main

import "fmt"
import "io"
import "bytes"
import "container/vector"
import "sort"
import . "github.com/zot/seq"
//import "reflect"

func add(i int, s Seq) Seq {
	return s.Map(func(el El)El {
		return i + el.(int)
	})
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func main() {
	d4 := add(1, AUpto(4))
	d6 := add(1, AUpto(6))
	d8 := add(1, AUpto(8))
	d10 := add(1, AUpto(10))
	names := map[interface{}]string{d4:"d4", d6:"d6", d8:"d8", d10:"d10"}
	dice := From(d4, d6, d8, d10)
	rank := map[Seq]int{d4:0, d6:1, d8:2, d10:3}
	sets := map[string]int{}
	//attempts is [[label, [score, ...]]...]
	attempts := Map(Filter(Product(From(dice, dice, dice)), func(d El)bool{
		oldRank := -1
		result := true
		Do(d.(Seq), func(set El){
			newRank := rank[set.(Seq)]
			result = result && newRank >= oldRank
			oldRank = newRank
		})
		return result
	}), func(el El)El{
		buf := bytes.NewBuffer(make([]byte, 0, 10))
		io.WriteString(buf, "<")
		Pretty(el.(Seq), names, buf)
		io.WriteString(buf, ">")
		return From(buf.String(), Map(Product(el.(Seq)), func(el El)El{
			return Fold(el.(Seq), 0, func(acc, el El)El{return max(acc.(int), el.(int))})
		}))

	})
	println("#sets:", len(sets))
	fmt.Println("#Attempts:", Len(attempts))
	println("results...")
	Do(attempts, func(el El){
		label, rolls := First2(el.(Seq))
		fmt.Printf("%s: %d\n", label, Len(rolls.(Seq)))
	})
	Do(CFlatMap(attempts, func(el El) Seq {
		label, sc := First2(el.(Seq))
		return CMap(attempts, func(del El) El {
			rolls, wins := 0, 0
			margins := map[int]int{}
			dlabel, dsc := First2(del.(Seq))
			Do(Product(From(sc,dsc)), func(rel El){
				rolls++
				attack, defense := First2(rel.(Seq))
				margin := attack.(int) - defense.(int)
				if margin > 0 {
					wins++
					margins[margin]++
				}

			})
			return From(label, dlabel, rolls, wins, margins)
		})
	}), func(el El){
		l, d, r, w, m := First5(el.(Seq))
		printResult(l.(string), d.(string), r.(int), w.(int), m.(map[int]int))
	})
}

func printResult(label string, dlabel string, rolls int, wins int, margins map[int]int) {
	fmt.Printf("%s vs %s rolls: %d wins: %d margins:", label, dlabel, rolls, wins)
	for i := 1; i <= 9; i++ {
		v := margins[i]
		if v > 0 {
			fmt.Printf(" %d %.2f", v, float(v)*100/float(wins))
		}
	}
	println()
	dumpMargin(wins, margins)
}

func round(value float) int {
	floor := int(value)
	if value - float(floor) > 0.5 {
		if value > 0 {
			return floor + 1
		}
		return floor - 1
	}
	return floor
}

func dumpMargin(totMargin int, margins map[int]int) {
	for k := 1; k <= 9; k++ {
		v := margins[k]
		if v > 0 {
			percent := float(v)*100/float(totMargin)
			fmt.Printf("%d: %10d (%6.2f) ", k, v, percent)
			for i := 0; i < int(round(percent)); i++ {
				print(".")
			}
			println()
		}
	}
}
func dumpResults(totMargin int, margins map[string]map[int]int) {
	vec := vector.StringVector(make([]string, 0, 32))
	for k := range margins {
		vec.Push(k)
	}
	sort.StringArray(vec).Sort()
	for _, dice := range vec {
		println("Margins for", dice)
		dumpMargin(totMargin, margins[dice])
	}
}
<|endoftext|>


<commit_before>package httpreciver
import (
	"net/http"
	"errors"
	"log"
	"os"
	"io"
	"mime/multipart"
	"strconv"
	"crypto/rand"
	"fmt"
)
func sep() string {
	st := strconv.QuoteRune(os.PathSeparator)
	st = st[1 : len(st) - 1]
	return st
}

func genUid() string {
	b := make([]byte, 16)
	rand.Read(b)
	return fmt.Sprintf("%X-%X-%X-%X-%X", b[0:4], b[4:6], b[6:8], b[8:10], b[10:])
}

//import "godownloader/monitor"
const BufferSize = 1024 * 1024

type OnDownloadFunc func(path string) (error)

type FilesReceiver struct {
	buffer   []byte
	mprd     *multipart.Reader
	rw       http.ResponseWriter
	req      *http.Request
	onDown   OnDownloadFunc
	dataSize int64
	upPos    int64

}

func CreateReciver(rw_ http.ResponseWriter, req_ *http.Request, odf OnDownloadFunc) (*FilesReceiver, error) {
	if req_ == nil {
		return nil, errors.New("error: empty http request")
	}
	var fs FilesReceiver
	fs.rw = rw_
	fs.req = req_
	fs.onDown = odf
	return &fs, nil
}

func (fs *FilesReceiver) DoWork() (bool, error) {
	defer log.Printf("info: uploaded data %v percent \n", (fs.upPos * 100) / fs.dataSize)
	p, err := fs.mprd.NextPart()
	if err != nil {
		log.Println("info: part is out of")
		return true, nil
	}

	fn := p.FormName()
	switch fn {
	case "files":{
		//if f, er := os.Create(os.TempDir() + sep() + genUid()); er != nil {
		if f, er := os.Create("/home/andrew/Desktop/da/" + sep() + genUid() + ".jpeg"); er != nil {
			defer f.Close()
			log.Println("error: can't create temp file")
			return false, er

		}else {
			for {
				if count, e := p.Read(fs.buffer); e == io.EOF {
					log.Printf("info: file %v writed to disk \n", p.FileName())
					return false, nil
				}else {
					f.Write(fs.buffer[0:count])
					fs.upPos += int64(count)
				}

			}
		}
	}
	case "server":{
		log.Println("detect server form")
	}
	case "port":{
		log.Println("detect port form")
	}
	case "aetitle":{
		log.Println("detect aetitle form")
	}
	default:{
		log.Println("warning: unsupported form")
	}
	}

	fs.upPos = fs.dataSize
	return false, nil
}
func (fs *FilesReceiver) GetProgress() interface{} {
	return (fs.upPos * 100) / fs.dataSize
}
func (fs *FilesReceiver) BeforeRun() error {
	var err error
	fs.mprd, err = fs.req.MultipartReader()
	fs.dataSize = fs.req.ContentLength
	if err != nil {
		return err
	}
	fs.buffer = make([]byte, BufferSize)
	return nil
}
func (fs *FilesReceiver) AfterStop() error {
	return nil
}<commit_msg>cosmetic fix<commit_after>package httpreciver
import (
	"net/http"
	"errors"
	"log"
	"os"
	"io"
	"mime/multipart"
	"strconv"
	"crypto/rand"
	"fmt"
)
func sep() string {
	st := strconv.QuoteRune(os.PathSeparator)
	st = st[1 : len(st) - 1]
	return st
}

func genUid() string {
	b := make([]byte, 16)
	rand.Read(b)
	return fmt.Sprintf("%X-%X-%X-%X-%X", b[0:4], b[4:6], b[6:8], b[8:10], b[10:])
}

//import "godownloader/monitor"
const BufferSize = 1024 * 1024

type OnDownloadFunc func(path string) (error)

type DicomReceiver struct {
	buffer   []byte
	mprd     *multipart.Reader
	rw       http.ResponseWriter
	req      *http.Request
	onDown   OnDownloadFunc
	dataSize int64
	upPos    int64

}

func CreateReciver(rw_ http.ResponseWriter, req_ *http.Request, odf OnDownloadFunc) (*DicomReceiver, error) {
	if req_ == nil {
		return nil, errors.New("error: empty http request")
	}
	var fs DicomReceiver
	fs.rw = rw_
	fs.req = req_
	fs.onDown = odf
	return &fs, nil
}

func (fs *DicomReceiver) DoWork() (bool, error) {
	defer log.Printf("info: uploaded data %v percent \n", (fs.upPos * 100) / fs.dataSize)
	p, err := fs.mprd.NextPart()
	if err != nil {
		log.Println("info: part is out of")
		return true, nil
	}

	fn := p.FormName()
	switch fn {
	case "files":{
		//if f, er := os.Create(os.TempDir() + sep() + genUid()); er != nil {
		if f, er := os.Create("/home/andrew/Desktop/da/" + sep() + genUid() + ".jpeg"); er != nil {
			defer f.Close()
			log.Println("error: can't create temp file")
			return false, er

		}else {
			for {
				if count, e := p.Read(fs.buffer); e == io.EOF {
					log.Printf("info: file %v writed to disk \n", p.FileName())
					return false, nil
				}else {
					f.Write(fs.buffer[0:count])
					fs.upPos += int64(count)
				}

			}
		}
	}
	case "server":{
		log.Println("detect server form")
	}
	case "port":{
		log.Println("detect port form")
	}
	case "aetitle":{
		log.Println("detect aetitle form")
	}
	default:{
		log.Println("warning: unsupported form")
	}
	}

	fs.upPos = fs.dataSize
	return false, nil
}
func (fs *DicomReceiver) GetProgress() interface{} {
	return (fs.upPos * 100) / fs.dataSize
}
func (fs *DicomReceiver) BeforeRun() error {
	var err error
	fs.mprd, err = fs.req.MultipartReader()
	fs.dataSize = fs.req.ContentLength
	if err != nil {
		return err
	}
	fs.buffer = make([]byte, BufferSize)
	return nil
}
func (fs *DicomReceiver) AfterStop() error {
	return nil
}<|endoftext|>


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

package imagetest

import (
	"testing"

	"github.com/GoogleCloudPlatform/compute-image-tools/daisy"
)

func TestAddStartStep(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	step, err := twf.addStartStep("stepname", "vmname")
	if err != nil {
		t.Errorf("failed to add start step to test workflow: %v", err)
	}
	if step.StartInstances == nil {
		t.Fatal("StartInstances step is missing")
	}
	if len(step.StartInstances.Instances) != 1 {
		t.Error("StartInstances step is malformed")
	}
	if step.StartInstances.Instances[0] != "vmname" {
		t.Error("StartInstances step is malformed")
	}
	if stepFromWF, ok := twf.wf.Steps["start-stepname"]; !ok || step != stepFromWF {
		t.Error("Step was not correctly added to workflow")
	}
}

func TestAddStopStep(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	step, err := twf.addStopStep("stepname", "vmname")
	if err != nil {
		t.Errorf("failed to add stop step to test workflow: %v", err)
	}
	if step.StopInstances == nil {
		t.Fatal("StopInstances step is missing")
	}
	if len(step.StopInstances.Instances) != 1 {
		t.Error("StopInstances step is malformed")
	}
	if step.StopInstances.Instances[0] != "vmname" {
		t.Error("StopInstances step is malformed")
	}
	if stepFromWF, ok := twf.wf.Steps["stop-stepname"]; !ok || step != stepFromWF {
		t.Error("step was not correctly added to workflow")
	}
}

func TestAddWaitStep(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	step, err := twf.addWaitStep("stepname", "vmname", false)
	if err != nil {
		t.Errorf("failed to add wait step to test workflow: %v", err)
	}
	if step.WaitForInstancesSignal == nil {
		t.Fatal("WaitForInstancesSignal step is missing")
	}
	instancesSignal := []*daisy.InstanceSignal(*step.WaitForInstancesSignal)
	if len(instancesSignal) != 1 {
		t.Error("waitInstances step is malformed")
	}
	if instancesSignal[0].Name != "vmname" {
		t.Error("waitInstances step is malformed")
	}
	if instancesSignal[0].SerialOutput.SuccessMatch != successMatch {
		t.Error("waitInstances step is malformed")
	}
	if instancesSignal[0].Stopped {
		t.Error("waitInstances step is malformed")
	}
	if stepFromWF, ok := twf.wf.Steps["wait-stepname"]; !ok || step != stepFromWF {
		t.Error("step was not correctly added to workflow")
	}
}

func TestAddWaitStoppedStep(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	step, err := twf.addWaitStep("stepname", "vmname", true)
	if err != nil {
		t.Errorf("failed to add wait step to test workflow: %v", err)
	}
	if step.WaitForInstancesSignal == nil {
		t.Fatal("WaitForInstancesSignal step is missing")
	}
	instancesSignal := []*daisy.InstanceSignal(*step.WaitForInstancesSignal)
	if len(instancesSignal) != 1 {
		t.Error("waitInstances step is malformed")
	}
	if instancesSignal[0].Name != "vmname" {
		t.Error("waitInstances step is malformed")
	}
	if instancesSignal[0].SerialOutput != nil {
		t.Error("waitInstances step is malformed")
	}
	if !instancesSignal[0].Stopped {
		t.Error("waitInstances step is malformed")
	}
	if stepFromWF, ok := twf.wf.Steps["wait-stepname"]; !ok || step != stepFromWF {
		t.Error("step was not correctly added to workflow")
	}
}

func TestAppendCreateDisksStep(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	step, err := twf.appendCreateDisksStep("diskname")
	if err != nil {
		t.Errorf("failed to add wait step to test workflow: %v", err)
	}
	if step.CreateDisks == nil {
		t.Fatal("CreateDisks step is missing")
	}
	disks := []*daisy.Disk(*step.CreateDisks)
	if len(disks) != 1 {
		t.Error("CreateDisks step is malformed")
	}
	if disks[0].Name != "diskname" {
		t.Error("CreateDisks step is malformed")
	}
	if disks[0].SourceImage != "image" {
		t.Error("CreateDisks step is malformed")
	}
	stepFromWF, ok := twf.wf.Steps["create-disks"]
	if !ok || step != stepFromWF {
		t.Error("step was not correctly added to workflow")
	}
	step2, err := twf.appendCreateDisksStep("diskname2")
	if err != nil {
		t.Fatalf("failed to add wait step to test workflow: %v", err)
	}
	if step2 != stepFromWF {
		t.Fatal("CreateDisks step was not appended")
	}
	disks = []*daisy.Disk(*step2.CreateDisks)
	if len(disks) != 2 {
		t.Fatal("CreateDisks step was not appended")
	}
	if disks[1].Name != "diskname2" {
		t.Error("CreateDisks step is malformed")
	}
}

func TestAppendCreateVMStep(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	if _, ok := twf.wf.Steps["create-disks"]; ok {
		t.Fatal("create-disks step already exists")
	}
	step, err := twf.appendCreateVMStep("vmname", "")
	if err != nil {
		t.Errorf("failed to add wait step to test workflow: %v", err)
	}
	if step.CreateInstances == nil {
		t.Fatal("CreateDisks step is missing")
	}
	instances := step.CreateInstances.Instances
	if len(instances) != 1 {
		t.Error("CreateInstances step is malformed")
	}
	if instances[0].Name != "vmname" {
		t.Error("CreateInstances step is malformed")
	}
	stepFromWF, ok := twf.wf.Steps["create-vms"]
	if !ok || step != stepFromWF {
		t.Error("step was not correctly added to workflow")
	}
	step2, err := twf.appendCreateVMStep("vmname2", "")
	if err != nil {
		t.Fatalf("failed to add wait step to test workflow: %v", err)
	}
	if step2 != stepFromWF {
		t.Fatal("CreateDisks step was not appended")
	}
	instances = step.CreateInstances.Instances
	if len(instances) != 2 {
		t.Fatal("CreateDisks step was not appended")
	}
	if instances[1].Name != "vmname2" {
		t.Error("CreateInstances step is malformed")
	}
}

func TestAppendCreateVMStepCustomHostname(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	if _, ok := twf.wf.Steps["create-disks"]; ok {
		t.Fatal("create-disks step already exists")
	}
	step, err := twf.appendCreateVMStep("vmname", "vmname.example.com")
	if err != nil {
		t.Errorf("failed to add wait step to test workflow: %v", err)
	}
	if step.CreateInstances == nil {
		t.Fatal("CreateDisks step is missing")
	}
	instances := step.CreateInstances.Instances
	if len(instances) != 1 {
		t.Error("CreateInstances step is malformed")
	}
	if instances[0].Name != "vmname" {
		t.Error("CreateInstances step is malformed")
	}
	if instances[0].Hostname != "vmname.example.com" {
		t.Error("CreateInstances step is malformed")
	}
}

func TestNewTestWorkflow(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	if len(twf.wf.Steps) != 0 {
		t.Error("test workflow has initial steps")
	}
}
<commit_msg>Image test unit test (#157)<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package imagetest

import (
	"testing"

	"github.com/GoogleCloudPlatform/compute-image-tools/daisy"
)

func TestAddStartStep(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	step, err := twf.addStartStep("stepname", "vmname")
	if err != nil {
		t.Errorf("failed to add start step to test workflow: %v", err)
	}
	if step.StartInstances == nil {
		t.Fatal("StartInstances step is missing")
	}
	if len(step.StartInstances.Instances) != 1 {
		t.Error("StartInstances step is malformed")
	}
	if step.StartInstances.Instances[0] != "vmname" {
		t.Error("StartInstances step is malformed")
	}
	if stepFromWF, ok := twf.wf.Steps["start-stepname"]; !ok || step != stepFromWF {
		t.Error("Step was not correctly added to workflow")
	}
}

func TestAddStopStep(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	step, err := twf.addStopStep("stepname", "vmname")
	if err != nil {
		t.Errorf("failed to add stop step to test workflow: %v", err)
	}
	if step.StopInstances == nil {
		t.Fatal("StopInstances step is missing")
	}
	if len(step.StopInstances.Instances) != 1 {
		t.Error("StopInstances step is malformed")
	}
	if step.StopInstances.Instances[0] != "vmname" {
		t.Error("StopInstances step is malformed")
	}
	if stepFromWF, ok := twf.wf.Steps["stop-stepname"]; !ok || step != stepFromWF {
		t.Error("step was not correctly added to workflow")
	}
}

func TestAddWaitStep(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	step, err := twf.addWaitStep("stepname", "vmname", false)
	if err != nil {
		t.Errorf("failed to add wait step to test workflow: %v", err)
	}
	if step.WaitForInstancesSignal == nil {
		t.Fatal("WaitForInstancesSignal step is missing")
	}
	instancesSignal := []*daisy.InstanceSignal(*step.WaitForInstancesSignal)
	if len(instancesSignal) != 1 {
		t.Error("waitInstances step is malformed")
	}
	if instancesSignal[0].Name != "vmname" {
		t.Error("waitInstances step is malformed")
	}
	if instancesSignal[0].SerialOutput.SuccessMatch != successMatch {
		t.Error("waitInstances step is malformed")
	}
	if instancesSignal[0].Stopped {
		t.Error("waitInstances step is malformed")
	}
	if stepFromWF, ok := twf.wf.Steps["wait-stepname"]; !ok || step != stepFromWF {
		t.Error("step was not correctly added to workflow")
	}
}

func TestAddWaitStoppedStep(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	step, err := twf.addWaitStep("stepname", "vmname", true)
	if err != nil {
		t.Errorf("failed to add wait step to test workflow: %v", err)
	}
	if step.WaitForInstancesSignal == nil {
		t.Fatal("WaitForInstancesSignal step is missing")
	}
	instancesSignal := []*daisy.InstanceSignal(*step.WaitForInstancesSignal)
	if len(instancesSignal) != 1 {
		t.Error("waitInstances step is malformed")
	}
	if instancesSignal[0].Name != "vmname" {
		t.Error("waitInstances step is malformed")
	}
	if instancesSignal[0].SerialOutput != nil {
		t.Error("waitInstances step is malformed")
	}
	if !instancesSignal[0].Stopped {
		t.Error("waitInstances step is malformed")
	}
	if stepFromWF, ok := twf.wf.Steps["wait-stepname"]; !ok || step != stepFromWF {
		t.Error("step was not correctly added to workflow")
	}
}

func TestAppendCreateDisksStep(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	step, err := twf.appendCreateDisksStep("diskname")
	if err != nil {
		t.Errorf("failed to add wait step to test workflow: %v", err)
	}
	if step.CreateDisks == nil {
		t.Fatal("CreateDisks step is missing")
	}
	disks := []*daisy.Disk(*step.CreateDisks)
	if len(disks) != 1 {
		t.Error("CreateDisks step is malformed")
	}
	if disks[0].Name != "diskname" {
		t.Error("CreateDisks step is malformed")
	}
	if disks[0].SourceImage != "image" {
		t.Error("CreateDisks step is malformed")
	}
	stepFromWF, ok := twf.wf.Steps["create-disks"]
	if !ok || step != stepFromWF {
		t.Error("step was not correctly added to workflow")
	}
	step2, err := twf.appendCreateDisksStep("diskname2")
	if err != nil {
		t.Fatalf("failed to add wait step to test workflow: %v", err)
	}
	if step2 != stepFromWF {
		t.Fatal("CreateDisks step was not appended")
	}
	disks = []*daisy.Disk(*step2.CreateDisks)
	if len(disks) != 2 {
		t.Fatal("CreateDisks step was not appended")
	}
	if disks[1].Name != "diskname2" {
		t.Error("CreateDisks step is malformed")
	}
}

func TestAppendCreateVMStep(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	if _, ok := twf.wf.Steps["create-disks"]; ok {
		t.Fatal("create-disks step already exists")
	}
	step, err := twf.appendCreateVMStep("vmname", "")
	if err != nil {
		t.Errorf("failed to add wait step to test workflow: %v", err)
	}
	if step.CreateInstances == nil {
		t.Fatal("CreateDisks step is missing")
	}
	instances := step.CreateInstances.Instances
	if len(instances) != 1 {
		t.Error("CreateInstances step is malformed")
	}
	if instances[0].Name != "vmname" {
		t.Error("CreateInstances step is malformed")
	}
	stepFromWF, ok := twf.wf.Steps["create-vms"]
	if !ok || step != stepFromWF {
		t.Error("step was not correctly added to workflow")
	}
	step2, err := twf.appendCreateVMStep("vmname2", "")
	if err != nil {
		t.Fatalf("failed to add wait step to test workflow: %v", err)
	}
	if step2 != stepFromWF {
		t.Fatal("CreateDisks step was not appended")
	}
	instances = step.CreateInstances.Instances
	if len(instances) != 2 {
		t.Fatal("CreateDisks step was not appended")
	}
	if instances[1].Name != "vmname2" {
		t.Error("CreateInstances step is malformed")
	}
}

func TestAppendCreateVMStepCustomHostname(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	if _, ok := twf.wf.Steps["create-disks"]; ok {
		t.Fatal("create-disks step already exists")
	}
	step, err := twf.appendCreateVMStep("vmname", "vmname.example.com")
	if err != nil {
		t.Errorf("failed to add wait step to test workflow: %v", err)
	}
	if step.CreateInstances == nil {
		t.Fatal("CreateDisks step is missing")
	}
	instances := step.CreateInstances.Instances
	if len(instances) != 1 {
		t.Error("CreateInstances step is malformed")
	}
	if instances[0].Name != "vmname" {
		t.Error("CreateInstances step is malformed")
	}
	if instances[0].Hostname != "vmname.example.com" {
		t.Error("CreateInstances step is malformed")
	}
}

func TestNewTestWorkflow(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if twf.wf == nil {
		t.Fatal("test workflow is malformed")
	}
	if len(twf.wf.Steps) != 0 {
		t.Error("test workflow has initial steps")
	}
}

func TestGetLastStepForVM(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	if _, err = twf.CreateTestVM("vm"); err != nil {
		t.Errorf("failed to create test vm: %v", err)
	}
	step, err := twf.getLastStepForVM("vm")
	if err != nil {
		t.Errorf("failed to get last step for vm: %v", err)
	}
	if step.WaitForInstancesSignal == nil {
		t.Error("not wait step")
	}
	if twf.wf.Steps["wait-vm"] != step {
		t.Error("not wait-vm step")
	}
}

func TestGetLastStepForVMWhenReboot(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	tvm, err := twf.CreateTestVM("vm")
	if err != nil {
		t.Errorf("failed to create test vm: %v", err)
	}
	if err := tvm.Reboot(); err != nil {
		t.Errorf("failed to reboot: %v", err)
	}
	step, err := twf.getLastStepForVM("vm")
	if err != nil {
		t.Errorf("failed to get last step for vm: %v", err)
	}
	if step.WaitForInstancesSignal == nil {
		t.Error("not wait step")
	}
	if twf.wf.Steps["wait-started-vm-1"] != step {
		t.Error("not wait-started-vm-1 step")
	}
}

func TestGetLastStepForVMWhenMultipleReboot(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	tvm, err := twf.CreateTestVM("vm")
	if err != nil {
		t.Errorf("failed to create test vm: %v", err)
	}
	if err := tvm.Reboot(); err != nil {
		t.Errorf("failed to reboot: %v", err)
	}
	if err := tvm.Reboot(); err != nil {
		t.Errorf("failed to reboot: %v", err)
	}
	step, err := twf.getLastStepForVM("vm")
	if err != nil {
		t.Errorf("failed to get last step for vm: %v", err)
	}
	if step.WaitForInstancesSignal == nil {
		t.Error("not wait step")
	}
	if twf.wf.Steps["wait-started-vm-2"] != step {
		t.Error("not wait-started-vm-2 step")
	}
}

func TestResizeDiskAndReboot(t *testing.T) {
	twf, err := NewTestWorkflow("name", "image")
	if err != nil {
		t.Errorf("failed to create test workflow: %v", err)
	}
	tvm, err := twf.CreateTestVM("vm")
	if err != nil {
		t.Errorf("failed to create test vm: %v", err)
	}
	if err := tvm.ResizeDiskAndReboot("vm", 200); err != nil {
		t.Errorf("failed to reboot: %v", err)
	}
	step, err := twf.getLastStepForVM("vm")
	if err != nil {
		t.Errorf("failed to get last step for vm: %v", err)
	}
	if step.WaitForInstancesSignal == nil {
		t.Error("not wait step")
	}
	if twf.wf.Steps["wait-started-vm-2"] != step {
		t.Error("not wait-started-vm-2 step")
	}
}
<|endoftext|>


<commit_before>package main

import "github.com/flydb/flydb"

func main() {
    db, _ := flydb.Open("db.json")

    // get
    println db.MustGetString("users.0.username")

    // set
    db.Set("users.0.star", 100)

    // filter
    db.MustGet("users").Filter(func(n *db.Node) bool {
        return true
    })

}<commit_msg>fix example<commit_after>package main

import "github.com/flydb/flydb"

func main() {
    db, _ := flydb.Open("db.json")

    // get
    println(db.Root().MustGet("users.0.username").MustValue().MustString())

    // set
    db.Root().Set("users.0.star", 100)

    println(db.Root().MustGet("users.0.star").MustValue().MustInt())

    // filter
    // db.MustGet("users").Filter(func(n *db.Node) bool {
    //     return true
    // })

}<|endoftext|>