<commit_before>package main

import (
	"bufio"
	"errors"
	"io"
	"strings"

	"golang.org/x/crypto/ssh"
	"golang.org/x/term"
)

type ptyRequestPayload struct {
	Term                                   string
	Width, Height, PixelWidth, PixelHeight uint32
	Modes                                  string
}

func (request ptyRequestPayload) reply() []byte {
	return nil
}
func (request ptyRequestPayload) logEntry(channelID int) logEntry {
	return ptyLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Terminal: request.Term,
		Width:    request.Width,
		Height:   request.Height,
	}
}

type shellRequestPayload struct{}

func (request shellRequestPayload) reply() []byte {
	return nil
}
func (request shellRequestPayload) logEntry(channelID int) logEntry {
	return shellLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
	}
}

type x11RequestPayload struct {
	SingleConnection         bool
	AuthProtocol, AuthCookie string
	ScreenNumber             uint32
}

func (request x11RequestPayload) reply() []byte {
	return nil
}
func (request x11RequestPayload) logEntry(channelID int) logEntry {
	return x11Log{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Screen: request.ScreenNumber,
	}
}

type envRequestPayload struct {
	Name, Value string
}

func (request envRequestPayload) reply() []byte {
	return nil
}
func (request envRequestPayload) logEntry(channelID int) logEntry {
	return envLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Name:  request.Name,
		Value: request.Value,
	}
}

type execRequestPayload struct {
	Command string
}

func (request execRequestPayload) reply() []byte {
	return nil
}
func (request execRequestPayload) logEntry(channelID int) logEntry {
	return execLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Command: request.Command,
	}
}

type subsystemRequestPayload struct {
	Subsystem string
}

func (request subsystemRequestPayload) reply() []byte {
	return nil
}
func (request subsystemRequestPayload) logEntry(channelID int) logEntry {
	return subsystemLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Subsystem: request.Subsystem,
	}
}

type windowChangeRequestPayload struct {
	Width, Height, PixelWidth, PixelHeight uint32
}

func (request windowChangeRequestPayload) reply() []byte {
	return nil
}
func (request windowChangeRequestPayload) logEntry(channelID int) logEntry {
	return windowChangeLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Width:  request.Width,
		Height: request.Height,
	}
}

type sessionContext struct {
	channelContext
	ssh.Channel
	inputChan chan string
	errorChan chan error
	active    bool
	pty       bool
}

type scannerReadLiner struct {
	scanner   *bufio.Scanner
	inputChan chan<- string
}

func (r scannerReadLiner) ReadLine() (string, error) {
	if !r.scanner.Scan() {
		if err := r.scanner.Err(); err != nil {
			return "", err
		}
		return "", io.EOF
	}
	line := r.scanner.Text()
	r.inputChan <- line
	return line, nil
}

type terminalReadLiner struct {
	terminal  *term.Terminal
	inputChan chan<- string
}

type clientEOFError struct{}

var clientEOF = clientEOFError{}

func (clientEOFError) Error() string {
	return "Client EOF"
}

func (r terminalReadLiner) ReadLine() (string, error) {
	line, err := r.terminal.ReadLine()
	if err == nil || line != "" {
		r.inputChan <- line
	}
	if err == io.EOF {
		return line, clientEOF
	}
	return line, err
}

func (context *sessionContext) handleProgram(program []string) {
	context.active = true
	var stdin readLiner
	var stdout, stderr io.Writer
	if context.pty {
		terminal := term.NewTerminal(context, "")
		stdin = terminalReadLiner{terminal, context.inputChan}
		stdout = terminal
		stderr = terminal
	} else {
		stdin = scannerReadLiner{bufio.NewScanner(context), context.inputChan}
		stdout = context
		stderr = context.Stderr()
	}
	go func() {
		defer close(context.inputChan)
		defer close(context.errorChan)
		result, err := executeProgram(commandContext{program, stdin, stdout, stderr, context.pty, context.User()})
		shouldSendEOW := err == nil || (context.pty && err == clientEOF)
		shouldSendCRLF := context.pty && err == clientEOF
		if err == clientEOF || err == io.EOF {
			err = nil
		}
		if err == nil && shouldSendCRLF {
			_, err = context.Write([]byte("\r\n"))
		}
		if err == nil {
			_, err = context.SendRequest("exit-status", false, ssh.Marshal(struct {
				ExitStatus uint32
			}{result}))
		}
		if err == nil && shouldSendEOW {
			_, err = context.SendRequest("eow@openssh.com", false, nil)
		}
		if err == nil {
			err = context.CloseWrite()
		}
		if err == nil {
			err = context.Close()
		}
		context.errorChan <- err
	}()
}

func (context *sessionContext) handleRequest(request *ssh.Request) error {
	switch request.Type {
	case "pty-req":
		if !context.active {
			if context.pty {
				return errors.New("a pty is already requested")
			}
			payload := &ptyRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.pty = true
			return nil
		}
	case "shell":
		if !context.active {
			if len(request.Payload) != 0 {
				return errors.New("invalid request payload")
			}
			payload := &shellRequestPayload{}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.active = true
			context.handleProgram(shellProgram)
			return nil
		}
	case "x11-req":
		if !context.active {
			payload := &x11RequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			return request.Reply(true, payload.reply())
		}
	case "env":
		if !context.active {
			payload := &envRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			return request.Reply(true, payload.reply())
		}
	case "exec":
		if !context.active {
			payload := &execRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.active = true
			context.handleProgram(strings.Fields(payload.Command))
			return nil
		}
	case "subsystem":
		if !context.active {
			payload := &subsystemRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.active = true
			context.handleProgram(strings.Fields(payload.Subsystem))
		}
	case "window-change":
		payload := &windowChangeRequestPayload{}
		if err := ssh.Unmarshal(request.Payload, payload); err != nil {
			return err
		}
		context.logEvent(payload.logEntry(context.channelID))
		return request.Reply(true, payload.reply())
	}
	warningLogger.Printf("Rejected session request: %s", request.Type)
	return request.Reply(false, nil)
}

func handleSessionChannel(newChannel ssh.NewChannel, context channelContext) error {
	if context.noMoreSessions {
		return errors.New("nore more sessions were supposed to be requested")
	}
	if len(newChannel.ExtraData()) != 0 {
		return errors.New("invalid channel data")
	}
	channel, requests, err := newChannel.Accept()
	if err != nil {
		return err
	}
	context.logEvent(sessionLog{
		channelLog: channelLog{
			ChannelID: context.channelID,
		},
	})
	defer context.logEvent(sessionCloseLog{
		channelLog: channelLog{
			ChannelID: context.channelID,
		},
	})

	inputChan := make(chan string)
	errorChan := make(chan error)
	session := sessionContext{context, channel, inputChan, errorChan, false, false}

	for inputChan != nil || errorChan != nil || requests != nil {
		select {
		case input, ok := <-inputChan:
			if !ok {
				inputChan = nil
				continue
			}
			context.logEvent(sessionInputLog{
				channelLog: channelLog{
					ChannelID: context.channelID,
				},
				Input: input,
			})
		case err, ok := <-errorChan:
			if !ok {
				errorChan = nil
				continue
			}
			if err != nil {
				return err
			}
		case request, ok := <-requests:
			if !ok {
				requests = nil
				if !session.active {
					close(inputChan)
					close(errorChan)
				}
				continue
			}
			if err := session.handleRequest(request); err != nil {
				return err
			}
		}
	}

	return nil
}
<commit_msg>Improve session error handling, prevent program errors from closing the whole connection<commit_after>package main

import (
	"bufio"
	"errors"
	"io"
	"strings"

	"golang.org/x/crypto/ssh"
	"golang.org/x/term"
)

type ptyRequestPayload struct {
	Term                                   string
	Width, Height, PixelWidth, PixelHeight uint32
	Modes                                  string
}

func (request ptyRequestPayload) reply() []byte {
	return nil
}
func (request ptyRequestPayload) logEntry(channelID int) logEntry {
	return ptyLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Terminal: request.Term,
		Width:    request.Width,
		Height:   request.Height,
	}
}

type shellRequestPayload struct{}

func (request shellRequestPayload) reply() []byte {
	return nil
}
func (request shellRequestPayload) logEntry(channelID int) logEntry {
	return shellLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
	}
}

type x11RequestPayload struct {
	SingleConnection         bool
	AuthProtocol, AuthCookie string
	ScreenNumber             uint32
}

func (request x11RequestPayload) reply() []byte {
	return nil
}
func (request x11RequestPayload) logEntry(channelID int) logEntry {
	return x11Log{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Screen: request.ScreenNumber,
	}
}

type envRequestPayload struct {
	Name, Value string
}

func (request envRequestPayload) reply() []byte {
	return nil
}
func (request envRequestPayload) logEntry(channelID int) logEntry {
	return envLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Name:  request.Name,
		Value: request.Value,
	}
}

type execRequestPayload struct {
	Command string
}

func (request execRequestPayload) reply() []byte {
	return nil
}
func (request execRequestPayload) logEntry(channelID int) logEntry {
	return execLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Command: request.Command,
	}
}

type subsystemRequestPayload struct {
	Subsystem string
}

func (request subsystemRequestPayload) reply() []byte {
	return nil
}
func (request subsystemRequestPayload) logEntry(channelID int) logEntry {
	return subsystemLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Subsystem: request.Subsystem,
	}
}

type windowChangeRequestPayload struct {
	Width, Height, PixelWidth, PixelHeight uint32
}

func (request windowChangeRequestPayload) reply() []byte {
	return nil
}
func (request windowChangeRequestPayload) logEntry(channelID int) logEntry {
	return windowChangeLog{
		channelLog: channelLog{
			ChannelID: channelID,
		},
		Width:  request.Width,
		Height: request.Height,
	}
}

type sessionContext struct {
	channelContext
	ssh.Channel
	inputChan chan string
	active    bool
	pty       bool
}

type scannerReadLiner struct {
	scanner   *bufio.Scanner
	inputChan chan<- string
}

func (r scannerReadLiner) ReadLine() (string, error) {
	if !r.scanner.Scan() {
		if err := r.scanner.Err(); err != nil {
			return "", err
		}
		return "", io.EOF
	}
	line := r.scanner.Text()
	r.inputChan <- line
	return line, nil
}

type terminalReadLiner struct {
	terminal  *term.Terminal
	inputChan chan<- string
}

type clientEOFError struct{}

var clientEOF = clientEOFError{}

func (clientEOFError) Error() string {
	return "Client EOF"
}

func (r terminalReadLiner) ReadLine() (string, error) {
	line, err := r.terminal.ReadLine()
	if err == nil || line != "" {
		r.inputChan <- line
	}
	if err == io.EOF {
		return line, clientEOF
	}
	return line, err
}

func (context *sessionContext) handleProgram(program []string) {
	context.active = true
	var stdin readLiner
	var stdout, stderr io.Writer
	if context.pty {
		terminal := term.NewTerminal(context, "")
		stdin = terminalReadLiner{terminal, context.inputChan}
		stdout = terminal
		stderr = terminal
	} else {
		stdin = scannerReadLiner{bufio.NewScanner(context), context.inputChan}
		stdout = context
		stderr = context.Stderr()
	}
	go func() {
		defer close(context.inputChan)

		result, err := executeProgram(commandContext{program, stdin, stdout, stderr, context.pty, context.User()})
		if err != nil && err != io.EOF && err != clientEOF {
			warningLogger.Printf("Error executing program: %s", err)
			return
		}

		if err == clientEOF && context.pty {
			if _, err := context.Write([]byte("\r\n")); err != nil {
				warningLogger.Printf("Error sending CRLF: %s", err)
				return
			}
		}

		if _, err := context.SendRequest("exit-status", false, ssh.Marshal(struct {
			ExitStatus uint32
		}{result})); err != nil {
			warningLogger.Printf("Error sending exit status: %s", err)
			return
		}

		if (context.pty && err == clientEOF) || err == nil {
			if _, err := context.SendRequest("eow@openssh.com", false, nil); err != nil {
				warningLogger.Printf("Error sending EOW: %s", err)
				return
			}
		}

		if err := context.CloseWrite(); err != nil {
			warningLogger.Printf("Error sending EOF: %s", err)
			return
		}

		if err := context.Close(); err != nil {
			warningLogger.Printf("Error closing channel: %s", err)
			return
		}
	}()
}

func (context *sessionContext) handleRequest(request *ssh.Request) error {
	switch request.Type {
	case "pty-req":
		if !context.active {
			if context.pty {
				return errors.New("a pty is already requested")
			}
			payload := &ptyRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.pty = true
			return nil
		}
	case "shell":
		if !context.active {
			if len(request.Payload) != 0 {
				return errors.New("invalid request payload")
			}
			payload := &shellRequestPayload{}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.active = true
			context.handleProgram(shellProgram)
			return nil
		}
	case "x11-req":
		if !context.active {
			payload := &x11RequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			return request.Reply(true, payload.reply())
		}
	case "env":
		if !context.active {
			payload := &envRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			return request.Reply(true, payload.reply())
		}
	case "exec":
		if !context.active {
			payload := &execRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.active = true
			context.handleProgram(strings.Fields(payload.Command))
			return nil
		}
	case "subsystem":
		if !context.active {
			payload := &subsystemRequestPayload{}
			if err := ssh.Unmarshal(request.Payload, payload); err != nil {
				return err
			}
			context.logEvent(payload.logEntry(context.channelID))
			if err := request.Reply(true, payload.reply()); err != nil {
				return err
			}
			context.active = true
			context.handleProgram(strings.Fields(payload.Subsystem))
		}
	case "window-change":
		payload := &windowChangeRequestPayload{}
		if err := ssh.Unmarshal(request.Payload, payload); err != nil {
			return err
		}
		context.logEvent(payload.logEntry(context.channelID))
		return request.Reply(true, payload.reply())
	}
	warningLogger.Printf("Rejected session request: %s", request.Type)
	return request.Reply(false, nil)
}

func handleSessionChannel(newChannel ssh.NewChannel, context channelContext) error {
	if context.noMoreSessions {
		return errors.New("nore more sessions were supposed to be requested")
	}
	if len(newChannel.ExtraData()) != 0 {
		return errors.New("invalid channel data")
	}
	channel, requests, err := newChannel.Accept()
	if err != nil {
		return err
	}
	context.logEvent(sessionLog{
		channelLog: channelLog{
			ChannelID: context.channelID,
		},
	})
	defer context.logEvent(sessionCloseLog{
		channelLog: channelLog{
			ChannelID: context.channelID,
		},
	})

	inputChan := make(chan string)
	session := sessionContext{context, channel, inputChan, false, false}

	for inputChan != nil || requests != nil {
		select {
		case input, ok := <-inputChan:
			if !ok {
				inputChan = nil
				continue
			}
			context.logEvent(sessionInputLog{
				channelLog: channelLog{
					ChannelID: context.channelID,
				},
				Input: input,
			})
		case request, ok := <-requests:
			if !ok {
				requests = nil
				if !session.active {
					close(inputChan)
				}
				continue
			}
			if err := session.handleRequest(request); err != nil {
				return err
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2013 The ql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSES/QL-LICENSE file.

// 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 tidb

import (
	"encoding/json"
	"fmt"
	"sync/atomic"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/field"
	"github.com/pingcap/tidb/kv"
	mysql "github.com/pingcap/tidb/mysqldef"
	"github.com/pingcap/tidb/rset"
	"github.com/pingcap/tidb/sessionctx"
	"github.com/pingcap/tidb/sessionctx/db"
	"github.com/pingcap/tidb/sessionctx/forupdate"
	"github.com/pingcap/tidb/sessionctx/variable"
	"github.com/pingcap/tidb/stmt"
	"github.com/pingcap/tidb/stmt/stmts"
	"github.com/pingcap/tidb/util/errors2"
)

// Session context
type Session interface {
	Status() uint16                               // Flag of current status, such as autocommit
	LastInsertID() uint64                         // Last inserted auto_increment id
	AffectedRows() uint64                         // Affected rows by lastest executed stmt
	Execute(sql string) ([]rset.Recordset, error) // Execute a sql statement
	SetUsername(name string)                      // Current user name
	String() string                               // For debug
	FinishTxn(rollback bool) error
	// For execute prepare statement in binary protocol
	PrepareStmt(sql string) (stmtID uint32, paramCount int, fields []*field.ResultField, err error)
	// Execute a prepared statement
	ExecutePreparedStmt(stmtID uint32, param ...interface{}) (rset.Recordset, error)
	DropPreparedStmt(stmtID uint32) error
	SetClientCapability(uint32) // Set client capability flags
	Close() error
	Retry() error
}

var (
	_         Session = (*session)(nil)
	sessionID int64
)

type stmtRecord struct {
	stmtID uint32
	st     stmt.Statement
	params []interface{}
}

type stmtHistory struct {
	history []*stmtRecord
}

func (h *stmtHistory) add(stmtID uint32, st stmt.Statement, params ...interface{}) {
	s := &stmtRecord{
		stmtID: stmtID,
		st:     st,
		params: append(([]interface{})(nil), params...),
	}

	h.history = append(h.history, s)
}

func (h *stmtHistory) reset() {
	if len(h.history) > 0 {
		h.history = h.history[:0]
	}
}

func (h *stmtHistory) clone() *stmtHistory {
	nh := *h
	nh.history = make([]*stmtRecord, len(h.history))
	copy(nh.history, h.history)
	return &nh
}

type session struct {
	txn      kv.Transaction // Current transaction
	userName string
	args     []interface{} // Statment execution args, this should be cleaned up after exec
	values   map[fmt.Stringer]interface{}
	store    kv.Storage
	sid      int64
	history  stmtHistory
}

func (s *session) Status() uint16 {
	return variable.GetSessionVars(s).Status
}

func (s *session) LastInsertID() uint64 {
	return variable.GetSessionVars(s).LastInsertID
}

func (s *session) AffectedRows() uint64 {
	return variable.GetSessionVars(s).AffectedRows
}

func (s *session) SetUsername(name string) {
	s.userName = name
}

func (s *session) resetHistory() {
	s.ClearValue(forupdate.ForUpdateKey)
	s.history.reset()
}

func (s *session) SetClientCapability(capability uint32) {
	variable.GetSessionVars(s).ClientCapability = capability
}

func (s *session) FinishTxn(rollback bool) error {
	// transaction has already been committed or rolled back
	if s.txn == nil {
		return nil
	}
	defer func() {
		s.txn = nil
	}()

	if rollback {
		return s.txn.Rollback()
	}

	err := s.txn.Commit()
	if err != nil {
		log.Warnf("txn:%s, %v", s.txn, err)
		return errors.Trace(err)
	}

	s.resetHistory()
	return nil
}

func (s *session) String() string {
	// TODO: how to print binded context in values appropriately?
	data := map[string]interface{}{
		"userName":   s.userName,
		"currDBName": db.GetCurrentSchema(s),
		"sid":        s.sid,
	}

	if s.txn != nil {
		// if txn is committed or rolled back, txn is nil.
		data["txn"] = s.txn.String()
	}

	b, _ := json.MarshalIndent(data, "", "  ")
	return string(b)
}

func needRetry(st stmt.Statement) bool {
	switch st.(type) {
	case *stmts.PreparedStmt, *stmts.ShowStmt, *stmts.DoStmt:
		return false
	default:
		return true
	}
}

func isPreparedStmt(st stmt.Statement) bool {
	switch st.(type) {
	case *stmts.PreparedStmt:
		return true
	default:
		return false
	}
}

func (s *session) Retry() error {
	nh := s.history.clone()
	defer func() {
		s.history.history = nh.history
	}()

	if forUpdate := s.Value(forupdate.ForUpdateKey); forUpdate != nil {
		return errors.Errorf("can not retry select for update statement")
	}

	var err error
	for {
		s.resetHistory()
		s.FinishTxn(true)
		success := true
		for _, sr := range nh.history {
			st := sr.st
			// Skip prepare statement
			if !needRetry(st) {
				continue
			}
			log.Warnf("Retry %s", st.OriginText())
			_, err = runStmt(s, st)
			if err != nil {
				if errors2.ErrorEqual(err, kv.ErrConditionNotMatch) {
					success = false
					break
				}
				log.Warnf("session:%v, err:%v", s, err)
				return errors.Trace(err)
			}
		}
		if success {
			return nil
		}
	}

	return nil
}

func (s *session) Execute(sql string) ([]rset.Recordset, error) {
	statements, err := Compile(sql)
	if err != nil {
		log.Errorf("Syntax error: %s", sql)
		log.Errorf("Error occurs at %s.", err)
		return nil, errors.Trace(err)
	}

	var rs []rset.Recordset

	for _, st := range statements {
		r, err := runStmt(s, st)
		if err != nil {
			log.Warnf("session:%v, err:%v", s, err)
			return nil, errors.Trace(err)
		}

		// Record executed query
		if isPreparedStmt(st) {
			ps := st.(*stmts.PreparedStmt)
			s.history.add(ps.ID, st)
		} else {
			s.history.add(0, st)
		}

		if r != nil {
			rs = append(rs, r)
		}
	}

	return rs, nil
}

// For execute prepare statement in binary protocol
func (s *session) PrepareStmt(sql string) (stmtID uint32, paramCount int, fields []*field.ResultField, err error) {
	return prepareStmt(s, sql)
}

// checkArgs make sure all the arguments's type is known and can be handled.
// integer types are converted to int64 and uint64, time.Time is converted to mysql.Time.
// time.Duration. is converted to mysql.Duration, other known types is leaved as it is.
func checkArgs(args ...interface{}) error {
	for i, v := range args {
		switch x := v.(type) {
		case bool:
			if x {
				args[i] = int64(1)
			} else {
				args[i] = int64(0)
			}
		case int8:
			args[i] = int64(x)
		case int16:
			args[i] = int64(x)
		case int32:
			args[i] = int64(x)
		case int:
			args[i] = int64(x)
		case uint8:
			args[i] = uint64(x)
		case uint16:
			args[i] = uint64(x)
		case uint32:
			args[i] = uint64(x)
		case uint:
			args[i] = uint64(x)
		case int64:
		case uint64:
		case float32:
		case float64:
		case string:
		case []byte:
		case time.Duration:
			args[i] = mysql.Duration{Duration: x}
		case time.Time:
			args[i] = mysql.Time{Time: x, Type: mysql.TypeDatetime}
		case nil:
		default:
			return errors.Errorf("cannot use arg[%d] (type %T):unsupported type", i, v)
		}
	}
	return nil
}

// Execute a prepared statement
func (s *session) ExecutePreparedStmt(stmtID uint32, args ...interface{}) (rset.Recordset, error) {
	err := checkArgs(args...)
	if err != nil {
		return nil, err
	}

	st := &stmts.ExecuteStmt{ID: stmtID}
	s.history.add(stmtID, st, args...)
	return runStmt(s, st, args...)
}

func (s *session) DropPreparedStmt(stmtID uint32) error {
	return dropPreparedStmt(s, stmtID)
}

// If forceNew is true, GetTxn() must return a new transaction.
// In this situation, if current transaction is still in progress,
// there will be an implicit commit and create a new transaction.
func (s *session) GetTxn(forceNew bool) (kv.Transaction, error) {
	var err error
	if s.txn == nil {
		s.resetHistory()
		s.txn, err = s.store.Begin()
		if err != nil {
			return nil, err
		}

		log.Infof("New txn:%s in session:%d", s.txn, s.sid)
		return s.txn, nil
	}
	if forceNew {
		err = s.txn.Commit()
		if err != nil {
			return nil, err
		}
		s.resetHistory()
		s.txn, err = s.store.Begin()
		if err != nil {
			return nil, err
		}
		log.Warnf("Force new txn:%s in session:%d", s.txn, s.sid)
	}
	return s.txn, nil
}

func (s *session) SetValue(key fmt.Stringer, value interface{}) {
	s.values[key] = value
}

func (s *session) Value(key fmt.Stringer) interface{} {
	value := s.values[key]
	return value
}

func (s *session) ClearValue(key fmt.Stringer) {
	delete(s.values, key)
}

// Close function does some clean work when session end.
func (s *session) Close() error {
	return s.FinishTxn(true)
}

// CreateSession creates a new session environment.
func CreateSession(store kv.Storage) (Session, error) {
	s := &session{
		values: make(map[fmt.Stringer]interface{}),
		store:  store,
		sid:    atomic.AddInt64(&sessionID, 1),
	}
	domain, err := domap.Get(store)
	if err != nil {
		return nil, err
	}
	sessionctx.BindDomain(s, domain)

	variable.BindSessionVars(s)
	variable.GetSessionVars(s).SetStatus(mysql.ServerStatusAutocommit)
	return s, nil
}
<commit_msg>tidb: checkArgs fix comment syntax.<commit_after>// Copyright 2013 The ql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSES/QL-LICENSE file.

// 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 tidb

import (
	"encoding/json"
	"fmt"
	"sync/atomic"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/field"
	"github.com/pingcap/tidb/kv"
	mysql "github.com/pingcap/tidb/mysqldef"
	"github.com/pingcap/tidb/rset"
	"github.com/pingcap/tidb/sessionctx"
	"github.com/pingcap/tidb/sessionctx/db"
	"github.com/pingcap/tidb/sessionctx/forupdate"
	"github.com/pingcap/tidb/sessionctx/variable"
	"github.com/pingcap/tidb/stmt"
	"github.com/pingcap/tidb/stmt/stmts"
	"github.com/pingcap/tidb/util/errors2"
)

// Session context
type Session interface {
	Status() uint16                               // Flag of current status, such as autocommit
	LastInsertID() uint64                         // Last inserted auto_increment id
	AffectedRows() uint64                         // Affected rows by lastest executed stmt
	Execute(sql string) ([]rset.Recordset, error) // Execute a sql statement
	SetUsername(name string)                      // Current user name
	String() string                               // For debug
	FinishTxn(rollback bool) error
	// For execute prepare statement in binary protocol
	PrepareStmt(sql string) (stmtID uint32, paramCount int, fields []*field.ResultField, err error)
	// Execute a prepared statement
	ExecutePreparedStmt(stmtID uint32, param ...interface{}) (rset.Recordset, error)
	DropPreparedStmt(stmtID uint32) error
	SetClientCapability(uint32) // Set client capability flags
	Close() error
	Retry() error
}

var (
	_         Session = (*session)(nil)
	sessionID int64
)

type stmtRecord struct {
	stmtID uint32
	st     stmt.Statement
	params []interface{}
}

type stmtHistory struct {
	history []*stmtRecord
}

func (h *stmtHistory) add(stmtID uint32, st stmt.Statement, params ...interface{}) {
	s := &stmtRecord{
		stmtID: stmtID,
		st:     st,
		params: append(([]interface{})(nil), params...),
	}

	h.history = append(h.history, s)
}

func (h *stmtHistory) reset() {
	if len(h.history) > 0 {
		h.history = h.history[:0]
	}
}

func (h *stmtHistory) clone() *stmtHistory {
	nh := *h
	nh.history = make([]*stmtRecord, len(h.history))
	copy(nh.history, h.history)
	return &nh
}

type session struct {
	txn      kv.Transaction // Current transaction
	userName string
	args     []interface{} // Statment execution args, this should be cleaned up after exec
	values   map[fmt.Stringer]interface{}
	store    kv.Storage
	sid      int64
	history  stmtHistory
}

func (s *session) Status() uint16 {
	return variable.GetSessionVars(s).Status
}

func (s *session) LastInsertID() uint64 {
	return variable.GetSessionVars(s).LastInsertID
}

func (s *session) AffectedRows() uint64 {
	return variable.GetSessionVars(s).AffectedRows
}

func (s *session) SetUsername(name string) {
	s.userName = name
}

func (s *session) resetHistory() {
	s.ClearValue(forupdate.ForUpdateKey)
	s.history.reset()
}

func (s *session) SetClientCapability(capability uint32) {
	variable.GetSessionVars(s).ClientCapability = capability
}

func (s *session) FinishTxn(rollback bool) error {
	// transaction has already been committed or rolled back
	if s.txn == nil {
		return nil
	}
	defer func() {
		s.txn = nil
	}()

	if rollback {
		return s.txn.Rollback()
	}

	err := s.txn.Commit()
	if err != nil {
		log.Warnf("txn:%s, %v", s.txn, err)
		return errors.Trace(err)
	}

	s.resetHistory()
	return nil
}

func (s *session) String() string {
	// TODO: how to print binded context in values appropriately?
	data := map[string]interface{}{
		"userName":   s.userName,
		"currDBName": db.GetCurrentSchema(s),
		"sid":        s.sid,
	}

	if s.txn != nil {
		// if txn is committed or rolled back, txn is nil.
		data["txn"] = s.txn.String()
	}

	b, _ := json.MarshalIndent(data, "", "  ")
	return string(b)
}

func needRetry(st stmt.Statement) bool {
	switch st.(type) {
	case *stmts.PreparedStmt, *stmts.ShowStmt, *stmts.DoStmt:
		return false
	default:
		return true
	}
}

func isPreparedStmt(st stmt.Statement) bool {
	switch st.(type) {
	case *stmts.PreparedStmt:
		return true
	default:
		return false
	}
}

func (s *session) Retry() error {
	nh := s.history.clone()
	defer func() {
		s.history.history = nh.history
	}()

	if forUpdate := s.Value(forupdate.ForUpdateKey); forUpdate != nil {
		return errors.Errorf("can not retry select for update statement")
	}

	var err error
	for {
		s.resetHistory()
		s.FinishTxn(true)
		success := true
		for _, sr := range nh.history {
			st := sr.st
			// Skip prepare statement
			if !needRetry(st) {
				continue
			}
			log.Warnf("Retry %s", st.OriginText())
			_, err = runStmt(s, st)
			if err != nil {
				if errors2.ErrorEqual(err, kv.ErrConditionNotMatch) {
					success = false
					break
				}
				log.Warnf("session:%v, err:%v", s, err)
				return errors.Trace(err)
			}
		}
		if success {
			return nil
		}
	}

	return nil
}

func (s *session) Execute(sql string) ([]rset.Recordset, error) {
	statements, err := Compile(sql)
	if err != nil {
		log.Errorf("Syntax error: %s", sql)
		log.Errorf("Error occurs at %s.", err)
		return nil, errors.Trace(err)
	}

	var rs []rset.Recordset

	for _, st := range statements {
		r, err := runStmt(s, st)
		if err != nil {
			log.Warnf("session:%v, err:%v", s, err)
			return nil, errors.Trace(err)
		}

		// Record executed query
		if isPreparedStmt(st) {
			ps := st.(*stmts.PreparedStmt)
			s.history.add(ps.ID, st)
		} else {
			s.history.add(0, st)
		}

		if r != nil {
			rs = append(rs, r)
		}
	}

	return rs, nil
}

// For execute prepare statement in binary protocol
func (s *session) PrepareStmt(sql string) (stmtID uint32, paramCount int, fields []*field.ResultField, err error) {
	return prepareStmt(s, sql)
}

// checkArgs makes sure all the arguments' types are known and can be handled.
// integer types are converted to int64 and uint64, time.Time is converted to mysql.Time.
// time.Duration is converted to mysql.Duration, other known types are leaved as it is.
func checkArgs(args ...interface{}) error {
	for i, v := range args {
		switch x := v.(type) {
		case bool:
			if x {
				args[i] = int64(1)
			} else {
				args[i] = int64(0)
			}
		case int8:
			args[i] = int64(x)
		case int16:
			args[i] = int64(x)
		case int32:
			args[i] = int64(x)
		case int:
			args[i] = int64(x)
		case uint8:
			args[i] = uint64(x)
		case uint16:
			args[i] = uint64(x)
		case uint32:
			args[i] = uint64(x)
		case uint:
			args[i] = uint64(x)
		case int64:
		case uint64:
		case float32:
		case float64:
		case string:
		case []byte:
		case time.Duration:
			args[i] = mysql.Duration{Duration: x}
		case time.Time:
			args[i] = mysql.Time{Time: x, Type: mysql.TypeDatetime}
		case nil:
		default:
			return errors.Errorf("cannot use arg[%d] (type %T):unsupported type", i, v)
		}
	}
	return nil
}

// Execute a prepared statement
func (s *session) ExecutePreparedStmt(stmtID uint32, args ...interface{}) (rset.Recordset, error) {
	err := checkArgs(args...)
	if err != nil {
		return nil, err
	}

	st := &stmts.ExecuteStmt{ID: stmtID}
	s.history.add(stmtID, st, args...)
	return runStmt(s, st, args...)
}

func (s *session) DropPreparedStmt(stmtID uint32) error {
	return dropPreparedStmt(s, stmtID)
}

// If forceNew is true, GetTxn() must return a new transaction.
// In this situation, if current transaction is still in progress,
// there will be an implicit commit and create a new transaction.
func (s *session) GetTxn(forceNew bool) (kv.Transaction, error) {
	var err error
	if s.txn == nil {
		s.resetHistory()
		s.txn, err = s.store.Begin()
		if err != nil {
			return nil, err
		}

		log.Infof("New txn:%s in session:%d", s.txn, s.sid)
		return s.txn, nil
	}
	if forceNew {
		err = s.txn.Commit()
		if err != nil {
			return nil, err
		}
		s.resetHistory()
		s.txn, err = s.store.Begin()
		if err != nil {
			return nil, err
		}
		log.Warnf("Force new txn:%s in session:%d", s.txn, s.sid)
	}
	return s.txn, nil
}

func (s *session) SetValue(key fmt.Stringer, value interface{}) {
	s.values[key] = value
}

func (s *session) Value(key fmt.Stringer) interface{} {
	value := s.values[key]
	return value
}

func (s *session) ClearValue(key fmt.Stringer) {
	delete(s.values, key)
}

// Close function does some clean work when session end.
func (s *session) Close() error {
	return s.FinishTxn(true)
}

// CreateSession creates a new session environment.
func CreateSession(store kv.Storage) (Session, error) {
	s := &session{
		values: make(map[fmt.Stringer]interface{}),
		store:  store,
		sid:    atomic.AddInt64(&sessionID, 1),
	}
	domain, err := domap.Get(store)
	if err != nil {
		return nil, err
	}
	sessionctx.BindDomain(s, domain)

	variable.BindSessionVars(s)
	variable.GetSessionVars(s).SetStatus(mysql.ServerStatusAutocommit)
	return s, nil
}
<|endoftext|>


<commit_before>package gorethink

import (
	"code.google.com/p/goprotobuf/proto"
	p "github.com/dancannon/gorethink/ql2"
	"sync/atomic"
	"time"
)

type Session struct {
	token      int64
	address    string
	database   string
	timeout    time.Duration
	authkey    string
	timeFormat string

	// Pool configuration options
	maxIdle     int
	maxActive   int
	idleTimeout time.Duration

	closed bool

	pool *Pool
}

func newSession(args map[string]interface{}) *Session {
	s := &Session{}

	if token, ok := args["token"]; ok {
		s.token = token.(int64)
	}
	if address, ok := args["address"]; ok {
		s.address = address.(string)
	}
	if database, ok := args["database"]; ok {
		s.database = database.(string)
	}
	if timeout, ok := args["timeout"]; ok {
		s.timeout = timeout.(time.Duration)
	}
	if authkey, ok := args["authkey"]; ok {
		s.authkey = authkey.(string)
	}

	// Pool configuration options
	if maxIdle, ok := args["maxIdle"]; ok {
		s.maxIdle = maxIdle.(int)
	} else {
		s.maxIdle = 1
	}
	if maxActive, ok := args["maxActive"]; ok {
		s.maxActive = maxActive.(int)
	} else {
		s.maxActive = 1
	}
	if idleTimeout, ok := args["idleTimeout"]; ok {
		s.idleTimeout = idleTimeout.(time.Duration)
	}

	return s
}

// Connect creates a new database session.
func Connect(args map[string]interface{}) (*Session, error) {
	s := newSession(args)
	err := s.Reconnect()

	return s, err
}

// Reconnect closes and re-opens a session.
func (s *Session) Reconnect() error {
	if err := s.Close(); err != nil {
		return err
	}

	s.closed = false
	if s.pool == nil {
		s.pool = &Pool{
			Session:     s,
			MaxIdle:     s.maxIdle,
			MaxActive:   s.maxActive,
			IdleTimeout: s.idleTimeout,
		}
	}
	return nil
}

// Close closes the session
func (s *Session) Close() error {
	if s.closed {
		return nil
	}

	var err error
	if s.pool != nil {
		err = s.pool.Close()
	}
	s.closed = true

	return err
}

// Use changes the default database used
func (s *Session) Use(database string) {
	s.database = database
}

// SetTimeout causes any future queries that are run on this session to timeout
// after the given duration, returning a timeout error.  Set to zero to disable.
func (s *Session) SetTimeout(timeout time.Duration) {
	s.timeout = timeout
}

// SetMaxIdleConns sets the maximum number of connections in the idle
// connection pool.
//
// If MaxOpenConns is greater than 0 but less than the new MaxIdleConns
// then the new MaxIdleConns will be reduced to match the MaxOpenConns limit
//
// If n <= 0, no idle connections are retained.
func (s *Session) SetMaxIdleConns(n int) {
	s.pool.MaxIdle = n
}

// SetMaxOpenConns sets the maximum number of open connections to the database.
//
// If MaxIdleConns is greater than 0 and the new MaxOpenConns is less than
// MaxIdleConns, then MaxIdleConns will be reduced to match the new
// MaxOpenConns limit
//
// If n <= 0, then there is no limit on the number of open connections.
// The default is 0 (unlimited).
func (s *Session) SetMaxOpenConns(n int) {
	s.pool.MaxActive = n
}

// getToken generates the next query token, used to number requests and match
// responses with requests.
func (s *Session) nextToken() int64 {
	return atomic.AddInt64(&s.token, 1)
}

// startQuery creates a query from the term given and sends it to the server.
// The result from the server is returned as ResultRows
func (s *Session) startQuery(t RqlTerm, opts map[string]interface{}) (*ResultRows, error) {
	token := s.nextToken()

	// Build query tree
	pt := t.build()

	// Build global options
	globalOpts := []*p.Query_AssocPair{}
	for k, v := range opts {
		if k == "db" {
			globalOpts = append(globalOpts, &p.Query_AssocPair{
				Key: proto.String("db"),
				Val: Db(v).build(),
			})
		} else if k == "use_outdated" {
			globalOpts = append(globalOpts, &p.Query_AssocPair{
				Key: proto.String("use_outdated"),
				Val: Expr(v).build(),
			})
		} else if k == "noreply" {
			globalOpts = append(globalOpts, &p.Query_AssocPair{
				Key: proto.String("noreply"),
				Val: Expr(v).build(),
			})
		}
	}
	// If no DB option was set default to the value set in the connection
	if _, ok := opts["db"]; !ok {
		globalOpts = append(globalOpts, &p.Query_AssocPair{
			Key: proto.String("db"),
			Val: Db(s.database).build(),
		})
	}

	// Construct query
	query := &p.Query{
		Type:          p.Query_START.Enum(),
		Token:         proto.Int64(token),
		Query:         pt,
		GlobalOptargs: globalOpts,
	}

	conn := s.pool.Get()
	defer conn.Close()

	return conn.SendQuery(s, query, t, opts)
}

// continueQuery continues a previously run query.
func (s *Session) continueQuery(q *p.Query, t RqlTerm, opts map[string]interface{}) (*ResultRows, error) {
	nq := &p.Query{
		Type:  p.Query_CONTINUE.Enum(),
		Token: q.Token,
	}

	conn := s.pool.Get()
	defer conn.Close()

	return conn.SendQuery(s, nq, t, opts)
}

// stopQuery sends closes a query by sending Query_STOP to the server.
func (s *Session) stopQuery(q *p.Query, t RqlTerm, opts map[string]interface{}) (*ResultRows, error) {
	nq := &p.Query{
		Type:  p.Query_STOP.Enum(),
		Token: q.Token,
	}

	conn := s.pool.Get()
	defer conn.Close()

	return conn.SendQuery(s, nq, t, opts)
}
<commit_msg>Added documentation for Connect() to show supported arguments.<commit_after>package gorethink

import (
	"code.google.com/p/goprotobuf/proto"
	p "github.com/dancannon/gorethink/ql2"
	"sync/atomic"
	"time"
)

type Session struct {
	token      int64
	address    string
	database   string
	timeout    time.Duration
	authkey    string
	timeFormat string

	// Pool configuration options
	maxIdle     int
	maxActive   int
	idleTimeout time.Duration

	closed bool

	pool *Pool
}

func newSession(args map[string]interface{}) *Session {
	s := &Session{}

	if token, ok := args["token"]; ok {
		s.token = token.(int64)
	}
	if address, ok := args["address"]; ok {
		s.address = address.(string)
	}
	if database, ok := args["database"]; ok {
		s.database = database.(string)
	}
	if timeout, ok := args["timeout"]; ok {
		s.timeout = timeout.(time.Duration)
	}
	if authkey, ok := args["authkey"]; ok {
		s.authkey = authkey.(string)
	}

	// Pool configuration options
	if maxIdle, ok := args["maxIdle"]; ok {
		s.maxIdle = maxIdle.(int)
	} else {
		s.maxIdle = 1
	}
	if maxActive, ok := args["maxActive"]; ok {
		s.maxActive = maxActive.(int)
	} else {
		s.maxActive = 1
	}
	if idleTimeout, ok := args["idleTimeout"]; ok {
		s.idleTimeout = idleTimeout.(time.Duration)
	}

	return s
}

// Connect creates a new database session.
//
// Supported arguments include token, address, database, timeout, authkey,
// and timeFormat. Pool options (when using the included connection pool manager)
// include maxIdle, maxActive and idleTimeout.
//
// Example usage:
//
// var session *r.Session
// session, err := r.Connect(map[string]interface{}{
//        	"address":  "localhost:28015",
//       	"database": "test",
//        	"authkey":  "14daak1cad13dj",
//    })
func Connect(args map[string]interface{}) (*Session, error) {
	s := newSession(args)
	err := s.Reconnect()

	return s, err
}

// Reconnect closes and re-opens a session.
func (s *Session) Reconnect() error {
	if err := s.Close(); err != nil {
		return err
	}

	s.closed = false
	if s.pool == nil {
		s.pool = &Pool{
			Session:     s,
			MaxIdle:     s.maxIdle,
			MaxActive:   s.maxActive,
			IdleTimeout: s.idleTimeout,
		}
	}
	return nil
}

// Close closes the session
func (s *Session) Close() error {
	if s.closed {
		return nil
	}

	var err error
	if s.pool != nil {
		err = s.pool.Close()
	}
	s.closed = true

	return err
}

// Use changes the default database used
func (s *Session) Use(database string) {
	s.database = database
}

// SetTimeout causes any future queries that are run on this session to timeout
// after the given duration, returning a timeout error.  Set to zero to disable.
func (s *Session) SetTimeout(timeout time.Duration) {
	s.timeout = timeout
}

// SetMaxIdleConns sets the maximum number of connections in the idle
// connection pool.
//
// If MaxOpenConns is greater than 0 but less than the new MaxIdleConns
// then the new MaxIdleConns will be reduced to match the MaxOpenConns limit
//
// If n <= 0, no idle connections are retained.
func (s *Session) SetMaxIdleConns(n int) {
	s.pool.MaxIdle = n
}

// SetMaxOpenConns sets the maximum number of open connections to the database.
//
// If MaxIdleConns is greater than 0 and the new MaxOpenConns is less than
// MaxIdleConns, then MaxIdleConns will be reduced to match the new
// MaxOpenConns limit
//
// If n <= 0, then there is no limit on the number of open connections.
// The default is 0 (unlimited).
func (s *Session) SetMaxOpenConns(n int) {
	s.pool.MaxActive = n
}

// getToken generates the next query token, used to number requests and match
// responses with requests.
func (s *Session) nextToken() int64 {
	return atomic.AddInt64(&s.token, 1)
}

// startQuery creates a query from the term given and sends it to the server.
// The result from the server is returned as ResultRows
func (s *Session) startQuery(t RqlTerm, opts map[string]interface{}) (*ResultRows, error) {
	token := s.nextToken()

	// Build query tree
	pt := t.build()

	// Build global options
	globalOpts := []*p.Query_AssocPair{}
	for k, v := range opts {
		if k == "db" {
			globalOpts = append(globalOpts, &p.Query_AssocPair{
				Key: proto.String("db"),
				Val: Db(v).build(),
			})
		} else if k == "use_outdated" {
			globalOpts = append(globalOpts, &p.Query_AssocPair{
				Key: proto.String("use_outdated"),
				Val: Expr(v).build(),
			})
		} else if k == "noreply" {
			globalOpts = append(globalOpts, &p.Query_AssocPair{
				Key: proto.String("noreply"),
				Val: Expr(v).build(),
			})
		}
	}
	// If no DB option was set default to the value set in the connection
	if _, ok := opts["db"]; !ok {
		globalOpts = append(globalOpts, &p.Query_AssocPair{
			Key: proto.String("db"),
			Val: Db(s.database).build(),
		})
	}

	// Construct query
	query := &p.Query{
		Type:          p.Query_START.Enum(),
		Token:         proto.Int64(token),
		Query:         pt,
		GlobalOptargs: globalOpts,
	}

	conn := s.pool.Get()
	defer conn.Close()

	return conn.SendQuery(s, query, t, opts)
}

// continueQuery continues a previously run query.
func (s *Session) continueQuery(q *p.Query, t RqlTerm, opts map[string]interface{}) (*ResultRows, error) {
	nq := &p.Query{
		Type:  p.Query_CONTINUE.Enum(),
		Token: q.Token,
	}

	conn := s.pool.Get()
	defer conn.Close()

	return conn.SendQuery(s, nq, t, opts)
}

// stopQuery sends closes a query by sending Query_STOP to the server.
func (s *Session) stopQuery(q *p.Query, t RqlTerm, opts map[string]interface{}) (*ResultRows, error) {
	nq := &p.Query{
		Type:  p.Query_STOP.Enum(),
		Token: q.Token,
	}

	conn := s.pool.Get()
	defer conn.Close()

	return conn.SendQuery(s, nq, t, opts)
}
<|endoftext|>


<commit_before>package transmission

type (
	// SessionCommand is the root command to interact with Transmission via RPC
	SessionCommand struct {
		Method  string  `json:"method,omitempty"`
		Session Session `json:"arguments,omitempty"`
		Result  string  `json:"result,omitempty"`
	}

	// Session information about the current transmission session
	Session struct {
		AltSpeedDown    int  `json:"alt-speed-down"`
		AltSpeedEnabled bool `json:"alt-speed-enabled"`
		//Alt_speed_time_begin   int    `json:"alt-speed-time-begin"`
		//Alt_speed_time_day     int    `json:"alt-speed-time-day"`
		//Alt_speed_time_enabled bool   `json:"alt-speed-time-enabled"`
		//Alt_speed_time_end     int    `json:"alt-speed-time-end"`
		AltSpeedUp int `json:"alt-speed-up"`
		//Blocklist_enabled       bool   `json:"blocklist-enabled"`
		//Blocklist_size          int    `json:"blocklist-size"`
		//Blocklist_url           string `json:"blocklist-url"`
		CacheSizeMB int `json:"cache-size-mb"`
		//Config_dir                   string `json:"config-dir"`
		//Dht_enabled                  bool   `json:"dht-enabled"`
		DownloadDir          string `json:"download-dir"`
		DownloadDirFreeSpace int    `json:"download-dir-free-space"`
		DownloadQueueEnabled bool   `json:"download-queue-enabled"`
		DownloadQueueSize    int    `json:"download-queue-size"`
		//Encryption                 string `json:"encryption"`
		//Idle_seeding_limit           int    `json:"idle-seeding-limit"`
		//Idle_seeding_limit_enabled   bool   `json:"idle-seeding-limit-enabled"`
		IncompleteDir string `json:"incomplete-dir"`
		//Incomplete_dir_enabled       bool   `json:"incomplete-dir-enabled"`
		//Lpd_enabled                  bool   `json:"lpd-enabled"`
		PeerLimitGlobal     int `json:"peer-limit-global"`
		PeerLimitPerTorrent int `json:"peer-limit-per-torrent"`
		//Peer_port                 int    `json:"peer-port"`
		//Peer_port_random_on_start bool   `json:"peer-port-random-on-start"`
		//Pex_enabled               bool   `json:"pex-enabled"`
		//Port_forwarding_enabled   bool   `json:"port-forwarding-enabled"`
		//Queue_stalled_enabled     bool   `json:"queue-stalled-enabled"`
		//Queue_stalled_minutes     int    `json:"queue-stalled-minutes"`
		//Rename_partial_files      bool   `json:"rename-partial-files"`
		//RPC_version               int    `json:"rpc-version"`
		//RPC_version_minimum          int    `json:"rpc-version-minimum"`
		//Script_torrent_done_enabled  bool   `json:"script-torrent-done-enabled"`
		//Script_torrent_done_filename string `json:"script-torrent-done-filename"`
		SeedQueueEnabled      bool    `json:"seed-queue-enabled"`
		SeedQueueSize         int     `json:"seed-queue-size"`
		SeedRatioLimit        float64 `json:"seedRatioLimit"`
		SeedRatioLimited      bool    `json:"seedRatioLimited"`
		SpeedLimitDown        int     `json:"speed-limit-down"`
		SpeedLimitDownEnabled bool    `json:"speed-limit-down-enabled"`
		SpeedLimitUp          int     `json:"speed-limit-up"`
		SpeedLimitUpEnabled   bool    `json:"speed-limit-up-enabled"`
		//Start_added_torrents         bool   `json:"start-added-torrents"`
		//Trash_original_torrent_files bool   `json:"trash-original-torrent-files"`
		//Utp_enabled                  bool   `json:"utp-enabled"`
		Version string `json:"version"`
	}
)
<commit_msg>Switches DownloadDirFreeSpace to uint64<commit_after>package transmission

type (
	// SessionCommand is the root command to interact with Transmission via RPC
	SessionCommand struct {
		Method  string  `json:"method,omitempty"`
		Session Session `json:"arguments,omitempty"`
		Result  string  `json:"result,omitempty"`
	}

	// Session information about the current transmission session
	Session struct {
		AltSpeedDown    int  `json:"alt-speed-down"`
		AltSpeedEnabled bool `json:"alt-speed-enabled"`
		//Alt_speed_time_begin   int    `json:"alt-speed-time-begin"`
		//Alt_speed_time_day     int    `json:"alt-speed-time-day"`
		//Alt_speed_time_enabled bool   `json:"alt-speed-time-enabled"`
		//Alt_speed_time_end     int    `json:"alt-speed-time-end"`
		AltSpeedUp int `json:"alt-speed-up"`
		//Blocklist_enabled       bool   `json:"blocklist-enabled"`
		//Blocklist_size          int    `json:"blocklist-size"`
		//Blocklist_url           string `json:"blocklist-url"`
		CacheSizeMB int `json:"cache-size-mb"`
		//Config_dir                   string `json:"config-dir"`
		//Dht_enabled                  bool   `json:"dht-enabled"`
		DownloadDir          string `json:"download-dir"`
		DownloadDirFreeSpace uint64 `json:"download-dir-free-space"`
		DownloadQueueEnabled bool   `json:"download-queue-enabled"`
		DownloadQueueSize    int    `json:"download-queue-size"`
		//Encryption                 string `json:"encryption"`
		//Idle_seeding_limit           int    `json:"idle-seeding-limit"`
		//Idle_seeding_limit_enabled   bool   `json:"idle-seeding-limit-enabled"`
		IncompleteDir string `json:"incomplete-dir"`
		//Incomplete_dir_enabled       bool   `json:"incomplete-dir-enabled"`
		//Lpd_enabled                  bool   `json:"lpd-enabled"`
		PeerLimitGlobal     int `json:"peer-limit-global"`
		PeerLimitPerTorrent int `json:"peer-limit-per-torrent"`
		//Peer_port                 int    `json:"peer-port"`
		//Peer_port_random_on_start bool   `json:"peer-port-random-on-start"`
		//Pex_enabled               bool   `json:"pex-enabled"`
		//Port_forwarding_enabled   bool   `json:"port-forwarding-enabled"`
		//Queue_stalled_enabled     bool   `json:"queue-stalled-enabled"`
		//Queue_stalled_minutes     int    `json:"queue-stalled-minutes"`
		//Rename_partial_files      bool   `json:"rename-partial-files"`
		//RPC_version               int    `json:"rpc-version"`
		//RPC_version_minimum          int    `json:"rpc-version-minimum"`
		//Script_torrent_done_enabled  bool   `json:"script-torrent-done-enabled"`
		//Script_torrent_done_filename string `json:"script-torrent-done-filename"`
		SeedQueueEnabled      bool    `json:"seed-queue-enabled"`
		SeedQueueSize         int     `json:"seed-queue-size"`
		SeedRatioLimit        float64 `json:"seedRatioLimit"`
		SeedRatioLimited      bool    `json:"seedRatioLimited"`
		SpeedLimitDown        int     `json:"speed-limit-down"`
		SpeedLimitDownEnabled bool    `json:"speed-limit-down-enabled"`
		SpeedLimitUp          int     `json:"speed-limit-up"`
		SpeedLimitUpEnabled   bool    `json:"speed-limit-up-enabled"`
		//Start_added_torrents         bool   `json:"start-added-torrents"`
		//Trash_original_torrent_files bool   `json:"trash-original-torrent-files"`
		//Utp_enabled                  bool   `json:"utp-enabled"`
		Version string `json:"version"`
	}
)
<|endoftext|>


<commit_before>// Copyright 2013 Beego Authors
// Copyright 2014 Unknwon
//
// 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 session a middleware that provides the session manager of Macaron.
package session

// NOTE: last sync 000033e on Nov 4, 2014.

import (
	"encoding/hex"
	"fmt"
	"net/http"
	"net/url"
	"time"

	"github.com/Unknwon/macaron"
)

const _VERSION = "0.1.1"

func Version() string {
	return _VERSION
}

// RawStore is the interface that operates the session data.
type RawStore interface {
	// Set sets value to given key in session.
	Set(key, value interface{}) error
	// Get gets value by given key in session.
	Get(key interface{}) interface{}
	// Delete delete a key from session.
	Delete(key interface{}) error
	// ID returns current session ID.
	ID() string
	// Release releases session resource and save data to provider.
	Release() error
	// Flush deletes all session data.
	Flush() error
}

// Store is the interface that contains all data for one session process with specific ID.
type Store interface {
	RawStore
	// Read returns raw session store by session ID.
	Read(sid string) (RawStore, error)
	// Destory deletes a session.
	Destory(*macaron.Context) error
	// RegenerateId regenerates a session store from old session ID to new one.
	RegenerateId(*macaron.Context) (RawStore, error)
	// Count counts and returns number of sessions.
	Count() int
	// GC calls GC to clean expired sessions.
	GC()
}

type store struct {
	RawStore
	*Manager
}

var _ Store = &store{}

// Options represents a struct for specifying configuration options for the session middleware.
type Options struct {
	// Name of provider. Default is "memory".
	Provider       string
	ProviderConfig string `json:"providerConfig"`
	CookieName     string `json:"cookieName"`
	CookiePath     string `json:"cookiePath"`
	Gclifetime     int64  `json:"gclifetime"`
	Maxlifetime    int64  `json:"maxLifetime"`
	Secure         bool   `json:"secure"`
	CookieLifeTime int    `json:"cookieLifeTime"`
	Domain         string `json:"domain"`
	IDLength       int    `json:"sessionIdLength"`
	// Configuration section name. Default is "session".
	Section string
}

func prepareOptions(options []Options) Options {
	var opt Options
	if len(options) > 0 {
		opt = options[0]
	}

	if len(opt.Section) == 0 {
		opt.Section = "session"
	}
	sec := macaron.Config().Section(opt.Section)

	if len(opt.Provider) == 0 {
		opt.Provider = sec.Key("PROVIDER").MustString("memory")
	}
	if len(opt.ProviderConfig) == 0 && opt.Provider == "file" {
		opt.ProviderConfig = sec.Key("PROVIDER_CONFIG").MustString("data/sessions")
	}
	if len(opt.CookieName) == 0 {
		opt.CookieName = sec.Key("COOKIE_NAME").MustString("MacaronSession")
	}
	if len(opt.CookiePath) == 0 {
		opt.CookiePath = sec.Key("COOKIE_PATH").MustString("/")
	}
	if opt.Gclifetime == 0 {
		opt.Gclifetime = sec.Key("GC_INTERVAL_TIME").MustInt64(3600)
	}
	if opt.Maxlifetime == 0 {
		opt.Maxlifetime = sec.Key("MAX_LIFE_TIME").MustInt64(opt.Gclifetime)
	}
	if !opt.Secure {
		opt.Secure = sec.Key("SECURE").MustBool()
	}
	if opt.CookieLifeTime == 0 {
		opt.CookieLifeTime = sec.Key("COOKIE_LIFE_TIME").MustInt()
	}
	if len(opt.Domain) == 0 {
		opt.Domain = sec.Key("DOMAIN").String()
	}
	if opt.IDLength == 0 {
		opt.IDLength = sec.Key("ID_LENGTH").MustInt(16)
	}

	return opt
}

// Sessioner is a middleware that maps a session.SessionStore service into the Macaron handler chain.
// An single variadic session.Options struct can be optionally provided to configure.
func Sessioner(options ...Options) macaron.Handler {
	opt := prepareOptions(options)
	manager, err := NewManager(opt.Provider, opt)
	if err != nil {
		panic(err)
	}
	go manager.startGC()

	return func(ctx *macaron.Context) {
		sess, err := manager.Start(ctx)
		if err != nil {
			panic("session: " + err.Error())
		}

		// Get flash.
		vals, _ := url.ParseQuery(ctx.GetCookie("macaron_flash"))
		if len(vals) > 0 {
			f := &Flash{Values: vals}
			f.ErrorMsg = f.Get("error")
			f.SuccessMsg = f.Get("success")
			f.InfoMsg = f.Get("info")
			f.WarningMsg = f.Get("warning")
			ctx.Data["Flash"] = f
			ctx.SetCookie("macaron_flash", "", -1, opt.CookiePath)
		}

		f := &Flash{ctx, url.Values{}, "", "", "", ""}
		ctx.Resp.Before(func(macaron.ResponseWriter) {
			if flash := f.Encode(); len(flash) > 0 {
				ctx.SetCookie("macaron_flash", flash, 0, opt.CookiePath)
			}
		})

		ctx.Map(f)
		s := store{
			RawStore: sess,
			Manager:  manager,
		}

		ctx.MapTo(s, (*Store)(nil))

		ctx.Next()

		if sess.Release() != nil {
			panic("session: " + err.Error())
		}
	}
}

// Provider is the interface that provides session manipulations.
type Provider interface {
	// Init initializes session provider.
	Init(gclifetime int64, config string) error
	// Read returns raw session store by session ID.
	Read(sid string) (RawStore, error)
	// Exist returns true if session with given ID exists.
	Exist(sid string) bool
	// Destory deletes a session by session ID.
	Destory(sid string) error
	// Regenerate regenerates a session store from old session ID to new one.
	Regenerate(oldsid, sid string) (RawStore, error)
	// Count counts and returns number of sessions.
	Count() int
	// GC calls GC to clean expired sessions.
	GC()
}

var providers = make(map[string]Provider)

// Register registers a provider.
func Register(name string, provider Provider) {
	if provider == nil {
		panic("session: cannot register provider with nil value")
	}
	if _, dup := providers[name]; dup {
		panic(fmt.Errorf("session: cannot register provider '%s' twice", name))
	}
	providers[name] = provider
}

//    _____
//   /     \ _____    ____ _____     ____   ___________
//  /  \ /  \\__  \  /    \\__  \   / ___\_/ __ \_  __ \
// /    Y    \/ __ \|   |  \/ __ \_/ /_/  >  ___/|  | \/
// \____|__  (____  /___|  (____  /\___  / \___  >__|
//         \/     \/     \/     \//_____/      \/

// Manager represents a struct that contains session provider and its configuration.
type Manager struct {
	provider Provider
	opt      Options
}

// NewManager creates and returns a new session manager by given provider name and configuration.
// It panics when given provider isn't registered.
func NewManager(name string, opt Options) (*Manager, error) {
	p, ok := providers[name]
	if !ok {
		return nil, fmt.Errorf("session: unknown provider ‘%q’(forgotten import?)", name)
	}
	if err := p.Init(opt.Maxlifetime, opt.ProviderConfig); err != nil {
		return nil, err
	}
	return &Manager{p, opt}, nil
}

// sessionId generates a new session ID with rand string, unix nano time, remote addr by hash function.
func (m *Manager) sessionId() string {
	return hex.EncodeToString(generateRandomKey(m.opt.IDLength))
}

// Start starts a session by generating new one
// or retrieve existence one by reading session ID from HTTP request if it's valid.
func (m *Manager) Start(ctx *macaron.Context) (RawStore, error) {
	sid := ctx.GetCookie(m.opt.CookieName)
	if len(sid) > 0 && m.provider.Exist(sid) {
		return m.provider.Read(sid)
	}

	sid = m.sessionId()
	sess, err := m.provider.Read(sid)
	if err != nil {
		return nil, err
	}

	cookie := &http.Cookie{
		Name:     m.opt.CookieName,
		Value:    sid,
		Path:     m.opt.CookiePath,
		HttpOnly: true,
		Secure:   m.opt.Secure,
		Domain:   m.opt.Domain,
	}
	if m.opt.CookieLifeTime >= 0 {
		cookie.MaxAge = m.opt.CookieLifeTime
	}
	http.SetCookie(ctx.Resp, cookie)
	ctx.Req.AddCookie(cookie)
	return sess, nil
}

// Read returns raw session store by session ID.
func (m *Manager) Read(sid string) (RawStore, error) {
	return m.provider.Read(sid)
}

// Destory deletes a session by given ID.
func (m *Manager) Destory(ctx *macaron.Context) error {
	sid := ctx.GetCookie(m.opt.CookieName)
	if len(sid) == 0 {
		return nil
	}

	if err := m.provider.Destory(sid); err != nil {
		return err
	}
	cookie := &http.Cookie{
		Name:     m.opt.CookieName,
		Path:     m.opt.CookiePath,
		HttpOnly: true,
		Expires:  time.Now(),
		MaxAge:   -1,
	}
	http.SetCookie(ctx.Resp, cookie)
	return nil
}

// RegenerateId regenerates a session store from old session ID to new one.
func (m *Manager) RegenerateId(ctx *macaron.Context) (sess RawStore, err error) {
	sid := m.sessionId()
	oldsid := ctx.GetCookie(m.opt.CookieName)
	if len(oldsid) == 0 {
		sess, err = m.provider.Read(oldsid)
		if err != nil {
			return nil, err
		}
	} else {
		sess, err = m.provider.Regenerate(oldsid, sid)
		if err != nil {
			return nil, err
		}
	}
	ck := &http.Cookie{
		Name:     m.opt.CookieName,
		Value:    sid,
		Path:     m.opt.CookiePath,
		HttpOnly: true,
		Secure:   m.opt.Secure,
		Domain:   m.opt.Domain,
	}
	if m.opt.CookieLifeTime >= 0 {
		ck.MaxAge = m.opt.CookieLifeTime
	}
	http.SetCookie(ctx.Resp, ck)
	ctx.Req.AddCookie(ck)
	return sess, nil
}

// Count counts and returns number of sessions.
func (m *Manager) Count() int {
	return m.provider.Count()
}

// GC starts GC job in a certain period.
func (m *Manager) GC() {
	m.provider.GC()
}

// startGC starts GC job in a certain period.
func (m *Manager) startGC() {
	m.GC()
	time.AfterFunc(time.Duration(m.opt.Gclifetime)*time.Second, func() { m.startGC() })
}

// SetSecure indicates whether to set cookie with HTTPS or not.
func (m *Manager) SetSecure(secure bool) {
	m.opt.Secure = secure
}

// ___________.____       _____    _________ ___ ___
// \_   _____/|    |     /  _  \  /   _____//   |   \
//  |    __)  |    |    /  /_\  \ \_____  \/    ~    \
//  |     \   |    |___/    |    \/        \    Y    /
//  \___  /   |_______ \____|__  /_______  /\___|_  /
//      \/            \/       \/        \/       \/

type Flash struct {
	ctx *macaron.Context
	url.Values
	ErrorMsg, WarningMsg, InfoMsg, SuccessMsg string
}

func (f *Flash) set(name, msg string, current ...bool) {
	isShow := false
	if (len(current) == 0 && macaron.FlashNow) ||
		(len(current) > 0 && current[0]) {
		isShow = true
	}

	if isShow {
		f.ctx.Data["Flash"] = f
	} else {
		f.Set(name, msg)
	}
}

func (f *Flash) Error(msg string, current ...bool) {
	f.ErrorMsg = msg
	f.set("error", msg, current...)
}

func (f *Flash) Warning(msg string, current ...bool) {
	f.WarningMsg = msg
	f.set("warning", msg, current...)
}

func (f *Flash) Info(msg string, current ...bool) {
	f.InfoMsg = msg
	f.set("info", msg, current...)
}

func (f *Flash) Success(msg string, current ...bool) {
	f.SuccessMsg = msg
	f.set("success", msg, current...)
}
<commit_msg>clean code and add comments<commit_after>// Copyright 2013 Beego Authors
// Copyright 2014 Unknwon
//
// 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 session a middleware that provides the session manager of Macaron.
package session

// NOTE: last sync 000033e on Nov 4, 2014.

import (
	"encoding/hex"
	"fmt"
	"net/http"
	"net/url"
	"time"

	"github.com/Unknwon/macaron"
)

const _VERSION = "0.1.1"

func Version() string {
	return _VERSION
}

// RawStore is the interface that operates the session data.
type RawStore interface {
	// Set sets value to given key in session.
	Set(key, value interface{}) error
	// Get gets value by given key in session.
	Get(key interface{}) interface{}
	// Delete delete a key from session.
	Delete(key interface{}) error
	// ID returns current session ID.
	ID() string
	// Release releases session resource and save data to provider.
	Release() error
	// Flush deletes all session data.
	Flush() error
}

// Store is the interface that contains all data for one session process with specific ID.
type Store interface {
	RawStore
	// Read returns raw session store by session ID.
	Read(sid string) (RawStore, error)
	// Destory deletes a session.
	Destory(*macaron.Context) error
	// RegenerateId regenerates a session store from old session ID to new one.
	RegenerateId(*macaron.Context) (RawStore, error)
	// Count counts and returns number of sessions.
	Count() int
	// GC calls GC to clean expired sessions.
	GC()
}

type store struct {
	RawStore
	*Manager
}

var _ Store = &store{}

// Options represents a struct for specifying configuration options for the session middleware.
type Options struct {
	// Name of provider. Default is "memory".
	Provider string
	// Provider configuration, it's corresponding to provider.
	ProviderConfig string
	// Cookie name to save session ID. Default is "MacaronSession".
	CookieName string
	// Cookie path to store. Default is "/".
	CookiePath string
	// GC interval time in seconds. Default is 3600.
	Gclifetime int64
	// Max life time in seconds. Default is whatever GC interval time is.
	Maxlifetime int64
	// Use HTTPS only. Default is false.
	Secure bool
	// Cookie life time. Default is 0.
	CookieLifeTime int
	// Cookie domain name. Default is empty.
	Domain string
	// Session ID length. Default is 16.
	IDLength int
	// Configuration section name. Default is "session".
	Section string
}

func prepareOptions(options []Options) Options {
	var opt Options
	if len(options) > 0 {
		opt = options[0]
	}
	if len(opt.Section) == 0 {
		opt.Section = "session"
	}
	sec := macaron.Config().Section(opt.Section)

	if len(opt.Provider) == 0 {
		opt.Provider = sec.Key("PROVIDER").MustString("memory")
	}
	if len(opt.ProviderConfig) == 0 && opt.Provider == "file" {
		opt.ProviderConfig = sec.Key("PROVIDER_CONFIG").MustString("data/sessions")
	}
	if len(opt.CookieName) == 0 {
		opt.CookieName = sec.Key("COOKIE_NAME").MustString("MacaronSession")
	}
	if len(opt.CookiePath) == 0 {
		opt.CookiePath = sec.Key("COOKIE_PATH").MustString("/")
	}
	if opt.Gclifetime == 0 {
		opt.Gclifetime = sec.Key("GC_INTERVAL_TIME").MustInt64(3600)
	}
	if opt.Maxlifetime == 0 {
		opt.Maxlifetime = sec.Key("MAX_LIFE_TIME").MustInt64(opt.Gclifetime)
	}
	if !opt.Secure {
		opt.Secure = sec.Key("SECURE").MustBool()
	}
	if opt.CookieLifeTime == 0 {
		opt.CookieLifeTime = sec.Key("COOKIE_LIFE_TIME").MustInt()
	}
	if len(opt.Domain) == 0 {
		opt.Domain = sec.Key("DOMAIN").String()
	}
	if opt.IDLength == 0 {
		opt.IDLength = sec.Key("ID_LENGTH").MustInt(16)
	}

	return opt
}

// Sessioner is a middleware that maps a session.SessionStore service into the Macaron handler chain.
// An single variadic session.Options struct can be optionally provided to configure.
func Sessioner(options ...Options) macaron.Handler {
	opt := prepareOptions(options)
	manager, err := NewManager(opt.Provider, opt)
	if err != nil {
		panic(err)
	}
	go manager.startGC()

	return func(ctx *macaron.Context) {
		sess, err := manager.Start(ctx)
		if err != nil {
			panic("session: " + err.Error())
		}

		// Get flash.
		vals, _ := url.ParseQuery(ctx.GetCookie("macaron_flash"))
		if len(vals) > 0 {
			f := &Flash{Values: vals}
			f.ErrorMsg = f.Get("error")
			f.SuccessMsg = f.Get("success")
			f.InfoMsg = f.Get("info")
			f.WarningMsg = f.Get("warning")
			ctx.Data["Flash"] = f
			ctx.SetCookie("macaron_flash", "", -1, opt.CookiePath)
		}

		f := &Flash{ctx, url.Values{}, "", "", "", ""}
		ctx.Resp.Before(func(macaron.ResponseWriter) {
			if flash := f.Encode(); len(flash) > 0 {
				ctx.SetCookie("macaron_flash", flash, 0, opt.CookiePath)
			}
		})

		ctx.Map(f)
		s := store{
			RawStore: sess,
			Manager:  manager,
		}

		ctx.MapTo(s, (*Store)(nil))

		ctx.Next()

		if sess.Release() != nil {
			panic("session: " + err.Error())
		}
	}
}

// Provider is the interface that provides session manipulations.
type Provider interface {
	// Init initializes session provider.
	Init(gclifetime int64, config string) error
	// Read returns raw session store by session ID.
	Read(sid string) (RawStore, error)
	// Exist returns true if session with given ID exists.
	Exist(sid string) bool
	// Destory deletes a session by session ID.
	Destory(sid string) error
	// Regenerate regenerates a session store from old session ID to new one.
	Regenerate(oldsid, sid string) (RawStore, error)
	// Count counts and returns number of sessions.
	Count() int
	// GC calls GC to clean expired sessions.
	GC()
}

var providers = make(map[string]Provider)

// Register registers a provider.
func Register(name string, provider Provider) {
	if provider == nil {
		panic("session: cannot register provider with nil value")
	}
	if _, dup := providers[name]; dup {
		panic(fmt.Errorf("session: cannot register provider '%s' twice", name))
	}
	providers[name] = provider
}

//    _____
//   /     \ _____    ____ _____     ____   ___________
//  /  \ /  \\__  \  /    \\__  \   / ___\_/ __ \_  __ \
// /    Y    \/ __ \|   |  \/ __ \_/ /_/  >  ___/|  | \/
// \____|__  (____  /___|  (____  /\___  / \___  >__|
//         \/     \/     \/     \//_____/      \/

// Manager represents a struct that contains session provider and its configuration.
type Manager struct {
	provider Provider
	opt      Options
}

// NewManager creates and returns a new session manager by given provider name and configuration.
// It panics when given provider isn't registered.
func NewManager(name string, opt Options) (*Manager, error) {
	p, ok := providers[name]
	if !ok {
		return nil, fmt.Errorf("session: unknown provider ‘%q’(forgotten import?)", name)
	}
	if err := p.Init(opt.Maxlifetime, opt.ProviderConfig); err != nil {
		return nil, err
	}
	return &Manager{p, opt}, nil
}

// sessionId generates a new session ID with rand string, unix nano time, remote addr by hash function.
func (m *Manager) sessionId() string {
	return hex.EncodeToString(generateRandomKey(m.opt.IDLength))
}

// Start starts a session by generating new one
// or retrieve existence one by reading session ID from HTTP request if it's valid.
func (m *Manager) Start(ctx *macaron.Context) (RawStore, error) {
	sid := ctx.GetCookie(m.opt.CookieName)
	if len(sid) > 0 && m.provider.Exist(sid) {
		return m.provider.Read(sid)
	}

	sid = m.sessionId()
	sess, err := m.provider.Read(sid)
	if err != nil {
		return nil, err
	}

	cookie := &http.Cookie{
		Name:     m.opt.CookieName,
		Value:    sid,
		Path:     m.opt.CookiePath,
		HttpOnly: true,
		Secure:   m.opt.Secure,
		Domain:   m.opt.Domain,
	}
	if m.opt.CookieLifeTime >= 0 {
		cookie.MaxAge = m.opt.CookieLifeTime
	}
	http.SetCookie(ctx.Resp, cookie)
	ctx.Req.AddCookie(cookie)
	return sess, nil
}

// Read returns raw session store by session ID.
func (m *Manager) Read(sid string) (RawStore, error) {
	return m.provider.Read(sid)
}

// Destory deletes a session by given ID.
func (m *Manager) Destory(ctx *macaron.Context) error {
	sid := ctx.GetCookie(m.opt.CookieName)
	if len(sid) == 0 {
		return nil
	}

	if err := m.provider.Destory(sid); err != nil {
		return err
	}
	cookie := &http.Cookie{
		Name:     m.opt.CookieName,
		Path:     m.opt.CookiePath,
		HttpOnly: true,
		Expires:  time.Now(),
		MaxAge:   -1,
	}
	http.SetCookie(ctx.Resp, cookie)
	return nil
}

// RegenerateId regenerates a session store from old session ID to new one.
func (m *Manager) RegenerateId(ctx *macaron.Context) (sess RawStore, err error) {
	sid := m.sessionId()
	oldsid := ctx.GetCookie(m.opt.CookieName)
	if len(oldsid) == 0 {
		sess, err = m.provider.Read(oldsid)
		if err != nil {
			return nil, err
		}
	} else {
		sess, err = m.provider.Regenerate(oldsid, sid)
		if err != nil {
			return nil, err
		}
	}
	ck := &http.Cookie{
		Name:     m.opt.CookieName,
		Value:    sid,
		Path:     m.opt.CookiePath,
		HttpOnly: true,
		Secure:   m.opt.Secure,
		Domain:   m.opt.Domain,
	}
	if m.opt.CookieLifeTime >= 0 {
		ck.MaxAge = m.opt.CookieLifeTime
	}
	http.SetCookie(ctx.Resp, ck)
	ctx.Req.AddCookie(ck)
	return sess, nil
}

// Count counts and returns number of sessions.
func (m *Manager) Count() int {
	return m.provider.Count()
}

// GC starts GC job in a certain period.
func (m *Manager) GC() {
	m.provider.GC()
}

// startGC starts GC job in a certain period.
func (m *Manager) startGC() {
	m.GC()
	time.AfterFunc(time.Duration(m.opt.Gclifetime)*time.Second, func() { m.startGC() })
}

// SetSecure indicates whether to set cookie with HTTPS or not.
func (m *Manager) SetSecure(secure bool) {
	m.opt.Secure = secure
}

// ___________.____       _____    _________ ___ ___
// \_   _____/|    |     /  _  \  /   _____//   |   \
//  |    __)  |    |    /  /_\  \ \_____  \/    ~    \
//  |     \   |    |___/    |    \/        \    Y    /
//  \___  /   |_______ \____|__  /_______  /\___|_  /
//      \/            \/       \/        \/       \/

type Flash struct {
	ctx *macaron.Context
	url.Values
	ErrorMsg, WarningMsg, InfoMsg, SuccessMsg string
}

func (f *Flash) set(name, msg string, current ...bool) {
	isShow := false
	if (len(current) == 0 && macaron.FlashNow) ||
		(len(current) > 0 && current[0]) {
		isShow = true
	}

	if isShow {
		f.ctx.Data["Flash"] = f
	} else {
		f.Set(name, msg)
	}
}

func (f *Flash) Error(msg string, current ...bool) {
	f.ErrorMsg = msg
	f.set("error", msg, current...)
}

func (f *Flash) Warning(msg string, current ...bool) {
	f.WarningMsg = msg
	f.set("warning", msg, current...)
}

func (f *Flash) Info(msg string, current ...bool) {
	f.InfoMsg = msg
	f.set("info", msg, current...)
}

func (f *Flash) Success(msg string, current ...bool) {
	f.SuccessMsg = msg
	f.set("success", msg, current...)
}
<|endoftext|>


<commit_before>package widget

import (
	"fmt"
	"html/template"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/qor/admin"
	"github.com/qor/qor"
	"github.com/qor/qor/resource"
	"github.com/qor/qor/utils"
	"github.com/qor/serializable_meta"
)

// QorWidgetSettingInterface qor widget setting interface
type QorWidgetSettingInterface interface {
	GetWidgetName() string
	SetWidgetName(string)
	GetGroupName() string
	SetGroupName(string)
	GetScope() string
	SetScope(string)
	GetTemplate() string
	SetTemplate(string)
	GetSourceType() string
	SetSourceType(string)
	GetSourceID() string
	SetSourceID(string)
	GetShared() bool
	SetShared(bool)
	serializable_meta.SerializableMetaInterface
}

// QorWidgetSetting default qor widget setting struct
type QorWidgetSetting struct {
	Name        string `gorm:"primary_key"`
	Scope       string `gorm:"primary_key;size:128;default:'default'"`
	SourceType  string `gorm:"primary_key;default:''"`
	SourceID    string `gorm:"primary_key;default:''"`
	Description string
	Shared      bool
	WidgetType  string
	GroupName   string
	Template    string
	serializable_meta.SerializableMeta
	CreatedAt time.Time
	UpdatedAt time.Time
}

// ResourceName get widget setting's resource name
func (widgetSetting *QorWidgetSetting) ResourceName() string {
	return "Widget Content"
}

// GetSerializableArgumentKind get serializable kind
func (widgetSetting *QorWidgetSetting) GetSerializableArgumentKind() string {
	if widgetSetting.WidgetType != "" {
		return widgetSetting.WidgetType
	}
	return widgetSetting.Kind
}

// SetSerializableArgumentKind set serializable kind
func (widgetSetting *QorWidgetSetting) SetSerializableArgumentKind(name string) {
	widgetSetting.WidgetType = name
	widgetSetting.Kind = name
}

// GetWidgetName get widget setting's group name
func (widgetSetting QorWidgetSetting) GetWidgetName() string {
	return widgetSetting.Name
}

// SetWidgetName set widget setting's group name
func (widgetSetting *QorWidgetSetting) SetWidgetName(name string) {
	widgetSetting.Name = name
}

// GetGroupName get widget setting's group name
func (widgetSetting QorWidgetSetting) GetGroupName() string {
	return widgetSetting.GroupName
}

// SetGroupName set widget setting's group name
func (widgetSetting *QorWidgetSetting) SetGroupName(groupName string) {
	widgetSetting.GroupName = groupName
}

// GetScope get widget's scope
func (widgetSetting QorWidgetSetting) GetScope() string {
	return widgetSetting.Scope
}

// SetScope set widget setting's scope
func (widgetSetting *QorWidgetSetting) SetScope(scope string) {
	widgetSetting.Scope = scope
}

// GetSourceType get widget's source type
func (widgetSetting QorWidgetSetting) GetSourceType() string {
	return widgetSetting.SourceID
}

// SetSourceType set widget setting's souce type
func (widgetSetting *QorWidgetSetting) SetSourceType(sourceType string) {
	widgetSetting.SourceType = sourceType
}

// GetSourceID get widget's source ID
func (widgetSetting QorWidgetSetting) GetSourceID() string {
	return widgetSetting.SourceID
}

// SetSourceID set widget setting's source id
func (widgetSetting *QorWidgetSetting) SetSourceID(sourceID string) {
	widgetSetting.SourceID = sourceID
}

// GetShared get widget's source ID
func (widgetSetting QorWidgetSetting) GetShared() bool {
	return widgetSetting.Shared
}

// SetShared set widget setting's source id
func (widgetSetting *QorWidgetSetting) SetShared(shared bool) {
	widgetSetting.Shared = shared
}

// GetTemplate get used widget template
func (widgetSetting QorWidgetSetting) GetTemplate() string {
	if widget := GetWidget(widgetSetting.GetSerializableArgumentKind()); widget != nil {
		for _, value := range widget.Templates {
			if value == widgetSetting.Template {
				return value
			}
		}

		// return first value of defined widget templates
		for _, value := range widget.Templates {
			return value
		}
	}
	return ""
}

// SetTemplate set used widget's template
func (widgetSetting *QorWidgetSetting) SetTemplate(template string) {
	widgetSetting.Template = template
}

// GetSerializableArgumentResource get setting's argument's resource
func (widgetSetting *QorWidgetSetting) GetSerializableArgumentResource() *admin.Resource {
	widget := GetWidget(widgetSetting.GetSerializableArgumentKind())
	if widget != nil {
		return widget.Setting
	}
	return nil
}

// ConfigureQorResource a method used to config Widget for qor admin
func (widgetSetting *QorWidgetSetting) ConfigureQorResource(res resource.Resourcer) {
	if res, ok := res.(*admin.Resource); ok {
		if res.GetMeta("Name") == nil {
			res.Meta(&admin.Meta{Name: "Name"})
		}

		if res.GetMeta("PreviewIcon") == nil {
			res.Meta(&admin.Meta{Name: "PreviewIcon", Valuer: func(result interface{}, context *qor.Context) interface{} {
				if setting, ok := result.(QorWidgetSettingInterface); ok {
					if widget := GetWidget(setting.GetSerializableArgumentKind()); widget != nil {
						return template.HTML(fmt.Sprintf("", widget.PreviewIcon))
					}
				}
				return ""
			}})
		}

		if res.GetMeta("DisplayName") == nil {
			res.Meta(&admin.Meta{Name: "DisplayName", Label: "Name", Type: "readonly", FieldName: "Name"})
		}

		if res.GetMeta("Scope") == nil {
			res.Meta(&admin.Meta{
				Name: "Scope",
				Type: "hidden",
				Valuer: func(result interface{}, context *qor.Context) interface{} {
					if scope := context.Request.URL.Query().Get("widget_scope"); scope != "" {
						return scope
					}

					if setting, ok := result.(QorWidgetSettingInterface); ok {
						if scope := setting.GetScope(); scope != "" {
							return scope
						}
					}

					return "default"
				},
				Setter: func(result interface{}, metaValue *resource.MetaValue, context *qor.Context) {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						setting.SetScope(utils.ToString(metaValue.Value))
					}
				},
			})
		}

		if res.GetMeta("SourceType") == nil {
			res.Meta(&admin.Meta{
				Name: "SourceType",
				Type: "hidden",
				Valuer: func(result interface{}, context *qor.Context) interface{} {
					if sourceType := context.Request.URL.Query().Get("widget_source_type"); sourceType != "" {
						return sourceType
					}

					if setting, ok := result.(QorWidgetSettingInterface); ok {
						if sourceType := setting.GetSourceType(); sourceType != "" {
							return sourceType
						}
					}
					return ""
				},
				Setter: func(result interface{}, metaValue *resource.MetaValue, context *qor.Context) {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						setting.SetSourceType(utils.ToString(metaValue.Value))
					}
				},
			})
		}

		if res.GetMeta("SourceID") == nil {
			res.Meta(&admin.Meta{
				Name: "SourceID",
				Type: "hidden",
				Valuer: func(result interface{}, context *qor.Context) interface{} {
					if sourceID := context.Request.URL.Query().Get("widget_source_id"); sourceID != "" {
						return sourceID
					}

					if setting, ok := result.(QorWidgetSettingInterface); ok {
						if sourceID := setting.GetSourceID(); sourceID != "" {
							return sourceID
						}
					}
					return ""
				},
				Setter: func(result interface{}, metaValue *resource.MetaValue, context *qor.Context) {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						setting.SetSourceID(utils.ToString(metaValue.Value))
					}
				},
			})
		}

		if res.GetMeta("Widgets") == nil {
			res.Meta(&admin.Meta{
				Name: "Widgets",
				Type: "select_one",
				Valuer: func(result interface{}, context *qor.Context) interface{} {
					if typ := context.Request.URL.Query().Get("widget_type"); typ != "" {
						return typ
					}

					if setting, ok := result.(QorWidgetSettingInterface); ok {
						widget := GetWidget(setting.GetSerializableArgumentKind())
						if widget == nil {
							return ""
						}
						return widget.Name
					}

					return ""
				},
				Collection: func(result interface{}, context *qor.Context) (results [][]string) {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						if setting.GetWidgetName() == "" {
							for _, widget := range registeredWidgets {
								results = append(results, []string{widget.Name, widget.Name})
							}
						} else {
							groupName := setting.GetGroupName()
							for _, group := range registeredWidgetsGroup {
								if group.Name == groupName {
									for _, widget := range group.Widgets {
										results = append(results, []string{widget, widget})
									}
								}
							}
						}

						if len(results) == 0 {
							results = append(results, []string{setting.GetSerializableArgumentKind(), setting.GetSerializableArgumentKind()})
						}
					}
					return
				},
				Setter: func(result interface{}, metaValue *resource.MetaValue, context *qor.Context) {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						setting.SetSerializableArgumentKind(utils.ToString(metaValue.Value))
					}
				},
			})
		}

		if res.GetMeta("Template") == nil {
			res.Meta(&admin.Meta{
				Name: "Template",
				Type: "select_one",
				Valuer: func(result interface{}, context *qor.Context) interface{} {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						return setting.GetTemplate()
					}
					return ""
				},
				Collection: func(result interface{}, context *qor.Context) (results [][]string) {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						if widget := GetWidget(setting.GetSerializableArgumentKind()); widget != nil {
							for _, value := range widget.Templates {
								results = append(results, []string{value, value})
							}
						}
					}
					return
				},
				Setter: func(result interface{}, metaValue *resource.MetaValue, context *qor.Context) {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						setting.SetTemplate(utils.ToString(metaValue.Value))
					}
				},
			})
		}

		if res.GetMeta("Shared") == nil {
			res.Meta(&admin.Meta{
				Name:  "Shared",
				Label: "This widget is shared",
			})
		}

		res.Scope(&admin.Scope{
			Name:  "Shared",
			Label: "Shared Widgets",
			Handle: func(db *gorm.DB, _ *qor.Context) *gorm.DB {
				return db.Where("shared = ?", true)
			},
		})

		res.Action(&admin.Action{
			Name: "Preview",
			URL: func(record interface{}, context *admin.Context) string {
				return fmt.Sprintf("%v/%v/%v/!preview", context.Admin.GetRouter().Prefix, res.ToParam(), record.(QorWidgetSettingInterface).GetWidgetName())
			},
			Modes: []string{"edit", "menu_item"},
		})

		res.AddProcessor(func(value interface{}, metaValues *resource.MetaValues, context *qor.Context) error {
			if widgetSetting, ok := value.(QorWidgetSettingInterface); ok {
				if widgetSetting.GetShared() {
					widgetSetting.SetSourceType("")
					widgetSetting.SetSourceID("")
				}
			}
			return nil
		})

		res.UseTheme("widget")

		res.IndexAttrs("PreviewIcon", "Name", "Description", "CreatedAt", "UpdatedAt")
		res.ShowAttrs("PreviewIcon", "Name", "Scope", "WidgetType", "Template", "Description", "Value", "CreatedAt", "UpdatedAt")
		res.EditAttrs(
			"DisplayName", "Description", "Scope", "Widgets", "Template",
			&admin.Section{
				Title: "Settings",
				Rows:  [][]string{{"Kind"}, {"SerializableMeta"}},
			},
			"Shared", "SourceType", "SourceID",
		)
		res.NewAttrs("Name", "Description", "Scope", "Widgets", "Template",
			&admin.Section{
				Title: "Settings",
				Rows:  [][]string{{"Kind"}, {"SerializableMeta"}},
			},
			"Shared", "SourceType", "SourceID",
		)

		searchHandler := res.SearchHandler
		res.SearchHandler = func(keyword string, context *qor.Context) *gorm.DB {
			context.SetDB(context.GetDB().Where("source_type = ?", ""))
			return searchHandler(keyword, context)
		}
	}
}
<commit_msg>Fix typo when get source type<commit_after>package widget

import (
	"fmt"
	"html/template"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/qor/admin"
	"github.com/qor/qor"
	"github.com/qor/qor/resource"
	"github.com/qor/qor/utils"
	"github.com/qor/serializable_meta"
)

// QorWidgetSettingInterface qor widget setting interface
type QorWidgetSettingInterface interface {
	GetWidgetName() string
	SetWidgetName(string)
	GetGroupName() string
	SetGroupName(string)
	GetScope() string
	SetScope(string)
	GetTemplate() string
	SetTemplate(string)
	GetSourceType() string
	SetSourceType(string)
	GetSourceID() string
	SetSourceID(string)
	GetShared() bool
	SetShared(bool)
	serializable_meta.SerializableMetaInterface
}

// QorWidgetSetting default qor widget setting struct
type QorWidgetSetting struct {
	Name        string `gorm:"primary_key"`
	Scope       string `gorm:"primary_key;size:128;default:'default'"`
	SourceType  string `gorm:"primary_key;default:''"`
	SourceID    string `gorm:"primary_key;default:''"`
	Description string
	Shared      bool
	WidgetType  string
	GroupName   string
	Template    string
	serializable_meta.SerializableMeta
	CreatedAt time.Time
	UpdatedAt time.Time
}

// ResourceName get widget setting's resource name
func (widgetSetting *QorWidgetSetting) ResourceName() string {
	return "Widget Content"
}

// GetSerializableArgumentKind get serializable kind
func (widgetSetting *QorWidgetSetting) GetSerializableArgumentKind() string {
	if widgetSetting.WidgetType != "" {
		return widgetSetting.WidgetType
	}
	return widgetSetting.Kind
}

// SetSerializableArgumentKind set serializable kind
func (widgetSetting *QorWidgetSetting) SetSerializableArgumentKind(name string) {
	widgetSetting.WidgetType = name
	widgetSetting.Kind = name
}

// GetWidgetName get widget setting's group name
func (widgetSetting QorWidgetSetting) GetWidgetName() string {
	return widgetSetting.Name
}

// SetWidgetName set widget setting's group name
func (widgetSetting *QorWidgetSetting) SetWidgetName(name string) {
	widgetSetting.Name = name
}

// GetGroupName get widget setting's group name
func (widgetSetting QorWidgetSetting) GetGroupName() string {
	return widgetSetting.GroupName
}

// SetGroupName set widget setting's group name
func (widgetSetting *QorWidgetSetting) SetGroupName(groupName string) {
	widgetSetting.GroupName = groupName
}

// GetScope get widget's scope
func (widgetSetting QorWidgetSetting) GetScope() string {
	return widgetSetting.Scope
}

// SetScope set widget setting's scope
func (widgetSetting *QorWidgetSetting) SetScope(scope string) {
	widgetSetting.Scope = scope
}

// GetSourceType get widget's source type
func (widgetSetting QorWidgetSetting) GetSourceType() string {
	return widgetSetting.SourceType
}

// SetSourceType set widget setting's souce type
func (widgetSetting *QorWidgetSetting) SetSourceType(sourceType string) {
	widgetSetting.SourceType = sourceType
}

// GetSourceID get widget's source ID
func (widgetSetting QorWidgetSetting) GetSourceID() string {
	return widgetSetting.SourceID
}

// SetSourceID set widget setting's source id
func (widgetSetting *QorWidgetSetting) SetSourceID(sourceID string) {
	widgetSetting.SourceID = sourceID
}

// GetShared get widget's source ID
func (widgetSetting QorWidgetSetting) GetShared() bool {
	return widgetSetting.Shared
}

// SetShared set widget setting's source id
func (widgetSetting *QorWidgetSetting) SetShared(shared bool) {
	widgetSetting.Shared = shared
}

// GetTemplate get used widget template
func (widgetSetting QorWidgetSetting) GetTemplate() string {
	if widget := GetWidget(widgetSetting.GetSerializableArgumentKind()); widget != nil {
		for _, value := range widget.Templates {
			if value == widgetSetting.Template {
				return value
			}
		}

		// return first value of defined widget templates
		for _, value := range widget.Templates {
			return value
		}
	}
	return ""
}

// SetTemplate set used widget's template
func (widgetSetting *QorWidgetSetting) SetTemplate(template string) {
	widgetSetting.Template = template
}

// GetSerializableArgumentResource get setting's argument's resource
func (widgetSetting *QorWidgetSetting) GetSerializableArgumentResource() *admin.Resource {
	widget := GetWidget(widgetSetting.GetSerializableArgumentKind())
	if widget != nil {
		return widget.Setting
	}
	return nil
}

// ConfigureQorResource a method used to config Widget for qor admin
func (widgetSetting *QorWidgetSetting) ConfigureQorResource(res resource.Resourcer) {
	if res, ok := res.(*admin.Resource); ok {
		if res.GetMeta("Name") == nil {
			res.Meta(&admin.Meta{Name: "Name"})
		}

		if res.GetMeta("PreviewIcon") == nil {
			res.Meta(&admin.Meta{Name: "PreviewIcon", Valuer: func(result interface{}, context *qor.Context) interface{} {
				if setting, ok := result.(QorWidgetSettingInterface); ok {
					if widget := GetWidget(setting.GetSerializableArgumentKind()); widget != nil {
						return template.HTML(fmt.Sprintf("", widget.PreviewIcon))
					}
				}
				return ""
			}})
		}

		if res.GetMeta("DisplayName") == nil {
			res.Meta(&admin.Meta{Name: "DisplayName", Label: "Name", Type: "readonly", FieldName: "Name"})
		}

		if res.GetMeta("Scope") == nil {
			res.Meta(&admin.Meta{
				Name: "Scope",
				Type: "hidden",
				Valuer: func(result interface{}, context *qor.Context) interface{} {
					if scope := context.Request.URL.Query().Get("widget_scope"); scope != "" {
						return scope
					}

					if setting, ok := result.(QorWidgetSettingInterface); ok {
						if scope := setting.GetScope(); scope != "" {
							return scope
						}
					}

					return "default"
				},
				Setter: func(result interface{}, metaValue *resource.MetaValue, context *qor.Context) {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						setting.SetScope(utils.ToString(metaValue.Value))
					}
				},
			})
		}

		if res.GetMeta("SourceType") == nil {
			res.Meta(&admin.Meta{
				Name: "SourceType",
				Type: "hidden",
				Valuer: func(result interface{}, context *qor.Context) interface{} {
					if sourceType := context.Request.URL.Query().Get("widget_source_type"); sourceType != "" {
						return sourceType
					}

					if setting, ok := result.(QorWidgetSettingInterface); ok {
						if sourceType := setting.GetSourceType(); sourceType != "" {
							return sourceType
						}
					}
					return ""
				},
				Setter: func(result interface{}, metaValue *resource.MetaValue, context *qor.Context) {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						setting.SetSourceType(utils.ToString(metaValue.Value))
					}
				},
			})
		}

		if res.GetMeta("SourceID") == nil {
			res.Meta(&admin.Meta{
				Name: "SourceID",
				Type: "hidden",
				Valuer: func(result interface{}, context *qor.Context) interface{} {
					if sourceID := context.Request.URL.Query().Get("widget_source_id"); sourceID != "" {
						return sourceID
					}

					if setting, ok := result.(QorWidgetSettingInterface); ok {
						if sourceID := setting.GetSourceID(); sourceID != "" {
							return sourceID
						}
					}
					return ""
				},
				Setter: func(result interface{}, metaValue *resource.MetaValue, context *qor.Context) {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						setting.SetSourceID(utils.ToString(metaValue.Value))
					}
				},
			})
		}

		if res.GetMeta("Widgets") == nil {
			res.Meta(&admin.Meta{
				Name: "Widgets",
				Type: "select_one",
				Valuer: func(result interface{}, context *qor.Context) interface{} {
					if typ := context.Request.URL.Query().Get("widget_type"); typ != "" {
						return typ
					}

					if setting, ok := result.(QorWidgetSettingInterface); ok {
						widget := GetWidget(setting.GetSerializableArgumentKind())
						if widget == nil {
							return ""
						}
						return widget.Name
					}

					return ""
				},
				Collection: func(result interface{}, context *qor.Context) (results [][]string) {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						if setting.GetWidgetName() == "" {
							for _, widget := range registeredWidgets {
								results = append(results, []string{widget.Name, widget.Name})
							}
						} else {
							groupName := setting.GetGroupName()
							for _, group := range registeredWidgetsGroup {
								if group.Name == groupName {
									for _, widget := range group.Widgets {
										results = append(results, []string{widget, widget})
									}
								}
							}
						}

						if len(results) == 0 {
							results = append(results, []string{setting.GetSerializableArgumentKind(), setting.GetSerializableArgumentKind()})
						}
					}
					return
				},
				Setter: func(result interface{}, metaValue *resource.MetaValue, context *qor.Context) {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						setting.SetSerializableArgumentKind(utils.ToString(metaValue.Value))
					}
				},
			})
		}

		if res.GetMeta("Template") == nil {
			res.Meta(&admin.Meta{
				Name: "Template",
				Type: "select_one",
				Valuer: func(result interface{}, context *qor.Context) interface{} {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						return setting.GetTemplate()
					}
					return ""
				},
				Collection: func(result interface{}, context *qor.Context) (results [][]string) {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						if widget := GetWidget(setting.GetSerializableArgumentKind()); widget != nil {
							for _, value := range widget.Templates {
								results = append(results, []string{value, value})
							}
						}
					}
					return
				},
				Setter: func(result interface{}, metaValue *resource.MetaValue, context *qor.Context) {
					if setting, ok := result.(QorWidgetSettingInterface); ok {
						setting.SetTemplate(utils.ToString(metaValue.Value))
					}
				},
			})
		}

		if res.GetMeta("Shared") == nil {
			res.Meta(&admin.Meta{
				Name:  "Shared",
				Label: "This widget is shared",
			})
		}

		res.Scope(&admin.Scope{
			Name:  "Shared",
			Label: "Shared Widgets",
			Handle: func(db *gorm.DB, _ *qor.Context) *gorm.DB {
				return db.Where("shared = ?", true)
			},
		})

		res.Action(&admin.Action{
			Name: "Preview",
			URL: func(record interface{}, context *admin.Context) string {
				return fmt.Sprintf("%v/%v/%v/!preview", context.Admin.GetRouter().Prefix, res.ToParam(), record.(QorWidgetSettingInterface).GetWidgetName())
			},
			Modes: []string{"edit", "menu_item"},
		})

		res.AddProcessor(func(value interface{}, metaValues *resource.MetaValues, context *qor.Context) error {
			if widgetSetting, ok := value.(QorWidgetSettingInterface); ok {
				if widgetSetting.GetShared() {
					widgetSetting.SetSourceType("")
					widgetSetting.SetSourceID("")
				}
			}
			return nil
		})

		res.UseTheme("widget")

		res.IndexAttrs("PreviewIcon", "Name", "Description", "CreatedAt", "UpdatedAt")
		res.ShowAttrs("PreviewIcon", "Name", "Scope", "WidgetType", "Template", "Description", "Value", "CreatedAt", "UpdatedAt")
		res.EditAttrs(
			"DisplayName", "Description", "Scope", "Widgets", "Template",
			&admin.Section{
				Title: "Settings",
				Rows:  [][]string{{"Kind"}, {"SerializableMeta"}},
			},
			"Shared", "SourceType", "SourceID",
		)
		res.NewAttrs("Name", "Description", "Scope", "Widgets", "Template",
			&admin.Section{
				Title: "Settings",
				Rows:  [][]string{{"Kind"}, {"SerializableMeta"}},
			},
			"Shared", "SourceType", "SourceID",
		)

		searchHandler := res.SearchHandler
		res.SearchHandler = func(keyword string, context *qor.Context) *gorm.DB {
			context.SetDB(context.GetDB().Where("source_type = ?", ""))
			return searchHandler(keyword, context)
		}
	}
}
<|endoftext|>


<commit_before><commit_msg>Go fmt<commit_after><|endoftext|>