<commit_before>package cwl

// Graphs represents "$graph" field in CWL.
type Graphs []*Root

// Graph represents an element of "steps"
type Graph struct {
	Run *Root
}

// New constructs "Graphs" from interface.
func (_ Graphs) New(i interface{}) Graphs {
	dest := Graphs{}
	switch x := i.(type) {
	case []interface{}:
		for _, v := range x {
			g := new(Root)
			g = g.AsStep(v)
			dest = append(dest, g)
		}
	}
	return dest
}

// New constructs "Step" from interface.
func (graph Graph) New(i interface{}) Graph {
	dest := Graph{}
	return dest
}
<commit_msg>Refactoring remove unused code<commit_after>package cwl

// Graphs represents "$graph" field in CWL.
type Graphs []*Root

// Graph represents an element of "steps"
type Graph struct {
	Run *Root
}

// New constructs "Graphs" from interface.
func (_ Graphs) New(i interface{}) Graphs {
	dest := Graphs{}
	switch x := i.(type) {
	case []interface{}:
		for _, v := range x {
			g := new(Root)
			g = g.AsStep(v)
			dest = append(dest, g)
		}
	}
	return dest
}
<|endoftext|>


<commit_before>// Copyright (c) Alex Ellis 2017. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

package metrics

import (
	"encoding/json"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"time"

	"log"

	"github.com/openfaas/faas/gateway/requests"
	"github.com/prometheus/client_golang/prometheus"
)

// Exporter is a prometheus exporter
type Exporter struct {
	metricOptions MetricOptions
	services      []requests.Function
}

// NewExporter creates a new exporter for the OpenFaaS gateway metrics
func NewExporter(options MetricOptions) *Exporter {
	return &Exporter{
		metricOptions: options,
		services:      []requests.Function{},
	}
}

// Describe is to describe the metrics for Prometheus
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	e.metricOptions.GatewayFunctionInvocation.Describe(ch)
	e.metricOptions.GatewayFunctionsHistogram.Describe(ch)
	e.metricOptions.ServiceReplicasGauge.Describe(ch)
}

// Collect collects data to be consumed by prometheus
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	e.metricOptions.GatewayFunctionInvocation.Collect(ch)
	e.metricOptions.GatewayFunctionsHistogram.Collect(ch)

	e.metricOptions.ServiceReplicasGauge.Reset()
	for _, service := range e.services {
		e.metricOptions.ServiceReplicasGauge.
			WithLabelValues(service.Name).
			Set(float64(service.Replicas))
	}
	e.metricOptions.ServiceReplicasGauge.Collect(ch)
}

// StartServiceWatcher starts a ticker and collects service replica counts to expose to prometheus
func (e *Exporter) StartServiceWatcher(endpointURL url.URL, metricsOptions MetricOptions, label string, interval time.Duration) {
	ticker := time.NewTicker(interval)
	quit := make(chan struct{})
	proxyClient := http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			DialContext: (&net.Dialer{
				Timeout:   3 * time.Second,
				KeepAlive: 0,
			}).DialContext,
			MaxIdleConns:          1,
			DisableKeepAlives:     true,
			IdleConnTimeout:       120 * time.Millisecond,
			ExpectContinueTimeout: 1500 * time.Millisecond,
		},
	}

	go func() {
		for {
			select {
			case <-ticker.C:

				get, _ := http.NewRequest(http.MethodGet, endpointURL.String()+"system/functions", nil)

				services := []requests.Function{}
				res, err := proxyClient.Do(get)
				if err != nil {
					log.Println(err)
					continue
				}
				bytesOut, readErr := ioutil.ReadAll(res.Body)
				if readErr != nil {
					log.Println(err)
					continue
				}
				unmarshalErr := json.Unmarshal(bytesOut, &services)
				if unmarshalErr != nil {
					log.Println(err)
					continue
				}

				e.services = services

				break
			case <-quit:
				return
			}
		}
	}()
}
<commit_msg>Updated copyright<commit_after>// Copyright (c) Alex Ellis 2017
// Copyright (c) 2018 OpenFaaS Author(s)
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

package metrics

import (
	"encoding/json"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"time"

	"log"

	"github.com/openfaas/faas/gateway/requests"
	"github.com/prometheus/client_golang/prometheus"
)

// Exporter is a prometheus exporter
type Exporter struct {
	metricOptions MetricOptions
	services      []requests.Function
}

// NewExporter creates a new exporter for the OpenFaaS gateway metrics
func NewExporter(options MetricOptions) *Exporter {
	return &Exporter{
		metricOptions: options,
		services:      []requests.Function{},
	}
}

// Describe is to describe the metrics for Prometheus
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	e.metricOptions.GatewayFunctionInvocation.Describe(ch)
	e.metricOptions.GatewayFunctionsHistogram.Describe(ch)
	e.metricOptions.ServiceReplicasGauge.Describe(ch)
}

// Collect collects data to be consumed by prometheus
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	e.metricOptions.GatewayFunctionInvocation.Collect(ch)
	e.metricOptions.GatewayFunctionsHistogram.Collect(ch)

	e.metricOptions.ServiceReplicasGauge.Reset()
	for _, service := range e.services {
		e.metricOptions.ServiceReplicasGauge.
			WithLabelValues(service.Name).
			Set(float64(service.Replicas))
	}
	e.metricOptions.ServiceReplicasGauge.Collect(ch)
}

// StartServiceWatcher starts a ticker and collects service replica counts to expose to prometheus
func (e *Exporter) StartServiceWatcher(endpointURL url.URL, metricsOptions MetricOptions, label string, interval time.Duration) {
	ticker := time.NewTicker(interval)
	quit := make(chan struct{})
	proxyClient := http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			DialContext: (&net.Dialer{
				Timeout:   3 * time.Second,
				KeepAlive: 0,
			}).DialContext,
			MaxIdleConns:          1,
			DisableKeepAlives:     true,
			IdleConnTimeout:       120 * time.Millisecond,
			ExpectContinueTimeout: 1500 * time.Millisecond,
		},
	}

	go func() {
		for {
			select {
			case <-ticker.C:

				get, _ := http.NewRequest(http.MethodGet, endpointURL.String()+"system/functions", nil)

				services := []requests.Function{}
				res, err := proxyClient.Do(get)
				if err != nil {
					log.Println(err)
					continue
				}
				bytesOut, readErr := ioutil.ReadAll(res.Body)
				if readErr != nil {
					log.Println(err)
					continue
				}
				unmarshalErr := json.Unmarshal(bytesOut, &services)
				if unmarshalErr != nil {
					log.Println(err)
					continue
				}

				e.services = services

				break
			case <-quit:
				return
			}
		}
	}()
}
<|endoftext|>


<commit_before>// This file was auto-generated using createmock. See the following page for
// more information:
//
//     https://github.com/jacobsa/oglemock
//

package mock_gcs

import (
	fmt "fmt"
	gcs "github.com/jacobsa/gcloud/gcs"
	oglemock "github.com/jacobsa/oglemock"
	context "golang.org/x/net/context"
	storage "google.golang.org/cloud/storage"
	io "io"
	runtime "runtime"
	unsafe "unsafe"
)

type MockBucket interface {
	gcs.Bucket
	oglemock.MockObject
}

type mockBucket struct {
	controller  oglemock.Controller
	description string
}

func NewMockBucket(
	c oglemock.Controller,
	desc string) MockBucket {
	return &mockBucket{
		controller:  c,
		description: desc,
	}
}

func (m *mockBucket) Oglemock_Id() uintptr {
	return uintptr(unsafe.Pointer(m))
}

func (m *mockBucket) Oglemock_Description() string {
	return m.description
}

func (m *mockBucket) CreateObject(p0 context.Context, p1 *gcs.CreateObjectRequest) (o0 *storage.Object, o1 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"CreateObject",
		file,
		line,
		[]interface{}{p0, p1})

	if len(retVals) != 2 {
		panic(fmt.Sprintf("mockBucket.CreateObject: invalid return values: %v", retVals))
	}

	// o0 *storage.Object
	if retVals[0] != nil {
		o0 = retVals[0].(*storage.Object)
	}

	// o1 error
	if retVals[1] != nil {
		o1 = retVals[1].(error)
	}

	return
}

func (m *mockBucket) DeleteObject(p0 context.Context, p1 string) (o0 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"DeleteObject",
		file,
		line,
		[]interface{}{p0, p1})

	if len(retVals) != 1 {
		panic(fmt.Sprintf("mockBucket.DeleteObject: invalid return values: %v", retVals))
	}

	// o0 error
	if retVals[0] != nil {
		o0 = retVals[0].(error)
	}

	return
}

func (m *mockBucket) ListObjects(p0 context.Context, p1 *storage.Query) (o0 *storage.Objects, o1 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"ListObjects",
		file,
		line,
		[]interface{}{p0, p1})

	if len(retVals) != 2 {
		panic(fmt.Sprintf("mockBucket.ListObjects: invalid return values: %v", retVals))
	}

	// o0 *storage.Objects
	if retVals[0] != nil {
		o0 = retVals[0].(*storage.Objects)
	}

	// o1 error
	if retVals[1] != nil {
		o1 = retVals[1].(error)
	}

	return
}

func (m *mockBucket) Name() (o0 string) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"Name",
		file,
		line,
		[]interface{}{})

	if len(retVals) != 1 {
		panic(fmt.Sprintf("mockBucket.Name: invalid return values: %v", retVals))
	}

	// o0 string
	if retVals[0] != nil {
		o0 = retVals[0].(string)
	}

	return
}

func (m *mockBucket) NewReader(p0 context.Context, p1 string) (o0 io.ReadCloser, o1 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"NewReader",
		file,
		line,
		[]interface{}{p0, p1})

	if len(retVals) != 2 {
		panic(fmt.Sprintf("mockBucket.NewReader: invalid return values: %v", retVals))
	}

	// o0 io.ReadCloser
	if retVals[0] != nil {
		o0 = retVals[0].(io.ReadCloser)
	}

	// o1 error
	if retVals[1] != nil {
		o1 = retVals[1].(error)
	}

	return
}

func (m *mockBucket) UpdateObject(p0 context.Context, p1 *gcs.UpdateObjectRequest) (o0 *storage.Object, o1 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"UpdateObject",
		file,
		line,
		[]interface{}{p0, p1})

	if len(retVals) != 2 {
		panic(fmt.Sprintf("mockBucket.UpdateObject: invalid return values: %v", retVals))
	}

	// o0 *storage.Object
	if retVals[0] != nil {
		o0 = retVals[0].(*storage.Object)
	}

	// o1 error
	if retVals[1] != nil {
		o1 = retVals[1].(error)
	}

	return
}
<commit_msg>Added mockBucket.StatObject.<commit_after>// This file was auto-generated using createmock. See the following page for
// more information:
//
//     https://github.com/jacobsa/oglemock
//

package mock_gcs

import (
	fmt "fmt"
	gcs "github.com/jacobsa/gcloud/gcs"
	oglemock "github.com/jacobsa/oglemock"
	context "golang.org/x/net/context"
	storage "google.golang.org/cloud/storage"
	io "io"
	runtime "runtime"
	unsafe "unsafe"
)

type MockBucket interface {
	gcs.Bucket
	oglemock.MockObject
}

type mockBucket struct {
	controller  oglemock.Controller
	description string
}

func NewMockBucket(
	c oglemock.Controller,
	desc string) MockBucket {
	return &mockBucket{
		controller:  c,
		description: desc,
	}
}

func (m *mockBucket) Oglemock_Id() uintptr {
	return uintptr(unsafe.Pointer(m))
}

func (m *mockBucket) Oglemock_Description() string {
	return m.description
}

func (m *mockBucket) CreateObject(p0 context.Context, p1 *gcs.CreateObjectRequest) (o0 *storage.Object, o1 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"CreateObject",
		file,
		line,
		[]interface{}{p0, p1})

	if len(retVals) != 2 {
		panic(fmt.Sprintf("mockBucket.CreateObject: invalid return values: %v", retVals))
	}

	// o0 *storage.Object
	if retVals[0] != nil {
		o0 = retVals[0].(*storage.Object)
	}

	// o1 error
	if retVals[1] != nil {
		o1 = retVals[1].(error)
	}

	return
}

func (m *mockBucket) DeleteObject(p0 context.Context, p1 string) (o0 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"DeleteObject",
		file,
		line,
		[]interface{}{p0, p1})

	if len(retVals) != 1 {
		panic(fmt.Sprintf("mockBucket.DeleteObject: invalid return values: %v", retVals))
	}

	// o0 error
	if retVals[0] != nil {
		o0 = retVals[0].(error)
	}

	return
}

func (m *mockBucket) ListObjects(p0 context.Context, p1 *storage.Query) (o0 *storage.Objects, o1 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"ListObjects",
		file,
		line,
		[]interface{}{p0, p1})

	if len(retVals) != 2 {
		panic(fmt.Sprintf("mockBucket.ListObjects: invalid return values: %v", retVals))
	}

	// o0 *storage.Objects
	if retVals[0] != nil {
		o0 = retVals[0].(*storage.Objects)
	}

	// o1 error
	if retVals[1] != nil {
		o1 = retVals[1].(error)
	}

	return
}

func (m *mockBucket) Name() (o0 string) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"Name",
		file,
		line,
		[]interface{}{})

	if len(retVals) != 1 {
		panic(fmt.Sprintf("mockBucket.Name: invalid return values: %v", retVals))
	}

	// o0 string
	if retVals[0] != nil {
		o0 = retVals[0].(string)
	}

	return
}

func (m *mockBucket) NewReader(p0 context.Context, p1 string) (o0 io.ReadCloser, o1 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"NewReader",
		file,
		line,
		[]interface{}{p0, p1})

	if len(retVals) != 2 {
		panic(fmt.Sprintf("mockBucket.NewReader: invalid return values: %v", retVals))
	}

	// o0 io.ReadCloser
	if retVals[0] != nil {
		o0 = retVals[0].(io.ReadCloser)
	}

	// o1 error
	if retVals[1] != nil {
		o1 = retVals[1].(error)
	}

	return
}

func (m *mockBucket) StatObject(p0 context.Context, p1 *gcs.StatObjectRequest) (o0 *storage.Object, o1 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"StatObject",
		file,
		line,
		[]interface{}{p0, p1})

	if len(retVals) != 2 {
		panic(fmt.Sprintf("mockBucket.StatObject: invalid return values: %v", retVals))
	}

	// o0 *storage.Object
	if retVals[0] != nil {
		o0 = retVals[0].(*storage.Object)
	}

	// o1 error
	if retVals[1] != nil {
		o1 = retVals[1].(error)
	}

	return
}

func (m *mockBucket) UpdateObject(p0 context.Context, p1 *gcs.UpdateObjectRequest) (o0 *storage.Object, o1 error) {
	// Get a file name and line number for the caller.
	_, file, line, _ := runtime.Caller(1)

	// Hand the call off to the controller, which does most of the work.
	retVals := m.controller.HandleMethodCall(
		m,
		"UpdateObject",
		file,
		line,
		[]interface{}{p0, p1})

	if len(retVals) != 2 {
		panic(fmt.Sprintf("mockBucket.UpdateObject: invalid return values: %v", retVals))
	}

	// o0 *storage.Object
	if retVals[0] != nil {
		o0 = retVals[0].(*storage.Object)
	}

	// o1 error
	if retVals[1] != nil {
		o1 = retVals[1].(error)
	}

	return
}
<|endoftext|>


<commit_before>package unfurl

import (
	"context"
	"regexp"
	"sync"

	"github.com/mvdan/xurls"

	"github.com/keybase/client/go/chat/types"
	"github.com/keybase/client/go/chat/utils"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
)

type ExtractorHitTyp int

const (
	ExtractorHitUnfurl ExtractorHitTyp = iota
	ExtractorHitPrompt
)

type ExtractorHit struct {
	URL string
	Typ ExtractorHitTyp
}

type Extractor struct {
	utils.DebugLabeler

	urlRegexp      *regexp.Regexp
	quoteRegexp    *regexp.Regexp
	maxHits        int
	exemptionsLock sync.Mutex
	exemptions     map[string]*WhitelistExemptionList
}

func NewExtractor(log logger.Logger) *Extractor {
	return &Extractor{
		DebugLabeler: utils.NewDebugLabeler(log, "Extractor", false),
		urlRegexp:    xurls.Strict(),
		quoteRegexp:  regexp.MustCompile("`[^`]*`"),
		exemptions:   make(map[string]*WhitelistExemptionList),
		maxHits:      5,
	}
}

func (e *Extractor) getExemptionList(uid gregor1.UID) (res *WhitelistExemptionList) {
	e.exemptionsLock.Lock()
	defer e.exemptionsLock.Unlock()
	var ok bool
	res, ok = e.exemptions[uid.String()]
	if !ok {
		res = NewWhitelistExemptionList()
		e.exemptions[uid.String()] = res
	}
	return res
}

func (e *Extractor) isWhitelistHit(ctx context.Context, convID chat1.ConversationID, msgID chat1.MessageID,
	hit string, whitelist map[string]bool, exemptions *WhitelistExemptionList) bool {
	domain, err := GetDomain(hit)
	if err != nil {
		e.Debug(ctx, "isWhitelistHit: failed to get domain: %s", err)
		return false
	}
	if whitelist[domain] {
		return true
	}
	// Check exemptions
	if exemptions.Use(convID, msgID, domain) {
		e.Debug(ctx, "isWhitelistHit: hit exemption for domain, letting through")
		return true
	}
	return false
}

func (e *Extractor) Extract(ctx context.Context, uid gregor1.UID, convID chat1.ConversationID,
	msgID chat1.MessageID, body string, userSettings *Settings) (res []ExtractorHit, err error) {
	defer e.Trace(ctx, func() error { return err }, "Extract")()
	settings, err := userSettings.Get(ctx, uid)
	if err != nil {
		return res, err
	}
	if settings.Mode == chat1.UnfurlMode_NEVER {
		return res, nil
	}
	body = e.quoteRegexp.ReplaceAllString(body, "")
	hits := e.urlRegexp.FindAllString(body, -1)
	for _, h := range hits {
		ehit := ExtractorHit{
			URL: h,
			Typ: ExtractorHitPrompt,
		}
		switch settings.Mode {
		case chat1.UnfurlMode_ALWAYS:
			ehit.Typ = ExtractorHitUnfurl
		case chat1.UnfurlMode_WHITELISTED:
			if e.isWhitelistHit(ctx, convID, msgID, h, settings.Whitelist, e.getExemptionList(uid)) {
				ehit.Typ = ExtractorHitUnfurl
			}
		}
		res = append(res, ehit)
		if len(res) >= e.maxHits {
			e.Debug(ctx, "Extract: max hits reached, aborting")
			break
		}
	}
	return res, nil
}

func (e *Extractor) AddWhitelistExemption(ctx context.Context, uid gregor1.UID,
	exemption types.WhitelistExemption) {
	defer e.Trace(ctx, func() error { return nil }, "AddWhitelistExemption")()
	e.getExemptionList(uid).Add(exemption)
}
<commit_msg>only fetch unfurl settings if relevant for message<commit_after>package unfurl

import (
	"context"
	"regexp"
	"sync"

	"github.com/mvdan/xurls"

	"github.com/keybase/client/go/chat/types"
	"github.com/keybase/client/go/chat/utils"
	"github.com/keybase/client/go/logger"
	"github.com/keybase/client/go/protocol/chat1"
	"github.com/keybase/client/go/protocol/gregor1"
)

type ExtractorHitTyp int

const (
	ExtractorHitUnfurl ExtractorHitTyp = iota
	ExtractorHitPrompt
)

type ExtractorHit struct {
	URL string
	Typ ExtractorHitTyp
}

type Extractor struct {
	utils.DebugLabeler

	urlRegexp      *regexp.Regexp
	quoteRegexp    *regexp.Regexp
	maxHits        int
	exemptionsLock sync.Mutex
	exemptions     map[string]*WhitelistExemptionList
}

func NewExtractor(log logger.Logger) *Extractor {
	return &Extractor{
		DebugLabeler: utils.NewDebugLabeler(log, "Extractor", false),
		urlRegexp:    xurls.Strict(),
		quoteRegexp:  regexp.MustCompile("`[^`]*`"),
		exemptions:   make(map[string]*WhitelistExemptionList),
		maxHits:      5,
	}
}

func (e *Extractor) getExemptionList(uid gregor1.UID) (res *WhitelistExemptionList) {
	e.exemptionsLock.Lock()
	defer e.exemptionsLock.Unlock()
	var ok bool
	res, ok = e.exemptions[uid.String()]
	if !ok {
		res = NewWhitelistExemptionList()
		e.exemptions[uid.String()] = res
	}
	return res
}

func (e *Extractor) isWhitelistHit(ctx context.Context, convID chat1.ConversationID, msgID chat1.MessageID,
	hit string, whitelist map[string]bool, exemptions *WhitelistExemptionList) bool {
	domain, err := GetDomain(hit)
	if err != nil {
		e.Debug(ctx, "isWhitelistHit: failed to get domain: %s", err)
		return false
	}
	if whitelist[domain] {
		return true
	}
	// Check exemptions
	if exemptions.Use(convID, msgID, domain) {
		e.Debug(ctx, "isWhitelistHit: hit exemption for domain, letting through")
		return true
	}
	return false
}

func (e *Extractor) Extract(ctx context.Context, uid gregor1.UID, convID chat1.ConversationID,
	msgID chat1.MessageID, body string, userSettings *Settings) (res []ExtractorHit, err error) {
	defer e.Trace(ctx, func() error { return err }, "Extract")()
	body = e.quoteRegexp.ReplaceAllString(body, "")
	hits := e.urlRegexp.FindAllString(body, -1)
	if len(hits) == 0 {
		return res, nil
	}
	settings, err := userSettings.Get(ctx, uid)
	if err != nil {
		return res, err
	}
	if settings.Mode == chat1.UnfurlMode_NEVER {
		return res, nil
	}
	for _, h := range hits {
		ehit := ExtractorHit{
			URL: h,
			Typ: ExtractorHitPrompt,
		}
		switch settings.Mode {
		case chat1.UnfurlMode_ALWAYS:
			ehit.Typ = ExtractorHitUnfurl
		case chat1.UnfurlMode_WHITELISTED:
			if e.isWhitelistHit(ctx, convID, msgID, h, settings.Whitelist, e.getExemptionList(uid)) {
				ehit.Typ = ExtractorHitUnfurl
			}
		}
		res = append(res, ehit)
		if len(res) >= e.maxHits {
			e.Debug(ctx, "Extract: max hits reached, aborting")
			break
		}
	}
	return res, nil
}

func (e *Extractor) AddWhitelistExemption(ctx context.Context, uid gregor1.UID,
	exemption types.WhitelistExemption) {
	defer e.Trace(ctx, func() error { return nil }, "AddWhitelistExemption")()
	e.getExemptionList(uid).Add(exemption)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"

	"github.com/danmane/abalone/go/api"
	"github.com/danmane/abalone/go/game"
)

var (
	playAgainstHuman = flag.Bool("playAgainstHuman", false, "play against human on frontend rather than AI vs AI")
	humanPort        = flag.String("humanPort", "1337", "port for javascript frontend")
	aiPort1          = flag.String("aiPort1", "3423", "port for first ai")
	aiPort2          = flag.String("aiPort2", "3424", "port for second ai (if present)")
)

func main() {
	flag.Parse()
	if err := run(); err != nil {
		log.Fatal(err)
	}
}

func run() error {
	whiteAI := api.Player{}
	blackAI := api.Player{}
	whiteAgent := PlayerInstance{Player: whiteAI, Port: *aiPort1}
	blackAgent := PlayerInstance{Player: blackAI, Port: *aiPort2}
	start := game.Standard
	result := playAIGame(whiteAgent, blackAgent, start)
	fmt.Println(result)
	return nil
}

type PlayerInstance struct {
	Player api.Player
	Port   string
}

func gameFromAI(port string, state *game.State) (*game.State, error) {
	var buf bytes.Buffer
	if err := json.NewEncoder(&buf).Encode(state); err != nil {
		return nil, err
	}
	resp, err := http.Post("http://localhost:"+port+"/move", "application/json", &buf)
	if err != nil {
		return nil, err
	}
	responseGame := &game.State{}
	if err := json.NewDecoder(resp.Body).Decode(responseGame); err != nil {
		return nil, err
	}
	resp.Body.Close()
	if !state.ValidFuture(responseGame) {
		return nil, fmt.Errorf("game parsed correctly, but isn't a valid future")
	}
	return responseGame, nil
}

func playAIGame(whiteAgent, blackAgent PlayerInstance, startState game.State) api.GameResult {
	states := []game.State{startState}
	currentGame := &startState
	victory := api.NoVictory
	outcome := game.NullOutcome
	for !currentGame.GameOver() {
		var nextAI PlayerInstance
		if currentGame.NextPlayer == game.White {
			nextAI = whiteAgent
		} else {
			nextAI = blackAgent
		}
		futureGame, err := gameFromAI(currentGame, nextAI.Port)
		if err != nil {
			fmt.Println(err)
			victory = api.InvalidResponse
			outcome = currentGame.NextPlayer.Loses()
			return api.GameResult{
				White:         whiteAgent.Player,
				Black:         blackAgent.Player,
				Outcome:       outcome,
				VictoryReason: victory,
				States:        states,
			}
		}
		currentGame = futureGame
		states = append(states, *currentGame)
	}

	outcome = currentGame.Outcome()
	loser := outcome.Loser()
	if loser != game.NullPlayer && currentGame.NumPieces(loser) <= currentGame.LossThreshold {
		victory = api.StonesDepleted
	} else {
		victory = api.MovesDepleted
	}
	return api.GameResult{
		White:         whiteAgent.Player,
		Black:         blackAgent.Player,
		Outcome:       outcome,
		VictoryReason: victory,
		States:        states,
	}

}
<commit_msg>Fix compilation error<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"

	"github.com/danmane/abalone/go/api"
	"github.com/danmane/abalone/go/game"
)

var (
	playAgainstHuman = flag.Bool("playAgainstHuman", false, "play against human on frontend rather than AI vs AI")
	humanPort        = flag.String("humanPort", "1337", "port for javascript frontend")
	aiPort1          = flag.String("aiPort1", "3423", "port for first ai")
	aiPort2          = flag.String("aiPort2", "3424", "port for second ai (if present)")
)

func main() {
	flag.Parse()
	if err := run(); err != nil {
		log.Fatal(err)
	}
}

func run() error {
	whiteAI := api.Player{}
	blackAI := api.Player{}
	whiteAgent := PlayerInstance{Player: whiteAI, Port: *aiPort1}
	blackAgent := PlayerInstance{Player: blackAI, Port: *aiPort2}
	start := game.Standard
	result := playAIGame(whiteAgent, blackAgent, start)
	fmt.Println(result)
	return nil
}

type PlayerInstance struct {
	Player api.Player
	Port   string
}

func gameFromAI(port string, state *game.State) (*game.State, error) {
	var buf bytes.Buffer
	if err := json.NewEncoder(&buf).Encode(state); err != nil {
		return nil, err
	}
	resp, err := http.Post("http://localhost:"+port+"/move", "application/json", &buf)
	if err != nil {
		return nil, err
	}
	responseGame := &game.State{}
	if err := json.NewDecoder(resp.Body).Decode(responseGame); err != nil {
		return nil, err
	}
	resp.Body.Close()
	if !state.ValidFuture(responseGame) {
		return nil, fmt.Errorf("game parsed correctly, but isn't a valid future")
	}
	return responseGame, nil
}

func playAIGame(whiteAgent, blackAgent PlayerInstance, startState game.State) api.GameResult {
	states := []game.State{startState}
	currentGame := &startState
	victory := api.NoVictory
	outcome := game.NullOutcome
	for !currentGame.GameOver() {
		var nextAI PlayerInstance
		if currentGame.NextPlayer == game.White {
			nextAI = whiteAgent
		} else {
			nextAI = blackAgent
		}
		futureGame, err := gameFromAI(nextAI.Port, currentGame)
		if err != nil {
			fmt.Println(err)
			victory = api.InvalidResponse
			outcome = currentGame.NextPlayer.Loses()
			return api.GameResult{
				White:         whiteAgent.Player,
				Black:         blackAgent.Player,
				Outcome:       outcome,
				VictoryReason: victory,
				States:        states,
			}
		}
		currentGame = futureGame
		states = append(states, *currentGame)
	}

	outcome = currentGame.Outcome()
	loser := outcome.Loser()
	if loser != game.NullPlayer && currentGame.NumPieces(loser) <= currentGame.LossThreshold {
		victory = api.StonesDepleted
	} else {
		victory = api.MovesDepleted
	}
	return api.GameResult{
		White:         whiteAgent.Player,
		Black:         blackAgent.Player,
		Outcome:       outcome,
		VictoryReason: victory,
		States:        states,
	}

}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// +build !production

package externals

import (
	"testing"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/pvlsource"
	"github.com/keybase/client/go/uidmap"
)

func SetupTest(tb testing.TB, name string, depth int) (tc libkb.TestContext) {
	ret := libkb.SetupTest(tb, name, depth+1)

	ret.G.SetServices(GetServices())
	ret.G.SetUIDMapper(uidmap.NewUIDMap(10000))
	pvlsource.NewPvlSourceAndInstall(ret.G)
	return ret
}
<commit_msg>Don't use real UIDMapper for now, breaks test (will debug separately)<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// +build !production

package externals

import (
	"testing"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/pvlsource"
)

func SetupTest(tb testing.TB, name string, depth int) (tc libkb.TestContext) {
	ret := libkb.SetupTest(tb, name, depth+1)

	ret.G.SetServices(GetServices())
	// XXX this breaks TestUPAKDeadlock
	// ret.G.SetUIDMapper(uidmap.NewUIDMap(10000))
	pvlsource.NewPvlSourceAndInstall(ret.G)
	return ret
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Tigera, Inc. All rights reserved.
//
// 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 rules

import (
	"github.com/projectcalico/felix/go/felix/hashutils"
	. "github.com/projectcalico/felix/go/felix/iptables"
	"github.com/projectcalico/felix/go/felix/proto"
)

func (r *ruleRenderer) WorkloadDispatchChains(endpoints map[proto.WorkloadEndpointID]*proto.WorkloadEndpoint) []*Chain {
	toEndpointRules := make([]Rule, 0, len(endpoints)+1)
	fromEndpointRules := make([]Rule, 0, len(endpoints)+1)
	for _, endpoint := range endpoints {
		fromEndpointRules = append(fromEndpointRules, Rule{
			Match: Match().InInterface(endpoint.Name),
			Action: GotoAction{
				Target: WorkloadEndpointChainName(WorkloadFromEndpointPfx, endpoint),
			},
		})
		toEndpointRules = append(toEndpointRules, Rule{
			Match: Match().OutInterface(endpoint.Name),
			Action: GotoAction{
				Target: WorkloadEndpointChainName(WorkloadToEndpointPfx, endpoint),
			},
		})
	}

	fromEndpointRules = append(fromEndpointRules, Rule{
		Action: DropAction{},
	})
	toEndpointRules = append(toEndpointRules, Rule{
		Action: DropAction{},
	})

	fromEndpointDispatchChain := Chain{
		Name:  DispatchFromWorkloadEndpoint,
		Rules: fromEndpointRules,
	}
	toEndpointDispatchChain := Chain{
		Name:  DispatchToWorkloadEndpoint,
		Rules: toEndpointRules,
	}

	return []*Chain{&toEndpointDispatchChain, &fromEndpointDispatchChain}
}

func (r *ruleRenderer) WorkloadEndpointToIptablesChains(epID *proto.WorkloadEndpointID, endpoint *proto.WorkloadEndpoint) []*Chain {
	inRules := []Rule{}
	outRules := []Rule{}

	// Start by ensuring that the accept mark bit is clear, policies set that bit to indicate
	// that they accepted the packet.
	inRules = append(inRules, Rule{
		Action: ClearMarkAction{
			Mark: r.IptablesMarkAccept,
		},
	})
	outRules = append(outRules, Rule{
		Action: ClearMarkAction{
			Mark: r.IptablesMarkAccept,
		},
	})

	// TODO(smc) Police the MAC?

	for _, tier := range endpoint.Tiers {
		// For each tier,  clear the "accepted by tier" mark.
		inRules = append(inRules, Rule{
			Comment: "Start of tier " + tier.Name,
			Action: ClearMarkAction{
				Mark: r.IptablesMarkNextTier,
			},
		})
		outRules = append(outRules, Rule{
			Comment: "Start of tier " + tier.Name,
			Action: ClearMarkAction{
				Mark: r.IptablesMarkNextTier,
			},
		})
		// Then, jump to each policy in turn.
		for _, polID := range tier.Policies {
			inPolChainName := PolicyChainName(
				PolicyInboundPfx,
				&proto.PolicyID{Tier: tier.Name, Name: polID},
			)
			inRules = append(inRules,
				Rule{
					Match:  Match().MarkClear(r.IptablesMarkNextTier),
					Action: JumpAction{Target: inPolChainName},
				},
				// If policy marked packet as accepted, it returns, setting the
				// accept mark bit.  If that is set, return from this chain.
				Rule{
					Match:   Match().MarkSet(r.IptablesMarkAccept),
					Action:  ReturnAction{},
					Comment: "Return if policy accepted",
				})
			outPolChainName := PolicyChainName(
				PolicyOutboundPfx,
				&proto.PolicyID{Tier: tier.Name, Name: polID},
			)
			outRules = append(outRules,
				Rule{
					Match:  Match().MarkClear(r.IptablesMarkNextTier),
					Action: JumpAction{Target: outPolChainName},
				},
				// If policy marked packet as accepted, it returns, setting the
				// accept mark bit.  If that is set, return from this chain.
				Rule{
					Match:   Match().MarkSet(r.IptablesMarkAccept),
					Action:  ReturnAction{},
					Comment: "Return if policy accepted",
				})
		}
		// If no policy in the tier marked the packet as next-tier, drop the packet.
		inRules = append(inRules, r.DropRules(Match().MarkClear(r.IptablesMarkNextTier), "Drop if no policies passed packet")...)
		outRules = append(outRules, r.DropRules(Match().MarkClear(r.IptablesMarkNextTier), "Drop if no policies passed packet")...)
	}

	// Then, jump to each profile in turn.
	for _, profileID := range endpoint.ProfileIds {
		inProfChainName := ProfileChainName(PolicyInboundPfx, &proto.ProfileID{Name: profileID})
		outProfChainName := ProfileChainName(PolicyOutboundPfx, &proto.ProfileID{Name: profileID})
		inRules = append(inRules,
			Rule{Action: JumpAction{Target: inProfChainName}},
			// If policy marked packet as accepted, it returns, setting the
			// accept mark bit.  If that is set, return from this chain.
			Rule{
				Match:   Match().MarkSet(r.IptablesMarkAccept),
				Action:  ReturnAction{},
				Comment: "Return if profile accepted",
			})
		outRules = append(outRules,
			Rule{Action: JumpAction{Target: outProfChainName}},
			// If policy marked packet as accepted, it returns, setting the
			// accept mark bit.  If that is set, return from this chain.
			Rule{
				Match:   Match().MarkSet(r.IptablesMarkAccept),
				Action:  ReturnAction{},
				Comment: "Return if profile accepted",
			})
	}

	inRules = append(inRules, r.DropRules(Match(), "Drop if no profiles matched")...)
	outRules = append(outRules, r.DropRules(Match(), "Drop if no profiles matched")...)

	toEndpointChain := Chain{
		Name:  WorkloadEndpointChainName(WorkloadToEndpointPfx, endpoint),
		Rules: inRules,
	}
	fromEndpointChain := Chain{
		Name:  WorkloadEndpointChainName(WorkloadFromEndpointPfx, endpoint),
		Rules: outRules,
	}
	return []*Chain{&toEndpointChain, &fromEndpointChain}
}

func (r *ruleRenderer) HostDispatchChains(map[proto.HostEndpointID]*proto.HostEndpoint) []*Chain {
	panic("Not implemented")
	return nil
}

func (r *ruleRenderer) HostEndpointToIptablesChains(epID *proto.HostEndpointID, endpoint *proto.HostEndpoint) []*Chain {
	panic("Not implemented")

	// TODO(smc) Failsafe chains

	return nil
}

func WorkloadEndpointChainName(prefix string, endpoint *proto.WorkloadEndpoint) string {
	return hashutils.GetLengthLimitedID(
		prefix,
		endpoint.Name,
		MaxChainNameLength,
	)
}
<commit_msg>Make WorkloadEndpointChainName usable for host endpoints also<commit_after>// Copyright (c) 2016 Tigera, Inc. All rights reserved.
//
// 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 rules

import (
	"github.com/projectcalico/felix/go/felix/hashutils"
	. "github.com/projectcalico/felix/go/felix/iptables"
	"github.com/projectcalico/felix/go/felix/proto"
)

func (r *ruleRenderer) WorkloadDispatchChains(endpoints map[proto.WorkloadEndpointID]*proto.WorkloadEndpoint) []*Chain {
	toEndpointRules := make([]Rule, 0, len(endpoints)+1)
	fromEndpointRules := make([]Rule, 0, len(endpoints)+1)
	for _, endpoint := range endpoints {
		fromEndpointRules = append(fromEndpointRules, Rule{
			Match: Match().InInterface(endpoint.Name),
			Action: GotoAction{
				Target: EndpointChainName(WorkloadFromEndpointPfx, endpoint.Name),
			},
		})
		toEndpointRules = append(toEndpointRules, Rule{
			Match: Match().OutInterface(endpoint.Name),
			Action: GotoAction{
				Target: EndpointChainName(WorkloadToEndpointPfx, endpoint.Name),
			},
		})
	}

	fromEndpointRules = append(fromEndpointRules, Rule{
		Action: DropAction{},
	})
	toEndpointRules = append(toEndpointRules, Rule{
		Action: DropAction{},
	})

	fromEndpointDispatchChain := Chain{
		Name:  DispatchFromWorkloadEndpoint,
		Rules: fromEndpointRules,
	}
	toEndpointDispatchChain := Chain{
		Name:  DispatchToWorkloadEndpoint,
		Rules: toEndpointRules,
	}

	return []*Chain{&toEndpointDispatchChain, &fromEndpointDispatchChain}
}

func (r *ruleRenderer) WorkloadEndpointToIptablesChains(epID *proto.WorkloadEndpointID, endpoint *proto.WorkloadEndpoint) []*Chain {
	inRules := []Rule{}
	outRules := []Rule{}

	// Start by ensuring that the accept mark bit is clear, policies set that bit to indicate
	// that they accepted the packet.
	inRules = append(inRules, Rule{
		Action: ClearMarkAction{
			Mark: r.IptablesMarkAccept,
		},
	})
	outRules = append(outRules, Rule{
		Action: ClearMarkAction{
			Mark: r.IptablesMarkAccept,
		},
	})

	// TODO(smc) Police the MAC?

	for _, tier := range endpoint.Tiers {
		// For each tier,  clear the "accepted by tier" mark.
		inRules = append(inRules, Rule{
			Comment: "Start of tier " + tier.Name,
			Action: ClearMarkAction{
				Mark: r.IptablesMarkNextTier,
			},
		})
		outRules = append(outRules, Rule{
			Comment: "Start of tier " + tier.Name,
			Action: ClearMarkAction{
				Mark: r.IptablesMarkNextTier,
			},
		})
		// Then, jump to each policy in turn.
		for _, polID := range tier.Policies {
			inPolChainName := PolicyChainName(
				PolicyInboundPfx,
				&proto.PolicyID{Tier: tier.Name, Name: polID},
			)
			inRules = append(inRules,
				Rule{
					Match:  Match().MarkClear(r.IptablesMarkNextTier),
					Action: JumpAction{Target: inPolChainName},
				},
				// If policy marked packet as accepted, it returns, setting the
				// accept mark bit.  If that is set, return from this chain.
				Rule{
					Match:   Match().MarkSet(r.IptablesMarkAccept),
					Action:  ReturnAction{},
					Comment: "Return if policy accepted",
				})
			outPolChainName := PolicyChainName(
				PolicyOutboundPfx,
				&proto.PolicyID{Tier: tier.Name, Name: polID},
			)
			outRules = append(outRules,
				Rule{
					Match:  Match().MarkClear(r.IptablesMarkNextTier),
					Action: JumpAction{Target: outPolChainName},
				},
				// If policy marked packet as accepted, it returns, setting the
				// accept mark bit.  If that is set, return from this chain.
				Rule{
					Match:   Match().MarkSet(r.IptablesMarkAccept),
					Action:  ReturnAction{},
					Comment: "Return if policy accepted",
				})
		}
		// If no policy in the tier marked the packet as next-tier, drop the packet.
		inRules = append(inRules, r.DropRules(Match().MarkClear(r.IptablesMarkNextTier), "Drop if no policies passed packet")...)
		outRules = append(outRules, r.DropRules(Match().MarkClear(r.IptablesMarkNextTier), "Drop if no policies passed packet")...)
	}

	// Then, jump to each profile in turn.
	for _, profileID := range endpoint.ProfileIds {
		inProfChainName := ProfileChainName(PolicyInboundPfx, &proto.ProfileID{Name: profileID})
		outProfChainName := ProfileChainName(PolicyOutboundPfx, &proto.ProfileID{Name: profileID})
		inRules = append(inRules,
			Rule{Action: JumpAction{Target: inProfChainName}},
			// If policy marked packet as accepted, it returns, setting the
			// accept mark bit.  If that is set, return from this chain.
			Rule{
				Match:   Match().MarkSet(r.IptablesMarkAccept),
				Action:  ReturnAction{},
				Comment: "Return if profile accepted",
			})
		outRules = append(outRules,
			Rule{Action: JumpAction{Target: outProfChainName}},
			// If policy marked packet as accepted, it returns, setting the
			// accept mark bit.  If that is set, return from this chain.
			Rule{
				Match:   Match().MarkSet(r.IptablesMarkAccept),
				Action:  ReturnAction{},
				Comment: "Return if profile accepted",
			})
	}

	inRules = append(inRules, r.DropRules(Match(), "Drop if no profiles matched")...)
	outRules = append(outRules, r.DropRules(Match(), "Drop if no profiles matched")...)

	toEndpointChain := Chain{
		Name:  EndpointChainName(WorkloadToEndpointPfx, endpoint.Name),
		Rules: inRules,
	}
	fromEndpointChain := Chain{
		Name:  EndpointChainName(WorkloadFromEndpointPfx, endpoint.Name),
		Rules: outRules,
	}
	return []*Chain{&toEndpointChain, &fromEndpointChain}
}

func (r *ruleRenderer) HostDispatchChains(map[proto.HostEndpointID]*proto.HostEndpoint) []*Chain {
	panic("Not implemented")
	return nil
}

func (r *ruleRenderer) HostEndpointToIptablesChains(epID *proto.HostEndpointID, endpoint *proto.HostEndpoint) []*Chain {
	panic("Not implemented")

	// TODO(smc) Failsafe chains

	return nil
}

func EndpointChainName(prefix string, name string) string {
	return hashutils.GetLengthLimitedID(
		prefix,
		name,
		MaxChainNameLength,
	)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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

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

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

package schema

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestParseAlterTableOptions(t *testing.T) {
	type expect struct {
		schema, table, options string
	}
	tests := map[string]expect{
		"add column i int, drop column d":                               {schema: "", table: "", options: "add column i int, drop column d"},
		"  add column i int, drop column d  ":                           {schema: "", table: "", options: "add column i int, drop column d"},
		"alter table t add column i int, drop column d":                 {schema: "", table: "t", options: "add column i int, drop column d"},
		"alter    table   t      add column i int, drop column d":       {schema: "", table: "t", options: "add column i int, drop column d"},
		"alter table `t` add column i int, drop column d":               {schema: "", table: "t", options: "add column i int, drop column d"},
		"alter table `scm`.`t` add column i int, drop column d":         {schema: "scm", table: "t", options: "add column i int, drop column d"},
		"alter table `scm`.t add column i int, drop column d":           {schema: "scm", table: "t", options: "add column i int, drop column d"},
		"alter table scm.`t` add column i int, drop column d":           {schema: "scm", table: "t", options: "add column i int, drop column d"},
		"alter table scm.t add column i int, drop column d":             {schema: "scm", table: "t", options: "add column i int, drop column d"},
		"  alter       table   scm.`t` add column i int, drop column d": {schema: "scm", table: "t", options: "add column i int, drop column d"},
		"ALTER  table scm.t ADD COLUMN i int, DROP COLUMN d":            {schema: "scm", table: "t", options: "ADD COLUMN i int, DROP COLUMN d"},
		"ALTER TABLE scm.t ADD COLUMN i int, DROP COLUMN d":             {schema: "scm", table: "t", options: "ADD COLUMN i int, DROP COLUMN d"},
	}
	for query, expect := range tests {
		schema, table, options := ParseAlterTableOptions(query)
		assert.Equal(t, expect.schema, schema)
		assert.Equal(t, expect.table, table)
		assert.Equal(t, expect.options, options)
	}
}

func TestNormalizeOnlineDDL(t *testing.T) {
	type expect struct {
		sqls    []string
		isError bool
	}
	tests := map[string]expect{
		"alter table t add column i int, drop column d": {sqls: []string{"alter table t add column i int, drop column d"}},
		"create table t(id int primary key)":            {sqls: []string{"create table t(id int primary key)"}},
		"drop table t":                                  {sqls: []string{"drop table t"}},
		"drop table if exists t":                        {sqls: []string{"drop table if exists t"}},
		"drop table t1, t2, t3":                         {sqls: []string{"drop table t1", "drop table t2", "drop table t3"}},
		"drop table if exists t1, t2, t3":               {sqls: []string{"drop table if exists t1", "drop table if exists t2", "drop table if exists t3"}},
		"create index i_idx on t(id)":                   {sqls: []string{"alter table t add index i_idx (id)"}},
		"create index i_idx on t(name(12))":             {sqls: []string{"alter table t add index i_idx (`name`(12))"}},
		"create index i_idx on t(id, `ts`, name(12))":   {sqls: []string{"alter table t add index i_idx (id, ts, `name`(12))"}},
		"create unique index i_idx on t(id)":            {sqls: []string{"alter table t add unique index i_idx (id)"}},
		"create index i_idx using btree on t(id)":       {sqls: []string{"alter table t add index i_idx using btree (id)"}},
		"create index with syntax error i_idx on t(id)": {isError: true},
		"select * from t":                               {isError: true},
		"drop database t":                               {isError: true},
	}
	for query, expect := range tests {
		normalized, err := NormalizeOnlineDDL(query)
		if expect.isError {
			assert.Error(t, err)
		} else {
			assert.NoError(t, err)
			sqls := []string{}
			for _, n := range normalized {
				sqls = append(sqls, n.SQL)
			}
			assert.Equal(t, expect.sqls, sqls)
		}
	}
}
<commit_msg>fix parser_test unit test<commit_after>/*
Copyright 2019 The Vitess Authors.

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

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

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

package schema

import (
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestParseAlterTableOptions(t *testing.T) {
	type expect struct {
		schema, table, options string
	}
	tests := map[string]expect{
		"add column i int, drop column d":                               {schema: "", table: "", options: "add column i int, drop column d"},
		"  add column i int, drop column d  ":                           {schema: "", table: "", options: "add column i int, drop column d"},
		"alter table t add column i int, drop column d":                 {schema: "", table: "t", options: "add column i int, drop column d"},
		"alter    table   t      add column i int, drop column d":       {schema: "", table: "t", options: "add column i int, drop column d"},
		"alter table `t` add column i int, drop column d":               {schema: "", table: "t", options: "add column i int, drop column d"},
		"alter table `scm`.`t` add column i int, drop column d":         {schema: "scm", table: "t", options: "add column i int, drop column d"},
		"alter table `scm`.t add column i int, drop column d":           {schema: "scm", table: "t", options: "add column i int, drop column d"},
		"alter table scm.`t` add column i int, drop column d":           {schema: "scm", table: "t", options: "add column i int, drop column d"},
		"alter table scm.t add column i int, drop column d":             {schema: "scm", table: "t", options: "add column i int, drop column d"},
		"  alter       table   scm.`t` add column i int, drop column d": {schema: "scm", table: "t", options: "add column i int, drop column d"},
		"ALTER  table scm.t ADD COLUMN i int, DROP COLUMN d":            {schema: "scm", table: "t", options: "ADD COLUMN i int, DROP COLUMN d"},
		"ALTER TABLE scm.t ADD COLUMN i int, DROP COLUMN d":             {schema: "scm", table: "t", options: "ADD COLUMN i int, DROP COLUMN d"},
	}
	for query, expect := range tests {
		schema, table, options := ParseAlterTableOptions(query)
		assert.Equal(t, expect.schema, schema)
		assert.Equal(t, expect.table, table)
		assert.Equal(t, expect.options, options)
	}
}

func TestNormalizeOnlineDDL(t *testing.T) {
	type expect struct {
		sqls    []string
		isError bool
	}
	tests := map[string]expect{
		"alter table t add column i int, drop column d": {sqls: []string{"alter table t add column i int, drop column d"}},
		"create table t(id int primary key)":            {sqls: []string{"create table t(id int primary key)"}},
		"drop table t":                                  {sqls: []string{"drop table t"}},
		"drop table if exists t":                        {sqls: []string{"drop table if exists t"}},
		"drop table t1, t2, t3":                         {sqls: []string{"drop table t1", "drop table t2", "drop table t3"}},
		"drop table if exists t1, t2, t3":               {sqls: []string{"drop table if exists t1", "drop table if exists t2", "drop table if exists t3"}},
		"create index i_idx on t(id)":                   {sqls: []string{"alter table t add index i_idx (id)"}},
		"create index i_idx on t(name(12))":             {sqls: []string{"alter table t add index i_idx (`name`(12))"}},
		"create index i_idx on t(id, `ts`, name(12))":   {sqls: []string{"alter table t add index i_idx (id, ts, `name`(12))"}},
		"create unique index i_idx on t(id)":            {sqls: []string{"alter table t add unique index i_idx (id)"}},
		"create index i_idx using btree on t(id)":       {sqls: []string{"alter table t add index i_idx (id) using btree"}},
		"create index with syntax error i_idx on t(id)": {isError: true},
		"select * from t":                               {isError: true},
		"drop database t":                               {isError: true},
	}
	for query, expect := range tests {
		t.Run(query, func(t *testing.T) {
			normalized, err := NormalizeOnlineDDL(query)
			if expect.isError {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
				sqls := []string{}
				for _, n := range normalized {
					sqls = append(sqls, n.SQL)
				}
				assert.Equal(t, expect.sqls, sqls)
			}
		})
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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

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

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

package sqlparser

// analyzer.go contains utility analysis functions.

import (
	"fmt"
	"strings"
	"unicode"

	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/vterrors"

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

// StatementType encodes the type of a SQL statement
type StatementType int

// These constants are used to identify the SQL statement type.
// Changing this list will require reviewing all calls to Preview.
const (
	StmtSelect StatementType = iota
	StmtStream
	StmtInsert
	StmtReplace
	StmtUpdate
	StmtDelete
	StmtDDL
	StmtBegin
	StmtCommit
	StmtRollback
	StmtSet
	StmtShow
	StmtUse
	StmtOther
	StmtUnknown
	StmtComment
	StmtPriv
	StmtExplain
	StmtSavepoint
	StmtSRollback
	StmtRelease
	StmtVStream
	StmtLockTables
	StmtUnlockTables
	StmtFlush
	StmtCallProc
)

//ASTToStatementType returns a StatementType from an AST stmt
func ASTToStatementType(stmt Statement) StatementType {
	switch stmt.(type) {
	case *Select, *Union:
		return StmtSelect
	case *Insert:
		return StmtInsert
	case *Update:
		return StmtUpdate
	case *Delete:
		return StmtDelete
	case *Set:
		return StmtSet
	case *Show:
		return StmtShow
	case DDLStatement, DBDDLStatement, *AlterVschema:
		return StmtDDL
	case *Use:
		return StmtUse
	case *OtherRead, *OtherAdmin, *Load:
		return StmtOther
	case Explain:
		return StmtExplain
	case *Begin:
		return StmtBegin
	case *Commit:
		return StmtCommit
	case *Rollback:
		return StmtRollback
	case *Savepoint:
		return StmtSavepoint
	case *SRollback:
		return StmtSRollback
	case *Release:
		return StmtRelease
	case *LockTables:
		return StmtLockTables
	case *UnlockTables:
		return StmtUnlockTables
	case *Flush:
		return StmtFlush
	case *CallProc:
		return StmtCallProc
	default:
		return StmtUnknown
	}
}

//CanNormalize takes Statement and returns if the statement can be normalized.
func CanNormalize(stmt Statement) bool {
	switch stmt.(type) {
	case *Select, *Union, *Insert, *Update, *Delete, *Set, *CallProc: // TODO: we could merge this logic into ASTrewriter
		return true
	}
	return false
}

// CachePlan takes Statement and returns true if the query plan should be cached
func CachePlan(stmt Statement) bool {
	switch stmt.(type) {
	case *Select, *Union, *ParenSelect,
		*Insert, *Update, *Delete:
		return true
	}
	return false
}

//MustRewriteAST takes Statement and returns if the statement is set statement.
func MustRewriteAST(stmt Statement) bool {
	switch node := stmt.(type) {
	case *Set:
		return true
	case *Show:
		switch node.Internal.(type) {
		case *ShowBasic:
			return true
		}
		return false
	}
	return false
}

// Preview analyzes the beginning of the query using a simpler and faster
// textual comparison to identify the statement type.
func Preview(sql string) StatementType {
	trimmed := StripLeadingComments(sql)

	if strings.Index(trimmed, "/*!") == 0 {
		return StmtComment
	}

	isNotLetter := func(r rune) bool { return !unicode.IsLetter(r) }
	firstWord := strings.TrimLeftFunc(trimmed, isNotLetter)

	if end := strings.IndexFunc(firstWord, unicode.IsSpace); end != -1 {
		firstWord = firstWord[:end]
	}
	// Comparison is done in order of priority.
	loweredFirstWord := strings.ToLower(firstWord)
	switch loweredFirstWord {
	case "select":
		return StmtSelect
	case "stream":
		return StmtStream
	case "vstream":
		return StmtVStream
	case "insert":
		return StmtInsert
	case "replace":
		return StmtReplace
	case "update":
		return StmtUpdate
	case "delete":
		return StmtDelete
	case "savepoint":
		return StmtSavepoint
	case "lock":
		return StmtLockTables
	case "unlock":
		return StmtUnlockTables
	}
	// For the following statements it is not sufficient to rely
	// on loweredFirstWord. This is because they are not statements
	// in the grammar and we are relying on Preview to parse them.
	// For instance, we don't want: "BEGIN JUNK" to be parsed
	// as StmtBegin.
	trimmedNoComments, _ := SplitMarginComments(trimmed)
	switch strings.ToLower(trimmedNoComments) {
	case "begin", "start transaction":
		return StmtBegin
	case "commit":
		return StmtCommit
	case "rollback":
		return StmtRollback
	}
	switch loweredFirstWord {
	case "create", "alter", "rename", "drop", "truncate":
		return StmtDDL
	case "flush":
		return StmtFlush
	case "set":
		return StmtSet
	case "show":
		return StmtShow
	case "use":
		return StmtUse
	case "describe", "desc", "explain":
		return StmtExplain
	case "analyze", "repair", "optimize":
		return StmtOther
	case "grant", "revoke":
		return StmtPriv
	case "release":
		return StmtRelease
	case "rollback":
		return StmtSRollback
	}
	return StmtUnknown
}

func (s StatementType) String() string {
	switch s {
	case StmtSelect:
		return "SELECT"
	case StmtStream:
		return "STREAM"
	case StmtVStream:
		return "VSTREAM"
	case StmtInsert:
		return "INSERT"
	case StmtReplace:
		return "REPLACE"
	case StmtUpdate:
		return "UPDATE"
	case StmtDelete:
		return "DELETE"
	case StmtDDL:
		return "DDL"
	case StmtBegin:
		return "BEGIN"
	case StmtCommit:
		return "COMMIT"
	case StmtRollback:
		return "ROLLBACK"
	case StmtSet:
		return "SET"
	case StmtShow:
		return "SHOW"
	case StmtUse:
		return "USE"
	case StmtOther:
		return "OTHER"
	case StmtPriv:
		return "PRIV"
	case StmtExplain:
		return "EXPLAIN"
	case StmtSavepoint:
		return "SAVEPOINT"
	case StmtSRollback:
		return "SAVEPOINT_ROLLBACK"
	case StmtRelease:
		return "RELEASE"
	case StmtLockTables:
		return "LOCK_TABLES"
	case StmtUnlockTables:
		return "UNLOCK_TABLES"
	case StmtFlush:
		return "FLUSH"
	case StmtCallProc:
		return "CALL_PROC"
	default:
		return "UNKNOWN"
	}
}

// IsDML returns true if the query is an INSERT, UPDATE or DELETE statement.
func IsDML(sql string) bool {
	switch Preview(sql) {
	case StmtInsert, StmtReplace, StmtUpdate, StmtDelete:
		return true
	}
	return false
}

//IsDMLStatement returns true if the query is an INSERT, UPDATE or DELETE statement.
func IsDMLStatement(stmt Statement) bool {
	switch stmt.(type) {
	case *Insert, *Update, *Delete:
		return true
	}

	return false
}

// SplitAndExpression breaks up the Expr into AND-separated conditions
// and appends them to filters. Outer parenthesis are removed. Precedence
// should be taken into account if expressions are recombined.
func SplitAndExpression(filters []Expr, node Expr) []Expr {
	if node == nil {
		return filters
	}
	switch node := node.(type) {
	case *AndExpr:
		filters = SplitAndExpression(filters, node.Left)
		return SplitAndExpression(filters, node.Right)
	}
	return append(filters, node)
}

// TableFromStatement returns the qualified table name for the query.
// This works only for select statements.
func TableFromStatement(sql string) (TableName, error) {
	stmt, err := Parse(sql)
	if err != nil {
		return TableName{}, err
	}
	sel, ok := stmt.(*Select)
	if !ok {
		return TableName{}, fmt.Errorf("unrecognized statement: %s", sql)
	}
	if len(sel.From) != 1 {
		return TableName{}, fmt.Errorf("table expression is complex")
	}
	aliased, ok := sel.From[0].(*AliasedTableExpr)
	if !ok {
		return TableName{}, fmt.Errorf("table expression is complex")
	}
	tableName, ok := aliased.Expr.(TableName)
	if !ok {
		return TableName{}, fmt.Errorf("table expression is complex")
	}
	return tableName, nil
}

// GetTableName returns the table name from the SimpleTableExpr
// only if it's a simple expression. Otherwise, it returns "".
func GetTableName(node SimpleTableExpr) TableIdent {
	if n, ok := node.(TableName); ok && n.Qualifier.IsEmpty() {
		return n.Name
	}
	// sub-select or '.' expression
	return NewTableIdent("")
}

// IsColName returns true if the Expr is a *ColName.
func IsColName(node Expr) bool {
	_, ok := node.(*ColName)
	return ok
}

// IsValue returns true if the Expr is a string, integral or value arg.
// NULL is not considered to be a value.
func IsValue(node Expr) bool {
	switch v := node.(type) {
	case Argument:
		return true
	case *Literal:
		switch v.Type {
		case StrVal, HexVal, IntVal:
			return true
		}
	}
	return false
}

// IsNull returns true if the Expr is SQL NULL
func IsNull(node Expr) bool {
	switch node.(type) {
	case *NullVal:
		return true
	}
	return false
}

// IsSimpleTuple returns true if the Expr is a ValTuple that
// contains simple values or if it's a list arg.
func IsSimpleTuple(node Expr) bool {
	switch vals := node.(type) {
	case ValTuple:
		for _, n := range vals {
			if !IsValue(n) {
				return false
			}
		}
		return true
	case ListArg:
		return true
	}
	// It's a subquery
	return false
}

// NewPlanValue builds a sqltypes.PlanValue from an Expr.
func NewPlanValue(node Expr) (sqltypes.PlanValue, error) {
	switch node := node.(type) {
	case Argument:
		return sqltypes.PlanValue{Key: string(node[1:])}, nil
	case *Literal:
		switch node.Type {
		case IntVal:
			n, err := sqltypes.NewIntegral(string(node.Val))
			if err != nil {
				return sqltypes.PlanValue{}, err
			}
			return sqltypes.PlanValue{Value: n}, nil
		case FloatVal:
			return sqltypes.PlanValue{Value: sqltypes.MakeTrusted(sqltypes.Float64, node.Val)}, nil
		case StrVal:
			return sqltypes.PlanValue{Value: sqltypes.MakeTrusted(sqltypes.VarBinary, node.Val)}, nil
		case HexVal:
			v, err := node.HexDecode()
			if err != nil {
				return sqltypes.PlanValue{}, err
			}
			return sqltypes.PlanValue{Value: sqltypes.MakeTrusted(sqltypes.VarBinary, v)}, nil
		}
	case ListArg:
		return sqltypes.PlanValue{ListKey: string(node[2:])}, nil
	case ValTuple:
		pv := sqltypes.PlanValue{
			Values: make([]sqltypes.PlanValue, 0, len(node)),
		}
		for _, val := range node {
			innerpv, err := NewPlanValue(val)
			if err != nil {
				return sqltypes.PlanValue{}, err
			}
			if innerpv.ListKey != "" || innerpv.Values != nil {
				return sqltypes.PlanValue{}, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: nested lists")
			}
			pv.Values = append(pv.Values, innerpv)
		}
		return pv, nil
	case *NullVal:
		return sqltypes.PlanValue{}, nil
	case *UnaryExpr:
		switch node.Operator {
		case UBinaryOp, Utf8mb4Op, Utf8Op, Latin1Op: // for some charset introducers, we can just ignore them
			return NewPlanValue(node.Expr)
		}
	}
	return sqltypes.PlanValue{}, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "expression is too complex '%v'", String(node))
}

//IsLockingFunc returns true for all functions that are used to work with mysql advisory locks
func IsLockingFunc(node Expr) bool {
	switch p := node.(type) {
	case *FuncExpr:
		_, found := lockingFunctions[p.Name.Lowered()]
		return found
	}
	return false
}

var lockingFunctions = map[string]interface{}{
	"get_lock":          nil,
	"is_free_lock":      nil,
	"is_used_lock":      nil,
	"release_all_locks": nil,
	"release_lock":      nil,
}
<commit_msg>fix comment<commit_after>/*
Copyright 2019 The Vitess Authors.

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

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

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

package sqlparser

// analyzer.go contains utility analysis functions.

import (
	"fmt"
	"strings"
	"unicode"

	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/vterrors"

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

// StatementType encodes the type of a SQL statement
type StatementType int

// These constants are used to identify the SQL statement type.
// Changing this list will require reviewing all calls to Preview.
const (
	StmtSelect StatementType = iota
	StmtStream
	StmtInsert
	StmtReplace
	StmtUpdate
	StmtDelete
	StmtDDL
	StmtBegin
	StmtCommit
	StmtRollback
	StmtSet
	StmtShow
	StmtUse
	StmtOther
	StmtUnknown
	StmtComment
	StmtPriv
	StmtExplain
	StmtSavepoint
	StmtSRollback
	StmtRelease
	StmtVStream
	StmtLockTables
	StmtUnlockTables
	StmtFlush
	StmtCallProc
)

//ASTToStatementType returns a StatementType from an AST stmt
func ASTToStatementType(stmt Statement) StatementType {
	switch stmt.(type) {
	case *Select, *Union:
		return StmtSelect
	case *Insert:
		return StmtInsert
	case *Update:
		return StmtUpdate
	case *Delete:
		return StmtDelete
	case *Set:
		return StmtSet
	case *Show:
		return StmtShow
	case DDLStatement, DBDDLStatement, *AlterVschema:
		return StmtDDL
	case *Use:
		return StmtUse
	case *OtherRead, *OtherAdmin, *Load:
		return StmtOther
	case Explain:
		return StmtExplain
	case *Begin:
		return StmtBegin
	case *Commit:
		return StmtCommit
	case *Rollback:
		return StmtRollback
	case *Savepoint:
		return StmtSavepoint
	case *SRollback:
		return StmtSRollback
	case *Release:
		return StmtRelease
	case *LockTables:
		return StmtLockTables
	case *UnlockTables:
		return StmtUnlockTables
	case *Flush:
		return StmtFlush
	case *CallProc:
		return StmtCallProc
	default:
		return StmtUnknown
	}
}

//CanNormalize takes Statement and returns if the statement can be normalized.
func CanNormalize(stmt Statement) bool {
	switch stmt.(type) {
	case *Select, *Union, *Insert, *Update, *Delete, *Set, *CallProc: // TODO: we could merge this logic into ASTrewriter
		return true
	}
	return false
}

// CachePlan takes Statement and returns true if the query plan should be cached
func CachePlan(stmt Statement) bool {
	switch stmt.(type) {
	case *Select, *Union, *ParenSelect,
		*Insert, *Update, *Delete:
		return true
	}
	return false
}

//MustRewriteAST takes Statement and returns true if RewriteAST must run on it for correct execution irrespective of user flags.
func MustRewriteAST(stmt Statement) bool {
	switch node := stmt.(type) {
	case *Set:
		return true
	case *Show:
		switch node.Internal.(type) {
		case *ShowBasic:
			return true
		}
		return false
	}
	return false
}

// Preview analyzes the beginning of the query using a simpler and faster
// textual comparison to identify the statement type.
func Preview(sql string) StatementType {
	trimmed := StripLeadingComments(sql)

	if strings.Index(trimmed, "/*!") == 0 {
		return StmtComment
	}

	isNotLetter := func(r rune) bool { return !unicode.IsLetter(r) }
	firstWord := strings.TrimLeftFunc(trimmed, isNotLetter)

	if end := strings.IndexFunc(firstWord, unicode.IsSpace); end != -1 {
		firstWord = firstWord[:end]
	}
	// Comparison is done in order of priority.
	loweredFirstWord := strings.ToLower(firstWord)
	switch loweredFirstWord {
	case "select":
		return StmtSelect
	case "stream":
		return StmtStream
	case "vstream":
		return StmtVStream
	case "insert":
		return StmtInsert
	case "replace":
		return StmtReplace
	case "update":
		return StmtUpdate
	case "delete":
		return StmtDelete
	case "savepoint":
		return StmtSavepoint
	case "lock":
		return StmtLockTables
	case "unlock":
		return StmtUnlockTables
	}
	// For the following statements it is not sufficient to rely
	// on loweredFirstWord. This is because they are not statements
	// in the grammar and we are relying on Preview to parse them.
	// For instance, we don't want: "BEGIN JUNK" to be parsed
	// as StmtBegin.
	trimmedNoComments, _ := SplitMarginComments(trimmed)
	switch strings.ToLower(trimmedNoComments) {
	case "begin", "start transaction":
		return StmtBegin
	case "commit":
		return StmtCommit
	case "rollback":
		return StmtRollback
	}
	switch loweredFirstWord {
	case "create", "alter", "rename", "drop", "truncate":
		return StmtDDL
	case "flush":
		return StmtFlush
	case "set":
		return StmtSet
	case "show":
		return StmtShow
	case "use":
		return StmtUse
	case "describe", "desc", "explain":
		return StmtExplain
	case "analyze", "repair", "optimize":
		return StmtOther
	case "grant", "revoke":
		return StmtPriv
	case "release":
		return StmtRelease
	case "rollback":
		return StmtSRollback
	}
	return StmtUnknown
}

func (s StatementType) String() string {
	switch s {
	case StmtSelect:
		return "SELECT"
	case StmtStream:
		return "STREAM"
	case StmtVStream:
		return "VSTREAM"
	case StmtInsert:
		return "INSERT"
	case StmtReplace:
		return "REPLACE"
	case StmtUpdate:
		return "UPDATE"
	case StmtDelete:
		return "DELETE"
	case StmtDDL:
		return "DDL"
	case StmtBegin:
		return "BEGIN"
	case StmtCommit:
		return "COMMIT"
	case StmtRollback:
		return "ROLLBACK"
	case StmtSet:
		return "SET"
	case StmtShow:
		return "SHOW"
	case StmtUse:
		return "USE"
	case StmtOther:
		return "OTHER"
	case StmtPriv:
		return "PRIV"
	case StmtExplain:
		return "EXPLAIN"
	case StmtSavepoint:
		return "SAVEPOINT"
	case StmtSRollback:
		return "SAVEPOINT_ROLLBACK"
	case StmtRelease:
		return "RELEASE"
	case StmtLockTables:
		return "LOCK_TABLES"
	case StmtUnlockTables:
		return "UNLOCK_TABLES"
	case StmtFlush:
		return "FLUSH"
	case StmtCallProc:
		return "CALL_PROC"
	default:
		return "UNKNOWN"
	}
}

// IsDML returns true if the query is an INSERT, UPDATE or DELETE statement.
func IsDML(sql string) bool {
	switch Preview(sql) {
	case StmtInsert, StmtReplace, StmtUpdate, StmtDelete:
		return true
	}
	return false
}

//IsDMLStatement returns true if the query is an INSERT, UPDATE or DELETE statement.
func IsDMLStatement(stmt Statement) bool {
	switch stmt.(type) {
	case *Insert, *Update, *Delete:
		return true
	}

	return false
}

// SplitAndExpression breaks up the Expr into AND-separated conditions
// and appends them to filters. Outer parenthesis are removed. Precedence
// should be taken into account if expressions are recombined.
func SplitAndExpression(filters []Expr, node Expr) []Expr {
	if node == nil {
		return filters
	}
	switch node := node.(type) {
	case *AndExpr:
		filters = SplitAndExpression(filters, node.Left)
		return SplitAndExpression(filters, node.Right)
	}
	return append(filters, node)
}

// TableFromStatement returns the qualified table name for the query.
// This works only for select statements.
func TableFromStatement(sql string) (TableName, error) {
	stmt, err := Parse(sql)
	if err != nil {
		return TableName{}, err
	}
	sel, ok := stmt.(*Select)
	if !ok {
		return TableName{}, fmt.Errorf("unrecognized statement: %s", sql)
	}
	if len(sel.From) != 1 {
		return TableName{}, fmt.Errorf("table expression is complex")
	}
	aliased, ok := sel.From[0].(*AliasedTableExpr)
	if !ok {
		return TableName{}, fmt.Errorf("table expression is complex")
	}
	tableName, ok := aliased.Expr.(TableName)
	if !ok {
		return TableName{}, fmt.Errorf("table expression is complex")
	}
	return tableName, nil
}

// GetTableName returns the table name from the SimpleTableExpr
// only if it's a simple expression. Otherwise, it returns "".
func GetTableName(node SimpleTableExpr) TableIdent {
	if n, ok := node.(TableName); ok && n.Qualifier.IsEmpty() {
		return n.Name
	}
	// sub-select or '.' expression
	return NewTableIdent("")
}

// IsColName returns true if the Expr is a *ColName.
func IsColName(node Expr) bool {
	_, ok := node.(*ColName)
	return ok
}

// IsValue returns true if the Expr is a string, integral or value arg.
// NULL is not considered to be a value.
func IsValue(node Expr) bool {
	switch v := node.(type) {
	case Argument:
		return true
	case *Literal:
		switch v.Type {
		case StrVal, HexVal, IntVal:
			return true
		}
	}
	return false
}

// IsNull returns true if the Expr is SQL NULL
func IsNull(node Expr) bool {
	switch node.(type) {
	case *NullVal:
		return true
	}
	return false
}

// IsSimpleTuple returns true if the Expr is a ValTuple that
// contains simple values or if it's a list arg.
func IsSimpleTuple(node Expr) bool {
	switch vals := node.(type) {
	case ValTuple:
		for _, n := range vals {
			if !IsValue(n) {
				return false
			}
		}
		return true
	case ListArg:
		return true
	}
	// It's a subquery
	return false
}

// NewPlanValue builds a sqltypes.PlanValue from an Expr.
func NewPlanValue(node Expr) (sqltypes.PlanValue, error) {
	switch node := node.(type) {
	case Argument:
		return sqltypes.PlanValue{Key: string(node[1:])}, nil
	case *Literal:
		switch node.Type {
		case IntVal:
			n, err := sqltypes.NewIntegral(string(node.Val))
			if err != nil {
				return sqltypes.PlanValue{}, err
			}
			return sqltypes.PlanValue{Value: n}, nil
		case FloatVal:
			return sqltypes.PlanValue{Value: sqltypes.MakeTrusted(sqltypes.Float64, node.Val)}, nil
		case StrVal:
			return sqltypes.PlanValue{Value: sqltypes.MakeTrusted(sqltypes.VarBinary, node.Val)}, nil
		case HexVal:
			v, err := node.HexDecode()
			if err != nil {
				return sqltypes.PlanValue{}, err
			}
			return sqltypes.PlanValue{Value: sqltypes.MakeTrusted(sqltypes.VarBinary, v)}, nil
		}
	case ListArg:
		return sqltypes.PlanValue{ListKey: string(node[2:])}, nil
	case ValTuple:
		pv := sqltypes.PlanValue{
			Values: make([]sqltypes.PlanValue, 0, len(node)),
		}
		for _, val := range node {
			innerpv, err := NewPlanValue(val)
			if err != nil {
				return sqltypes.PlanValue{}, err
			}
			if innerpv.ListKey != "" || innerpv.Values != nil {
				return sqltypes.PlanValue{}, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: nested lists")
			}
			pv.Values = append(pv.Values, innerpv)
		}
		return pv, nil
	case *NullVal:
		return sqltypes.PlanValue{}, nil
	case *UnaryExpr:
		switch node.Operator {
		case UBinaryOp, Utf8mb4Op, Utf8Op, Latin1Op: // for some charset introducers, we can just ignore them
			return NewPlanValue(node.Expr)
		}
	}
	return sqltypes.PlanValue{}, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "expression is too complex '%v'", String(node))
}

//IsLockingFunc returns true for all functions that are used to work with mysql advisory locks
func IsLockingFunc(node Expr) bool {
	switch p := node.(type) {
	case *FuncExpr:
		_, found := lockingFunctions[p.Name.Lowered()]
		return found
	}
	return false
}

var lockingFunctions = map[string]interface{}{
	"get_lock":          nil,
	"is_free_lock":      nil,
	"is_used_lock":      nil,
	"release_all_locks": nil,
	"release_lock":      nil,
}
<|endoftext|>


<commit_before>package memo

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"sync"
	"testing"
	"time"
)

// BPS is emulation download speed for url
const BPS = 10

type slowReader struct {
	delay time.Duration
	r     io.Reader
}

func (sr slowReader) Read(p []byte) (int, error) {
	time.Sleep(sr.delay)
	return sr.r.Read(p[:1])
}

func newReader(r io.Reader, bps int) io.Reader {
	delay := time.Second / time.Duration(bps)
	return slowReader{r: r, delay: delay}
}

// httpGetBodyMock emulates time-bound read functions.
func httpGetBodyMock(str string, done <-chan struct{}) (interface{}, error) {
	s := strings.NewReader(str)
	r := newReader(s, BPS)
	return ioutil.ReadAll(r)
}

func httpGetBody(url string, done <-chan struct{}) (interface{}, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return ioutil.ReadAll(resp.Body)
}

var HTTPGetBody = httpGetBodyMock

func incomingURLs() <-chan string {
	ch := make(chan string)
	go func() {
		for _, url := range []string{
			"https://golang.org",
			"https://godoc.org",
			"https://play.golang.org",
			"http://gopl.io",
			"https://golang.org",
			"https://godoc.org",
			"https://play.golang.org",
			"http://gopl.io",
		} {
			ch <- url
		}
		close(ch)
	}()
	return ch
}

type M interface {
	Get(key string, done <-chan struct{}) (interface{}, error)
}

func Sequential(t *testing.T, m M) {
	for url := range incomingURLs() {
		start := time.Now()
		value, err := m.Get(url, nil)
		if err != nil {
			log.Print(err)
			continue
		}
		fmt.Printf("%s, %s, %d bytes\n",
			url, time.Since(start), len(value.([]byte)))
	}
}

func Concurrent(t *testing.T, m M) {
	var n sync.WaitGroup
	for url := range incomingURLs() {
		n.Add(1)
		go func(url string) {
			defer n.Done()
			start := time.Now()
			value, err := m.Get(url, nil)
			if err != nil {
				log.Print(err)
				return
			}
			fmt.Printf("%s, %s, %d bytes\n",
				url, time.Since(start), len(value.([]byte)))
		}(url)
	}
	n.Wait()
}

func TestSequential(t *testing.T) {
	m := New(httpGetBodyMock)
	Sequential(t, m)
}

func TestConcurrent(t *testing.T) {
	m := New(httpGetBodyMock)
	Concurrent(t, m)
}
<commit_msg>[9.3] add memo.Close to tests<commit_after>package memo

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"sync"
	"testing"
	"time"
)

// BPS is emulation download speed for url
const BPS = 10

type slowReader struct {
	delay time.Duration
	r     io.Reader
}

func (sr slowReader) Read(p []byte) (int, error) {
	time.Sleep(sr.delay)
	return sr.r.Read(p[:1])
}

func newReader(r io.Reader, bps int) io.Reader {
	delay := time.Second / time.Duration(bps)
	return slowReader{r: r, delay: delay}
}

// httpGetBodyMock emulates time-bound read functions.
func httpGetBodyMock(str string, done <-chan struct{}) (interface{}, error) {
	s := strings.NewReader(str)
	r := newReader(s, BPS)
	return ioutil.ReadAll(r)
}

func httpGetBody(url string, done <-chan struct{}) (interface{}, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return ioutil.ReadAll(resp.Body)
}

var HTTPGetBody = httpGetBodyMock

func incomingURLs() <-chan string {
	ch := make(chan string)
	go func() {
		for _, url := range []string{
			"https://golang.org",
			"https://godoc.org",
			"https://play.golang.org",
			"http://gopl.io",
			"https://golang.org",
			"https://godoc.org",
			"https://play.golang.org",
			"http://gopl.io",
		} {
			ch <- url
		}
		close(ch)
	}()
	return ch
}

type M interface {
	Get(key string, done <-chan struct{}) (interface{}, error)
}

func Sequential(t *testing.T, m M) {
	for url := range incomingURLs() {
		start := time.Now()
		value, err := m.Get(url, nil)
		if err != nil {
			log.Print(err)
			continue
		}
		fmt.Printf("%s, %s, %d bytes\n",
			url, time.Since(start), len(value.([]byte)))
	}
}

func Concurrent(t *testing.T, m M) {
	var n sync.WaitGroup
	for url := range incomingURLs() {
		n.Add(1)
		go func(url string) {
			defer n.Done()
			start := time.Now()
			value, err := m.Get(url, nil)
			if err != nil {
				log.Print(err)
				return
			}
			fmt.Printf("%s, %s, %d bytes\n",
				url, time.Since(start), len(value.([]byte)))
		}(url)
	}
	n.Wait()
}

func TestSequential(t *testing.T) {
	m := New(httpGetBodyMock)
	defer m.Close()
	Sequential(t, m)
}

func TestConcurrent(t *testing.T) {
	m := New(httpGetBodyMock)
	defer m.Close()
	Concurrent(t, m)
}
<|endoftext|>


<commit_before>// +build functional

package cri_containerd

import (
	"context"
	"testing"

	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/runtime/v1alpha2"
)

func Test_SandboxStats_Single_LCOW(t *testing.T) {
	pullRequiredLcowImages(t, []string{imageLcowK8sPause})

	request := &runtime.RunPodSandboxRequest{
		Config: &runtime.PodSandboxConfig{
			Metadata: &runtime.PodSandboxMetadata{
				Name:      t.Name(),
				Namespace: testNamespace,
			},
		},
		RuntimeHandler: lcowRuntimeHandler,
	}

	client := newTestRuntimeClient(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	podID := runPodSandbox(t, client, ctx, request)
	defer stopPodSandbox(t, client, ctx, podID)
	defer removePodSandbox(t, client, ctx, podID)

	statsRequest := &runtime.ContainerStatsRequest{
		ContainerId: podID,
	}

	stats, err := client.ContainerStats(ctx, statsRequest)
	if err != nil {
		t.Fatal(err)
	}

	stat := stats.Stats
	if stat.Cpu.Timestamp == 0 {
		t.Fatalf("expected cpu stat timestamp != 0 but got %d", stat.Cpu.Timestamp)
	}
	if stat.Cpu.UsageCoreNanoSeconds.Value == 0 {
		t.Fatalf("expected cpu usage != 0 but got %d", stat.Cpu.UsageCoreNanoSeconds.Value)
	}
	if stat.Memory.Timestamp == 0 {
		t.Fatalf("expected memory stat timestamp != 0 but got %d", stat.Memory.Timestamp)
	}
	if stat.Memory.WorkingSetBytes.Value == 0 {
		t.Fatalf("expected memory usage != 0 but got %d", stat.Memory.WorkingSetBytes.Value)
	}
}

func Test_SandboxStats_List_ContainerID_LCOW(t *testing.T) {
	pullRequiredLcowImages(t, []string{imageLcowK8sPause})

	request := &runtime.RunPodSandboxRequest{
		Config: &runtime.PodSandboxConfig{
			Metadata: &runtime.PodSandboxMetadata{
				Name:      t.Name(),
				Namespace: testNamespace,
			},
		},
		RuntimeHandler: lcowRuntimeHandler,
	}

	client := newTestRuntimeClient(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	podID := runPodSandbox(t, client, ctx, request)
	defer stopPodSandbox(t, client, ctx, podID)
	defer removePodSandbox(t, client, ctx, podID)

	statsRequest := &runtime.ListContainerStatsRequest{
		Filter: &runtime.ContainerStatsFilter{
			Id: podID,
		},
	}

	stats, err := client.ListContainerStats(ctx, statsRequest)
	if err != nil {
		t.Fatal(err)
	}

	if len(stats.Stats) != 1 {
		t.Fatalf("expected 1 stats result but got %d", len(stats.Stats))
	}
	stat := stats.Stats[0]
	if stat.Cpu.Timestamp == 0 {
		t.Fatalf("expected cpu stat timestamp != 0 but got %d", stat.Cpu.Timestamp)
	}
	if stat.Cpu.UsageCoreNanoSeconds.Value == 0 {
		t.Fatalf("expected cpu usage != 0 but got %d", stat.Cpu.UsageCoreNanoSeconds.Value)
	}
	if stat.Memory.Timestamp == 0 {
		t.Fatalf("expected memory stat timestamp != 0 but got %d", stat.Memory.Timestamp)
	}
	if stat.Memory.WorkingSetBytes.Value == 0 {
		t.Fatalf("expected memory usage != 0 but got %d", stat.Memory.WorkingSetBytes.Value)
	}
}

func Test_SandboxStats_List_PodID_LCOW(t *testing.T) {
	pullRequiredLcowImages(t, []string{imageLcowK8sPause})

	request := &runtime.RunPodSandboxRequest{
		Config: &runtime.PodSandboxConfig{
			Metadata: &runtime.PodSandboxMetadata{
				Name:      t.Name(),
				Namespace: testNamespace,
			},
		},
		RuntimeHandler: lcowRuntimeHandler,
	}

	client := newTestRuntimeClient(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	podID := runPodSandbox(t, client, ctx, request)
	defer stopPodSandbox(t, client, ctx, podID)
	defer removePodSandbox(t, client, ctx, podID)

	statsRequest := &runtime.ListContainerStatsRequest{
		Filter: &runtime.ContainerStatsFilter{
			PodSandboxId: podID,
		},
	}

	stats, err := client.ListContainerStats(ctx, statsRequest)
	if err != nil {
		t.Fatal(err)
	}

	if len(stats.Stats) != 1 {
		t.Fatalf("expected 1 stats result but got %d", len(stats.Stats))
	}
	stat := stats.Stats[0]
	if stat.Cpu.Timestamp == 0 {
		t.Fatalf("expected cpu stat timestamp != 0 but got %d", stat.Cpu.Timestamp)
	}
	if stat.Cpu.UsageCoreNanoSeconds.Value == 0 {
		t.Fatalf("expected cpu usage != 0 but got %d", stat.Cpu.UsageCoreNanoSeconds.Value)
	}
	if stat.Memory.Timestamp == 0 {
		t.Fatalf("expected memory stat timestamp != 0 but got %d", stat.Memory.Timestamp)
	}
	if stat.Memory.WorkingSetBytes.Value == 0 {
		t.Fatalf("expected memory usage != 0 but got %d", stat.Memory.WorkingSetBytes.Value)
	}
}
<commit_msg>Fix test bug in defer order<commit_after>// +build functional

package cri_containerd

import (
	"context"
	"testing"

	runtime "k8s.io/kubernetes/pkg/kubelet/apis/cri/runtime/v1alpha2"
)

func Test_SandboxStats_Single_LCOW(t *testing.T) {
	pullRequiredLcowImages(t, []string{imageLcowK8sPause})

	request := &runtime.RunPodSandboxRequest{
		Config: &runtime.PodSandboxConfig{
			Metadata: &runtime.PodSandboxMetadata{
				Name:      t.Name(),
				Namespace: testNamespace,
			},
		},
		RuntimeHandler: lcowRuntimeHandler,
	}

	client := newTestRuntimeClient(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	podID := runPodSandbox(t, client, ctx, request)
	defer removePodSandbox(t, client, ctx, podID)
	defer stopPodSandbox(t, client, ctx, podID)

	statsRequest := &runtime.ContainerStatsRequest{
		ContainerId: podID,
	}

	stats, err := client.ContainerStats(ctx, statsRequest)
	if err != nil {
		t.Fatal(err)
	}

	stat := stats.Stats
	if stat.Cpu.Timestamp == 0 {
		t.Fatalf("expected cpu stat timestamp != 0 but got %d", stat.Cpu.Timestamp)
	}
	if stat.Cpu.UsageCoreNanoSeconds.Value == 0 {
		t.Fatalf("expected cpu usage != 0 but got %d", stat.Cpu.UsageCoreNanoSeconds.Value)
	}
	if stat.Memory.Timestamp == 0 {
		t.Fatalf("expected memory stat timestamp != 0 but got %d", stat.Memory.Timestamp)
	}
	if stat.Memory.WorkingSetBytes.Value == 0 {
		t.Fatalf("expected memory usage != 0 but got %d", stat.Memory.WorkingSetBytes.Value)
	}
}

func Test_SandboxStats_List_ContainerID_LCOW(t *testing.T) {
	pullRequiredLcowImages(t, []string{imageLcowK8sPause})

	request := &runtime.RunPodSandboxRequest{
		Config: &runtime.PodSandboxConfig{
			Metadata: &runtime.PodSandboxMetadata{
				Name:      t.Name(),
				Namespace: testNamespace,
			},
		},
		RuntimeHandler: lcowRuntimeHandler,
	}

	client := newTestRuntimeClient(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	podID := runPodSandbox(t, client, ctx, request)
	defer removePodSandbox(t, client, ctx, podID)
	defer stopPodSandbox(t, client, ctx, podID)

	statsRequest := &runtime.ListContainerStatsRequest{
		Filter: &runtime.ContainerStatsFilter{
			Id: podID,
		},
	}

	stats, err := client.ListContainerStats(ctx, statsRequest)
	if err != nil {
		t.Fatal(err)
	}

	if len(stats.Stats) != 1 {
		t.Fatalf("expected 1 stats result but got %d", len(stats.Stats))
	}
	stat := stats.Stats[0]
	if stat.Cpu.Timestamp == 0 {
		t.Fatalf("expected cpu stat timestamp != 0 but got %d", stat.Cpu.Timestamp)
	}
	if stat.Cpu.UsageCoreNanoSeconds.Value == 0 {
		t.Fatalf("expected cpu usage != 0 but got %d", stat.Cpu.UsageCoreNanoSeconds.Value)
	}
	if stat.Memory.Timestamp == 0 {
		t.Fatalf("expected memory stat timestamp != 0 but got %d", stat.Memory.Timestamp)
	}
	if stat.Memory.WorkingSetBytes.Value == 0 {
		t.Fatalf("expected memory usage != 0 but got %d", stat.Memory.WorkingSetBytes.Value)
	}
}

func Test_SandboxStats_List_PodID_LCOW(t *testing.T) {
	pullRequiredLcowImages(t, []string{imageLcowK8sPause})

	request := &runtime.RunPodSandboxRequest{
		Config: &runtime.PodSandboxConfig{
			Metadata: &runtime.PodSandboxMetadata{
				Name:      t.Name(),
				Namespace: testNamespace,
			},
		},
		RuntimeHandler: lcowRuntimeHandler,
	}

	client := newTestRuntimeClient(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	podID := runPodSandbox(t, client, ctx, request)
	defer removePodSandbox(t, client, ctx, podID)
	defer stopPodSandbox(t, client, ctx, podID)

	statsRequest := &runtime.ListContainerStatsRequest{
		Filter: &runtime.ContainerStatsFilter{
			PodSandboxId: podID,
		},
	}

	stats, err := client.ListContainerStats(ctx, statsRequest)
	if err != nil {
		t.Fatal(err)
	}

	if len(stats.Stats) != 1 {
		t.Fatalf("expected 1 stats result but got %d", len(stats.Stats))
	}
	stat := stats.Stats[0]
	if stat.Cpu.Timestamp == 0 {
		t.Fatalf("expected cpu stat timestamp != 0 but got %d", stat.Cpu.Timestamp)
	}
	if stat.Cpu.UsageCoreNanoSeconds.Value == 0 {
		t.Fatalf("expected cpu usage != 0 but got %d", stat.Cpu.UsageCoreNanoSeconds.Value)
	}
	if stat.Memory.Timestamp == 0 {
		t.Fatalf("expected memory stat timestamp != 0 but got %d", stat.Memory.Timestamp)
	}
	if stat.Memory.WorkingSetBytes.Value == 0 {
		t.Fatalf("expected memory usage != 0 but got %d", stat.Memory.WorkingSetBytes.Value)
	}
}
<|endoftext|>


<commit_before>package condor

import (
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"tes/config"
	pbe "tes/ga4gh"
	"tes/logger"
	sched "tes/scheduler"
	pbr "tes/server/proto"
	"text/template"
)

var log = logger.New("condor")

// prefix is a string prefixed to condor worker IDs, so that condor
// workers can be identified by ShouldStartWorker() below.
// TODO move to worker metadata to be consistent with GCE
const prefix = "condor-"

// NewScheduler returns a new HTCondor Scheduler instance.
func NewScheduler(conf config.Config) sched.Scheduler {
	return &scheduler{conf}
}

type scheduler struct {
	conf config.Config
}

// Schedule schedules a job on the HTCondor queue and returns a corresponding Offer.
func (s *scheduler) Schedule(j *pbe.Job) *sched.Offer {
	log.Debug("Running condor scheduler")

	// TODO could we call condor_submit --dry-run to test if a job would succeed?
	w := &pbr.Worker{
		Id: prefix + sched.GenWorkerID(),
	}
	return sched.NewOffer(w, j, sched.Scores{})
}

func (s *scheduler) ShouldStartWorker(w *pbr.Worker) bool {
	return strings.HasPrefix(w.Id, prefix) &&
		w.State == pbr.WorkerState_Uninitialized
}

// StartWorker submits a job via "condor_submit" to start a new worker.
func (s *scheduler) StartWorker(w *pbr.Worker) error {
	log.Debug("Starting condor worker")

	// TODO document that these working dirs need manual cleanup
	workdir := path.Join(s.conf.WorkDir, w.Id)
	workdir, _ = filepath.Abs(workdir)
	os.MkdirAll(workdir, 0755)

	c := s.conf.Worker
	c.ID = w.Id
	c.Timeout = 0
	c.Storage = s.conf.Storage

	confPath := path.Join(workdir, "worker.conf.yml")
	c.ToYamlFile(confPath)

	workerPath := sched.DetectWorkerPath()

	submitPath := path.Join(workdir, "condor.submit")
	f, _ := os.Create(submitPath)

	submitTpl, _ := template.New("condor.submit").Parse(`
		universe    = vanilla
		executable  = {{.Executable}}
		arguments   = -config worker.conf.yml
		environment = "PATH=/usr/bin"
		log         = {{.WorkDir}}/condor-event-log
		error       = {{.WorkDir}}/tes-worker-stderr
		output      = {{.WorkDir}}/tes-worker-stdout
    input       = {{.Config}}
    should_transfer_files   = YES
    when_to_transfer_output = ON_EXIT
		queue
	`)
	submitTpl.Execute(f, map[string]string{
		"Executable": workerPath,
		"WorkDir":    workdir,
		"Config":     confPath,
	})
	f.Close()

	cmd := exec.Command("condor_submit")
	stdin, _ := os.Open(submitPath)
	cmd.Stdin = stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Run()

	// TODO better error checking
	return nil
}
<commit_msg>pass requirements to condor worker<commit_after>package condor

import (
	"fmt"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"tes/config"
	pbe "tes/ga4gh"
	"tes/logger"
	sched "tes/scheduler"
	pbr "tes/server/proto"
	"text/template"
)

var log = logger.New("condor")

// prefix is a string prefixed to condor worker IDs, so that condor
// workers can be identified by ShouldStartWorker() below.
// TODO move to worker metadata to be consistent with GCE
const prefix = "condor-"

// NewScheduler returns a new HTCondor Scheduler instance.
func NewScheduler(conf config.Config) sched.Scheduler {
	return &scheduler{conf}
}

type scheduler struct {
	conf config.Config
}

// Schedule schedules a job on the HTCondor queue and returns a corresponding Offer.
func (s *scheduler) Schedule(j *pbe.Job) *sched.Offer {
	log.Debug("Running condor scheduler")

	// TODO could we call condor_submit --dry-run to test if a job would succeed?
	w := &pbr.Worker{
		Id: prefix + sched.GenWorkerID(),
		Resources: &pbr.Resources{
			Cpus: j.Task.Resources.MinimumCpuCores,
			Ram: j.Task.Resources.MinimumRamGb,
		},
	}
	return sched.NewOffer(w, j, sched.Scores{})
}

func (s *scheduler) ShouldStartWorker(w *pbr.Worker) bool {
	return strings.HasPrefix(w.Id, prefix) &&
		w.State == pbr.WorkerState_Uninitialized
}

// StartWorker submits a job via "condor_submit" to start a new worker.
func (s *scheduler) StartWorker(w *pbr.Worker) error {
	log.Debug("Starting condor worker")

	// TODO document that these working dirs need manual cleanup
	workdir := path.Join(s.conf.WorkDir, w.Id)
	workdir, _ = filepath.Abs(workdir)
	os.MkdirAll(workdir, 0755)

	c := s.conf.Worker
	c.ID = w.Id
	c.Timeout = 0
	c.Storage = s.conf.Storage

	confPath := path.Join(workdir, "worker.conf.yml")
	c.ToYamlFile(confPath)

	workerPath := sched.DetectWorkerPath()

	submitPath := path.Join(workdir, "condor.submit")
	f, _ := os.Create(submitPath)

	submitTpl, _ := template.New("condor.submit").Parse(`
		universe       = vanilla
		executable     = {{.Executable}}
		arguments      = -config worker.conf.yml
		environment    = "PATH=/usr/bin"
		log            = {{.WorkDir}}/condor-event-log
		error          = {{.WorkDir}}/tes-worker-stderr
		output         = {{.WorkDir}}/tes-worker-stdout
    input          = {{.Config}}
    request_cpus   = {{.CPU}}
    request_memory = {{.RAM}}
    should_transfer_files   = YES
    when_to_transfer_output = ON_EXIT
		queue
	`)
	submitTpl.Execute(f, map[string]string{
		"Executable": workerPath,
		"WorkDir":    workdir,
		"Config":     confPath,
		"request_cpus": fmt.Sprintf("%s", w.Resources.Cpus),
		"request_memory": fmt.Sprintf("%s GB", w.Resources.Ram),
	})
	f.Close()

	cmd := exec.Command("condor_submit")
	stdin, _ := os.Open(submitPath)
	cmd.Stdin = stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Run()

	// TODO better error checking
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2013 Mikio Hara. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE.

package ipaddr

import (
	"encoding/binary"
	"math"
	"net"
)

func npop64(bs uint64) byte {
	bs = bs&0x5555555555555555 + bs>>1&0x5555555555555555
	bs = bs&0x3333333333333333 + bs>>2&0x3333333333333333
	bs = bs&0x0f0f0f0f0f0f0f0f + bs>>4&0x0f0f0f0f0f0f0f0f
	bs = bs&0x00ff00ff00ff00ff + bs>>8&0x00ff00ff00ff00ff
	bs = bs&0x0000ffff0000ffff + bs>>16&0x0000ffff0000ffff
	return byte(bs&0x00000000ffffffff + bs>>32&0x00000000ffffffff)
}

func npop32(bs uint32) byte {
	bs = bs&0x55555555 + bs>>1&0x55555555
	bs = bs&0x33333333 + bs>>2&0x33333333
	bs = bs&0x0f0f0f0f + bs>>4&0x0f0f0f0f
	bs = bs&0x00ff00ff + bs>>8&0x00ff00ff
	return byte(bs&0x0000ffff + bs>>16&0x0000ffff)
}

func nlz64(bs uint64) byte {
	bs |= bs >> 1
	bs |= bs >> 2
	bs |= bs >> 4
	bs |= bs >> 8
	bs |= bs >> 16
	bs |= bs >> 32
	return npop64(^bs)
}

func nlz32(bs uint32) byte {
	bs |= bs >> 1
	bs |= bs >> 2
	bs |= bs >> 4
	bs |= bs >> 8
	bs |= bs >> 16
	return npop32(^bs)
}

func mask64(nbits byte) uint64 {
	return -uint64(1 << uint(64-nbits))
}

func mask32(nbits byte) uint32 {
	return -uint32(1 << uint(32-nbits))
}

type ipv4Int uint32

func (i ipv4Int) IP() net.IP {
	var b [net.IPv4len]byte
	binary.BigEndian.PutUint32(b[:], uint32(i))
	return net.IP(b[:])
}

func (i ipv4Int) encodeNLRI(b []byte, nbits byte) int {
	l := int(nbits) / 8
	if nbits%8 > 0 {
		l++
	}
	b[0] = nbits
	l++
	b = b[1:]
	for n := 0; n < l; n++ {
		b[n] = byte(i >> uint(32-8*(n+1)))
	}
	return l
}

func (p *IPv4) decodeNLRI(b []byte) error {
	p.nbits = b[0]
	l := len(b) - 1
	b = b[1:]
	for n := 0; n < l; n++ {
		p.addr |= ipv4Int(b[n]) << uint(32-8*(n+1))
	}
	return nil
}

func ipToIPv4Int(ip []byte) ipv4Int {
	return ipv4Int(binary.BigEndian.Uint32(ip))
}

type ipv6Int [2]uint64

func (i *ipv6Int) lshift(nbits byte) {
	i[0] = i[0]<<uint(nbits) | i[1]>>uint(64-nbits) | i[1]<<uint(nbits-64)
	i[1] = i[1] << uint(nbits)
}

func (i *ipv6Int) rshift(nbits byte) {
	i[1] = i[1]>>uint(nbits) | i[0]<<uint(64-nbits) | i[0]>>uint(nbits-64)
	i[0] = i[0] >> uint(nbits)
}

func (i *ipv6Int) incr() {
	if i[1] == math.MaxUint64 {
		i[0]++
		i[1] = 0
	} else {
		i[1]++
	}
}

func (i *ipv6Int) setHostmask(nbits byte) {
	if nbits > 64 {
		i[0], i[1] = ^mask64(64), ^mask64(nbits-64)
	} else {
		i[0], i[1] = ^mask64(nbits), mask64(64)
	}
}

func (i *ipv6Int) setNetmask(nbits byte) {
	if nbits > 64 {
		i[0], i[1] = mask64(64), mask64(nbits-64)
	} else {
		i[0], i[1] = mask64(nbits), 0
	}
}

func (i *ipv6Int) IP() net.IP {
	var b [net.IPv6len]byte
	binary.BigEndian.PutUint64(b[:8], i[0])
	binary.BigEndian.PutUint64(b[8:16], i[1])
	return net.IP(b[:])
}

func (i *ipv6Int) encodeNLRI(b []byte, nbits byte) int {
	l := int(nbits) / 8
	if nbits%8 > 0 {
		l++
	}
	b[0] = nbits
	l++
	b = b[1:]
	for n := 0; n < l; n++ {
		switch {
		case n < 8:
			b[n] = byte(i[0] >> uint(64-(8*(n+1))))
		case n >= 8 && n < 16:
			b[n] = byte(i[1] >> uint(128-(8*(n+1))))
		}
	}
	return l
}

func (p *IPv6) decodeNLRI(b []byte) error {
	p.nbits = b[0]
	l := len(b) - 1
	b = b[1:]
	for n := 0; n < l; n++ {
		switch {
		case n < 8:
			p.addr[0] |= uint64(b[n]) << uint(64-8*(n+1))
		case n >= 8 && n < 16:
			p.addr[1] |= uint64(b[n]) << uint(128-8*(n+1))
		}
	}
	return nil
}

func ipToIPv6Int(ip []byte) ipv6Int {
	return ipv6Int{binary.BigEndian.Uint64(ip[:8]), binary.BigEndian.Uint64(ip[8:16])}
}
<commit_msg>ipaddr: stylistic change<commit_after>// Copyright 2013 Mikio Hara. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE.

package ipaddr

import (
	"encoding/binary"
	"math"
	"net"
)

func npop64(bs uint64) byte {
	bs = bs&0x5555555555555555 + bs>>1&0x5555555555555555
	bs = bs&0x3333333333333333 + bs>>2&0x3333333333333333
	bs = bs&0x0f0f0f0f0f0f0f0f + bs>>4&0x0f0f0f0f0f0f0f0f
	bs = bs&0x00ff00ff00ff00ff + bs>>8&0x00ff00ff00ff00ff
	bs = bs&0x0000ffff0000ffff + bs>>16&0x0000ffff0000ffff
	return byte(bs&0x00000000ffffffff + bs>>32&0x00000000ffffffff)
}

func npop32(bs uint32) byte {
	bs = bs&0x55555555 + bs>>1&0x55555555
	bs = bs&0x33333333 + bs>>2&0x33333333
	bs = bs&0x0f0f0f0f + bs>>4&0x0f0f0f0f
	bs = bs&0x00ff00ff + bs>>8&0x00ff00ff
	return byte(bs&0x0000ffff + bs>>16&0x0000ffff)
}

func nlz64(bs uint64) byte {
	bs |= bs >> 1
	bs |= bs >> 2
	bs |= bs >> 4
	bs |= bs >> 8
	bs |= bs >> 16
	bs |= bs >> 32
	return npop64(^bs)
}

func nlz32(bs uint32) byte {
	bs |= bs >> 1
	bs |= bs >> 2
	bs |= bs >> 4
	bs |= bs >> 8
	bs |= bs >> 16
	return npop32(^bs)
}

func mask64(nbits byte) uint64 {
	return -uint64(1 << uint(64-nbits))
}

func mask32(nbits byte) uint32 {
	return -uint32(1 << uint(32-nbits))
}

type ipv4Int uint32

func (i ipv4Int) IP() net.IP {
	var b [net.IPv4len]byte
	binary.BigEndian.PutUint32(b[:], uint32(i))
	return net.IP(b[:])
}

func (i ipv4Int) encodeNLRI(b []byte, nbits byte) int {
	l := int(nbits) / 8
	if nbits%8 > 0 {
		l++
	}
	b[0] = nbits
	l++
	b = b[1:]
	for n := 0; n < l; n++ {
		b[n] = byte(i >> uint(32-8*(n+1)))
	}
	return l
}

func (p *IPv4) decodeNLRI(b []byte) error {
	p.nbits = b[0]
	l := len(b) - 1
	b = b[1:]
	for n := 0; n < l; n++ {
		p.addr |= ipv4Int(b[n]) << uint(32-8*(n+1))
	}
	return nil
}

func ipToIPv4Int(ip []byte) ipv4Int {
	return ipv4Int(binary.BigEndian.Uint32(ip))
}

type ipv6Int [2]uint64

func (i *ipv6Int) lshift(nbits byte) {
	i[0] = i[0]<<uint(nbits) | i[1]>>uint(64-nbits) | i[1]<<uint(nbits-64)
	i[1] = i[1] << uint(nbits)
}

func (i *ipv6Int) rshift(nbits byte) {
	i[1] = i[1]>>uint(nbits) | i[0]<<uint(64-nbits) | i[0]>>uint(nbits-64)
	i[0] = i[0] >> uint(nbits)
}

func (i *ipv6Int) incr() {
	if i[1] == math.MaxUint64 {
		i[0]++
		i[1] = 0
	} else {
		i[1]++
	}
}

func (i *ipv6Int) setHostmask(nbits byte) {
	if nbits > 64 {
		i[0], i[1] = ^mask64(64), ^mask64(nbits-64)
	} else {
		i[0], i[1] = ^mask64(nbits), mask64(64)
	}
}

func (i *ipv6Int) setNetmask(nbits byte) {
	if nbits > 64 {
		i[0], i[1] = mask64(64), mask64(nbits-64)
	} else {
		i[0], i[1] = mask64(nbits), 0
	}
}

func (i *ipv6Int) IP() net.IP {
	var b [net.IPv6len]byte
	binary.BigEndian.PutUint64(b[:8], i[0])
	binary.BigEndian.PutUint64(b[8:16], i[1])
	return net.IP(b[:])
}

func (i *ipv6Int) encodeNLRI(b []byte, nbits byte) int {
	l := int(nbits) / 8
	if nbits%8 > 0 {
		l++
	}
	b[0] = nbits
	l++
	b = b[1:]
	for n := 0; n < l; n++ {
		switch {
		case n < 8:
			b[n] = byte(i[0] >> uint(64-(8*(n+1))))
		case 8 <= n && n < 16:
			b[n] = byte(i[1] >> uint(128-(8*(n+1))))
		}
	}
	return l
}

func (p *IPv6) decodeNLRI(b []byte) error {
	p.nbits = b[0]
	l := len(b) - 1
	b = b[1:]
	for n := 0; n < l; n++ {
		switch {
		case n < 8:
			p.addr[0] |= uint64(b[n]) << uint(64-8*(n+1))
		case 8 <= n && n < 16:
			p.addr[1] |= uint64(b[n]) << uint(128-8*(n+1))
		}
	}
	return nil
}

func ipToIPv6Int(ip []byte) ipv6Int {
	return ipv6Int{binary.BigEndian.Uint64(ip[:8]), binary.BigEndian.Uint64(ip[8:16])}
}
<|endoftext|>