<commit_before>package exchange

import (
	"bytes"
	"errors"
	"testing"
	"time"

	"github.com/jinzhu/gorm"
	_ "github.com/mattn/go-sqlite3"
	"github.com/qor/qor"
	"github.com/qor/qor/resource"
)

type User struct {
	Id           int64
	Name         string
	Age          int
	CellPhone    Phone
	Addresses    []Address
	OldAddresses []OldAddress

	// TODO
	// Pointer type in exchange.Export, for example
	CellPhonePtr *Phone
	// AddressesPtr []*Address
}

type Address struct {
	Id      int64
	UserId  int64
	Name    string
	Country string

	Phone Phone
}

type OldAddress struct {
	Id      int64
	UserId  int64
	Name    string
	Country string
}

type Phone struct {
	Id        int64
	UserId    int64
	AddressId int64
	Num       string
	CreatedAt time.Time
}

var (
	testdb = func() *gorm.DB {
		db, err := gorm.Open("sqlite3", "/tmp/qor_exchange_test.db")
		if err != nil {
			panic(err)
		}

		return &db
	}()
)

func cleanup() {
	testdb.DropTable(&User{})
	testdb.AutoMigrate(&User{})
	testdb.DropTable(&Address{})
	testdb.AutoMigrate(&Address{})
	testdb.DropTable(&OldAddress{})
	testdb.AutoMigrate(&OldAddress{})
	testdb.DropTable(&Phone{})
	testdb.AutoMigrate(&Phone{})
}

func TestImportSimple(t *testing.T) {
	cleanup()

	useres := NewResource(User{})
	useres.RegisterMeta(&resource.Meta{Name: "Name", Label: "Name"})
	useres.RegisterMeta(&resource.Meta{Name: "Age", Label: "Age"})
	ex := New(useres)

	f, err := NewXLSXFile("simple.xlsx")
	if err != nil {
		t.Fatal(err)
	}

	var buf bytes.Buffer
	err = ex.Import(f, &buf, &qor.Context{Config: &qor.Config{DB: testdb}})
	if err != nil {
		t.Error(err)
	}

	var users []User
	testdb.Find(&users)
	if len(users) != 3 {
		t.Errorf("should get 3 records, but got %d", len(users))
	}
}

func TestImportNested(t *testing.T) {
	cleanup()

	useres := NewResource(User{})
	useres.RegisterMeta(&resource.Meta{Name: "Name", Label: "Name"})
	useres.RegisterMeta(&resource.Meta{Name: "Age", Label: "Age"})
	addres := NewResource(Address{})
	addres.HasSequentialColumns = true
	useres.RegisterMeta(&resource.Meta{Name: "Addresses", Resource: addres})
	addres.RegisterMeta(&resource.Meta{Name: "Country", Label: "Country"})
	ex := New(useres)

	f, err := NewXLSXFile("nested_resource.xlsx")
	if err != nil {
		t.Fatal(err)
	}

	var buf bytes.Buffer
	err = ex.Import(f, &buf, &qor.Context{Config: &qor.Config{DB: testdb}})
	if err != nil {
		t.Error(err)
	}

	var users []User
	testdb.Find(&users)
	if len(users) != 3 {
		t.Errorf("should get 3 users, but got %d", len(users))
	}
	var addresses []Address
	testdb.Find(&addresses)
	if len(addresses) != 6 {
		t.Errorf("should get 6 addresses, but got %d", len(addresses))
	}
}

type FullMarathon struct {
	RunningLevel float64
	Min1500      int
	Sec1500      int
}

func TestImportNormalizeHeader(t *testing.T) {
	testdb.DropTable(FullMarathon{})
	testdb.AutoMigrate(FullMarathon{})

	marathon := NewResource(FullMarathon{})
	marathon.RegisterMeta(&resource.Meta{Name: "RunningLevel", Label: "Running Level"})
	marathon.RegisterMeta(&resource.Meta{Name: "Min1500", Label: "1500M Min"})
	marathon.RegisterMeta(&resource.Meta{Name: "Sec1500", Label: "1500M Sec"})
	ex := New(marathon)
	ex.JobThrottle = 10
	ex.DataStartAt = 3
	ex.NormalizeHeaders = func(f File) (headers []string) {
		if f.TotalLines() < 3 {
			return
		}

		topHeaders, subHeaders := f.Line(1), f.Line(2)
		if len(topHeaders) != len(subHeaders) {
			return
		}
		for i, subHeader := range subHeaders {
			var prefix string
			topSec := topHeaders[:i+1]
			lenSec := len(topSec)
			for j := lenSec - 1; j >= 0; j-- {
				if topHeader := topSec[j]; topHeader != "" {
					prefix = topHeader
					break
				}
			}

			headers = append(headers, prefix+" "+subHeader)
		}

		return
	}
	f, err := NewXLSXFile("headers.xlsx")
	if err != nil {
		t.Error(err)
	}

	var buf bytes.Buffer
	err = ex.Import(f, &buf, &qor.Context{Config: &qor.Config{DB: testdb}})
	if err != nil {
		t.Error(err)
	}

	var marathones []FullMarathon
	testdb.Find(&marathones)
	if len(marathones) != 12 {
		t.Errorf("should get 12 records, but got %d", len(marathones))
	}
	if marathones[1].RunningLevel != 28.1 {
		t.Errorf("should get 28.1, but got %f", marathones[1].RunningLevel)
	}
	if marathones[1].Min1500 != 8 {
		t.Errorf("should get 8, but got %f", marathones[1].Min1500)
	}
	if marathones[1].Sec1500 != 26 {
		t.Errorf("should get 26, but got %f", marathones[1].Sec1500)
	}
}

func TestImportError(t *testing.T) {
	cleanup()

	useres := NewResource(User{})
	useres.RegisterMeta(&resource.Meta{Name: "Name", Label: "Name"})
	useres.RegisterMeta(&resource.Meta{Name: "Age", Label: "Age"})
	ex := New(useres)
	ferr := errors.New("an validator error in the second line")
	var i int
	useres.AddValidator(func(rel interface{}, mvs *resource.MetaValues, ctx *qor.Context) error {
		if i++; i == 2 {
			return ferr
		}
		return nil
	})

	f, err := NewXLSXFile("simple.xlsx")
	if err != nil {
		t.Error(err)
	}
	var buf bytes.Buffer
	err = ex.Import(f, &buf, &qor.Context{Config: &qor.Config{DB: testdb}})
	if err == nil {
		t.Error("Should encouter error")
	}

	logs := "2/4: 1 Van 24  \n3/4: an validator error in the second line; \n4/4: 3 Kate 25  \n"
	if buf.String() != logs {
		t.Errorf(`Expect log %q but got %q`, logs, buf.String())
	}

	var users []User
	testdb.Find(&users)
	if len(users) != 0 {
		t.Errorf("should get 0 records, but got %d", len(users))
	}
}

func TestExport(t *testing.T) {
	cleanup()
	records := []User{
		User{
			Name: "Van",
			Age:  24,
			Addresses: []Address{
				{Country: "China", Phone: Phone{Num: "xxx-xxx-xxx-0"}},
				{Country: "Japan", Phone: Phone{Num: "zzz-zzz-zzz-0"}},
				{Country: "New Zealand", Phone: Phone{Num: "kkk-kkk-kkk-0"}},
			},
			OldAddresses: []OldAddress{
				{Country: "Africa"},
			},
			CellPhone:    Phone{Num: "yyy-yyy-yyy-0"},
			CellPhonePtr: &Phone{Num: "yyy-yyy-yyy-0"},
		},
		User{
			Name: "Jane",
			Age:  26,
			Addresses: []Address{
				{Country: "USA", Phone: Phone{Num: "xxx-xxx-xxx-1"}},
			},
			OldAddresses: []OldAddress{
				{Country: "Africa"},
				{Country: "Brazil"},
			},
			CellPhone: Phone{Num: "yyy-yyy-yyy-1"},
		},
	}

	for _, record := range records {
		testdb.Create(&record)
	}

	var (
		user       = NewResource(User{})
		address    = NewResource(Address{})
		oldAddress = NewResource(Address{})
		phone      = NewResource(Phone{})
		cellphone  = NewResource(Phone{})
	)
	phone.HasSequentialColumns = true
	address.HasSequentialColumns = true
	oldAddress.MultiDelimiter = ","

	user.RegisterMeta(&resource.Meta{Name: "Name", Label: "Name"})
	user.RegisterMeta(&resource.Meta{Name: "Age", Label: "Age"})
	user.RegisterMeta(&resource.Meta{Name: "CellPhone", Resource: cellphone})
	user.RegisterMeta(&resource.Meta{Name: "Addresses", Resource: address})
	user.RegisterMeta(&resource.Meta{Name: "OldAddresses", Resource: oldAddress})
	address.RegisterMeta(&resource.Meta{Name: "Country", Label: "Country"})
	address.RegisterMeta(&resource.Meta{Name: "Phone", Resource: phone})
	phone.RegisterMeta(&resource.Meta{Name: "Num", Label: "Phone"})
	cellphone.RegisterMeta(&resource.Meta{Name: "Num", Label: "CellPhone"})
	oldAddress.RegisterMeta(&resource.Meta{Name: "Country", Label: "Old Countries"})

	ex := New(user)
	var buf bytes.Buffer
	records = []User{}

	db := testdb.Find(&records)
	for i, record := range records {
		testdb.Model(&record).Related(&record.CellPhone)
		testdb.Model(&record).Related(&record.Addresses)
		testdb.Model(&record).Related(&record.OldAddresses)
		for j, addr := range record.Addresses {
			testdb.Model(&addr).Related(&addr.Phone)
			record.Addresses[j] = addr
		}
		records[i] = record
	}
	if db.Error != nil {
		t.Fatal(db.Error)
	}
	ex.Export(records, &buf, &qor.Context{Config: &qor.Config{DB: testdb}})
	expect := `Name,Age,CellPhone,Country 01,Country 02,Country 03,Phone 01,Phone 02,Phone 03,Old Countries
Van,24,yyy-yyy-yyy-0,China,Japan,New Zealand,xxx-xxx-xxx-0,zzz-zzz-zzz-0,kkk-kkk-kkk-0,Africa
Jane,26,yyy-yyy-yyy-1,USA,,,xxx-xxx-xxx-1,,,"Africa,Brazil"
`

	if buf.String() != expect {
		t.Errorf("expect: %s\ngot: %s", expect, buf.String())
	}
}
<commit_msg>temporally ignore pointer type tests<commit_after>package exchange

import (
	"bytes"
	"errors"
	"testing"
	"time"

	"github.com/jinzhu/gorm"
	_ "github.com/mattn/go-sqlite3"
	"github.com/qor/qor"
	"github.com/qor/qor/resource"
)

type User struct {
	Id           int64
	Name         string
	Age          int
	CellPhone    Phone
	Addresses    []Address
	OldAddresses []OldAddress

	// TODO
	// Pointer type in exchange.Export, for example
	// CellPhonePtr *Phone
	// AddressesPtr []*Address
}

type Address struct {
	Id      int64
	UserId  int64
	Name    string
	Country string

	Phone Phone
}

type OldAddress struct {
	Id      int64
	UserId  int64
	Name    string
	Country string
}

type Phone struct {
	Id        int64
	UserId    int64
	AddressId int64
	Num       string
	CreatedAt time.Time
}

var (
	testdb = func() *gorm.DB {
		db, err := gorm.Open("sqlite3", "/tmp/qor_exchange_test.db")
		if err != nil {
			panic(err)
		}

		return &db
	}()
)

func cleanup() {
	testdb.DropTable(&User{})
	testdb.AutoMigrate(&User{})
	testdb.DropTable(&Address{})
	testdb.AutoMigrate(&Address{})
	testdb.DropTable(&OldAddress{})
	testdb.AutoMigrate(&OldAddress{})
	testdb.DropTable(&Phone{})
	testdb.AutoMigrate(&Phone{})
}

func TestImportSimple(t *testing.T) {
	cleanup()

	useres := NewResource(User{})
	useres.RegisterMeta(&resource.Meta{Name: "Name", Label: "Name"})
	useres.RegisterMeta(&resource.Meta{Name: "Age", Label: "Age"})
	ex := New(useres)

	f, err := NewXLSXFile("simple.xlsx")
	if err != nil {
		t.Fatal(err)
	}

	var buf bytes.Buffer
	err = ex.Import(f, &buf, &qor.Context{Config: &qor.Config{DB: testdb}})
	if err != nil {
		t.Error(err)
	}

	var users []User
	testdb.Find(&users)
	if len(users) != 3 {
		t.Errorf("should get 3 records, but got %d", len(users))
	}
}

func TestImportNested(t *testing.T) {
	cleanup()

	useres := NewResource(User{})
	useres.RegisterMeta(&resource.Meta{Name: "Name", Label: "Name"})
	useres.RegisterMeta(&resource.Meta{Name: "Age", Label: "Age"})
	addres := NewResource(Address{})
	addres.HasSequentialColumns = true
	useres.RegisterMeta(&resource.Meta{Name: "Addresses", Resource: addres})
	addres.RegisterMeta(&resource.Meta{Name: "Country", Label: "Country"})
	ex := New(useres)

	f, err := NewXLSXFile("nested_resource.xlsx")
	if err != nil {
		t.Fatal(err)
	}

	var buf bytes.Buffer
	err = ex.Import(f, &buf, &qor.Context{Config: &qor.Config{DB: testdb}})
	if err != nil {
		t.Error(err)
	}

	var users []User
	testdb.Find(&users)
	if len(users) != 3 {
		t.Errorf("should get 3 users, but got %d", len(users))
	}
	var addresses []Address
	testdb.Find(&addresses)
	if len(addresses) != 6 {
		t.Errorf("should get 6 addresses, but got %d", len(addresses))
	}
}

type FullMarathon struct {
	RunningLevel float64
	Min1500      int
	Sec1500      int
}

func TestImportNormalizeHeader(t *testing.T) {
	testdb.DropTable(FullMarathon{})
	testdb.AutoMigrate(FullMarathon{})

	marathon := NewResource(FullMarathon{})
	marathon.RegisterMeta(&resource.Meta{Name: "RunningLevel", Label: "Running Level"})
	marathon.RegisterMeta(&resource.Meta{Name: "Min1500", Label: "1500M Min"})
	marathon.RegisterMeta(&resource.Meta{Name: "Sec1500", Label: "1500M Sec"})
	ex := New(marathon)
	ex.JobThrottle = 10
	ex.DataStartAt = 3
	ex.NormalizeHeaders = func(f File) (headers []string) {
		if f.TotalLines() < 3 {
			return
		}

		topHeaders, subHeaders := f.Line(1), f.Line(2)
		if len(topHeaders) != len(subHeaders) {
			return
		}
		for i, subHeader := range subHeaders {
			var prefix string
			topSec := topHeaders[:i+1]
			lenSec := len(topSec)
			for j := lenSec - 1; j >= 0; j-- {
				if topHeader := topSec[j]; topHeader != "" {
					prefix = topHeader
					break
				}
			}

			headers = append(headers, prefix+" "+subHeader)
		}

		return
	}
	f, err := NewXLSXFile("headers.xlsx")
	if err != nil {
		t.Error(err)
	}

	var buf bytes.Buffer
	err = ex.Import(f, &buf, &qor.Context{Config: &qor.Config{DB: testdb}})
	if err != nil {
		t.Error(err)
	}

	var marathones []FullMarathon
	testdb.Find(&marathones)
	if len(marathones) != 12 {
		t.Errorf("should get 12 records, but got %d", len(marathones))
	}
	if marathones[1].RunningLevel != 28.1 {
		t.Errorf("should get 28.1, but got %f", marathones[1].RunningLevel)
	}
	if marathones[1].Min1500 != 8 {
		t.Errorf("should get 8, but got %f", marathones[1].Min1500)
	}
	if marathones[1].Sec1500 != 26 {
		t.Errorf("should get 26, but got %f", marathones[1].Sec1500)
	}
}

func TestImportError(t *testing.T) {
	cleanup()

	useres := NewResource(User{})
	useres.RegisterMeta(&resource.Meta{Name: "Name", Label: "Name"})
	useres.RegisterMeta(&resource.Meta{Name: "Age", Label: "Age"})
	ex := New(useres)
	ferr := errors.New("an validator error in the second line")
	var i int
	useres.AddValidator(func(rel interface{}, mvs *resource.MetaValues, ctx *qor.Context) error {
		if i++; i == 2 {
			return ferr
		}
		return nil
	})

	f, err := NewXLSXFile("simple.xlsx")
	if err != nil {
		t.Error(err)
	}
	var buf bytes.Buffer
	err = ex.Import(f, &buf, &qor.Context{Config: &qor.Config{DB: testdb}})
	if err == nil {
		t.Error("Should encouter error")
	}

	logs := "2/4: 1 Van 24  \n3/4: an validator error in the second line; \n4/4: 3 Kate 25  \n"
	if buf.String() != logs {
		t.Errorf(`Expect log %q but got %q`, logs, buf.String())
	}

	var users []User
	testdb.Find(&users)
	if len(users) != 0 {
		t.Errorf("should get 0 records, but got %d", len(users))
	}
}

func TestExport(t *testing.T) {
	cleanup()
	records := []User{
		User{
			Name: "Van",
			Age:  24,
			Addresses: []Address{
				{Country: "China", Phone: Phone{Num: "xxx-xxx-xxx-0"}},
				{Country: "Japan", Phone: Phone{Num: "zzz-zzz-zzz-0"}},
				{Country: "New Zealand", Phone: Phone{Num: "kkk-kkk-kkk-0"}},
			},
			OldAddresses: []OldAddress{
				{Country: "Africa"},
			},
			CellPhone: Phone{Num: "yyy-yyy-yyy-0"},
			// CellPhonePtr: &Phone{Num: "yyy-yyy-yyy-0"},
		},
		User{
			Name: "Jane",
			Age:  26,
			Addresses: []Address{
				{Country: "USA", Phone: Phone{Num: "xxx-xxx-xxx-1"}},
			},
			OldAddresses: []OldAddress{
				{Country: "Africa"},
				{Country: "Brazil"},
			},
			CellPhone: Phone{Num: "yyy-yyy-yyy-1"},
		},
	}

	for _, record := range records {
		testdb.Create(&record)
	}

	var (
		user       = NewResource(User{})
		address    = NewResource(Address{})
		oldAddress = NewResource(Address{})
		phone      = NewResource(Phone{})
		cellphone  = NewResource(Phone{})
	)
	phone.HasSequentialColumns = true
	address.HasSequentialColumns = true
	oldAddress.MultiDelimiter = ","

	user.RegisterMeta(&resource.Meta{Name: "Name", Label: "Name"})
	user.RegisterMeta(&resource.Meta{Name: "Age", Label: "Age"})
	user.RegisterMeta(&resource.Meta{Name: "CellPhone", Resource: cellphone})
	user.RegisterMeta(&resource.Meta{Name: "Addresses", Resource: address})
	user.RegisterMeta(&resource.Meta{Name: "OldAddresses", Resource: oldAddress})
	address.RegisterMeta(&resource.Meta{Name: "Country", Label: "Country"})
	address.RegisterMeta(&resource.Meta{Name: "Phone", Resource: phone})
	phone.RegisterMeta(&resource.Meta{Name: "Num", Label: "Phone"})
	cellphone.RegisterMeta(&resource.Meta{Name: "Num", Label: "CellPhone"})
	oldAddress.RegisterMeta(&resource.Meta{Name: "Country", Label: "Old Countries"})

	ex := New(user)
	var buf bytes.Buffer
	records = []User{}

	db := testdb.Find(&records)
	for i, record := range records {
		testdb.Model(&record).Related(&record.CellPhone)
		testdb.Model(&record).Related(&record.Addresses)
		testdb.Model(&record).Related(&record.OldAddresses)
		for j, addr := range record.Addresses {
			testdb.Model(&addr).Related(&addr.Phone)
			record.Addresses[j] = addr
		}
		records[i] = record
	}
	if db.Error != nil {
		t.Fatal(db.Error)
	}
	ex.Export(records, &buf, &qor.Context{Config: &qor.Config{DB: testdb}})
	expect := `Name,Age,CellPhone,Country 01,Country 02,Country 03,Phone 01,Phone 02,Phone 03,Old Countries
Van,24,yyy-yyy-yyy-0,China,Japan,New Zealand,xxx-xxx-xxx-0,zzz-zzz-zzz-0,kkk-kkk-kkk-0,Africa
Jane,26,yyy-yyy-yyy-1,USA,,,xxx-xxx-xxx-1,,,"Africa,Brazil"
`

	if buf.String() != expect {
		t.Errorf("expect: %s\ngot: %s", expect, buf.String())
	}
}
<|endoftext|>


<commit_before>// Copyright 2010 The Walk 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 walk

import (
	"fmt"
	"syscall"
	"unsafe"
)

import . "github.com/lxn/go-winapi"

type FileDialog struct {
	Title          string
	FilePath       string
	InitialDirPath string
	Filter         string
	FilterIndex    int
}

func (dlg *FileDialog) show(owner RootWidget, fun func(ofn *OPENFILENAME) bool) (accepted bool, err error) {
	ofn := &OPENFILENAME{}

	ofn.LStructSize = uint32(unsafe.Sizeof(*ofn))
	ofn.HwndOwner = owner.BaseWidget().hWnd

	filter := make([]uint16, len(dlg.Filter)+2)
	copy(filter, syscall.StringToUTF16(dlg.Filter))
	// Replace '|' with the expected '\0'.
	for i, c := range filter {
		if byte(c) == '|' {
			filter[i] = uint16(0)
		}
	}
	ofn.LpstrFilter = &filter[0]
	ofn.NFilterIndex = uint32(dlg.FilterIndex)

	filePath := make([]uint16, 1024)
	copy(filePath, syscall.StringToUTF16(dlg.FilePath))
	ofn.LpstrFile = &filePath[0]
	ofn.NMaxFile = uint32(len(filePath))

	ofn.LpstrInitialDir = syscall.StringToUTF16Ptr(dlg.InitialDirPath)
	ofn.LpstrTitle = syscall.StringToUTF16Ptr(dlg.Title)
	ofn.Flags = OFN_FILEMUSTEXIST

	if !fun(ofn) {
		errno := CommDlgExtendedError()
		if errno != 0 {
			err = newError(fmt.Sprintf("Error %d", errno))
		}
		return
	}

	dlg.FilePath = syscall.UTF16ToString(filePath)

	accepted = true

	return
}

func (dlg *FileDialog) ShowOpen(owner RootWidget) (accepted bool, err error) {
	return dlg.show(owner, GetOpenFileName)
}

func (dlg *FileDialog) ShowSave(owner RootWidget) (accepted bool, err error) {
	return dlg.show(owner, GetSaveFileName)
}
<commit_msg>FileDialog: Add ShowBrowseFolder, fixes #37, #43<commit_after>// Copyright 2010 The Walk 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 walk

import (
	"fmt"
	"syscall"
	"unsafe"
)

import . "github.com/lxn/go-winapi"

type FileDialog struct {
	Title          string
	FilePath       string
	InitialDirPath string
	Filter         string
	FilterIndex    int
}

func (dlg *FileDialog) show(owner RootWidget, fun func(ofn *OPENFILENAME) bool) (accepted bool, err error) {
	ofn := new(OPENFILENAME)

	ofn.LStructSize = uint32(unsafe.Sizeof(*ofn))
	if owner != nil {
		ofn.HwndOwner = owner.Handle()
	}

	filter := make([]uint16, len(dlg.Filter)+2)
	copy(filter, syscall.StringToUTF16(dlg.Filter))
	// Replace '|' with the expected '\0'.
	for i, c := range filter {
		if byte(c) == '|' {
			filter[i] = uint16(0)
		}
	}
	ofn.LpstrFilter = &filter[0]
	ofn.NFilterIndex = uint32(dlg.FilterIndex)

	filePath := make([]uint16, 1024)
	copy(filePath, syscall.StringToUTF16(dlg.FilePath))
	ofn.LpstrFile = &filePath[0]
	ofn.NMaxFile = uint32(len(filePath))

	ofn.LpstrInitialDir = syscall.StringToUTF16Ptr(dlg.InitialDirPath)
	ofn.LpstrTitle = syscall.StringToUTF16Ptr(dlg.Title)
	ofn.Flags = OFN_FILEMUSTEXIST

	if !fun(ofn) {
		errno := CommDlgExtendedError()
		if errno != 0 {
			err = newError(fmt.Sprintf("Error %d", errno))
		}
		return
	}

	dlg.FilePath = syscall.UTF16ToString(filePath)

	accepted = true

	return
}

func (dlg *FileDialog) ShowOpen(owner RootWidget) (accepted bool, err error) {
	return dlg.show(owner, GetOpenFileName)
}

func (dlg *FileDialog) ShowSave(owner RootWidget) (accepted bool, err error) {
	return dlg.show(owner, GetSaveFileName)
}

func (dlg *FileDialog) ShowBrowseFolder(owner RootWidget) (accepted bool, err error) {
	// Calling OleInitialize (or similar) is required for BIF_NEWDIALOGSTYLE.
	if hr := OleInitialize(); hr != S_OK && hr != S_FALSE {
		return false, newError(fmt.Sprint("OleInitialize Error: ", hr))
	}
	defer OleUninitialize()

	pathFromPIDL := func(pidl uintptr) (string, error) {
		var path [MAX_PATH]uint16
		if !SHGetPathFromIDList(pidl, &path[0]) {
			return "", newError("SHGetPathFromIDList failed")
		}

		return syscall.UTF16ToString(path[:]), nil
	}

	// We use this callback to disable the OK button in case of "invalid"
	// selections.
	callback := func(hwnd HWND, msg uint32, lp, wp uintptr) int32 {
		const BFFM_SELCHANGED = 2
		if msg == BFFM_SELCHANGED {
			_, err := pathFromPIDL(lp)
			var enabled uintptr
			if err == nil {
				enabled = 1
			}

			const BFFM_ENABLEOK = WM_USER + 101

			SendMessage(hwnd, BFFM_ENABLEOK, 0, enabled)
		}

		return 0
	}

	var ownerHwnd HWND
	if owner != nil {
		ownerHwnd = owner.Handle()
	}

	// We need to put the initial path into a buffer of at least MAX_LENGTH
	// length, or we may get random crashes.
	var buf [MAX_PATH]uint16
	copy(buf[:], syscall.StringToUTF16(dlg.InitialDirPath))

	const BIF_NEWDIALOGSTYLE = 0x00000040

	bi := BROWSEINFO{
		HwndOwner:      ownerHwnd,
		PszDisplayName: &buf[0],
		LpszTitle:      syscall.StringToUTF16Ptr(dlg.Title),
		UlFlags:        BIF_NEWDIALOGSTYLE,
		Lpfn:           syscall.NewCallback(callback),
	}

	pidl := SHBrowseForFolder(&bi)
	if pidl == 0 {
		return false, nil
	}
	defer CoTaskMemFree(pidl)

	dlg.FilePath, err = pathFromPIDL(pidl)
	accepted = dlg.FilePath != ""
	return
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/howeyc/fsnotify"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"syscall"
	"time"
)

// Milliseconds to wait for the next job to begin after a file change
const WorkDelay = 900

// Default pattern to match files which trigger a build
const FilePattern = `(.+\.go|.+\.c)$`

var (
	flag_directory = flag.String("directory", ".", "Directory to watch for changes")
	flag_pattern   = flag.String("pattern", FilePattern, "Pattern of watched files")
	flag_command   = flag.String("command", "", "Command to run and restart after build")
	flag_recursive = flag.Bool("recursive", true, "Watch all dirs. recursively")
	flag_build     = flag.String("build", "go build", "Command to rebuild after changes")
)

// Run `go build` and print the output if something's gone wrong.
func build() bool {
	log.Println("Running build command!")

	args := strings.Split(*flag_build, " ")
	if len(args) == 0 {
		// If the user has specified and empty then we are done.
		return true
	}

	cmd := exec.Command(args[0], args[1:]...)

	cmd.Dir = *flag_directory

	output, err := cmd.CombinedOutput()

	if err == nil {
		log.Println("Build ok.")
	} else {
		log.Println("Error while building:\n", string(output))
	}

	return err == nil
}

func matchesPattern(pattern *regexp.Regexp, file string) bool {
	return pattern.MatchString(file)
}

// Accept build jobs and start building when there are no jobs rushing in.
// The inrush protection is WorkDelay milliseconds long, in this period
// every incoming job will reset the timer.
func builder(jobs <-chan string, buildDone chan<- bool) {
	createThreshold := func() <-chan time.Time {
		return time.After(time.Duration(WorkDelay * time.Millisecond))
	}

	threshold := createThreshold()

	for {
		select {
		case <-jobs:
			threshold = createThreshold()
		case <-threshold:
			if build() {
				buildDone <- true
			}
		}
	}
}

func logger(stdoutChan <-chan io.ReadCloser) {
	dumper := func(pipe io.ReadCloser, prefix string) {
		reader := bufio.NewReader(pipe)

	readloop:
		for {
			line, err := reader.ReadString('\n')

			if err != nil {
				break readloop
			}

			log.Print(prefix, " ", line)
		}
	}

	for {
		pipe := <-stdoutChan

		go dumper(pipe, "stdout:")

		pipe = <-stdoutChan

		go dumper(pipe, "stderr:")
	}
}

// Run the command in the given string and restart it after
// a message was received on the buildDone channel.
func runner(command string, buildDone chan bool) {
	var currentProcess *os.Process

	stdoutChan := make(chan io.ReadCloser)

	go logger(stdoutChan)

	for {
		<-buildDone

		args := strings.Split(command, " ")
		cmd := exec.Command(args[0], args[1:]...)

		if currentProcess != nil {
			err := currentProcess.Kill()

			if err != nil {
				log.Fatal("Could not kill child process. Aborting due to danger of infinite forks.")
			}
		}

		log.Println("Restarting the given command.")

		pipe, err := cmd.StdoutPipe()

		if err != nil {
			log.Fatal("Can't get stdout pipe for command:", err)
		}

		stdoutChan <- pipe

		pipe, err = cmd.StderrPipe()

		if err != nil {
			log.Fatal("Can't get stderr pipe for command:", err)
		}

		stdoutChan <- pipe

		err = cmd.Start()

		if err != nil {
			log.Println("Error while running command:", err)
		}

		currentProcess = cmd.Process
	}
}

func flusher(buildDone <-chan bool) {
	for {
		<-buildDone
	}
}

func main() {
	flag.Parse()

	if *flag_directory == "" {
		fmt.Fprintf(os.Stderr, "-directory=... is required.\n")
		os.Exit(1)
	}

	watcher, err := fsnotify.NewWatcher()

	if err != nil {
		log.Fatal(err)
	}

	defer watcher.Close()

	if *flag_recursive == true {
		err = filepath.Walk(*flag_directory, func(path string, info os.FileInfo, err error) error {
			if err == nil && info.IsDir() {
				return watcher.Watch(path)
			}
			return err
		})

		if err != nil {
			log.Fatal("filepath.Walk():", err)
		}

	} else {
		err := watcher.Watch(*flag_directory)

		if err != nil {
			log.Fatal("watcher.Watch():", err)
		}
	}

	pattern := regexp.MustCompile(*flag_pattern)
	jobs := make(chan string)
	buildDone := make(chan bool)

	go builder(jobs, buildDone)

	if *flag_command != "" {
		go runner(*flag_command, buildDone)
	} else {
		go flusher(buildDone)
	}

	for {
		select {
		case ev := <-watcher.Event:
			if ev.Name != "" && matchesPattern(pattern, ev.Name) {
				jobs <- ev.Name
			}

		case err := <-watcher.Error:
			if v, ok := err.(*os.SyscallError); ok {
				if v.Err == syscall.EINTR {
					continue
				}
				log.Fatal("watcher.Error: SyscallError:", v)
			}
			log.Fatal("watcher.Error:", err)
		}
	}
}
<commit_msg>Read process exit status. Prevent zombies. The command processes have been accumulating as zombies because their exit statuses have not been read.<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/howeyc/fsnotify"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"syscall"
	"time"
)

// Milliseconds to wait for the next job to begin after a file change
const WorkDelay = 900

// Default pattern to match files which trigger a build
const FilePattern = `(.+\.go|.+\.c)$`

var (
	flag_directory = flag.String("directory", ".", "Directory to watch for changes")
	flag_pattern   = flag.String("pattern", FilePattern, "Pattern of watched files")
	flag_command   = flag.String("command", "", "Command to run and restart after build")
	flag_recursive = flag.Bool("recursive", true, "Watch all dirs. recursively")
	flag_build     = flag.String("build", "go build", "Command to rebuild after changes")
)

// Run `go build` and print the output if something's gone wrong.
func build() bool {
	log.Println("Running build command!")

	args := strings.Split(*flag_build, " ")
	if len(args) == 0 {
		// If the user has specified and empty then we are done.
		return true
	}

	cmd := exec.Command(args[0], args[1:]...)

	cmd.Dir = *flag_directory

	output, err := cmd.CombinedOutput()

	if err == nil {
		log.Println("Build ok.")
	} else {
		log.Println("Error while building:\n", string(output))
	}

	return err == nil
}

func matchesPattern(pattern *regexp.Regexp, file string) bool {
	return pattern.MatchString(file)
}

// Accept build jobs and start building when there are no jobs rushing in.
// The inrush protection is WorkDelay milliseconds long, in this period
// every incoming job will reset the timer.
func builder(jobs <-chan string, buildDone chan<- bool) {
	createThreshold := func() <-chan time.Time {
		return time.After(time.Duration(WorkDelay * time.Millisecond))
	}

	threshold := createThreshold()

	for {
		select {
		case <-jobs:
			threshold = createThreshold()
		case <-threshold:
			if build() {
				buildDone <- true
			}
		}
	}
}

func logger(stdoutChan <-chan io.ReadCloser) {
	dumper := func(pipe io.ReadCloser, prefix string) {
		reader := bufio.NewReader(pipe)

	readloop:
		for {
			line, err := reader.ReadString('\n')

			if err != nil {
				break readloop
			}

			log.Print(prefix, " ", line)
		}
	}

	for {
		pipe := <-stdoutChan

		go dumper(pipe, "stdout:")

		pipe = <-stdoutChan

		go dumper(pipe, "stderr:")
	}
}

// Run the command in the given string and restart it after
// a message was received on the buildDone channel.
func runner(command string, buildDone chan bool) {
	var currentProcess *os.Process

	stdoutChan := make(chan io.ReadCloser)

	go logger(stdoutChan)

	for {
		<-buildDone

		args := strings.Split(command, " ")
		cmd := exec.Command(args[0], args[1:]...)

		if currentProcess != nil {
			err := currentProcess.Kill()

			if err != nil {
				log.Fatal("Could not kill child process. Aborting due to danger of infinite forks.")
			}

			// Leave a thread around to read the exit status, this keeps the process from staying
			// as a zombie. It looks like currentProcess.Release() should do this, but it isn't happening for me.
			go func(p *os.Process) {
				// If the process has ignored the KILL, we block here forever. Might want a timeout
				_, err := p.Wait()
				if err != nil {
					log.Fatal("Failed to wait for killed command.")
				}
			}(currentProcess)
		}

		log.Println("Restarting the given command.")

		pipe, err := cmd.StdoutPipe()

		if err != nil {
			log.Fatal("Can't get stdout pipe for command:", err)
		}

		stdoutChan <- pipe

		pipe, err = cmd.StderrPipe()

		if err != nil {
			log.Fatal("Can't get stderr pipe for command:", err)
		}

		stdoutChan <- pipe

		err = cmd.Start()

		if err != nil {
			log.Println("Error while running command:", err)
		}

		currentProcess = cmd.Process
	}
}

func flusher(buildDone <-chan bool) {
	for {
		<-buildDone
	}
}

func main() {
	flag.Parse()

	if *flag_directory == "" {
		fmt.Fprintf(os.Stderr, "-directory=... is required.\n")
		os.Exit(1)
	}

	watcher, err := fsnotify.NewWatcher()

	if err != nil {
		log.Fatal(err)
	}

	defer watcher.Close()

	if *flag_recursive == true {
		err = filepath.Walk(*flag_directory, func(path string, info os.FileInfo, err error) error {
			if err == nil && info.IsDir() {
				return watcher.Watch(path)
			}
			return err
		})

		if err != nil {
			log.Fatal("filepath.Walk():", err)
		}

	} else {
		err := watcher.Watch(*flag_directory)

		if err != nil {
			log.Fatal("watcher.Watch():", err)
		}
	}

	pattern := regexp.MustCompile(*flag_pattern)
	jobs := make(chan string)
	buildDone := make(chan bool)

	go builder(jobs, buildDone)

	if *flag_command != "" {
		go runner(*flag_command, buildDone)
	} else {
		go flusher(buildDone)
	}

	for {
		select {
		case ev := <-watcher.Event:
			if ev.Name != "" && matchesPattern(pattern, ev.Name) {
				jobs <- ev.Name
			}

		case err := <-watcher.Error:
			if v, ok := err.(*os.SyscallError); ok {
				if v.Err == syscall.EINTR {
					continue
				}
				log.Fatal("watcher.Error: SyscallError:", v)
			}
			log.Fatal("watcher.Error:", err)
		}
	}
}
<|endoftext|>


<commit_before>package miniprogram

import (
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"encoding/json"
	"errors"
)

var (
	// ErrAppIDNotMatch appid不匹配
	ErrAppIDNotMatch = errors.New("app id not match")
	// ErrInvalidBlockSize block size不合法
	ErrInvalidBlockSize = errors.New("invalid block size")
	// ErrInvalidPKCS7Data PKCS7数据不合法
	ErrInvalidPKCS7Data = errors.New("invalid PKCS7 data")
	// ErrInvalidPKCS7Padding 输入padding失败
	ErrInvalidPKCS7Padding = errors.New("invalid padding on input")
)

// UserInfo 用户信息
type UserInfo struct {
	OpenID    string `json:"openId"`
	UnionID   string `json:"unionId"`
	NickName  string `json:"nickName"`
	Gender    int    `json:"gender"`
	City      string `json:"city"`
	Province  string `json:"province"`
	Country   string `json:"country"`
	AvatarURL string `json:"avatarUrl"`
	Language  string `json:"language"`
	Watermark struct {
		Timestamp int64  `json:"timestamp"`
		AppID     string `json:"appid"`
	} `json:"watermark"`
}

// pkcs7Unpad returns slice of the original data without padding
func pkcs7Unpad(data []byte, blockSize int) ([]byte, error) {
	if blockSize <= 0 {
		return nil, ErrInvalidBlockSize
	}
	if len(data)%blockSize != 0 || len(data) == 0 {
		return nil, ErrInvalidPKCS7Data
	}
	c := data[len(data)-1]
	n := int(c)
	if n == 0 || n > len(data) {
		return nil, ErrInvalidPKCS7Padding
	}
	for i := 0; i < n; i++ {
		if data[len(data)-n+i] != c {
			return nil, ErrInvalidPKCS7Padding
		}
	}
	return data[:len(data)-n], nil
}

// Decrypt 解密数据
func (wxa *MiniProgram) Decrypt(sessionKey, encryptedData, iv string) (*UserInfo, error) {
	aesKey, err := base64.StdEncoding.DecodeString(sessionKey)
	if err != nil {
		return nil, err
	}
	cipherText, err := base64.StdEncoding.DecodeString(encryptedData)
	if err != nil {
		return nil, err
	}
	ivBytes, err := base64.StdEncoding.DecodeString(iv)
	if err != nil {
		return nil, err
	}
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return nil, err
	}
	mode := cipher.NewCBCDecrypter(block, ivBytes)
	mode.CryptBlocks(cipherText, cipherText)
	cipherText, err = pkcs7Unpad(cipherText, block.BlockSize())
	if err != nil {
		return nil, err
	}
	var userInfo UserInfo
	err = json.Unmarshal(cipherText, &userInfo)
	if err != nil {
		return nil, err
	}
	if userInfo.Watermark.AppID != wxa.AppID {
		return nil, ErrAppIDNotMatch
	}
	return &userInfo, nil
}
<commit_msg>小程序数据解密，增加手机号解密，兼容之前的用户信息解密<commit_after>package miniprogram

import (
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"encoding/json"
	"errors"
)

var (
	// ErrAppIDNotMatch appid不匹配
	ErrAppIDNotMatch = errors.New("app id not match")
	// ErrInvalidBlockSize block size不合法
	ErrInvalidBlockSize = errors.New("invalid block size")
	// ErrInvalidPKCS7Data PKCS7数据不合法
	ErrInvalidPKCS7Data = errors.New("invalid PKCS7 data")
	// ErrInvalidPKCS7Padding 输入padding失败
	ErrInvalidPKCS7Padding = errors.New("invalid padding on input")
)

// UserInfo 用户信息
type UserInfo struct {
	OpenID    string `json:"openId"`
	UnionID   string `json:"unionId"`
	NickName  string `json:"nickName"`
	Gender    int    `json:"gender"`
	City      string `json:"city"`
	Province  string `json:"province"`
	Country   string `json:"country"`
	AvatarURL string `json:"avatarUrl"`
	Language  string `json:"language"`
	Watermark struct {
		Timestamp int64  `json:"timestamp"`
		AppID     string `json:"appid"`
	} `json:"watermark"`
}

// 用户手机号
type PhoneInfo struct {
	PhoneNumber     string `json:"phoneNumber"`
	PurePhoneNumber string `json:"purePhoneNumber"`
	CountryCode     string `json:"countryCode"`
	Watermark       struct {
		Timestamp int64  `json:"timestamp"`
		AppID     string `json:"appid"`
	} `json:"watermark"`
}

// pkcs7Unpad returns slice of the original data without padding
func pkcs7Unpad(data []byte, blockSize int) ([]byte, error) {
	if blockSize <= 0 {
		return nil, ErrInvalidBlockSize
	}
	if len(data)%blockSize != 0 || len(data) == 0 {
		return nil, ErrInvalidPKCS7Data
	}
	c := data[len(data)-1]
	n := int(c)
	if n == 0 || n > len(data) {
		return nil, ErrInvalidPKCS7Padding
	}
	for i := 0; i < n; i++ {
		if data[len(data)-n+i] != c {
			return nil, ErrInvalidPKCS7Padding
		}
	}
	return data[:len(data)-n], nil
}

// get cipherText
func getCipherText(sessionKey, encryptedData, iv string) ([]byte, error) {
	aesKey, err := base64.StdEncoding.DecodeString(sessionKey)
	if err != nil {
		return nil, err
	}
	cipherText, err := base64.StdEncoding.DecodeString(encryptedData)
	if err != nil {
		return nil, err
	}
	ivBytes, err := base64.StdEncoding.DecodeString(iv)
	if err != nil {
		return nil, err
	}
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return nil, err
	}
	mode := cipher.NewCBCDecrypter(block, ivBytes)
	mode.CryptBlocks(cipherText, cipherText)
	cipherText, err = pkcs7Unpad(cipherText, block.BlockSize())
	if err != nil {
		return nil, err
	}
	return cipherText, nil
}

// Decrypt 解密数据(用户)
func (wxa *MiniProgram) Decrypt(sessionKey, encryptedData, iv string) (*UserInfo, error) {
	cipherText, err := getCipherText(sessionKey, encryptedData, iv)
	if err != nil {
		return nil, err
	}
	var userInfo UserInfo
	err = json.Unmarshal(cipherText, &userInfo)
	if err != nil {
		return nil, err
	}
	if userInfo.Watermark.AppID != wxa.AppID {
		return nil, ErrAppIDNotMatch
	}
	return &userInfo, nil
}

// DecryptPhone 解密数据(手机)
func (wxa *MiniProgram) DecryptPHone(sessionKey, encryptedData, iv string) (*PhoneInfo, error) {
	cipherText, err := getCipherText(sessionKey, encryptedData, iv)
	if err != nil {
		return nil, err
	}
	var phoneInfo PhoneInfo
	err = json.Unmarshal(cipherText, &phoneInfo)
	if err != nil {
		return nil, err
	}
	if phoneInfo.Watermark.AppID != wxa.AppID {
		return nil, ErrAppIDNotMatch
	}
	return &phoneInfo, nil
}
<|endoftext|>


<commit_before>/* Document management and index maintenance. */
package db

import (
	"encoding/json"
	"fmt"
	"github.com/HouzuoGuo/tiedot/tdlog"
	"math/rand"
)

// Resolve the attribute(s) in the document structure along the given path.
func GetIn(doc interface{}, path []string) (ret []interface{}) {
	docMap, ok := doc.(map[string]interface{})
	if !ok {
		return
	}
	var thing interface{} = docMap
	// Get into each path segment
	for i, seg := range path {
		if aMap, ok := thing.(map[string]interface{}); ok {
			thing = aMap[seg]
		} else if anArray, ok := thing.([]interface{}); ok {
			for _, element := range anArray {
				ret = append(ret, GetIn(element, path[i:])...)
			}
			return ret
		} else {
			return nil
		}
	}
	switch thing.(type) {
	case []interface{}:
		return append(ret, thing.([]interface{})...)
	default:
		return append(ret, thing)
	}
}

// Hash a string using sdbm algorithm.
func StrHash(thing interface{}) int {
	var hash int
	for _, c := range fmt.Sprint(thing) {
		hash = int(c) + (hash << 6) + (hash << 16) - hash
	}
	if hash < 0 {
		return -hash
	}
	return hash
}

// Put a document on all user-created indexes.
func (col *Col) indexDoc(id int, doc map[string]interface{}) {
	for idxName, idxPath := range col.indexPaths {
		for _, idxVal := range GetIn(doc, idxPath) {
			if idxVal != nil {
				hashKey := StrHash(fmt.Sprint(idxVal))
				partNum := hashKey % col.db.numParts
				ht := col.hts[partNum][idxName]
				ht.Lock.Lock()
				ht.Put(hashKey, id)
				ht.Lock.Unlock()
			}
		}
	}
}

// Remove a document from all user-created indexes.
func (col *Col) unindexDoc(id int, doc map[string]interface{}) {
	for idxName, idxPath := range col.indexPaths {
		for _, idxVal := range GetIn(doc, idxPath) {
			if idxVal != nil {
				hashKey := StrHash(fmt.Sprint(idxVal))
				partNum := hashKey % col.db.numParts
				ht := col.hts[partNum][idxName]
				ht.Lock.Lock()
				ht.Remove(hashKey, id)
				ht.Lock.Unlock()
			}
		}
	}
}

// Insert a document with the specified ID into the collection (incl. index). Does not place partition/schema lock.
func (col *Col) InsertRecovery(id int, doc map[string]interface{}) (err error) {
	docJS, err := json.Marshal(doc)
	if err != nil {
		return
	}
	partNum := id % col.db.numParts
	part := col.parts[partNum]
	// Put document data into collection
	if _, err = part.Insert(id, []byte(docJS)); err != nil {
		return
	}
	// Index the document
	col.indexDoc(id, doc)
	return
}

// Insert a document into the collection.
func (col *Col) Insert(doc map[string]interface{}) (id int, err error) {
	docJS, err := json.Marshal(doc)
	if err != nil {
		return
	}
	id = rand.Int()
	partNum := id % col.db.numParts
	col.db.schemaLock.RLock()
	part := col.parts[partNum]
	// Put document data into collection
	part.Lock.Lock()
	if _, err = part.Insert(id, []byte(docJS)); err != nil {
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return
	}
	// If another thread is updating the document in the meanwhile, let it take over index maintenance
	if err = part.LockUpdate(id); err != nil {
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return id, nil
	}
	part.Lock.Unlock()
	// Index the document
	col.indexDoc(id, doc)
	part.Lock.Lock()
	part.UnlockUpdate(id)
	part.Lock.Unlock()
	col.db.schemaLock.RUnlock()
	return
}

func (col *Col) read(id int, placeSchemaLock bool) (doc map[string]interface{}, err error) {
	if placeSchemaLock {
		col.db.schemaLock.RLock()
	}
	part := col.parts[id%col.db.numParts]
	part.Lock.RLock()
	docB, err := part.Read(id)
	part.Lock.RUnlock()
	if err != nil {
		if placeSchemaLock {
			col.db.schemaLock.RUnlock()
		}
		return
	}
	err = json.Unmarshal(docB, &doc)
	if placeSchemaLock {
		col.db.schemaLock.RUnlock()
	}
	return
}

// Find and retrieve a document by ID.
func (col *Col) Read(id int) (doc map[string]interface{}, err error) {
	return col.read(id, true)
}

// Update a document.
func (col *Col) Update(id int, doc map[string]interface{}) error {
	if doc == nil {
		return fmt.Errorf("Updating %d: input doc may not be nil", id)
	}
	docJS, err := json.Marshal(doc)
	if err != nil {
		return err
	}
	col.db.schemaLock.RLock()
	part := col.parts[id%col.db.numParts]
	part.Lock.Lock()
	// Place lock, read back original document and update
	if err := part.LockUpdate(id); err != nil {
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return fmt.Errorf("Document %d is locked for update, please try again later", id)
	}
	originalB, err := part.Read(id)
	if err != nil {
		part.UnlockUpdate(id)
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return fmt.Errorf("Cannot update %d: cannot read back original document - %v", id, err)
	}
	var original map[string]interface{}
	if err = json.Unmarshal(originalB, &original); err != nil {
		tdlog.Noticef("Will not attempt to unindex document %d during update", id)
	}
	if err = part.Update(id, []byte(docJS)); err != nil {
		part.UnlockUpdate(id)
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return err
	}
	// Done with the collection data, next is to maintain indexed values
	part.Lock.Unlock()
	if original != nil {
		col.unindexDoc(id, original)
	}
	col.indexDoc(id, doc)
	// Done with the document
	part.Lock.Lock()
	part.UnlockUpdate(id)
	part.Lock.Unlock()
	col.db.schemaLock.RUnlock()
	return nil
}

// Delete a document.
func (col *Col) Delete(id int) error {
	col.db.schemaLock.RLock()
	part := col.parts[id%col.db.numParts]
	part.Lock.Lock()
	// Place lock, read back original document and delete document
	if err := part.LockUpdate(id); err != nil {
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return fmt.Errorf("Document %d is locked for update, please try again later", id)
	}
	originalB, err := part.Read(id)
	if err != nil {
		part.UnlockUpdate(id)
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return fmt.Errorf("Cannot delete %d: %v", id, err)
	}
	var original map[string]interface{}
	if err = json.Unmarshal(originalB, &original); err != nil {
		tdlog.Noticef("Will not attempt to unindex document %d during delete", id)
	}
	if err = part.Delete(id); err != nil {
		part.UnlockUpdate(id)
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return err
	}
	// Done with the collection data, next is to remove indexed values
	part.Lock.Unlock()
	if original != nil {
		col.unindexDoc(id, original)
	}
	part.Lock.Lock()
	part.UnlockUpdate(id)
	part.Lock.Unlock()
	col.db.schemaLock.RUnlock()
	return nil
}
<commit_msg>remove unnecessary fmt.Sprint - a nice little performance bump<commit_after>/* Document management and index maintenance. */
package db

import (
	"encoding/json"
	"fmt"
	"github.com/HouzuoGuo/tiedot/tdlog"
	"math/rand"
)

// Resolve the attribute(s) in the document structure along the given path.
func GetIn(doc interface{}, path []string) (ret []interface{}) {
	docMap, ok := doc.(map[string]interface{})
	if !ok {
		return
	}
	var thing interface{} = docMap
	// Get into each path segment
	for i, seg := range path {
		if aMap, ok := thing.(map[string]interface{}); ok {
			thing = aMap[seg]
		} else if anArray, ok := thing.([]interface{}); ok {
			for _, element := range anArray {
				ret = append(ret, GetIn(element, path[i:])...)
			}
			return ret
		} else {
			return nil
		}
	}
	switch thing.(type) {
	case []interface{}:
		return append(ret, thing.([]interface{})...)
	default:
		return append(ret, thing)
	}
}

// Hash a string using sdbm algorithm.
func StrHash(str string) int {
	var hash int
	for _, c := range str {
		hash = int(c) + (hash << 6) + (hash << 16) - hash
	}
	if hash < 0 {
		return -hash
	}
	return hash
}

// Put a document on all user-created indexes.
func (col *Col) indexDoc(id int, doc map[string]interface{}) {
	for idxName, idxPath := range col.indexPaths {
		for _, idxVal := range GetIn(doc, idxPath) {
			if idxVal != nil {
				hashKey := StrHash(fmt.Sprint(idxVal))
				partNum := hashKey % col.db.numParts
				ht := col.hts[partNum][idxName]
				ht.Lock.Lock()
				ht.Put(hashKey, id)
				ht.Lock.Unlock()
			}
		}
	}
}

// Remove a document from all user-created indexes.
func (col *Col) unindexDoc(id int, doc map[string]interface{}) {
	for idxName, idxPath := range col.indexPaths {
		for _, idxVal := range GetIn(doc, idxPath) {
			if idxVal != nil {
				hashKey := StrHash(fmt.Sprint(idxVal))
				partNum := hashKey % col.db.numParts
				ht := col.hts[partNum][idxName]
				ht.Lock.Lock()
				ht.Remove(hashKey, id)
				ht.Lock.Unlock()
			}
		}
	}
}

// Insert a document with the specified ID into the collection (incl. index). Does not place partition/schema lock.
func (col *Col) InsertRecovery(id int, doc map[string]interface{}) (err error) {
	docJS, err := json.Marshal(doc)
	if err != nil {
		return
	}
	partNum := id % col.db.numParts
	part := col.parts[partNum]
	// Put document data into collection
	if _, err = part.Insert(id, []byte(docJS)); err != nil {
		return
	}
	// Index the document
	col.indexDoc(id, doc)
	return
}

// Insert a document into the collection.
func (col *Col) Insert(doc map[string]interface{}) (id int, err error) {
	docJS, err := json.Marshal(doc)
	if err != nil {
		return
	}
	id = rand.Int()
	partNum := id % col.db.numParts
	col.db.schemaLock.RLock()
	part := col.parts[partNum]
	// Put document data into collection
	part.Lock.Lock()
	if _, err = part.Insert(id, []byte(docJS)); err != nil {
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return
	}
	// If another thread is updating the document in the meanwhile, let it take over index maintenance
	if err = part.LockUpdate(id); err != nil {
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return id, nil
	}
	part.Lock.Unlock()
	// Index the document
	col.indexDoc(id, doc)
	part.Lock.Lock()
	part.UnlockUpdate(id)
	part.Lock.Unlock()
	col.db.schemaLock.RUnlock()
	return
}

func (col *Col) read(id int, placeSchemaLock bool) (doc map[string]interface{}, err error) {
	if placeSchemaLock {
		col.db.schemaLock.RLock()
	}
	part := col.parts[id%col.db.numParts]
	part.Lock.RLock()
	docB, err := part.Read(id)
	part.Lock.RUnlock()
	if err != nil {
		if placeSchemaLock {
			col.db.schemaLock.RUnlock()
		}
		return
	}
	err = json.Unmarshal(docB, &doc)
	if placeSchemaLock {
		col.db.schemaLock.RUnlock()
	}
	return
}

// Find and retrieve a document by ID.
func (col *Col) Read(id int) (doc map[string]interface{}, err error) {
	return col.read(id, true)
}

// Update a document.
func (col *Col) Update(id int, doc map[string]interface{}) error {
	if doc == nil {
		return fmt.Errorf("Updating %d: input doc may not be nil", id)
	}
	docJS, err := json.Marshal(doc)
	if err != nil {
		return err
	}
	col.db.schemaLock.RLock()
	part := col.parts[id%col.db.numParts]
	part.Lock.Lock()
	// Place lock, read back original document and update
	if err := part.LockUpdate(id); err != nil {
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return fmt.Errorf("Document %d is locked for update, please try again later", id)
	}
	originalB, err := part.Read(id)
	if err != nil {
		part.UnlockUpdate(id)
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return fmt.Errorf("Cannot update %d: cannot read back original document - %v", id, err)
	}
	var original map[string]interface{}
	if err = json.Unmarshal(originalB, &original); err != nil {
		tdlog.Noticef("Will not attempt to unindex document %d during update", id)
	}
	if err = part.Update(id, []byte(docJS)); err != nil {
		part.UnlockUpdate(id)
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return err
	}
	// Done with the collection data, next is to maintain indexed values
	part.Lock.Unlock()
	if original != nil {
		col.unindexDoc(id, original)
	}
	col.indexDoc(id, doc)
	// Done with the document
	part.Lock.Lock()
	part.UnlockUpdate(id)
	part.Lock.Unlock()
	col.db.schemaLock.RUnlock()
	return nil
}

// Delete a document.
func (col *Col) Delete(id int) error {
	col.db.schemaLock.RLock()
	part := col.parts[id%col.db.numParts]
	part.Lock.Lock()
	// Place lock, read back original document and delete document
	if err := part.LockUpdate(id); err != nil {
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return fmt.Errorf("Document %d is locked for update, please try again later", id)
	}
	originalB, err := part.Read(id)
	if err != nil {
		part.UnlockUpdate(id)
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return fmt.Errorf("Cannot delete %d: %v", id, err)
	}
	var original map[string]interface{}
	if err = json.Unmarshal(originalB, &original); err != nil {
		tdlog.Noticef("Will not attempt to unindex document %d during delete", id)
	}
	if err = part.Delete(id); err != nil {
		part.UnlockUpdate(id)
		part.Lock.Unlock()
		col.db.schemaLock.RUnlock()
		return err
	}
	// Done with the collection data, next is to remove indexed values
	part.Lock.Unlock()
	if original != nil {
		col.unindexDoc(id, original)
	}
	part.Lock.Lock()
	part.UnlockUpdate(id)
	part.Lock.Unlock()
	col.db.schemaLock.RUnlock()
	return nil
}
<|endoftext|>


<commit_before>package db

import (
	"database/sql"
	"github.com/Cepave/hbs/model"
	commonModel "github.com/Cepave/common/model"
	"log"
	"time"
)

// Inserts a new agent or updates existing one
func RefreshAgentInfo(agent *model.NqmAgent) (err error) {
	/**
	 * Update the data
	 */
	if _, err = DB.Exec(
		`
		INSERT INTO nqm_agent(ag_connection_id, ag_hostname, ag_ip_address)
		VALUES(?, ?, ?)
		ON DUPLICATE KEY UPDATE
			ag_hostname = VALUES(ag_hostname),
			ag_ip_address = VALUES(ag_ip_address)
		`,
		agent.ConnectionId(),
		agent.Hostname(),
		[]byte(agent.IpAddress),
	)
		err != nil {

		log.Printf("Cannot refresh agent: [%v]. Error: %v", agent.ConnectionId(), err)
		return
	}
	// :~)

	/**
	 * Loads id from auto-generated PK
	 */
	if err = DB.QueryRow(
		`
		SELECT ag_id
		FROM nqm_agent
		WHERE ag_connection_id = ?
		`,
		agent.ConnectionId(),
	).Scan(&agent.Id)
		err != nil {
		return
	}
	// :~)

	return
}

// Gets the data of agent for RPC
//
// If there is no need to perform ping task, this method would return nil as result.
//
// Reasons for not doing ping task:
// 1) No ping task configuration
// 2) The period is overed yet
func GetAndRefreshNeedPingAgentForRpc(agentId int, checkedTime time.Time) (result *commonModel.NqmAgent, err error) {
	result = &commonModel.NqmAgent {
		Id: agentId,
	}

	err = inTx(func(tx *sql.Tx) (err error) {
		var dbAgentName sql.NullString

		if err = tx.QueryRow(
			/**
			 * Gets one row if the executing of ping task is needed
			 *
			 * Gets no row if:
			 * 1) No ping task configuration
			 * 2) The period is overed yet
			 * 3) The agent is disabled
			 */
			`
			SELECT ag_name,
				isp_id, isp_name,
				pv_id, pv_name,
				ct_id, ct_name
			FROM nqm_agent AS ag
				INNER JOIN
				(
					SELECT pt_ag_id
					FROM nqm_ping_task
					WHERE pt_ag_id = ?
						/* Check if the period is overed since last time executed */
						AND TIMESTAMPDIFF(
							MINUTE,
							IFNULL(pt_time_last_execute, FROM_UNIXTIME(0)), /* Use the very first time */
							FROM_UNIXTIME(?)
						) >= pt_period
						/* :~) */
				) AS pt
				ON ag.ag_id = pt.pt_ag_id
					AND ag.ag_status & b'00000001' = b'00000001'
				INNER JOIN
				owl_isp AS isp
				ON ag.ag_isp_id = isp.isp_id
				INNER JOIN
				owl_province AS pv
				ON ag.ag_pv_id = pv.pv_id
				INNER JOIN
				owl_city AS ct
				ON ag.ag_ct_id = ct.ct_id
			`,
			// :~)
			agentId, checkedTime.Unix(),
		).Scan(
			&dbAgentName,
			&result.IspId, &result.IspName,
			&result.ProvinceId, &result.ProvinceName,
			&result.CityId, &result.CityName,
		)
			err != nil {

			/**
			 * There is no need to perform ping task
			 */
			result = nil
			if err == sql.ErrNoRows {
				err = nil
			}
			// :~)

			return
		}

		/**
		 * Loads name of agent
		 */
		if dbAgentName.Valid {
			result.Name = dbAgentName.String
		} else {
			result.Name = commonModel.UNDEFINED_STRING
		}
		// :~)

		/**
		 * Updates the last execute
		 */
		if _, err = DB.Exec(
			`
			UPDATE nqm_ping_task
			SET pt_time_last_execute = FROM_UNIXTIME(?)
			WHERE pt_ag_id = ?
			`,
			checkedTime.Unix(), agentId,
		)
			err != nil {

			result = nil
		}
		// :~)

		return
	})

	return
}

const (
	NO_PING_TASK = 0
	HAS_PING_TASK_WITH_FILTER = 1
	HAS_PING_TASK_ALL_MATCHING = 2
)

// Gets the targets(to be probed) for RPC
func GetTargetsByAgentForRpc(agentId int) (targets []commonModel.NqmTarget, err error) {
	var taskState int

	if taskState, err = getPingTaskState(agentId)
		err != nil {
		return
	}

	var rows *sql.Rows

	switch taskState {
	case NO_PING_TASK:
		targets = make([]commonModel.NqmTarget, 0)
		return
	case HAS_PING_TASK_WITH_FILTER:
		rows, err = loadTargetsByFilter(agentId)
	case HAS_PING_TASK_ALL_MATCHING:
		rows, err = loadAllTargets()
	}

	if err != nil { return }

	/**
	 * Converts the data to NQM targets for RPC
	 */
	defer rows.Close()
	targets = make([]commonModel.NqmTarget, 0, 4)
	for rows.Next() {
		targets = append(targets, commonModel.NqmTarget{})
		var currentTarget *commonModel.NqmTarget = &targets[len(targets) - 1]
		currentTarget.NameTag = commonModel.UNDEFINED_STRING

		var nameTag sql.NullString

		rows.Scan(
			&currentTarget.Id,
			&currentTarget.Host,
			&currentTarget.IspId,
			&currentTarget.IspName,
			&currentTarget.ProvinceId,
			&currentTarget.ProvinceName,
			&currentTarget.CityId,
			&currentTarget.CityName,
			&nameTag,
		)

		if nameTag.Valid {
			currentTarget.NameTag = nameTag.String
		}
	}
	// :~)

	return
}

func loadAllTargets() (*sql.Rows, error) {
	return DB.Query(
		`
		SELECT tg_id, tg_host,
			isp_id, isp_name,
			pv_id, pv_name,
			ct_id, ct_name, tg.tg_name_tag
		FROM nqm_target AS tg
			INNER JOIN
			owl_isp AS isp
			ON tg.tg_isp_id = isp.isp_id
			INNER JOIN
			owl_province AS pv
			ON tg.tg_pv_id = pv.pv_id
			INNER JOIN
			owl_city AS ct
			ON tg.tg_ct_id = ct.ct_id
		`,
	)
}
func loadTargetsByFilter(agentId int) (*sql.Rows, error) {
	return DB.Query(
		`
		SELECT
			tg_id, tg_host,
			isp_id, isp_name,
			pv_id, pv_name,
			ct_id, ct_name, tg.tg_name_tag
		FROM (
			/* Matched target by ISP */
			SELECT tg_id, tg_host, tg_isp_id, tg_pv_id, tg_ct_id, tg_name_tag
			FROM nqm_target tg
				INNER JOIN
				nqm_pt_target_filter_isp AS tfisp
				ON tfisp.tfisp_pt_ag_id = ?
					AND tg.tg_isp_id = tfisp.tfisp_isp_id
			/* :~) */
			UNION
			/* Matched target by province */
			SELECT tg_id, tg_host, tg_isp_id, tg_pv_id, tg_ct_id, tg_name_tag
			FROM nqm_target tg
				INNER JOIN
				nqm_pt_target_filter_province AS tfpv
				ON tfpv.tfpv_pt_ag_id = ?
					AND tg.tg_pv_id = tfpv.tfpv_pv_id
			/* :~) */
			UNION
			/* Matched target by city */
			SELECT tg_id, tg_host, tg_isp_id, tg_pv_id, tg_ct_id, tg_name_tag
			FROM nqm_target tg
				INNER JOIN
				nqm_pt_target_filter_city AS tfct
				ON tfct.tfct_pt_ag_id = ?
					AND tg.tg_ct_id = tfct.tfct_ct_id
			/* :~) */
			UNION
			/* Matched target by name tag */
			SELECT tg_id, tg_host, tg_isp_id, tg_pv_id, tg_ct_id, tg_name_tag
			FROM nqm_target tg
				INNER JOIN
				nqm_pt_target_filter_name_tag AS tfnt
				ON tfnt.tfnt_pt_ag_id = ?
					AND tg.tg_name_tag = tfnt.tfnt_name_tag
			/* :~) */
			UNION
			/* Matched target which to be probed by all */
			SELECT tg_id, tg_host, tg_isp_id, tg_pv_id, tg_ct_id, tg_name_tag
			FROM nqm_target tg
			WHERE tg_probed_by_all = true
			/* :~) */
		) AS tg
		INNER JOIN
		owl_isp AS isp
		ON tg.tg_isp_id = isp.isp_id
		INNER JOIN
		owl_province AS pv
		ON tg.tg_pv_id = pv.pv_id
		INNER JOIN
		owl_city AS ct
		ON tg.tg_ct_id = ct.ct_id
		`,
		agentId, agentId, agentId, agentId,
	)
}

func getPingTaskState(agentId int) (result int, err error) {
	result = NO_PING_TASK

	if err = DB.QueryRow(
		`
		SELECT COUNT(pt_ag_id)
		FROM nqm_ping_task
		WHERE pt_ag_id = ?
		`,
		agentId,
	).Scan(&result)
		err != nil {
		return
	}

	/**
	 * Without ping task
	 */
	if result == NO_PING_TASK {
		return
	}
	// :~)

	/**
	 * Checks whether there is any filter configuration for the ping task
	 */
	var hasFilter int
	if err = DB.QueryRow(
		`
		SELECT COUNT(*)
		FROM (
			/* Get any condition of name tag filter  */
			SELECT tfnt.tfnt_pt_ag_id
			FROM nqm_pt_target_filter_name_tag AS tfnt
			WHERE tfnt.tfnt_pt_ag_id = ?
			LIMIT 1
			/* :~) */
			UNION
			/* Get any condition of ISP filter  */
			SELECT tfisp.tfisp_pt_ag_id
			FROM nqm_pt_target_filter_isp AS tfisp
			WHERE tfisp.tfisp_pt_ag_id = ?
			LIMIT 1
			/* :~) */
			UNION
			/* Get any condition of province filter  */
			SELECT tfpv.tfpv_pt_ag_id
			FROM nqm_pt_target_filter_province AS tfpv
			WHERE tfpv.tfpv_pt_ag_id = ?
			LIMIT 1
			/* :~) */
			UNION
			/* Get any condition of city filter  */
			SELECT tfct.tfct_pt_ag_id
			FROM nqm_pt_target_filter_city AS tfct
			WHERE tfct.tfct_pt_ag_id = ?
			LIMIT 1
			/* :~) */
		) AS filter
		`,
		agentId, agentId, agentId, agentId,
	).Scan(&hasFilter)
		err != nil {
		return
	}
	// :~)

	if hasFilter == 1 {
		result = HAS_PING_TASK_WITH_FILTER
	} else {
		result = HAS_PING_TASK_ALL_MATCHING
	}

	return
}
<commit_msg>[OWL-451] Modify some SQL queries<commit_after>package db

import (
	"database/sql"
	"github.com/Cepave/hbs/model"
	commonModel "github.com/Cepave/common/model"
	"log"
	"time"
)

// Inserts a new agent or updates existing one
func RefreshAgentInfo(agent *model.NqmAgent) (err error) {
	/**
	 * Update the data
	 */
	if _, err = DB.Exec(
		`
		INSERT INTO nqm_agent(ag_connection_id, ag_hostname, ag_ip_address, ag_last_heartbeat)
		VALUES(?, ?, ?, ?)
		ON DUPLICATE KEY UPDATE
			ag_hostname = VALUES(ag_hostname),
			ag_ip_address = VALUES(ag_ip_address),
			ag_last_heartbeat = VALUES(ag_last_heartbeat)
		`,
		agent.ConnectionId(),
		agent.Hostname(),
		[]byte(agent.IpAddress),
		time.Now(),
	)
		err != nil {

		log.Printf("Cannot refresh agent: [%v]. Error: %v", agent.ConnectionId(), err)
		return
	}
	// :~)

	/**
	 * Loads id from auto-generated PK
	 */
	if err = DB.QueryRow(
		`
		SELECT ag_id
		FROM nqm_agent
		WHERE ag_connection_id = ?
		`,
		agent.ConnectionId(),
	).Scan(&agent.Id)
		err != nil {
		return
	}
	// :~)

	return
}

// Gets the data of agent for RPC
//
// If there is no need to perform ping task, this method would return nil as result.
//
// Reasons for not doing ping task:
// 1) No ping task configuration
// 2) The period is overed yet
func GetAndRefreshNeedPingAgentForRpc(agentId int, checkedTime time.Time) (result *commonModel.NqmAgent, err error) {
	result = &commonModel.NqmAgent {
		Id: agentId,
	}

	err = inTx(func(tx *sql.Tx) (err error) {
		var dbAgentName sql.NullString

		if err = tx.QueryRow(
			/**
			 * Gets one row if the executing of ping task is needed
			 *
			 * Gets no row if:
			 * 1) No ping task configuration
			 * 2) The period is overed yet
			 * 3) The agent is disabled
			 */
			`
			SELECT ag_name,
				isp_id, isp_name,
				pv_id, pv_name,
				ct_id, ct_name
			FROM nqm_agent AS ag
				INNER JOIN
				(
					SELECT pt_ag_id
					FROM nqm_ping_task
					WHERE pt_ag_id = ?
						/* Check if the period is overed since last time executed */
						AND TIMESTAMPDIFF(
							MINUTE,
							IFNULL(pt_time_last_execute, FROM_UNIXTIME(0)), /* Use the very first time */
							FROM_UNIXTIME(?)
						) >= pt_period
						/* :~) */
				) AS pt
				ON ag.ag_id = pt.pt_ag_id
					AND ag.ag_status & b'00000001' = b'00000001'
				INNER JOIN
				owl_isp AS isp
				ON ag.ag_isp_id = isp.isp_id
				INNER JOIN
				owl_province AS pv
				ON ag.ag_pv_id = pv.pv_id
				INNER JOIN
				owl_city AS ct
				ON ag.ag_ct_id = ct.ct_id
			`,
			// :~)
			agentId, checkedTime.Unix(),
		).Scan(
			&dbAgentName,
			&result.IspId, &result.IspName,
			&result.ProvinceId, &result.ProvinceName,
			&result.CityId, &result.CityName,
		)
			err != nil {

			/**
			 * There is no need to perform ping task
			 */
			result = nil
			if err == sql.ErrNoRows {
				err = nil
			}
			// :~)

			return
		}

		/**
		 * Loads name of agent
		 */
		if dbAgentName.Valid {
			result.Name = dbAgentName.String
		} else {
			result.Name = commonModel.UNDEFINED_STRING
		}
		// :~)

		/**
		 * Updates the last execute
		 */
		if _, err = DB.Exec(
			`
			UPDATE nqm_ping_task
			SET pt_time_last_execute = FROM_UNIXTIME(?)
			WHERE pt_ag_id = ?
			`,
			checkedTime.Unix(), agentId,
		)
			err != nil {

			result = nil
		}
		// :~)

		return
	})

	return
}

const (
	NO_PING_TASK = 0
	HAS_PING_TASK_WITH_FILTER = 1
	HAS_PING_TASK_ALL_MATCHING = 2
)

// Gets the targets(to be probed) for RPC
func GetTargetsByAgentForRpc(agentId int) (targets []commonModel.NqmTarget, err error) {
	var taskState int

	if taskState, err = getPingTaskState(agentId)
		err != nil {
		return
	}

	var rows *sql.Rows

	switch taskState {
	case NO_PING_TASK:
		targets = make([]commonModel.NqmTarget, 0)
		return
	case HAS_PING_TASK_WITH_FILTER:
		rows, err = loadTargetsByFilter(agentId)
	case HAS_PING_TASK_ALL_MATCHING:
		rows, err = loadAllTargets()
	}

	if err != nil { return }

	/**
	 * Converts the data to NQM targets for RPC
	 */
	defer rows.Close()
	targets = make([]commonModel.NqmTarget, 0, 4)
	for rows.Next() {
		targets = append(targets, commonModel.NqmTarget{})
		var currentTarget *commonModel.NqmTarget = &targets[len(targets) - 1]
		currentTarget.NameTag = commonModel.UNDEFINED_STRING

		var nameTag sql.NullString

		rows.Scan(
			&currentTarget.Id,
			&currentTarget.Host,
			&currentTarget.IspId,
			&currentTarget.IspName,
			&currentTarget.ProvinceId,
			&currentTarget.ProvinceName,
			&currentTarget.CityId,
			&currentTarget.CityName,
			&nameTag,
		)

		if nameTag.Valid {
			currentTarget.NameTag = nameTag.String
		}
	}
	// :~)

	return
}

func loadAllTargets() (*sql.Rows, error) {
	return DB.Query(
		`
		SELECT tg_id, tg_host,
			isp_id, isp_name,
			pv_id, pv_name,
			ct_id, ct_name, tg.tg_name_tag
		FROM nqm_target AS tg
			INNER JOIN
			owl_isp AS isp
			ON tg.tg_isp_id = isp.isp_id
			INNER JOIN
			owl_province AS pv
			ON tg.tg_pv_id = pv.pv_id
			INNER JOIN
			owl_city AS ct
			ON tg.tg_ct_id = ct.ct_id
		WHERE
			tg.tg_status=true
			AND
			tg.tg_available=true;
		`,
	)
}
func loadTargetsByFilter(agentId int) (*sql.Rows, error) {
	return DB.Query(
		`
		SELECT
			tg_id, tg_host,
			isp_id, isp_name,
			pv_id, pv_name,
			ct_id, ct_name, tg.tg_name_tag
		FROM (
			/* Matched target by ISP */
			SELECT tg_id, tg_host, tg_isp_id, tg_pv_id, tg_ct_id, tg_name_tag
			FROM nqm_target tg
				INNER JOIN
				nqm_pt_target_filter_isp AS tfisp
				ON tfisp.tfisp_pt_ag_id = ?
					AND tg.tg_isp_id = tfisp.tfisp_isp_id
					AND tg.tg_status = true
					AND tg.tg_available = true
			/* :~) */
			UNION
			/* Matched target by province */
			SELECT tg_id, tg_host, tg_isp_id, tg_pv_id, tg_ct_id, tg_name_tag
			FROM nqm_target tg
				INNER JOIN
				nqm_pt_target_filter_province AS tfpv
				ON tfpv.tfpv_pt_ag_id = ?
					AND tg.tg_pv_id = tfpv.tfpv_pv_id
					AND tg.tg_status = true
					AND tg.tg_available = true
			/* :~) */
			UNION
			/* Matched target by city */
			SELECT tg_id, tg_host, tg_isp_id, tg_pv_id, tg_ct_id, tg_name_tag
			FROM nqm_target tg
				INNER JOIN
				nqm_pt_target_filter_city AS tfct
				ON tfct.tfct_pt_ag_id = ?
					AND tg.tg_ct_id = tfct.tfct_ct_id
					AND tg.tg_status = true
					AND tg.tg_available = true
			/* :~) */
			UNION
			/* Matched target by name tag */
			SELECT tg_id, tg_host, tg_isp_id, tg_pv_id, tg_ct_id, tg_name_tag
			FROM nqm_target tg
				INNER JOIN
				nqm_pt_target_filter_name_tag AS tfnt
				ON tfnt.tfnt_pt_ag_id = ?
					AND tg.tg_name_tag = tfnt.tfnt_name_tag
					AND tg.tg_status = true
					AND tg.tg_available = true
			/* :~) */
			UNION
			/* Matched target which to be probed by all */
			SELECT tg_id, tg_host, tg_isp_id, tg_pv_id, tg_ct_id, tg_name_tag
			FROM nqm_target tg
			WHERE tg_probed_by_all = true
				AND tg.tg_status = true
				AND tg.tg_available = true
			/* :~) */
		) AS tg
		INNER JOIN
		owl_isp AS isp
		ON tg.tg_isp_id = isp.isp_id
		INNER JOIN
		owl_province AS pv
		ON tg.tg_pv_id = pv.pv_id
		INNER JOIN
		owl_city AS ct
		ON tg.tg_ct_id = ct.ct_id
		`,
		agentId, agentId, agentId, agentId,
	)
}

func getPingTaskState(agentId int) (result int, err error) {
	result = NO_PING_TASK

	if err = DB.QueryRow(
		`
		SELECT COUNT(pt_ag_id)
		FROM nqm_ping_task
		WHERE pt_ag_id = ?
		`,
		agentId,
	).Scan(&result)
		err != nil {
		return
	}

	/**
	 * Without ping task
	 */
	if result == NO_PING_TASK {
		return
	}
	// :~)

	/**
	 * Checks whether there is any filter configuration for the ping task
	 */
	var hasFilter int
	if err = DB.QueryRow(
		`
		SELECT COUNT(*)
		FROM (
			/* Get any condition of name tag filter  */
			SELECT tfnt.tfnt_pt_ag_id
			FROM nqm_pt_target_filter_name_tag AS tfnt
			WHERE tfnt.tfnt_pt_ag_id = ?
			LIMIT 1
			/* :~) */
			UNION
			/* Get any condition of ISP filter  */
			SELECT tfisp.tfisp_pt_ag_id
			FROM nqm_pt_target_filter_isp AS tfisp
			WHERE tfisp.tfisp_pt_ag_id = ?
			LIMIT 1
			/* :~) */
			UNION
			/* Get any condition of province filter  */
			SELECT tfpv.tfpv_pt_ag_id
			FROM nqm_pt_target_filter_province AS tfpv
			WHERE tfpv.tfpv_pt_ag_id = ?
			LIMIT 1
			/* :~) */
			UNION
			/* Get any condition of city filter  */
			SELECT tfct.tfct_pt_ag_id
			FROM nqm_pt_target_filter_city AS tfct
			WHERE tfct.tfct_pt_ag_id = ?
			LIMIT 1
			/* :~) */
		) AS filter
		`,
		agentId, agentId, agentId, agentId,
	).Scan(&hasFilter)
		err != nil {
		return
	}
	// :~)

	if hasFilter == 1 {
		result = HAS_PING_TASK_WITH_FILTER
	} else {
		result = HAS_PING_TASK_ALL_MATCHING
	}

	return
}
<|endoftext|>


<commit_before>// +build !windows

package oci

import (
	"context"
	"path"
	"sync"

	"github.com/containerd/containerd/containers"
	"github.com/containerd/containerd/mount"
	"github.com/containerd/containerd/namespaces"
	"github.com/containerd/containerd/oci"
	"github.com/containerd/continuity/fs"
	"github.com/mitchellh/hashstructure"
	"github.com/moby/buildkit/executor"
	"github.com/moby/buildkit/snapshot"
	"github.com/moby/buildkit/util/network"
	specs "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
)

// Ideally we don't have to import whole containerd just for the default spec

// GenerateSpec generates spec using containerd functionality.
func GenerateSpec(ctx context.Context, meta executor.Meta, mounts []executor.Mount, id, resolvConf, hostsFile string, namespace network.Namespace, opts ...oci.SpecOpts) (*specs.Spec, func(), error) {
	c := &containers.Container{
		ID: id,
	}
	_, ok := namespaces.Namespace(ctx)
	if !ok {
		ctx = namespaces.WithNamespace(ctx, "buildkit")
	}

	// Note that containerd.GenerateSpec is namespaced so as to make
	// specs.Linux.CgroupsPath namespaced
	s, err := oci.GenerateSpec(ctx, nil, c, opts...)
	if err != nil {
		return nil, nil, err
	}
	// set the networking information on the spec
	namespace.Set(s)

	s.Process.Args = meta.Args
	s.Process.Env = meta.Env
	s.Process.Cwd = meta.Cwd
	s.Process.Rlimits = nil // reset open files limit
	s.Hostname = "buildkitsandbox"

	s.Mounts = GetMounts(ctx,
		withROBind(resolvConf, "/etc/resolv.conf"),
		withROBind(hostsFile, "/etc/hosts"),
	)

	s.Mounts = append(s.Mounts, specs.Mount{
		Destination: "/sys/fs/cgroup",
		Type:        "cgroup",
		Source:      "cgroup",
		Options:     []string{"ro", "nosuid", "noexec", "nodev"},
	})

	// TODO: User

	sm := &submounts{}

	var releasers []func() error
	releaseAll := func() {
		sm.cleanup()
		for _, f := range releasers {
			f()
		}
	}

	for _, m := range mounts {
		if m.Src == nil {
			return nil, nil, errors.Errorf("mount %s has no source", m.Dest)
		}
		mountable, err := m.Src.Mount(ctx, m.Readonly)
		if err != nil {
			releaseAll()
			return nil, nil, errors.Wrapf(err, "failed to mount %s", m.Dest)
		}
		mounts, err := mountable.Mount()
		if err != nil {
			releaseAll()
			return nil, nil, errors.WithStack(err)
		}
		releasers = append(releasers, mountable.Release)
		for _, mount := range mounts {
			mount, err = sm.subMount(mount, m.Selector)
			if err != nil {
				releaseAll()
				return nil, nil, err
			}
			s.Mounts = append(s.Mounts, specs.Mount{
				Destination: m.Dest,
				Type:        mount.Type,
				Source:      mount.Source,
				Options:     mount.Options,
			})
		}
	}

	return s, releaseAll, nil
}

type mountRef struct {
	mount   mount.Mount
	unmount func() error
}

type submounts struct {
	m map[uint64]mountRef
}

func (s *submounts) subMount(m mount.Mount, subPath string) (mount.Mount, error) {
	if path.Join("/", subPath) == "/" {
		return m, nil
	}
	if s.m == nil {
		s.m = map[uint64]mountRef{}
	}
	h, err := hashstructure.Hash(m, nil)
	if err != nil {
		return mount.Mount{}, nil
	}
	if mr, ok := s.m[h]; ok {
		sm, err := sub(mr.mount, subPath)
		if err != nil {
			return mount.Mount{}, nil
		}
		return sm, nil
	}

	lm := snapshot.LocalMounterWithMounts([]mount.Mount{m})

	mp, err := lm.Mount()
	if err != nil {
		return mount.Mount{}, err
	}

	opts := []string{"rbind"}
	for _, opt := range m.Options {
		if opt == "ro" {
			opts = append(opts, opt)
		}
	}

	s.m[h] = mountRef{
		mount: mount.Mount{
			Source:  mp,
			Type:    "bind",
			Options: opts,
		},
		unmount: lm.Unmount,
	}

	sm, err := sub(s.m[h].mount, subPath)
	if err != nil {
		return mount.Mount{}, err
	}
	return sm, nil
}

func (s *submounts) cleanup() {
	var wg sync.WaitGroup
	wg.Add(len(s.m))
	for _, m := range s.m {
		func(m mountRef) {
			go func() {
				m.unmount()
				wg.Done()
			}()
		}(m)
	}
	wg.Wait()
}

func sub(m mount.Mount, subPath string) (mount.Mount, error) {
	src, err := fs.RootPath(m.Source, subPath)
	if err != nil {
		return mount.Mount{}, err
	}
	m.Source = src
	return m, nil
}
<commit_msg>oci: avoid nonewprivileges in default spec<commit_after>// +build !windows

package oci

import (
	"context"
	"path"
	"sync"

	"github.com/containerd/containerd/containers"
	"github.com/containerd/containerd/mount"
	"github.com/containerd/containerd/namespaces"
	"github.com/containerd/containerd/oci"
	"github.com/containerd/continuity/fs"
	"github.com/mitchellh/hashstructure"
	"github.com/moby/buildkit/executor"
	"github.com/moby/buildkit/snapshot"
	"github.com/moby/buildkit/util/network"
	specs "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
)

// Ideally we don't have to import whole containerd just for the default spec

// GenerateSpec generates spec using containerd functionality.
func GenerateSpec(ctx context.Context, meta executor.Meta, mounts []executor.Mount, id, resolvConf, hostsFile string, namespace network.Namespace, opts ...oci.SpecOpts) (*specs.Spec, func(), error) {
	c := &containers.Container{
		ID: id,
	}
	_, ok := namespaces.Namespace(ctx)
	if !ok {
		ctx = namespaces.WithNamespace(ctx, "buildkit")
	}

	// Note that containerd.GenerateSpec is namespaced so as to make
	// specs.Linux.CgroupsPath namespaced
	s, err := oci.GenerateSpec(ctx, nil, c, opts...)
	if err != nil {
		return nil, nil, err
	}
	// set the networking information on the spec
	namespace.Set(s)

	s.Process.Args = meta.Args
	s.Process.Env = meta.Env
	s.Process.Cwd = meta.Cwd
	s.Process.Rlimits = nil           // reset open files limit
	s.Process.NoNewPrivileges = false // reset nonewprivileges
	s.Hostname = "buildkitsandbox"

	s.Mounts = GetMounts(ctx,
		withROBind(resolvConf, "/etc/resolv.conf"),
		withROBind(hostsFile, "/etc/hosts"),
	)

	s.Mounts = append(s.Mounts, specs.Mount{
		Destination: "/sys/fs/cgroup",
		Type:        "cgroup",
		Source:      "cgroup",
		Options:     []string{"ro", "nosuid", "noexec", "nodev"},
	})

	// TODO: User

	sm := &submounts{}

	var releasers []func() error
	releaseAll := func() {
		sm.cleanup()
		for _, f := range releasers {
			f()
		}
	}

	for _, m := range mounts {
		if m.Src == nil {
			return nil, nil, errors.Errorf("mount %s has no source", m.Dest)
		}
		mountable, err := m.Src.Mount(ctx, m.Readonly)
		if err != nil {
			releaseAll()
			return nil, nil, errors.Wrapf(err, "failed to mount %s", m.Dest)
		}
		mounts, err := mountable.Mount()
		if err != nil {
			releaseAll()
			return nil, nil, errors.WithStack(err)
		}
		releasers = append(releasers, mountable.Release)
		for _, mount := range mounts {
			mount, err = sm.subMount(mount, m.Selector)
			if err != nil {
				releaseAll()
				return nil, nil, err
			}
			s.Mounts = append(s.Mounts, specs.Mount{
				Destination: m.Dest,
				Type:        mount.Type,
				Source:      mount.Source,
				Options:     mount.Options,
			})
		}
	}

	return s, releaseAll, nil
}

type mountRef struct {
	mount   mount.Mount
	unmount func() error
}

type submounts struct {
	m map[uint64]mountRef
}

func (s *submounts) subMount(m mount.Mount, subPath string) (mount.Mount, error) {
	if path.Join("/", subPath) == "/" {
		return m, nil
	}
	if s.m == nil {
		s.m = map[uint64]mountRef{}
	}
	h, err := hashstructure.Hash(m, nil)
	if err != nil {
		return mount.Mount{}, nil
	}
	if mr, ok := s.m[h]; ok {
		sm, err := sub(mr.mount, subPath)
		if err != nil {
			return mount.Mount{}, nil
		}
		return sm, nil
	}

	lm := snapshot.LocalMounterWithMounts([]mount.Mount{m})

	mp, err := lm.Mount()
	if err != nil {
		return mount.Mount{}, err
	}

	opts := []string{"rbind"}
	for _, opt := range m.Options {
		if opt == "ro" {
			opts = append(opts, opt)
		}
	}

	s.m[h] = mountRef{
		mount: mount.Mount{
			Source:  mp,
			Type:    "bind",
			Options: opts,
		},
		unmount: lm.Unmount,
	}

	sm, err := sub(s.m[h].mount, subPath)
	if err != nil {
		return mount.Mount{}, err
	}
	return sm, nil
}

func (s *submounts) cleanup() {
	var wg sync.WaitGroup
	wg.Add(len(s.m))
	for _, m := range s.m {
		func(m mountRef) {
			go func() {
				m.unmount()
				wg.Done()
			}()
		}(m)
	}
	wg.Wait()
}

func sub(m mount.Mount, subPath string) (mount.Mount, error) {
	src, err := fs.RootPath(m.Source, subPath)
	if err != nil {
		return mount.Mount{}, err
	}
	m.Source = src
	return m, nil
}
<|endoftext|>


<commit_before>package expr

import (
	"github.com/grafana/metrictank/api/models"
	"github.com/grafana/metrictank/schema"
)

type FuncGroupByNodes struct {
	in         GraphiteFunc
	aggregator string
	nodes      []expr
	byNode     bool
}

func NewGroupByNodesConstructor(groupByNode bool) func() GraphiteFunc {
	return func() GraphiteFunc {
		return &FuncGroupByNodes{byNode: groupByNode}
	}
}

func (s *FuncGroupByNodes) Signature() ([]Arg, []Arg) {
	// This function supports both groupByNode and groupByNodes,
	// but the signatures of them are different in Graphite.
	if !s.byNode {
		return []Arg{
			ArgSeriesList{val: &s.in},
			ArgString{val: &s.aggregator, validator: []Validator{IsAggFunc}},
			ArgStringsOrInts{val: &s.nodes},
		}, []Arg{ArgSeries{}}
	}
	// groupByNode accepts either an integer node num or a string tag key
	s.nodes = append(s.nodes, expr{})
	return []Arg{
		ArgSeriesList{val: &s.in},
		ArgStringOrInt{val: &s.nodes[0]},
		ArgString{val: &s.aggregator, validator: []Validator{IsAggFunc}},
	}, []Arg{ArgSeries{}}
}

func (s *FuncGroupByNodes) Context(context Context) Context {
	context.PNGroup = 0
	return context
}

func (s *FuncGroupByNodes) Exec(dataMap DataMap) ([]models.Series, error) {
	series, err := s.in.Exec(dataMap)
	if err != nil {
		return nil, err
	}

	if len(series) == 0 {
		return series, nil
	}

	type Group struct {
		s []models.Series
		m models.SeriesMeta
	}
	groups := make(map[string]Group)
	// Group series by nodes, this is mostly similar to GroupByTags,
	// except that the group keys are different.
	for _, serie := range series {
		key := aggKey(serie, s.nodes)
		group := groups[key]
		group.s = append(group.s, serie)
		group.m = group.m.Merge(serie.Meta)
		groups[key] = group
	}
	// Similar to FuncGroupByTags, apply aggregate functions to each group
	output := make([]models.Series, 0, len(groups))
	aggFunc := getCrossSeriesAggFunc(s.aggregator)

	for key, group := range groups {
		consolidator, queryConsolidator := summarizeCons(group.s)
		outSeries := models.Series{
			Target:       key,
			QueryPatt:    key,
			Consolidator: consolidator,
			QueryCons:    queryConsolidator,
			QueryFrom:    group.s[0].QueryFrom,
			QueryTo:      group.s[0].QueryTo,
			QueryMDP:     group.s[0].QueryMDP,
			QueryPNGroup: group.s[0].QueryPNGroup,
			Meta:         group.m,
		}
		group.s = Normalize(dataMap, group.s)
		outSeries.Interval = group.s[0].Interval
		outSeries.SetTags()
		outSeries.Datapoints = pointSlicePool.Get().([]schema.Point)
		aggFunc(group.s, &outSeries.Datapoints)
		dataMap.Add(Req{}, outSeries)
		output = append(output, outSeries)
	}
	return output, nil
}
<commit_msg>groupByNode[s]: respect order. fix #1767<commit_after>package expr

import (
	"github.com/grafana/metrictank/api/models"
	"github.com/grafana/metrictank/schema"
)

type FuncGroupByNodes struct {
	in         GraphiteFunc
	aggregator string
	nodes      []expr
	byNode     bool
}

func NewGroupByNodesConstructor(groupByNode bool) func() GraphiteFunc {
	return func() GraphiteFunc {
		return &FuncGroupByNodes{byNode: groupByNode}
	}
}

func (s *FuncGroupByNodes) Signature() ([]Arg, []Arg) {
	// This function supports both groupByNode and groupByNodes,
	// but the signatures of them are different in Graphite.
	if !s.byNode {
		return []Arg{
			ArgSeriesList{val: &s.in},
			ArgString{val: &s.aggregator, validator: []Validator{IsAggFunc}},
			ArgStringsOrInts{val: &s.nodes},
		}, []Arg{ArgSeries{}}
	}
	// groupByNode accepts either an integer node num or a string tag key
	s.nodes = append(s.nodes, expr{})
	return []Arg{
		ArgSeriesList{val: &s.in},
		ArgStringOrInt{val: &s.nodes[0]},
		ArgString{val: &s.aggregator, validator: []Validator{IsAggFunc}},
	}, []Arg{ArgSeries{}}
}

func (s *FuncGroupByNodes) Context(context Context) Context {
	context.PNGroup = 0
	return context
}

func (s *FuncGroupByNodes) Exec(dataMap DataMap) ([]models.Series, error) {
	series, err := s.in.Exec(dataMap)
	if err != nil {
		return nil, err
	}

	if len(series) == 0 {
		return series, nil
	}

	type Group struct {
		s []models.Series
		m models.SeriesMeta
	}

	groups := make(map[string]Group)

	// list of aggregation keys in order they were seen.
	// we need to return series in this order
	var keyList []string

	// Group series by nodes, this is mostly similar to GroupByTags,
	// except that the group keys are different.
	for _, serie := range series {
		key := aggKey(serie, s.nodes)
		group, ok := groups[key]
		if !ok {
			keyList = append(keyList, key)
		}
		group.s = append(group.s, serie)
		group.m = group.m.Merge(serie.Meta)
		groups[key] = group
	}
	// Similar to FuncGroupByTags, apply aggregate functions to each group
	output := make([]models.Series, 0, len(groups))
	aggFunc := getCrossSeriesAggFunc(s.aggregator)

	for _, key := range keyList {
		group := groups[key]
		consolidator, queryConsolidator := summarizeCons(group.s)
		outSeries := models.Series{
			Target:       key,
			QueryPatt:    key,
			Consolidator: consolidator,
			QueryCons:    queryConsolidator,
			QueryFrom:    group.s[0].QueryFrom,
			QueryTo:      group.s[0].QueryTo,
			QueryMDP:     group.s[0].QueryMDP,
			QueryPNGroup: group.s[0].QueryPNGroup,
			Meta:         group.m,
		}
		group.s = Normalize(dataMap, group.s)
		outSeries.Interval = group.s[0].Interval
		outSeries.SetTags()
		outSeries.Datapoints = pointSlicePool.Get().([]schema.Point)
		aggFunc(group.s, &outSeries.Datapoints)
		dataMap.Add(Req{}, outSeries)
		output = append(output, outSeries)
	}
	return output, nil
}
<|endoftext|>


<commit_before>package msp

import (
	"encoding/json"
	"github.com/gorilla/mux"
	"net/http"
	"fmt"
)

type Consumer interface {
	RequestStatus() ([]Mongod, *Error)
	EstablishMongodState(m Mongod) *Error
}

type Listener struct {
	port PortNumber
	listener Consumer
	router   *mux.Router
}

func NewServer(listener Consumer, port PortNumber) *Listener {
	s := new(Listener)
	s.listener = listener
	s.port = port

	s.router = mux.NewRouter().StrictSlash(true)
	s.router.Methods("GET").Path("/msp/status").Name("RequestStatus").HandlerFunc(s.handleRequestStatus)
	s.router.Methods("POST").Path("/msp/establishMongodState").Name("EstablishMongodState").HandlerFunc(s.handleMspEstablishMongodState)

	return s
}

func (s Listener) handleRequestStatus(w http.ResponseWriter, r *http.Request) {
	status, err := s.listener.RequestStatus()
	if status != nil {
		json.NewEncoder(w).Encode(status)
	} else {
		w.WriteHeader(http.StatusInternalServerError)
		json.NewEncoder(w).Encode(err)
	}
}

func (s Listener) handleMspEstablishMongodState(w http.ResponseWriter, r *http.Request) {
	var mongodState Mongod
	json.NewDecoder(r.Body).Decode(&mongodState) //TODO Check decode error
	err := s.listener.EstablishMongodState(mongodState)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		json.NewEncoder(w).Encode(err)
	}
}

func (s Listener) Run() {
	http.ListenAndServe(fmt.Sprintf(":%d", s.port), s.router)
}
<commit_msg>FIX: format<commit_after>package msp

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	"net/http"
)

type Consumer interface {
	RequestStatus() ([]Mongod, *Error)
	EstablishMongodState(m Mongod) *Error
}

type Listener struct {
	port     PortNumber
	listener Consumer
	router   *mux.Router
}

func NewServer(listener Consumer, port PortNumber) *Listener {
	s := new(Listener)
	s.listener = listener
	s.port = port

	s.router = mux.NewRouter().StrictSlash(true)
	s.router.Methods("GET").Path("/msp/status").Name("RequestStatus").HandlerFunc(s.handleRequestStatus)
	s.router.Methods("POST").Path("/msp/establishMongodState").Name("EstablishMongodState").HandlerFunc(s.handleMspEstablishMongodState)

	return s
}

func (s Listener) handleRequestStatus(w http.ResponseWriter, r *http.Request) {
	status, err := s.listener.RequestStatus()
	if status != nil {
		json.NewEncoder(w).Encode(status)
	} else {
		w.WriteHeader(http.StatusInternalServerError)
		json.NewEncoder(w).Encode(err)
	}
}

func (s Listener) handleMspEstablishMongodState(w http.ResponseWriter, r *http.Request) {
	var mongodState Mongod
	json.NewDecoder(r.Body).Decode(&mongodState) //TODO Check decode error
	err := s.listener.EstablishMongodState(mongodState)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		json.NewEncoder(w).Encode(err)
	}
}

func (s Listener) Run() {
	http.ListenAndServe(fmt.Sprintf(":%d", s.port), s.router)
}
<|endoftext|>


<commit_before>package config

import (
	"bytes"
	"errors"
	"io"
	"log"
	"os"
)

type Config struct {
	Path string
}

const (
	NO_BUF    = iota // no buffer
	VAR_NAME  = iota // variable name
	VAR_VALUE = iota // string variable value
)

type parserState struct {
	lastByte     byte          // previous byte
	commentLevel uint8         // comment level
	inComment    bool          // true if in a comment
	escaped      bool          // true if the current byte is escaped
	buffer       *bytes.Buffer // current buffer
	buffType     uint8         // type of buffer
}

// increase block comment level
func (state *parserState) increaseCommentLevel() {
	state.removeLastCharacter()
	state.commentLevel++
	state.inComment = true
}

// decrease block comment level
func (state *parserState) decreaseCommentLevel() {
	state.removeLastCharacter()
	if state.commentLevel <= 0 {
		return
	}
	state.commentLevel--
	state.inComment = state.commentLevel != 0
}

func (state *parserState) removeLastCharacter() {
	if state.buffer == nil || state.buffer.Len() < 1 {
		return
	}
	state.buffer.Truncate(state.buffer.Len() - 1)
}

// start a new buffer
func (state *parserState) startBuffer(t uint8) {
	state.buffer = new(bytes.Buffer)
	state.buffType = t
}

// destroy a buffer
func (state *parserState) endBuffer() string {
	str := state.buffer.String()
	state.buffer = nil
	state.buffType = NO_BUF
	return str
}

func New(path string) *Config {
	return &Config{Path: path}
}

func (conf *Config) Parse() error {

	// open the config
	file, err := os.Open(conf.Path)
	defer file.Close()
	if err != nil {
		return err
	}

	state := &parserState{}
	for {
		b := make([]byte, 1)
		_, err := file.Read(b)

		// eof
		if err == io.EOF {
			break
		}

		// some other error
		if err != nil {
			return err
		}

		// handle the character
		err = conf.handleCharacter(state, b[0])

		// character error
		if err != nil {
			return err
		}
	}

	return nil
}

func (conf *Config) handleCharacter(state *parserState, b byte) error {

	// this character is escaped
	if state.escaped {
		b = 0
		state.escaped = false
	}

BYTE:
	switch b {

	// comment entrance
	case '*':
		if state.lastByte == '/' {
			state.increaseCommentLevel()
			break
		}
		fallthrough

	// comment closure
	case '/':
		if state.lastByte == '*' && state.inComment {
			state.decreaseCommentLevel()
			break
		}
		fallthrough

	// escape
	case '\\':
		state.escaped = true

	// start of a variable
	case '@':

		// we're already in a variable name
		if state.buffType == VAR_NAME {
			return errors.New("Invalid @ in variable name @" + state.buffer.String())
		}

		// we're not in a value, so this starts a variable
		if state.buffType == 0 {
			log.Println("Starting variable name")
			state.startBuffer(VAR_NAME)
			break
		}

		fallthrough

	// end of variable name
	case ':':

		// we're in the var name, so terminate it
		if state.buffType == VAR_NAME {
			name := state.endBuffer()
			log.Println("Got variable name: @" + name)
			state.startBuffer(VAR_VALUE)
			break
		}

		fallthrough

	// end of a variable definition
	case ';':

		switch state.buffType {

		// terminating a boolean
		case VAR_NAME:
			state.endBuffer()
			log.Println("Got boolean")
			break BYTE

		// terminating a string
		case VAR_VALUE:
			str := state.endBuffer()
			log.Println("Got string: " + str)
			break BYTE
		}

		fallthrough

	default:

		// we're in a comment; ignore this
		if state.inComment {
			break
		}

		// otherwise, write this to the current buffer
		if state.buffer != nil {
			state.buffer.WriteByte(b)
		}
	}

	state.lastByte = b
	return nil
}

func (conf *Config) Get(varName string) string {
	return ""
}

func (conf *Config) GetBool(varName string) bool {
	if str := conf.Get(varName); str == "" {
		return false
	}
	return true
}
<commit_msg>realDefault<commit_after>package config

import (
	"bytes"
	"errors"
	"io"
	"log"
	"os"
)

type Config struct {
	Path string
}

const (
	NO_BUF    = iota // no buffer
	VAR_NAME  = iota // variable name
	VAR_VALUE = iota // string variable value
)

type parserState struct {
	lastByte     byte          // previous byte
	commentLevel uint8         // comment level
	inComment    bool          // true if in a comment
	escaped      bool          // true if the current byte is escaped
	buffer       *bytes.Buffer // current buffer
	buffType     uint8         // type of buffer
}

// increase block comment level
func (state *parserState) increaseCommentLevel() {
	state.removeLastCharacter()
	state.commentLevel++
	state.inComment = true
}

// decrease block comment level
func (state *parserState) decreaseCommentLevel() {
	state.removeLastCharacter()
	if state.commentLevel <= 0 {
		return
	}
	state.commentLevel--
	state.inComment = state.commentLevel != 0
}

func (state *parserState) removeLastCharacter() {
	if state.buffer == nil || state.buffer.Len() < 1 {
		return
	}
	state.buffer.Truncate(state.buffer.Len() - 1)
}

// start a new buffer
func (state *parserState) startBuffer(t uint8) {
	state.buffer = new(bytes.Buffer)
	state.buffType = t
}

// destroy a buffer
func (state *parserState) endBuffer() string {
	str := state.buffer.String()
	state.buffer = nil
	state.buffType = NO_BUF
	return str
}

func New(path string) *Config {
	return &Config{Path: path}
}

func (conf *Config) Parse() error {

	// open the config
	file, err := os.Open(conf.Path)
	defer file.Close()
	if err != nil {
		return err
	}

	state := &parserState{}
	for {
		b := make([]byte, 1)
		_, err := file.Read(b)

		// eof
		if err == io.EOF {
			break
		}

		// some other error
		if err != nil {
			return err
		}

		// handle the character
		err = conf.handleCharacter(state, b[0])
        state.lastByte = b[0]

		// character error
		if err != nil {
			return err
		}
	}

	return nil
}

func (conf *Config) handleCharacter(state *parserState, b byte) error {

	// this character is escaped
	if state.escaped {
		b = 0
		state.escaped = false
	}

	switch b {

	// comment entrance
	case '*':
		if state.lastByte != '/' {
			goto realDefault
		}
		log.Println("Increasing comment level")
		state.increaseCommentLevel()

	// comment closure
	case '/':
		if state.lastByte != '*' || !state.inComment {
			goto realDefault
		}
		log.Println("Decreasing comment level")
		state.decreaseCommentLevel()

	// escape
	case '\\':
		state.escaped = true

	// start of a variable
	case '@':

		// we're already in a variable name
		if state.buffType == VAR_NAME {
			return errors.New("Invalid @ in variable name @" + state.buffer.String())
		}

		// we're not in a value, so this starts a variable
		if state.buffType != NO_BUF {
			goto realDefault
		}
		log.Println("Starting variable name")
		state.startBuffer(VAR_NAME)

	// end of variable name
	case ':':

		// we're in the var name, so terminate it
		if state.buffType != VAR_NAME {
			goto realDefault
		}
		name := state.endBuffer()
		log.Println("Got variable name: @" + name)
		state.startBuffer(VAR_VALUE)

	// end of a variable definition
	case ';':

		// terminating a boolean
		if state.buffType == VAR_NAME {
			state.endBuffer()
			log.Println("Got boolean")
		} else if state.buffType == VAR_VALUE {
			str := state.endBuffer()
			log.Println("Got string: " + str)
		} else {
			goto realDefault
		}
		break

	default:
		goto realDefault
	}

    // this is skipped if going to realDefault
	return nil

realDefault:

	// we're in a comment; ignore this
	if state.inComment {
		return nil
	}

	// otherwise, write this to the current buffer
	if state.buffer != nil {
		state.buffer.WriteByte(b)
	}

	return nil
}

func (conf *Config) Get(varName string) string {
	return ""
}

func (conf *Config) GetBool(varName string) bool {
	if str := conf.Get(varName); str == "" {
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package config

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/user"
	"path"
	"text/tabwriter"
)

type AccountConfig struct {
	URI        string
	NameFormat *string `json:",omitempty"`
	UserAgent  *string `json:",omitempty"`
}

func (cfg AccountConfig) String() string {
	var b bytes.Buffer
	w := tabwriter.NewWriter(&b, 11, 8, 0, '\t', 0)
	fmt.Fprintf(w, "URI\t\"%s\"\n", cfg.URI)
	if cfg.NameFormat != nil {
		fmt.Fprintf(w, "Name Format\t\"%s\"\n", *cfg.NameFormat)
	}
	if cfg.UserAgent != nil {
		fmt.Fprintf(w, "User Agent\t\"%s\"\n", *cfg.UserAgent)
	}
	w.Flush()
	return b.String()
}

type GrueConfig struct {
	path        string
	Recipient   string
	FromAddress string
	NameFormat  string
	UserAgent   string
	SmtpUser    *string
	SmtpPass    *string
	SmtpServer  *string
	LogLevel    *string
	Accounts    map[string]AccountConfig
}

func (conf *GrueConfig) Lock() error {
	lock, err := os.OpenFile(conf.path+".lock", os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0666)
	switch {
	case err == nil:
		defer lock.Close()
		_, err = lock.WriteString(string(os.Getpid()))
		return err
	case os.IsExist(err):
		return fmt.Errorf("Aborting due to existing lock on %s\n", conf.path)
	default:
		return err
	}
}

func (conf *GrueConfig) Unlock() error {
	return os.Remove(conf.path + ".lock")
}

func (conf *GrueConfig) String() string {
	b, err := json.Marshal(conf)
	if err != nil {
		panic("Can't Marshal config")
	}
	return string(b)
}

func defaultFrom() (string, error) {
	cur, err := user.Current()
	if err != nil {
		return "", err
	}
	host, err := os.Hostname()
	if err != nil {
		return "", err
	}
	from := cur.Username + "@" + host
	return from, nil
}

func (conf *GrueConfig) encode(file *os.File) error {
	enc := json.NewEncoder(file)
	enc.SetIndent("", "  ")
	return enc.Encode(conf)
}

func (conf *GrueConfig) save() error {
	tmpfile, err := ioutil.TempFile(path.Dir(conf.path), path.Base(conf.path))
	if err != nil {
		return err
	}
	err = conf.encode(tmpfile)
	tmpfile.Close()
	if err != nil {
		os.Remove(tmpfile.Name())
		return err
	}
	return os.Rename(tmpfile.Name(), conf.path)
}
func makeDefConfig() (*GrueConfig, error) {
	from, err := defaultFrom()
	if err != nil {
		return nil, err
	}
	var conf = &GrueConfig{FromAddress: from, NameFormat: "{name}: {title}", UserAgent: "grue/{version}"}
	return conf, nil
}

func writeDefConfig(path string) (*GrueConfig, error) {
	conf, err := makeDefConfig()
	if err != nil {
		return nil, err
	}
	conf.path = path
	return conf, conf.save()
}

func getConfigPath() string {
	cfgPath := os.Getenv("XDG_CONFIG_HOME")
	if cfgPath == "" {
		home := os.Getenv("HOME")
		if home == "" {
			panic("Can't find path to data directory")
		}
		return path.Join(os.Getenv("HOME"), ".config", "grue.cfg")
	}
	return path.Join(cfgPath, "grue.cfg")
}

func ReadConfig() (*GrueConfig, error) {
	var conf *GrueConfig = new(GrueConfig)
	conf.path = getConfigPath()
	err := conf.Lock()
	if err != nil {
		return nil, err
	}
	file, err := os.Open(conf.path)
	if os.IsNotExist(err) {
		return writeDefConfig(conf.path)
	} else if err != nil {
		return nil, err
	}
	defer file.Close()
	dec := json.NewDecoder(file)
	err = dec.Decode(conf)
	if err != nil {
		return nil, err
	}
	return conf, nil
}

func (conf *GrueConfig) AddAccount(name, uri string) error {
	if conf.Accounts == nil {
		conf.Accounts = make(map[string]AccountConfig)
	}
	conf.Accounts[name] = AccountConfig{URI: uri}
	return conf.save()
}

func (conf *GrueConfig) DeleteAccount(name string) error {
	if conf.Accounts == nil {
		return nil
	}
	if _, ok := conf.Accounts[name]; !ok {
		return fmt.Errorf("%s: account does not exist", name)
	}
	delete(conf.Accounts, name)
	return conf.save()
}

func (conf *GrueConfig) RenameAccount(old, new string) error {
	if conf.Accounts == nil {
		return fmt.Errorf("%s: account does not exist", old)
	}
	if _, ok := conf.Accounts[old]; !ok {
		return fmt.Errorf("%s: account does not exist", old)
	}
	if _, ok := conf.Accounts[new]; ok {
		return fmt.Errorf("%s: account already exists", new)
	}
	conf.Accounts[new] = conf.Accounts[old]
	delete(conf.Accounts, old)
	return conf.save()
}
<commit_msg>config: Use fmt.Sprintf instead of string() for pid -> string<commit_after>package config

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/user"
	"path"
	"text/tabwriter"
)

type AccountConfig struct {
	URI        string
	NameFormat *string `json:",omitempty"`
	UserAgent  *string `json:",omitempty"`
}

func (cfg AccountConfig) String() string {
	var b bytes.Buffer
	w := tabwriter.NewWriter(&b, 11, 8, 0, '\t', 0)
	fmt.Fprintf(w, "URI\t\"%s\"\n", cfg.URI)
	if cfg.NameFormat != nil {
		fmt.Fprintf(w, "Name Format\t\"%s\"\n", *cfg.NameFormat)
	}
	if cfg.UserAgent != nil {
		fmt.Fprintf(w, "User Agent\t\"%s\"\n", *cfg.UserAgent)
	}
	w.Flush()
	return b.String()
}

type GrueConfig struct {
	path        string
	Recipient   string
	FromAddress string
	NameFormat  string
	UserAgent   string
	SmtpUser    *string
	SmtpPass    *string
	SmtpServer  *string
	LogLevel    *string
	Accounts    map[string]AccountConfig
}

func (conf *GrueConfig) Lock() error {
	lock, err := os.OpenFile(conf.path+".lock", os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0666)
	switch {
	case err == nil:
		defer lock.Close()
		_, err = lock.WriteString(fmt.Sprint(os.Getpid()))
		return err
	case os.IsExist(err):
		return fmt.Errorf("Aborting due to existing lock on %s\n", conf.path)
	default:
		return err
	}
}

func (conf *GrueConfig) Unlock() error {
	return os.Remove(conf.path + ".lock")
}

func (conf *GrueConfig) String() string {
	b, err := json.Marshal(conf)
	if err != nil {
		panic("Can't Marshal config")
	}
	return string(b)
}

func defaultFrom() (string, error) {
	cur, err := user.Current()
	if err != nil {
		return "", err
	}
	host, err := os.Hostname()
	if err != nil {
		return "", err
	}
	from := cur.Username + "@" + host
	return from, nil
}

func (conf *GrueConfig) encode(file *os.File) error {
	enc := json.NewEncoder(file)
	enc.SetIndent("", "  ")
	return enc.Encode(conf)
}

func (conf *GrueConfig) save() error {
	tmpfile, err := ioutil.TempFile(path.Dir(conf.path), path.Base(conf.path))
	if err != nil {
		return err
	}
	err = conf.encode(tmpfile)
	tmpfile.Close()
	if err != nil {
		os.Remove(tmpfile.Name())
		return err
	}
	return os.Rename(tmpfile.Name(), conf.path)
}
func makeDefConfig() (*GrueConfig, error) {
	from, err := defaultFrom()
	if err != nil {
		return nil, err
	}
	var conf = &GrueConfig{FromAddress: from, NameFormat: "{name}: {title}", UserAgent: "grue/{version}"}
	return conf, nil
}

func writeDefConfig(path string) (*GrueConfig, error) {
	conf, err := makeDefConfig()
	if err != nil {
		return nil, err
	}
	conf.path = path
	return conf, conf.save()
}

func getConfigPath() string {
	cfgPath := os.Getenv("XDG_CONFIG_HOME")
	if cfgPath == "" {
		home := os.Getenv("HOME")
		if home == "" {
			panic("Can't find path to data directory")
		}
		return path.Join(os.Getenv("HOME"), ".config", "grue.cfg")
	}
	return path.Join(cfgPath, "grue.cfg")
}

func ReadConfig() (*GrueConfig, error) {
	var conf *GrueConfig = new(GrueConfig)
	conf.path = getConfigPath()
	err := conf.Lock()
	if err != nil {
		return nil, err
	}
	file, err := os.Open(conf.path)
	if os.IsNotExist(err) {
		return writeDefConfig(conf.path)
	} else if err != nil {
		return nil, err
	}
	defer file.Close()
	dec := json.NewDecoder(file)
	err = dec.Decode(conf)
	if err != nil {
		return nil, err
	}
	return conf, nil
}

func (conf *GrueConfig) AddAccount(name, uri string) error {
	if conf.Accounts == nil {
		conf.Accounts = make(map[string]AccountConfig)
	}
	conf.Accounts[name] = AccountConfig{URI: uri}
	return conf.save()
}

func (conf *GrueConfig) DeleteAccount(name string) error {
	if conf.Accounts == nil {
		return nil
	}
	if _, ok := conf.Accounts[name]; !ok {
		return fmt.Errorf("%s: account does not exist", name)
	}
	delete(conf.Accounts, name)
	return conf.save()
}

func (conf *GrueConfig) RenameAccount(old, new string) error {
	if conf.Accounts == nil {
		return fmt.Errorf("%s: account does not exist", old)
	}
	if _, ok := conf.Accounts[old]; !ok {
		return fmt.Errorf("%s: account does not exist", old)
	}
	if _, ok := conf.Accounts[new]; ok {
		return fmt.Errorf("%s: account already exists", new)
	}
	conf.Accounts[new] = conf.Accounts[old]
	delete(conf.Accounts, old)
	return conf.save()
}
<|endoftext|>