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

package app

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"sort"

	"golang.org/x/net/context"
	"golang.org/x/perf/storage/benchfmt"
	"golang.org/x/perf/storage/db"
)

// upload is the handler for the /upload endpoint. It serves a form on
// GET requests and processes files in a multipart/x-form-data POST
// request.
func (a *App) upload(w http.ResponseWriter, r *http.Request) {
	ctx := requestContext(r)

	user, err := a.Auth(w, r)
	switch {
	case err == ErrResponseWritten:
		return
	case err != nil:
		http.Error(w, err.Error(), 500)
	}

	if r.Method == http.MethodGet {
		http.ServeFile(w, r, "static/upload.html")
		return
	}
	if r.Method != http.MethodPost {
		http.Error(w, "/upload must be called as a POST request", http.StatusMethodNotAllowed)
		return
	}

	// We use r.MultipartReader instead of r.ParseForm to avoid
	// storing uploaded data in memory.
	mr, err := r.MultipartReader()
	if err != nil {
		errorf(ctx, "%v", err)
		http.Error(w, err.Error(), 500)
		return
	}

	result, err := a.processUpload(ctx, user, mr)
	if err != nil {
		errorf(ctx, "%v", err)
		http.Error(w, err.Error(), 500)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(result); err != nil {
		errorf(ctx, "%v", err)
		http.Error(w, err.Error(), 500)
		return
	}
}

// uploadStatus is the response to an /upload POST served as JSON.
type uploadStatus struct {
	// UploadID is the upload ID assigned to the upload.
	UploadID string `json:"uploadid"`
	// FileIDs is the list of file IDs assigned to the files in the upload.
	FileIDs []string `json:"fileids"`
	// ViewURL is a URL that can be used to interactively view the upload.
	ViewURL string `json:"viewurl,omitempty"`
}

// processUpload takes one or more files from a multipart.Reader,
// writes them to the filesystem, and indexes their content.
func (a *App) processUpload(ctx context.Context, user string, mr *multipart.Reader) (*uploadStatus, error) {
	var upload *db.Upload
	var fileids []string

	for i := 0; ; i++ {
		p, err := mr.NextPart()
		if err == io.EOF {
			break
		}

		name := p.FormName()
		if name != "file" {
			return nil, fmt.Errorf("unexpected field %q", name)
		}

		if upload == nil {
			var err error
			upload, err = a.DB.NewUpload(ctx)
			if err != nil {
				return nil, err
			}
			defer func() {
				if upload != nil {
					upload.Abort()
				}
			}()
		}

		// The incoming file needs to be stored in Cloud
		// Storage and it also needs to be indexed. If the file
		// is invalid (contains no valid records) it needs to
		// be rejected and the Cloud Storage upload aborted.

		meta := fileMetadata(user, upload.ID, i)

		// We need to do two things with the incoming data:
		// - Write it to permanent storage via a.FS
		// - Write index records to a.DB
		// AND if anything fails, attempt to clean up both the
		// FS and the index records.

		if err := a.indexFile(ctx, upload, p, meta); err != nil {
			return nil, err
		}

		fileids = append(fileids, meta["fileid"])
	}

	if upload == nil {
		return nil, errors.New("no files processed")
	}
	if err := upload.Commit(); err != nil {
		return nil, err
	}

	status := &uploadStatus{UploadID: upload.ID, FileIDs: fileids}
	if a.ViewURLBase != "" {
		status.ViewURL = a.ViewURLBase + url.QueryEscape(upload.ID)
	}

	upload = nil

	return status, nil
}

func (a *App) indexFile(ctx context.Context, upload *db.Upload, p io.Reader, meta map[string]string) (err error) {
	fw, err := a.FS.NewWriter(ctx, fmt.Sprintf("uploads/%s.txt", meta["fileid"]), meta)
	if err != nil {
		return err
	}
	defer func() {
		if err != nil {
			fw.CloseWithError(err)
		} else {
			err = fw.Close()
		}
	}()
	var keys []string
	for k := range meta {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		if _, err := fmt.Fprintf(fw, "%s: %s\n", k, meta[k]); err != nil {
			return err
		}
	}

	// TODO(quentin): Add a separate goroutine and buffer for writes to fw?
	tr := io.TeeReader(p, fw)
	br := benchfmt.NewReader(tr)
	br.AddLabels(meta)
	i := 0
	for {
		result, err := br.Next()
		if err != nil {
			if err != io.EOF {
				return err
			}
			if i == 0 {
				return errors.New("no valid benchmark lines found")
			}
			return nil
		}
		i++
		if err := upload.InsertRecord(result); err != nil {
			return err
		}
	}
}

// fileMetadata returns the extra metadata fields associated with an
// uploaded file.
func fileMetadata(user string, uploadid string, filenum int) map[string]string {
	// TODO(quentin): Add the upload time.
	// TODO(quentin): Add other fields?
	m := map[string]string{
		"uploadid": uploadid,
		"fileid":   fmt.Sprintf("%s/%d", uploadid, filenum),
	}
	if user != "" {
		m["by"] = user
	}
	return m
}
<commit_msg>storage/app: add additional metadata keys<commit_after>// Copyright 2016 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package app

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/perf/storage/benchfmt"
	"golang.org/x/perf/storage/db"
)

// upload is the handler for the /upload endpoint. It serves a form on
// GET requests and processes files in a multipart/x-form-data POST
// request.
func (a *App) upload(w http.ResponseWriter, r *http.Request) {
	ctx := requestContext(r)

	user, err := a.Auth(w, r)
	switch {
	case err == ErrResponseWritten:
		return
	case err != nil:
		http.Error(w, err.Error(), 500)
	}

	if r.Method == http.MethodGet {
		http.ServeFile(w, r, "static/upload.html")
		return
	}
	if r.Method != http.MethodPost {
		http.Error(w, "/upload must be called as a POST request", http.StatusMethodNotAllowed)
		return
	}

	// We use r.MultipartReader instead of r.ParseForm to avoid
	// storing uploaded data in memory.
	mr, err := r.MultipartReader()
	if err != nil {
		errorf(ctx, "%v", err)
		http.Error(w, err.Error(), 500)
		return
	}

	result, err := a.processUpload(ctx, user, mr)
	if err != nil {
		errorf(ctx, "%v", err)
		http.Error(w, err.Error(), 500)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(result); err != nil {
		errorf(ctx, "%v", err)
		http.Error(w, err.Error(), 500)
		return
	}
}

// uploadStatus is the response to an /upload POST served as JSON.
type uploadStatus struct {
	// UploadID is the upload ID assigned to the upload.
	UploadID string `json:"uploadid"`
	// FileIDs is the list of file IDs assigned to the files in the upload.
	FileIDs []string `json:"fileids"`
	// ViewURL is a URL that can be used to interactively view the upload.
	ViewURL string `json:"viewurl,omitempty"`
}

// processUpload takes one or more files from a multipart.Reader,
// writes them to the filesystem, and indexes their content.
func (a *App) processUpload(ctx context.Context, user string, mr *multipart.Reader) (*uploadStatus, error) {
	var upload *db.Upload
	var fileids []string

	uploadtime := time.Now().UTC().Format(time.RFC3339)

	for i := 0; ; i++ {
		p, err := mr.NextPart()
		if err == io.EOF {
			break
		}

		name := p.FormName()
		if name != "file" {
			return nil, fmt.Errorf("unexpected field %q", name)
		}

		if upload == nil {
			var err error
			upload, err = a.DB.NewUpload(ctx)
			if err != nil {
				return nil, err
			}
			defer func() {
				if upload != nil {
					upload.Abort()
				}
			}()
		}

		// The incoming file needs to be stored in Cloud
		// Storage and it also needs to be indexed. If the file
		// is invalid (contains no valid records) it needs to
		// be rejected and the Cloud Storage upload aborted.

		meta := map[string]string{
			"uploadid":   upload.ID,
			"fileid":     fmt.Sprintf("%s/%d", upload.ID, filenum),
			"uploadtime": uploadtime,
		}
		name := p.FileName()
		if slash := strings.LastIndexAny(name, `/\`); slash >= 0 {
			name = name[slash+1:]
		}
		if name != "" {
			m["uploadfile"] = name
		}
		if user != "" {
			m["by"] = user
		}

		// We need to do two things with the incoming data:
		// - Write it to permanent storage via a.FS
		// - Write index records to a.DB
		// AND if anything fails, attempt to clean up both the
		// FS and the index records.

		if err := a.indexFile(ctx, upload, p, meta); err != nil {
			return nil, err
		}

		fileids = append(fileids, meta["fileid"])
	}

	if upload == nil {
		return nil, errors.New("no files processed")
	}
	if err := upload.Commit(); err != nil {
		return nil, err
	}

	status := &uploadStatus{UploadID: upload.ID, FileIDs: fileids}
	if a.ViewURLBase != "" {
		status.ViewURL = a.ViewURLBase + url.QueryEscape(upload.ID)
	}

	upload = nil

	return status, nil
}

func (a *App) indexFile(ctx context.Context, upload *db.Upload, p io.Reader, meta map[string]string) (err error) {
	fw, err := a.FS.NewWriter(ctx, fmt.Sprintf("uploads/%s.txt", meta["fileid"]), meta)
	if err != nil {
		return err
	}
	defer func() {
		if err != nil {
			fw.CloseWithError(err)
		} else {
			err = fw.Close()
		}
	}()
	var keys []string
	for k := range meta {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		if _, err := fmt.Fprintf(fw, "%s: %s\n", k, meta[k]); err != nil {
			return err
		}
	}
	// Write a blank line to separate metadata from user-generated content.
	fmt.Fprintf(fw, "\n")

	// TODO(quentin): Add a separate goroutine and buffer for writes to fw?
	tr := io.TeeReader(p, fw)
	br := benchfmt.NewReader(tr)
	br.AddLabels(meta)
	i := 0
	for {
		result, err := br.Next()
		if err != nil {
			if err != io.EOF {
				return err
			}
			if i == 0 {
				return errors.New("no valid benchmark lines found")
			}
			return nil
		}
		i++
		if err := upload.InsertRecord(result); err != nil {
			return err
		}
	}
}
<|endoftext|>


<commit_before>/*

bolt provides a bolt-backed database that implements both
boardgame.StorageManager and boardgame/server.StorageManager.

*/
package mysql

import (
	"database/sql"
	"errors"
	"github.com/go-gorp/gorp"
	"github.com/go-sql-driver/mysql"
	"github.com/jkomoros/boardgame"
	"github.com/jkomoros/boardgame/server/api/users"
	"log"
	"strings"
)

const (
	TableGames       = "games"
	TableUsers       = "users"
	TableStates      = "states"
	TableCookies     = "cookies"
	TablePlayers     = "players"
	TableAgentStates = "agentstates"
)

type StorageManager struct {
	db    *sql.DB
	dbMap *gorp.DbMap
	//If in test mode we'll... mock stuff, I guess?
	testMode bool
	//The config string that we were provided in connect.
	config string
}

func NewStorageManager(testMode bool) *StorageManager {
	//We actually don't do much; we do more of our work in Connect()
	return &StorageManager{
		testMode: testMode,
	}
}

func getDSN(config string) (string, error) {

	//Substantially recreated in boardgame-mysqL-admin

	parsedDSN, err := mysql.ParseDSN(config)

	if err != nil {
		return "", errors.New("config provided was not valid DSN: " + err.Error())
	}

	parsedDSN.Collation = "utf8mb4_unicode_ci"
	parsedDSN.MultiStatements = true

	return parsedDSN.FormatDSN(), nil
}

func (s *StorageManager) Connect(config string) error {

	configToUse, err := getDSN(config)

	if err != nil {
		return err
	}

	db, err := sql.Open("mysql", configToUse)
	if err != nil {
		return errors.New("Unable to open database: " + err.Error())
	}

	s.config = config

	s.db = db

	s.dbMap = &gorp.DbMap{
		Db: db,
		Dialect: gorp.MySQLDialect{
			Engine: "InnoDB",
			//the mb4 is necessary to support e.g. emojis
			Encoding: "utf8mb4",
		},
	}

	s.dbMap.AddTableWithName(UserStorageRecord{}, TableUsers).SetKeys(false, "Id")
	s.dbMap.AddTableWithName(GameStorageRecord{}, TableGames).SetKeys(false, "Id")
	s.dbMap.AddTableWithName(StateStorageRecord{}, TableStates).SetKeys(true, "Id")
	s.dbMap.AddTableWithName(CookieStorageRecord{}, TableCookies).SetKeys(false, "Cookie")
	s.dbMap.AddTableWithName(PlayerStorageRecord{}, TablePlayers).SetKeys(true, "Id")
	s.dbMap.AddTableWithName(AgentStateStorageRecord{}, TableAgentStates).SetKeys(true, "Id")

	//TODO: sanity check that the tables exist

	_, err = s.dbMap.SelectInt("select count(*) from " + TableGames)

	if err != nil {
		return errors.New("Sanity check failed for db. Have you used the admin tool to migrate it up? " + err.Error())
	}

	return nil

}

func (s *StorageManager) Close() {
	if s.db == nil {
		return
	}
	s.db.Close()
	s.db = nil
}

//Cleanup will only drop tables if we're in test mode, and the config string
//used to open the database talked about a test database on localhost (as
//sanity check).
func (s *StorageManager) CleanUp() {

	if !s.testMode {
		return
	}

	if !strings.Contains(s.config, "_test") {
		log.Println("Sanity check on boardgame config before cleanup didn't find _test")
		return
	}

	if !strings.Contains(s.config, "localhost") {
		log.Println("Sanity check on boardgame config before cleanup didn't find localhost")
	}

	if s.db == nil {
		log.Println("Couldn't clean up; db already closed")
		return
	}

	log.Println("Sanity checks passed. Dropping tables to cleanup...")

	if err := s.dbMap.DropTables(); err != nil {
		log.Println("Error dropping tables:", err)
		return
	}
}

func (s *StorageManager) Name() string {
	return "mysql"
}

func (s *StorageManager) State(gameId string, version int) (boardgame.StateStorageRecord, error) {
	var state StateStorageRecord

	err := s.dbMap.SelectOne(&state, "select * from "+TableStates+" where GameId=? and Version=?", gameId, version)

	if err == sql.ErrNoRows {
		return nil, errors.New("No such state")
	}

	if err != nil {
		return nil, errors.New("Unexpected error: " + err.Error())
	}

	return (&state).ToStorageRecord(), nil
}

func (s *StorageManager) Game(id string) (*boardgame.GameStorageRecord, error) {
	var game GameStorageRecord

	err := s.dbMap.SelectOne(&game, "select * from "+TableGames+" where Id=?", id)

	if err == sql.ErrNoRows {
		return nil, errors.New("No such game")
	}

	if err != nil {
		return nil, errors.New("Unexpected error: " + err.Error())
	}

	return (&game).ToStorageRecord(), nil
}

func (s *StorageManager) SaveGameAndCurrentState(game *boardgame.GameStorageRecord, state boardgame.StateStorageRecord) error {

	version := game.Version

	gameRecord := NewGameStorageRecord(game)
	stateRecord := NewStateStorageRecord(game.Id, version, state)

	count, _ := s.dbMap.SelectInt("select count(*) from "+TableGames+" where Id=?", game.Id)

	if count < 1 {
		//Need to insert
		err := s.dbMap.Insert(gameRecord)

		if err != nil {

			return errors.New("Couldn't update game: " + err.Error())

		}

	} else {
		//Need to update
		_, err := s.dbMap.Update(gameRecord)

		if err != nil {
			return errors.New("Couldn't insert game: " + err.Error())
		}
	}

	err := s.dbMap.Insert(stateRecord)

	if err != nil {
		return errors.New("Couldn't insert state: " + err.Error())
	}

	return nil
}

func (s *StorageManager) AgentState(gameId string, player boardgame.PlayerIndex) ([]byte, error) {

	var agent AgentStateStorageRecord

	err := s.dbMap.SelectOne(&agent, "select * from "+TableAgentStates+" where GameId=? and PlayerIndex=? order by Id desc limit 1", gameId, int64(player))

	if err == sql.ErrNoRows {
		return nil, nil
	}

	if err != nil {
		return nil, err
	}

	return agent.ToStorageRecord(), nil

}

func (s *StorageManager) SaveAgentState(gameId string, player boardgame.PlayerIndex, state []byte) error {
	record := NewAgentStateStorageRecord(gameId, player, state)

	err := s.dbMap.Insert(record)

	if err != nil {
		return errors.New("Couldn't save record: " + err.Error())
	}

	return nil
}

func (s *StorageManager) ListGames(max int) []*boardgame.GameStorageRecord {
	var games []GameStorageRecord

	if max < 1 {
		max = 100
	}

	if _, err := s.dbMap.Select(&games, "select * from "+TableGames+" limit ?", max); err != nil {
		return nil
	}

	result := make([]*boardgame.GameStorageRecord, len(games))

	for i, record := range games {
		result[i] = (&record).ToStorageRecord()
	}

	return result
}

func (s *StorageManager) SetPlayerForGame(gameId string, playerIndex boardgame.PlayerIndex, userId string) error {

	game, err := s.Game(gameId)

	if err != nil {
		return errors.New("Couldn't get game: " + err.Error())
	}

	if game == nil {
		return errors.New("No game returned")
	}

	if playerIndex < 0 || int(playerIndex) >= int(game.NumPlayers) {
		return errors.New("Invalid player index")
	}

	//TODO: should we validate that this is a real userId?

	var player PlayerStorageRecord

	err = s.dbMap.SelectOne(&player, "select * from "+TablePlayers+" where GameId=? and PlayerIndex=?", game.Id, int(playerIndex))

	if err == sql.ErrNoRows {
		// Insert the row

		player = PlayerStorageRecord{
			GameId:      game.Id,
			PlayerIndex: int64(playerIndex),
			UserId:      userId,
		}

		err = s.dbMap.Insert(&player)

		if err != nil {
			return errors.New("Couldn't insert new player line: " + err.Error())
		}

		return nil
	}

	//Update the row, if it wasn't an error.

	if err != nil {
		return errors.New("Failed to retrieve existing Player line: " + err.Error())
	}

	player.UserId = userId

	_, err = s.dbMap.Update(player)

	if err != nil {
		return errors.New("Couldn't update player line: " + err.Error())
	}

	return nil

}

func (s *StorageManager) UserIdsForGame(gameId string) []string {

	game, err := s.Game(gameId)

	if err != nil {
		log.Println("Couldn't get game: " + err.Error())
		return nil
	}

	if game == nil {
		log.Println("No game returned.")
		return nil
	}

	var players []PlayerStorageRecord

	_, err = s.dbMap.Select(&players, "select * from "+TablePlayers+" where GameId=? order by PlayerIndex desc", game.Id)

	result := make([]string, game.NumPlayers)

	if err == sql.ErrNoRows {
		return result
	}

	if err != nil {
		log.Println("Couldn't get rows: ", err.Error())
		return result
	}

	for _, rec := range players {
		index := int(rec.PlayerIndex)

		if index < 0 || index >= len(result) {
			log.Println("Invalid index", rec)
			continue
		}

		result[index] = rec.UserId
	}

	return result

}

func (s *StorageManager) UpdateUser(user *users.StorageRecord) error {
	userRecord := NewUserStorageRecord(user)

	existingRecord, _ := s.dbMap.SelectInt("select count(*) from "+TableUsers+" where Id=?", user.Id)

	if existingRecord < 1 {
		//Need to insert
		err := s.dbMap.Insert(userRecord)

		if err != nil {
			return errors.New("Couldn't insert user: " + err.Error())
		}
	} else {
		//Need to update
		//TODO: I wonder if this will fail if the user is not yet in the database.
		count, err := s.dbMap.Update(userRecord)
		if err != nil {
			return errors.New("Couldn't update user: " + err.Error())
		}

		if count < 1 {
			return errors.New("Row could not be updated.")
		}
	}

	return nil
}

func (s *StorageManager) GetUserById(uid string) *users.StorageRecord {
	var user UserStorageRecord

	err := s.dbMap.SelectOne(&user, "select * from "+TableUsers+" where Id=?", uid)

	if err == sql.ErrNoRows {
		//Normal
		return nil
	}

	if err != nil {
		log.Println("Unexpected error getting user:", err)
		return nil
	}

	return (&user).ToStorageRecord()
}

func (s *StorageManager) GetUserByCookie(cookie string) *users.StorageRecord {

	var cookieRecord CookieStorageRecord

	err := s.dbMap.SelectOne(&cookieRecord, "select * from "+TableCookies+" where Cookie=?", cookie)

	if err == sql.ErrNoRows {
		//No user
		return nil
	}

	if err != nil {
		log.Println("Unexpected error getting user by cookie: " + err.Error())
		return nil
	}

	return s.GetUserById(cookieRecord.UserId)

}

func (s *StorageManager) ConnectCookieToUser(cookie string, user *users.StorageRecord) error {
	//If user is nil, then delete any records with that cookie.
	if user == nil {

		var cookieRecord CookieStorageRecord

		err := s.dbMap.SelectOne(&cookieRecord, "select * from "+TableCookies+" where Cookie=?", cookie)

		if err == sql.ErrNoRows {
			//We're fine, because it wasn't in the table any way!
			return nil
		}

		if err != nil {
			return errors.New("Unexpected error: " + err.Error())
		}

		//It was there, so we need to delete it.

		count, err := s.dbMap.Delete(&cookieRecord)

		if count < 1 && err != nil {
			return errors.New("Couldnt' delete cookie record when instructed to: " + err.Error())
		}

		return nil
	}

	//If user does not yet exist in database, put them in.
	otherUser := s.GetUserById(user.Id)

	if otherUser == nil {

		//Have to save the user for the first time
		if err := s.UpdateUser(user); err != nil {
			return errors.New("Couldn't add a new user to the database when connecting to cookie: " + err.Error())
		}

		return nil
	}

	record := &CookieStorageRecord{
		Cookie: cookie,
		UserId: user.Id,
	}

	if err := s.dbMap.Insert(record); err != nil {
		return errors.New("Failed to insert cookie pointer record: " + err.Error())
	}
	return nil
}
<commit_msg>Fixed a broken test with a somewhat-hacky answer. Part of #273.<commit_after>/*

bolt provides a bolt-backed database that implements both
boardgame.StorageManager and boardgame/server.StorageManager.

*/
package mysql

import (
	"database/sql"
	"errors"
	"github.com/go-gorp/gorp"
	"github.com/go-sql-driver/mysql"
	"github.com/jkomoros/boardgame"
	"github.com/jkomoros/boardgame/server/api/users"
	"log"
	"strings"
)

const (
	TableGames       = "games"
	TableUsers       = "users"
	TableStates      = "states"
	TableCookies     = "cookies"
	TablePlayers     = "players"
	TableAgentStates = "agentstates"
)

type StorageManager struct {
	db    *sql.DB
	dbMap *gorp.DbMap
	//If in test mode we'll... mock stuff, I guess?
	testMode bool
	//The config string that we were provided in connect.
	config string
}

func NewStorageManager(testMode bool) *StorageManager {
	//We actually don't do much; we do more of our work in Connect()
	return &StorageManager{
		testMode: testMode,
	}
}

func getDSN(config string) (string, error) {

	//Substantially recreated in boardgame-mysqL-admin

	parsedDSN, err := mysql.ParseDSN(config)

	if err != nil {
		return "", errors.New("config provided was not valid DSN: " + err.Error())
	}

	parsedDSN.Collation = "utf8mb4_unicode_ci"
	parsedDSN.MultiStatements = true

	return parsedDSN.FormatDSN(), nil
}

func (s *StorageManager) Connect(config string) error {

	configToUse, err := getDSN(config)

	if err != nil {
		return err
	}

	db, err := sql.Open("mysql", configToUse)
	if err != nil {
		return errors.New("Unable to open database: " + err.Error())
	}

	s.config = config

	s.db = db

	s.dbMap = &gorp.DbMap{
		Db: db,
		Dialect: gorp.MySQLDialect{
			Engine: "InnoDB",
			//the mb4 is necessary to support e.g. emojis
			Encoding: "utf8mb4",
		},
	}

	s.dbMap.AddTableWithName(UserStorageRecord{}, TableUsers).SetKeys(false, "Id")
	s.dbMap.AddTableWithName(GameStorageRecord{}, TableGames).SetKeys(false, "Id")
	s.dbMap.AddTableWithName(StateStorageRecord{}, TableStates).SetKeys(true, "Id")
	s.dbMap.AddTableWithName(CookieStorageRecord{}, TableCookies).SetKeys(false, "Cookie")
	s.dbMap.AddTableWithName(PlayerStorageRecord{}, TablePlayers).SetKeys(true, "Id")
	s.dbMap.AddTableWithName(AgentStateStorageRecord{}, TableAgentStates).SetKeys(true, "Id")

	if s.testMode {
		//TODO: it's weird that tests exercse a code path that we don't use in
		//general. Ideally we'd run the exact same migrations on test db.
		s.dbMap.CreateTablesIfNotExists()
	}

	_, err = s.dbMap.SelectInt("select count(*) from " + TableGames)

	if err != nil {
		return errors.New("Sanity check failed for db. Have you used the admin tool to migrate it up? " + err.Error())
	}

	return nil

}

func (s *StorageManager) Close() {
	if s.db == nil {
		return
	}
	s.db.Close()
	s.db = nil
}

//Cleanup will only drop tables if we're in test mode, and the config string
//used to open the database talked about a test database on localhost (as
//sanity check).
func (s *StorageManager) CleanUp() {

	if !s.testMode {
		return
	}

	if !strings.Contains(s.config, "_test") {
		log.Println("Sanity check on boardgame config before cleanup didn't find _test")
		return
	}

	if !strings.Contains(s.config, "localhost") {
		log.Println("Sanity check on boardgame config before cleanup didn't find localhost")
	}

	if s.db == nil {
		log.Println("Couldn't clean up; db already closed")
		return
	}

	log.Println("Sanity checks passed. Dropping tables to cleanup...")

	if err := s.dbMap.DropTables(); err != nil {
		log.Println("Error dropping tables:", err)
		return
	}
}

func (s *StorageManager) Name() string {
	return "mysql"
}

func (s *StorageManager) State(gameId string, version int) (boardgame.StateStorageRecord, error) {
	var state StateStorageRecord

	err := s.dbMap.SelectOne(&state, "select * from "+TableStates+" where GameId=? and Version=?", gameId, version)

	if err == sql.ErrNoRows {
		return nil, errors.New("No such state")
	}

	if err != nil {
		return nil, errors.New("Unexpected error: " + err.Error())
	}

	return (&state).ToStorageRecord(), nil
}

func (s *StorageManager) Game(id string) (*boardgame.GameStorageRecord, error) {
	var game GameStorageRecord

	err := s.dbMap.SelectOne(&game, "select * from "+TableGames+" where Id=?", id)

	if err == sql.ErrNoRows {
		return nil, errors.New("No such game")
	}

	if err != nil {
		return nil, errors.New("Unexpected error: " + err.Error())
	}

	return (&game).ToStorageRecord(), nil
}

func (s *StorageManager) SaveGameAndCurrentState(game *boardgame.GameStorageRecord, state boardgame.StateStorageRecord) error {

	version := game.Version

	gameRecord := NewGameStorageRecord(game)
	stateRecord := NewStateStorageRecord(game.Id, version, state)

	count, _ := s.dbMap.SelectInt("select count(*) from "+TableGames+" where Id=?", game.Id)

	if count < 1 {
		//Need to insert
		err := s.dbMap.Insert(gameRecord)

		if err != nil {

			return errors.New("Couldn't update game: " + err.Error())

		}

	} else {
		//Need to update
		_, err := s.dbMap.Update(gameRecord)

		if err != nil {
			return errors.New("Couldn't insert game: " + err.Error())
		}
	}

	err := s.dbMap.Insert(stateRecord)

	if err != nil {
		return errors.New("Couldn't insert state: " + err.Error())
	}

	return nil
}

func (s *StorageManager) AgentState(gameId string, player boardgame.PlayerIndex) ([]byte, error) {

	var agent AgentStateStorageRecord

	err := s.dbMap.SelectOne(&agent, "select * from "+TableAgentStates+" where GameId=? and PlayerIndex=? order by Id desc limit 1", gameId, int64(player))

	if err == sql.ErrNoRows {
		return nil, nil
	}

	if err != nil {
		return nil, err
	}

	return agent.ToStorageRecord(), nil

}

func (s *StorageManager) SaveAgentState(gameId string, player boardgame.PlayerIndex, state []byte) error {
	record := NewAgentStateStorageRecord(gameId, player, state)

	err := s.dbMap.Insert(record)

	if err != nil {
		return errors.New("Couldn't save record: " + err.Error())
	}

	return nil
}

func (s *StorageManager) ListGames(max int) []*boardgame.GameStorageRecord {
	var games []GameStorageRecord

	if max < 1 {
		max = 100
	}

	if _, err := s.dbMap.Select(&games, "select * from "+TableGames+" limit ?", max); err != nil {
		return nil
	}

	result := make([]*boardgame.GameStorageRecord, len(games))

	for i, record := range games {
		result[i] = (&record).ToStorageRecord()
	}

	return result
}

func (s *StorageManager) SetPlayerForGame(gameId string, playerIndex boardgame.PlayerIndex, userId string) error {

	game, err := s.Game(gameId)

	if err != nil {
		return errors.New("Couldn't get game: " + err.Error())
	}

	if game == nil {
		return errors.New("No game returned")
	}

	if playerIndex < 0 || int(playerIndex) >= int(game.NumPlayers) {
		return errors.New("Invalid player index")
	}

	//TODO: should we validate that this is a real userId?

	var player PlayerStorageRecord

	err = s.dbMap.SelectOne(&player, "select * from "+TablePlayers+" where GameId=? and PlayerIndex=?", game.Id, int(playerIndex))

	if err == sql.ErrNoRows {
		// Insert the row

		player = PlayerStorageRecord{
			GameId:      game.Id,
			PlayerIndex: int64(playerIndex),
			UserId:      userId,
		}

		err = s.dbMap.Insert(&player)

		if err != nil {
			return errors.New("Couldn't insert new player line: " + err.Error())
		}

		return nil
	}

	//Update the row, if it wasn't an error.

	if err != nil {
		return errors.New("Failed to retrieve existing Player line: " + err.Error())
	}

	player.UserId = userId

	_, err = s.dbMap.Update(player)

	if err != nil {
		return errors.New("Couldn't update player line: " + err.Error())
	}

	return nil

}

func (s *StorageManager) UserIdsForGame(gameId string) []string {

	game, err := s.Game(gameId)

	if err != nil {
		log.Println("Couldn't get game: " + err.Error())
		return nil
	}

	if game == nil {
		log.Println("No game returned.")
		return nil
	}

	var players []PlayerStorageRecord

	_, err = s.dbMap.Select(&players, "select * from "+TablePlayers+" where GameId=? order by PlayerIndex desc", game.Id)

	result := make([]string, game.NumPlayers)

	if err == sql.ErrNoRows {
		return result
	}

	if err != nil {
		log.Println("Couldn't get rows: ", err.Error())
		return result
	}

	for _, rec := range players {
		index := int(rec.PlayerIndex)

		if index < 0 || index >= len(result) {
			log.Println("Invalid index", rec)
			continue
		}

		result[index] = rec.UserId
	}

	return result

}

func (s *StorageManager) UpdateUser(user *users.StorageRecord) error {
	userRecord := NewUserStorageRecord(user)

	existingRecord, _ := s.dbMap.SelectInt("select count(*) from "+TableUsers+" where Id=?", user.Id)

	if existingRecord < 1 {
		//Need to insert
		err := s.dbMap.Insert(userRecord)

		if err != nil {
			return errors.New("Couldn't insert user: " + err.Error())
		}
	} else {
		//Need to update
		//TODO: I wonder if this will fail if the user is not yet in the database.
		count, err := s.dbMap.Update(userRecord)
		if err != nil {
			return errors.New("Couldn't update user: " + err.Error())
		}

		if count < 1 {
			return errors.New("Row could not be updated.")
		}
	}

	return nil
}

func (s *StorageManager) GetUserById(uid string) *users.StorageRecord {
	var user UserStorageRecord

	err := s.dbMap.SelectOne(&user, "select * from "+TableUsers+" where Id=?", uid)

	if err == sql.ErrNoRows {
		//Normal
		return nil
	}

	if err != nil {
		log.Println("Unexpected error getting user:", err)
		return nil
	}

	return (&user).ToStorageRecord()
}

func (s *StorageManager) GetUserByCookie(cookie string) *users.StorageRecord {

	var cookieRecord CookieStorageRecord

	err := s.dbMap.SelectOne(&cookieRecord, "select * from "+TableCookies+" where Cookie=?", cookie)

	if err == sql.ErrNoRows {
		//No user
		return nil
	}

	if err != nil {
		log.Println("Unexpected error getting user by cookie: " + err.Error())
		return nil
	}

	return s.GetUserById(cookieRecord.UserId)

}

func (s *StorageManager) ConnectCookieToUser(cookie string, user *users.StorageRecord) error {
	//If user is nil, then delete any records with that cookie.
	if user == nil {

		var cookieRecord CookieStorageRecord

		err := s.dbMap.SelectOne(&cookieRecord, "select * from "+TableCookies+" where Cookie=?", cookie)

		if err == sql.ErrNoRows {
			//We're fine, because it wasn't in the table any way!
			return nil
		}

		if err != nil {
			return errors.New("Unexpected error: " + err.Error())
		}

		//It was there, so we need to delete it.

		count, err := s.dbMap.Delete(&cookieRecord)

		if count < 1 && err != nil {
			return errors.New("Couldnt' delete cookie record when instructed to: " + err.Error())
		}

		return nil
	}

	//If user does not yet exist in database, put them in.
	otherUser := s.GetUserById(user.Id)

	if otherUser == nil {

		//Have to save the user for the first time
		if err := s.UpdateUser(user); err != nil {
			return errors.New("Couldn't add a new user to the database when connecting to cookie: " + err.Error())
		}

		return nil
	}

	record := &CookieStorageRecord{
		Cookie: cookie,
		UserId: user.Id,
	}

	if err := s.dbMap.Insert(record); err != nil {
		return errors.New("Failed to insert cookie pointer record: " + err.Error())
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"
	"time"
)

const (
	OpGet         = "/get"
	OpSet         = "/post"
	OpInfo        = "._info"
	OpDelete      = "/delete"
	OpCursorFirst = "_first"
	OpCursorLast  = "_last"
	OpCursorNext  = "._next"
	OpCursorPrev  = "._prev"
	OpPrefixMatch = "._match"
)

func httpGet(w http.ResponseWriter, op Operation) *HTTPError {
	value := op.Database.Unjar(op.Key)
	// Check if the item existed
	if value == nil {
		return &HTTPError{Code: 404, Message: "These aren't your ghosts."}
	}

	// Send value
	w.Write(value)
	return nil
}

func httpSet(w http.ResponseWriter, op Operation, r *http.Request) *HTTPError {
	value, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return &HTTPError{Code: 500, Message: "Your post body is messed up!"}
	}

	// Check if value already existed
	exists := op.Database.Exists(op.Key)

	res := op.Database.Jar(op.Key, value)
	if res == 0 {
		// Status 201 if created, 200 if updated
		if exists {
			w.WriteHeader(200)
		} else {
			w.WriteHeader(201)
		}
		fmt.Fprintf(w, "無駄")
	} else {
		return &HTTPError{Code: 500, Message: "Something went horribly wrong..."}
	}

	// Try to set expiration, if provided
	if eep, ok := r.Header["X-Olegdb-Use-By"]; ok {
		ep, err := strconv.Atoi(eep[0])
		if err != nil {
			return &HTTPError{Code: 500, Message: "The expiration format is wrong!"}
		}
		date := time.Unix(int64(ep), 0)
		op.Database.Spoil(op.Key, date)
		// fmt.Fprintf(w, "\r\nThe jar is spoiling!")
	}

	return nil
}

func httpInfo(w http.ResponseWriter, op Operation) *HTTPError {
	// Does it even exists?
	if !op.Database.Exists(op.Key) {
		return &HTTPError{Code: 404, Message: "Key not found in database"}
	}

	// Get and set Expiration
	res, doesExpire := op.Database.Expiration(op.Key)
	if doesExpire {
		w.Header().Add("Expires", strconv.Itoa(int(res.UTC().Unix())))
	}

	// Add Record count
	w.Header().Add("X-Olegdb-Rcrd-Cnt", strconv.Itoa(int(*op.Database.RecordCount)))

	// Send empty body
	fmt.Fprintf(w, "\r\n")
	return nil
}

func httpDelete(w http.ResponseWriter, op Operation) *HTTPError {
	res := op.Database.Scoop(op.Key)
	if res != 0 {
		return &HTTPError{Code: 500, Message: "Something went horribly wrong..."}
	}

	fmt.Fprintf(w, "Key deleted successfully!")
	return nil
}

func httpMatch(w http.ResponseWriter, op Operation) *HTTPError {
	has, res := op.Database.PrefixMatch(op.Key)
	if !has {
		return &HTTPError{Code: 404, Message: "No matches found"}
	}
	w.Header().Add("X-Olegdb-Num-Matches", strconv.Itoa(len(res)))
	fmt.Fprintf(w, strings.Join(res, "\n"))
	return nil
}

func httpCurFirst(w http.ResponseWriter, op Operation) *HTTPError {
	has, key, data := op.Database.First()
	if !has {
		return &HTTPError{Code: 404, Message: "No records found"}
	}
	w.Header().Add("X-Olegdb-Key", key)
	w.Write(data)
	return nil
}

func httpCurLast(w http.ResponseWriter, op Operation) *HTTPError {
	has, key, data := op.Database.Last()
	if !has {
		return &HTTPError{Code: 404, Message: "No records found"}
	}
	w.Header().Add("X-Olegdb-Key", key)
	w.Write(data)
	return nil
}

func httpCurNext(w http.ResponseWriter, op Operation) *HTTPError {
	has, key, data := op.Database.Next(op.Key)
	if !has {
		return &HTTPError{Code: 404, Message: "No records found"}
	}
	w.Header().Add("X-Olegdb-Key", key)
	w.Write(data)
	return nil
}

func httpCurPrev(w http.ResponseWriter, op Operation) *HTTPError {
	has, key, data := op.Database.Prev(op.Key)
	if !has {
		return &HTTPError{Code: 404, Message: "No records found"}
	}
	w.Header().Add("X-Olegdb-Key", key)
	w.Write(data)
	return nil
}
<commit_msg>Revert "Revert "Set content-length header on prefix matching.""<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"
	"time"
)

const (
	OpGet         = "/get"
	OpSet         = "/post"
	OpInfo        = "._info"
	OpDelete      = "/delete"
	OpCursorFirst = "_first"
	OpCursorLast  = "_last"
	OpCursorNext  = "._next"
	OpCursorPrev  = "._prev"
	OpPrefixMatch = "._match"
)

func httpGet(w http.ResponseWriter, op Operation) *HTTPError {
	value := op.Database.Unjar(op.Key)
	// Check if the item existed
	if value == nil {
		return &HTTPError{Code: 404, Message: "These aren't your ghosts."}
	}

	// Send value
	w.Write(value)
	return nil
}

func httpSet(w http.ResponseWriter, op Operation, r *http.Request) *HTTPError {
	value, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return &HTTPError{Code: 500, Message: "Your post body is messed up!"}
	}

	// Check if value already existed
	exists := op.Database.Exists(op.Key)

	res := op.Database.Jar(op.Key, value)
	if res == 0 {
		// Status 201 if created, 200 if updated
		if exists {
			w.WriteHeader(200)
		} else {
			w.WriteHeader(201)
		}
		fmt.Fprintf(w, "無駄")
	} else {
		return &HTTPError{Code: 500, Message: "Something went horribly wrong..."}
	}

	// Try to set expiration, if provided
	if eep, ok := r.Header["X-Olegdb-Use-By"]; ok {
		ep, err := strconv.Atoi(eep[0])
		if err != nil {
			return &HTTPError{Code: 500, Message: "The expiration format is wrong!"}
		}
		date := time.Unix(int64(ep), 0)
		op.Database.Spoil(op.Key, date)
		// fmt.Fprintf(w, "\r\nThe jar is spoiling!")
	}

	return nil
}

func httpInfo(w http.ResponseWriter, op Operation) *HTTPError {
	// Does it even exists?
	if !op.Database.Exists(op.Key) {
		return &HTTPError{Code: 404, Message: "Key not found in database"}
	}

	// Get and set Expiration
	res, doesExpire := op.Database.Expiration(op.Key)
	if doesExpire {
		w.Header().Add("Expires", strconv.Itoa(int(res.UTC().Unix())))
	}

	// Add Record count
	w.Header().Add("X-Olegdb-Rcrd-Cnt", strconv.Itoa(int(*op.Database.RecordCount)))

	// Send empty body
	fmt.Fprintf(w, "\r\n")
	return nil
}

func httpDelete(w http.ResponseWriter, op Operation) *HTTPError {
	res := op.Database.Scoop(op.Key)
	if res != 0 {
		return &HTTPError{Code: 500, Message: "Something went horribly wrong..."}
	}

	fmt.Fprintf(w, "Key deleted successfully!")
	return nil
}

func httpMatch(w http.ResponseWriter, op Operation) *HTTPError {
	has, res := op.Database.PrefixMatch(op.Key)
	if !has {
		return &HTTPError{Code: 404, Message: "No matches found"}
	}
	w.Header().Add("X-Olegdb-Num-Matches", strconv.Itoa(len(res)))
	content := strings.Join(res, "\n")
	w.Header().Add("Content-Length", strconv.Itoa(len(content)))
	fmt.Fprintf(w, content)
	return nil
}

func httpCurFirst(w http.ResponseWriter, op Operation) *HTTPError {
	has, key, data := op.Database.First()
	if !has {
		return &HTTPError{Code: 404, Message: "No records found"}
	}
	w.Header().Add("X-Olegdb-Key", key)
	w.Write(data)
	return nil
}

func httpCurLast(w http.ResponseWriter, op Operation) *HTTPError {
	has, key, data := op.Database.Last()
	if !has {
		return &HTTPError{Code: 404, Message: "No records found"}
	}
	w.Header().Add("X-Olegdb-Key", key)
	w.Write(data)
	return nil
}

func httpCurNext(w http.ResponseWriter, op Operation) *HTTPError {
	has, key, data := op.Database.Next(op.Key)
	if !has {
		return &HTTPError{Code: 404, Message: "No records found"}
	}
	w.Header().Add("X-Olegdb-Key", key)
	w.Write(data)
	return nil
}

func httpCurPrev(w http.ResponseWriter, op Operation) *HTTPError {
	has, key, data := op.Database.Prev(op.Key)
	if !has {
		return &HTTPError{Code: 404, Message: "No records found"}
	}
	w.Header().Add("X-Olegdb-Key", key)
	w.Write(data)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strings"

	"github.com/litl/galaxy/log"
	"github.com/litl/galaxy/shuttle/client"

	"github.com/gorilla/mux"
)

func getConfig(w http.ResponseWriter, r *http.Request) {
	w.Write(marshal(Registry.Config()))
}

func getStats(w http.ResponseWriter, r *http.Request) {
	if len(Registry.Config()) == 0 {
		w.WriteHeader(503)
	}
	w.Write(marshal(Registry.Stats()))
}

func getService(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	serviceStats, err := Registry.ServiceStats(vars["service"])
	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}

	w.Write(marshal(serviceStats))
}

// Update a service and/or backends.
// Adding a `backends_only` query parameter will prevent the service from being
// shutdown and replaced if the ServiceConfig is not identical..
func postService(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Errorln(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer r.Body.Close()

	svcCfg := client.ServiceConfig{Name: vars["service"]}
	err = json.Unmarshal(body, &svcCfg)
	if err != nil {
		log.Errorln(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if Registry.GetService(svcCfg.Name) == nil {
		if e := Registry.AddService(svcCfg); e != nil {
			http.Error(w, e.Error(), http.StatusInternalServerError)
			return
		}
	}

	if e := Registry.UpdateService(svcCfg); e != nil {
		log.Errorln("Unable to update service %s", svcCfg.Name)
		http.Error(w, e.Error(), http.StatusInternalServerError)
		return
	}

	vhosts := make(map[string][]string)
	for _, svcCfg := range Registry.Config() {
		for _, vhost := range svcCfg.VirtualHosts {
			for _, backend := range svcCfg.Backends {
				if backend.Addr == "" {
					log.Warnf("No address specifed for %s for %s. Skipping.", backend.Name, svcCfg.Name)
					continue
				}
				addr := "http://" + backend.Addr
				httpRouter.AddBackend(svcCfg.Name, vhost, addr)
				vhosts[vhost] = append(vhosts[vhost], addr)
			}

		}
	}
	for vhost, addrs := range vhosts {
		httpRouter.RemoveBackends(vhost, addrs)
	}

	go writeStateConfig()
	w.Write(marshal(Registry.Config()))
}

func deleteService(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	service := Registry.GetService(vars["service"])
	if service != nil {
		for _, vhost := range service.VirtualHosts {
			httpRouter.RemoveRouter(vhost)
		}
	}

	err := Registry.RemoveService(vars["service"])
	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}
	go writeStateConfig()
	w.Write(marshal(Registry.Config()))
}

func getBackend(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	serviceName := vars["service"]
	backendName := vars["backend"]

	backend, err := Registry.BackendStats(serviceName, backendName)
	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}

	w.Write(marshal(backend))
}

func postBackend(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Errorln(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer r.Body.Close()

	backendName := vars["backend"]
	serviceName := vars["service"]

	backendCfg := client.BackendConfig{Name: backendName}
	err = json.Unmarshal(body, &backendCfg)
	if err != nil {
		log.Errorln(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if err := Registry.AddBackend(serviceName, backendCfg); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	service := Registry.GetService(vars["service"])
	if service != nil {
		for _, vhost := range service.VirtualHosts {
			httpRouter.AddBackend(backendCfg.Name, vhost, "http://"+backendCfg.Addr)
		}
	}

	go writeStateConfig()
	w.Write(marshal(Registry.Config()))
}

func deleteBackend(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	serviceName := vars["service"]
	backendName := vars["backend"]

	service := Registry.GetService(vars["service"])
	if service != nil {
		for _, vhost := range service.VirtualHosts {
			backend := service.get(backendName)
			httpRouter.RemoveBackend(vhost, "http://"+backend.Addr)
		}
	}

	if err := Registry.RemoveBackend(serviceName, backendName); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	go writeStateConfig()
	w.Write(marshal(Registry.Config()))
}

func addHandlers() {
	r := mux.NewRouter()
	r.HandleFunc("/", getStats).Methods("GET")
	r.HandleFunc("/_config", getConfig).Methods("GET")
	r.HandleFunc("/{service}", getService).Methods("GET")
	r.HandleFunc("/{service}", postService).Methods("PUT", "POST")
	r.HandleFunc("/{service}", deleteService).Methods("DELETE")
	r.HandleFunc("/{service}/{backend}", getBackend).Methods("GET")
	r.HandleFunc("/{service}/{backend}", postBackend).Methods("PUT", "POST")
	r.HandleFunc("/{service}/{backend}", deleteBackend).Methods("DELETE")
	http.Handle("/", r)
}

func startAdminHTTPServer() {
	defer wg.Done()
	addHandlers()
	log.Println("Admin server listening on", adminListenAddr)

	netw := "tcp"

	if strings.HasPrefix(adminListenAddr, "/") {
		netw = "unix"

		// remove our old socket if we left it lying around
		if stats, err := os.Stat(adminListenAddr); err == nil {
			if stats.Mode()&os.ModeSocket != 0 {
				os.Remove(adminListenAddr)
			}
		}

		defer os.Remove(adminListenAddr)
	}

	listener, err := net.Listen(netw, adminListenAddr)
	if err != nil {
		log.Fatalln(err)
	}

	http.Serve(listener, nil)
}
<commit_msg>Return bad request if trying to setup listener on shuttle ports<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strings"

	"github.com/litl/galaxy/log"
	"github.com/litl/galaxy/shuttle/client"

	"github.com/gorilla/mux"
)

func getConfig(w http.ResponseWriter, r *http.Request) {
	w.Write(marshal(Registry.Config()))
}

func getStats(w http.ResponseWriter, r *http.Request) {
	if len(Registry.Config()) == 0 {
		w.WriteHeader(503)
	}
	w.Write(marshal(Registry.Stats()))
}

func getService(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	serviceStats, err := Registry.ServiceStats(vars["service"])
	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}

	w.Write(marshal(serviceStats))
}

// Update a service and/or backends.
// Adding a `backends_only` query parameter will prevent the service from being
// shutdown and replaced if the ServiceConfig is not identical..
func postService(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Errorln(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer r.Body.Close()

	svcCfg := client.ServiceConfig{Name: vars["service"]}
	err = json.Unmarshal(body, &svcCfg)
	if err != nil {
		log.Errorln(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	invalidPorts := []string{
		listenAddr[strings.Index(listenAddr, ":")+1:],
		adminListenAddr[strings.Index(adminListenAddr, ":")+1:],
	}

	for _, port := range invalidPorts {
		if strings.HasSuffix(svcCfg.Addr, port) {
			log.Errorf("Cannot use shuttle port: %s for %s service listener. Shuttle is using it.", port, svcCfg.Name)
			http.Error(w, fmt.Sprintf("cannot use %s for listener port", port), http.StatusBadRequest)
			return
		}
	}

	if Registry.GetService(svcCfg.Name) == nil {
		if e := Registry.AddService(svcCfg); e != nil {
			log.Errorln(err)
			http.Error(w, e.Error(), http.StatusInternalServerError)
			return
		}
	}

	if e := Registry.UpdateService(svcCfg); e != nil {
		log.Errorln("Unable to update service %s", svcCfg.Name)
		http.Error(w, e.Error(), http.StatusInternalServerError)
		return
	}

	vhosts := make(map[string][]string)
	for _, svcCfg := range Registry.Config() {
		for _, vhost := range svcCfg.VirtualHosts {
			for _, backend := range svcCfg.Backends {
				if backend.Addr == "" {
					log.Warnf("No address specifed for %s for %s. Skipping.", backend.Name, svcCfg.Name)
					continue
				}
				addr := "http://" + backend.Addr
				httpRouter.AddBackend(svcCfg.Name, vhost, addr)
				vhosts[vhost] = append(vhosts[vhost], addr)
			}

		}
	}
	for vhost, addrs := range vhosts {
		httpRouter.RemoveBackends(vhost, addrs)
	}

	go writeStateConfig()
	w.Write(marshal(Registry.Config()))
}

func deleteService(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	service := Registry.GetService(vars["service"])
	if service != nil {
		for _, vhost := range service.VirtualHosts {
			httpRouter.RemoveRouter(vhost)
		}
	}

	err := Registry.RemoveService(vars["service"])
	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}
	go writeStateConfig()
	w.Write(marshal(Registry.Config()))
}

func getBackend(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	serviceName := vars["service"]
	backendName := vars["backend"]

	backend, err := Registry.BackendStats(serviceName, backendName)
	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}

	w.Write(marshal(backend))
}

func postBackend(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Errorln(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer r.Body.Close()

	backendName := vars["backend"]
	serviceName := vars["service"]

	backendCfg := client.BackendConfig{Name: backendName}
	err = json.Unmarshal(body, &backendCfg)
	if err != nil {
		log.Errorln(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	if err := Registry.AddBackend(serviceName, backendCfg); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	service := Registry.GetService(vars["service"])
	if service != nil {
		for _, vhost := range service.VirtualHosts {
			httpRouter.AddBackend(backendCfg.Name, vhost, "http://"+backendCfg.Addr)
		}
	}

	go writeStateConfig()
	w.Write(marshal(Registry.Config()))
}

func deleteBackend(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	serviceName := vars["service"]
	backendName := vars["backend"]

	service := Registry.GetService(vars["service"])
	if service != nil {
		for _, vhost := range service.VirtualHosts {
			backend := service.get(backendName)
			httpRouter.RemoveBackend(vhost, "http://"+backend.Addr)
		}
	}

	if err := Registry.RemoveBackend(serviceName, backendName); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	go writeStateConfig()
	w.Write(marshal(Registry.Config()))
}

func addHandlers() {
	r := mux.NewRouter()
	r.HandleFunc("/", getStats).Methods("GET")
	r.HandleFunc("/_config", getConfig).Methods("GET")
	r.HandleFunc("/{service}", getService).Methods("GET")
	r.HandleFunc("/{service}", postService).Methods("PUT", "POST")
	r.HandleFunc("/{service}", deleteService).Methods("DELETE")
	r.HandleFunc("/{service}/{backend}", getBackend).Methods("GET")
	r.HandleFunc("/{service}/{backend}", postBackend).Methods("PUT", "POST")
	r.HandleFunc("/{service}/{backend}", deleteBackend).Methods("DELETE")
	http.Handle("/", r)
}

func startAdminHTTPServer() {
	defer wg.Done()
	addHandlers()
	log.Println("Admin server listening on", adminListenAddr)

	netw := "tcp"

	if strings.HasPrefix(adminListenAddr, "/") {
		netw = "unix"

		// remove our old socket if we left it lying around
		if stats, err := os.Stat(adminListenAddr); err == nil {
			if stats.Mode()&os.ModeSocket != 0 {
				os.Remove(adminListenAddr)
			}
		}

		defer os.Remove(adminListenAddr)
	}

	listener, err := net.Listen(netw, adminListenAddr)
	if err != nil {
		log.Fatalln(err)
	}

	http.Serve(listener, nil)
}
<|endoftext|>


<commit_before>package adodb

import (
	"database/sql"
	"database/sql/driver"
	"github.com/mattn/go-ole"
	"github.com/mattn/go-ole/oleutil"
	"io"
	"math/big"
	"time"
	"unsafe"
)

func init() {
	ole.CoInitialize(0)
	sql.Register("adodb", &AdodbDriver{})
}

type AdodbDriver struct {

}

type AdodbConn struct {
	db *ole.IDispatch
}

type AdodbTx struct {
	c *AdodbConn
}

func (tx *AdodbTx) Commit() error {
	_, err := oleutil.CallMethod(tx.c.db, "CommitTrans")
	if err != nil {
		return err
	}
	return nil
}

func (tx *AdodbTx) Rollback() error {
	_, err := oleutil.CallMethod(tx.c.db, "Rollback")
	if err != nil {
		return err
	}
	return nil
}

func (c *AdodbConn) exec(cmd string) error {
	_, err := oleutil.CallMethod(c.db, "Execute", cmd)
	return err
}

func (c *AdodbConn) Begin() (driver.Tx, error) {
	_, err := oleutil.CallMethod(c.db, "BeginTrans")
	if err != nil {
		return nil, err
	}
	return &AdodbTx{c}, nil
}

func (d *AdodbDriver) Open(dsn string) (driver.Conn, error) {
	unknown, err := oleutil.CreateObject("ADODB.Connection")
	if err != nil {
		return nil, err
	}
	db, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.CallMethod(db, "Open", dsn)
	if err != nil {
		return nil, err
	}
	return &AdodbConn{db}, nil
}

func (c *AdodbConn) Close() error {
	_, err := oleutil.CallMethod(c.db, "Close")
	if err != nil {
		return err
	}
	c.db = nil
	return nil
}

type AdodbStmt struct {
	c  *AdodbConn
	s  *ole.IDispatch
	ps *ole.IDispatch
	b  []string
}

func (c *AdodbConn) Prepare(query string) (driver.Stmt, error) {
	unknown, err := oleutil.CreateObject("ADODB.Command")
	if err != nil {
		return nil, err
	}
	s, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "ActiveConnection", c.db)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandText", query)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandType", 1)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "Prepared", true)
	if err != nil {
		return nil, err
	}
	val, err := oleutil.GetProperty(s, "Parameters")
	if err != nil {
		return nil, err
	}
	return &AdodbStmt{c, s, val.ToIDispatch(), nil}, nil
}

func (s *AdodbStmt) Bind(bind []string) error {
	s.b = bind
	return nil
}

func (s *AdodbStmt) Close() error {
	s.s.Release()
	return nil
}

func (s *AdodbStmt) NumInput() int {
	if s.b != nil {
		return len(s.b)
	}
	_, err := oleutil.CallMethod(s.ps, "Refresh")
	if err != nil {
		return -1
	}
	val, err := oleutil.GetProperty(s.ps, "Count")
	if err != nil {
		return -1
	}
	c := int(val.Val)
	return c
}

func (s *AdodbStmt) bind(args []driver.Value) error {
	if s.b != nil {
		for i, v := range args {
			var b string = "?"
			if len(s.b) < i {
				b = s.b[i]
			}
			unknown, err := oleutil.CallMethod(s.s, "CreateParameter", b, 12, 1)
			if err != nil {
				return err
			}
			param := unknown.ToIDispatch()
			defer param.Release()
			_, err = oleutil.PutProperty(param, "Value", v)
			if err != nil {
				return err
			}
			_, err = oleutil.CallMethod(s.ps, "Append", param)
			if err != nil {
				return err
			}
		}
	} else {
		for i, v := range args {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(s.ps, "Item", &varval)
			if err != nil {
				return err
			}
			item := val.ToIDispatch()
			defer item.Release()
			_, err = oleutil.PutProperty(item, "Value", v)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *AdodbStmt) Query(args []driver.Value) (driver.Rows, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	rc, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return &AdodbRows{s, rc.ToIDispatch(), -1, nil}, nil
}

func (s *AdodbStmt) Exec(args []driver.Value) (driver.Result, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	_, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return driver.ResultNoRows, nil
}

type AdodbRows struct {
	s    *AdodbStmt
	rc   *ole.IDispatch
	nc   int
	cols []string
}

func (rc *AdodbRows) Close() error {
	_, err := oleutil.CallMethod(rc.rc, "Close")
	if err != nil {
		return err
	}
	return nil
}

func (rc *AdodbRows) Columns() []string {
	if rc.nc != len(rc.cols) {
		unknown, err := oleutil.GetProperty(rc.rc, "Fields")
		if err != nil {
			return []string{}
		}
		fields := unknown.ToIDispatch()
		defer fields.Release()
		val, err := oleutil.GetProperty(fields, "Count")
		if err != nil {
			return []string{}
		}
		rc.nc = int(val.Val)
		rc.cols = make([]string, rc.nc)
		for i := 0; i < rc.nc; i++ {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(fields, "Item", &varval)
			if err != nil {
				return []string{}
			}
			item := val.ToIDispatch()
			if err != nil {
				return []string{}
			}
			name, err := oleutil.GetProperty(item, "Name")
			if err != nil {
				return []string{}
			}
			rc.cols[i] = name.ToString()
			item.Release()
		}
	}
	return rc.cols
}

func (rc *AdodbRows) Next(dest []driver.Value) error {
	unknown, err := oleutil.GetProperty(rc.rc, "EOF")
	if err != nil {
		return io.EOF
	}
	if unknown.Val != 0 {
		return io.EOF
	}
	unknown, err = oleutil.GetProperty(rc.rc, "Fields")
	if err != nil {
		return err
	}
	fields := unknown.ToIDispatch()
	defer fields.Release()
	for i := range dest {
		var varval ole.VARIANT
		varval.VT = ole.VT_I4
		varval.Val = int64(i)
		val, err := oleutil.CallMethod(fields, "Item", &varval)
		if err != nil {
			return err
		}
		field := val.ToIDispatch()
		defer field.Release()
		typ, err := oleutil.GetProperty(field, "Type")
		if err != nil {
			return err
		}
		val, err = oleutil.GetProperty(field, "Value")
		if err != nil {
			return err
		}
		field.Release()
		switch typ.Val {
		case 0: // ADEMPTY
			// TODO
		case 2: // ADSMALLINT
			dest[i] = int16(val.Val)
		case 3: // ADINTEGER
			dest[i] = int32(val.Val)
		case 4: // ADSINGLE
			dest[i] = float32(val.Val)
		case 5: // ADDOUBLE
			dest[i] = val.Val
		case 6: // ADCURRENCY
			dest[i] = float64(val.Val)
		case 7: // ADDATE
			// TODO
		case 8: // ADBSTR
			dest[i] = val.ToString()
		case 9: // ADIDISPATCH
			dest[i] = val.ToIDispatch()
		case 10: // ADERROR
			// TODO
		case 11: // ADBOOLEAN
			if val.Val != 0 {
				dest[i] = true
			} else {
				dest[i] = false
			}
		case 12: // ADVARIANT
			dest[i] = val
		case 13: // ADIUNKNOWN
			dest[i] = val.ToIUnknown()
		case 14: // ADDECIMAL
			dest[i] = float64(val.Val)
		case 16: // ADTINYINT
			dest[i] = int8(val.Val)
		case 17: // ADUNSIGNEDTINYINT
			dest[i] = uint8(val.Val)
		case 18: // ADUNSIGNEDSMALLINT
			dest[i] = uint16(val.Val)
		case 19: // ADUNSIGNEDINT
			dest[i] = uint32(val.Val)
		case 20: // ADBIGINT
			dest[i] = big.NewInt(val.Val)
		case 21: // ADUNSIGNEDBIGINT
			// TODO
		case 72: // ADGUID
			// TODO
		case 128: // ADBINARY
			sa := *(**ole.SAFEARRAY)(unsafe.Pointer(&val.Val))
			dest[i] = (*[1 << 30]byte)(unsafe.Pointer(uintptr(sa.PvData)))[0:sa.CbElements]
		case 129: // ADCHAR
			dest[i] = val.ToString()//uint8(val.Val)
		case 130: // ADWCHAR
			dest[i] = val.ToString()//uint16(val.Val)
		case 131: // ADNUMERIC
			dest[i] = val.Val
		case 132: // ADUSERDEFINED
			dest[i] = uintptr(val.Val)
		case 133: // ADDBDATE
			dest[i] = time.Unix(0, val.Val).UTC()
		case 134: // ADDBTIME
			dest[i] = time.Unix(0, val.Val).UTC()
		case 135: // ADDBTIMESTAMP
			dest[i] = time.Unix(0, val.Val).UTC()
		case 136: // ADCHAPTER
			dest[i] = val.ToString()
		case 200: // ADVARCHAR
			dest[i] = val.ToString()
		case 201: // ADLONGVARCHAR
			dest[i] = val.ToString()
		case 202: // ADVARWCHAR
			dest[i] = val.ToString()
		case 203: // ADLONGVARWCHAR
			dest[i] = val.ToString()
		case 204: // ADVARBINARY
			// TODO
		case 205: // ADLONGVARBINARY
			// TODO
		}
	}
	_, err = oleutil.CallMethod(rc.rc, "MoveNext")
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>Add support for GUID, single, double, and LongVarBinary types.<commit_after>package adodb

import (
	"database/sql"
	"database/sql/driver"
	"github.com/mattn/go-ole"
	"github.com/mattn/go-ole/oleutil"
	"io"
	"math"
	"math/big"
	"time"
	"unsafe"
)

func init() {
	ole.CoInitialize(0)
	sql.Register("adodb", &AdodbDriver{})
}

type AdodbDriver struct {
}

type AdodbConn struct {
	db *ole.IDispatch
}

type AdodbTx struct {
	c *AdodbConn
}

func (tx *AdodbTx) Commit() error {
	_, err := oleutil.CallMethod(tx.c.db, "CommitTrans")
	if err != nil {
		return err
	}
	return nil
}

func (tx *AdodbTx) Rollback() error {
	_, err := oleutil.CallMethod(tx.c.db, "Rollback")
	if err != nil {
		return err
	}
	return nil
}

func (c *AdodbConn) exec(cmd string) error {
	_, err := oleutil.CallMethod(c.db, "Execute", cmd)
	return err
}

func (c *AdodbConn) Begin() (driver.Tx, error) {
	_, err := oleutil.CallMethod(c.db, "BeginTrans")
	if err != nil {
		return nil, err
	}
	return &AdodbTx{c}, nil
}

func (d *AdodbDriver) Open(dsn string) (driver.Conn, error) {
	unknown, err := oleutil.CreateObject("ADODB.Connection")
	if err != nil {
		return nil, err
	}
	db, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.CallMethod(db, "Open", dsn)
	if err != nil {
		return nil, err
	}
	return &AdodbConn{db}, nil
}

func (c *AdodbConn) Close() error {
	_, err := oleutil.CallMethod(c.db, "Close")
	if err != nil {
		return err
	}
	c.db = nil
	return nil
}

type AdodbStmt struct {
	c  *AdodbConn
	s  *ole.IDispatch
	ps *ole.IDispatch
	b  []string
}

func (c *AdodbConn) Prepare(query string) (driver.Stmt, error) {
	unknown, err := oleutil.CreateObject("ADODB.Command")
	if err != nil {
		return nil, err
	}
	s, err := unknown.QueryInterface(ole.IID_IDispatch)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "ActiveConnection", c.db)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandText", query)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "CommandType", 1)
	if err != nil {
		return nil, err
	}
	_, err = oleutil.PutProperty(s, "Prepared", true)
	if err != nil {
		return nil, err
	}
	val, err := oleutil.GetProperty(s, "Parameters")
	if err != nil {
		return nil, err
	}
	return &AdodbStmt{c, s, val.ToIDispatch(), nil}, nil
}

func (s *AdodbStmt) Bind(bind []string) error {
	s.b = bind
	return nil
}

func (s *AdodbStmt) Close() error {
	s.s.Release()
	return nil
}

func (s *AdodbStmt) NumInput() int {
	if s.b != nil {
		return len(s.b)
	}
	_, err := oleutil.CallMethod(s.ps, "Refresh")
	if err != nil {
		return -1
	}
	val, err := oleutil.GetProperty(s.ps, "Count")
	if err != nil {
		return -1
	}
	c := int(val.Val)
	return c
}

func (s *AdodbStmt) bind(args []driver.Value) error {
	if s.b != nil {
		for i, v := range args {
			var b string = "?"
			if len(s.b) < i {
				b = s.b[i]
			}
			unknown, err := oleutil.CallMethod(s.s, "CreateParameter", b, 12, 1)
			if err != nil {
				return err
			}
			param := unknown.ToIDispatch()
			defer param.Release()
			_, err = oleutil.PutProperty(param, "Value", v)
			if err != nil {
				return err
			}
			_, err = oleutil.CallMethod(s.ps, "Append", param)
			if err != nil {
				return err
			}
		}
	} else {
		for i, v := range args {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(s.ps, "Item", &varval)
			if err != nil {
				return err
			}
			item := val.ToIDispatch()
			defer item.Release()
			_, err = oleutil.PutProperty(item, "Value", v)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (s *AdodbStmt) Query(args []driver.Value) (driver.Rows, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	rc, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return &AdodbRows{s, rc.ToIDispatch(), -1, nil}, nil
}

func (s *AdodbStmt) Exec(args []driver.Value) (driver.Result, error) {
	if err := s.bind(args); err != nil {
		return nil, err
	}
	_, err := oleutil.CallMethod(s.s, "Execute")
	if err != nil {
		return nil, err
	}
	return driver.ResultNoRows, nil
}

type AdodbRows struct {
	s    *AdodbStmt
	rc   *ole.IDispatch
	nc   int
	cols []string
}

func (rc *AdodbRows) Close() error {
	_, err := oleutil.CallMethod(rc.rc, "Close")
	if err != nil {
		return err
	}
	return nil
}

func (rc *AdodbRows) Columns() []string {
	if rc.nc != len(rc.cols) {
		unknown, err := oleutil.GetProperty(rc.rc, "Fields")
		if err != nil {
			return []string{}
		}
		fields := unknown.ToIDispatch()
		defer fields.Release()
		val, err := oleutil.GetProperty(fields, "Count")
		if err != nil {
			return []string{}
		}
		rc.nc = int(val.Val)
		rc.cols = make([]string, rc.nc)
		for i := 0; i < rc.nc; i++ {
			var varval ole.VARIANT
			varval.VT = ole.VT_I4
			varval.Val = int64(i)
			val, err := oleutil.CallMethod(fields, "Item", &varval)
			if err != nil {
				return []string{}
			}
			item := val.ToIDispatch()
			if err != nil {
				return []string{}
			}
			name, err := oleutil.GetProperty(item, "Name")
			if err != nil {
				return []string{}
			}
			rc.cols[i] = name.ToString()
			item.Release()
		}
	}
	return rc.cols
}

func (rc *AdodbRows) Next(dest []driver.Value) error {
	unknown, err := oleutil.GetProperty(rc.rc, "EOF")
	if err != nil {
		return io.EOF
	}
	if unknown.Val != 0 {
		return io.EOF
	}
	unknown, err = oleutil.GetProperty(rc.rc, "Fields")
	if err != nil {
		return err
	}
	fields := unknown.ToIDispatch()
	defer fields.Release()
	for i := range dest {
		var varval ole.VARIANT
		varval.VT = ole.VT_I4
		varval.Val = int64(i)
		val, err := oleutil.CallMethod(fields, "Item", &varval)
		if err != nil {
			return err
		}
		field := val.ToIDispatch()
		defer field.Release()
		typ, err := oleutil.GetProperty(field, "Type")
		if err != nil {
			return err
		}
		val, err = oleutil.GetProperty(field, "Value")
		if err != nil {
			return err
		}
		field.Release()
		switch typ.Val {
		case 0: // ADEMPTY
			dest[i] = nil
		case 2: // ADSMALLINT
			dest[i] = int64(int16(val.Val))
		case 3: // ADINTEGER
			dest[i] = int64(int32(val.Val))
		case 4: // ADSINGLE
			dest[i] = float64(math.Float32frombits(uint32(val.Val)))
		case 5: // ADDOUBLE
			dest[i] = math.Float64frombits(uint64(val.Val))
		case 6: // ADCURRENCY
			dest[i] = float64(val.Val)
		case 7: // ADDATE
			// TODO
		case 8: // ADBSTR
			dest[i] = val.ToString()
		case 9: // ADIDISPATCH
			dest[i] = val.ToIDispatch()
		case 10: // ADERROR
			// TODO
		case 11: // ADBOOLEAN
			if val.Val != 0 {
				dest[i] = true
			} else {
				dest[i] = false
			}
		case 12: // ADVARIANT
			dest[i] = val
		case 13: // ADIUNKNOWN
			dest[i] = val.ToIUnknown()
		case 14: // ADDECIMAL
			dest[i] = float64(val.Val)
		case 16: // ADTINYINT
			dest[i] = int8(val.Val)
		case 17: // ADUNSIGNEDTINYINT
			dest[i] = uint8(val.Val)
		case 18: // ADUNSIGNEDSMALLINT
			dest[i] = uint16(val.Val)
		case 19: // ADUNSIGNEDINT
			dest[i] = uint32(val.Val)
		case 20: // ADBIGINT
			dest[i] = big.NewInt(val.Val)
		case 21: // ADUNSIGNEDBIGINT
			// TODO
		case 72: // ADGUID
			dest[i] = val.ToString()
		case 128: // ADBINARY
			sa := *(**ole.SAFEARRAY)(unsafe.Pointer(&val.Val))
			dest[i] = (*[1 << 30]byte)(unsafe.Pointer(uintptr(sa.PvData)))[0:sa.CbElements]
		case 129: // ADCHAR
			dest[i] = val.ToString() //uint8(val.Val)
		case 130: // ADWCHAR
			dest[i] = val.ToString() //uint16(val.Val)
		case 131: // ADNUMERIC
			// TODO: handle numbers that aren't positive integers.
			dest[i] = val.Val
		case 132: // ADUSERDEFINED
			dest[i] = uintptr(val.Val)
		case 133: // ADDBDATE
			dest[i] = time.Unix(0, val.Val).UTC()
		case 134: // ADDBTIME
			dest[i] = time.Unix(0, val.Val).UTC()
		case 135: // ADDBTIMESTAMP
			dest[i] = time.Unix(0, val.Val).UTC()
		case 136: // ADCHAPTER
			dest[i] = val.ToString()
		case 200: // ADVARCHAR
			dest[i] = val.ToString()
		case 201: // ADLONGVARCHAR
			dest[i] = val.ToString()
		case 202: // ADVARWCHAR
			dest[i] = val.ToString()
		case 203: // ADLONGVARWCHAR
			dest[i] = val.ToString()
		case 204: // ADVARBINARY
			// TODO
		case 205: // ADLONGVARBINARY
			sa := (*ole.SAFEARRAY)(unsafe.Pointer(uintptr(val.Val)))
			dest[i] = (*[1 << 30]byte)(unsafe.Pointer(uintptr(sa.PvData)))[0:sa.RgsaBound.CElements]
		}
	}
	_, err = oleutil.CallMethod(rc.rc, "MoveNext")
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package algoholic

// A trie is a data structure which stores words in a tree by letter, when the letters form a
// valid string, the node is marked terminal and a value can be stored.

type Trie struct {
	Char     rune
	Children map[rune]*Trie
	Terminal bool
	Value    interface{}
}

// A value indicating that the node is the root node, thus it doesn't denote a character
// itself.
const RootTrie = rune(0)

// Create a new trie with the specified character.
func NewTrie(chr rune) *Trie {
	return &Trie{chr, make(map[rune]*Trie), false, nil}
}

func NewRootTrie() *Trie {
	return NewTrie(RootTrie)
}

// Create a new trie with strings mapped to specified values.
func NewTrieFromMap(strMap map[string]interface{}) *Trie {
	ret := NewRootTrie()

	for str, val := range strMap {
		ret.Insert(str, val)
	}

	return ret
}

// Create a new trie with strings whose values we don't care about.
func NewTrieFromStrings(strs []string) *Trie {
	ret := NewRootTrie()

	for _, str := range strs {
		ret.Insert(str, nil)
	}

	return ret
}

// Find the specified string and return its trie node.
// O(m) worst-case where m is the length of the string searched for.
// Note this returns non-terminal nodes.
func (trie *Trie) FindTrie(str string) *Trie {
	if len(str) == 0 {
		return trie
	}

	if next := trie.Children[rune(str[0])]; next != nil {
		return next.FindTrie(str[1:])
	}

	return nil
}

// Find the specified string and return its value.
// O(m) worst-case where m is the length of the string searched for.
func (trie *Trie) Find(str string) (val interface{}, has bool) {
	ret := trie.FindTrie(str)

	if ret == nil || !ret.Terminal {
		// Not found.
		return
	}

	has = true
	val = ret.Value

	return
}

// Find all valid strings that consist of suffixes of the input prefix.
// O(m) worst-case where m is the length of the longest returned string.
func (trie *Trie) FindSuffixes(prefix string) []string {
	trie = trie.FindTrie(prefix)

	if trie == nil {
		return nil
	}

	var ret []string

	for suffix := range trie.Walk() {
		ret = append(ret, prefix+suffix[1:])
	}

	return ret
}

// Insert string, value pair into the specified trie.
// O(m) worst-case where m is the length of the inserted string.
func (trie *Trie) Insert(str string, val interface{}) {
	var (
		i   int
		chr rune
	)

	// Search through existing nodes.
	for i, chr = range str {
		if next, has := trie.Children[chr]; has {
			trie = next
		} else {
			break
		}
	}

	// Insert nodes as necessary.
	for _, chr = range str[i:] {
		next := NewTrie(chr)
		trie.Children[chr] = next
		trie = next
	}
	trie.Terminal = true
	trie.Value = val
}

// Recursively walk through all children of the input trie, adding string, value pairs to
// trieMap as the walk is performed. Trie is traversed in pre-order.
// O(n) where n is the number of nodes in the input trie.
func (trie *Trie) doWalk(trieMap map[string]interface{}, prev []rune) {
	// TODO: Use something other than a hash for map to allow alphabetical output ordering.

	if trie.Char != RootTrie {
		prev = append(prev, trie.Char)
	}

	if trie.Terminal {
		str := string(prev)
		trieMap[str] = trie.Value
	}

	for _, child := range trie.Children {
		child.doWalk(trieMap, prev)
	}
}

// Recursively walk through all children of the input trie, returning a map of string, value
// pairs.
// O(n) where n is the number of nodes in the input trie.
func (trie *Trie) Walk() map[string]interface{} {
	ret := make(map[string]interface{})
	trie.doWalk(ret, nil)
	return ret
}
<commit_msg>Also store parents in trie nodes.<commit_after>package algoholic

// A trie is a data structure which stores words in a tree by letter, when the letters form a
// valid string, the node is marked terminal and a value can be stored.

type Trie struct {
	Char     rune
	Parent   *Trie
	Children map[rune]*Trie
	Terminal bool
	Value    interface{}
}

// A value indicating that the node is the root node, thus it doesn't denote a character
// itself.
const RootTrie = rune(0)

// Create a new trie with the specified character and parent.
func NewTrie(parent *Trie, chr rune) *Trie {
	return &Trie{chr, parent, make(map[rune]*Trie), false, nil}
}

func NewRootTrie() *Trie {
	return NewTrie(nil, RootTrieChar)
}

// Create a new trie with strings mapped to specified values.
func NewTrieFromMap(strMap map[string]interface{}) *Trie {
	ret := NewRootTrie()

	for str, val := range strMap {
		ret.Insert(str, val)
	}

	return ret
}

// Create a new trie with strings whose values we don't care about.
func NewTrieFromStrings(strs []string) *Trie {
	ret := NewRootTrie()

	for _, str := range strs {
		ret.Insert(str, nil)
	}

	return ret
}

// Find the specified string and return its trie node.
// O(m) worst-case where m is the length of the string searched for.
// Note this returns non-terminal nodes.
func (trie *Trie) FindTrie(str string) *Trie {
	if len(str) == 0 {
		return trie
	}

	if next := trie.Children[rune(str[0])]; next != nil {
		return next.FindTrie(str[1:])
	}

	return nil
}

// Find the specified string and return its value.
// O(m) worst-case where m is the length of the string searched for.
func (trie *Trie) Find(str string) (val interface{}, has bool) {
	ret := trie.FindTrie(str)

	if ret == nil || !ret.Terminal {
		// Not found.
		return
	}

	has = true
	val = ret.Value

	return
}

// Find all valid strings that consist of suffixes of the input prefix.
// O(m) worst-case where m is the length of the longest returned string.
func (trie *Trie) FindSuffixes(prefix string) []string {
	trie = trie.FindTrie(prefix)

	if trie == nil {
		return nil
	}

	var ret []string

	for suffix := range trie.Walk() {
		ret = append(ret, prefix+suffix[1:])
	}

	return ret
}

// Insert string, value pair into the specified trie.
// O(m) worst-case where m is the length of the inserted string.
func (trie *Trie) Insert(str string, val interface{}) {
	var (
		i   int
		chr rune
	)

	// Search through existing nodes.
	for i, chr = range str {
		if next, has := trie.Children[chr]; has {
			trie = next
		} else {
			break
		}
	}

	// Insert nodes as necessary.
	for _, chr = range str[i:] {
		next := NewTrie(trie, chr)
		trie.Children[chr] = next
		trie = next
	}
	trie.Terminal = true
	trie.Value = val
}

// Recursively walk through all children of the input trie, adding string, value pairs to
// trieMap as the walk is performed. Trie is traversed in pre-order.
// O(n) where n is the number of nodes in the input trie.
func (trie *Trie) doWalk(trieMap map[string]interface{}, prev []rune) {
	// TODO: Use something other than a hash for map to allow alphabetical output ordering.

	if trie.Char != RootTrie {
		prev = append(prev, trie.Char)
	}

	if trie.Terminal {
		str := string(prev)
		trieMap[str] = trie.Value
	}

	for _, child := range trie.Children {
		child.doWalk(trieMap, prev)
	}
}

// Recursively walk through all children of the input trie, returning a map of string, value
// pairs.
// O(n) where n is the number of nodes in the input trie.
func (trie *Trie) Walk() map[string]interface{} {
	ret := make(map[string]interface{})
	trie.doWalk(ret, nil)
	return ret
}
<|endoftext|>


<commit_before>package gocd

import (
	"context"
	"fmt"
	"net/url"
)

type AgentsService service

//go:generate gocd-response-links -type=AgentsLinks,AgentLinks
type AgentsLinks struct {
	Self *url.URL
	Doc  *url.URL
}

type AgentLinks struct {
	Self *url.URL
	Doc  *url.URL
	Find *url.URL
}

type AgentsResponse struct {
	Links    AgentsLinks `json:"_links,omitempty"`
	Embedded struct {
		Agents []*Agent `json:"agents"`
	} `json:"_embedded"`
}

type Agent struct {
	Uuid             string        `json:"uuid"`
	Hostname         string        `json:"hostname"`
	ElasticAgentId   string        `json:"elastic_agent_id"`
	ElasticPluginId  string        `json:"elastic_plugin_id"`
	IpAddress        string        `json:"ip_address"`
	Sandbox          string        `json:"sandbox"`
	OperatingSystem  string        `json:"operating_system"`
	FreeSpace        int64         `json:"free_space"`
	AgentConfigState string        `json:"agent_config_state"`
	AgentState       string        `json:"agent_state"`
	Resources        []string      `json:"resources"`
	Environments     []string      `json:"environments"`
	BuildState       string        `json:"build_state"`
	BuildDetails     *BuildDetails `json:"build_details"`
	Links            *AgentLinks   `json:"_links"`
}

type BuildDetails struct {
	Links    *BuildDetailsLinks `json:"_links"`
	Pipeline string             `json:"pipeline"`
	Stage    string             `json:"stage"`
	Job      string             `json:"job"`
}

//go:generate gocd-response-links -type=BuildDetailsLinks
type BuildDetailsLinks struct {
	Job      *url.URL
	Stage    *url.URL
	Pipeline *url.URL
}

func (s *AgentsService) List(ctx context.Context) ([]*Agent, *APIResponse, error) {
	u, err := addOptions("agents")

	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", u, nil, apiV4)
	if err != nil {
		return nil, nil, err
	}

	r := AgentsResponse{}
	resp, err := s.client.Do(ctx, req, &r)
	if err != nil {
		return nil, resp, err
	}

	//return &r.Embedded, resp, nil
	return r.Embedded.Agents, resp, nil
}

func (s *AgentsService) Get(ctx context.Context, uuid string) (*Agent, *APIResponse, error) {
	return s.handleAgentRequest(ctx, "GET", uuid, nil)
}

func (s *AgentsService) Update(ctx context.Context, uuid string, agent *Agent) (*Agent, *APIResponse, error) {
	return s.handleAgentRequest(ctx, "PATCH", uuid, agent)
}

func (s *AgentsService) handleAgentRequest(ctx context.Context, action string, uuid string, body *Agent) (*Agent, *APIResponse, error) {
	u, err := addOptions(fmt.Sprintf("agents/%s", uuid))
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest(action, u, body, apiV4)
	if err != nil {
		return nil, nil, err
	}

	a := Agent{}
	resp, err := s.client.Do(ctx, req, &a)
	if err != nil {
		return nil, resp, err
	}

	return &a, resp, nil
}
<commit_msg>Autocommit for 'gofmt -w'.<commit_after>package gocd

import (
	"context"
	"fmt"
	"net/url"
)

type AgentsService service

//go:generate gocd-response-links -type=AgentsLinks,AgentLinks
type AgentsLinks struct {
	Self *url.URL
	Doc  *url.URL
}

type AgentLinks struct {
	Self *url.URL
	Doc  *url.URL
	Find *url.URL
}

type AgentsResponse struct {
	Links    AgentsLinks `json:"_links,omitempty"`
	Embedded struct {
		Agents []*Agent `json:"agents"`
	} `json:"_embedded"`
}

type Agent struct {
	Uuid             string        `json:"uuid"`
	Hostname         string        `json:"hostname"`
	ElasticAgentId   string        `json:"elastic_agent_id"`
	ElasticPluginId  string        `json:"elastic_plugin_id"`
	IpAddress        string        `json:"ip_address"`
	Sandbox          string        `json:"sandbox"`
	OperatingSystem  string        `json:"operating_system"`
	FreeSpace        int64         `json:"free_space"`
	AgentConfigState string        `json:"agent_config_state"`
	AgentState       string        `json:"agent_state"`
	Resources        []string      `json:"resources"`
	Environments     []string      `json:"environments"`
	BuildState       string        `json:"build_state"`
	BuildDetails     *BuildDetails `json:"build_details"`
	Links            *AgentLinks   `json:"_links"`
}

type BuildDetails struct {
	Links    *BuildDetailsLinks `json:"_links"`
	Pipeline string             `json:"pipeline"`
	Stage    string             `json:"stage"`
	Job      string             `json:"job"`
}

//go:generate gocd-response-links -type=BuildDetailsLinks
type BuildDetailsLinks struct {
	Job      *url.URL `json:"job"`
	Stage    *url.URL `json:"stage"`
	Pipeline *url.URL `json:"pipeline"`
}

func (s *AgentsService) List(ctx context.Context) ([]*Agent, *APIResponse, error) {
	u, err := addOptions("agents")

	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest("GET", u, nil, apiV4)
	if err != nil {
		return nil, nil, err
	}

	r := AgentsResponse{}
	resp, err := s.client.Do(ctx, req, &r)
	if err != nil {
		return nil, resp, err
	}

	//return &r.Embedded, resp, nil
	return r.Embedded.Agents, resp, nil
}

func (s *AgentsService) Get(ctx context.Context, uuid string) (*Agent, *APIResponse, error) {
	return s.handleAgentRequest(ctx, "GET", uuid, nil)
}

func (s *AgentsService) Update(ctx context.Context, uuid string, agent *Agent) (*Agent, *APIResponse, error) {
	return s.handleAgentRequest(ctx, "PATCH", uuid, agent)
}

func (s *AgentsService) handleAgentRequest(ctx context.Context, action string, uuid string, body *Agent) (*Agent, *APIResponse, error) {
	u, err := addOptions(fmt.Sprintf("agents/%s", uuid))
	if err != nil {
		return nil, nil, err
	}

	req, err := s.client.NewRequest(action, u, body, apiV4)
	if err != nil {
		return nil, nil, err
	}

	a := Agent{}
	resp, err := s.client.Do(ctx, req, &a)
	if err != nil {
		return nil, resp, err
	}

	return &a, resp, nil
}
<|endoftext|>


<commit_before>package amigo

import (
	"os"

	"github.com/pelletier/go-toml"
)

const EnvMapKey = "envmap"

type Config struct {
	env      map[string]string
	confFile *toml.TomlTree
}

// Associate a given config file key with an environment var
func (c *Config) Env(confKey string, envKey string) {
	if val := os.Getenv(envKey); val != "" {
		c.env[confKey] = val
	}
}

func (c *Config) Get(key string) interface{} {
	if val, ok := c.env[key]; ok {
		return val
	}
	return c.confFile.Get(key)
}

func New(filepath string) (*Config, error) {
	file, err := toml.LoadFile(filepath)

	if err != nil {
		return nil, err
	}

	c := &Config{}
	c.confFile = file

	if file.Has(EnvMapKey) {
		envmap := file.Get(EnvMapKey).(*toml.TomlTree)
		for _, confKey := range envmap.Keys() {
			envKey := envmap.Get(confKey).(string)
			c.Env(confKey, envKey)
		}
	}
	return c, nil
}
<commit_msg>some comments<commit_after>package amigo

import (
	"os"

	"github.com/pelletier/go-toml"
)

const EnvMapKey = "envmap"

type Config struct {
	env      map[string]string
	confFile *toml.TomlTree
}

// Associate a given config file key with an environment var
func (c *Config) Env(confKey string, envKey string) {
	if val := os.Getenv(envKey); val != "" {
		c.env[confKey] = val
	}
}

func (c *Config) Get(key string) interface{} {
	if val, ok := c.env[key]; ok {
		return val
	}
	return c.confFile.Get(key)
}

// Return a new configuration object for use by library consumers
func New(filepath string) (*Config, error) {
	file, err := toml.LoadFile(filepath)

	if err != nil {
		return nil, err
	}

	c := &Config{}
	c.confFile = file

	// if a envmap table is defined, associate the specified keys with the env
	// vars defiend there
	if file.Has(EnvMapKey) {
		envmap := file.Get(EnvMapKey).(*toml.TomlTree)
		for _, confKey := range envmap.Keys() {
			envKey := envmap.Get(confKey).(string)
			c.Env(confKey, envKey)
		}
	}
	return c, nil
}
<|endoftext|>


<commit_before>package parse

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"sort"
	"strings"

	"github.com/thought-machine/please/rules"
	"github.com/thought-machine/please/rules/bazel"
	"github.com/thought-machine/please/src/cli"
	"github.com/thought-machine/please/src/core"
	"github.com/thought-machine/please/src/parse/asp"
)

// InitParser initialises the parser engine. This is guaranteed to be called exactly once before any calls to Parse().
func InitParser(state *core.BuildState) *core.BuildState {
	if state.Parser == nil {
		p := &aspParser{parser: newAspParser(state), init: make(chan struct{})}
		state.Parser = p
		go p.preloadSubincludes(state)
	}
	return state
}

// aspParser implements the core.Parser interface around our parser package.
type aspParser struct {
	parser *asp.Parser
	init   chan struct{}
}

// newAspParser returns a asp.Parser object with all the builtins loaded
func newAspParser(state *core.BuildState) *asp.Parser {
	p := asp.NewParser(state)
	log.Debug("Loading built-in build rules...")
	dir, _ := rules.AllAssets(state.ExcludedBuiltinRules())
	sort.Strings(dir)
	for _, filename := range dir {
		src, _ := rules.ReadAsset(filename)
		p.MustLoadBuiltins(filename, src)
	}

	for _, preload := range state.Config.Parse.PreloadBuildDefs {
		log.Debug("Preloading build defs from %s...", preload)
		p.MustLoadBuiltins(preload, nil)
	}

	if state.Config.Bazel.Compatibility {
		// Add a subrepo for @bazel_tools which appears to be one of their builtins.
		// Mostly we only include build defs in there.
		createBazelSubrepo(state)
	}

	log.Debug("parser initialised")
	return p
}

// NewParser creates a new parser for the state
func (p *aspParser) NewParser(state *core.BuildState) {
	// TODO(jpoole): remove this once we refactor core so it can depend on this package and call this itself
	state.Parser = nil
	InitParser(state)
}

func (p *aspParser) WaitForInit() {
	<-p.init
}

func (p *aspParser) preloadSubincludes(state *core.BuildState) {
	includes := state.Config.Parse.PreloadSubincludes
	if state.RepoConfig != nil {
		// TODO(jpoole): is this the right thing to do?
		includes = append(includes, state.RepoConfig.Parse.PreloadSubincludes...)
	}
	for _, inc := range includes {
		if inc.IsPseudoTarget() {
			log.Fatalf("Can't preload pseudotarget %v", inc)
		}
		// Queue them up asynchronously to feed the queues as quickly as possible
		go func(inc core.BuildLabel) {
			state.WaitForBuiltTarget(inc, core.OriginalTarget)
		}(inc)
	}

	// Preload them in order to avoid non-deterministic errors when the subincludes depend on each other
	for _, inc := range includes {
		if err := p.parser.SubincludeTarget(state, state.WaitForTargetAndEnsureDownload(inc, core.OriginalTarget)); err != nil {
			log.Fatalf("%v", err)
		}
	}
	p.parser.Finalise()
	close(p.init)
}

func (p *aspParser) ParseFile(pkg *core.Package, filename string) error {
	return p.parser.ParseFile(pkg, filename)
}

func (p *aspParser) ParseReader(pkg *core.Package, reader io.ReadSeeker) error {
	_, err := p.parser.ParseReader(pkg, reader)
	return err
}

func (p *aspParser) RunPreBuildFunction(threadID int, state *core.BuildState, target *core.BuildTarget) error {
	return p.runBuildFunction(threadID, state, target, "pre", func() error {
		return target.PreBuildFunction.Call(target)
	})
}

func (p *aspParser) RunPostBuildFunction(threadID int, state *core.BuildState, target *core.BuildTarget, output string) error {
	return p.runBuildFunction(threadID, state, target, "post", func() error {
		log.Debug("Running post-build function for %s. Build output:\n%s", target.Label, output)
		return target.PostBuildFunction.Call(target, output)
	})
}

// BuildRuleArgOrder returns a map of the arguments to build rule and the order they appear in the source file
func (p *aspParser) BuildRuleArgOrder() map[string]int {
	return p.parser.BuildRuleArgOrder()
}

// runBuildFunction runs either the pre- or post-build function.
func (p *aspParser) runBuildFunction(tid int, state *core.BuildState, target *core.BuildTarget, callbackType string, f func() error) error {
	state.LogBuildResult(tid, target, core.PackageParsing, fmt.Sprintf("Running %s-build function for %s", callbackType, target.Label))
	state.SyncParsePackage(target.Label)
	if err := f(); err != nil {
		state.LogBuildError(tid, target.Label, core.ParseFailed, err, "Failed %s-build function for %s", callbackType, target.Label)
		return err
	}
	state.LogBuildResult(tid, target, core.TargetBuilding, fmt.Sprintf("Finished %s-build function for %s", callbackType, target.Label))
	return nil
}

func createBazelSubrepo(state *core.BuildState) {
	if sr := state.Graph.Subrepo("bazel_tools"); sr != nil {
		return
	}
	dir := path.Join(core.OutDir, "bazel_tools")
	state.Graph.AddSubrepo(&core.Subrepo{
		Name:  "bazel_tools",
		Root:  dir,
		State: state,
		Arch:  cli.HostArch(),
	})
	// TODO(peterebden): This is a bit yuck... would be nice if we could avoid hardcoding all
	//                   this upfront and add a build target to do it for us.
	dir = path.Join(dir, "tools/build_defs/repo")
	if err := os.MkdirAll(dir, core.DirPermissions); err != nil {
		log.Fatalf("%s", err)
	}
	for filename, data := range bazel.AllFiles() {
		if err := ioutil.WriteFile(path.Join(dir, strings.ReplaceAll(filename, ".build_defs", ".bzl")), data, 0644); err != nil {
			log.Fatalf("%s", err)
		}
	}
}
<commit_msg>Fix race while preloading (#2405)<commit_after>package parse

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"sort"
	"strings"
	"sync"

	"github.com/thought-machine/please/rules"
	"github.com/thought-machine/please/rules/bazel"
	"github.com/thought-machine/please/src/cli"
	"github.com/thought-machine/please/src/core"
	"github.com/thought-machine/please/src/parse/asp"
)

// InitParser initialises the parser engine. This is guaranteed to be called exactly once before any calls to Parse().
func InitParser(state *core.BuildState) *core.BuildState {
	if state.Parser == nil {
		p := &aspParser{parser: newAspParser(state), init: make(chan struct{})}
		state.Parser = p
		go p.preloadSubincludes(state)
	}
	return state
}

// aspParser implements the core.Parser interface around our parser package.
type aspParser struct {
	parser *asp.Parser
	init   chan struct{}
}

// newAspParser returns a asp.Parser object with all the builtins loaded
func newAspParser(state *core.BuildState) *asp.Parser {
	p := asp.NewParser(state)
	log.Debug("Loading built-in build rules...")
	dir, _ := rules.AllAssets(state.ExcludedBuiltinRules())
	sort.Strings(dir)
	for _, filename := range dir {
		src, _ := rules.ReadAsset(filename)
		p.MustLoadBuiltins(filename, src)
	}

	for _, preload := range state.Config.Parse.PreloadBuildDefs {
		log.Debug("Preloading build defs from %s...", preload)
		p.MustLoadBuiltins(preload, nil)
	}

	if state.Config.Bazel.Compatibility {
		// Add a subrepo for @bazel_tools which appears to be one of their builtins.
		// Mostly we only include build defs in there.
		createBazelSubrepo(state)
	}

	log.Debug("parser initialised")
	return p
}

// NewParser creates a new parser for the state
func (p *aspParser) NewParser(state *core.BuildState) {
	// TODO(jpoole): remove this once we refactor core so it can depend on this package and call this itself
	state.Parser = nil
	InitParser(state)
}

func (p *aspParser) WaitForInit() {
	<-p.init
}

func (p *aspParser) preloadSubincludes(state *core.BuildState) {
	includes := state.Config.Parse.PreloadSubincludes
	if state.RepoConfig != nil {
		// TODO(jpoole): is this the right thing to do?
		includes = append(includes, state.RepoConfig.Parse.PreloadSubincludes...)
	}
	wg := sync.WaitGroup{}
	for _, inc := range includes {
		if inc.IsPseudoTarget() {
			log.Fatalf("Can't preload pseudotarget %v", inc)
		}
		wg.Add(1)
		// Queue them up asynchronously to feed the queues as quickly as possible
		go func(inc core.BuildLabel) {
			state.WaitForBuiltTarget(inc, core.OriginalTarget)
			wg.Done()
		}(inc)
	}

	// We must wait for all the subinclude targets to be built otherwise updating the locals might race with parsing
	// a package
	wg.Wait()

	// Preload them in order to avoid non-deterministic errors when the subincludes depend on each other
	for _, inc := range includes {
		if err := p.parser.SubincludeTarget(state, state.WaitForTargetAndEnsureDownload(inc, core.OriginalTarget)); err != nil {
			log.Fatalf("%v", err)
		}
	}
	p.parser.Finalise()
	close(p.init)
}

func (p *aspParser) ParseFile(pkg *core.Package, filename string) error {
	return p.parser.ParseFile(pkg, filename)
}

func (p *aspParser) ParseReader(pkg *core.Package, reader io.ReadSeeker) error {
	_, err := p.parser.ParseReader(pkg, reader)
	return err
}

func (p *aspParser) RunPreBuildFunction(threadID int, state *core.BuildState, target *core.BuildTarget) error {
	return p.runBuildFunction(threadID, state, target, "pre", func() error {
		return target.PreBuildFunction.Call(target)
	})
}

func (p *aspParser) RunPostBuildFunction(threadID int, state *core.BuildState, target *core.BuildTarget, output string) error {
	return p.runBuildFunction(threadID, state, target, "post", func() error {
		log.Debug("Running post-build function for %s. Build output:\n%s", target.Label, output)
		return target.PostBuildFunction.Call(target, output)
	})
}

// BuildRuleArgOrder returns a map of the arguments to build rule and the order they appear in the source file
func (p *aspParser) BuildRuleArgOrder() map[string]int {
	return p.parser.BuildRuleArgOrder()
}

// runBuildFunction runs either the pre- or post-build function.
func (p *aspParser) runBuildFunction(tid int, state *core.BuildState, target *core.BuildTarget, callbackType string, f func() error) error {
	state.LogBuildResult(tid, target, core.PackageParsing, fmt.Sprintf("Running %s-build function for %s", callbackType, target.Label))
	state.SyncParsePackage(target.Label)
	if err := f(); err != nil {
		state.LogBuildError(tid, target.Label, core.ParseFailed, err, "Failed %s-build function for %s", callbackType, target.Label)
		return err
	}
	state.LogBuildResult(tid, target, core.TargetBuilding, fmt.Sprintf("Finished %s-build function for %s", callbackType, target.Label))
	return nil
}

func createBazelSubrepo(state *core.BuildState) {
	if sr := state.Graph.Subrepo("bazel_tools"); sr != nil {
		return
	}
	dir := path.Join(core.OutDir, "bazel_tools")
	state.Graph.AddSubrepo(&core.Subrepo{
		Name:  "bazel_tools",
		Root:  dir,
		State: state,
		Arch:  cli.HostArch(),
	})
	// TODO(peterebden): This is a bit yuck... would be nice if we could avoid hardcoding all
	//                   this upfront and add a build target to do it for us.
	dir = path.Join(dir, "tools/build_defs/repo")
	if err := os.MkdirAll(dir, core.DirPermissions); err != nil {
		log.Fatalf("%s", err)
	}
	for filename, data := range bazel.AllFiles() {
		if err := ioutil.WriteFile(path.Join(dir, strings.ReplaceAll(filename, ".build_defs", ".bzl")), data, 0644); err != nil {
			log.Fatalf("%s", err)
		}
	}
}
<|endoftext|>