<commit_before>package bot

import (
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/lnxjedi/readline"
	"github.com/lnxjedi/robot"
)

func init() {
	RegisterPreload("connectors/terminal.so")
	RegisterConnector("terminal", Initialize)
}

// Global persistent map of user name to user index
var userIDMap = make(map[string]int)
var userMap = make(map[string]int)

type termUser struct {
	Name                                        string // username / handle
	InternalID                                  string // connector internal identifier
	Email, FullName, FirstName, LastName, Phone string
}

type termconfig struct {
	StartChannel string // the initial channel
	StartUser    string // the initial userid
	EOF          string // command to send on EOF (ctrl-D), default ";quit"
	Abort        string // command to send on ctrl-c
	Users        []termUser
	Channels     []string
}

// termConnector holds all the relevant data about a connection
type termConnector struct {
	currentChannel string             // The current channel for the user
	currentUser    string             // The current userid
	eof            string             // command to send on ctrl-d (EOF)
	abort          string             // command to send on ctrl-c (interrupt)
	running        bool               // set on call to Run
	width          int                // width of terminal
	users          []termUser         // configured users
	channels       []string           // the channels the robot is in
	heard          chan string        // when the user speaks
	reader         *readline.Instance // readline for speaking
	robot.Handler                     // bot API for connectors
	sync.RWMutex                      // shared mutex for locking connector data structures
}

var exit = struct {
	kbquit, robotexit bool
	waitchan          chan struct{}
	sync.Mutex
}{
	false, false,
	make(chan struct{}),
	sync.Mutex{},
}

var quitTimeout = 4 * time.Second

var lock sync.Mutex // package var lock
var started bool    // set when connector is started

// Initialize sets up the connector and returns a connector object
func Initialize(handler robot.Handler, l *log.Logger) robot.Connector {
	lock.Lock()
	if started {
		lock.Unlock()
		return nil
	}
	started = true
	lock.Unlock()

	var c termconfig

	err := handler.GetProtocolConfig(&c)
	if err != nil {
		handler.Log(robot.Fatal, "Unable to retrieve protocol configuration: %v", err)
	}
	eof := ";quit"
	abort := ";abort"
	if len(c.EOF) > 0 {
		eof = c.EOF
	}
	if len(c.Abort) > 0 {
		abort = c.Abort
	}
	found := false
	for i, u := range c.Users {
		userMap[u.Name] = i
		userIDMap[u.InternalID] = i
		if c.StartUser == u.Name {
			found = true
		}
	}
	if !found {
		handler.Log(robot.Fatal, "Start user \"%s\" not listed in Users array", c.StartUser)
	}

	found = false
	for _, ch := range c.Channels {
		if c.StartChannel == ch {
			found = true
		}
	}
	if !found {
		handler.Log(robot.Fatal, "Start channel \"%s\" not listed in Channels array", c.StartChannel)
	}

	var histfile string
	home := os.Getenv("HOME")
	if len(home) == 0 {
		home = os.Getenv("USERPROFILE")
	}
	if len(home) > 0 {
		histfile = path.Join(home, ".gopherbot_history")
	}

	rl, err := readline.NewEx(&readline.Config{
		Prompt:            fmt.Sprintf("c:%s/u:%s -> ", c.StartChannel, c.StartUser),
		HistoryFile:       histfile,
		HistorySearchFold: true,
		InterruptPrompt:   "abort",
		EOFPrompt:         "exit",
	})
	if err != nil {
		panic(err)
	}

	tc := &termConnector{
		currentChannel: c.StartChannel,
		currentUser:    c.StartUser,
		eof:            eof,
		abort:          abort,
		channels:       c.Channels,
		running:        false,
		width:          readline.GetScreenWidth(),
		users:          c.Users,
		heard:          make(chan string),
		reader:         rl,
	}

	tc.Handler = handler
	tc.SetTerminalWriter(tc.reader)
	return robot.Connector(tc)
}

func (tc *termConnector) Run(stop <-chan struct{}) {
	tc.Lock()
	// This should never happen, just a bit of defensive coding
	if tc.running {
		tc.Unlock()
		return
	}
	tc.running = true
	tc.Unlock()
	defer func() {
	}()

	// listen loop
	go func(tc *termConnector) {
	readloop:
		for {
			line, err := tc.reader.Readline()
			exit.Lock()
			robotexit := exit.robotexit
			if robotexit {
				exit.Unlock()
				tc.heard <- ""
				break readloop
			}
			kbquit := false
			if err == io.EOF {
				tc.heard <- tc.eof
				kbquit = true
			} else if err == readline.ErrInterrupt {
				tc.heard <- tc.abort
				kbquit = true
			} else if err == nil {
				tc.heard <- line
				line = strings.TrimSpace(line)
				if line == tc.eof || line == tc.abort {
					kbquit = true
				}
			}
			if kbquit {
				exit.kbquit = true
				exit.Unlock()
				select {
				case <-exit.waitchan:
					break readloop
				case <-time.After(quitTimeout):
					exit.Lock()
					exit.kbquit = false
					exit.Unlock()
					tc.reader.Write([]byte("(timed out waiting for robot to exit; check terminal connector settings 'EOF' and 'Abort')\n"))
				}
			} else {
				exit.Unlock()
			}
		}
	}(tc)

	tc.reader.Write([]byte("Terminal connector running; Type '|c?' to list channels, '|u?' to list users\n"))

	kbquit := false

loop:
	// Main loop and prompting
	for {
		select {
		case <-stop:
			tc.Log(robot.Info, "Received stop in connector")
			exit.Lock()
			kbquit = exit.kbquit
			exit.robotexit = true
			exit.Unlock()
			if kbquit {
				exit.waitchan <- struct{}{}
			} else {
				tc.reader.Write([]byte("Exiting (press <enter> ...)\n"))
			}
			break loop
		case input := <-tc.heard:
			if len(input) == 0 {
				evs := tc.GetEventStrings()
				if len(*evs) > 0 {
					tc.reader.Write([]byte(fmt.Sprintf("Events gathered: %s\n", strings.Join(*evs, ", "))))
				}
				continue
			}
			if input[0] == '|' {
				if len(input) == 1 {
					continue
				}
				switch input[1] {
				case 'C', 'c':
					exists := false
					newchan := input[2:]
					if newchan == "?" {
						chanlist := []string{"Available channels:", "(direct message); type: '|c'"}
						for _, channel := range tc.channels {
							chanlist = append(chanlist, fmt.Sprintf("'%s'; type: '|c%s'", channel, channel))
						}
						tc.reader.Write([]byte(strings.Join(chanlist, "\n")))
						tc.reader.Write([]byte("\n"))
						continue
					}
					tc.Lock()
					if newchan == "" {
						tc.currentChannel = ""
						tc.reader.SetPrompt(fmt.Sprintf("c:(direct)/u:%s -> ", tc.currentUser))
						tc.reader.Write([]byte("Changed current channel to: direct message\n"))
					} else {
						for _, ch := range tc.channels {
							if ch == newchan {
								exists = true
								break
							}
						}
						if exists {
							tc.currentChannel = newchan
							tc.reader.SetPrompt(fmt.Sprintf("c:%s/u:%s -> ", tc.currentChannel, tc.currentUser))
							tc.reader.Write([]byte(fmt.Sprintf("Changed current channel to: %s\n", newchan)))
						} else {
							tc.reader.Write([]byte("Invalid channel\n"))
						}
					}
					tc.Unlock()
				case 'U', 'u':
					exists := false
					newuser := input[2:]
					if newuser == "?" {
						userlist := []string{"Available users:"}
						for _, user := range tc.users {
							userlist = append(userlist, fmt.Sprintf("'%s'; type: '|u%s'", user.Name, user.Name))
						}
						tc.reader.Write([]byte(strings.Join(userlist, "\n")))
						tc.reader.Write([]byte("\n"))
						continue
					}
					tc.Lock()
					if newuser == "" {
						tc.reader.Write([]byte("Invalid 0-length user\n"))
					} else {
						for _, u := range tc.users {
							if u.Name == newuser {
								exists = true
							}
						}
						if exists {
							tc.currentUser = newuser
							tc.reader.SetPrompt(fmt.Sprintf("c:%s/u:%s -> ", tc.currentChannel, tc.currentUser))
							tc.reader.Write([]byte(fmt.Sprintf("Changed current user to: %s\n", newuser)))
						} else {
							tc.reader.Write([]byte("Invalid user\n"))
						}
					}
					tc.Unlock()
				default:
					tc.reader.Write([]byte("Invalid terminal connector command\n"))
				}
			} else {
				var channelID string
				direct := false
				if len(tc.currentChannel) > 0 {
					channelID = "#" + tc.currentChannel
				} else {
					direct = true
				}
				i := userMap[tc.currentUser]
				ui := tc.users[i]
				botMsg := &robot.ConnectorMessage{
					Protocol:      "terminal",
					UserName:      tc.currentUser,
					UserID:        ui.InternalID,
					ChannelName:   tc.currentChannel,
					ChannelID:     channelID,
					MessageText:   input,
					DirectMessage: direct,
				}
				tc.RLock()
				tc.IncomingMessage(botMsg)
				tc.RUnlock()
			}
		}
	}
	if !kbquit {
		<-tc.heard
	}
	tc.reader.Write([]byte("Terminal connector finished\n"))
	tc.reader.Close()
}

func (tc *termConnector) MessageHeard(u, c string) {
	return
}

func (tc *termConnector) getUserInfo(u string) (*termUser, bool) {
	var i int
	var exists bool
	if id, ok := tc.ExtractID(u); ok {
		i, exists = userIDMap[id]
	} else {
		i, exists = userMap[u]
	}
	if exists {
		return &tc.users[i], true
	}
	return nil, false
}

func (tc *termConnector) getChannel(c string) string {
	if ch, ok := tc.ExtractID(c); ok {
		return strings.TrimPrefix(ch, "#")
	}
	return c
}

// SetUserMap lets Gopherbot provide a mapping of usernames to user IDs
func (tc *termConnector) SetUserMap(map[string]string) {
	return
}

// GetUserAttribute returns a string attribute or nil if slack doesn't
// have that information
func (tc *termConnector) GetProtocolUserAttribute(u, attr string) (value string, ret robot.RetVal) {
	var user *termUser
	var exists bool
	if user, exists = tc.getUserInfo(u); !exists {
		return "", robot.UserNotFound
	}
	switch attr {
	case "email":
		return user.Email, robot.Ok
	case "internalid":
		return user.InternalID, robot.Ok
	case "realname", "fullname", "real name", "full name":
		return user.FullName, robot.Ok
	case "firstname", "first name":
		return user.FirstName, robot.Ok
	case "lastname", "last name":
		return user.LastName, robot.Ok
	case "phone":
		return user.Phone, robot.Ok
	// that's all the attributes we can currently get from slack
	default:
		return "", robot.AttributeNotFound
	}
}

// SendProtocolChannelMessage sends a message to a channel
func (tc *termConnector) SendProtocolChannelMessage(ch string, msg string, f robot.MessageFormat) (ret robot.RetVal) {
	channel := tc.getChannel(ch)
	return tc.sendMessage(channel, msg, f)
}

// SendProtocolChannelMessage sends a message to a channel
func (tc *termConnector) SendProtocolUserChannelMessage(uid, uname, ch, msg string, f robot.MessageFormat) (ret robot.RetVal) {
	channel := tc.getChannel(ch)
	msg = "@" + uname + " " + msg
	return tc.sendMessage(channel, msg, f)
}

// SendProtocolUserMessage sends a direct message to a user
func (tc *termConnector) SendProtocolUserMessage(u string, msg string, f robot.MessageFormat) (ret robot.RetVal) {
	var user *termUser
	var exists bool
	if user, exists = tc.getUserInfo(u); !exists {
		return robot.UserNotFound
	}
	return tc.sendMessage(fmt.Sprintf("(dm:%s)", user.Name), msg, f)
}

// JoinChannel joins a channel given it's human-readable name, e.g. "general"
// Only useful for connectors that require it, a noop otherwise
func (tc *termConnector) JoinChannel(c string) (ret robot.RetVal) {
	return robot.Ok
}
<commit_msg>Less fancy, more correct - fixes #225<commit_after>package bot

import (
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/lnxjedi/readline"
	"github.com/lnxjedi/robot"
)

func init() {
	RegisterPreload("connectors/terminal.so")
	RegisterConnector("terminal", Initialize)
}

// Global persistent map of user name to user index
var userIDMap = make(map[string]int)
var userMap = make(map[string]int)

type termUser struct {
	Name                                        string // username / handle
	InternalID                                  string // connector internal identifier
	Email, FullName, FirstName, LastName, Phone string
}

type termconfig struct {
	StartChannel string // the initial channel
	StartUser    string // the initial userid
	EOF          string // command to send on EOF (ctrl-D), default ";quit"
	Abort        string // command to send on ctrl-c
	Users        []termUser
	Channels     []string
}

// termConnector holds all the relevant data about a connection
type termConnector struct {
	currentChannel string             // The current channel for the user
	currentUser    string             // The current userid
	eof            string             // command to send on ctrl-d (EOF)
	abort          string             // command to send on ctrl-c (interrupt)
	running        bool               // set on call to Run
	width          int                // width of terminal
	users          []termUser         // configured users
	channels       []string           // the channels the robot is in
	heard          chan string        // when the user speaks
	reader         *readline.Instance // readline for speaking
	robot.Handler                     // bot API for connectors
	sync.RWMutex                      // shared mutex for locking connector data structures
}

var exit = struct {
	kbquit, robotexit bool
	waitchan          chan struct{}
	sync.Mutex
}{
	false, false,
	make(chan struct{}),
	sync.Mutex{},
}

var quitTimeout = 4 * time.Second

var lock sync.Mutex // package var lock
var started bool    // set when connector is started

// Initialize sets up the connector and returns a connector object
func Initialize(handler robot.Handler, l *log.Logger) robot.Connector {
	lock.Lock()
	if started {
		lock.Unlock()
		return nil
	}
	started = true
	lock.Unlock()

	var c termconfig

	err := handler.GetProtocolConfig(&c)
	if err != nil {
		handler.Log(robot.Fatal, "Unable to retrieve protocol configuration: %v", err)
	}
	eof := ";quit"
	abort := ";abort"
	if len(c.EOF) > 0 {
		eof = c.EOF
	}
	if len(c.Abort) > 0 {
		abort = c.Abort
	}
	found := false
	for i, u := range c.Users {
		userMap[u.Name] = i
		userIDMap[u.InternalID] = i
		if c.StartUser == u.Name {
			found = true
		}
	}
	if !found {
		handler.Log(robot.Fatal, "Start user \"%s\" not listed in Users array", c.StartUser)
	}

	found = false
	for _, ch := range c.Channels {
		if c.StartChannel == ch {
			found = true
		}
	}
	if !found {
		handler.Log(robot.Fatal, "Start channel \"%s\" not listed in Channels array", c.StartChannel)
	}

	var histfile string
	home := os.Getenv("HOME")
	if len(home) == 0 {
		home = os.Getenv("USERPROFILE")
	}
	if len(home) > 0 {
		histfile = path.Join(home, ".gopherbot_history")
	}

	rl, err := readline.NewEx(&readline.Config{
		Prompt:            fmt.Sprintf("c:%s/u:%s -> ", c.StartChannel, c.StartUser),
		HistoryFile:       histfile,
		HistorySearchFold: true,
		InterruptPrompt:   "abort",
		EOFPrompt:         "exit",
	})
	if err != nil {
		panic(err)
	}

	tc := &termConnector{
		currentChannel: c.StartChannel,
		currentUser:    c.StartUser,
		eof:            eof,
		abort:          abort,
		channels:       c.Channels,
		running:        false,
		width:          readline.GetScreenWidth(),
		users:          c.Users,
		heard:          make(chan string),
		reader:         rl,
	}

	tc.Handler = handler
	tc.SetTerminalWriter(tc.reader)
	return robot.Connector(tc)
}

func (tc *termConnector) Run(stop <-chan struct{}) {
	tc.Lock()
	// This should never happen, just a bit of defensive coding
	if tc.running {
		tc.Unlock()
		return
	}
	tc.running = true
	tc.Unlock()
	defer func() {
	}()

	// listen loop
	go func(tc *termConnector) {
	readloop:
		for {
			line, err := tc.reader.Readline()
			exit.Lock()
			robotexit := exit.robotexit
			if robotexit {
				exit.Unlock()
				tc.heard <- ""
				break readloop
			}
			kbquit := false
			if err == io.EOF {
				tc.heard <- tc.eof
				kbquit = true
			} else if err == readline.ErrInterrupt {
				tc.heard <- tc.abort
				kbquit = true
			} else if err == nil {
				tc.heard <- line
				line = strings.TrimSpace(line)
				if line == tc.eof || line == tc.abort {
					kbquit = true
				}
			}
			if kbquit {
				exit.kbquit = true
				exit.Unlock()
				select {
				case <-exit.waitchan:
					break readloop
				case <-time.After(quitTimeout):
					exit.Lock()
					exit.kbquit = false
					exit.Unlock()
					tc.reader.Write([]byte("(timed out waiting for robot to exit; check terminal connector settings 'EOF' and 'Abort')\n"))
				}
			} else {
				exit.Unlock()
			}
		}
	}(tc)

	tc.reader.Write([]byte("Terminal connector running; Type '|c?' to list channels, '|u?' to list users\n"))

	kbquit := false

loop:
	// Main loop and prompting
	for {
		select {
		case <-stop:
			tc.Log(robot.Info, "Received stop in connector")
			exit.Lock()
			kbquit = exit.kbquit
			exit.robotexit = true
			exit.Unlock()
			if kbquit {
				exit.waitchan <- struct{}{}
			} else {
				tc.reader.Write([]byte("Exiting (press <enter> ...)\n"))
			}
			break loop
		case input := <-tc.heard:
			if len(input) == 0 {
				evs := tc.GetEventStrings()
				if len(*evs) > 0 {
					tc.reader.Write([]byte(fmt.Sprintf("Events gathered: %s\n", strings.Join(*evs, ", "))))
				}
				continue
			}
			if input[0] == '|' {
				if len(input) == 1 {
					continue
				}
				switch input[1] {
				case 'C', 'c':
					exists := false
					newchan := input[2:]
					if newchan == "?" {
						tc.reader.Write([]byte("Available channels:\n"))
						tc.reader.Write([]byte("(direct message); type: '|c'\n"))
						for _, channel := range tc.channels {
							tc.reader.Write([]byte(fmt.Sprintf("'%s'; type: '|c%s'\n", channel, channel)))
						}
						continue
					}
					tc.Lock()
					if newchan == "" {
						tc.currentChannel = ""
						tc.reader.SetPrompt(fmt.Sprintf("c:(direct)/u:%s -> ", tc.currentUser))
						tc.reader.Write([]byte("Changed current channel to: direct message\n"))
					} else {
						for _, ch := range tc.channels {
							if ch == newchan {
								exists = true
								break
							}
						}
						if exists {
							tc.currentChannel = newchan
							tc.reader.SetPrompt(fmt.Sprintf("c:%s/u:%s -> ", tc.currentChannel, tc.currentUser))
							tc.reader.Write([]byte(fmt.Sprintf("Changed current channel to: %s\n", newchan)))
						} else {
							tc.reader.Write([]byte("Invalid channel\n"))
						}
					}
					tc.Unlock()
				case 'U', 'u':
					exists := false
					newuser := input[2:]
					if newuser == "?" {
						tc.reader.Write([]byte("Available users:\n"))
						for _, user := range tc.users {
							tc.reader.Write([]byte(fmt.Sprintf("'%s'; type: '|u%s'\n", user.Name, user.Name)))
						}
						continue
					}
					tc.Lock()
					if newuser == "" {
						tc.reader.Write([]byte("Invalid 0-length user\n"))
					} else {
						for _, u := range tc.users {
							if u.Name == newuser {
								exists = true
							}
						}
						if exists {
							tc.currentUser = newuser
							tc.reader.SetPrompt(fmt.Sprintf("c:%s/u:%s -> ", tc.currentChannel, tc.currentUser))
							tc.reader.Write([]byte(fmt.Sprintf("Changed current user to: %s\n", newuser)))
						} else {
							tc.reader.Write([]byte("Invalid user\n"))
						}
					}
					tc.Unlock()
				default:
					tc.reader.Write([]byte("Invalid terminal connector command\n"))
				}
			} else {
				var channelID string
				direct := false
				if len(tc.currentChannel) > 0 {
					channelID = "#" + tc.currentChannel
				} else {
					direct = true
				}
				i := userMap[tc.currentUser]
				ui := tc.users[i]
				botMsg := &robot.ConnectorMessage{
					Protocol:      "terminal",
					UserName:      tc.currentUser,
					UserID:        ui.InternalID,
					ChannelName:   tc.currentChannel,
					ChannelID:     channelID,
					MessageText:   input,
					DirectMessage: direct,
				}
				tc.RLock()
				tc.IncomingMessage(botMsg)
				tc.RUnlock()
			}
		}
	}
	if !kbquit {
		<-tc.heard
	}
	tc.reader.Write([]byte("Terminal connector finished\n"))
	tc.reader.Close()
}

func (tc *termConnector) MessageHeard(u, c string) {
	return
}

func (tc *termConnector) getUserInfo(u string) (*termUser, bool) {
	var i int
	var exists bool
	if id, ok := tc.ExtractID(u); ok {
		i, exists = userIDMap[id]
	} else {
		i, exists = userMap[u]
	}
	if exists {
		return &tc.users[i], true
	}
	return nil, false
}

func (tc *termConnector) getChannel(c string) string {
	if ch, ok := tc.ExtractID(c); ok {
		return strings.TrimPrefix(ch, "#")
	}
	return c
}

// SetUserMap lets Gopherbot provide a mapping of usernames to user IDs
func (tc *termConnector) SetUserMap(map[string]string) {
	return
}

// GetUserAttribute returns a string attribute or nil if slack doesn't
// have that information
func (tc *termConnector) GetProtocolUserAttribute(u, attr string) (value string, ret robot.RetVal) {
	var user *termUser
	var exists bool
	if user, exists = tc.getUserInfo(u); !exists {
		return "", robot.UserNotFound
	}
	switch attr {
	case "email":
		return user.Email, robot.Ok
	case "internalid":
		return user.InternalID, robot.Ok
	case "realname", "fullname", "real name", "full name":
		return user.FullName, robot.Ok
	case "firstname", "first name":
		return user.FirstName, robot.Ok
	case "lastname", "last name":
		return user.LastName, robot.Ok
	case "phone":
		return user.Phone, robot.Ok
	// that's all the attributes we can currently get from slack
	default:
		return "", robot.AttributeNotFound
	}
}

// SendProtocolChannelMessage sends a message to a channel
func (tc *termConnector) SendProtocolChannelMessage(ch string, msg string, f robot.MessageFormat) (ret robot.RetVal) {
	channel := tc.getChannel(ch)
	return tc.sendMessage(channel, msg, f)
}

// SendProtocolChannelMessage sends a message to a channel
func (tc *termConnector) SendProtocolUserChannelMessage(uid, uname, ch, msg string, f robot.MessageFormat) (ret robot.RetVal) {
	channel := tc.getChannel(ch)
	msg = "@" + uname + " " + msg
	return tc.sendMessage(channel, msg, f)
}

// SendProtocolUserMessage sends a direct message to a user
func (tc *termConnector) SendProtocolUserMessage(u string, msg string, f robot.MessageFormat) (ret robot.RetVal) {
	var user *termUser
	var exists bool
	if user, exists = tc.getUserInfo(u); !exists {
		return robot.UserNotFound
	}
	return tc.sendMessage(fmt.Sprintf("(dm:%s)", user.Name), msg, f)
}

// JoinChannel joins a channel given it's human-readable name, e.g. "general"
// Only useful for connectors that require it, a noop otherwise
func (tc *termConnector) JoinChannel(c string) (ret robot.RetVal) {
	return robot.Ok
}
<|endoftext|>


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

package executor

import (
	"fmt"
	"math"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/juju/errors"
	"github.com/pingcap/tidb/ast"
	"github.com/pingcap/tidb/config"
	"github.com/pingcap/tidb/context"
	"github.com/pingcap/tidb/infoschema"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/model"
	"github.com/pingcap/tidb/mysql"
	"github.com/pingcap/tidb/plan"
	"github.com/pingcap/tidb/sessionctx/variable"
	"github.com/pingcap/tidb/terror"
	"github.com/pingcap/tidb/util/logutil"
)

type processinfoSetter interface {
	SetProcessInfo(string)
}

// recordSet wraps an executor, implements ast.RecordSet interface
type recordSet struct {
	fields      []*ast.ResultField
	executor    Executor
	stmt        *ExecStmt
	processinfo processinfoSetter
	lastErr     error
}

func (a *recordSet) Fields() []*ast.ResultField {
	if len(a.fields) == 0 {
		for _, col := range a.executor.Schema().Columns {
			dbName := col.DBName.O
			if dbName == "" && col.TblName.L != "" {
				dbName = a.stmt.ctx.GetSessionVars().CurrentDB
			}
			rf := &ast.ResultField{
				ColumnAsName: col.ColName,
				TableAsName:  col.TblName,
				DBName:       model.NewCIStr(dbName),
				Table:        &model.TableInfo{Name: col.OrigTblName},
				Column: &model.ColumnInfo{
					FieldType: *col.RetType,
					Name:      col.ColName,
				},
			}
			a.fields = append(a.fields, rf)
		}
	}
	return a.fields
}

func (a *recordSet) Next() (*ast.Row, error) {
	row, err := a.executor.Next()
	if err != nil {
		a.lastErr = err
		return nil, errors.Trace(err)
	}
	if row == nil {
		if a.stmt != nil {
			a.stmt.ctx.GetSessionVars().LastFoundRows = a.stmt.ctx.GetSessionVars().StmtCtx.FoundRows()
		}
		return nil, nil
	}

	if a.stmt != nil {
		a.stmt.ctx.GetSessionVars().StmtCtx.AddFoundRows(1)
	}
	return &ast.Row{Data: row}, nil
}

func (a *recordSet) Close() error {
	err := a.executor.Close()
	a.stmt.logSlowQuery(a.lastErr == nil)
	if a.processinfo != nil {
		a.processinfo.SetProcessInfo("")
	}
	return errors.Trace(err)
}

// ExecStmt implements the ast.Statement interface, it builds a plan.Plan to an ast.Statement.
type ExecStmt struct {
	// InfoSchema stores a reference to the schema information.
	InfoSchema infoschema.InfoSchema
	// Plan stores a reference to the final physical plan.
	Plan plan.Plan
	// Expensive represents whether this query is an expensive one.
	Expensive bool
	// Cacheable represents whether the physical plan can be cached.
	Cacheable bool
	// Text represents the origin query text.
	Text string

	ctx            context.Context
	startTime      time.Time
	isPreparedStmt bool

	// ReadOnly represents the statement is read-only.
	ReadOnly bool
}

// OriginText implements ast.Statement interface.
func (a *ExecStmt) OriginText() string {
	return a.Text
}

// IsPrepared implements ast.Statement interface.
func (a *ExecStmt) IsPrepared() bool {
	return a.isPreparedStmt
}

// IsReadOnly implements ast.Statement interface.
func (a *ExecStmt) IsReadOnly() bool {
	return a.ReadOnly
}

// Exec implements the ast.Statement Exec interface.
// This function builds an Executor from a plan. If the Executor doesn't return result,
// like the INSERT, UPDATE statements, it executes in this function, if the Executor returns
// result, execution is done after this function returns, in the returned ast.RecordSet Next method.
func (a *ExecStmt) Exec(ctx context.Context) (ast.RecordSet, error) {
	a.startTime = time.Now()
	a.ctx = ctx

	if _, ok := a.Plan.(*plan.Analyze); ok && ctx.GetSessionVars().InRestrictedSQL {
		oriStats := ctx.GetSessionVars().Systems[variable.TiDBBuildStatsConcurrency]
		oriScan := ctx.GetSessionVars().DistSQLScanConcurrency
		oriIndex := ctx.GetSessionVars().IndexSerialScanConcurrency
		oriIso := ctx.GetSessionVars().Systems[variable.TxnIsolation]
		ctx.GetSessionVars().Systems[variable.TiDBBuildStatsConcurrency] = "1"
		ctx.GetSessionVars().DistSQLScanConcurrency = 1
		ctx.GetSessionVars().IndexSerialScanConcurrency = 1
		ctx.GetSessionVars().Systems[variable.TxnIsolation] = ast.ReadCommitted
		defer func() {
			ctx.GetSessionVars().Systems[variable.TiDBBuildStatsConcurrency] = oriStats
			ctx.GetSessionVars().DistSQLScanConcurrency = oriScan
			ctx.GetSessionVars().IndexSerialScanConcurrency = oriIndex
			ctx.GetSessionVars().Systems[variable.TxnIsolation] = oriIso
		}()
	}

	e, err := a.buildExecutor(ctx)
	if err != nil {
		return nil, errors.Trace(err)
	}

	if err := e.Open(); err != nil {
		return nil, errors.Trace(err)
	}

	var pi processinfoSetter
	if raw, ok := ctx.(processinfoSetter); ok {
		pi = raw
		sql := a.OriginText()
		if simple, ok := a.Plan.(*plan.Simple); ok && simple.Statement != nil {
			if ss, ok := simple.Statement.(ast.SensitiveStmtNode); ok {
				// Use SecureText to avoid leak password information.
				sql = ss.SecureText()
			}
		}
		// Update processinfo, ShowProcess() will use it.
		pi.SetProcessInfo(sql)
	}
	// Fields or Schema are only used for statements that return result set.
	if e.Schema().Len() == 0 {
		return a.handleNoDelayExecutor(e, ctx, pi)
	}

	return &recordSet{
		executor:    e,
		stmt:        a,
		processinfo: pi,
	}, nil
}

func (a *ExecStmt) handleNoDelayExecutor(e Executor, ctx context.Context, pi processinfoSetter) (ast.RecordSet, error) {
	// Check if "tidb_snapshot" is set for the write executors.
	// In history read mode, we can not do write operations.
	switch e.(type) {
	case *DeleteExec, *InsertExec, *UpdateExec, *ReplaceExec, *LoadData, *DDLExec:
		snapshotTS := ctx.GetSessionVars().SnapshotTS
		if snapshotTS != 0 {
			return nil, errors.New("can not execute write statement when 'tidb_snapshot' is set")
		}
	}

	var err error
	defer func() {
		if pi != nil {
			pi.SetProcessInfo("")
		}
		terror.Log(errors.Trace(e.Close()))
		a.logSlowQuery(err == nil)
	}()
	for {
		var row Row
		row, err = e.Next()
		if err != nil {
			return nil, errors.Trace(err)
		}
		// Even though there isn't any result set, the row is still used to indicate if there is
		// more work to do.
		// For example, the UPDATE statement updates a single row on a Next call, we keep calling Next until
		// There is no more rows to update.
		if row == nil {
			return nil, nil
		}
	}
}

// buildExecutor build a executor from plan, prepared statement may need additional procedure.
func (a *ExecStmt) buildExecutor(ctx context.Context) (Executor, error) {
	priority := kv.PriorityNormal
	if _, ok := a.Plan.(*plan.Execute); !ok {
		// Do not sync transaction for Execute statement, because the real optimization work is done in
		// "ExecuteExec.Build".
		var err error
		isPointGet := IsPointGetWithPKOrUniqueKeyByAutoCommit(ctx, a.Plan)
		if isPointGet {
			log.Debugf("[%d][InitTxnWithStartTS] %s", ctx.GetSessionVars().ConnectionID, a.Text)
			err = ctx.InitTxnWithStartTS(math.MaxUint64)
		} else {
			log.Debugf("[%d][ActivePendingTxn] %s", ctx.GetSessionVars().ConnectionID, a.Text)
			err = ctx.ActivePendingTxn()
		}
		if err != nil {
			return nil, errors.Trace(err)
		}

		if stmtPri := ctx.GetSessionVars().StmtCtx.Priority; stmtPri != mysql.NoPriority {
			priority = int(stmtPri)
		} else {
			switch {
			case isPointGet:
				priority = kv.PriorityHigh
			case a.Expensive:
				priority = kv.PriorityLow
			}
		}
	}
	if _, ok := a.Plan.(*plan.Analyze); ok && ctx.GetSessionVars().InRestrictedSQL {
		priority = kv.PriorityLow
	}

	b := newExecutorBuilder(ctx, a.InfoSchema, priority)
	e := b.build(a.Plan)
	if b.err != nil {
		return nil, errors.Trace(b.err)
	}

	// ExecuteExec is not a real Executor, we only use it to build another Executor from a prepared statement.
	if executorExec, ok := e.(*ExecuteExec); ok {
		err := executorExec.Build()
		if err != nil {
			return nil, errors.Trace(err)
		}
		a.Text = executorExec.Stmt.Text()
		a.isPreparedStmt = true
		a.Plan = executorExec.Plan
		e = executorExec.StmtExec
	}
	return e, nil
}

func (a *ExecStmt) logSlowQuery(succ bool) {
	cfg := config.GetGlobalConfig()
	costTime := time.Since(a.startTime)
	sql := a.Text
	if len(sql) > cfg.Log.QueryLogMaxLen {
		sql = fmt.Sprintf("%.*q(len:%d)", cfg.Log.QueryLogMaxLen, sql, len(a.Text))
	}
	connID := a.ctx.GetSessionVars().ConnectionID
	currentDB := a.ctx.GetSessionVars().CurrentDB
	logEntry := log.NewEntry(logutil.SlowQueryLogger)
	logEntry.Data = log.Fields{
		"connectionId": connID,
		"costTime":     costTime,
		"database":     currentDB,
		"sql":          sql,
	}
	if costTime < time.Duration(cfg.Log.SlowThreshold)*time.Millisecond {
		logEntry.WithField("type", "query").WithField("succ", succ).Debugf("query")
	} else {
		logEntry.WithField("type", "slow-query").WithField("succ", succ).Warnf("slow-query")
	}
}

// IsPointGetWithPKOrUniqueKeyByAutoCommit returns true when meets following conditions:
//  1. ctx is auto commit tagged
//  2. txn is nil
//  2. plan is point get by pk or unique key
func IsPointGetWithPKOrUniqueKeyByAutoCommit(ctx context.Context, p plan.Plan) bool {
	// check auto commit
	if !ctx.GetSessionVars().IsAutocommit() {
		return false
	}

	// check txn
	if ctx.Txn() != nil {
		return false
	}

	// check plan
	if proj, ok := p.(*plan.Projection); ok {
		if len(proj.Children()) != 1 {
			return false
		}
		p = proj.Children()[0]
	}

	switch v := p.(type) {
	case *plan.PhysicalIndexScan:
		return v.IsPointGetByUniqueKey(ctx.GetSessionVars().StmtCtx)
	case *plan.PhysicalIndexReader:
		indexScan := v.IndexPlans[0].(*plan.PhysicalIndexScan)
		return indexScan.IsPointGetByUniqueKey(ctx.GetSessionVars().StmtCtx)
	case *plan.PhysicalIndexLookUpReader:
		indexScan := v.IndexPlans[0].(*plan.PhysicalIndexScan)
		return indexScan.IsPointGetByUniqueKey(ctx.GetSessionVars().StmtCtx)
	case *plan.PhysicalTableScan:
		return len(v.Ranges) == 1 && v.Ranges[0].IsPoint()
	case *plan.PhysicalTableReader:
		tableScan := v.TablePlans[0].(*plan.PhysicalTableScan)
		return len(tableScan.Ranges) == 1 && tableScan.Ranges[0].IsPoint()
	default:
		return false
	}
}
<commit_msg>executor: Add debug logs (#5083)<commit_after>// Copyright 2015 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package executor

import (
	"fmt"
	"math"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/juju/errors"
	"github.com/pingcap/tidb/ast"
	"github.com/pingcap/tidb/config"
	"github.com/pingcap/tidb/context"
	"github.com/pingcap/tidb/infoschema"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/model"
	"github.com/pingcap/tidb/mysql"
	"github.com/pingcap/tidb/plan"
	"github.com/pingcap/tidb/sessionctx/variable"
	"github.com/pingcap/tidb/terror"
	"github.com/pingcap/tidb/util/logutil"
)

type processinfoSetter interface {
	SetProcessInfo(string)
}

// recordSet wraps an executor, implements ast.RecordSet interface
type recordSet struct {
	fields      []*ast.ResultField
	executor    Executor
	stmt        *ExecStmt
	processinfo processinfoSetter
	lastErr     error
	txnStartTS  uint64
}

func (a *recordSet) Fields() []*ast.ResultField {
	if len(a.fields) == 0 {
		for _, col := range a.executor.Schema().Columns {
			dbName := col.DBName.O
			if dbName == "" && col.TblName.L != "" {
				dbName = a.stmt.ctx.GetSessionVars().CurrentDB
			}
			rf := &ast.ResultField{
				ColumnAsName: col.ColName,
				TableAsName:  col.TblName,
				DBName:       model.NewCIStr(dbName),
				Table:        &model.TableInfo{Name: col.OrigTblName},
				Column: &model.ColumnInfo{
					FieldType: *col.RetType,
					Name:      col.ColName,
				},
			}
			a.fields = append(a.fields, rf)
		}
	}
	return a.fields
}

func (a *recordSet) Next() (*ast.Row, error) {
	row, err := a.executor.Next()
	if err != nil {
		a.lastErr = err
		return nil, errors.Trace(err)
	}
	if row == nil {
		if a.stmt != nil {
			a.stmt.ctx.GetSessionVars().LastFoundRows = a.stmt.ctx.GetSessionVars().StmtCtx.FoundRows()
		}
		return nil, nil
	}

	if a.stmt != nil {
		a.stmt.ctx.GetSessionVars().StmtCtx.AddFoundRows(1)
	}
	return &ast.Row{Data: row}, nil
}

func (a *recordSet) Close() error {
	err := a.executor.Close()
	a.stmt.logSlowQuery(a.txnStartTS, a.lastErr == nil)
	if a.processinfo != nil {
		a.processinfo.SetProcessInfo("")
	}
	return errors.Trace(err)
}

// ExecStmt implements the ast.Statement interface, it builds a plan.Plan to an ast.Statement.
type ExecStmt struct {
	// InfoSchema stores a reference to the schema information.
	InfoSchema infoschema.InfoSchema
	// Plan stores a reference to the final physical plan.
	Plan plan.Plan
	// Expensive represents whether this query is an expensive one.
	Expensive bool
	// Cacheable represents whether the physical plan can be cached.
	Cacheable bool
	// Text represents the origin query text.
	Text string

	ctx            context.Context
	startTime      time.Time
	isPreparedStmt bool

	// ReadOnly represents the statement is read-only.
	ReadOnly bool
}

// OriginText implements ast.Statement interface.
func (a *ExecStmt) OriginText() string {
	return a.Text
}

// IsPrepared implements ast.Statement interface.
func (a *ExecStmt) IsPrepared() bool {
	return a.isPreparedStmt
}

// IsReadOnly implements ast.Statement interface.
func (a *ExecStmt) IsReadOnly() bool {
	return a.ReadOnly
}

// Exec implements the ast.Statement Exec interface.
// This function builds an Executor from a plan. If the Executor doesn't return result,
// like the INSERT, UPDATE statements, it executes in this function, if the Executor returns
// result, execution is done after this function returns, in the returned ast.RecordSet Next method.
func (a *ExecStmt) Exec(ctx context.Context) (ast.RecordSet, error) {
	a.startTime = time.Now()
	a.ctx = ctx

	if _, ok := a.Plan.(*plan.Analyze); ok && ctx.GetSessionVars().InRestrictedSQL {
		oriStats := ctx.GetSessionVars().Systems[variable.TiDBBuildStatsConcurrency]
		oriScan := ctx.GetSessionVars().DistSQLScanConcurrency
		oriIndex := ctx.GetSessionVars().IndexSerialScanConcurrency
		oriIso := ctx.GetSessionVars().Systems[variable.TxnIsolation]
		ctx.GetSessionVars().Systems[variable.TiDBBuildStatsConcurrency] = "1"
		ctx.GetSessionVars().DistSQLScanConcurrency = 1
		ctx.GetSessionVars().IndexSerialScanConcurrency = 1
		ctx.GetSessionVars().Systems[variable.TxnIsolation] = ast.ReadCommitted
		defer func() {
			ctx.GetSessionVars().Systems[variable.TiDBBuildStatsConcurrency] = oriStats
			ctx.GetSessionVars().DistSQLScanConcurrency = oriScan
			ctx.GetSessionVars().IndexSerialScanConcurrency = oriIndex
			ctx.GetSessionVars().Systems[variable.TxnIsolation] = oriIso
		}()
	}

	e, err := a.buildExecutor(ctx)
	if err != nil {
		return nil, errors.Trace(err)
	}

	if err := e.Open(); err != nil {
		return nil, errors.Trace(err)
	}

	var pi processinfoSetter
	if raw, ok := ctx.(processinfoSetter); ok {
		pi = raw
		sql := a.OriginText()
		if simple, ok := a.Plan.(*plan.Simple); ok && simple.Statement != nil {
			if ss, ok := simple.Statement.(ast.SensitiveStmtNode); ok {
				// Use SecureText to avoid leak password information.
				sql = ss.SecureText()
			}
		}
		// Update processinfo, ShowProcess() will use it.
		pi.SetProcessInfo(sql)
	}
	// Fields or Schema are only used for statements that return result set.
	if e.Schema().Len() == 0 {
		return a.handleNoDelayExecutor(e, ctx, pi)
	}

	return &recordSet{
		executor:    e,
		stmt:        a,
		processinfo: pi,
		txnStartTS:  ctx.Txn().StartTS(),
	}, nil
}

func (a *ExecStmt) handleNoDelayExecutor(e Executor, ctx context.Context, pi processinfoSetter) (ast.RecordSet, error) {
	// Check if "tidb_snapshot" is set for the write executors.
	// In history read mode, we can not do write operations.
	switch e.(type) {
	case *DeleteExec, *InsertExec, *UpdateExec, *ReplaceExec, *LoadData, *DDLExec:
		snapshotTS := ctx.GetSessionVars().SnapshotTS
		if snapshotTS != 0 {
			return nil, errors.New("can not execute write statement when 'tidb_snapshot' is set")
		}
	}

	var err error
	defer func() {
		if pi != nil {
			pi.SetProcessInfo("")
		}
		terror.Log(errors.Trace(e.Close()))
		txnTS := uint64(0)
		if ctx.Txn() != nil {
			txnTS = ctx.Txn().StartTS()
		}
		a.logSlowQuery(txnTS, err == nil)
	}()
	for {
		var row Row
		row, err = e.Next()
		if err != nil {
			return nil, errors.Trace(err)
		}
		// Even though there isn't any result set, the row is still used to indicate if there is
		// more work to do.
		// For example, the UPDATE statement updates a single row on a Next call, we keep calling Next until
		// There is no more rows to update.
		if row == nil {
			return nil, nil
		}
	}
}

// buildExecutor build a executor from plan, prepared statement may need additional procedure.
func (a *ExecStmt) buildExecutor(ctx context.Context) (Executor, error) {
	priority := kv.PriorityNormal
	if _, ok := a.Plan.(*plan.Execute); !ok {
		// Do not sync transaction for Execute statement, because the real optimization work is done in
		// "ExecuteExec.Build".
		var err error
		isPointGet := IsPointGetWithPKOrUniqueKeyByAutoCommit(ctx, a.Plan)
		if isPointGet {
			log.Debugf("[%d][InitTxnWithStartTS] %s", ctx.GetSessionVars().ConnectionID, a.Text)
			err = ctx.InitTxnWithStartTS(math.MaxUint64)
		} else {
			log.Debugf("[%d][ActivePendingTxn] %s", ctx.GetSessionVars().ConnectionID, a.Text)
			err = ctx.ActivePendingTxn()
		}
		if err != nil {
			return nil, errors.Trace(err)
		}

		if stmtPri := ctx.GetSessionVars().StmtCtx.Priority; stmtPri != mysql.NoPriority {
			priority = int(stmtPri)
		} else {
			switch {
			case isPointGet:
				priority = kv.PriorityHigh
			case a.Expensive:
				priority = kv.PriorityLow
			}
		}
	}
	if _, ok := a.Plan.(*plan.Analyze); ok && ctx.GetSessionVars().InRestrictedSQL {
		priority = kv.PriorityLow
	}

	b := newExecutorBuilder(ctx, a.InfoSchema, priority)
	e := b.build(a.Plan)
	if b.err != nil {
		return nil, errors.Trace(b.err)
	}

	// ExecuteExec is not a real Executor, we only use it to build another Executor from a prepared statement.
	if executorExec, ok := e.(*ExecuteExec); ok {
		err := executorExec.Build()
		if err != nil {
			return nil, errors.Trace(err)
		}
		a.Text = executorExec.Stmt.Text()
		a.isPreparedStmt = true
		a.Plan = executorExec.Plan
		e = executorExec.StmtExec
	}
	return e, nil
}

func (a *ExecStmt) logSlowQuery(txnTS uint64, succ bool) {
	cfg := config.GetGlobalConfig()
	costTime := time.Since(a.startTime)
	sql := a.Text
	if len(sql) > cfg.Log.QueryLogMaxLen {
		sql = fmt.Sprintf("%.*q(len:%d)", cfg.Log.QueryLogMaxLen, sql, len(a.Text))
	}
	connID := a.ctx.GetSessionVars().ConnectionID
	currentDB := a.ctx.GetSessionVars().CurrentDB
	logEntry := log.NewEntry(logutil.SlowQueryLogger)
	logEntry.Data = log.Fields{
		"connectionId": connID,
		"costTime":     costTime,
		"database":     currentDB,
		"sql":          sql,
		"txnStartTS":   txnTS,
	}
	if costTime < time.Duration(cfg.Log.SlowThreshold)*time.Millisecond {
		logEntry.WithField("type", "query").WithField("succ", succ).Debugf("query")
	} else {
		logEntry.WithField("type", "slow-query").WithField("succ", succ).Warnf("slow-query")
	}
}

// IsPointGetWithPKOrUniqueKeyByAutoCommit returns true when meets following conditions:
//  1. ctx is auto commit tagged
//  2. txn is nil
//  2. plan is point get by pk or unique key
func IsPointGetWithPKOrUniqueKeyByAutoCommit(ctx context.Context, p plan.Plan) bool {
	// check auto commit
	if !ctx.GetSessionVars().IsAutocommit() {
		return false
	}

	// check txn
	if ctx.Txn() != nil {
		return false
	}

	// check plan
	if proj, ok := p.(*plan.Projection); ok {
		if len(proj.Children()) != 1 {
			return false
		}
		p = proj.Children()[0]
	}

	switch v := p.(type) {
	case *plan.PhysicalIndexScan:
		return v.IsPointGetByUniqueKey(ctx.GetSessionVars().StmtCtx)
	case *plan.PhysicalIndexReader:
		indexScan := v.IndexPlans[0].(*plan.PhysicalIndexScan)
		return indexScan.IsPointGetByUniqueKey(ctx.GetSessionVars().StmtCtx)
	case *plan.PhysicalIndexLookUpReader:
		indexScan := v.IndexPlans[0].(*plan.PhysicalIndexScan)
		return indexScan.IsPointGetByUniqueKey(ctx.GetSessionVars().StmtCtx)
	case *plan.PhysicalTableScan:
		return len(v.Ranges) == 1 && v.Ranges[0].IsPoint()
	case *plan.PhysicalTableReader:
		tableScan := v.TablePlans[0].(*plan.PhysicalTableScan)
		return len(tableScan.Ranges) == 1 && tableScan.Ranges[0].IsPoint()
	default:
		return false
	}
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 sdb

import (
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestGet(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type GetTest struct {
	domainTest

	item          ItemName
	constistentRead bool
	names       []string

	attributes []Attribute
	err error
}

func (t *GetTest) SetUp(i *TestInfo) {
	// Call common setup code.
	t.domainTest.SetUp(i)

	// Make the request legal by default.
	t.item = "foo"
}

func (t *GetTest) callDomain() {
	t.attributes, t.err = t.domain.GetAttributes(t.item, t.constistentRead, t.names)
}

func init() { RegisterTestSuite(&GetTest{}) }

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *GetTest) ItemNameEmpty() {
	ExpectEq("TODO", "")
}

func (t *GetTest) ItemNameInvalid() {
	ExpectEq("TODO", "")
}

func (t *GetTest) OneAttributeNameEmpty() {
	ExpectEq("TODO", "")
}

func (t *GetTest) OneAttributeNameInvalid() {
	ExpectEq("TODO", "")
}

func (t *GetTest) InconsistentReadWithNoAttributeNames() {
	ExpectEq("TODO", "")
}

func (t *GetTest) ConsistentRead() {
	ExpectEq("TODO", "")
}

func (t *GetTest) SomeAttributeNames() {
	ExpectEq("TODO", "")
}

func (t *GetTest) ConnReturnsError() {
	ExpectEq("TODO", "")
}

func (t *GetTest) ConnReturnsJunk() {
	ExpectEq("TODO", "")
}

func (t *GetTest) NoAttributesInResponse() {
	ExpectEq("TODO", "")
}

func (t *GetTest) SomeAttributesInResponse() {
	ExpectEq("TODO", "")
}
<commit_msg>GetTest.ItemNameEmpty<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 sdb

import (
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestGet(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type GetTest struct {
	domainTest

	item          string
	constistentRead bool
	names       []string

	attributes []Attribute
	err error
}

func (t *GetTest) SetUp(i *TestInfo) {
	// Call common setup code.
	t.domainTest.SetUp(i)

	// Make the request legal by default.
	t.item = "foo"
}

func (t *GetTest) callDomain() {
	t.attributes, t.err = t.domain.GetAttributes(
		ItemName(t.item),
		t.constistentRead,
		t.names)
}

func init() { RegisterTestSuite(&GetTest{}) }

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *GetTest) ItemNameEmpty() {
	t.item = ""

	// Call
	t.callDomain()

	ExpectThat(t.err, Error(HasSubstr("item")))
	ExpectThat(t.err, Error(HasSubstr("name")))
	ExpectThat(t.err, Error(HasSubstr("empty")))
}

func (t *GetTest) ItemNameInvalid() {
	ExpectEq("TODO", "")
}

func (t *GetTest) OneAttributeNameEmpty() {
	ExpectEq("TODO", "")
}

func (t *GetTest) OneAttributeNameInvalid() {
	ExpectEq("TODO", "")
}

func (t *GetTest) InconsistentReadWithNoAttributeNames() {
	ExpectEq("TODO", "")
}

func (t *GetTest) ConsistentRead() {
	ExpectEq("TODO", "")
}

func (t *GetTest) SomeAttributeNames() {
	ExpectEq("TODO", "")
}

func (t *GetTest) ConnReturnsError() {
	ExpectEq("TODO", "")
}

func (t *GetTest) ConnReturnsJunk() {
	ExpectEq("TODO", "")
}

func (t *GetTest) NoAttributesInResponse() {
	ExpectEq("TODO", "")
}

func (t *GetTest) SomeAttributesInResponse() {
	ExpectEq("TODO", "")
}
<|endoftext|>


<commit_before>// Copyright 2013 The Go Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd.

// This file implements an http.Client with request timeouts set by command
// line flags.

package main

import (
	"flag"
	"log"
	"net"
	"net/http"
	"time"
)

var (
	dialTimeout    = flag.Duration("dial_timeout", 5*time.Second, "Timeout for dialing an HTTP connection.")
	requestTimeout = flag.Duration("request_timeout", 20*time.Second, "Time out for roundtripping an HTTP request.")
)

type timeoutConn struct {
	net.Conn
}

func (c timeoutConn) Read(p []byte) (int, error) {
	n, err := c.Conn.Read(p)
	c.Conn.SetReadDeadline(time.Time{})
	return n, err
}

func timeoutDial(network, addr string) (net.Conn, error) {
	c, err := net.DialTimeout(network, addr, *dialTimeout)
	if err != nil {
		return c, err
	}
	// The net/http transport CancelRequest feature does not work until after
	// the TLS handshake is complete. To help catch hangs during the TLS
	// handshake, we set a deadline on the connection here and clear the
	// deadline when the first read on the connection completes. This is not
	// perfect, but it does catch the case where the server accepts and ignores
	// a connection.
	c.SetDeadline(time.Now().Add(*requestTimeout))
	return timeoutConn{c}, nil
}

type transport struct {
	t http.Transport
}

func (t *transport) RoundTrip(req *http.Request) (*http.Response, error) {
	timer := time.AfterFunc(*requestTimeout, func() {
		t.t.CancelRequest(req)
		log.Printf("Canceled request for %s", req.URL)
	})
	defer timer.Stop()
	if req.URL.Host == "api.github.com" && gitHubCredentials != "" {
		if req.URL.RawQuery == "" {
			req.URL.RawQuery = gitHubCredentials
		} else {
			req.URL.RawQuery += "&" + gitHubCredentials
		}
	}
	if userAgent != "" {
		req.Header.Set("User-Agent", userAgent)
	}
	return t.t.RoundTrip(req)
}

var httpClient = &http.Client{Transport: &transport{
	t: http.Transport{
		Dial: timeoutDial,
		ResponseHeaderTimeout: *requestTimeout / 2,
	}}}
<commit_msg>Load proxy variables from the environment<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd.

// This file implements an http.Client with request timeouts set by command
// line flags.

package main

import (
	"flag"
	"log"
	"net"
	"net/http"
	"time"
)

var (
	dialTimeout    = flag.Duration("dial_timeout", 5*time.Second, "Timeout for dialing an HTTP connection.")
	requestTimeout = flag.Duration("request_timeout", 20*time.Second, "Time out for roundtripping an HTTP request.")
)

type timeoutConn struct {
	net.Conn
}

func (c timeoutConn) Read(p []byte) (int, error) {
	n, err := c.Conn.Read(p)
	c.Conn.SetReadDeadline(time.Time{})
	return n, err
}

func timeoutDial(network, addr string) (net.Conn, error) {
	c, err := net.DialTimeout(network, addr, *dialTimeout)
	if err != nil {
		return c, err
	}
	// The net/http transport CancelRequest feature does not work until after
	// the TLS handshake is complete. To help catch hangs during the TLS
	// handshake, we set a deadline on the connection here and clear the
	// deadline when the first read on the connection completes. This is not
	// perfect, but it does catch the case where the server accepts and ignores
	// a connection.
	c.SetDeadline(time.Now().Add(*requestTimeout))
	return timeoutConn{c}, nil
}

type transport struct {
	t http.Transport
}

func (t *transport) RoundTrip(req *http.Request) (*http.Response, error) {
	timer := time.AfterFunc(*requestTimeout, func() {
		t.t.CancelRequest(req)
		log.Printf("Canceled request for %s", req.URL)
	})
	defer timer.Stop()
	if req.URL.Host == "api.github.com" && gitHubCredentials != "" {
		if req.URL.RawQuery == "" {
			req.URL.RawQuery = gitHubCredentials
		} else {
			req.URL.RawQuery += "&" + gitHubCredentials
		}
	}
	if userAgent != "" {
		req.Header.Set("User-Agent", userAgent)
	}
	return t.t.RoundTrip(req)
}

var httpClient = &http.Client{Transport: &transport{
	t: http.Transport{
		Proxy: http.ProxyFromEnvironment,
		Dial:  timeoutDial,
		ResponseHeaderTimeout: *requestTimeout / 2,
	}}}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/qedus/osmpbf"
	"io"
	"log"
	"math"
	"os"
	"runtime"
	"strconv"
	"strings"
)

func main() {
	debug := false // switch to true for a bit of extra information

	// grok the command line args
	filename, dist, lat, lon, pattern := grok_args(os.Args)

	if debug {
		fmt.Printf("# d,lat,lon,'pattern': %.3f, %f,%f,'%s'\n", dist, lat, lon, pattern)
	}

	// start reading the file
	f, err := os.Open(filename)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	d := osmpbf.NewDecoder(f)
	err = d.Start(runtime.GOMAXPROCS(-1))
	if err != nil {
		log.Fatal(err)
	}

	var nc, wc, rc uint64
	for {
		if v, err := d.Decode(); err == io.EOF {
			break
		} else if err != nil {
			log.Fatal(err)
		} else {
			switch v := v.(type) {
			case *osmpbf.Node:
				handleNode(*v, dist, lat, lon, pattern)
				//handleNode(*v, lat, lon, dist)
				nc++
			case *osmpbf.Way:
				// Process Way v.
				wc++
			case *osmpbf.Relation:
				// Process Relation v.
				rc++
			default:
				log.Fatalf("unknown type %T\n", v)
			}
		}
	}
	if debug {
		fmt.Printf("# Nodes: %d, Ways: %d, Relations: %d\n", nc, wc, rc)
	}
}

func grok_args(args []string) (filename string, dist float64, lat float64, lon float64, pattern string) {

	var err error

	// default values
	filename = ""
	lat = 0.0
	lon = 0.0
	dist = -1.0 // negative distance means: no within-distance checking
	pattern = ""

	// show help if not enough args
	if len(os.Args) < 2 {
		fmt.Fprintf(os.Stderr, `Usage:

%s osm-file [-d max-dist lat lon] [pattern]

    eg. %s england-latest.osm.pbf mexican
    eg. %s central-america-latest.osm.pbf -d 10 12.1166 -68.9333 > willemstad10k.csv

The unit for maximum distance is km.
`, os.Args[0], os.Args[0], os.Args[0])
		os.Exit(1)
	}

	// first arg is filename
	filename = os.Args[1]
	if !fileExists(filename) {
		fmt.Fprintf(os.Stderr, "File does not exist: %s\n", os.Args[1])
		os.Exit(1)
	}

	// distance,lat,lon args
	pattern_idx := 2
	if args[2] == "-d" {
		dist, err = strconv.ParseFloat(args[3], 64)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Illegal value for max-distance: %s\n", os.Args[3])
			os.Exit(1)
		}

		lat, err = strconv.ParseFloat(args[4], 64)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Illegal value for latitude: %s\n", os.Args[4])
			os.Exit(1)
		}

		lon, err = strconv.ParseFloat(args[5], 64)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Illegal value for longitude: %s\n", os.Args[5])
			os.Exit(1)
		}
		pattern_idx = 6
	}

	// pattern arg
	if len(args) >= (pattern_idx + 1) {
		pattern = strings.ToLower(args[pattern_idx]) // SEARCH IN LOWER CASE!
	}
	return
}

/* approximately calculate the distance between 2 points
   from: http://www.movable-type.co.uk/scripts/latlong.html
   note: φ=lat λ=lon  in RADIANS!
   var x = (λ2-λ1) * Math.cos((φ1+φ2)/2);
   var y = (φ2-φ1);
   var d = Math.sqrt(x*x + y*y) * R;
*/
func rough_distance(lat1, lon1, lat2, lon2 float64) float64 {

	// convert to radians
	lat1 = lat1 * math.Pi / 180.0
	lon1 = lon1 * math.Pi / 180.0
	lat2 = lat2 * math.Pi / 180.0
	lon2 = lon2 * math.Pi / 180.0

	r := 6371.0 // km
	x := (lon2 - lon1) * math.Cos((lat1+lat2)/2)
	y := (lat2 - lat1)
	d := math.Sqrt(x*x+y*y) * r
	return d
}

func fileExists(name string) bool {
	if _, err := os.Stat(name); err != nil {
		if os.IsNotExist(err) {
			return false
		}
	}
	return true
}

func handleNode(nd osmpbf.Node, dist float64, lat float64, lon float64, contains_pattern string) {
	estim_distance := 0.0
	bingo_distance := true
	if dist >= 0.0 { // do we need to examine the distance?
		estim_distance = rough_distance(lat, lon, nd.Lat, nd.Lon)
		bingo_distance = (estim_distance < dist)
	}

	bingo_pattern := true
	if len(contains_pattern) > 0 { // do we need to examine the pattern?
		bingo_pattern = false
		for k, v := range nd.Tags {
			bingo_pattern = strings.Contains(strings.ToLower(k), contains_pattern) ||
				strings.Contains(strings.ToLower(v), contains_pattern)
			if bingo_pattern {
				break // out of the loop
			}
		}
	}

	if bingo_distance && bingo_pattern {
		// turn the Tags map into a k:v string
		tgs := ""
		for k, v := range nd.Tags {
			tgs = tgs + " " + k + ":" + v
		}

		if dist >= 0.0 {
			fmt.Printf("%f, %f, %s #,%.2f\n", nd.Lat, nd.Lon, tgs, estim_distance)
		} else {
			fmt.Printf("%f, %f, %s\n", nd.Lat, nd.Lon, tgs)
		}
	}
}
<commit_msg>drop empties, replace commas<commit_after>package main

import (
	"fmt"
	"github.com/qedus/osmpbf"
	"io"
	"log"
	"math"
	"os"
	"runtime"
	"strconv"
	"strings"
)

func main() {
	debug := false // switch to true for a bit of extra information

	// grok the command line args
	filename, dist, lat, lon, pattern := grok_args(os.Args)

	if debug {
		fmt.Printf("# d,lat,lon,'pattern': %.3f, %f,%f,'%s'\n", dist, lat, lon, pattern)
	}

	// start reading the file
	f, err := os.Open(filename)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()

	d := osmpbf.NewDecoder(f)
	err = d.Start(runtime.GOMAXPROCS(-1))
	if err != nil {
		log.Fatal(err)
	}

	var nc, wc, rc uint64
	for {
		if v, err := d.Decode(); err == io.EOF {
			break
		} else if err != nil {
			log.Fatal(err)
		} else {
			switch v := v.(type) {
			case *osmpbf.Node:
				handleNode(*v, dist, lat, lon, pattern)
				//handleNode(*v, lat, lon, dist)
				nc++
			case *osmpbf.Way:
				// Process Way v.
				wc++
			case *osmpbf.Relation:
				// Process Relation v.
				rc++
			default:
				log.Fatalf("unknown type %T\n", v)
			}
		}
	}
	if debug {
		fmt.Printf("# Nodes: %d, Ways: %d, Relations: %d\n", nc, wc, rc)
	}
}

func grok_args(args []string) (filename string, dist float64, lat float64, lon float64, pattern string) {

	var err error

	// default values
	filename = ""
	lat = 0.0
	lon = 0.0
	dist = -1.0 // negative distance means: no within-distance checking
	pattern = ""

	// show help if not enough args
	if len(os.Args) < 2 {
		fmt.Fprintf(os.Stderr, `Usage:

%s osm-file [-d max-dist lat lon] [pattern]

    eg. %s england-latest.osm.pbf mexican
    eg. %s central-america-latest.osm.pbf -d 10 12.1166 -68.9333 > willemstad10k.csv

The unit for maximum distance is km.
`, os.Args[0], os.Args[0], os.Args[0])
		os.Exit(1)
	}

	// first arg is filename
	filename = os.Args[1]
	if !fileExists(filename) {
		fmt.Fprintf(os.Stderr, "File does not exist: %s\n", os.Args[1])
		os.Exit(1)
	}

	// distance,lat,lon args
	pattern_idx := 2
	if args[2] == "-d" {
		dist, err = strconv.ParseFloat(args[3], 64)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Illegal value for max-distance: %s\n", os.Args[3])
			os.Exit(1)
		}

		lat, err = strconv.ParseFloat(args[4], 64)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Illegal value for latitude: %s\n", os.Args[4])
			os.Exit(1)
		}

		lon, err = strconv.ParseFloat(args[5], 64)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Illegal value for longitude: %s\n", os.Args[5])
			os.Exit(1)
		}
		pattern_idx = 6
	}

	// pattern arg
	if len(args) >= (pattern_idx + 1) {
		pattern = strings.ToLower(args[pattern_idx]) // SEARCH IN LOWER CASE!
	}
	return
}

/* approximately calculate the distance between 2 points
   from: http://www.movable-type.co.uk/scripts/latlong.html
   note: φ=lat λ=lon  in RADIANS!
   var x = (λ2-λ1) * Math.cos((φ1+φ2)/2);
   var y = (φ2-φ1);
   var d = Math.sqrt(x*x + y*y) * R;
*/
func rough_distance(lat1, lon1, lat2, lon2 float64) float64 {

	// convert to radians
	lat1 = lat1 * math.Pi / 180.0
	lon1 = lon1 * math.Pi / 180.0
	lat2 = lat2 * math.Pi / 180.0
	lon2 = lon2 * math.Pi / 180.0

	r := 6371.0 // km
	x := (lon2 - lon1) * math.Cos((lat1+lat2)/2)
	y := (lat2 - lat1)
	d := math.Sqrt(x*x+y*y) * r
	return d
}

func fileExists(name string) bool {
	if _, err := os.Stat(name); err != nil {
		if os.IsNotExist(err) {
			return false
		}
	}
	return true
}

func handleNode(nd osmpbf.Node, dist float64, lat float64, lon float64, contains_pattern string) {
	estim_distance := 0.0
	bingo_distance := true
	if dist >= 0.0 { // do we need to examine the distance?
		estim_distance = rough_distance(lat, lon, nd.Lat, nd.Lon)
		bingo_distance = (estim_distance < dist)
	}

	bingo_pattern := true
	if len(contains_pattern) > 0 { // do we need to examine the pattern?
		bingo_pattern = false
		for k, v := range nd.Tags {
			bingo_pattern = strings.Contains(strings.ToLower(k), contains_pattern) ||
				strings.Contains(strings.ToLower(v), contains_pattern)
			if bingo_pattern {
				break // out of the loop
			}
		}
	}

	if bingo_distance && bingo_pattern {
		// turn the Tags map into a k:v string
		tgs := ""
		for k, v := range nd.Tags {
			tgs = tgs + " " + k + ":" + v
		}
		desc := strings.TrimSpace(tgs)

		// if we don't have tags/description, then don't print
		if len(desc) == 0 {
			return
		}

		// replace the comma's in the description
		desc = strings.Replace(desc, ",", "", -1)

		if dist >= 0.0 {
			fmt.Printf("%f, %f, %s #,%.2f\n", nd.Lat, nd.Lon, desc, estim_distance)
		} else {
			fmt.Printf("%f, %f, %s\n", nd.Lat, nd.Lon, desc)
		}
	}
}
<|endoftext|>


<commit_before>package cloud

import (
	"bytes"
	"text/template"

	api "github.com/appscode/pharmer/apis/v1alpha1"
	"github.com/ghodss/yaml"
	"github.com/hashicorp/go-version"
	kubeadmapi "k8s.io/kubernetes/cmd/kubeadm/app/apis/kubeadm/v1alpha1"
)

type TemplateData struct {
	KubernetesVersion   string
	KubeadmVersion      string
	KubeadmToken        string
	CAKey               string
	FrontProxyKey       string
	APIServerAddress    string
	APIBindPort         int32
	ExtraDomains        string
	NetworkProvider     string
	CloudConfig         string
	Provider            string
	ExternalProvider    bool
	ConfigurationBucket string

	MasterConfiguration *kubeadmapi.MasterConfiguration
	KubeletExtraArgs    map[string]string
}

func (td TemplateData) MasterConfigurationYAML() (string, error) {
	if td.MasterConfiguration == nil {
		return "", nil
	}
	cb, err := yaml.Marshal(td.MasterConfiguration)
	return string(cb), err
}

func (td TemplateData) IsPreReleaseVersion() bool {
	if v, err := version.NewVersion(td.KubeadmVersion); err == nil && v.Prerelease() != "" {
		return true
	}
	return false
}

func (td TemplateData) KubeletExtraArgsStr() string {
	var buf bytes.Buffer
	for k, v := range td.KubeletExtraArgs {
		buf.WriteString("--")
		buf.WriteString(k)
		buf.WriteRune('=')
		buf.WriteString(v)
		buf.WriteRune(' ')
	}
	return buf.String()
}

func (td TemplateData) KubeletExtraArgsWithoutCloudProviderStr() string {
	var buf bytes.Buffer
	for k, v := range td.KubeletExtraArgs {
		if k == "cloud-config" {
			continue
		}
		if k == "cloud-provider" {
			v = ""
		}
		buf.WriteString("--")
		buf.WriteString(k)
		buf.WriteRune('=')
		buf.WriteString(v)
		buf.WriteRune(' ')
	}
	return buf.String()
}

var (
	StartupScriptTemplate = template.Must(template.New(api.RoleMaster).Parse(`#!/bin/bash
set -x
set -o errexit
set -o nounset
set -o pipefail

# log to /var/log/startup-script.log
exec > >(tee -a /var/log/startup-script.log)
exec 2>&1

# kill apt processes (E: Unable to lock directory /var/lib/apt/lists/)
kill $(ps aux | grep '[a]pt' | awk '{print $2}') || true

{{ template "prepare-host" . }}

apt-get update -y
apt-get install -y apt-transport-https curl ca-certificates

curl -fSsL https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
echo 'deb http://apt.kubernetes.io/ kubernetes-xenial main' > /etc/apt/sources.list.d/kubernetes.list

add-apt-repository -y ppa:gluster/glusterfs-3.10

apt-get update -y
apt-get install -y \
	socat \
	ebtables \
	git \
	haveged \
	nfs-common \
	cron \
	glusterfs-client \
	kubectl \
	kubelet \
	{{ if not .IsPreReleaseVersion }}kubeadm{{ if .KubeadmVersion }}={{ .KubeadmVersion }}{{ end }}{{ end }} \
	cloud-utils \
	docker.io || true

{{ if .IsPreReleaseVersion }}
curl -Lo kubeadm https://dl.k8s.io/release/{{ .KubeadmVersion }}/bin/linux/amd64/kubeadm \
    && chmod +x kubeadm \
	&& mv kubeadm /usr/bin/
{{ end }}

curl -Lo pre-k https://cdn.appscode.com/binaries/pre-k/0.1.0-alpha.5/pre-k-linux-amd64 \
	&& chmod +x pre-k \
	&& mv pre-k /usr/bin/

systemctl enable docker
systemctl start docker

cat > /etc/systemd/system/kubelet.service.d/20-pharmer.conf <<EOF
[Service]
Environment="KUBELET_EXTRA_ARGS={{ if .ExternalProvider }}{{ .KubeletExtraArgsWithoutCloudProviderStr }}{{ else }}{{ .KubeletExtraArgsStr }}{{ end }}"
EOF

systemctl daemon-reload
systemctl restart kubelet

kubeadm reset

{{ template "setup-certs" . }}

{{ if .CloudConfig }}
cat > /etc/kubernetes/cloud-config <<EOF
{{ .CloudConfig }}
EOF
{{ end }}

mkdir -p /etc/kubernetes/kubeadm

{{ if .MasterConfiguration }}
cat > /etc/kubernetes/kubeadm/config.yaml <<EOF
{{ .MasterConfigurationYAML }}
EOF
{{ end }}

pre-k merge master-config \
	--config=/etc/kubernetes/kubeadm/config.yaml \
	--apiserver-bind-port={{ .APIBindPort }} \
	--apiserver-advertise-address=$(pre-k get public-ips --all=false) \
	--apiserver-cert-extra-sans=$(pre-k get public-ips --routable) \
	--apiserver-cert-extra-sans=$(pre-k get private-ips) \
	--apiserver-cert-extra-sans={{ .ExtraDomains }} \
	--kubernetes-version={{ .KubernetesVersion }} \
	> /etc/kubernetes/kubeadm/config.yaml
kubeadm init --config=/etc/kubernetes/kubeadm/config.yaml --skip-token-print

{{ if eq .NetworkProvider "flannel" }}
{{ template "flannel" . }}
{{ else if eq .NetworkProvider "calico" }}
{{ template "calico" . }}
{{ end }}

mkdir -p ~/.kube
sudo cp -i /etc/kubernetes/admin.conf ~/.kube/config
sudo chown $(id -u):$(id -g) ~/.kube/config

{{ if .ExternalProvider }}
{{ template "ccm" . }}
{{end}}

kubectl apply \
  -f https://raw.githubusercontent.com/appscode/pharmer/master/addons/kubeadm-probe/ds.yaml \
  --kubeconfig /etc/kubernetes/admin.conf
`))

	_ = template.Must(StartupScriptTemplate.New(api.RoleNode).Parse(`#!/bin/bash
set -x
set -o errexit
set -o nounset
set -o pipefail

# log to /var/log/startup-script.log
exec > >(tee -a /var/log/startup-script.log)
exec 2>&1

# kill apt processes (E: Unable to lock directory /var/lib/apt/lists/)
kill $(ps aux | grep '[a]pt' | awk '{print $2}') || true

{{ template "prepare-host" . }}

apt-get update -y
apt-get install -y apt-transport-https curl ca-certificates

curl -fSsL https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
echo 'deb http://apt.kubernetes.io/ kubernetes-xenial main' > /etc/apt/sources.list.d/kubernetes.list

add-apt-repository -y ppa:gluster/glusterfs-3.10

apt-get update -y
apt-get install -y \
	socat \
	ebtables \
	git \
	haveged \
	nfs-common \
	cron \
	glusterfs-client \
	kubelet \
	kubectl \
	{{ if not .IsPreReleaseVersion }}kubeadm{{ if .KubeadmVersion }}={{ .KubeadmVersion }}{{ end }}{{ end }} \
	docker.io || true

{{ if .IsPreReleaseVersion }}
curl -Lo kubeadm https://dl.k8s.io/release/{{ .KubeadmVersion }}/bin/linux/amd64/kubeadm \
    && chmod +x kubeadm \
	&& mv kubeadm /usr/bin/
{{ end }}

systemctl enable docker
systemctl start docker

{{ if .CloudConfig }}
cat > /etc/kubernetes/cloud-config <<EOF
{{ .CloudConfig }}
EOF
{{ end }}

cat > /etc/systemd/system/kubelet.service.d/20-pharmer.conf <<EOF
[Service]
Environment="KUBELET_EXTRA_ARGS={{ .KubeletExtraArgsStr }}"
EOF

systemctl daemon-reload
systemctl restart kubelet

kubeadm reset
{{ if .ConfigurationBucket }}
 {{ .ConfigurationBucket  }}
 source /etc/kubernetes/config.sh
 kubeadm join --token=${KUBEADM_TOKEN} {{ .APIServerAddress }}
{{ else }}
 kubeadm join --token={{ .KubeadmToken }} {{ .APIServerAddress }}
{{ end }}
`))

	_ = template.Must(StartupScriptTemplate.New("prepare-host").Parse(``))

	_ = template.Must(StartupScriptTemplate.New("setup-certs").Parse(`
mkdir -p /etc/kubernetes/pki

cat > /etc/kubernetes/pki/ca.key <<EOF
{{ .CAKey }}
EOF
pre-k get cacert --common-name=ca < /etc/kubernetes/pki/ca.key > /etc/kubernetes/pki/ca.crt

cat > /etc/kubernetes/pki/front-proxy-ca.key <<EOF
{{ .FrontProxyKey }}
EOF
pre-k get cacert --common-name=front-proxy-ca < /etc/kubernetes/pki/front-proxy-ca.key > /etc/kubernetes/pki/front-proxy-ca.crt

chmod 600 /etc/kubernetes/pki/ca.key /etc/kubernetes/pki/front-proxy-ca.key
`))

	_ = template.Must(StartupScriptTemplate.New("ccm").Parse(`
until [ $(kubectl get pods -n kube-system -l k8s-app=kube-dns -o jsonpath='{.items[0].status.phase}' --kubeconfig /etc/kubernetes/admin.conf) == "Running" ]
do
   echo '.'
   sleep 5
done

kubectl apply -f "https://raw.githubusercontent.com/appscode/pharmer/master/cloud/providers/{{ .Provider }}/cloud-control-manager.yaml" --kubeconfig /etc/kubernetes/admin.conf

until [ $(kubectl get pods -n kube-system -l app=cloud-controller-manager -o jsonpath='{.items[0].status.phase}' --kubeconfig /etc/kubernetes/admin.conf) == "Running" ]
do
   echo '.'
   sleep 5
done

cat > /etc/systemd/system/kubelet.service.d/20-pharmer.conf <<EOF
[Service]
Environment="KUBELET_EXTRA_ARGS={{ .KubeletExtraArgsStr }}"
EOF

NODE_NAME=$(uname -n)
kubectl taint nodes ${NODE_NAME} node.cloudprovider.kubernetes.io/uninitialized=true:NoSchedule --kubeconfig /etc/kubernetes/admin.conf

systemctl daemon-reload
systemctl restart kubelet

sleep 10
reboot
`))

	_ = template.Must(StartupScriptTemplate.New("calico").Parse(`
kubectl apply \
  -f http://docs.projectcalico.org/v2.3/getting-started/kubernetes/installation/hosted/kubeadm/1.6/calico.yaml \
  --kubeconfig /etc/kubernetes/admin.conf
`))

	_ = template.Must(StartupScriptTemplate.New("flannel").Parse(`
kubectl apply \
  -f https://raw.githubusercontent.com/coreos/flannel/v0.8.0/Documentation/kube-flannel.yml \
  --kubeconfig /etc/kubernetes/admin.conf
kubectl apply \
  -f https://raw.githubusercontent.com/coreos/flannel/v0.8.0/Documentation/kube-flannel-rbac.yml \
  --kubeconfig /etc/kubernetes/admin.conf
`))
)
<commit_msg>Install kubeadm-probe befor CCM<commit_after>package cloud

import (
	"bytes"
	"text/template"

	api "github.com/appscode/pharmer/apis/v1alpha1"
	"github.com/ghodss/yaml"
	"github.com/hashicorp/go-version"
	kubeadmapi "k8s.io/kubernetes/cmd/kubeadm/app/apis/kubeadm/v1alpha1"
)

type TemplateData struct {
	KubernetesVersion   string
	KubeadmVersion      string
	KubeadmToken        string
	CAKey               string
	FrontProxyKey       string
	APIServerAddress    string
	APIBindPort         int32
	ExtraDomains        string
	NetworkProvider     string
	CloudConfig         string
	Provider            string
	ExternalProvider    bool
	ConfigurationBucket string

	MasterConfiguration *kubeadmapi.MasterConfiguration
	KubeletExtraArgs    map[string]string
}

func (td TemplateData) MasterConfigurationYAML() (string, error) {
	if td.MasterConfiguration == nil {
		return "", nil
	}
	cb, err := yaml.Marshal(td.MasterConfiguration)
	return string(cb), err
}

func (td TemplateData) IsPreReleaseVersion() bool {
	if v, err := version.NewVersion(td.KubeadmVersion); err == nil && v.Prerelease() != "" {
		return true
	}
	return false
}

func (td TemplateData) KubeletExtraArgsStr() string {
	var buf bytes.Buffer
	for k, v := range td.KubeletExtraArgs {
		buf.WriteString("--")
		buf.WriteString(k)
		buf.WriteRune('=')
		buf.WriteString(v)
		buf.WriteRune(' ')
	}
	return buf.String()
}

func (td TemplateData) KubeletExtraArgsWithoutCloudProviderStr() string {
	var buf bytes.Buffer
	for k, v := range td.KubeletExtraArgs {
		if k == "cloud-config" {
			continue
		}
		if k == "cloud-provider" {
			v = ""
		}
		buf.WriteString("--")
		buf.WriteString(k)
		buf.WriteRune('=')
		buf.WriteString(v)
		buf.WriteRune(' ')
	}
	return buf.String()
}

var (
	StartupScriptTemplate = template.Must(template.New(api.RoleMaster).Parse(`#!/bin/bash
set -x
set -o errexit
set -o nounset
set -o pipefail

# log to /var/log/startup-script.log
exec > >(tee -a /var/log/startup-script.log)
exec 2>&1

# kill apt processes (E: Unable to lock directory /var/lib/apt/lists/)
kill $(ps aux | grep '[a]pt' | awk '{print $2}') || true

{{ template "prepare-host" . }}

apt-get update -y
apt-get install -y apt-transport-https curl ca-certificates

curl -fSsL https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
echo 'deb http://apt.kubernetes.io/ kubernetes-xenial main' > /etc/apt/sources.list.d/kubernetes.list

add-apt-repository -y ppa:gluster/glusterfs-3.10

apt-get update -y
apt-get install -y \
	socat \
	ebtables \
	git \
	haveged \
	nfs-common \
	cron \
	glusterfs-client \
	kubectl \
	kubelet \
	{{ if not .IsPreReleaseVersion }}kubeadm{{ if .KubeadmVersion }}={{ .KubeadmVersion }}{{ end }}{{ end }} \
	cloud-utils \
	docker.io || true

{{ if .IsPreReleaseVersion }}
curl -Lo kubeadm https://dl.k8s.io/release/{{ .KubeadmVersion }}/bin/linux/amd64/kubeadm \
    && chmod +x kubeadm \
	&& mv kubeadm /usr/bin/
{{ end }}

curl -Lo pre-k https://cdn.appscode.com/binaries/pre-k/0.1.0-alpha.5/pre-k-linux-amd64 \
	&& chmod +x pre-k \
	&& mv pre-k /usr/bin/

systemctl enable docker
systemctl start docker

cat > /etc/systemd/system/kubelet.service.d/20-pharmer.conf <<EOF
[Service]
Environment="KUBELET_EXTRA_ARGS={{ if .ExternalProvider }}{{ .KubeletExtraArgsWithoutCloudProviderStr }}{{ else }}{{ .KubeletExtraArgsStr }}{{ end }}"
EOF

systemctl daemon-reload
systemctl restart kubelet

kubeadm reset

{{ template "setup-certs" . }}

{{ if .CloudConfig }}
cat > /etc/kubernetes/cloud-config <<EOF
{{ .CloudConfig }}
EOF
{{ end }}

mkdir -p /etc/kubernetes/kubeadm

{{ if .MasterConfiguration }}
cat > /etc/kubernetes/kubeadm/config.yaml <<EOF
{{ .MasterConfigurationYAML }}
EOF
{{ end }}

pre-k merge master-config \
	--config=/etc/kubernetes/kubeadm/config.yaml \
	--apiserver-bind-port={{ .APIBindPort }} \
	--apiserver-advertise-address=$(pre-k get public-ips --all=false) \
	--apiserver-cert-extra-sans=$(pre-k get public-ips --routable) \
	--apiserver-cert-extra-sans=$(pre-k get private-ips) \
	--apiserver-cert-extra-sans={{ .ExtraDomains }} \
	--kubernetes-version={{ .KubernetesVersion }} \
	> /etc/kubernetes/kubeadm/config.yaml
kubeadm init --config=/etc/kubernetes/kubeadm/config.yaml --skip-token-print

{{ if eq .NetworkProvider "flannel" }}
{{ template "flannel" . }}
{{ else if eq .NetworkProvider "calico" }}
{{ template "calico" . }}
{{ end }}

kubectl apply \
  -f https://raw.githubusercontent.com/appscode/pharmer/master/addons/kubeadm-probe/ds.yaml \
  --kubeconfig /etc/kubernetes/admin.conf

mkdir -p ~/.kube
sudo cp -i /etc/kubernetes/admin.conf ~/.kube/config
sudo chown $(id -u):$(id -g) ~/.kube/config

{{ if .ExternalProvider }}
{{ template "ccm" . }}
{{end}}

`))

	_ = template.Must(StartupScriptTemplate.New(api.RoleNode).Parse(`#!/bin/bash
set -x
set -o errexit
set -o nounset
set -o pipefail

# log to /var/log/startup-script.log
exec > >(tee -a /var/log/startup-script.log)
exec 2>&1

# kill apt processes (E: Unable to lock directory /var/lib/apt/lists/)
kill $(ps aux | grep '[a]pt' | awk '{print $2}') || true

{{ template "prepare-host" . }}

apt-get update -y
apt-get install -y apt-transport-https curl ca-certificates

curl -fSsL https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
echo 'deb http://apt.kubernetes.io/ kubernetes-xenial main' > /etc/apt/sources.list.d/kubernetes.list

add-apt-repository -y ppa:gluster/glusterfs-3.10

apt-get update -y
apt-get install -y \
	socat \
	ebtables \
	git \
	haveged \
	nfs-common \
	cron \
	glusterfs-client \
	kubelet \
	kubectl \
	{{ if not .IsPreReleaseVersion }}kubeadm{{ if .KubeadmVersion }}={{ .KubeadmVersion }}{{ end }}{{ end }} \
	docker.io || true

{{ if .IsPreReleaseVersion }}
curl -Lo kubeadm https://dl.k8s.io/release/{{ .KubeadmVersion }}/bin/linux/amd64/kubeadm \
    && chmod +x kubeadm \
	&& mv kubeadm /usr/bin/
{{ end }}

systemctl enable docker
systemctl start docker

{{ if .CloudConfig }}
cat > /etc/kubernetes/cloud-config <<EOF
{{ .CloudConfig }}
EOF
{{ end }}

cat > /etc/systemd/system/kubelet.service.d/20-pharmer.conf <<EOF
[Service]
Environment="KUBELET_EXTRA_ARGS={{ .KubeletExtraArgsStr }}"
EOF

systemctl daemon-reload
systemctl restart kubelet

kubeadm reset
{{ if .ConfigurationBucket }}
 {{ .ConfigurationBucket  }}
 source /etc/kubernetes/config.sh
 kubeadm join --token=${KUBEADM_TOKEN} {{ .APIServerAddress }}
{{ else }}
 kubeadm join --token={{ .KubeadmToken }} {{ .APIServerAddress }}
{{ end }}
`))

	_ = template.Must(StartupScriptTemplate.New("prepare-host").Parse(``))

	_ = template.Must(StartupScriptTemplate.New("setup-certs").Parse(`
mkdir -p /etc/kubernetes/pki

cat > /etc/kubernetes/pki/ca.key <<EOF
{{ .CAKey }}
EOF
pre-k get cacert --common-name=ca < /etc/kubernetes/pki/ca.key > /etc/kubernetes/pki/ca.crt

cat > /etc/kubernetes/pki/front-proxy-ca.key <<EOF
{{ .FrontProxyKey }}
EOF
pre-k get cacert --common-name=front-proxy-ca < /etc/kubernetes/pki/front-proxy-ca.key > /etc/kubernetes/pki/front-proxy-ca.crt

chmod 600 /etc/kubernetes/pki/ca.key /etc/kubernetes/pki/front-proxy-ca.key
`))

	_ = template.Must(StartupScriptTemplate.New("ccm").Parse(`
until [ $(kubectl get pods -n kube-system -l k8s-app=kube-dns -o jsonpath='{.items[0].status.phase}' --kubeconfig /etc/kubernetes/admin.conf) == "Running" ]
do
   echo '.'
   sleep 5
done

kubectl apply -f "https://raw.githubusercontent.com/appscode/pharmer/master/cloud/providers/{{ .Provider }}/cloud-control-manager.yaml" --kubeconfig /etc/kubernetes/admin.conf

until [ $(kubectl get pods -n kube-system -l app=cloud-controller-manager -o jsonpath='{.items[0].status.phase}' --kubeconfig /etc/kubernetes/admin.conf) == "Running" ]
do
   echo '.'
   sleep 5
done

cat > /etc/systemd/system/kubelet.service.d/20-pharmer.conf <<EOF
[Service]
Environment="KUBELET_EXTRA_ARGS={{ .KubeletExtraArgsStr }}"
EOF

NODE_NAME=$(uname -n)
kubectl taint nodes ${NODE_NAME} node.cloudprovider.kubernetes.io/uninitialized=true:NoSchedule --kubeconfig /etc/kubernetes/admin.conf

systemctl daemon-reload
systemctl restart kubelet

# sleep 10
# reboot
`))

	_ = template.Must(StartupScriptTemplate.New("calico").Parse(`
kubectl apply \
  -f http://docs.projectcalico.org/v2.3/getting-started/kubernetes/installation/hosted/kubeadm/1.6/calico.yaml \
  --kubeconfig /etc/kubernetes/admin.conf
`))

	_ = template.Must(StartupScriptTemplate.New("flannel").Parse(`
kubectl apply \
  -f https://raw.githubusercontent.com/coreos/flannel/v0.8.0/Documentation/kube-flannel.yml \
  --kubeconfig /etc/kubernetes/admin.conf
kubectl apply \
  -f https://raw.githubusercontent.com/coreos/flannel/v0.8.0/Documentation/kube-flannel-rbac.yml \
  --kubeconfig /etc/kubernetes/admin.conf
`))
)
<|endoftext|>


<commit_before>package tools

import (
	"gopkg.in/mgo.v2"
	"fmt"
	"errors"
)

type SessionConf struct {
	Hosts    []string
	Database string
}

var (
	session *mgo.Session
)

func InitSession(conf SessionConf) *mgo.Session {
	var err error
	dialInfo := mgo.DialInfo{
		Addrs:    conf.Hosts,
		Database: conf.Database,
	}

	session, err = mgo.DialWithInfo(&dialInfo)

	if err != nil {
		panic(fmt.Sprintf("Failed to connect to DB server. %s", err))
	}
	return session
}

func GetSession() (*mgo.Session, error) {
	var err error
	if session == nil {
		err = errors.New("Session is not initialized")
	}
	return session, err
}
<commit_msg>Make SessionConf into shortcut type for mgo.DialInfo<commit_after>package tools

import (
	"gopkg.in/mgo.v2"
	"fmt"
	"errors"
)

type SessionConf mgo.DialInfo

var (
	session *mgo.Session
)

func InitSession(conf *SessionConf) *mgo.Session {
	var err error
	session, err = mgo.DialWithInfo((*mgo.DialInfo)(conf))

	if err != nil {
		panic(fmt.Sprintf("Failed to connect to DB server. %s", err))
	}
	return session
}

func GetSession() (*mgo.Session, error) {
	var err error
	if session == nil {
		err = errors.New("Session is not initialized")
	}
	return session, err
}
<|endoftext|>