<commit_before>package msg

import (
	"bytes"
	"encoding/json"
	"fmt"
	"unicode"

	"github.com/bitfinexcom/bitfinex-api-go/pkg/convert"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/balanceinfo"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/book"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/candle"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/event"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/fundingcredit"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/fundingoffer"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/order"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/status"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/ticker"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/trades"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/wallet"
)

type Msg struct {
	Data     []byte
	Err      error
	CID      int
	IsPublic bool
}

func (m Msg) IsEvent() bool {
	t := bytes.TrimLeftFunc(m.Data, unicode.IsSpace)
	return bytes.HasPrefix(t, []byte("{"))
}

func (m Msg) IsRaw() bool {
	t := bytes.TrimLeftFunc(m.Data, unicode.IsSpace)
	return bytes.HasPrefix(t, []byte("["))
}

func (m Msg) ProcessRaw(chanInfo map[int64]event.Info) (interface{}, error) {
	var raw []interface{}
	if err := json.Unmarshal(m.Data, &raw); err != nil {
		return nil, fmt.Errorf("parsing msg: %s, err: %s", m.Data, err)
	}
	// payload data is always last element of the slice
	pld := raw[len(raw)-1]
	// chanID is always 1st element of the slice
	chID := convert.I64ValOrZero(raw[0])
	// allocate channel name by id to know how to transform raw data
	inf, ok := chanInfo[chID]
	if !ok {
		return nil, fmt.Errorf("unrecognized chanId:%d", chID)
	}

	switch data := pld.(type) {
	case string:
		return event.Info{
			ChanID:    chID,
			Subscribe: event.Subscribe{Event: data},
		}, nil
	case []interface{}:
		switch inf.Channel {
		case "trades":
			return trades.FromWSRaw(inf.Symbol, raw, data)
		case "ticker":
			return ticker.FromWSRaw(inf.Symbol, data)
		case "book":
			return book.FromWSRaw(inf.Symbol, inf.Precision, data)
		case "candles":
			return candle.FromWSRaw(inf.Key, data)
		case "status":
			return status.FromWSRaw(inf.Key, data)
		}
	}

	return raw, nil
}

func (m Msg) ProcessPrivateRaw() (interface{}, error) {
	var raw []interface{}
	if err := json.Unmarshal(m.Data, &raw); err != nil {
		return nil, fmt.Errorf("parsing auth msg: %s, err: %s", m.Data, err)
	}
	// payload data is always last element of the slice
	pld := raw[len(raw)-1]
	// op name is 2nd element
	op := convert.SValOrEmpty(raw[1])

	switch data := pld.(type) {
	case string:
		return event.Info{
			ChanID:    convert.I64ValOrZero(raw[0]),
			Subscribe: event.Subscribe{Event: data},
		}, nil
	case []interface{}:
		switch op {
		case "os":
			return order.SnapshotFromRaw(data)
		case "on":
			return order.NewFromRaw(data)
		case "ou":
			return order.UpdateFromRaw(data)
		case "oc":
			return order.CancelFromRaw(data)
		case "ws":
			return wallet.SnapshotFromRaw(data, wallet.FromWsRaw)
		case "wu":
			return wallet.UpdateFromRaw(data)
		case "bu":
			return balanceinfo.UpdateFromRaw(data)
		case "foc":
			return fundingoffer.CancelFromRaw(data)
		case "fon":
			return fundingoffer.NewFromRaw(data)
		case "fos":
			return fundingoffer.SnapshotFromRaw(data)
		case "fcs":
			return fundingcredit.SnapshotFromRaw(data)
		case "fcn":
			return fundingcredit.NewFromRaw(raw)
		case "fcu":
			return fundingcredit.UpdateFromRaw(raw)
		case "fcc":
			return fundingcredit.CancelFromRaw(raw)
		case "ftu":
			return trades.AFTUFromRaw(raw)
		case "fte":
			return trades.AFTEFromRaw(raw)
		}
	}

	return raw, nil
}

func (m Msg) ProcessEvent() (i event.Info, err error) {
	if err = json.Unmarshal(m.Data, &i); err != nil {
		return i, fmt.Errorf("parsing msg: %s, err: %s", m.Data, err)
	}
	return
}
<commit_msg>adding more msg mappings and placeholders for missing ones<commit_after>package msg

import (
	"bytes"
	"encoding/json"
	"fmt"
	"unicode"

	"github.com/bitfinexcom/bitfinex-api-go/pkg/convert"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/balanceinfo"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/book"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/candle"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/event"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/fundingcredit"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/fundingloan"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/fundingoffer"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/order"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/position"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/status"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/ticker"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/trades"
	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/wallet"
)

type Msg struct {
	Data     []byte
	Err      error
	CID      int
	IsPublic bool
}

func (m Msg) IsEvent() bool {
	t := bytes.TrimLeftFunc(m.Data, unicode.IsSpace)
	return bytes.HasPrefix(t, []byte("{"))
}

func (m Msg) IsRaw() bool {
	t := bytes.TrimLeftFunc(m.Data, unicode.IsSpace)
	return bytes.HasPrefix(t, []byte("["))
}

func (m Msg) ProcessRaw(chanInfo map[int64]event.Info) (interface{}, error) {
	var raw []interface{}
	if err := json.Unmarshal(m.Data, &raw); err != nil {
		return nil, fmt.Errorf("parsing msg: %s, err: %s", m.Data, err)
	}
	// payload data is always last element of the slice
	pld := raw[len(raw)-1]
	// chanID is always 1st element of the slice
	chID := convert.I64ValOrZero(raw[0])
	// allocate channel name by id to know how to transform raw data
	inf, ok := chanInfo[chID]
	if !ok {
		return nil, fmt.Errorf("unrecognized chanId:%d", chID)
	}

	switch data := pld.(type) {
	case string:
		return event.Info{
			ChanID:    chID,
			Subscribe: event.Subscribe{Event: data},
		}, nil
	case []interface{}:
		switch inf.Channel {
		case "trades":
			return trades.FromWSRaw(inf.Symbol, raw, data)
		case "ticker":
			return ticker.FromWSRaw(inf.Symbol, data)
		case "book":
			return book.FromWSRaw(inf.Symbol, inf.Precision, data)
		case "candles":
			return candle.FromWSRaw(inf.Key, data)
		case "status":
			return status.FromWSRaw(inf.Key, data)
		}
	}

	return raw, nil
}

func (m Msg) ProcessPrivateRaw() (interface{}, error) {
	var raw []interface{}
	if err := json.Unmarshal(m.Data, &raw); err != nil {
		return nil, fmt.Errorf("parsing auth msg: %s, err: %s", m.Data, err)
	}
	// payload data is always last element of the slice
	pld := raw[len(raw)-1]
	// op name is 2nd element
	op := convert.SValOrEmpty(raw[1])

	switch data := pld.(type) {
	case string:
		return event.Info{
			ChanID:    convert.I64ValOrZero(raw[0]),
			Subscribe: event.Subscribe{Event: data},
		}, nil
	case []interface{}:
		switch op {
		case "bu":
			return balanceinfo.UpdateFromRaw(data)
		case "ps":
			return position.SnapshotFromRaw(data)
		case "pn":
			// TODO
		case "pu":
			// TODO
		case "pc":
			// TODO
		case "ws":
			return wallet.SnapshotFromRaw(data, wallet.FromWsRaw)
		case "wu":
			return wallet.UpdateFromRaw(data)
		case "os":
			return order.SnapshotFromRaw(data)
		case "on":
			return order.NewFromRaw(data)
		case "on-req":
			// TODO
		case "ou":
			return order.UpdateFromRaw(data)
		case "oc":
			return order.CancelFromRaw(data)
		case "oc-req":
			// TODO
		case "oc_multi-req":
			// TODO
		case "te":
			// TODO
		case "tu":
			// TODO
		case "fte":
			return trades.AFTEFromRaw(data)
		case "ftu":
			return trades.AFTUFromRaw(data)
		case "mis":
			// TODO
		case "miu":
			// TODO
		case "n":
			// TODO
		case "fos":
			return fundingoffer.SnapshotFromRaw(data)
		case "fon":
			return fundingoffer.NewFromRaw(data)
		case "fou":
			return fundingoffer.UpdateFromRaw(data)
		case "foc":
			return fundingoffer.CancelFromRaw(data)
		case "fcs":
			return fundingcredit.SnapshotFromRaw(data)
		case "fcn":
			return fundingcredit.NewFromRaw(data)
		case "fcu":
			return fundingcredit.UpdateFromRaw(data)
		case "fcc":
			return fundingcredit.CancelFromRaw(data)
		case "fls":
			return fundingloan.SnapshotFromRaw(data)
		case "fln":
			return fundingloan.NewFromRaw(data)
		case "flu":
			return fundingloan.UpdateFromRaw(data)
		case "flc":
			return fundingloan.CancelFromRaw(data)
		case "hfts":
			// TODO
		case "uac":
			// TODO
		}
	}

	return raw, nil
}

func (m Msg) ProcessEvent() (i event.Info, err error) {
	if err = json.Unmarshal(m.Data, &i); err != nil {
		return i, fmt.Errorf("parsing msg: %s, err: %s", m.Data, err)
	}
	return
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Vitess Authors.

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

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

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

package engine

import (
	"sync"

	"vitess.io/vitess/go/mysql"

	"vitess.io/vitess/go/sqltypes"
	querypb "vitess.io/vitess/go/vt/proto/query"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/vterrors"
)

// Concatenate Primitive is used to concatenate results from multiple sources.
var _ Primitive = (*Concatenate)(nil)

//Concatenate specified the parameter for concatenate primitive
type Concatenate struct {
	LHS, RHS Primitive
}

//RouteType returns a description of the query routing type used by the primitive
func (c *Concatenate) RouteType() string {
	return "Concatenate"
}

// GetKeyspaceName specifies the Keyspace that this primitive routes to
func (c *Concatenate) GetKeyspaceName() string {
	return formatTwoOptionsNicely(c.LHS.GetKeyspaceName(), c.RHS.GetKeyspaceName())
}

// GetTableName specifies the table that this primitive routes to.
func (c *Concatenate) GetTableName() string {
	return formatTwoOptionsNicely(c.LHS.GetTableName(), c.RHS.GetTableName())
}

func formatTwoOptionsNicely(a, b string) string {
	if a == b {
		return a
	}
	return a + "_" + b
}

// Execute performs a non-streaming exec.
func (c *Concatenate) Execute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
	lhs, rhs, err := c.execSources(vcursor, bindVars, wantfields)
	if err != nil {
		return nil, err
	}

	fields, err := c.getFields(lhs.Fields, rhs.Fields)
	if err != nil {
		return nil, err
	}

	if len(lhs.Rows) > 0 &&
		len(rhs.Rows) > 0 &&
		len(lhs.Rows[0]) != len(rhs.Rows[0]) {
		return nil, mysql.NewSQLError(mysql.ERWrongNumberOfColumnsInSelect, "21000", "The used SELECT statements have a different number of columns")
	}

	return &sqltypes.Result{
		Fields:       fields,
		RowsAffected: lhs.RowsAffected + rhs.RowsAffected,
		Rows:         append(lhs.Rows, rhs.Rows...),
	}, nil
}

func (c *Concatenate) getFields(a, b []*querypb.Field) ([]*querypb.Field, error) {
	switch {
	case a != nil && b != nil:
		err := compareFields(a, b)
		if err != nil {
			return nil, err
		}
		return a, nil
	case a != nil:
		return a, nil
	case b != nil:
		return b, nil
	}

	return nil, nil
}
func (c *Concatenate) execSources(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, *sqltypes.Result, error) {
	var lhs, rhs *sqltypes.Result
	g := vcursor.ErrorGroupCancellableContext()
	g.Go(func() error {
		result, err := c.LHS.Execute(vcursor, bindVars, wantfields)
		if err != nil {
			return err
		}
		lhs = result
		return nil
	})
	g.Go(func() error {
		result, err := c.RHS.Execute(vcursor, bindVars, wantfields)
		if err != nil {
			return err
		}
		rhs = result
		return nil
	})
	if err := g.Wait(); err != nil {
		return nil, nil, vterrors.Wrap(err, "Concatenate.Execute")
	}
	return lhs, rhs, nil
}

// StreamExecute performs a streaming exec.
func (c *Concatenate) StreamExecute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
	var seenFields []*querypb.Field
	var fieldset sync.WaitGroup

	g := vcursor.ErrorGroupCancellableContext()
	fieldset.Add(1)
	var mu sync.Mutex

	fieldsSent := false
	g.Go(func() error {
		err := c.LHS.StreamExecute(vcursor, bindVars, wantfields, func(resultChunk *sqltypes.Result) error {
			if !fieldsSent {
				defer fieldset.Done()
				seenFields = resultChunk.Fields
				fieldsSent = true
				// No other call can happen before this call.
				return callback(resultChunk)
			}
			// This to ensure only one send happens back to the client.
			mu.Lock()
			defer mu.Unlock()
			select {
			case <-vcursor.Context().Done():
				return nil
			default:
				return callback(resultChunk)
			}
		})
		// This is to ensure other streams complete if the first stream failed to unlock the wait.
		if !fieldsSent {
			fieldset.Done()
		}
		return err
	})
	g.Go(func() error {
		err := c.RHS.StreamExecute(vcursor, bindVars, wantfields, func(resultChunk *sqltypes.Result) error {
			fieldset.Wait()

			if resultChunk.Fields != nil {
				err := compareFields(seenFields, resultChunk.Fields)
				if err != nil {
					return err
				}
			}

			// This to ensure only one send happens back to the client.
			mu.Lock()
			defer mu.Unlock()
			select {
			case <-vcursor.Context().Done():
				return nil
			default:
				return callback(resultChunk)
			}
		})
		return err
	})
	if err := g.Wait(); err != nil {
		return err
	}
	return nil
}

// GetFields fetches the field info.
func (c *Concatenate) GetFields(vcursor VCursor, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
	lhs, err := c.LHS.GetFields(vcursor, bindVars)
	if err != nil {
		return nil, err
	}
	rhs, err := c.RHS.GetFields(vcursor, bindVars)
	if err != nil {
		return nil, err
	}
	err = compareFields(lhs.Fields, rhs.Fields)
	if err != nil {
		return nil, err
	}

	return lhs, nil
}

//NeedsTransaction returns whether a transaction is needed for this primitive
func (c *Concatenate) NeedsTransaction() bool {
	return c.LHS.NeedsTransaction() || c.RHS.NeedsTransaction()
}

// Inputs returns the input primitives for this
func (c *Concatenate) Inputs() []Primitive {
	return []Primitive{c.LHS, c.RHS}
}

func (c *Concatenate) description() PrimitiveDescription {
	return PrimitiveDescription{OperatorType: c.RouteType()}
}

func compareFields(fields1 []*querypb.Field, fields2 []*querypb.Field) error {
	if len(fields1) != len(fields2) {
		return mysql.NewSQLError(mysql.ERWrongNumberOfColumnsInSelect, "21000", "The used SELECT statements have a different number of columns")
	}
	for i, field2 := range fields2 {
		field1 := fields1[i]
		if field1.Type != field2.Type {
			return vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "column field type does not match for name: (%v, %v) types: (%v, %v)", field1.Name, field2.Name, field1.Type, field2.Type)
		}
	}
	return nil
}
<commit_msg>simplify StreamExecute<commit_after>/*
Copyright 2020 The Vitess Authors.

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

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

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

package engine

import (
	"sync"

	"vitess.io/vitess/go/mysql"

	"vitess.io/vitess/go/sqltypes"
	querypb "vitess.io/vitess/go/vt/proto/query"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/vterrors"
)

// Concatenate Primitive is used to concatenate results from multiple sources.
var _ Primitive = (*Concatenate)(nil)

//Concatenate specified the parameter for concatenate primitive
type Concatenate struct {
	LHS, RHS Primitive
}

//RouteType returns a description of the query routing type used by the primitive
func (c *Concatenate) RouteType() string {
	return "Concatenate"
}

// GetKeyspaceName specifies the Keyspace that this primitive routes to
func (c *Concatenate) GetKeyspaceName() string {
	return formatTwoOptionsNicely(c.LHS.GetKeyspaceName(), c.RHS.GetKeyspaceName())
}

// GetTableName specifies the table that this primitive routes to.
func (c *Concatenate) GetTableName() string {
	return formatTwoOptionsNicely(c.LHS.GetTableName(), c.RHS.GetTableName())
}

func formatTwoOptionsNicely(a, b string) string {
	if a == b {
		return a
	}
	return a + "_" + b
}

// Execute performs a non-streaming exec.
func (c *Concatenate) Execute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
	lhs, rhs, err := c.execSources(vcursor, bindVars, wantfields)
	if err != nil {
		return nil, err
	}

	fields, err := c.getFields(lhs.Fields, rhs.Fields)
	if err != nil {
		return nil, err
	}

	if len(lhs.Rows) > 0 &&
		len(rhs.Rows) > 0 &&
		len(lhs.Rows[0]) != len(rhs.Rows[0]) {
		return nil, mysql.NewSQLError(mysql.ERWrongNumberOfColumnsInSelect, "21000", "The used SELECT statements have a different number of columns")
	}

	return &sqltypes.Result{
		Fields:       fields,
		RowsAffected: lhs.RowsAffected + rhs.RowsAffected,
		Rows:         append(lhs.Rows, rhs.Rows...),
	}, nil
}

func (c *Concatenate) getFields(a, b []*querypb.Field) ([]*querypb.Field, error) {
	switch {
	case a != nil && b != nil:
		err := compareFields(a, b)
		if err != nil {
			return nil, err
		}
		return a, nil
	case a != nil:
		return a, nil
	case b != nil:
		return b, nil
	}

	return nil, nil
}
func (c *Concatenate) execSources(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, *sqltypes.Result, error) {
	var lhs, rhs *sqltypes.Result
	g := vcursor.ErrorGroupCancellableContext()
	g.Go(func() error {
		result, err := c.LHS.Execute(vcursor, bindVars, wantfields)
		if err != nil {
			return err
		}
		lhs = result
		return nil
	})
	g.Go(func() error {
		result, err := c.RHS.Execute(vcursor, bindVars, wantfields)
		if err != nil {
			return err
		}
		rhs = result
		return nil
	})
	if err := g.Wait(); err != nil {
		return nil, nil, vterrors.Wrap(err, "Concatenate.Execute")
	}
	return lhs, rhs, nil
}

// StreamExecute performs a streaming exec.
func (c *Concatenate) StreamExecute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
	var seenFields []*querypb.Field
	var fieldset sync.WaitGroup

	g := vcursor.ErrorGroupCancellableContext()
	fieldset.Add(1)
	var cbMu, visitFieldsMu sync.Mutex

	visitFields := func(fields []*querypb.Field) (unlocked bool, err error) {
		visitFieldsMu.Lock()
		defer visitFieldsMu.Unlock()
		if seenFields == nil {
			seenFields = fields
			return true, nil
		}
		return false, compareFields(fields, seenFields)
	}

	visitor := func(resultChunk *sqltypes.Result) error {
		if resultChunk.Fields != nil {
			var err error
			unlocked, err := visitFields(resultChunk.Fields)
			if err != nil {
				return err
			}
			if unlocked {
				defer fieldset.Done()
				// nothing else will be sent to the client until we send out this first one
				return callback(resultChunk)
			}
		}
		fieldset.Wait()

		// This to ensure only one send happens back to the client.
		cbMu.Lock()
		defer cbMu.Unlock()

		select {
		case <-vcursor.Context().Done():
			return nil
		default:
			return callback(resultChunk)
		}
	}
	g.Go(func() error {
		return c.LHS.StreamExecute(vcursor, bindVars, wantfields, visitor)
	})
	g.Go(func() error {
		return c.RHS.StreamExecute(vcursor, bindVars, wantfields, visitor)
	})
	if err := g.Wait(); err != nil {
		return err
	}
	return nil
}

// GetFields fetches the field info.
func (c *Concatenate) GetFields(vcursor VCursor, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
	lhs, err := c.LHS.GetFields(vcursor, bindVars)
	if err != nil {
		return nil, err
	}
	rhs, err := c.RHS.GetFields(vcursor, bindVars)
	if err != nil {
		return nil, err
	}
	err = compareFields(lhs.Fields, rhs.Fields)
	if err != nil {
		return nil, err
	}

	return lhs, nil
}

//NeedsTransaction returns whether a transaction is needed for this primitive
func (c *Concatenate) NeedsTransaction() bool {
	return c.LHS.NeedsTransaction() || c.RHS.NeedsTransaction()
}

// Inputs returns the input primitives for this
func (c *Concatenate) Inputs() []Primitive {
	return []Primitive{c.LHS, c.RHS}
}

func (c *Concatenate) description() PrimitiveDescription {
	return PrimitiveDescription{OperatorType: c.RouteType()}
}

func compareFields(fields1 []*querypb.Field, fields2 []*querypb.Field) error {
	if len(fields1) != len(fields2) {
		return mysql.NewSQLError(mysql.ERWrongNumberOfColumnsInSelect, "21000", "The used SELECT statements have a different number of columns")
	}
	for i, field2 := range fields2 {
		field1 := fields1[i]
		if field1.Type != field2.Type {
			return vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "column field type does not match for name: (%v, %v) types: (%v, %v)", field1.Name, field2.Name, field1.Type, field2.Type)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package oauth

import (
	"github.com/hellofresh/janus/pkg/cors"
	"github.com/hellofresh/janus/pkg/proxy"
)

// AccessRequestType is the type for OAuth param `grant_type`
type AccessRequestType string

// AuthorizeRequestType is the type for OAuth param `response_type`
type AuthorizeRequestType string

// Spec Holds an api definition and basic options
type Spec struct {
	*OAuth
	Manager Manager
}

// OAuth holds the configuration for oauth proxies
type OAuth struct {
	Name                   string                 `bson:"name" json:"name" valid:"required"`
	Endpoints              Endpoints              `bson:"oauth_endpoints" json:"oauth_endpoints"`
	ClientEndpoints        ClientEndpoints        `bson:"oauth_client_endpoints" json:"oauth_client_endpoints"`
	AllowedAccessTypes     []AccessRequestType    `bson:"allowed_access_types" json:"allowed_access_types"`
	AllowedAuthorizeTypes  []AuthorizeRequestType `bson:"allowed_authorize_types" json:"allowed_authorize_types"`
	AuthorizeLoginRedirect string                 `bson:"auth_login_redirect" json:"auth_login_redirect"`
	Secrets                map[string]string      `bson:"secrets" json:"secrets"`
	CorsMeta               cors.Meta              `bson:"cors_meta" json:"cors_meta" valid:"cors_meta"`
	TokenStrategy          TokenStrategy          `bson:"token_strategy" json:"token_strategy"`
}

// Endpoints defines the oauth endpoints that wil be proxied
type Endpoints struct {
	Authorize *proxy.Definition `bson:"authorize" json:"authorize"`
	Token     *proxy.Definition `bson:"token" json:"token"`
	Info      *proxy.Definition `bson:"info" json:"info"`
	Revoke    *proxy.Definition `bson:"revoke" json:"revoke"`
}

// ClientEndpoints defines the oauth client endpoints that wil be proxied
type ClientEndpoints struct {
	Create *proxy.Definition `bson:"create" json:"create"`
	Remove *proxy.Definition `bson:"remove" json:"remove"`
}

// TokenStrategy defines the token strategy fields
type TokenStrategy struct {
	Name     string                `bson:"name" json:"name"`
	Settings TokenStrategySettings `bson:"settings" json:"settings"`
}

// TokenStrategySettings represents the settings for the token strategy
type TokenStrategySettings map[string]string
<commit_msg>replicating the CORS metadata<commit_after>package oauth

import (
	"github.com/hellofresh/janus/pkg/proxy"
)

// AccessRequestType is the type for OAuth param `grant_type`
type AccessRequestType string

// AuthorizeRequestType is the type for OAuth param `response_type`
type AuthorizeRequestType string

// Spec Holds an api definition and basic options
type Spec struct {
	*OAuth
	Manager Manager
}

// OAuth holds the configuration for oauth proxies
type OAuth struct {
	Name                   string                 `bson:"name" json:"name" valid:"required"`
	Endpoints              Endpoints              `bson:"oauth_endpoints" json:"oauth_endpoints"`
	ClientEndpoints        ClientEndpoints        `bson:"oauth_client_endpoints" json:"oauth_client_endpoints"`
	AllowedAccessTypes     []AccessRequestType    `bson:"allowed_access_types" json:"allowed_access_types"`
	AllowedAuthorizeTypes  []AuthorizeRequestType `bson:"allowed_authorize_types" json:"allowed_authorize_types"`
	AuthorizeLoginRedirect string                 `bson:"auth_login_redirect" json:"auth_login_redirect"`
	Secrets                map[string]string      `bson:"secrets" json:"secrets"`
	CorsMeta               meta                   `bson:"cors_meta" json:"cors_meta" valid:"cors_meta"`
	TokenStrategy          TokenStrategy          `bson:"token_strategy" json:"token_strategy"`
}

// Endpoints defines the oauth endpoints that wil be proxied
type Endpoints struct {
	Authorize *proxy.Definition `bson:"authorize" json:"authorize"`
	Token     *proxy.Definition `bson:"token" json:"token"`
	Info      *proxy.Definition `bson:"info" json:"info"`
	Revoke    *proxy.Definition `bson:"revoke" json:"revoke"`
}

// ClientEndpoints defines the oauth client endpoints that wil be proxied
type ClientEndpoints struct {
	Create *proxy.Definition `bson:"create" json:"create"`
	Remove *proxy.Definition `bson:"remove" json:"remove"`
}

// TokenStrategy defines the token strategy fields
type TokenStrategy struct {
	Name     string                `bson:"name" json:"name"`
	Settings TokenStrategySettings `bson:"settings" json:"settings"`
}

// TokenStrategySettings represents the settings for the token strategy
type TokenStrategySettings map[string]string

type meta struct {
	Domains        []string `mapstructure:"domains" bson:"domains" json:"domains"`
	Methods        []string `mapstructure:"methods" bson:"methods" json:"methods"`
	RequestHeaders []string `mapstructure:"request_headers" bson:"request_headers" json:"request_headers"`
	ExposedHeaders []string `mapstructure:"exposed_headers" bson:"exposed_headers" json:"exposed_headers"`
	Enabled        bool     `bson:"enabled" json:"enabled"`
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Vitess Authors.

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

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

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

package evalengine

import (
	"fmt"

	"vitess.io/vitess/go/mysql/collations"
	querypb "vitess.io/vitess/go/vt/proto/query"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vterrors"
)

type (
	ConverterLookup interface {
		ColumnLookup(col *sqlparser.ColName) (int, error)
		CollationForExpr(expr sqlparser.Expr) collations.ID
		DefaultCollation() collations.ID
	}
)

var ErrConvertExprNotSupported = "expr cannot be converted, not supported"

func convertComparisonExpr(op sqlparser.ComparisonExprOperator, left, right sqlparser.Expr, lookup ConverterLookup) (Expr, error) {
	l, err := convertExpr(left, lookup)
	if err != nil {
		return nil, err
	}
	r, err := convertExpr(right, lookup)
	if err != nil {
		return nil, err
	}
	return convertComparisonExpr2(op, l, r)
}

func convertComparisonExpr2(op sqlparser.ComparisonExprOperator, left, right Expr) (Expr, error) {
	binaryExpr := BinaryExpr{
		Left:  left,
		Right: right,
	}

	if op == sqlparser.InOp || op == sqlparser.NotInOp {
		return &InExpr{
			BinaryExpr: binaryExpr,
			Negate:     op == sqlparser.NotInOp,
			Hashed:     nil,
		}, nil
	}

	switch op {
	case sqlparser.EqualOp:
		return &ComparisonExpr{binaryExpr, compareEQ{}}, nil
	case sqlparser.NotEqualOp:
		return &ComparisonExpr{binaryExpr, compareNE{}}, nil
	case sqlparser.LessThanOp:
		return &ComparisonExpr{binaryExpr, compareLT{}}, nil
	case sqlparser.LessEqualOp:
		return &ComparisonExpr{binaryExpr, compareLE{}}, nil
	case sqlparser.GreaterThanOp:
		return &ComparisonExpr{binaryExpr, compareGT{}}, nil
	case sqlparser.GreaterEqualOp:
		return &ComparisonExpr{binaryExpr, compareGE{}}, nil
	case sqlparser.NullSafeEqualOp:
		return &ComparisonExpr{binaryExpr, compareNullSafeEQ{}}, nil
	case sqlparser.LikeOp:
		return &LikeExpr{BinaryExpr: binaryExpr}, nil
	case sqlparser.NotLikeOp:
		return &LikeExpr{BinaryExpr: binaryExpr, Negate: true}, nil
	default:
		return nil, vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, op.ToString())
	}
}

func convertLogicalExpr(opname string, left, right sqlparser.Expr, lookup ConverterLookup) (Expr, error) {
	l, err := convertExpr(left, lookup)
	if err != nil {
		return nil, err
	}
	r, err := convertExpr(right, lookup)
	if err != nil {
		return nil, err
	}

	var logic func(l, r boolean) boolean
	switch opname {
	case "AND":
		logic = func(l, r boolean) boolean { return l.and(r) }
	case "OR":
		logic = func(l, r boolean) boolean { return l.or(r) }
	case "XOR":
		logic = func(l, r boolean) boolean { return l.xor(r) }
	default:
		panic("unexpected logical operator")
	}

	return &LogicalExpr{
		BinaryExpr: BinaryExpr{
			Left:  l,
			Right: r,
		},
		op:     logic,
		opname: opname,
	}, nil
}

func convertIsExpr(left sqlparser.Expr, op sqlparser.IsExprOperator, lookup ConverterLookup) (Expr, error) {
	expr, err := convertExpr(left, lookup)
	if err != nil {
		return nil, err
	}

	var check func(result *EvalResult) bool

	switch op {
	case sqlparser.IsNullOp:
		check = func(er *EvalResult) bool { return er.null() }
	case sqlparser.IsNotNullOp:
		check = func(er *EvalResult) bool { return !er.null() }
	case sqlparser.IsTrueOp:
		check = func(er *EvalResult) bool { return er.truthy() == boolTrue }
	case sqlparser.IsNotTrueOp:
		check = func(er *EvalResult) bool { return er.truthy() != boolTrue }
	case sqlparser.IsFalseOp:
		check = func(er *EvalResult) bool { return er.truthy() == boolFalse }
	case sqlparser.IsNotFalseOp:
		check = func(er *EvalResult) bool { return er.truthy() != boolFalse }
	}

	return &IsExpr{
		UnaryExpr: UnaryExpr{expr},
		Op:        op,
		Check:     check,
	}, nil
}

func getCollation(expr sqlparser.Expr, lookup ConverterLookup) collations.TypedCollation {
	collation := collations.TypedCollation{
		Collation:    lookup.CollationForExpr(expr),
		Coercibility: collations.CoerceCoercible,
		Repertoire:   collations.RepertoireUnicode,
	}
	if collation.Collation == collations.Unknown {
		collation.Collation = lookup.DefaultCollation()
	}
	return collation
}

func ConvertEx(e sqlparser.Expr, lookup ConverterLookup, simplify bool) (Expr, error) {
	expr, err := convertExpr(e, lookup)
	if err != nil {
		return nil, err
	}
	if simplify {
		var staticEnv ExpressionEnv
		if lookup != nil {
			staticEnv.DefaultCollation = lookup.DefaultCollation()
		} else {
			staticEnv.DefaultCollation = collations.Default()
		}
		expr, err = simplifyExpr(&staticEnv, expr)
	}
	return expr, err
}

// Convert converts between AST expressions and executable expressions
func Convert(e sqlparser.Expr, lookup ConverterLookup) (Expr, error) {
	return ConvertEx(e, lookup, true)
}

func convertExpr(e sqlparser.Expr, lookup ConverterLookup) (Expr, error) {
	switch node := e.(type) {
	case *sqlparser.ColName:
		if lookup == nil {
			return nil, vterrors.Wrap(convertNotSupported(e), "cannot lookup column")
		}
		idx, err := lookup.ColumnLookup(node)
		if err != nil {
			return nil, err
		}
		collation := getCollation(node, lookup)
		return NewColumn(idx, collation), nil
	case *sqlparser.ComparisonExpr:
		expr, err := convertComparisonExpr(node.Operator, node.Left, node.Right, lookup)
		if err != nil {
			return nil, convertNotSupported(e)
		}
		return expr, err
	case sqlparser.Argument:
		collation := getCollation(e, lookup)
		return NewBindVar(string(node), collation), nil
	case sqlparser.ListArg:
		collation := getCollation(e, lookup)
		return NewBindVar(string(node), collation), nil
	case *sqlparser.Literal:
		switch node.Type {
		case sqlparser.IntVal:
			return NewLiteralIntegralFromBytes(node.Bytes())
		case sqlparser.FloatVal:
			return NewLiteralFloatFromBytes(node.Bytes())
		case sqlparser.DecimalVal:
			return NewLiteralDecimalFromBytes(node.Bytes())
		case sqlparser.StrVal:
			collation := getCollation(e, lookup)
			return NewLiteralString(node.Bytes(), collation), nil
		case sqlparser.HexVal:
			return NewLiteralBinaryFromHex(node.Bytes())
		}
	case sqlparser.BoolVal:
		if node {
			return NewLiteralInt(1), nil
		}
		return NewLiteralInt(0), nil
	case *sqlparser.AndExpr:
		return convertLogicalExpr("AND", node.Left, node.Right, lookup)
	case *sqlparser.OrExpr:
		return convertLogicalExpr("OR", node.Left, node.Right, lookup)
	case *sqlparser.XorExpr:
		return convertLogicalExpr("XOR", node.Left, node.Right, lookup)
	case *sqlparser.NotExpr:
		inner, err := convertExpr(node.Expr, lookup)
		if err != nil {
			return nil, err
		}
		return &NotExpr{UnaryExpr{inner}}, nil
	case *sqlparser.BinaryExpr:
		var op ArithmeticOp
		switch node.Operator {
		case sqlparser.PlusOp:
			op = &OpAddition{}
		case sqlparser.MinusOp:
			op = &OpSubstraction{}
		case sqlparser.MultOp:
			op = &OpMultiplication{}
		case sqlparser.DivOp:
			op = &OpDivision{}
		default:
			return nil, convertNotSupported(e)
		}
		left, err := convertExpr(node.Left, lookup)
		if err != nil {
			return nil, err
		}
		right, err := convertExpr(node.Right, lookup)
		if err != nil {
			return nil, err
		}
		return &ArithmeticExpr{
			BinaryExpr: BinaryExpr{
				Left:  left,
				Right: right,
			},
			Op: op,
		}, nil
	case sqlparser.ValTuple:
		var exprs TupleExpr
		for _, expr := range node {
			convertedExpr, err := convertExpr(expr, lookup)
			if err != nil {
				return nil, err
			}
			exprs = append(exprs, convertedExpr)
		}
		return exprs, nil
	case *sqlparser.NullVal:
		return NullExpr, nil
	case *sqlparser.CollateExpr:
		expr, err := convertExpr(node.Expr, lookup)
		if err != nil {
			return nil, err
		}
		coll := collations.Local().LookupByName(node.Collation)
		if coll == nil {
			return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "Unknown collation: '%s'", node.Collation)
		}
		return &CollateExpr{
			UnaryExpr: UnaryExpr{expr},
			TypedCollation: collations.TypedCollation{
				Collation:    coll.ID(),
				Coercibility: collations.CoerceExplicit,
				Repertoire:   collations.RepertoireUnicode,
			},
		}, nil
	case *sqlparser.IntroducerExpr:
		expr, err := convertExpr(node.Expr, lookup)
		if err != nil {
			return nil, err
		}
		var collation collations.ID

		// TODO: case insensitive
		if node.CharacterSet == "_binary" {
			collation = collations.CollationBinaryID
		} else {
			defaultCollation := collations.Local().DefaultCollationForCharset(node.CharacterSet[1:])
			if defaultCollation == nil {
				panic(fmt.Sprintf("unknown character set: %s", node.CharacterSet))
			}
			collation = defaultCollation.ID()
		}

		switch lit := expr.(type) {
		case *Literal:
			switch collation {
			case collations.CollationBinaryID:
				lit.Val.type_ = querypb.Type_VARBINARY
				lit.Val.collation_ = collationBinary
			default:
				lit.Val.type_ = querypb.Type_VARCHAR
				lit.Val.replaceCollationID(collation)
			}
		case *BindVariable:
			switch collation {
			case collations.CollationBinaryID:
				lit.coerceType = querypb.Type_VARBINARY
				lit.coll = collationBinary
			default:
				lit.coerceType = querypb.Type_VARCHAR
				lit.coll.Collation = collation
			}
		default:
			panic("character set introducers are only supported for literals and arguments")
		}
		return expr, nil
	case *sqlparser.IsExpr:
		return convertIsExpr(node.Left, node.Right, lookup)
	case *sqlparser.FuncExpr:
		method := node.Name.Lowered()
		call, ok := builtinFunctions[method]
		if !ok {
			break
		}

		var args TupleExpr
		var aliases []sqlparser.ColIdent
		for _, expr := range node.Exprs {
			aliased, ok := expr.(*sqlparser.AliasedExpr)
			if !ok {
				return nil, convertNotSupported(e)
			}
			convertedExpr, err := convertExpr(aliased.Expr, lookup)
			if err != nil {
				return nil, err
			}
			args = append(args, convertedExpr)
			aliases = append(aliases, aliased.As)
		}
		return &CallExpression{
			Arguments: args,
			Aliases:   aliases,
			Method:    method,
			Call:      call,
		}, nil
	}
	return nil, convertNotSupported(e)
}

var builtinFunctions = map[string]func(*ExpressionEnv, []EvalResult, *EvalResult){
	"coalesce":  builtinFuncCoalesce,
	"greatest":  builtinFuncGreatest,
	"least":     builtinFuncLeast,
	"collation": builtinFuncCollation,
}

func convertNotSupported(e sqlparser.Expr) error {
	return vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "%s: %s", ErrConvertExprNotSupported, sqlparser.String(e))
}
<commit_msg>evalengine: support nil lookup functions<commit_after>/*
Copyright 2021 The Vitess Authors.

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

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

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

package evalengine

import (
	"fmt"

	"vitess.io/vitess/go/mysql/collations"
	querypb "vitess.io/vitess/go/vt/proto/query"
	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vterrors"
)

type (
	ConverterLookup interface {
		ColumnLookup(col *sqlparser.ColName) (int, error)
		CollationForExpr(expr sqlparser.Expr) collations.ID
		DefaultCollation() collations.ID
	}
)

var ErrConvertExprNotSupported = "expr cannot be converted, not supported"

func convertComparisonExpr(op sqlparser.ComparisonExprOperator, left, right sqlparser.Expr, lookup ConverterLookup) (Expr, error) {
	l, err := convertExpr(left, lookup)
	if err != nil {
		return nil, err
	}
	r, err := convertExpr(right, lookup)
	if err != nil {
		return nil, err
	}
	return convertComparisonExpr2(op, l, r)
}

func convertComparisonExpr2(op sqlparser.ComparisonExprOperator, left, right Expr) (Expr, error) {
	binaryExpr := BinaryExpr{
		Left:  left,
		Right: right,
	}

	if op == sqlparser.InOp || op == sqlparser.NotInOp {
		return &InExpr{
			BinaryExpr: binaryExpr,
			Negate:     op == sqlparser.NotInOp,
			Hashed:     nil,
		}, nil
	}

	switch op {
	case sqlparser.EqualOp:
		return &ComparisonExpr{binaryExpr, compareEQ{}}, nil
	case sqlparser.NotEqualOp:
		return &ComparisonExpr{binaryExpr, compareNE{}}, nil
	case sqlparser.LessThanOp:
		return &ComparisonExpr{binaryExpr, compareLT{}}, nil
	case sqlparser.LessEqualOp:
		return &ComparisonExpr{binaryExpr, compareLE{}}, nil
	case sqlparser.GreaterThanOp:
		return &ComparisonExpr{binaryExpr, compareGT{}}, nil
	case sqlparser.GreaterEqualOp:
		return &ComparisonExpr{binaryExpr, compareGE{}}, nil
	case sqlparser.NullSafeEqualOp:
		return &ComparisonExpr{binaryExpr, compareNullSafeEQ{}}, nil
	case sqlparser.LikeOp:
		return &LikeExpr{BinaryExpr: binaryExpr}, nil
	case sqlparser.NotLikeOp:
		return &LikeExpr{BinaryExpr: binaryExpr, Negate: true}, nil
	default:
		return nil, vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, op.ToString())
	}
}

func convertLogicalExpr(opname string, left, right sqlparser.Expr, lookup ConverterLookup) (Expr, error) {
	l, err := convertExpr(left, lookup)
	if err != nil {
		return nil, err
	}
	r, err := convertExpr(right, lookup)
	if err != nil {
		return nil, err
	}

	var logic func(l, r boolean) boolean
	switch opname {
	case "AND":
		logic = func(l, r boolean) boolean { return l.and(r) }
	case "OR":
		logic = func(l, r boolean) boolean { return l.or(r) }
	case "XOR":
		logic = func(l, r boolean) boolean { return l.xor(r) }
	default:
		panic("unexpected logical operator")
	}

	return &LogicalExpr{
		BinaryExpr: BinaryExpr{
			Left:  l,
			Right: r,
		},
		op:     logic,
		opname: opname,
	}, nil
}

func convertIsExpr(left sqlparser.Expr, op sqlparser.IsExprOperator, lookup ConverterLookup) (Expr, error) {
	expr, err := convertExpr(left, lookup)
	if err != nil {
		return nil, err
	}

	var check func(result *EvalResult) bool

	switch op {
	case sqlparser.IsNullOp:
		check = func(er *EvalResult) bool { return er.null() }
	case sqlparser.IsNotNullOp:
		check = func(er *EvalResult) bool { return !er.null() }
	case sqlparser.IsTrueOp:
		check = func(er *EvalResult) bool { return er.truthy() == boolTrue }
	case sqlparser.IsNotTrueOp:
		check = func(er *EvalResult) bool { return er.truthy() != boolTrue }
	case sqlparser.IsFalseOp:
		check = func(er *EvalResult) bool { return er.truthy() == boolFalse }
	case sqlparser.IsNotFalseOp:
		check = func(er *EvalResult) bool { return er.truthy() != boolFalse }
	}

	return &IsExpr{
		UnaryExpr: UnaryExpr{expr},
		Op:        op,
		Check:     check,
	}, nil
}

func getCollation(expr sqlparser.Expr, lookup ConverterLookup) collations.TypedCollation {
	collation := collations.TypedCollation{
		Coercibility: collations.CoerceCoercible,
		Repertoire:   collations.RepertoireUnicode,
	}
	if lookup != nil {
		collation.Collation = lookup.CollationForExpr(expr)
		if collation.Collation == collations.Unknown {
			collation.Collation = lookup.DefaultCollation()
		}
	} else {
		collation.Collation = collations.Default()
	}
	return collation
}

func ConvertEx(e sqlparser.Expr, lookup ConverterLookup, simplify bool) (Expr, error) {
	expr, err := convertExpr(e, lookup)
	if err != nil {
		return nil, err
	}
	if simplify {
		var staticEnv ExpressionEnv
		if lookup != nil {
			staticEnv.DefaultCollation = lookup.DefaultCollation()
		} else {
			staticEnv.DefaultCollation = collations.Default()
		}
		expr, err = simplifyExpr(&staticEnv, expr)
	}
	return expr, err
}

// Convert converts between AST expressions and executable expressions
func Convert(e sqlparser.Expr, lookup ConverterLookup) (Expr, error) {
	return ConvertEx(e, lookup, true)
}

func convertExpr(e sqlparser.Expr, lookup ConverterLookup) (Expr, error) {
	switch node := e.(type) {
	case *sqlparser.ColName:
		if lookup == nil {
			return nil, vterrors.Wrap(convertNotSupported(e), "cannot lookup column")
		}
		idx, err := lookup.ColumnLookup(node)
		if err != nil {
			return nil, err
		}
		collation := getCollation(node, lookup)
		return NewColumn(idx, collation), nil
	case *sqlparser.ComparisonExpr:
		expr, err := convertComparisonExpr(node.Operator, node.Left, node.Right, lookup)
		if err != nil {
			return nil, convertNotSupported(e)
		}
		return expr, err
	case sqlparser.Argument:
		collation := getCollation(e, lookup)
		return NewBindVar(string(node), collation), nil
	case sqlparser.ListArg:
		collation := getCollation(e, lookup)
		return NewBindVar(string(node), collation), nil
	case *sqlparser.Literal:
		switch node.Type {
		case sqlparser.IntVal:
			return NewLiteralIntegralFromBytes(node.Bytes())
		case sqlparser.FloatVal:
			return NewLiteralFloatFromBytes(node.Bytes())
		case sqlparser.DecimalVal:
			return NewLiteralDecimalFromBytes(node.Bytes())
		case sqlparser.StrVal:
			collation := getCollation(e, lookup)
			return NewLiteralString(node.Bytes(), collation), nil
		case sqlparser.HexVal:
			return NewLiteralBinaryFromHex(node.Bytes())
		}
	case sqlparser.BoolVal:
		if node {
			return NewLiteralInt(1), nil
		}
		return NewLiteralInt(0), nil
	case *sqlparser.AndExpr:
		return convertLogicalExpr("AND", node.Left, node.Right, lookup)
	case *sqlparser.OrExpr:
		return convertLogicalExpr("OR", node.Left, node.Right, lookup)
	case *sqlparser.XorExpr:
		return convertLogicalExpr("XOR", node.Left, node.Right, lookup)
	case *sqlparser.NotExpr:
		inner, err := convertExpr(node.Expr, lookup)
		if err != nil {
			return nil, err
		}
		return &NotExpr{UnaryExpr{inner}}, nil
	case *sqlparser.BinaryExpr:
		var op ArithmeticOp
		switch node.Operator {
		case sqlparser.PlusOp:
			op = &OpAddition{}
		case sqlparser.MinusOp:
			op = &OpSubstraction{}
		case sqlparser.MultOp:
			op = &OpMultiplication{}
		case sqlparser.DivOp:
			op = &OpDivision{}
		default:
			return nil, convertNotSupported(e)
		}
		left, err := convertExpr(node.Left, lookup)
		if err != nil {
			return nil, err
		}
		right, err := convertExpr(node.Right, lookup)
		if err != nil {
			return nil, err
		}
		return &ArithmeticExpr{
			BinaryExpr: BinaryExpr{
				Left:  left,
				Right: right,
			},
			Op: op,
		}, nil
	case sqlparser.ValTuple:
		var exprs TupleExpr
		for _, expr := range node {
			convertedExpr, err := convertExpr(expr, lookup)
			if err != nil {
				return nil, err
			}
			exprs = append(exprs, convertedExpr)
		}
		return exprs, nil
	case *sqlparser.NullVal:
		return NullExpr, nil
	case *sqlparser.CollateExpr:
		expr, err := convertExpr(node.Expr, lookup)
		if err != nil {
			return nil, err
		}
		coll := collations.Local().LookupByName(node.Collation)
		if coll == nil {
			return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "Unknown collation: '%s'", node.Collation)
		}
		return &CollateExpr{
			UnaryExpr: UnaryExpr{expr},
			TypedCollation: collations.TypedCollation{
				Collation:    coll.ID(),
				Coercibility: collations.CoerceExplicit,
				Repertoire:   collations.RepertoireUnicode,
			},
		}, nil
	case *sqlparser.IntroducerExpr:
		expr, err := convertExpr(node.Expr, lookup)
		if err != nil {
			return nil, err
		}
		var collation collations.ID

		// TODO: case insensitive
		if node.CharacterSet == "_binary" {
			collation = collations.CollationBinaryID
		} else {
			defaultCollation := collations.Local().DefaultCollationForCharset(node.CharacterSet[1:])
			if defaultCollation == nil {
				panic(fmt.Sprintf("unknown character set: %s", node.CharacterSet))
			}
			collation = defaultCollation.ID()
		}

		switch lit := expr.(type) {
		case *Literal:
			switch collation {
			case collations.CollationBinaryID:
				lit.Val.type_ = querypb.Type_VARBINARY
				lit.Val.collation_ = collationBinary
			default:
				lit.Val.type_ = querypb.Type_VARCHAR
				lit.Val.replaceCollationID(collation)
			}
		case *BindVariable:
			switch collation {
			case collations.CollationBinaryID:
				lit.coerceType = querypb.Type_VARBINARY
				lit.coll = collationBinary
			default:
				lit.coerceType = querypb.Type_VARCHAR
				lit.coll.Collation = collation
			}
		default:
			panic("character set introducers are only supported for literals and arguments")
		}
		return expr, nil
	case *sqlparser.IsExpr:
		return convertIsExpr(node.Left, node.Right, lookup)
	case *sqlparser.FuncExpr:
		method := node.Name.Lowered()
		call, ok := builtinFunctions[method]
		if !ok {
			break
		}

		var args TupleExpr
		var aliases []sqlparser.ColIdent
		for _, expr := range node.Exprs {
			aliased, ok := expr.(*sqlparser.AliasedExpr)
			if !ok {
				return nil, convertNotSupported(e)
			}
			convertedExpr, err := convertExpr(aliased.Expr, lookup)
			if err != nil {
				return nil, err
			}
			args = append(args, convertedExpr)
			aliases = append(aliases, aliased.As)
		}
		return &CallExpression{
			Arguments: args,
			Aliases:   aliases,
			Method:    method,
			Call:      call,
		}, nil
	}
	return nil, convertNotSupported(e)
}

var builtinFunctions = map[string]func(*ExpressionEnv, []EvalResult, *EvalResult){
	"coalesce":  builtinFuncCoalesce,
	"greatest":  builtinFuncGreatest,
	"least":     builtinFuncLeast,
	"collation": builtinFuncCollation,
}

func convertNotSupported(e sqlparser.Expr) error {
	return vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "%s: %s", ErrConvertExprNotSupported, sqlparser.String(e))
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package main

import (
	"fmt"
	"github.com/spf13/cobra"
	api "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/apis/kops/registry"
	"k8s.io/kops/pkg/client/simple"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup"
	"k8s.io/kops/upup/pkg/fi/cloudup/awsup"
	"k8s.io/kops/upup/pkg/kutil"
	"k8s.io/kops/util/pkg/tables"
	"k8s.io/kops/util/pkg/vfs"
	"os"
)

type DeleteClusterCmd struct {
	Yes        bool
	Region     string
	External   bool
	Unregister bool
}

var deleteCluster DeleteClusterCmd

func init() {
	cmd := &cobra.Command{
		Use:   "cluster CLUSTERNAME [--yes]",
		Short: "Delete cluster",
		Long:  `Deletes a k8s cluster.`,
		Run: func(cmd *cobra.Command, args []string) {
			err := deleteCluster.Run(args)
			if err != nil {
				exitWithError(err)
			}
		},
	}

	deleteCmd.AddCommand(cmd)

	cmd.Flags().BoolVar(&deleteCluster.Yes, "yes", false, "Delete without confirmation")
	cmd.Flags().BoolVar(&deleteCluster.Unregister, "unregister", false, "Don't delete cloud resources, just unregister the cluster")
	cmd.Flags().BoolVar(&deleteCluster.External, "external", false, "Delete an external cluster")

	cmd.Flags().StringVar(&deleteCluster.Region, "region", "", "region")
}

type getter func(o interface{}) interface{}

func (c *DeleteClusterCmd) Run(args []string) error {
	var configBase vfs.Path
	var clientset simple.Clientset

	err := rootCommand.ProcessArgs(args)
	if err != nil {
		return err
	}

	clusterName := rootCommand.clusterName
	if clusterName == "" {
		return fmt.Errorf("--name is required (for safety)")
	}

	var cloud fi.Cloud
	var cluster *api.Cluster

	if c.External {
		region := c.Region
		if region == "" {
			return fmt.Errorf("--region is required (when --external)")
		}

		tags := map[string]string{"KubernetesCluster": clusterName}
		cloud, err = awsup.NewAWSCloud(region, tags)
		if err != nil {
			return fmt.Errorf("error initializing AWS client: %v", err)
		}
	} else {
		clientset, err = rootCommand.Clientset()
		if err != nil {
			return err
		}

		cluster, err = clientset.Clusters().Get(clusterName)
		if err != nil {
			return err
		}

		if cluster == nil {
			return fmt.Errorf("cluster %q not found", clusterName)
		}

		if clusterName != cluster.ObjectMeta.Name {
			return fmt.Errorf("sanity check failed: cluster name mismatch")
		}

		configBase, err = registry.ConfigBase(cluster)
		if err != nil {
			return err
		}
	}

	wouldDeleteCloudResources := false

	if !c.Unregister {
		if cloud == nil {
			cloud, err = cloudup.BuildCloud(cluster)
			if err != nil {
				return err
			}
		}

		d := &kutil.DeleteCluster{}
		d.ClusterName = clusterName
		d.Cloud = cloud

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

		if len(resources) == 0 {
			fmt.Printf("No cloud resources to delete\n")
		} else {
			wouldDeleteCloudResources = true

			t := &tables.Table{}
			t.AddColumn("TYPE", func(r *kutil.ResourceTracker) string {
				return r.Type
			})
			t.AddColumn("ID", func(r *kutil.ResourceTracker) string {
				return r.ID
			})
			t.AddColumn("NAME", func(r *kutil.ResourceTracker) string {
				return r.Name
			})
			var l []*kutil.ResourceTracker
			for _, v := range resources {
				l = append(l, v)
			}

			err := t.Render(l, os.Stdout, "TYPE", "NAME", "ID")
			if err != nil {
				return err
			}

			if !c.Yes {
				return fmt.Errorf("Must specify --yes to delete")
			}

			fmt.Fprintf(os.Stdout, "\n")

			err = d.DeleteResources(resources)
			if err != nil {
				return err
			}
		}
	}

	if !c.External {
		if !c.Yes {
			if wouldDeleteCloudResources {
				fmt.Printf("\nMust specify --yes to delete cloud resources & unregister cluster\n")
			} else {
				fmt.Printf("\nMust specify --yes to unregister the cluster\n")
			}
			return nil
		}
		err := registry.DeleteAllClusterState(configBase)
		if err != nil {
			return fmt.Errorf("error removing cluster from state store: %v", err)
		}
	}

	b := kutil.NewKubeconfigBuilder()
	b.Context = cluster.ObjectMeta.Name
	b.DeleteKubeConfig()

	fmt.Printf("\nCluster deleted\n")
	return nil
}
<commit_msg>this can't inherit the flag from the parent if it declares its own version of the flag<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package main

import (
	"fmt"
	"os"

	"github.com/spf13/cobra"
	api "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/apis/kops/registry"
	"k8s.io/kops/pkg/client/simple"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup"
	"k8s.io/kops/upup/pkg/fi/cloudup/awsup"
	"k8s.io/kops/upup/pkg/kutil"
	"k8s.io/kops/util/pkg/tables"
	"k8s.io/kops/util/pkg/vfs"
)

type DeleteClusterCmd struct {
	Yes        bool
	Region     string
	External   bool
	Unregister bool
}

var deleteCluster DeleteClusterCmd

func init() {
	cmd := &cobra.Command{
		Use:   "cluster CLUSTERNAME [--yes]",
		Short: "Delete cluster",
		Long:  `Deletes a k8s cluster.`,
		Run: func(cmd *cobra.Command, args []string) {
			err := deleteCluster.Run(args)
			if err != nil {
				exitWithError(err)
			}
		},
	}

	deleteCmd.AddCommand(cmd)

	deleteCluster.Yes = confirmDelete
	cmd.Flags().BoolVar(&deleteCluster.Unregister, "unregister", false, "Don't delete cloud resources, just unregister the cluster")
	cmd.Flags().BoolVar(&deleteCluster.External, "external", false, "Delete an external cluster")

	cmd.Flags().StringVar(&deleteCluster.Region, "region", "", "region")
}

type getter func(o interface{}) interface{}

func (c *DeleteClusterCmd) Run(args []string) error {
	var configBase vfs.Path
	var clientset simple.Clientset

	err := rootCommand.ProcessArgs(args)
	if err != nil {
		return err
	}

	clusterName := rootCommand.clusterName
	if clusterName == "" {
		return fmt.Errorf("--name is required (for safety)")
	}

	var cloud fi.Cloud
	var cluster *api.Cluster

	if c.External {
		region := c.Region
		if region == "" {
			return fmt.Errorf("--region is required (when --external)")
		}

		tags := map[string]string{"KubernetesCluster": clusterName}
		cloud, err = awsup.NewAWSCloud(region, tags)
		if err != nil {
			return fmt.Errorf("error initializing AWS client: %v", err)
		}
	} else {
		clientset, err = rootCommand.Clientset()
		if err != nil {
			return err
		}

		cluster, err = clientset.Clusters().Get(clusterName)
		if err != nil {
			return err
		}

		if cluster == nil {
			return fmt.Errorf("cluster %q not found", clusterName)
		}

		if clusterName != cluster.ObjectMeta.Name {
			return fmt.Errorf("sanity check failed: cluster name mismatch")
		}

		configBase, err = registry.ConfigBase(cluster)
		if err != nil {
			return err
		}
	}

	wouldDeleteCloudResources := false

	if !c.Unregister {
		if cloud == nil {
			cloud, err = cloudup.BuildCloud(cluster)
			if err != nil {
				return err
			}
		}

		d := &kutil.DeleteCluster{}
		d.ClusterName = clusterName
		d.Cloud = cloud

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

		if len(resources) == 0 {
			fmt.Printf("No cloud resources to delete\n")
		} else {
			wouldDeleteCloudResources = true

			t := &tables.Table{}
			t.AddColumn("TYPE", func(r *kutil.ResourceTracker) string {
				return r.Type
			})
			t.AddColumn("ID", func(r *kutil.ResourceTracker) string {
				return r.ID
			})
			t.AddColumn("NAME", func(r *kutil.ResourceTracker) string {
				return r.Name
			})
			var l []*kutil.ResourceTracker
			for _, v := range resources {
				l = append(l, v)
			}

			err := t.Render(l, os.Stdout, "TYPE", "NAME", "ID")
			if err != nil {
				return err
			}

			if !c.Yes {
				return fmt.Errorf("Must specify --yes to delete")
			}

			fmt.Fprintf(os.Stdout, "\n")

			err = d.DeleteResources(resources)
			if err != nil {
				return err
			}
		}
	}

	if !c.External {
		if !c.Yes {
			if wouldDeleteCloudResources {
				fmt.Printf("\nMust specify --yes to delete cloud resources & unregister cluster\n")
			} else {
				fmt.Printf("\nMust specify --yes to unregister the cluster\n")
			}
			return nil
		}
		err := registry.DeleteAllClusterState(configBase)
		if err != nil {
			return fmt.Errorf("error removing cluster from state store: %v", err)
		}
	}

	b := kutil.NewKubeconfigBuilder()
	b.Context = cluster.ObjectMeta.Name
	b.DeleteKubeConfig()

	fmt.Printf("\nCluster deleted\n")
	return nil
}
<|endoftext|>


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

package relay

import (
	"bytes"
	"fmt"
	"github.com/prometheus/statsd_exporter/pkg/clock"
	"net"
	"strings"
	"time"

	"github.com/go-kit/log"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/statsd_exporter/pkg/level"
)

type Relay struct {
	addr          *net.UDPAddr
	bufferChannel chan []byte
	conn          *net.UDPConn
	logger        log.Logger
	packetLength  uint

	packetsTotal   prometheus.Counter
	longLinesTotal prometheus.Counter
}

var (
	relayPacketsTotal = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "statsd_exporter_relay_packets_total",
			Help: "The number of StatsD packets relayed.",
		},
		[]string{"target"},
	)
	relayLongLinesTotal = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "statsd_exporter_relay_long_lines_total",
			Help: "The number lines that were too long to relay.",
		},
		[]string{"target"},
	)
)

// NewRelay creates a statsd UDP relay. It can be used to send copies of statsd raw
// lines to a separate service.
func NewRelay(l log.Logger, target string, packetLength uint) (*Relay, error) {
	addr, err := net.ResolveUDPAddr("udp", target)
	if err != nil {
		return nil, fmt.Errorf("unable to resolve target %s, err: %w", target, err)
	}
	conn, err := net.ListenUDP("udp", nil)
	if err != nil {
		return nil, fmt.Errorf("unable to listen on UDP, err: %w", err)
	}

	c := make(chan []byte, 100)

	r := Relay{
		addr:          addr,
		bufferChannel: c,
		conn:          conn,
		logger:        l,
		packetLength:  packetLength,

		packetsTotal:   relayPacketsTotal.WithLabelValues(target),
		longLinesTotal: relayLongLinesTotal.WithLabelValues(target),
	}

	// Startup the UDP sender.
	go r.relayOutput()

	return &r, nil
}

// relayOutput buffers statsd lines and sends them to the relay target.
func (r *Relay) relayOutput() {
	var buffer bytes.Buffer
	var err error

	relayInterval := clock.NewTicker(1 * time.Second)
	defer relayInterval.Stop()

	for {
		select {
		case <-relayInterval.C:
			err = r.sendPacket(buffer.Bytes())
			if err != nil {
				level.Error(r.logger).Log("msg", "Error sending UDP packet", "error", err)
				return
			}
			// Clear out the buffer.
			buffer.Reset()
		case b := <-r.bufferChannel:
			if uint(len(b)+buffer.Len()) > r.packetLength {
				level.Debug(r.logger).Log("msg", "Buffer full, sending packet", "length", buffer.Len())
				err = r.sendPacket(buffer.Bytes())
				if err != nil {
					level.Error(r.logger).Log("msg", "Error sending UDP packet", "error", err)
					return
				}
				// Seed the new buffer with the new line.
				buffer.Reset()
				buffer.Write(b)
			} else {
				level.Debug(r.logger).Log("msg", "Adding line to buffer", "line", b)
				buffer.Write(b)
			}
		}
	}
}

// sendPacket sends a single relay line to the destination target.
func (r *Relay) sendPacket(buf []byte) error {
	if len(buf) == 0 {
		level.Debug(r.logger).Log("msg", "Empty buffer, nothing to send")
		return nil
	}
	level.Debug(r.logger).Log("msg", "Sending packet", "length", len(buf), "data", buf)
	_, err := r.conn.WriteToUDP(buf, r.addr)
	r.packetsTotal.Inc()
	return err
}

// RelayLine processes a single statsd line and forwards it to the relay target.
func (r *Relay) RelayLine(l string) {
	lineLength := uint(len(l))
	if lineLength == 0 {
		level.Debug(r.logger).Log("msg", "Empty line, not relaying")
		return
	}
	if lineLength > r.packetLength-1 {
		level.Warn(r.logger).Log("msg", "line too long, not relaying", "length", lineLength, "max", r.packetLength)
		r.longLinesTotal.Inc()
		return
	}
	level.Debug(r.logger).Log("msg", "Relaying line", "line", l)
	if !strings.HasSuffix(l, "\n") {
		l = l + "\n"
	}
	r.bufferChannel <- []byte(l)
}
<commit_msg>Adding new metric to relay<commit_after>// Copyright 2021 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package relay

import (
	"bytes"
	"fmt"
	"github.com/prometheus/statsd_exporter/pkg/clock"
	"net"
	"strings"
	"time"

	"github.com/go-kit/log"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/statsd_exporter/pkg/level"
)

type Relay struct {
	addr          *net.UDPAddr
	bufferChannel chan []byte
	conn          *net.UDPConn
	logger        log.Logger
	packetLength  uint

	packetsTotal      prometheus.Counter
	longLinesTotal    prometheus.Counter
	relayedLinesTotal prometheus.Counter
}

var (
	relayPacketsTotal = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "statsd_exporter_relay_packets_total",
			Help: "The number of StatsD packets relayed.",
		},
		[]string{"target"},
	)
	relayLongLinesTotal = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "statsd_exporter_relay_long_lines_total",
			Help: "The number lines that were too long to relay.",
		},
		[]string{"target"},
	)
	relayLinesRelayedTotal = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "statsd_exporter_relay_lines_relayed_total",
			Help: "The number of lines that were buffered to be relayed.",
		},
		[]string{"target"},
	)
)

// NewRelay creates a statsd UDP relay. It can be used to send copies of statsd raw
// lines to a separate service.
func NewRelay(l log.Logger, target string, packetLength uint) (*Relay, error) {
	addr, err := net.ResolveUDPAddr("udp", target)
	if err != nil {
		return nil, fmt.Errorf("unable to resolve target %s, err: %w", target, err)
	}
	conn, err := net.ListenUDP("udp", nil)
	if err != nil {
		return nil, fmt.Errorf("unable to listen on UDP, err: %w", err)
	}

	c := make(chan []byte, 100)

	r := Relay{
		addr:          addr,
		bufferChannel: c,
		conn:          conn,
		logger:        l,
		packetLength:  packetLength,

		packetsTotal:      relayPacketsTotal.WithLabelValues(target),
		longLinesTotal:    relayLongLinesTotal.WithLabelValues(target),
		relayedLinesTotal: relayLinesRelayedTotal.WithLabelValues(target),
	}

	// Startup the UDP sender.
	go r.relayOutput()

	return &r, nil
}

// relayOutput buffers statsd lines and sends them to the relay target.
func (r *Relay) relayOutput() {
	var buffer bytes.Buffer
	var err error

	relayInterval := clock.NewTicker(1 * time.Second)
	defer relayInterval.Stop()

	for {
		select {
		case <-relayInterval.C:
			err = r.sendPacket(buffer.Bytes())
			if err != nil {
				level.Error(r.logger).Log("msg", "Error sending UDP packet", "error", err)
				return
			}
			// Clear out the buffer.
			buffer.Reset()
		case b := <-r.bufferChannel:
			if uint(len(b)+buffer.Len()) > r.packetLength {
				level.Debug(r.logger).Log("msg", "Buffer full, sending packet", "length", buffer.Len())
				err = r.sendPacket(buffer.Bytes())
				if err != nil {
					level.Error(r.logger).Log("msg", "Error sending UDP packet", "error", err)
					return
				}
				// Seed the new buffer with the new line.
				buffer.Reset()
				buffer.Write(b)
			} else {
				level.Debug(r.logger).Log("msg", "Adding line to buffer", "line", b)
				buffer.Write(b)
			}
		}
	}
}

// sendPacket sends a single relay line to the destination target.
func (r *Relay) sendPacket(buf []byte) error {
	if len(buf) == 0 {
		level.Debug(r.logger).Log("msg", "Empty buffer, nothing to send")
		return nil
	}
	level.Debug(r.logger).Log("msg", "Sending packet", "length", len(buf), "data", buf)
	_, err := r.conn.WriteToUDP(buf, r.addr)
	r.packetsTotal.Inc()
	return err
}

// RelayLine processes a single statsd line and forwards it to the relay target.
func (r *Relay) RelayLine(l string) {
	lineLength := uint(len(l))
	if lineLength == 0 {
		level.Debug(r.logger).Log("msg", "Empty line, not relaying")
		return
	}
	if lineLength > r.packetLength-1 {
		level.Warn(r.logger).Log("msg", "line too long, not relaying", "length", lineLength, "max", r.packetLength)
		r.longLinesTotal.Inc()
		return
	}
	level.Debug(r.logger).Log("msg", "Relaying line", "line", l)
	if !strings.HasSuffix(l, "\n") {
		l = l + "\n"
	}
	r.relayedLinesTotal.Inc()
	r.bufferChannel <- []byte(l)
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Vitess Authors.

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

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

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

package semantics

import (
	"fmt"

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

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

type (
	// analyzer is a struct to work with analyzing the query.
	analyzer struct {
		si SchemaInformation

		Tables    []TableInfo
		scopes    []*scope
		exprDeps  map[sqlparser.Expr]TableSet
		err       error
		currentDb string

		selectScope map[*sqlparser.Select]*scope
	}
)

// newAnalyzer create the semantic analyzer
func newAnalyzer(dbName string, si SchemaInformation) *analyzer {
	return &analyzer{
		exprDeps:    map[sqlparser.Expr]TableSet{},
		selectScope: map[*sqlparser.Select]*scope{},
		currentDb:   dbName,
		si:          si,
	}
}

// Analyze analyzes the parsed query.
func Analyze(statement sqlparser.Statement, currentDb string, si SchemaInformation) (*SemTable, error) {
	analyzer := newAnalyzer(currentDb, si)
	// Initial scope
	err := analyzer.analyze(statement)
	if err != nil {
		return nil, err
	}
	return &SemTable{exprDependencies: analyzer.exprDeps, Tables: analyzer.Tables, selectScope: analyzer.selectScope}, nil
}

// analyzeDown pushes new scopes when we encounter sub queries,
// and resolves the table a column is using
func (a *analyzer) analyzeDown(cursor *sqlparser.Cursor) bool {
	current := a.currentScope()
	n := cursor.Node()
	switch node := n.(type) {
	case *sqlparser.Select:
		a.push(newScope(current))
		a.selectScope[node] = a.currentScope()
		if err := a.analyzeTableExprs(node.From); err != nil {
			a.err = err
			return false
		}
	case *sqlparser.DerivedTable:
		a.err = Gen4NotSupportedF("derived tables")
	case *sqlparser.TableExprs:
		// this has already been visited when we encountered the SELECT struct
		return false

	// we don't need to push new scope for sub queries since we do that for SELECT and UNION

	case *sqlparser.Union:
		a.push(newScope(current))
	case *sqlparser.ColName:
		t, err := a.resolveColumn(node, current)
		if err != nil {
			a.err = err
		} else {
			a.exprDeps[node] = t
		}
	}

	// this is the visitor going down the tree. Returning false here would just not visit the children
	// to the current node, but that is not what we want if we have encountered an error.
	// In order to abort the whole visitation, we have to return true here and then return false in the `analyzeUp` method
	return true
}

func (a *analyzer) resolveColumn(colName *sqlparser.ColName, current *scope) (TableSet, error) {
	var t TableInfo
	var err error
	if colName.Qualifier.IsEmpty() {
		t, err = a.resolveUnQualifiedColumn(current, colName)
	} else {
		t, err = a.resolveQualifiedColumn(current, colName)
	}
	if err != nil {
		return 0, err
	}
	return a.tableSetFor(t.GetExpr()), nil
}

func (a *analyzer) analyzeTableExprs(tablExprs sqlparser.TableExprs) error {
	for _, tableExpr := range tablExprs {
		if err := a.analyzeTableExpr(tableExpr); err != nil {
			return err
		}
	}
	return nil
}

func (a *analyzer) analyzeTableExpr(tableExpr sqlparser.TableExpr) error {
	switch table := tableExpr.(type) {
	case *sqlparser.AliasedTableExpr:
		return a.bindTable(table, table.Expr)
	case *sqlparser.JoinTableExpr:
		if table.Join != sqlparser.NormalJoinType {
			return Gen4NotSupportedF("join type %s", table.Join.ToString())
		}
		if err := a.analyzeTableExpr(table.LeftExpr); err != nil {
			return err
		}
		if err := a.analyzeTableExpr(table.RightExpr); err != nil {
			return err
		}
	case *sqlparser.ParenTableExpr:
		return a.analyzeTableExprs(table.Exprs)
	}
	return nil
}

// resolveQualifiedColumn handles `tabl.col` expressions
func (a *analyzer) resolveQualifiedColumn(current *scope, expr *sqlparser.ColName) (TableInfo, error) {
	// search up the scope stack until we find a match
	for current != nil {
		for _, table := range current.tables {
			if table.Matches(expr.Qualifier) {
				return table, nil
			}
		}
		current = current.parent
	}
	return nil, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.BadFieldError, "Unknown table referenced by '%s'", sqlparser.String(expr))
}

// resolveUnQualifiedColumn
func (a *analyzer) resolveUnQualifiedColumn(current *scope, expr *sqlparser.ColName) (TableInfo, error) {
	if len(current.tables) == 1 {
		return current.tables[0], nil
	}

	var tblInfo TableInfo
	for _, tbl := range current.tables {
		if !tbl.Authoritative() {
			return nil, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.NonUniqError, fmt.Sprintf("Column '%s' in field list is ambiguous", sqlparser.String(expr)))
		}
		for _, col := range tbl.GetColumns() {
			if expr.Name.String() == col.Name {
				if tblInfo != nil {
					return nil, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.NonUniqError, fmt.Sprintf("Column '%s' in field list is ambiguous", sqlparser.String(expr)))
				}
				tblInfo = tbl
			}
		}
	}
	return tblInfo, nil
}

func (a *analyzer) tableSetFor(t *sqlparser.AliasedTableExpr) TableSet {
	for i, t2 := range a.Tables {
		if t == t2.GetExpr() {
			return TableSet(1 << i)
		}
	}
	panic("unknown table")
}

func (a *analyzer) createTable(t sqlparser.TableName, alias *sqlparser.AliasedTableExpr, tbl *vindexes.Table) TableInfo {
	dbName := t.Qualifier.String()
	if dbName == "" {
		dbName = a.currentDb
	}
	if alias.As.IsEmpty() {
		return &RealTable{
			dbName:    dbName,
			tableName: t.Name.String(),
			ASTNode:   alias,
			Table:     tbl,
		}
	}
	return &AliasedTable{
		tableName: alias.As.String(),
		ASTNode:   alias,
		Table:     tbl,
	}
}

func (a *analyzer) bindTable(alias *sqlparser.AliasedTableExpr, expr sqlparser.SimpleTableExpr) error {
	switch t := expr.(type) {
	case *sqlparser.DerivedTable:
		return Gen4NotSupportedF("derived table")
	case sqlparser.TableName:
		tbl, vdx, _, _, _, err := a.si.FindTableOrVindex(t)
		if err != nil {
			return err
		}
		if tbl == nil && vdx != nil {
			return Gen4NotSupportedF("vindex in FROM")
		}
		scope := a.currentScope()
		tableInfo := a.createTable(t, alias, tbl)

		a.Tables = append(a.Tables, tableInfo)
		return scope.addTable(tableInfo)
	}
	return nil
}

func (a *analyzer) analyze(statement sqlparser.Statement) error {
	_ = sqlparser.Rewrite(statement, a.analyzeDown, a.analyzeUp)
	return a.err
}

func (a *analyzer) analyzeUp(cursor *sqlparser.Cursor) bool {
	switch cursor.Node().(type) {
	case *sqlparser.Union, *sqlparser.Select:
		a.popScope()
	}
	return a.shouldContinue()
}

func (a *analyzer) shouldContinue() bool {
	return a.err == nil
}

func (a *analyzer) push(s *scope) {
	a.scopes = append(a.scopes, s)
}

func (a *analyzer) popScope() {
	l := len(a.scopes) - 1
	a.scopes = a.scopes[:l]
}

func (a *analyzer) currentScope() *scope {
	size := len(a.scopes)
	if size == 0 {
		return nil
	}
	return a.scopes[size-1]
}

// Gen4NotSupportedF returns a common error for shortcomings in the gen4 planner
func Gen4NotSupportedF(format string, args ...interface{}) error {
	return vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "gen4 does not yet support: "+format, args...)
}
<commit_msg>Addition of write and read scopes<commit_after>/*
Copyright 2020 The Vitess Authors.

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

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

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

package semantics

import (
	"fmt"

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

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

type (
	// analyzer is a struct to work with analyzing the query.
	analyzer struct {
		si SchemaInformation

		Tables    []TableInfo
		scopes    []*scope
		exprDeps  map[sqlparser.Expr]TableSet
		err       error
		currentDb string

		rScope map[*sqlparser.Select]*scope
		wScope map[*sqlparser.Select]*scope
	}
)

// newAnalyzer create the semantic analyzer
func newAnalyzer(dbName string, si SchemaInformation) *analyzer {
	return &analyzer{
		exprDeps:  map[sqlparser.Expr]TableSet{},
		rScope:    map[*sqlparser.Select]*scope{},
		wScope:    map[*sqlparser.Select]*scope{},
		currentDb: dbName,
		si:        si,
	}
}

// Analyze analyzes the parsed query.
func Analyze(statement sqlparser.Statement, currentDb string, si SchemaInformation) (*SemTable, error) {
	analyzer := newAnalyzer(currentDb, si)
	// Initial scope
	err := analyzer.analyze(statement)
	if err != nil {
		return nil, err
	}
	return &SemTable{exprDependencies: analyzer.exprDeps, Tables: analyzer.Tables, selectScope: analyzer.rScope}, nil
}

// analyzeDown pushes new scopes when we encounter sub queries,
// and resolves the table a column is using
func (a *analyzer) analyzeDown(cursor *sqlparser.Cursor) bool {
	current := a.currentScope()
	n := cursor.Node()
	switch node := n.(type) {
	case *sqlparser.Select:
		a.push(newScope(current))
		a.rScope[node] = a.currentScope()
		if err := a.analyzeTableExprs(node.From); err != nil {
			a.err = err
			return false
		}
	case *sqlparser.DerivedTable:
		a.err = Gen4NotSupportedF("derived tables")
	case *sqlparser.TableExprs:
		// this has already been visited when we encountered the SELECT struct
		return false

	// we don't need to push new scope for sub queries since we do that for SELECT and UNION

	case *sqlparser.Union:
		a.push(newScope(current))
	case *sqlparser.ColName:
		t, err := a.resolveColumn(node, current)
		if err != nil {
			a.err = err
		} else {
			a.exprDeps[node] = t
		}
	}

	// this is the visitor going down the tree. Returning false here would just not visit the children
	// to the current node, but that is not what we want if we have encountered an error.
	// In order to abort the whole visitation, we have to return true here and then return false in the `analyzeUp` method
	return true
}

func (a *analyzer) resolveColumn(colName *sqlparser.ColName, current *scope) (TableSet, error) {
	var t TableInfo
	var err error
	if colName.Qualifier.IsEmpty() {
		t, err = a.resolveUnQualifiedColumn(current, colName)
	} else {
		t, err = a.resolveQualifiedColumn(current, colName)
	}
	if err != nil {
		return 0, err
	}
	return a.tableSetFor(t.GetExpr()), nil
}

func (a *analyzer) analyzeTableExprs(tablExprs sqlparser.TableExprs) error {
	for _, tableExpr := range tablExprs {
		if err := a.analyzeTableExpr(tableExpr); err != nil {
			return err
		}
	}
	return nil
}

func (a *analyzer) analyzeTableExpr(tableExpr sqlparser.TableExpr) error {
	switch table := tableExpr.(type) {
	case *sqlparser.AliasedTableExpr:
		return a.bindTable(table, table.Expr)
	case *sqlparser.JoinTableExpr:
		if table.Join != sqlparser.NormalJoinType {
			return Gen4NotSupportedF("join type %s", table.Join.ToString())
		}
		if err := a.analyzeTableExpr(table.LeftExpr); err != nil {
			return err
		}
		if err := a.analyzeTableExpr(table.RightExpr); err != nil {
			return err
		}
	case *sqlparser.ParenTableExpr:
		return a.analyzeTableExprs(table.Exprs)
	}
	return nil
}

// resolveQualifiedColumn handles `tabl.col` expressions
func (a *analyzer) resolveQualifiedColumn(current *scope, expr *sqlparser.ColName) (TableInfo, error) {
	// search up the scope stack until we find a match
	for current != nil {
		for _, table := range current.tables {
			if table.Matches(expr.Qualifier) {
				return table, nil
			}
		}
		current = current.parent
	}
	return nil, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.BadFieldError, "Unknown table referenced by '%s'", sqlparser.String(expr))
}

// resolveUnQualifiedColumn
func (a *analyzer) resolveUnQualifiedColumn(current *scope, expr *sqlparser.ColName) (TableInfo, error) {
	if len(current.tables) == 1 {
		return current.tables[0], nil
	}

	var tblInfo TableInfo
	for _, tbl := range current.tables {
		if !tbl.Authoritative() {
			return nil, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.NonUniqError, fmt.Sprintf("Column '%s' in field list is ambiguous", sqlparser.String(expr)))
		}
		for _, col := range tbl.GetColumns() {
			if expr.Name.String() == col.Name {
				if tblInfo != nil {
					return nil, vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.NonUniqError, fmt.Sprintf("Column '%s' in field list is ambiguous", sqlparser.String(expr)))
				}
				tblInfo = tbl
			}
		}
	}
	return tblInfo, nil
}

func (a *analyzer) tableSetFor(t *sqlparser.AliasedTableExpr) TableSet {
	for i, t2 := range a.Tables {
		if t == t2.GetExpr() {
			return TableSet(1 << i)
		}
	}
	panic("unknown table")
}

func (a *analyzer) createTable(t sqlparser.TableName, alias *sqlparser.AliasedTableExpr, tbl *vindexes.Table) TableInfo {
	dbName := t.Qualifier.String()
	if dbName == "" {
		dbName = a.currentDb
	}
	if alias.As.IsEmpty() {
		return &RealTable{
			dbName:    dbName,
			tableName: t.Name.String(),
			ASTNode:   alias,
			Table:     tbl,
		}
	}
	return &AliasedTable{
		tableName: alias.As.String(),
		ASTNode:   alias,
		Table:     tbl,
	}
}

func (a *analyzer) bindTable(alias *sqlparser.AliasedTableExpr, expr sqlparser.SimpleTableExpr) error {
	switch t := expr.(type) {
	case *sqlparser.DerivedTable:
		return Gen4NotSupportedF("derived table")
	case sqlparser.TableName:
		tbl, vdx, _, _, _, err := a.si.FindTableOrVindex(t)
		if err != nil {
			return err
		}
		if tbl == nil && vdx != nil {
			return Gen4NotSupportedF("vindex in FROM")
		}
		scope := a.currentScope()
		tableInfo := a.createTable(t, alias, tbl)

		a.Tables = append(a.Tables, tableInfo)
		return scope.addTable(tableInfo)
	}
	return nil
}

func (a *analyzer) analyze(statement sqlparser.Statement) error {
	_ = sqlparser.Rewrite(statement, a.analyzeDown, a.analyzeUp)
	return a.err
}

func (a *analyzer) analyzeUp(cursor *sqlparser.Cursor) bool {
	switch cursor.Node().(type) {
	case *sqlparser.Union, *sqlparser.Select:
		a.popScope()
	}
	return a.shouldContinue()
}

func (a *analyzer) shouldContinue() bool {
	return a.err == nil
}

func (a *analyzer) push(s *scope) {
	a.scopes = append(a.scopes, s)
}

func (a *analyzer) popScope() {
	l := len(a.scopes) - 1
	a.scopes = a.scopes[:l]
}

func (a *analyzer) currentScope() *scope {
	size := len(a.scopes)
	if size == 0 {
		return nil
	}
	return a.scopes[size-1]
}

// Gen4NotSupportedF returns a common error for shortcomings in the gen4 planner
func Gen4NotSupportedF(format string, args ...interface{}) error {
	return vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "gen4 does not yet support: "+format, args...)
}
<|endoftext|>