<commit_before>package commands

import (
	"fmt"
	"io/ioutil"
	"os"
	"strconv"

	"github.com/integralist/go-fastly-cli/common"
	"github.com/integralist/go-fastly-cli/flags"

	fastly "github.com/sethvargo/go-fastly"
)

// Upload takes specified list of files and creates new remote version
// if upload fails it'll attempt uploading over existing remote version
func Upload(f flags.Flags, client *fastly.Client) {
	checkIncorrectFlagConfiguration(f)
	configureSkipMatch(f)

	// store value rather than dereference pointer multiple times later
	fastlyServiceID = *f.Top.Service

	// the acquireVersion function checks if we should...
	//
	// 		A. clone the specified version before uploading files: `-clone`
	// 		B. upload files to the specified version: `-version`
	// 		C. upload files to the latest version: `-latest`
	// 		D. clone the latest version available
	selectedVersion, err := acquireVersion(f, client)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	processFiles(selectedVersion, uploadVCL, handleResponse, f, client)
}

func checkIncorrectFlagConfiguration(f flags.Flags) {
	if *f.Sub.CloneVersion != "" && *f.Sub.UploadVersion != "" {
		fmt.Println("Please do not provide both -clone-version and -upload-version flags")
		os.Exit(1)
	}
}

func cloneFromVersion(version int, client *fastly.Client) (*fastly.Version, error) {
	clonedVersion, err := client.CloneVersion(&fastly.CloneVersionInput{
		Service: fastlyServiceID,
		Version: version,
	})
	if err != nil {
		return nil, err
	}

	return clonedVersion, nil
}

func acquireVersion(f flags.Flags, client *fastly.Client) (int, error) {
	// clone from specified version and upload to that
	if *f.Sub.CloneVersion != "" {
		cloneVersion, err := strconv.Atoi(*f.Sub.CloneVersion)
		if err != nil {
			return 0, err
		}

		clonedVersion, err := cloneFromVersion(cloneVersion, client)
		if err != nil {
			return 0, err
		}

		fmt.Printf("Successfully created new version %d from existing version %s\n\n", clonedVersion.Number, *f.Sub.CloneVersion)
		return clonedVersion.Number, nil
	}

	// upload to the specified version (it can't be activated)
	if *f.Sub.UploadVersion != "" {
		uploadVersion, err := strconv.Atoi(*f.Sub.UploadVersion)
		if err != nil {
			return 0, err
		}

		getVersion, err := client.GetVersion(&fastly.GetVersionInput{
			Service: fastlyServiceID,
			Version: uploadVersion,
		})
		if err != nil {
			return 0, err
		}

		if getVersion.Active {
			return 0, fmt.Errorf("Sorry, the specified version is already activated")
		}

		return uploadVersion, nil
	}

	latestVersion, err := common.GetLatestVCLVersion(*f.Top.Service, client)
	if err != nil {
		return 0, err
	}

	// upload to the latest version
	// note: latest version must not be activated already
	if *f.Sub.UseLatestVersion {
		getVersion, err := client.GetVersion(&fastly.GetVersionInput{
			Service: fastlyServiceID,
			Version: latestVersion,
		})
		if err != nil {
			return 0, err
		}

		if getVersion.Active {
			return 0, fmt.Errorf("Sorry, the latest version is already activated")
		}

		return latestVersion, nil
	}

	// otherwise clone the latest version and upload to that
	clonedVersion, err := cloneFromVersion(latestVersion, client)
	if err != nil {
		return 0, err
	}

	fmt.Printf("Successfully created new version %d from latest version %d\n\n", clonedVersion.Number, latestVersion)
	return clonedVersion.Number, nil
}

func uploadVCL(selectedVersion int, path string, client *fastly.Client, ch chan vclResponse) {
	defer wg.Done()

	name := extractName(path)
	content, err := getLocalVCL(path)

	if err != nil {
		ch <- vclResponse{
			Path:    path,
			Name:    name,
			Content: fmt.Sprintf("get local vcl error: %s", err),
			Error:   true,
		}
	} else {
		vclFile, err := client.CreateVCL(&fastly.CreateVCLInput{
			Service: fastlyServiceID,
			Version: selectedVersion,
			Name:    name,
			Content: content,
		})

		if err != nil {
			fmt.Printf("There was an error creating the file '%s':\n%s\nWe'll now try updating this file instead of creating it\n\n", common.Yellow(name), common.Red(err))

			vclFileUpdate, updateErr := client.UpdateVCL(&fastly.UpdateVCLInput{
				Service: fastlyServiceID,
				Version: selectedVersion,
				Name:    name,
				Content: content,
			})
			if updateErr != nil {
				ch <- vclResponse{
					Path:    path,
					Name:    name,
					Content: fmt.Sprintf("error: %s", updateErr),
					Error:   true,
				}
			} else {
				ch <- vclResponse{
					Path:    path,
					Name:    name,
					Content: vclFileUpdate.Content,
					Error:   false,
				}
			}
		} else {
			ch <- vclResponse{
				Path:    path,
				Name:    name,
				Content: vclFile.Content,
				Error:   false,
			}
		}
	}
}

func getLocalVCL(path string) (string, error) {
	content, err := ioutil.ReadFile(path)
	if err != nil {
		return "", err
	}
	return string(content), nil
}

func handleResponse(vr vclResponse, debug bool, selectedVersion int) {
	if vr.Error {
		fmt.Printf("Whoops, the file '%s' didn't upload to version '%d' because of the following error:\n\t%s\n\n", common.Yellow(vr.Name), selectedVersion, common.Red(vr.Content))
	} else {
		fmt.Printf("Yay, the file '%s' in version '%s' was updated successfully\n", common.Green(vr.Name), common.Yellow(selectedVersion))
	}
}
<commit_msg>Fix Upload Workflow<commit_after>package commands

import (
	"fmt"
	"io/ioutil"
	"os"
	"strconv"

	"github.com/integralist/go-fastly-cli/common"
	"github.com/integralist/go-fastly-cli/flags"

	fastly "github.com/sethvargo/go-fastly"
)

// Upload takes specified list of files and creates new remote version
// if upload fails it'll attempt uploading over existing remote version
func Upload(f flags.Flags, client *fastly.Client) {
	checkIncorrectFlagConfiguration(f)
	configureSkipMatch(f)

	// store value rather than dereference pointer multiple times later
	fastlyServiceID = *f.Top.Service

	// the acquireVersion function checks if we should...
	//
	// 		A. clone the specified version before uploading files: `-clone`
	// 		B. upload files to the specified version: `-version`
	// 		C. upload files to the latest version: `-latest`
	// 		D. clone the latest version available
	selectedVersion, err := acquireVersion(f, client)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	processFiles(selectedVersion, uploadVCL, handleResponse, f, client)
}

func checkIncorrectFlagConfiguration(f flags.Flags) {
	if *f.Sub.CloneVersion != "" && *f.Sub.UploadVersion != "" {
		fmt.Println("Please do not provide both -clone-version and -upload-version flags")
		os.Exit(1)
	}
}

func cloneFromVersion(version int, client *fastly.Client) (*fastly.Version, error) {
	clonedVersion, err := client.CloneVersion(&fastly.CloneVersionInput{
		Service: fastlyServiceID,
		Version: version,
	})
	if err != nil {
		return nil, err
	}

	return clonedVersion, nil
}

func acquireVersion(f flags.Flags, client *fastly.Client) (int, error) {
	// clone from specified version and upload to that
	if *f.Sub.CloneVersion != "" {
		cloneVersion, err := strconv.Atoi(*f.Sub.CloneVersion)
		if err != nil {
			return 0, err
		}

		clonedVersion, err := cloneFromVersion(cloneVersion, client)
		if err != nil {
			return 0, err
		}

		fmt.Printf("Successfully created new version %d from existing version %s\n\n", clonedVersion.Number, *f.Sub.CloneVersion)
		return clonedVersion.Number, nil
	}

	// upload to the specified version (it can't be activated)
	if *f.Sub.UploadVersion != "" {
		uploadVersion, err := strconv.Atoi(*f.Sub.UploadVersion)
		if err != nil {
			return 0, err
		}

		getVersion, err := client.GetVersion(&fastly.GetVersionInput{
			Service: fastlyServiceID,
			Version: uploadVersion,
		})
		if err != nil {
			return 0, err
		}

		if getVersion.Active {
			return 0, fmt.Errorf("Sorry, the specified version is already activated")
		}

		return uploadVersion, nil
	}

	latestVersion, err := common.GetLatestVCLVersion(*f.Top.Service, client)
	if err != nil {
		return 0, err
	}

	// upload to the latest version
	// note: latest version must not be activated already
	if *f.Sub.UseLatestVersion {
		getVersion, err := client.GetVersion(&fastly.GetVersionInput{
			Service: fastlyServiceID,
			Version: latestVersion,
		})
		if err != nil {
			return 0, err
		}

		if getVersion.Active {
			return 0, fmt.Errorf("Sorry, the latest version is already activated")
		}

		return latestVersion, nil
	}

	// otherwise clone the latest version and upload to that
	clonedVersion, err := cloneFromVersion(latestVersion, client)
	if err != nil {
		return 0, err
	}

	fmt.Printf("Successfully created new version %d from latest version %d\n\n", clonedVersion.Number, latestVersion)
	return clonedVersion.Number, nil
}

func uploadVCL(selectedVersion int, path string, client *fastly.Client, ch chan vclResponse) {
	defer wg.Done()

	name := extractName(path)
	content, err := getLocalVCL(path)

	if err != nil {
		ch <- vclResponse{
			Path:    path,
			Name:    name,
			Content: fmt.Sprintf("get local vcl error: %s", err),
			Error:   true,
		}
	} else {
		// First check if the local file exists already on the remote
		_, err := client.GetVCL(&fastly.GetVCLInput{
			Service: fastlyServiceID,
			Version: selectedVersion,
			Name:    name,
		})

		if err != nil {
			// If the file DOESNT exist, then we'll create it
			fmt.Printf("\n%+v\n> will attempt to create the file %s.vcl\n\n", err, name)

			vclFile, err := client.CreateVCL(&fastly.CreateVCLInput{
				Service: fastlyServiceID,
				Version: selectedVersion,
				Name:    name,
				Content: content,
			})

			if err != nil {
				fmt.Printf("There was an error creating the file '%s':\n%s\n\n", common.Yellow(name), common.Red(err))
			} else {
				ch <- vclResponse{
					Path:    path,
					Name:    name,
					Content: vclFile.Content,
					Error:   false,
				}
			}
		} else {
			// If the file DOES exist, then we'll upload our version on top of it
			vclFileUpdate, updateErr := client.UpdateVCL(&fastly.UpdateVCLInput{
				Service: fastlyServiceID,
				Version: selectedVersion,
				Name:    name,
				Content: content,
			})
			if updateErr != nil {
				ch <- vclResponse{
					Path:    path,
					Name:    name,
					Content: fmt.Sprintf("error: %s", updateErr),
					Error:   true,
				}
			} else {
				ch <- vclResponse{
					Path:    path,
					Name:    name,
					Content: vclFileUpdate.Content,
					Error:   false,
				}
			}
		}
	}
}

func getLocalVCL(path string) (string, error) {
	content, err := ioutil.ReadFile(path)
	if err != nil {
		return "", err
	}
	return string(content), nil
}

func handleResponse(vr vclResponse, debug bool, selectedVersion int) {
	if vr.Error {
		fmt.Printf("The file '%s' didn't upload to version '%d' because of the following error:\n\t%s\n\n", common.Yellow(vr.Name), selectedVersion, common.Red(vr.Content))
	} else {
		fmt.Printf("The file '%s' in version '%s' was updated successfully\n", common.Green(vr.Name), common.Yellow(selectedVersion))
	}
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"os"

	"github.com/peterhellberg/neocities/api"
)

var cmdUpload = &Command{
	Run:   runUpload,
	Usage: "upload <filename> [<another filename>]",
	Short: "Upload files to Neocities",
	Long:  "Upload files to your Neocities account",
}

func init() {
	CmdRunner.Use(cmdUpload)
}

func runUpload(cmd *Command, args *Args) {
	if args.IsParamsEmpty() {
		cmd.PrintUsage()
		os.Exit(0)
	}

	cred, err := getCredentials()
	check(err)

	files := args.Params
	response, err := api.UploadFiles(cred, files)
	check(err)

	if os.Getenv("NEOCITIES_VERBOSE") != "false" {
		response.Print()
	}

	os.Exit(0)
}

func getCredentials() (*api.Credentials, error) {
	user := os.Getenv("NEOCITIES_USER")
	pass := os.Getenv("NEOCITIES_PASS")

	return &api.Credentials{User: user, Pass: pass}, nil
}

func check(err error) {
	if err != nil {
		fmt.Println("Error:", err)

		os.Exit(1)
	}
}
<commit_msg>Check for environment variables (user/pass)<commit_after>package commands

import (
	"errors"
	"fmt"
	"os"

	"github.com/peterhellberg/neocities/api"
)

var cmdUpload = &Command{
	Run:   runUpload,
	Usage: "upload <filename> [<another filename>]",
	Short: "Upload files to Neocities",
	Long:  "Upload files to your Neocities account",
}

func init() {
	CmdRunner.Use(cmdUpload)
}

func runUpload(cmd *Command, args *Args) {
	if args.IsParamsEmpty() {
		cmd.PrintUsage()
		os.Exit(0)
	}

	cred, err := getCredentials()
	check(err)

	files := args.Params
	response, err := api.UploadFiles(cred, files)
	check(err)

	if os.Getenv("NEOCITIES_VERBOSE") != "false" {
		response.Print()
	}

	os.Exit(0)
}

func getCredentials() (*api.Credentials, error) {
	user, err := getenv("NEOCITIES_USER")
	check(err)

	pass, err := getenv("NEOCITIES_PASS")
	check(err)

	return &api.Credentials{User: user, Pass: pass}, nil
}

func getenv(variable string) (string, error) {
	value := os.Getenv(variable)

	if value == "" {
		return value, errors.New("Missing environment variable " + variable)
	}

	return value, nil
}

func check(err error) {
	if err != nil {
		fmt.Println("Error:", err)

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


<commit_before>package influxdb

import (
	"fmt"
	"math"
	"sort"
	"sync"
	"time"

	"github.com/boltdb/bolt"
	"github.com/influxdb/influxdb/influxql"
)

// tx represents a transaction that spans multiple shard data stores.
// This transaction will open and close all data stores atomically.
type tx struct {
	mu     sync.Mutex
	server *Server
	opened bool
	now    time.Time

	itrs []*shardIterator // shard iterators

	// used by DecodeFields and FieldIDs. Only used in a raw query, which won't let you select from more than one measurement
	measurement *Measurement
	decoder     fieldDecoder
}

// newTx return a new initialized Tx.
func newTx(server *Server) *tx {
	return &tx{
		server: server,
		now:    time.Now(),
	}
}

// SetNow sets the current time for the transaction.
func (tx *tx) SetNow(now time.Time) { tx.now = now }

// Open opens a read-only transaction on all data stores atomically.
func (tx *tx) Open() error {
	tx.mu.Lock()
	defer tx.mu.Unlock()

	// Mark transaction as open.
	tx.opened = true

	// Open each iterator individually. If any fail close the transaction and error out
	for _, itr := range tx.itrs {
		if err := itr.open(); err != nil {
			_ = tx.close()
			return err
		}
	}

	return nil
}

// Close closes all data store transactions atomically.
func (tx *tx) Close() error {
	tx.mu.Lock()
	defer tx.mu.Unlock()
	return tx.close()
}

func (tx *tx) close() error {
	// Mark transaction as closed.
	tx.opened = false

	for _, itr := range tx.itrs {
		_ = itr.close()
	}

	return nil
}

// CreateIterators returns an iterator for a simple select statement.
func (tx *tx) CreateIterators(stmt *influxql.SelectStatement) ([]influxql.Iterator, error) {
	// Parse the source segments.
	database, policyName, measurement, err := splitIdent(stmt.Source.(*influxql.Measurement).Name)
	if err != nil {
		return nil, err
	}

	// Grab time range from statement.
	tmin, tmax := influxql.TimeRange(stmt.Condition)
	if tmin.IsZero() {
		tmin = time.Unix(0, 1)
	}
	if tmax.IsZero() {
		tmax = tx.now
	}

	// Find database and retention policy.
	db := tx.server.databases[database]
	if db == nil {
		return nil, ErrDatabaseNotFound
	}
	rp := db.policies[policyName]
	if rp == nil {
		return nil, ErrRetentionPolicyNotFound
	}

	// Find measurement.
	m, err := tx.server.measurement(database, measurement)
	if err != nil {
		return nil, err
	}
	if m == nil {
		return nil, ErrMeasurementNotFound
	}
	tx.measurement = m

	// Find shard groups within time range.
	var shardGroups []*ShardGroup
	for _, group := range rp.shardGroups {
		if group.Contains(tmin, tmax) {
			shardGroups = append(shardGroups, group)
		}
	}
	if len(shardGroups) == 0 {
		return nil, nil
	}

	// Normalize dimensions to extract the interval.
	_, dimensions, err := stmt.Dimensions.Normalize()
	if err != nil {
		return nil, err
	}

	// Find field.
	fieldName := stmt.Fields[0].Expr.(*influxql.VarRef).Val
	f := m.FieldByName(fieldName)
	if f == nil {
		return nil, fmt.Errorf("field not found: %s", fieldName)
	}
	tagSets := m.tagSets(stmt, dimensions)

	// Get a field decoder.
	d := NewFieldCodec(m)
	tx.decoder = d

	// if it's a raw data query we'll need to pass these to the series cursor
	var fieldIDs []uint8
	var fieldNames []string
	if stmt.RawQuery {
		fieldIDs, _ = tx.FieldIDs(stmt.Fields)
		fieldNames = tx.fieldNames(stmt.Fields)
	}

	// limit the number of series in this query if they specified a limit
	if stmt.Limit > 0 {
		if stmt.Offset > len(tagSets) {
			return nil, nil
		}

		limitSets := make(map[string]map[uint32]influxql.Expr)
		orderedSets := make([]string, 0, len(tagSets))
		for k, _ := range tagSets {
			orderedSets = append(orderedSets, k)
		}
		sort.Strings(orderedSets)

		if stmt.Offset+stmt.Limit > len(orderedSets) {
			stmt.Limit = len(orderedSets) - stmt.Offset
		}

		sets := orderedSets[stmt.Offset : stmt.Offset+stmt.Limit]
		for _, s := range sets {
			limitSets[s] = tagSets[s]
		}
		tagSets = limitSets
	}

	// Create an iterator for every shard.
	var itrs []influxql.Iterator
	for tag, set := range tagSets {
		for _, group := range shardGroups {
			// TODO: only create iterators for the shards we actually have to hit in a group
			for _, sh := range group.Shards {

				// create a series cursor for each unique series id
				cursors := make([]*seriesCursor, 0, len(set))
				for id, cond := range set {
					c := &seriesCursor{id: id, condition: cond, decoder: d, rawQuery: stmt.RawQuery}
					if stmt.RawQuery {
						c.fieldIDs = fieldIDs
						c.fieldNames = fieldNames
						c.tx = tx
					}
					cursors = append(cursors, c)
				}

				// create the shard iterator that will map over all series for the shard
				itr := &shardIterator{
					measurement: m,
					fieldName:   f.Name,
					fieldID:     f.ID,
					tags:        tag,
					db:          sh.store,
					cursors:     cursors,
					tmin:        tmin.UnixNano(),
					tmax:        tmax.UnixNano(),
				}

				// Add to tx so the bolt transaction can be opened/closed.
				tx.itrs = append(tx.itrs, itr)

				itrs = append(itrs, itr)
			}
		}
	}

	return itrs, nil
}

// DecodeValues is for use in a raw data query
func (tx *tx) DecodeValues(fieldIDs []uint8, timestamp int64, data []byte) []interface{} {
	vals := make([]interface{}, len(fieldIDs))
	for i, id := range fieldIDs {
		v, _ := tx.decoder.DecodeByID(id, data)
		vals[i] = v
	}
	return vals
}

// FieldIDs will take an array of fields and return the id associated with each
func (tx *tx) FieldIDs(fields []*influxql.Field) ([]uint8, error) {
	names := tx.fieldNames(fields)
	ids := make([]uint8, len(names))

	for i, n := range names {
		field := tx.measurement.FieldByName(n)
		if field == nil {
			return nil, ErrFieldNotFound
		}
		ids[i] = field.ID
	}

	return ids, nil
}

func (tx *tx) fieldNames(fields []*influxql.Field) []string {
	var a []string
	for _, f := range fields {
		if v, ok := f.Expr.(*influxql.VarRef); ok { // this is a raw query so we handle it differently
			a = append(a, v.Val)
		}
	}
	return a
}

// splitIdent splits an identifier into it's database, policy, and measurement parts.
func splitIdent(s string) (db, rp, m string, err error) {
	a, err := influxql.SplitIdent(s)
	if err != nil {
		return "", "", "", err
	} else if len(a) != 3 {
		return "", "", "", fmt.Errorf("invalid ident, expected 3 segments: %q", s)
	}
	return a[0], a[1], a[2], nil
}

// shardIterator represents an iterator for traversing over a single series.
type shardIterator struct {
	fieldName   string
	fieldID     uint8
	measurement *Measurement
	tags        string // encoded dimensional tag values
	cursors     []*seriesCursor
	keyValues   []keyValue
	db          *bolt.DB // data stores by shard id
	txn         *bolt.Tx // read transactions by shard id
	tmin, tmax  int64
}

func (i *shardIterator) open() error {
	// Open the data store
	txn, err := i.db.Begin(false)
	if err != nil {
		return err
	}
	i.txn = txn

	// Open cursors for each series id
	for _, c := range i.cursors {
		b := i.txn.Bucket(u32tob(c.id))
		if b == nil {
			continue
		}

		c.cur = b.Cursor()
	}

	i.keyValues = make([]keyValue, len(i.cursors))
	for j, cur := range i.cursors {
		i.keyValues[j].key, i.keyValues[j].data, i.keyValues[j].value = cur.Next(i.fieldName, i.fieldID, i.tmin, i.tmax)
	}

	return nil
}

func (i *shardIterator) close() error {
	_ = i.txn.Rollback()
	return nil
}

func (i *shardIterator) Tags() string { return i.tags }

func (i *shardIterator) Next() (key int64, data []byte, value interface{}) {
	min := -1
	minKey := int64(math.MaxInt64)

	for ind, kv := range i.keyValues {
		if kv.key != 0 && kv.key < i.tmax && kv.key < minKey {
			min = ind
			minKey = kv.key
		}
	}

	// if min is -1 we've exhausted all cursors for the given time range
	if min == -1 {
		return 0, nil, nil
	}

	kv := i.keyValues[min]
	key = kv.key
	data = kv.data
	value = kv.value

	i.keyValues[min].key, i.keyValues[min].data, i.keyValues[min].value = i.cursors[min].Next(i.fieldName, i.fieldID, i.tmin, i.tmax)
	return key, data, value
}

type keyValue struct {
	key   int64
	data  []byte
	value interface{}
}

type fieldDecoder interface {
	DecodeByID(fieldID uint8, b []byte) (interface{}, error)
}

type seriesCursor struct {
	id          uint32
	condition   influxql.Expr
	cur         *bolt.Cursor
	initialized bool
	decoder     fieldDecoder
	// these are for raw data queries
	tx         *tx
	rawQuery   bool
	fieldIDs   []uint8
	fieldNames []string
}

func (c *seriesCursor) Next(fieldName string, fieldID uint8, tmin, tmax int64) (key int64, data []byte, value interface{}) {
	// TODO: clean this up when we make it so series ids are only queried against the shards they exist in.
	//       Right now we query for all series ids on a query against each shard, even if that shard may not have the
	//       data, so cur could be nil.
	if c.cur == nil {
		return 0, nil, nil
	}

	for {
		var k, v []byte
		if !c.initialized {
			k, v = c.cur.Seek(u64tob(uint64(tmin)))
			c.initialized = true
		} else {
			k, v = c.cur.Next()
		}

		// Exit if there is no more data.
		if k == nil {
			return 0, nil, nil
		}

		// Marshal key & value.
		key := int64(btou64(k))

		if key > tmax {
			return 0, nil, nil
		}

		// if it's a raw query we handle things differently
		if c.rawQuery {
			// we'll need to marshal all the field values if the condition isn't nil
			if c.condition != nil {
				fieldValues := make(map[string]interface{})
				values := c.tx.DecodeValues(c.fieldIDs, 0, data)
				for i, val := range values {
					fieldValues[c.fieldNames[i]] = val
				}
				if ok, _ := influxql.Eval(c.condition, fieldValues).(bool); !ok {
					continue
				}
			}

			// no condition so yield all data by default
			return key, v, nil
		}

		value, err := c.decoder.DecodeByID(fieldID, v)
		if err != nil {
			continue
		}

		// Skip to the next if we don't have a field value for this field for this point
		if value == nil {
			continue
		}

		// Evaluate condition. Move to next key/value if non-true.
		if c.condition != nil {
			if ok, _ := influxql.Eval(c.condition, map[string]interface{}{fieldName: value}).(bool); !ok {
				continue
			}
		}

		return key, v, value
	}
}
<commit_msg>Skip timestamp while checking conditions<commit_after>package influxdb

import (
	"fmt"
	"math"
	"sort"
	"sync"
	"time"

	"github.com/boltdb/bolt"
	"github.com/influxdb/influxdb/influxql"
)

// tx represents a transaction that spans multiple shard data stores.
// This transaction will open and close all data stores atomically.
type tx struct {
	mu     sync.Mutex
	server *Server
	opened bool
	now    time.Time

	itrs []*shardIterator // shard iterators

	// used by DecodeFields and FieldIDs. Only used in a raw query, which won't let you select from more than one measurement
	measurement *Measurement
	decoder     fieldDecoder
}

// newTx return a new initialized Tx.
func newTx(server *Server) *tx {
	return &tx{
		server: server,
		now:    time.Now(),
	}
}

// SetNow sets the current time for the transaction.
func (tx *tx) SetNow(now time.Time) { tx.now = now }

// Open opens a read-only transaction on all data stores atomically.
func (tx *tx) Open() error {
	tx.mu.Lock()
	defer tx.mu.Unlock()

	// Mark transaction as open.
	tx.opened = true

	// Open each iterator individually. If any fail close the transaction and error out
	for _, itr := range tx.itrs {
		if err := itr.open(); err != nil {
			_ = tx.close()
			return err
		}
	}

	return nil
}

// Close closes all data store transactions atomically.
func (tx *tx) Close() error {
	tx.mu.Lock()
	defer tx.mu.Unlock()
	return tx.close()
}

func (tx *tx) close() error {
	// Mark transaction as closed.
	tx.opened = false

	for _, itr := range tx.itrs {
		_ = itr.close()
	}

	return nil
}

// CreateIterators returns an iterator for a simple select statement.
func (tx *tx) CreateIterators(stmt *influxql.SelectStatement) ([]influxql.Iterator, error) {
	// Parse the source segments.
	database, policyName, measurement, err := splitIdent(stmt.Source.(*influxql.Measurement).Name)
	if err != nil {
		return nil, err
	}

	// Grab time range from statement.
	tmin, tmax := influxql.TimeRange(stmt.Condition)
	if tmin.IsZero() {
		tmin = time.Unix(0, 1)
	}
	if tmax.IsZero() {
		tmax = tx.now
	}

	// Find database and retention policy.
	db := tx.server.databases[database]
	if db == nil {
		return nil, ErrDatabaseNotFound
	}
	rp := db.policies[policyName]
	if rp == nil {
		return nil, ErrRetentionPolicyNotFound
	}

	// Find measurement.
	m, err := tx.server.measurement(database, measurement)
	if err != nil {
		return nil, err
	}
	if m == nil {
		return nil, ErrMeasurementNotFound
	}
	tx.measurement = m

	// Find shard groups within time range.
	var shardGroups []*ShardGroup
	for _, group := range rp.shardGroups {
		if group.Contains(tmin, tmax) {
			shardGroups = append(shardGroups, group)
		}
	}
	if len(shardGroups) == 0 {
		return nil, nil
	}

	// Normalize dimensions to extract the interval.
	_, dimensions, err := stmt.Dimensions.Normalize()
	if err != nil {
		return nil, err
	}

	// Find field.
	fieldName := stmt.Fields[0].Expr.(*influxql.VarRef).Val
	f := m.FieldByName(fieldName)
	if f == nil {
		return nil, fmt.Errorf("field not found: %s", fieldName)
	}
	tagSets := m.tagSets(stmt, dimensions)

	// Get a field decoder.
	d := NewFieldCodec(m)
	tx.decoder = d

	// if it's a raw data query we'll need to pass these to the series cursor
	var fieldIDs []uint8
	var fieldNames []string
	if stmt.RawQuery {
		fieldIDs, _ = tx.FieldIDs(stmt.Fields)
		fieldNames = tx.fieldNames(stmt.Fields)
	}

	// limit the number of series in this query if they specified a limit
	if stmt.Limit > 0 {
		if stmt.Offset > len(tagSets) {
			return nil, nil
		}

		limitSets := make(map[string]map[uint32]influxql.Expr)
		orderedSets := make([]string, 0, len(tagSets))
		for k, _ := range tagSets {
			orderedSets = append(orderedSets, k)
		}
		sort.Strings(orderedSets)

		if stmt.Offset+stmt.Limit > len(orderedSets) {
			stmt.Limit = len(orderedSets) - stmt.Offset
		}

		sets := orderedSets[stmt.Offset : stmt.Offset+stmt.Limit]
		for _, s := range sets {
			limitSets[s] = tagSets[s]
		}
		tagSets = limitSets
	}

	// Create an iterator for every shard.
	var itrs []influxql.Iterator
	for tag, set := range tagSets {
		for _, group := range shardGroups {
			// TODO: only create iterators for the shards we actually have to hit in a group
			for _, sh := range group.Shards {

				// create a series cursor for each unique series id
				cursors := make([]*seriesCursor, 0, len(set))
				for id, cond := range set {
					c := &seriesCursor{id: id, condition: cond, decoder: d, rawQuery: stmt.RawQuery}
					if stmt.RawQuery {
						c.fieldIDs = fieldIDs
						c.fieldNames = fieldNames
						c.tx = tx
					}
					cursors = append(cursors, c)
				}

				// create the shard iterator that will map over all series for the shard
				itr := &shardIterator{
					measurement: m,
					fieldName:   f.Name,
					fieldID:     f.ID,
					tags:        tag,
					db:          sh.store,
					cursors:     cursors,
					tmin:        tmin.UnixNano(),
					tmax:        tmax.UnixNano(),
				}

				// Add to tx so the bolt transaction can be opened/closed.
				tx.itrs = append(tx.itrs, itr)

				itrs = append(itrs, itr)
			}
		}
	}

	return itrs, nil
}

// DecodeValues is for use in a raw data query
func (tx *tx) DecodeValues(fieldIDs []uint8, timestamp int64, data []byte) []interface{} {
	vals := make([]interface{}, len(fieldIDs)+1)
	vals[0] = timestamp
	for i, id := range fieldIDs {
		v, _ := tx.decoder.DecodeByID(id, data)
		vals[i+1] = v
	}
	return vals
}

// FieldIDs will take an array of fields and return the id associated with each
func (tx *tx) FieldIDs(fields []*influxql.Field) ([]uint8, error) {
	names := tx.fieldNames(fields)
	ids := make([]uint8, len(names))

	for i, n := range names {
		field := tx.measurement.FieldByName(n)
		if field == nil {
			return nil, ErrFieldNotFound
		}
		ids[i] = field.ID
	}

	return ids, nil
}

func (tx *tx) fieldNames(fields []*influxql.Field) []string {
	var a []string
	for _, f := range fields {
		if v, ok := f.Expr.(*influxql.VarRef); ok { // this is a raw query so we handle it differently
			a = append(a, v.Val)
		}
	}
	return a
}

// splitIdent splits an identifier into it's database, policy, and measurement parts.
func splitIdent(s string) (db, rp, m string, err error) {
	a, err := influxql.SplitIdent(s)
	if err != nil {
		return "", "", "", err
	} else if len(a) != 3 {
		return "", "", "", fmt.Errorf("invalid ident, expected 3 segments: %q", s)
	}
	return a[0], a[1], a[2], nil
}

// shardIterator represents an iterator for traversing over a single series.
type shardIterator struct {
	fieldName   string
	fieldID     uint8
	measurement *Measurement
	tags        string // encoded dimensional tag values
	cursors     []*seriesCursor
	keyValues   []keyValue
	db          *bolt.DB // data stores by shard id
	txn         *bolt.Tx // read transactions by shard id
	tmin, tmax  int64
}

func (i *shardIterator) open() error {
	// Open the data store
	txn, err := i.db.Begin(false)
	if err != nil {
		return err
	}
	i.txn = txn

	// Open cursors for each series id
	for _, c := range i.cursors {
		b := i.txn.Bucket(u32tob(c.id))
		if b == nil {
			continue
		}

		c.cur = b.Cursor()
	}

	i.keyValues = make([]keyValue, len(i.cursors))
	for j, cur := range i.cursors {
		i.keyValues[j].key, i.keyValues[j].data, i.keyValues[j].value = cur.Next(i.fieldName, i.fieldID, i.tmin, i.tmax)
	}

	return nil
}

func (i *shardIterator) close() error {
	_ = i.txn.Rollback()
	return nil
}

func (i *shardIterator) Tags() string { return i.tags }

func (i *shardIterator) Next() (key int64, data []byte, value interface{}) {
	min := -1
	minKey := int64(math.MaxInt64)

	for ind, kv := range i.keyValues {
		if kv.key != 0 && kv.key < i.tmax && kv.key < minKey {
			min = ind
			minKey = kv.key
		}
	}

	// if min is -1 we've exhausted all cursors for the given time range
	if min == -1 {
		return 0, nil, nil
	}

	kv := i.keyValues[min]
	key = kv.key
	data = kv.data
	value = kv.value

	i.keyValues[min].key, i.keyValues[min].data, i.keyValues[min].value = i.cursors[min].Next(i.fieldName, i.fieldID, i.tmin, i.tmax)
	return key, data, value
}

type keyValue struct {
	key   int64
	data  []byte
	value interface{}
}

type fieldDecoder interface {
	DecodeByID(fieldID uint8, b []byte) (interface{}, error)
}

type seriesCursor struct {
	id          uint32
	condition   influxql.Expr
	cur         *bolt.Cursor
	initialized bool
	decoder     fieldDecoder
	// these are for raw data queries
	tx         *tx
	rawQuery   bool
	fieldIDs   []uint8
	fieldNames []string
}

func (c *seriesCursor) Next(fieldName string, fieldID uint8, tmin, tmax int64) (key int64, data []byte, value interface{}) {
	// TODO: clean this up when we make it so series ids are only queried against the shards they exist in.
	//       Right now we query for all series ids on a query against each shard, even if that shard may not have the
	//       data, so cur could be nil.
	if c.cur == nil {
		return 0, nil, nil
	}

	for {
		var k, v []byte
		if !c.initialized {
			k, v = c.cur.Seek(u64tob(uint64(tmin)))
			c.initialized = true
		} else {
			k, v = c.cur.Next()
		}

		// Exit if there is no more data.
		if k == nil {
			return 0, nil, nil
		}

		// Marshal key & value.
		key := int64(btou64(k))

		if key > tmax {
			return 0, nil, nil
		}

		// if it's a raw query we handle things differently
		if c.rawQuery {
			// we'll need to marshal all the field values if the condition isn't nil
			if c.condition != nil {
				fieldValues := make(map[string]interface{})
				values := c.tx.DecodeValues(c.fieldIDs, 0, data)
				for i, val := range values[1:] { // Skip the timestamp.
					fieldValues[c.fieldNames[i]] = val
				}
				if ok, _ := influxql.Eval(c.condition, fieldValues).(bool); !ok {
					continue
				}
			}

			// no condition so yield all data by default
			return key, v, nil
		}

		value, err := c.decoder.DecodeByID(fieldID, v)
		if err != nil {
			continue
		}

		// Skip to the next if we don't have a field value for this field for this point
		if value == nil {
			continue
		}

		// Evaluate condition. Move to next key/value if non-true.
		if c.condition != nil {
			if ok, _ := influxql.Eval(c.condition, map[string]interface{}{fieldName: value}).(bool); !ok {
				continue
			}
		}

		return key, v, value
	}
}
<|endoftext|>


<commit_before>package deploystack

import (
	"bufio"
	"fmt"
	"os"
	"sort"
	"strings"
)

// ImageManage promps a user to select a disk type.
func ImageManage(project string) (string, error) {
	fmt.Println(Divider)
	fmt.Printf("There are a large number of machine images to choose from. For more infomration, \n")
	fmt.Printf("please refer to the following link for more infomation about machine images.\n")
	fmt.Printf("%shttps://cloud.google.com/compute/docs/images%s\n", TERMCYANB, TERMCLEAR)
	fmt.Println(Divider)

	colorPrintln("Choose an operating system.", TERMCYANB)
	ImageTypeProject := listSelect(DiskProjects, DefaultImageProject)

	fmt.Printf("Polling for %s images...\n", ImageTypeProject.Value)
	images, err := images(project, ImageTypeProject.Value)
	if err != nil {
		return "", err
	}

	families := getListOfImageFamilies(images)

	colorPrintln("Choose a disk family to use for this application.", TERMCYANB)
	family := listSelect(families, DefaultImageFamily)

	imagesByFam := getListOfImageTypesByFamily(images, ImageTypeProject.Value, family.Value)

	colorPrintln("Choose a disk type to use for this application.", TERMCYANB)
	result := listSelect(imagesByFam, imagesByFam[len(imagesByFam)-1].Value)

	return result.Value, nil
}

func colorPrintln(msg, color string) {
	fmt.Printf("%s%s %s\n", color, msg, TERMCLEAR)
}

func MachineTypeManage(project, zone string) (string, error) {
	fmt.Println(Divider)
	fmt.Printf("There are a large number of machine types to choose from. For more infomration, \n")
	fmt.Printf("please refer to the following link for more infomation about Machine types.\n")
	fmt.Printf("%shttps://cloud.google.com/compute/docs/machine-types%s\n", TERMCYANB, TERMCLEAR)
	fmt.Println(Divider)

	fmt.Printf("Polling for machine types...\n")
	types, err := machineTypes(project, zone)
	if err != nil {
		return "", fmt.Errorf("error polling for machine types : %s", err)
	}

	typefamilies := getListOfMachineTypeFamily(types)

	fmt.Printf("Choose an Machine Type Family\n")
	familyProject := listSelect(typefamilies, DefaultMachineType)

	filteredtypes := getListOfMachineTypeByFamily(types, familyProject.Value)

	fmt.Printf("%sChoose a machine type to use for this application. %s\n", TERMCYANB, TERMCLEAR)
	result := listSelect(filteredtypes, filteredtypes[0].Value)

	return result.Value, nil
}

type GCEInstanceConfig map[string]string

func (gce GCEInstanceConfig) Print(title string) {
	keys := []string{}
	for i := range gce {
		keys = append(keys, i)
	}

	longest := longestLength(toLabeledValueSlice(keys))

	colorPrintln(title, TERMCYANREV)
	exclude := []string{}

	if s, ok := gce["instance-name"]; ok && len(s) > 0 {
		printSetting("instance-name", s, longest)
		exclude = append(exclude, "instance-name")
	}

	if s, ok := gce["region"]; ok && len(s) > 0 {
		printSetting("region", s, longest)
		exclude = append(exclude, "region")
	}

	if s, ok := gce["zone"]; ok && len(s) > 0 {
		printSetting("zone", s, longest)
		exclude = append(exclude, "zone")
	}

	ordered := []string{}
	for i, v := range gce {
		if strings.Contains(strings.Join(exclude, " "), i) {
			continue
		}
		if len(v) < 1 {
			continue
		}

		ordered = append(ordered, i)
	}
	sort.Strings(ordered)

	for i := range ordered {
		key := ordered[i]
		printSetting(key, gce[key], longest)
	}
}

func GCEInstanceManage(project, basename string) (GCEInstanceConfig, error) {
	var err error
	configs := make(map[string]string)

	defaultImage, err := getLatestImage(project, DefaultImageProject, DefaultImageFamily)
	if err != nil {
		return configs, err
	}

	defaultConfig := GCEInstanceConfig{
		"instance-image":        defaultImage,
		"instance-disksize":     "200",
		"instance-disktype":     "pd-standard",
		"instance-tags":         "[http-server,https-server]",
		"instance-name":         fmt.Sprintf("%s-instance", basename),
		"region":                DefaultRegion,
		"zone":                  fmt.Sprintf("%s-a", DefaultRegion),
		"instance-machine-type": "n1-standard-1",
	}

	ClearScreen()
	fmt.Println(Divider)
	colorPrintln("Configure a Compute Engine Instance", TERMCYANB)
	fmt.Printf("Let's walk through configuring a Compute Engine Instance (Virtual Machine). \n")
	fmt.Printf("you can either accept a default configuration with settings that work for \n")
	fmt.Printf("trying out most use cases, or hand configure key settings. \n")
	fmt.Println(Divider)

	defaultConfig.Print("Default Configuration")

	chooseDefault := Custom{
		Name:        "choosedefault",
		Description: "Do you want to use the default? ('No' means custom)",
		Default:     "yes",
		Validation:  "yesorno",
	}

	if err := chooseDefault.Collect(); err != nil {
		return configs, err
	}

	if chooseDefault.Value == "yes" {
		return defaultConfig, nil
	}

	nameItem := Custom{
		Name:        "name",
		Description: "Enter the name of the instance",
		Default:     fmt.Sprintf("%s-instance", basename),
	}

	if err := nameItem.Collect(); err != nil {
		return configs, err
	}

	configs["instance-name"] = nameItem.Value

	configs["region"], err = RegionManage(project, "compute", DefaultRegion)
	if err != nil {
		return configs, err
	}

	configs["zone"], err = ZoneManage(project, configs["region"])
	if err != nil {
		return configs, err
	}

	configs["instance-machine-type"], err = MachineTypeManage(project, configs["zone"])
	if err != nil {
		return configs, err
	}
	configs["instance-image"], err = ImageManage(project)
	if err != nil {
		return configs, err
	}

	items := Customs{
		{Name: "instance-disksize", Description: "Enter the size of the boot disk you want in GB", Default: "200", Validation: "integer"},
		{Name: "instance-disktype", Description: "Enter the type of the boot disk you want", Default: "pd-standard", Options: []string{"pd-standard", "pd-balanced", "pd-ssd"}},
		{Name: "webserver", Description: "Do you want this to be a webserver (Expose ports 80 & 443)? ", Default: "no", Validation: "yesorno"},
	}

	if err := items.Collect(); err != nil {
		return configs, err
	}

	for _, v := range items {

		if v.Name == "webserver" {
			configs["instance-tags"] = "[]"
			if v.Value == "yes" {
				configs["instance-tags"] = "[http-server,https-server]"
			}
			continue
		}

		configs[v.Name] = v.Value

	}

	return configs, nil
}

// BillingAccountManage either grabs the users only BillingAccount or
// presents a list of BillingAccounts to select from.
func BillingAccountManage() (string, error) {
	accounts, err := billingAccounts()
	if err != nil {
		return "", fmt.Errorf("could not get list of billing accounts: %s", err)
	}

	labeled := []string{}

	for _, v := range accounts {
		labeled = append(labeled, fmt.Sprintf("%s (%s)", v.DisplayName, strings.ReplaceAll(v.Name, "billingAccounts/", "")))
	}

	if len(accounts) == 1 {
		fmt.Printf("\nOnly found one billing account. Using : %s%s%s.\n", TERMCYAN, accounts[0].DisplayName, TERMCLEAR)
		return extractAccount(labeled[0]), nil
	}

	fmt.Printf("\n%sPlease select one of your billing accounts to use with this project%s.\n", TERMCYAN, TERMCLEAR)
	result := listSelect(toLabeledValueSlice(labeled), labeled[0])

	return extractAccount(result.Value), nil
}

func extractAccount(s string) string {
	sl := strings.Split(s, "(")
	return strings.ReplaceAll(sl[1], ")", "")
}

// ProjectManage promps a user to select a project.
func ProjectManage() (string, string, error) {
	createString := "CREATE NEW PROJECT"
	project, err := ProjectID()
	if err != nil {
		return "", "", err
	}

	projects, err := projects()
	if err != nil {
		return "", "", err
	}

	lvs := LabeledValues{}

	for _, v := range projects {
		lv := LabeledValue{Label: v.Name, Value: v.ID}

		if !v.BillingEnabled {
			lv.Label = fmt.Sprintf("%s%s (Billing Disabled)%s", TERMGREY, v.Name, TERMCLEAR)
		}

		lvs = append(lvs, lv)
	}

	fmt.Printf("\n%sChoose a project to use for this application.%s\n\n", TERMCYANB, TERMCLEAR)
	fmt.Printf("%sNOTE:%s This app will make changes to the project. %s\n", TERMCYANREV, TERMCYAN, TERMCLEAR)
	fmt.Printf("While those changes are reverseable, it would be better to put it in a fresh new project. \n")

	lv := listSelect(lvs, project)
	project = lv.Value

	if project == createString {
		project, err = projectPrompt()
		if err != nil {
			return "", "", err
		}
	}

	if err := ProjectIDSet(project); err != nil {
		return lv.Value, lv.Label, fmt.Errorf("error: unable to set project (%s) in environment: %s", project, err)
	}

	return lv.Value, lv.Label, nil
}

// projectPrompt manages the interaction of creating a project, including prompts.
func projectPrompt() (string, error) {
	result := ""
	sec1 := NewSection("Creating the project")

	sec1.Open()
	fmt.Printf("Project IDs are immutable and can be set only during project \n")
	fmt.Printf("creation. They must start with a lowercase letter and can have \n")
	fmt.Printf("lowercase ASCII letters, digits or hyphens.  \n")
	fmt.Printf("Project IDs must be between 6 and 30 characters. \n")
	fmt.Printf("%sPlease enter a new project name to create: %s\n", TERMCYANB, TERMCLEAR)

	reader := bufio.NewReader(os.Stdin)

	for {
		fmt.Print("> ")
		text, _ := reader.ReadString('\n')
		text = strings.Replace(text, "\n", "", -1)

		if len(text) == 0 {
			fmt.Printf("%sPlease enter a new project name to create: %s\n", TERMCYANB, TERMCLEAR)
			continue
		}

		if err := projectCreate(text); err != nil {
			fmt.Printf("%sProject name could not be created %s\n", TERMREDREV, TERMCLEAR)
			fmt.Printf("%sReason: %s %s\n", TERMREDB, err, TERMCLEAR)
			fmt.Printf("%sPlease choose another. %s\n", TERMREDREV, TERMCLEAR)
			continue
		}

		fmt.Printf("Project Created\n")
		result = text
		break

	}
	sec1.Close()

	sec2 := NewSection("Activating Billing for the project")
	sec2.Open()
	account, err := BillingAccountManage()
	if err != nil {
		return "", fmt.Errorf("could not determine proper billing account: %s ", err)
	}

	if err := BillingAccountProjectAttach(result, account); err != nil {
		return "", fmt.Errorf("could not link billing account: %s ", err)
	}
	sec2.Close()
	return result, nil
}

// regions will return a list of regions depending on product type
func regions(project, product string) ([]string, error) {
	switch product {
	case "compute":
		return regionsCompute(project)
	case "functions":
		return regionsFunctions(project)
	case "run":
		return regionsRun(project)
	}

	return []string{}, fmt.Errorf("invalid product requested: %s", product)
}

// RegionManage promps a user to select a region.
func RegionManage(project, product, def string) (string, error) {
	fmt.Printf("Polling for regions...\n")
	regions, err := regions(project, product)
	if err != nil {
		return "", err
	}
	fmt.Printf("%sChoose a valid region to use for this application. %s\n", TERMCYANB, TERMCLEAR)
	region := listSelect(toLabeledValueSlice(regions), def)

	return region.Value, nil
}

// ZoneManage promps a user to select a zone.
func ZoneManage(project, region string) (string, error) {
	fmt.Printf("Polling for zones...\n")
	zones, err := zones(project, region)
	if err != nil {
		return "", err
	}

	fmt.Printf("%sChoose a valid zone to use for this application. %s\n", TERMCYANB, TERMCLEAR)
	zone := listSelect(toLabeledValueSlice(zones), zones[0])
	return zone.Value, nil
}

// Start presents a little documentation screen which also prevents the user
// from timing out the request to activate Cloud Shell
func Start() {
	fmt.Printf(Divider)
	colorPrintln("Deploystack", TERMCYANB)
	fmt.Printf("Deploystack will walk you through setting some options for the  \n")
	fmt.Printf("stack this solutions installs. \n")
	fmt.Printf("Most questions have a default that you can choose by hitting the Enter key  \n")
	fmt.Printf(Divider)
	colorPrintln("Press the Enter Key to continue", TERMCYANB)
	var input string
	fmt.Scanln(&input)
}
<commit_msg>fix: making Create Project to work again<commit_after>package deploystack

import (
	"bufio"
	"fmt"
	"os"
	"sort"
	"strings"
)

// ImageManage promps a user to select a disk type.
func ImageManage(project string) (string, error) {
	fmt.Println(Divider)
	fmt.Printf("There are a large number of machine images to choose from. For more infomration, \n")
	fmt.Printf("please refer to the following link for more infomation about machine images.\n")
	fmt.Printf("%shttps://cloud.google.com/compute/docs/images%s\n", TERMCYANB, TERMCLEAR)
	fmt.Println(Divider)

	colorPrintln("Choose an operating system.", TERMCYANB)
	ImageTypeProject := listSelect(DiskProjects, DefaultImageProject)

	fmt.Printf("Polling for %s images...\n", ImageTypeProject.Value)
	images, err := images(project, ImageTypeProject.Value)
	if err != nil {
		return "", err
	}

	families := getListOfImageFamilies(images)

	colorPrintln("Choose a disk family to use for this application.", TERMCYANB)
	family := listSelect(families, DefaultImageFamily)

	imagesByFam := getListOfImageTypesByFamily(images, ImageTypeProject.Value, family.Value)

	colorPrintln("Choose a disk type to use for this application.", TERMCYANB)
	result := listSelect(imagesByFam, imagesByFam[len(imagesByFam)-1].Value)

	return result.Value, nil
}

func colorPrintln(msg, color string) {
	fmt.Printf("%s%s %s\n", color, msg, TERMCLEAR)
}

func MachineTypeManage(project, zone string) (string, error) {
	fmt.Println(Divider)
	fmt.Printf("There are a large number of machine types to choose from. For more infomration, \n")
	fmt.Printf("please refer to the following link for more infomation about Machine types.\n")
	fmt.Printf("%shttps://cloud.google.com/compute/docs/machine-types%s\n", TERMCYANB, TERMCLEAR)
	fmt.Println(Divider)

	fmt.Printf("Polling for machine types...\n")
	types, err := machineTypes(project, zone)
	if err != nil {
		return "", fmt.Errorf("error polling for machine types : %s", err)
	}

	typefamilies := getListOfMachineTypeFamily(types)

	fmt.Printf("Choose an Machine Type Family\n")
	familyProject := listSelect(typefamilies, DefaultMachineType)

	filteredtypes := getListOfMachineTypeByFamily(types, familyProject.Value)

	fmt.Printf("%sChoose a machine type to use for this application. %s\n", TERMCYANB, TERMCLEAR)
	result := listSelect(filteredtypes, filteredtypes[0].Value)

	return result.Value, nil
}

type GCEInstanceConfig map[string]string

func (gce GCEInstanceConfig) Print(title string) {
	keys := []string{}
	for i := range gce {
		keys = append(keys, i)
	}

	longest := longestLength(toLabeledValueSlice(keys))

	colorPrintln(title, TERMCYANREV)
	exclude := []string{}

	if s, ok := gce["instance-name"]; ok && len(s) > 0 {
		printSetting("instance-name", s, longest)
		exclude = append(exclude, "instance-name")
	}

	if s, ok := gce["region"]; ok && len(s) > 0 {
		printSetting("region", s, longest)
		exclude = append(exclude, "region")
	}

	if s, ok := gce["zone"]; ok && len(s) > 0 {
		printSetting("zone", s, longest)
		exclude = append(exclude, "zone")
	}

	ordered := []string{}
	for i, v := range gce {
		if strings.Contains(strings.Join(exclude, " "), i) {
			continue
		}
		if len(v) < 1 {
			continue
		}

		ordered = append(ordered, i)
	}
	sort.Strings(ordered)

	for i := range ordered {
		key := ordered[i]
		printSetting(key, gce[key], longest)
	}
}

func GCEInstanceManage(project, basename string) (GCEInstanceConfig, error) {
	var err error
	configs := make(map[string]string)

	defaultImage, err := getLatestImage(project, DefaultImageProject, DefaultImageFamily)
	if err != nil {
		return configs, err
	}

	defaultConfig := GCEInstanceConfig{
		"instance-image":        defaultImage,
		"instance-disksize":     "200",
		"instance-disktype":     "pd-standard",
		"instance-tags":         "[http-server,https-server]",
		"instance-name":         fmt.Sprintf("%s-instance", basename),
		"region":                DefaultRegion,
		"zone":                  fmt.Sprintf("%s-a", DefaultRegion),
		"instance-machine-type": "n1-standard-1",
	}

	ClearScreen()
	fmt.Println(Divider)
	colorPrintln("Configure a Compute Engine Instance", TERMCYANB)
	fmt.Printf("Let's walk through configuring a Compute Engine Instance (Virtual Machine). \n")
	fmt.Printf("you can either accept a default configuration with settings that work for \n")
	fmt.Printf("trying out most use cases, or hand configure key settings. \n")
	fmt.Println(Divider)

	defaultConfig.Print("Default Configuration")

	chooseDefault := Custom{
		Name:        "choosedefault",
		Description: "Do you want to use the default? ('No' means custom)",
		Default:     "yes",
		Validation:  "yesorno",
	}

	if err := chooseDefault.Collect(); err != nil {
		return configs, err
	}

	if chooseDefault.Value == "yes" {
		return defaultConfig, nil
	}

	nameItem := Custom{
		Name:        "name",
		Description: "Enter the name of the instance",
		Default:     fmt.Sprintf("%s-instance", basename),
	}

	if err := nameItem.Collect(); err != nil {
		return configs, err
	}

	configs["instance-name"] = nameItem.Value

	configs["region"], err = RegionManage(project, "compute", DefaultRegion)
	if err != nil {
		return configs, err
	}

	configs["zone"], err = ZoneManage(project, configs["region"])
	if err != nil {
		return configs, err
	}

	configs["instance-machine-type"], err = MachineTypeManage(project, configs["zone"])
	if err != nil {
		return configs, err
	}
	configs["instance-image"], err = ImageManage(project)
	if err != nil {
		return configs, err
	}

	items := Customs{
		{Name: "instance-disksize", Description: "Enter the size of the boot disk you want in GB", Default: "200", Validation: "integer"},
		{Name: "instance-disktype", Description: "Enter the type of the boot disk you want", Default: "pd-standard", Options: []string{"pd-standard", "pd-balanced", "pd-ssd"}},
		{Name: "webserver", Description: "Do you want this to be a webserver (Expose ports 80 & 443)? ", Default: "no", Validation: "yesorno"},
	}

	if err := items.Collect(); err != nil {
		return configs, err
	}

	for _, v := range items {

		if v.Name == "webserver" {
			configs["instance-tags"] = "[]"
			if v.Value == "yes" {
				configs["instance-tags"] = "[http-server,https-server]"
			}
			continue
		}

		configs[v.Name] = v.Value

	}

	return configs, nil
}

// BillingAccountManage either grabs the users only BillingAccount or
// presents a list of BillingAccounts to select from.
func BillingAccountManage() (string, error) {
	accounts, err := billingAccounts()
	if err != nil {
		return "", fmt.Errorf("could not get list of billing accounts: %s", err)
	}

	labeled := []string{}

	for _, v := range accounts {
		labeled = append(labeled, fmt.Sprintf("%s (%s)", v.DisplayName, strings.ReplaceAll(v.Name, "billingAccounts/", "")))
	}

	if len(accounts) == 1 {
		fmt.Printf("\nOnly found one billing account. Using : %s%s%s.\n", TERMCYAN, accounts[0].DisplayName, TERMCLEAR)
		return extractAccount(labeled[0]), nil
	}

	fmt.Printf("\n%sPlease select one of your billing accounts to use with this project%s.\n", TERMCYAN, TERMCLEAR)
	result := listSelect(toLabeledValueSlice(labeled), labeled[0])

	return extractAccount(result.Value), nil
}

func extractAccount(s string) string {
	sl := strings.Split(s, "(")
	return strings.ReplaceAll(sl[1], ")", "")
}

// ProjectManage promps a user to select a project.
func ProjectManage() (string, string, error) {
	createString := "CREATE NEW PROJECT"
	project, err := ProjectID()
	if err != nil {
		return "", "", err
	}

	projects, err := projects()
	if err != nil {
		return "", "", err
	}

	lvs := LabeledValues{}

	for _, v := range projects {
		lv := LabeledValue{Label: v.Name, Value: v.ID}

		if !v.BillingEnabled {
			lv.Label = fmt.Sprintf("%s%s (Billing Disabled)%s", TERMGREY, v.Name, TERMCLEAR)
		}

		lvs = append(lvs, lv)
	}

	lvs = append([]LabeledValue{{createString, createString}}, lvs...)

	fmt.Printf("\n%sChoose a project to use for this application.%s\n\n", TERMCYANB, TERMCLEAR)
	fmt.Printf("%sNOTE:%s This app will make changes to the project. %s\n", TERMCYANREV, TERMCYAN, TERMCLEAR)
	fmt.Printf("While those changes are reverseable, it would be better to put it in a fresh new project. \n")

	lv := listSelect(lvs, project)
	project = lv.Value

	if project == createString {
		project, err = projectPrompt()
		if err != nil {
			return "", "", err
		}
		lv = LabeledValue{project, project}
	}

	if err := ProjectIDSet(project); err != nil {
		return lv.Value, lv.Label, fmt.Errorf("error: unable to set project (%s) in environment: %s", project, err)
	}

	return lv.Value, lv.Label, nil
}

// projectPrompt manages the interaction of creating a project, including prompts.
func projectPrompt() (string, error) {
	result := ""
	sec1 := NewSection("Creating the project")

	sec1.Open()
	fmt.Printf("Project IDs are immutable and can be set only during project \n")
	fmt.Printf("creation. They must start with a lowercase letter and can have \n")
	fmt.Printf("lowercase ASCII letters, digits or hyphens.  \n")
	fmt.Printf("Project IDs must be between 6 and 30 characters. \n")
	fmt.Printf("%sPlease enter a new project name to create: %s\n", TERMCYANB, TERMCLEAR)

	reader := bufio.NewReader(os.Stdin)

	for {
		fmt.Print("> ")
		text, _ := reader.ReadString('\n')
		text = strings.Replace(text, "\n", "", -1)

		if len(text) == 0 {
			fmt.Printf("%sPlease enter a new project name to create: %s\n", TERMCYANB, TERMCLEAR)
			continue
		}

		if err := projectCreate(text); err != nil {
			fmt.Printf("%sProject name could not be created %s\n", TERMREDREV, TERMCLEAR)
			fmt.Printf("%sReason: %s %s\n", TERMREDB, err, TERMCLEAR)
			fmt.Printf("%sPlease choose another. %s\n", TERMREDREV, TERMCLEAR)
			continue
		}

		fmt.Printf("Project Created\n")
		result = text
		break

	}
	sec1.Close()

	sec2 := NewSection("Activating Billing for the project")
	sec2.Open()
	account, err := BillingAccountManage()
	if err != nil {
		return "", fmt.Errorf("could not determine proper billing account: %s ", err)
	}

	if err := BillingAccountProjectAttach(result, account); err != nil {
		return "", fmt.Errorf("could not link billing account: %s ", err)
	}
	sec2.Close()
	return result, nil
}

// regions will return a list of regions depending on product type
func regions(project, product string) ([]string, error) {
	switch product {
	case "compute":
		return regionsCompute(project)
	case "functions":
		return regionsFunctions(project)
	case "run":
		return regionsRun(project)
	}

	return []string{}, fmt.Errorf("invalid product requested: %s", product)
}

// RegionManage promps a user to select a region.
func RegionManage(project, product, def string) (string, error) {
	fmt.Printf("Polling for regions...\n")
	regions, err := regions(project, product)
	if err != nil {
		return "", err
	}
	fmt.Printf("%sChoose a valid region to use for this application. %s\n", TERMCYANB, TERMCLEAR)
	region := listSelect(toLabeledValueSlice(regions), def)

	return region.Value, nil
}

// ZoneManage promps a user to select a zone.
func ZoneManage(project, region string) (string, error) {
	fmt.Printf("Polling for zones...\n")
	zones, err := zones(project, region)
	if err != nil {
		return "", err
	}

	fmt.Printf("%sChoose a valid zone to use for this application. %s\n", TERMCYANB, TERMCLEAR)
	zone := listSelect(toLabeledValueSlice(zones), zones[0])
	return zone.Value, nil
}

// Start presents a little documentation screen which also prevents the user
// from timing out the request to activate Cloud Shell
func Start() {
	fmt.Printf(Divider)
	colorPrintln("Deploystack", TERMCYANB)
	fmt.Printf("Deploystack will walk you through setting some options for the  \n")
	fmt.Printf("stack this solutions installs. \n")
	fmt.Printf("Most questions have a default that you can choose by hitting the Enter key  \n")
	fmt.Printf(Divider)
	colorPrintln("Press the Enter Key to continue", TERMCYANB)
	var input string
	fmt.Scanln(&input)
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2020 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package js_test

import (
	"context"
	"net/url"
	"testing"
	"time"

	"github.com/loadimpact/k6/js"
	"github.com/loadimpact/k6/js/common"
	"github.com/loadimpact/k6/js/internal/modules"
	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/testutils"
	"github.com/loadimpact/k6/loader"
	"github.com/loadimpact/k6/stats"
	"github.com/spf13/afero"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gopkg.in/guregu/null.v3"
)

type CheckModule struct {
	t             testing.TB
	initCtxCalled int
	vuCtxCalled   int
}

func (cm *CheckModule) InitCtx(ctx context.Context) {
	cm.initCtxCalled++
	assert.NotNil(cm.t, common.GetRuntime(ctx))
	assert.NotNil(cm.t, common.GetInitEnv(ctx))
	assert.Nil(cm.t, lib.GetState(ctx))
}

func (cm *CheckModule) VuCtx(ctx context.Context) {
	cm.vuCtxCalled++
	assert.NotNil(cm.t, common.GetRuntime(ctx))
	assert.Nil(cm.t, common.GetInitEnv(ctx))
	assert.NotNil(cm.t, lib.GetState(ctx))
}

func TestNewJSRunnerWithCustomModule(t *testing.T) {
	t.Skip()
	checkModule := &CheckModule{t: t}
	modules.Register("k6/check", checkModule)

	script := `
		var check = require("k6/check");
		check.initCtx();

		module.exports.options = { vus: 1, iterations: 1 };
		module.exports.default = function() {
			check.vuCtx();
		};
	`

	logger := testutils.NewLogger(t)
	rtOptions := lib.RuntimeOptions{CompatibilityMode: null.StringFrom("base")}
	runner, err := js.New(
		logger,
		&loader.SourceData{
			URL:  &url.URL{Path: "blah", Scheme: "file"},
			Data: []byte(script),
		},
		map[string]afero.Fs{"file": afero.NewMemMapFs(), "https": afero.NewMemMapFs()},
		rtOptions,
	)
	require.NoError(t, err)
	assert.Equal(t, checkModule.initCtxCalled, 1)
	assert.Equal(t, checkModule.vuCtxCalled, 0)

	vu, err := runner.NewVU(1, make(chan stats.SampleContainer, 100))
	require.NoError(t, err)
	assert.Equal(t, checkModule.initCtxCalled, 2)
	assert.Equal(t, checkModule.vuCtxCalled, 0)

	vuCtx, vuCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer vuCancel()

	activeVU := vu.Activate(&lib.VUActivationParams{RunContext: vuCtx})
	require.NoError(t, activeVU.RunOnce())
	assert.Equal(t, checkModule.initCtxCalled, 2)
	assert.Equal(t, checkModule.vuCtxCalled, 1)
	require.NoError(t, activeVU.RunOnce())
	assert.Equal(t, checkModule.initCtxCalled, 2)
	assert.Equal(t, checkModule.vuCtxCalled, 2)

	arc := runner.MakeArchive()
	assert.Equal(t, checkModule.initCtxCalled, 2) // shouldn't change, we're not executing the init context again
	assert.Equal(t, checkModule.vuCtxCalled, 2)

	runnerFromArc, err := js.NewFromArchive(logger, arc, rtOptions)
	require.NoError(t, err)
	assert.Equal(t, checkModule.initCtxCalled, 3) // changes because we need to get the exported functions
	assert.Equal(t, checkModule.vuCtxCalled, 2)
	vuFromArc, err := runnerFromArc.NewVU(2, make(chan stats.SampleContainer, 100))
	require.NoError(t, err)
	assert.Equal(t, checkModule.initCtxCalled, 4)
	assert.Equal(t, checkModule.vuCtxCalled, 2)
	activeVUFromArc := vuFromArc.Activate(&lib.VUActivationParams{RunContext: vuCtx})
	require.NoError(t, activeVUFromArc.RunOnce())
	assert.Equal(t, checkModule.initCtxCalled, 4)
	assert.Equal(t, checkModule.vuCtxCalled, 3)
	require.NoError(t, activeVUFromArc.RunOnce())
	assert.Equal(t, checkModule.initCtxCalled, 4)
	assert.Equal(t, checkModule.vuCtxCalled, 4)
}
<commit_msg>Reenable TestNewJSRunnerWithCustomModule<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2020 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package js_test

import (
	"context"
	"fmt"
	"net/url"
	"sync/atomic"
	"testing"
	"time"

	"github.com/loadimpact/k6/js"
	"github.com/loadimpact/k6/js/common"
	"github.com/loadimpact/k6/js/internal/modules"
	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/testutils"
	"github.com/loadimpact/k6/loader"
	"github.com/loadimpact/k6/stats"
	"github.com/spf13/afero"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gopkg.in/guregu/null.v3"
)

type CheckModule struct {
	t             testing.TB
	initCtxCalled int
	vuCtxCalled   int
}

func (cm *CheckModule) InitCtx(ctx context.Context) {
	cm.initCtxCalled++
	assert.NotNil(cm.t, common.GetRuntime(ctx))
	assert.NotNil(cm.t, common.GetInitEnv(ctx))
	assert.Nil(cm.t, lib.GetState(ctx))
}

func (cm *CheckModule) VuCtx(ctx context.Context) {
	cm.vuCtxCalled++
	assert.NotNil(cm.t, common.GetRuntime(ctx))
	assert.Nil(cm.t, common.GetInitEnv(ctx))
	assert.NotNil(cm.t, lib.GetState(ctx))
}

var uniqueModuleNumber int64 //nolint:gochecknoglobals

func TestNewJSRunnerWithCustomModule(t *testing.T) {
	t.Parallel()
	checkModule := &CheckModule{t: t}
	moduleName := fmt.Sprintf("k6/check-%d", atomic.AddInt64(&uniqueModuleNumber, 1))
	modules.Register(moduleName, checkModule)

	script := fmt.Sprintf(`
		var check = require("%s");
		check.initCtx();

		module.exports.options = { vus: 1, iterations: 1 };
		module.exports.default = function() {
			check.vuCtx();
		};
	`, moduleName)

	logger := testutils.NewLogger(t)
	rtOptions := lib.RuntimeOptions{CompatibilityMode: null.StringFrom("base")}
	runner, err := js.New(
		logger,
		&loader.SourceData{
			URL:  &url.URL{Path: "blah", Scheme: "file"},
			Data: []byte(script),
		},
		map[string]afero.Fs{"file": afero.NewMemMapFs(), "https": afero.NewMemMapFs()},
		rtOptions,
	)
	require.NoError(t, err)
	assert.Equal(t, checkModule.initCtxCalled, 1)
	assert.Equal(t, checkModule.vuCtxCalled, 0)

	vu, err := runner.NewVU(1, make(chan stats.SampleContainer, 100))
	require.NoError(t, err)
	assert.Equal(t, checkModule.initCtxCalled, 2)
	assert.Equal(t, checkModule.vuCtxCalled, 0)

	vuCtx, vuCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer vuCancel()

	activeVU := vu.Activate(&lib.VUActivationParams{RunContext: vuCtx})
	require.NoError(t, activeVU.RunOnce())
	assert.Equal(t, checkModule.initCtxCalled, 2)
	assert.Equal(t, checkModule.vuCtxCalled, 1)
	require.NoError(t, activeVU.RunOnce())
	assert.Equal(t, checkModule.initCtxCalled, 2)
	assert.Equal(t, checkModule.vuCtxCalled, 2)

	arc := runner.MakeArchive()
	assert.Equal(t, checkModule.initCtxCalled, 2) // shouldn't change, we're not executing the init context again
	assert.Equal(t, checkModule.vuCtxCalled, 2)

	runnerFromArc, err := js.NewFromArchive(logger, arc, rtOptions)
	require.NoError(t, err)
	assert.Equal(t, checkModule.initCtxCalled, 3) // changes because we need to get the exported functions
	assert.Equal(t, checkModule.vuCtxCalled, 2)
	vuFromArc, err := runnerFromArc.NewVU(2, make(chan stats.SampleContainer, 100))
	require.NoError(t, err)
	assert.Equal(t, checkModule.initCtxCalled, 4)
	assert.Equal(t, checkModule.vuCtxCalled, 2)
	activeVUFromArc := vuFromArc.Activate(&lib.VUActivationParams{RunContext: vuCtx})
	require.NoError(t, activeVUFromArc.RunOnce())
	assert.Equal(t, checkModule.initCtxCalled, 4)
	assert.Equal(t, checkModule.vuCtxCalled, 3)
	require.NoError(t, activeVUFromArc.RunOnce())
	assert.Equal(t, checkModule.initCtxCalled, 4)
	assert.Equal(t, checkModule.vuCtxCalled, 4)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/constant"
	"go/printer"
	"go/token"
	"go/types"
	"log"
	"os"
	"path/filepath"

	"golang.org/x/tools/go/ast/astutil"
	"golang.org/x/tools/go/buildutil"
	"golang.org/x/tools/go/loader"
)

var (
	fRecursive bool
	fOneLine   bool
	fJSON      bool
	fMinify    bool
)

func init() {
	flag.BoolVar(&fRecursive, "r", false, "keyify struct initializers recursively")
	flag.BoolVar(&fOneLine, "o", false, "print new struct initializer on a single line")
	flag.BoolVar(&fJSON, "json", false, "print new struct initializer as JSON")
	flag.BoolVar(&fMinify, "m", false, "omit fields that are set to their zero value")
}

func usage() {
	fmt.Printf("Usage: %s [flags] <position>\n\n", os.Args[0])
	flag.PrintDefaults()
}

func main() {
	log.SetFlags(0)
	flag.Usage = usage
	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(2)
	}
	pos := flag.Args()[0]
	name, start, _, err := parsePos(pos)
	if err != nil {
		log.Fatal(err)
	}
	eval, err := filepath.EvalSymlinks(name)
	if err != nil {
		log.Fatal(err)
	}
	name, err = filepath.Abs(eval)
	if err != nil {
		log.Fatal(err)
	}
	cwd, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	bpkg, err := buildutil.ContainingPackage(&build.Default, cwd, name)
	if err != nil {
		log.Fatal(err)
	}
	conf := &loader.Config{}
	conf.TypeCheckFuncBodies = func(s string) bool {
		return s == bpkg.ImportPath || s == bpkg.ImportPath+"_test"
	}
	conf.ImportWithTests(bpkg.ImportPath)
	lprog, err := conf.Load()
	if err != nil {
		log.Fatal(err)
	}
	var tf *token.File
	var af *ast.File
	pkg := lprog.InitialPackages()[0]
	for _, ff := range pkg.Files {
		file := lprog.Fset.File(ff.Pos())
		if file.Name() == name {
			af = ff
			tf = file
			break
		}
	}
	tstart, tend, err := fileOffsetToPos(tf, start, start)
	if err != nil {
		log.Fatal(err)
	}
	path, _ := astutil.PathEnclosingInterval(af, tstart, tend)
	var complit *ast.CompositeLit
	for _, p := range path {
		if p, ok := p.(*ast.CompositeLit); ok {
			complit = p
			break
		}
	}
	if complit == nil {
		log.Fatal("no composite literal found near point")
	}
	if len(complit.Elts) == 0 {
		printComplit(complit, complit, lprog.Fset, lprog.Fset)
		return
	}
	if _, ok := complit.Elts[0].(*ast.KeyValueExpr); ok {
		lit := complit
		if fOneLine {
			lit = copyExpr(complit, 1).(*ast.CompositeLit)
		}
		printComplit(complit, lit, lprog.Fset, lprog.Fset)
		return
	}
	_, ok := pkg.TypeOf(complit.Type).Underlying().(*types.Struct)
	if !ok {
		log.Fatal("not a struct initialiser")
		return
	}

	newComplit, lines := keyify(pkg, complit)
	newFset := token.NewFileSet()
	newFile := newFset.AddFile("", -1, lines)
	for i := 1; i <= lines; i++ {
		newFile.AddLine(i)
	}
	printComplit(complit, newComplit, lprog.Fset, newFset)
}

func keyify(
	pkg *loader.PackageInfo,
	complit *ast.CompositeLit,
) (*ast.CompositeLit, int) {
	var calcPos func(int) token.Pos
	if fOneLine {
		calcPos = func(int) token.Pos { return token.Pos(1) }
	} else {
		calcPos = func(i int) token.Pos { return token.Pos(2 + i) }
	}

	st, _ := pkg.TypeOf(complit.Type).Underlying().(*types.Struct)
	newComplit := &ast.CompositeLit{
		Type:   complit.Type,
		Lbrace: 1,
		Rbrace: token.Pos(st.NumFields() + 2),
	}
	if fOneLine {
		newComplit.Rbrace = 1
	}
	numLines := 2 + st.NumFields()
	n := 0
	for i := 0; i < st.NumFields(); i++ {
		field := st.Field(i)
		val := complit.Elts[i]
		if fRecursive {
			if val2, ok := val.(*ast.CompositeLit); ok {
				if _, ok := pkg.TypeOf(val2.Type).Underlying().(*types.Struct); ok {
					var lines int
					numLines += lines
					val, lines = keyify(pkg, val2)
				}
			}
		}
		_, isIface := st.Field(i).Type().Underlying().(*types.Interface)
		if fMinify && (isNil(val, pkg) || (!isIface && isZero(val, pkg))) {
			continue
		}
		elt := &ast.KeyValueExpr{
			Key:   &ast.Ident{NamePos: calcPos(n), Name: field.Name()},
			Value: copyExpr(val, calcPos(n)),
		}
		newComplit.Elts = append(newComplit.Elts, elt)
		n++
	}
	return newComplit, numLines
}

func isNil(val ast.Expr, pkg *loader.PackageInfo) bool {
	ident, ok := val.(*ast.Ident)
	if !ok {
		return false
	}
	if _, ok := pkg.ObjectOf(ident).(*types.Nil); ok {
		return true
	}
	if c, ok := pkg.ObjectOf(ident).(*types.Const); ok {
		if c.Val().Kind() != constant.Bool {
			return false
		}
		return !constant.BoolVal(c.Val())
	}
	return false
}

func isZero(val ast.Expr, pkg *loader.PackageInfo) bool {
	switch val := val.(type) {
	case *ast.BasicLit:
		switch val.Value {
		case `""`, "``", "0", "0.0", "0i", "0.":
			return true
		default:
			return false
		}
	case *ast.Ident:
		return isNil(val, pkg)
	case *ast.CompositeLit:
		typ := pkg.TypeOf(val.Type)
		if typ == nil {
			return false
		}
		isIface := false
		switch typ := typ.Underlying().(type) {
		case *types.Struct:
		case *types.Array:
			_, isIface = typ.Elem().Underlying().(*types.Interface)
		default:
			return false
		}
		for _, elt := range val.Elts {
			if isNil(elt, pkg) || (!isIface && !isZero(elt, pkg)) {
				return false
			}
		}
		return true
	}
	return false
}

func printComplit(oldlit, newlit *ast.CompositeLit, oldfset, newfset *token.FileSet) {
	buf := &bytes.Buffer{}
	cfg := printer.Config{Mode: printer.UseSpaces | printer.TabIndent, Tabwidth: 8}
	_ = cfg.Fprint(buf, newfset, newlit)
	if fJSON {
		output := struct {
			Start       int    `json:"start"`
			End         int    `json:"end"`
			Replacement string `json:"replacement"`
		}{
			oldfset.Position(oldlit.Pos()).Offset,
			oldfset.Position(oldlit.End()).Offset,
			buf.String(),
		}
		_ = json.NewEncoder(os.Stdout).Encode(output)
	} else {
		fmt.Println(buf.String())
	}
}

func copyExpr(expr ast.Expr, line token.Pos) ast.Expr {
	switch expr := expr.(type) {
	case *ast.BasicLit:
		cp := *expr
		cp.ValuePos = 0
		return &cp
	case *ast.BinaryExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.OpPos = 0
		cp.Y = copyExpr(cp.Y, line)
		return &cp
	case *ast.CallExpr:
		cp := *expr
		cp.Fun = copyExpr(cp.Fun, line)
		cp.Lparen = 0
		for i, v := range cp.Args {
			cp.Args[i] = copyExpr(v, line)
		}
		if cp.Ellipsis != 0 {
			cp.Ellipsis = line
		}
		cp.Rparen = 0
		return &cp
	case *ast.CompositeLit:
		cp := *expr
		cp.Type = copyExpr(cp.Type, line)
		cp.Lbrace = 0
		for i, v := range cp.Elts {
			cp.Elts[i] = copyExpr(v, line)
		}
		cp.Rbrace = 0
		return &cp
	case *ast.Ident:
		cp := *expr
		cp.NamePos = 0
		return &cp
	case *ast.IndexExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Lbrack = 0
		cp.Index = copyExpr(cp.Index, line)
		cp.Rbrack = 0
		return &cp
	case *ast.KeyValueExpr:
		cp := *expr
		cp.Key = copyExpr(cp.Key, line)
		cp.Colon = 0
		cp.Value = copyExpr(cp.Value, line)
		return &cp
	case *ast.ParenExpr:
		cp := *expr
		cp.Lparen = 0
		cp.X = copyExpr(cp.X, line)
		cp.Rparen = 0
		return &cp
	case *ast.SelectorExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Sel = copyExpr(cp.Sel, line).(*ast.Ident)
		return &cp
	case *ast.SliceExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Lbrack = 0
		cp.Low = copyExpr(cp.Low, line)
		cp.High = copyExpr(cp.High, line)
		cp.Max = copyExpr(cp.Max, line)
		cp.Rbrack = 0
		return &cp
	case *ast.StarExpr:
		cp := *expr
		cp.Star = 0
		cp.X = copyExpr(cp.X, line)
		return &cp
	case *ast.TypeAssertExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Lparen = 0
		cp.Type = copyExpr(cp.Type, line)
		cp.Rparen = 0
		return &cp
	case *ast.UnaryExpr:
		cp := *expr
		cp.OpPos = 0
		cp.X = copyExpr(cp.X, line)
		return &cp
	case *ast.MapType:
		cp := *expr
		cp.Map = 0
		cp.Key = copyExpr(cp.Key, line)
		cp.Value = copyExpr(cp.Value, line)
		return &cp
	case *ast.ArrayType:
		cp := *expr
		cp.Lbrack = 0
		cp.Len = copyExpr(cp.Len, line)
		cp.Elt = copyExpr(cp.Elt, line)
		return &cp
	case *ast.Ellipsis:
		cp := *expr
		cp.Elt = copyExpr(cp.Elt, line)
		cp.Ellipsis = line
		return &cp
	case *ast.InterfaceType:
		cp := *expr
		cp.Interface = 0
		return &cp
	case nil:
		return nil
	default:
		panic(fmt.Sprintf("shouldn't happen: unknown ast.Expr of type %T", expr))
	}
	return nil
}
<commit_msg>Don't crash when seeing anonymous struct<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/constant"
	"go/printer"
	"go/token"
	"go/types"
	"log"
	"os"
	"path/filepath"

	"golang.org/x/tools/go/ast/astutil"
	"golang.org/x/tools/go/buildutil"
	"golang.org/x/tools/go/loader"
)

var (
	fRecursive bool
	fOneLine   bool
	fJSON      bool
	fMinify    bool
)

func init() {
	flag.BoolVar(&fRecursive, "r", false, "keyify struct initializers recursively")
	flag.BoolVar(&fOneLine, "o", false, "print new struct initializer on a single line")
	flag.BoolVar(&fJSON, "json", false, "print new struct initializer as JSON")
	flag.BoolVar(&fMinify, "m", false, "omit fields that are set to their zero value")
}

func usage() {
	fmt.Printf("Usage: %s [flags] <position>\n\n", os.Args[0])
	flag.PrintDefaults()
}

func main() {
	log.SetFlags(0)
	flag.Usage = usage
	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(2)
	}
	pos := flag.Args()[0]
	name, start, _, err := parsePos(pos)
	if err != nil {
		log.Fatal(err)
	}
	eval, err := filepath.EvalSymlinks(name)
	if err != nil {
		log.Fatal(err)
	}
	name, err = filepath.Abs(eval)
	if err != nil {
		log.Fatal(err)
	}
	cwd, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	bpkg, err := buildutil.ContainingPackage(&build.Default, cwd, name)
	if err != nil {
		log.Fatal(err)
	}
	conf := &loader.Config{}
	conf.TypeCheckFuncBodies = func(s string) bool {
		return s == bpkg.ImportPath || s == bpkg.ImportPath+"_test"
	}
	conf.ImportWithTests(bpkg.ImportPath)
	lprog, err := conf.Load()
	if err != nil {
		log.Fatal(err)
	}
	var tf *token.File
	var af *ast.File
	pkg := lprog.InitialPackages()[0]
	for _, ff := range pkg.Files {
		file := lprog.Fset.File(ff.Pos())
		if file.Name() == name {
			af = ff
			tf = file
			break
		}
	}
	tstart, tend, err := fileOffsetToPos(tf, start, start)
	if err != nil {
		log.Fatal(err)
	}
	path, _ := astutil.PathEnclosingInterval(af, tstart, tend)
	var complit *ast.CompositeLit
	for _, p := range path {
		if p, ok := p.(*ast.CompositeLit); ok {
			complit = p
			break
		}
	}
	if complit == nil {
		log.Fatal("no composite literal found near point")
	}
	if len(complit.Elts) == 0 {
		printComplit(complit, complit, lprog.Fset, lprog.Fset)
		return
	}
	if _, ok := complit.Elts[0].(*ast.KeyValueExpr); ok {
		lit := complit
		if fOneLine {
			lit = copyExpr(complit, 1).(*ast.CompositeLit)
		}
		printComplit(complit, lit, lprog.Fset, lprog.Fset)
		return
	}
	_, ok := pkg.TypeOf(complit.Type).Underlying().(*types.Struct)
	if !ok {
		log.Fatal("not a struct initialiser")
		return
	}

	newComplit, lines := keyify(pkg, complit)
	newFset := token.NewFileSet()
	newFile := newFset.AddFile("", -1, lines)
	for i := 1; i <= lines; i++ {
		newFile.AddLine(i)
	}
	printComplit(complit, newComplit, lprog.Fset, newFset)
}

func keyify(
	pkg *loader.PackageInfo,
	complit *ast.CompositeLit,
) (*ast.CompositeLit, int) {
	var calcPos func(int) token.Pos
	if fOneLine {
		calcPos = func(int) token.Pos { return token.Pos(1) }
	} else {
		calcPos = func(i int) token.Pos { return token.Pos(2 + i) }
	}

	st, _ := pkg.TypeOf(complit.Type).Underlying().(*types.Struct)
	newComplit := &ast.CompositeLit{
		Type:   complit.Type,
		Lbrace: 1,
		Rbrace: token.Pos(st.NumFields() + 2),
	}
	if fOneLine {
		newComplit.Rbrace = 1
	}
	numLines := 2 + st.NumFields()
	n := 0
	for i := 0; i < st.NumFields(); i++ {
		field := st.Field(i)
		val := complit.Elts[i]
		if fRecursive {
			if val2, ok := val.(*ast.CompositeLit); ok {
				if _, ok := pkg.TypeOf(val2.Type).Underlying().(*types.Struct); ok {
					var lines int
					numLines += lines
					val, lines = keyify(pkg, val2)
				}
			}
		}
		_, isIface := st.Field(i).Type().Underlying().(*types.Interface)
		if fMinify && (isNil(val, pkg) || (!isIface && isZero(val, pkg))) {
			continue
		}
		elt := &ast.KeyValueExpr{
			Key:   &ast.Ident{NamePos: calcPos(n), Name: field.Name()},
			Value: copyExpr(val, calcPos(n)),
		}
		newComplit.Elts = append(newComplit.Elts, elt)
		n++
	}
	return newComplit, numLines
}

func isNil(val ast.Expr, pkg *loader.PackageInfo) bool {
	ident, ok := val.(*ast.Ident)
	if !ok {
		return false
	}
	if _, ok := pkg.ObjectOf(ident).(*types.Nil); ok {
		return true
	}
	if c, ok := pkg.ObjectOf(ident).(*types.Const); ok {
		if c.Val().Kind() != constant.Bool {
			return false
		}
		return !constant.BoolVal(c.Val())
	}
	return false
}

func isZero(val ast.Expr, pkg *loader.PackageInfo) bool {
	switch val := val.(type) {
	case *ast.BasicLit:
		switch val.Value {
		case `""`, "``", "0", "0.0", "0i", "0.":
			return true
		default:
			return false
		}
	case *ast.Ident:
		return isNil(val, pkg)
	case *ast.CompositeLit:
		typ := pkg.TypeOf(val.Type)
		if typ == nil {
			return false
		}
		isIface := false
		switch typ := typ.Underlying().(type) {
		case *types.Struct:
		case *types.Array:
			_, isIface = typ.Elem().Underlying().(*types.Interface)
		default:
			return false
		}
		for _, elt := range val.Elts {
			if isNil(elt, pkg) || (!isIface && !isZero(elt, pkg)) {
				return false
			}
		}
		return true
	}
	return false
}

func printComplit(oldlit, newlit *ast.CompositeLit, oldfset, newfset *token.FileSet) {
	buf := &bytes.Buffer{}
	cfg := printer.Config{Mode: printer.UseSpaces | printer.TabIndent, Tabwidth: 8}
	_ = cfg.Fprint(buf, newfset, newlit)
	if fJSON {
		output := struct {
			Start       int    `json:"start"`
			End         int    `json:"end"`
			Replacement string `json:"replacement"`
		}{
			oldfset.Position(oldlit.Pos()).Offset,
			oldfset.Position(oldlit.End()).Offset,
			buf.String(),
		}
		_ = json.NewEncoder(os.Stdout).Encode(output)
	} else {
		fmt.Println(buf.String())
	}
}

func copyExpr(expr ast.Expr, line token.Pos) ast.Expr {
	switch expr := expr.(type) {
	case *ast.BasicLit:
		cp := *expr
		cp.ValuePos = 0
		return &cp
	case *ast.BinaryExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.OpPos = 0
		cp.Y = copyExpr(cp.Y, line)
		return &cp
	case *ast.CallExpr:
		cp := *expr
		cp.Fun = copyExpr(cp.Fun, line)
		cp.Lparen = 0
		for i, v := range cp.Args {
			cp.Args[i] = copyExpr(v, line)
		}
		if cp.Ellipsis != 0 {
			cp.Ellipsis = line
		}
		cp.Rparen = 0
		return &cp
	case *ast.CompositeLit:
		cp := *expr
		cp.Type = copyExpr(cp.Type, line)
		cp.Lbrace = 0
		for i, v := range cp.Elts {
			cp.Elts[i] = copyExpr(v, line)
		}
		cp.Rbrace = 0
		return &cp
	case *ast.Ident:
		cp := *expr
		cp.NamePos = 0
		return &cp
	case *ast.IndexExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Lbrack = 0
		cp.Index = copyExpr(cp.Index, line)
		cp.Rbrack = 0
		return &cp
	case *ast.KeyValueExpr:
		cp := *expr
		cp.Key = copyExpr(cp.Key, line)
		cp.Colon = 0
		cp.Value = copyExpr(cp.Value, line)
		return &cp
	case *ast.ParenExpr:
		cp := *expr
		cp.Lparen = 0
		cp.X = copyExpr(cp.X, line)
		cp.Rparen = 0
		return &cp
	case *ast.SelectorExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Sel = copyExpr(cp.Sel, line).(*ast.Ident)
		return &cp
	case *ast.SliceExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Lbrack = 0
		cp.Low = copyExpr(cp.Low, line)
		cp.High = copyExpr(cp.High, line)
		cp.Max = copyExpr(cp.Max, line)
		cp.Rbrack = 0
		return &cp
	case *ast.StarExpr:
		cp := *expr
		cp.Star = 0
		cp.X = copyExpr(cp.X, line)
		return &cp
	case *ast.TypeAssertExpr:
		cp := *expr
		cp.X = copyExpr(cp.X, line)
		cp.Lparen = 0
		cp.Type = copyExpr(cp.Type, line)
		cp.Rparen = 0
		return &cp
	case *ast.UnaryExpr:
		cp := *expr
		cp.OpPos = 0
		cp.X = copyExpr(cp.X, line)
		return &cp
	case *ast.MapType:
		cp := *expr
		cp.Map = 0
		cp.Key = copyExpr(cp.Key, line)
		cp.Value = copyExpr(cp.Value, line)
		return &cp
	case *ast.ArrayType:
		cp := *expr
		cp.Lbrack = 0
		cp.Len = copyExpr(cp.Len, line)
		cp.Elt = copyExpr(cp.Elt, line)
		return &cp
	case *ast.Ellipsis:
		cp := *expr
		cp.Elt = copyExpr(cp.Elt, line)
		cp.Ellipsis = line
		return &cp
	case *ast.InterfaceType:
		cp := *expr
		cp.Interface = 0
		return &cp
	case *ast.StructType:
		cp := *expr
		cp.Struct = 0
		return &cp
	case nil:
		return nil
	default:
		panic(fmt.Sprintf("shouldn't happen: unknown ast.Expr of type %T", expr))
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2018 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package sqlbuilder

import (
	"context"
	"database/sql"

	"github.com/issue9/orm/core"
	"github.com/issue9/orm/fetch"
)

// SelectStmt 查询语句
type SelectStmt struct {
	engine    core.Engine
	table     string
	where     *WhereStmt
	cols      []string
	distinct  string
	forupdate bool

	joins  []*join
	orders *core.StringBuilder
	group  string

	havingQuery string
	havingVals  []interface{}

	limitQuery string
	limitVals  []interface{}
}

type join struct {
	typ   string
	on    string
	table string
}

// Select 声明一条 Select 语句
func Select(e core.Engine) *SelectStmt {
	return &SelectStmt{
		engine: e,
		where:  newWhereStmt(),
	}
}

// Distinct 声明一条 Select 语句的 Distinct
func (stmt *SelectStmt) Distinct(col string) *SelectStmt {
	stmt.distinct = col
	return stmt
}

// Reset 重置语句
func (stmt *SelectStmt) Reset() {
	stmt.table = ""
	stmt.where.Reset()
	stmt.cols = stmt.cols[:0]
	stmt.distinct = ""
	stmt.forupdate = false

	stmt.joins = stmt.joins[:]
	stmt.orders.Reset()
	stmt.group = ""

	stmt.havingQuery = ""
	stmt.havingVals = nil

	stmt.limitQuery = ""
	stmt.limitVals = nil
}

// SQL 获取 SQL 语句及对应的参数
func (stmt *SelectStmt) SQL() (string, []interface{}, error) {
	if stmt.table == "" {
		return "", nil, ErrTableIsEmpty
	}

	if len(stmt.cols) == 0 {
		return "", nil, ErrColumnsIsEmpty
	}

	buf := core.NewStringBuilder("SELECT ")
	args := make([]interface{}, 0, 10)

	if stmt.distinct != "" {
		buf.WriteString("DISTINCT ")
		buf.WriteString(stmt.distinct)
		buf.WriteByte(' ')
	}

	for _, c := range stmt.cols {
		buf.WriteString(c)
		buf.WriteByte(',')
	}
	buf.TruncateLast(1)

	buf.WriteString(" FROM ")
	buf.WriteString(stmt.table)

	// join
	if len(stmt.joins) > 0 {
		buf.WriteByte(' ')
		for _, join := range stmt.joins {
			buf.WriteString(join.typ)
			buf.WriteString(" JOIN ")
			buf.WriteString(join.table)
			buf.WriteString(" ON ")
			buf.WriteString(join.on)
			buf.WriteByte(',')
		}
		buf.TruncateLast(1)
	}

	// where
	wq, wa, err := stmt.where.SQL()
	if err != nil {
		return "", nil, err
	}
	buf.WriteString(wq)
	args = append(args, wa...)

	// group by
	if stmt.group != "" {
		buf.WriteString(stmt.group)
	}

	// having
	if stmt.havingQuery != "" {
		buf.WriteString(stmt.havingQuery)
		args = append(args, stmt.havingVals...)
	}

	// order by
	if stmt.orders != nil && stmt.orders.Len() > 0 {
		buf.WriteString(stmt.orders.String())
	}

	// limit
	if stmt.limitQuery != "" {
		buf.WriteString(stmt.limitQuery)
		args = append(args, stmt.limitVals...)
	}

	// for update
	if stmt.forupdate {
		buf.WriteString(" FOR UPDATE")
	}

	return buf.String(), args, nil
}

// Select 指定列名
func (stmt *SelectStmt) Select(cols ...string) *SelectStmt {
	if stmt.cols == nil {
		stmt.cols = cols
	} else {
		stmt.cols = append(stmt.cols, cols...)
	}
	return stmt
}

// From 指定表名
func (stmt *SelectStmt) From(table string) *SelectStmt {
	stmt.table = table

	return stmt
}

// Having 指定 having 语句
func (stmt *SelectStmt) Having(expr string, args ...interface{}) *SelectStmt {
	stmt.havingQuery = expr
	stmt.havingVals = args

	return stmt
}

// WhereStmt 实现 WhereStmter 接口
func (stmt *SelectStmt) WhereStmt() *WhereStmt {
	return stmt.where
}

// Where 指定 where 语句
func (stmt *SelectStmt) Where(cond string, args ...interface{}) *SelectStmt {
	return stmt.And(cond, args...)
}

// And 指定 where ... AND ... 语句
func (stmt *SelectStmt) And(cond string, args ...interface{}) *SelectStmt {
	stmt.where.And(cond, args...)
	return stmt
}

// Or 指定 where ... OR ... 语句
func (stmt *SelectStmt) Or(cond string, args ...interface{}) *SelectStmt {
	stmt.where.Or(cond, args...)
	return stmt
}

// Join 添加一条 Join 语句
func (stmt *SelectStmt) Join(typ, table, on string) *SelectStmt {
	if stmt.joins == nil {
		stmt.joins = make([]*join, 0, 5)
	}

	stmt.joins = append(stmt.joins, &join{typ: typ, table: table, on: on})
	return stmt
}

// Desc 倒序查询
func (stmt *SelectStmt) Desc(col ...string) *SelectStmt {
	return stmt.orderBy(false, col...)
}

// Asc 正序查询
func (stmt *SelectStmt) Asc(col ...string) *SelectStmt {
	return stmt.orderBy(true, col...)
}

func (stmt *SelectStmt) orderBy(asc bool, col ...string) *SelectStmt {
	if stmt.orders == nil {
		stmt.orders = core.NewStringBuilder(" ORDER BY ")
	} else {
		stmt.orders.WriteByte(',')
	}

	for _, c := range col {
		stmt.orders.WriteString(c)
		stmt.orders.WriteByte(',')
	}
	stmt.orders.TruncateLast(1)

	if asc {
		stmt.orders.WriteString(" ASC ")
	} else {
		stmt.orders.WriteString(" DESC ")
	}

	return stmt
}

// ForUpdate 添加 FOR UPDATE 语句部分
func (stmt *SelectStmt) ForUpdate() *SelectStmt {
	stmt.forupdate = true
	return stmt
}

// Group 添加 GROUP BY 语句
func (stmt *SelectStmt) Group(col string) *SelectStmt {
	stmt.group = " GROUP BY " + col + " "
	return stmt
}

// Limit 生成 SQL 的 Limit 语句
func (stmt *SelectStmt) Limit(limit int, offset ...int) *SelectStmt {
	query, vals := stmt.engine.Dialect().LimitSQL(limit, offset...)
	stmt.limitQuery = query
	stmt.limitVals = vals
	return stmt
}

// Query 查询
func (stmt *SelectStmt) Query() (*sql.Rows, error) {
	return query(stmt.engine, stmt)
}

// QueryContext 查询
func (stmt *SelectStmt) QueryContext(ctx context.Context) (*sql.Rows, error) {
	return queryContext(ctx, stmt.engine, stmt)
}

// QueryObj 将符合当前条件的所有记录依次写入 objs 中。
func (stmt *SelectStmt) QueryObj(objs interface{}) (int, error) {
	rows, err := stmt.Query()
	if err != nil {
		return 0, err
	}
	defer rows.Close()

	return fetch.Obj(objs, rows)
}

// Prepare 预编译
func (stmt *SelectStmt) Prepare() (*sql.Stmt, error) {
	return prepare(stmt.engine, stmt)
}

// PrepareContext 预编译
func (stmt *SelectStmt) PrepareContext(ctx context.Context) (*sql.Stmt, error) {
	return prepareContext(ctx, stmt.engine, stmt)
}
<commit_msg>添加 Count 函数<commit_after>// Copyright 2018 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package sqlbuilder

import (
	"context"
	"database/sql"

	"github.com/issue9/orm/core"
	"github.com/issue9/orm/fetch"
)

// SelectStmt 查询语句
type SelectStmt struct {
	engine    core.Engine
	table     string
	where     *WhereStmt
	cols      []string
	distinct  string
	forupdate bool

	joins  []*join
	orders *core.StringBuilder
	group  string

	havingQuery string
	havingVals  []interface{}

	limitQuery string
	limitVals  []interface{}
}

type join struct {
	typ   string
	on    string
	table string
}

// Select 声明一条 Select 语句
func Select(e core.Engine) *SelectStmt {
	return &SelectStmt{
		engine: e,
		where:  newWhereStmt(),
	}
}

// Distinct 声明一条 Select 语句的 Distinct
func (stmt *SelectStmt) Distinct(col string) *SelectStmt {
	stmt.distinct = col
	return stmt
}

// Reset 重置语句
func (stmt *SelectStmt) Reset() {
	stmt.table = ""
	stmt.where.Reset()
	stmt.cols = stmt.cols[:0]
	stmt.distinct = ""
	stmt.forupdate = false

	stmt.joins = stmt.joins[:]
	stmt.orders.Reset()
	stmt.group = ""

	stmt.havingQuery = ""
	stmt.havingVals = nil

	stmt.limitQuery = ""
	stmt.limitVals = nil
}

// SQL 获取 SQL 语句及对应的参数
func (stmt *SelectStmt) SQL() (string, []interface{}, error) {
	if stmt.table == "" {
		return "", nil, ErrTableIsEmpty
	}

	if len(stmt.cols) == 0 {
		return "", nil, ErrColumnsIsEmpty
	}

	buf := core.NewStringBuilder("SELECT ")
	args := make([]interface{}, 0, 10)

	if stmt.distinct != "" {
		buf.WriteString("DISTINCT ")
		buf.WriteString(stmt.distinct)
		buf.WriteByte(' ')
	}

	for _, c := range stmt.cols {
		buf.WriteString(c)
		buf.WriteByte(',')
	}
	buf.TruncateLast(1)

	buf.WriteString(" FROM ")
	buf.WriteString(stmt.table)

	// join
	if len(stmt.joins) > 0 {
		buf.WriteByte(' ')
		for _, join := range stmt.joins {
			buf.WriteString(join.typ)
			buf.WriteString(" JOIN ")
			buf.WriteString(join.table)
			buf.WriteString(" ON ")
			buf.WriteString(join.on)
			buf.WriteByte(',')
		}
		buf.TruncateLast(1)
	}

	// where
	wq, wa, err := stmt.where.SQL()
	if err != nil {
		return "", nil, err
	}
	buf.WriteString(wq)
	args = append(args, wa...)

	// group by
	if stmt.group != "" {
		buf.WriteString(stmt.group)
	}

	// having
	if stmt.havingQuery != "" {
		buf.WriteString(stmt.havingQuery)
		args = append(args, stmt.havingVals...)
	}

	// order by
	if stmt.orders != nil && stmt.orders.Len() > 0 {
		buf.WriteString(stmt.orders.String())
	}

	// limit
	if stmt.limitQuery != "" {
		buf.WriteString(stmt.limitQuery)
		args = append(args, stmt.limitVals...)
	}

	// for update
	if stmt.forupdate {
		buf.WriteString(" FOR UPDATE")
	}

	return buf.String(), args, nil
}

// Select 指定列名
func (stmt *SelectStmt) Select(cols ...string) *SelectStmt {
	if stmt.cols == nil {
		stmt.cols = cols
	} else {
		stmt.cols = append(stmt.cols, cols...)
	}
	return stmt
}

// From 指定表名
func (stmt *SelectStmt) From(table string) *SelectStmt {
	stmt.table = table

	return stmt
}

// Having 指定 having 语句
func (stmt *SelectStmt) Having(expr string, args ...interface{}) *SelectStmt {
	stmt.havingQuery = expr
	stmt.havingVals = args

	return stmt
}

// WhereStmt 实现 WhereStmter 接口
func (stmt *SelectStmt) WhereStmt() *WhereStmt {
	return stmt.where
}

// Where 指定 where 语句
func (stmt *SelectStmt) Where(cond string, args ...interface{}) *SelectStmt {
	return stmt.And(cond, args...)
}

// And 指定 where ... AND ... 语句
func (stmt *SelectStmt) And(cond string, args ...interface{}) *SelectStmt {
	stmt.where.And(cond, args...)
	return stmt
}

// Or 指定 where ... OR ... 语句
func (stmt *SelectStmt) Or(cond string, args ...interface{}) *SelectStmt {
	stmt.where.Or(cond, args...)
	return stmt
}

// Join 添加一条 Join 语句
func (stmt *SelectStmt) Join(typ, table, on string) *SelectStmt {
	if stmt.joins == nil {
		stmt.joins = make([]*join, 0, 5)
	}

	stmt.joins = append(stmt.joins, &join{typ: typ, table: table, on: on})
	return stmt
}

// Desc 倒序查询
func (stmt *SelectStmt) Desc(col ...string) *SelectStmt {
	return stmt.orderBy(false, col...)
}

// Asc 正序查询
func (stmt *SelectStmt) Asc(col ...string) *SelectStmt {
	return stmt.orderBy(true, col...)
}

func (stmt *SelectStmt) orderBy(asc bool, col ...string) *SelectStmt {
	if stmt.orders == nil {
		stmt.orders = core.NewStringBuilder(" ORDER BY ")
	} else {
		stmt.orders.WriteByte(',')
	}

	for _, c := range col {
		stmt.orders.WriteString(c)
		stmt.orders.WriteByte(',')
	}
	stmt.orders.TruncateLast(1)

	if asc {
		stmt.orders.WriteString(" ASC ")
	} else {
		stmt.orders.WriteString(" DESC ")
	}

	return stmt
}

// ForUpdate 添加 FOR UPDATE 语句部分
func (stmt *SelectStmt) ForUpdate() *SelectStmt {
	stmt.forupdate = true
	return stmt
}

// Group 添加 GROUP BY 语句
func (stmt *SelectStmt) Group(col string) *SelectStmt {
	stmt.group = " GROUP BY " + col + " "
	return stmt
}

// Limit 生成 SQL 的 Limit 语句
func (stmt *SelectStmt) Limit(limit int, offset ...int) *SelectStmt {
	query, vals := stmt.engine.Dialect().LimitSQL(limit, offset...)
	stmt.limitQuery = query
	stmt.limitVals = vals
	return stmt
}

// Count 用于统计符合当前条件的数量。
//
// NOTE: 会去掉 Limit 限制，其它不变，如果要分页，
// 在 Count 之后，需要再次调用 Limit 函数.
func (stmt *SelectStmt) Count() *SelectStmt {
	stmt.limitQuery = stmt.limitQuery[:0]
	stmt.limitQuery = ""
	return stmt
}

// Prepare 预编译
func (stmt *SelectStmt) Prepare() (*sql.Stmt, error) {
	return prepare(stmt.engine, stmt)
}

// PrepareContext 预编译
func (stmt *SelectStmt) PrepareContext(ctx context.Context) (*sql.Stmt, error) {
	return prepareContext(ctx, stmt.engine, stmt)
}

// Query 查询
func (stmt *SelectStmt) Query() (*sql.Rows, error) {
	return query(stmt.engine, stmt)
}

// QueryContext 查询
func (stmt *SelectStmt) QueryContext(ctx context.Context) (*sql.Rows, error) {
	return queryContext(ctx, stmt.engine, stmt)
}

// QueryObj 将符合当前条件的所有记录依次写入 objs 中。
func (stmt *SelectStmt) QueryObj(objs interface{}) (int, error) {
	rows, err := stmt.Query()
	if err != nil {
		return 0, err
	}
	defer rows.Close()

	return fetch.Obj(objs, rows)
}
<|endoftext|>