<commit_before>package main

// TODO:
// - to store icons, see: http://blogs.msdn.com/b/oldnewthing/archive/2012/07/20/10331787.aspx
//   - also need to first read and split input ico file

import (
	"debug/pe"
	"encoding/binary"
	"flag"
	"fmt"
	"os"
	"reflect"
	"regexp"
	"strconv"

	"github.com/akavel/rsrc/ico"
)

type Dir struct { // struct IMAGE_RESOURCE_DIRECTORY
	Characteristics      uint32
	TimeDateStamp        uint32
	MajorVersion         uint16
	MinorVersion         uint16
	NumberOfNamedEntries uint16
	NumberOfIdEntries    uint16
	DirEntries
	Dirs
}

type DirEntries []DirEntry
type Dirs []Dir

type DirEntry struct { // struct IMAGE_RESOURCE_DIRECTORY_ENTRY
	NameOrId     uint32
	OffsetToData uint32
}

type DataEntry struct { // struct IMAGE_RESOURCE_DATA_ENTRY
	OffsetToData uint32
	Size1        uint32
	CodePage     uint32
	Reserved     uint32
}

type RelocationEntry struct {
	RVA         uint32 // "offset within the Section's raw data where the address starts."
	SymbolIndex uint32 // "(zero based) index in the Symbol table to which the reference refers."
	Type        uint16
}

type Symbol struct {
	Name           [8]byte
	Value          uint32
	SectionNumber  uint16
	Type           uint16
	StorageClass   uint8
	AuxiliaryCount uint8
}

type StringsHeader struct {
	Length uint32
}

const (
	MASK_SUBDIRECTORY = 1 << 31

	RT_ICON       = 3
	RT_GROUP_ICON = 3 + 11
	RT_MANIFEST   = 24
)

var (
	STRING_RSRC = [8]byte{'.', 'r', 's', 'r', 'c', 0, 0, 0}
	LANG_ENTRY  = DirEntries{{NameOrId: 0x0409}} //FIXME: language; what value should be here?
)

type GRPICONDIR struct {
	ico.ICONDIR
	Entries []GRPICONDIRENTRY
}

type GRPICONDIRENTRY struct {
	ico.IconDirEntryCommon
	Id uint16
}

func main() {
	//TODO: allow in options advanced specification of multiple resources, as a tree (json?)
	var fnamein, fnameico, fnameout string
	flags := flag.NewFlagSet("", flag.ContinueOnError)
	flags.StringVar(&fnamein, "manifest", "", "path to a Windows manifest file to embed")
	flags.StringVar(&fnameico, "ico", "", "UNSUPPORTED: path to ICO file to embed")
	flags.StringVar(&fnameout, "o", "rsrc.syso", "name of output COFF (.res or .syso) file")
	_ = flags.Parse(os.Args[1:])
	if fnamein == "" {
		fmt.Fprintf(os.Stderr, "USAGE: %s -manifest FILE.exe.manifest [-o FILE.syso]\n"+
			"Generates a .syso file with specified resources embedded in .rsrc section,\n"+
			"aimed for consumption by Go linker when building Win32 excecutables.\n"+
			"OPTIONS:\n",
			os.Args[0])
		flags.PrintDefaults()
		os.Exit(1)
	}

	err := run(fnamein, fnameico, fnameout)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

type Coff struct {
	pe.FileHeader
	pe.SectionHeader32

	Dir
	DataEntries []DataEntry
	Data        []interface{}

	Relocations []RelocationEntry
	Symbols     []Symbol
	StringsHeader
}

func run(fnamein, fnameico, fnameout string) error {
	manifest, err := SizedOpen(fnamein)
	if err != nil {
		return fmt.Errorf("Error opening manifest file '%s': %s", fnamein, err)
	}
	defer manifest.Close()

	var icons []ico.ICONDIRENTRY
	var iconsf *os.File
	if fnameico != "" {
		iconsf, err = os.Open(fnameico)
		if err != nil {
			return err
		}
		defer iconsf.Close()
		icons, err = ico.DecodeHeaders(iconsf)
		if err != nil {
			return err
		}
	}

	newid := make(chan uint16)
	go func() {
		for i := uint16(1); ; i++ {
			newid <- i
		}
	}()

	out, err := os.Create(fnameout)
	if err != nil {
		return err
	}
	defer out.Close()
	w := Writer{W: out}

	coff := Coff{
		pe.FileHeader{
			Machine:              0x014c, //FIXME: find out how to differentiate this value, or maybe not necessary for Go
			NumberOfSections:     1,      // .rsrc
			TimeDateStamp:        0,      // was also 0 in sample data from MinGW's windres.exe
			NumberOfSymbols:      1,
			SizeOfOptionalHeader: 0,
			Characteristics:      0x0104, //FIXME: copied from windres.exe output, find out what should be here and why
		},
		pe.SectionHeader32{
			Name:                STRING_RSRC,
			NumberOfRelocations: 1,
			Characteristics:     0x40000040, // "INITIALIZED_DATA MEM_READ" ?
		},

		// now, build "directory hierarchy" of .rsrc section: first type, then id/name, then language
		Dir{
			NumberOfIdEntries: 1,
			DirEntries:        DirEntries{{NameOrId: RT_MANIFEST}},
			Dirs: Dirs{{
				NumberOfIdEntries: 1,
				DirEntries:        DirEntries{{NameOrId: uint32(<-newid)}}, // resource ID
				Dirs: Dirs{{
					NumberOfIdEntries: 1,
					DirEntries:        LANG_ENTRY,
				}},
			}},
		},
		[]DataEntry{
			DataEntry{
				Size1:    uint32(manifest.Size()),
				CodePage: 0, //FIXME: what value here? for now just tried 0
			},
		},
		[]interface{}{
			manifest,
		},

		[]RelocationEntry{RelocationEntry{
			SymbolIndex: 0, // "(zero based) index in the Symbol table to which the reference refers. Once you have loaded the COFF file into memory and know where each symbol is, you find the new updated address for the given symbol and update the reference accordingly."
			Type:        7, // according to ldpe.c, this decodes to: IMAGE_REL_I386_DIR32NB
		}},

		[]Symbol{Symbol{
			Name:           STRING_RSRC,
			Value:          0,
			SectionNumber:  1,
			Type:           0, // FIXME: wtf?
			StorageClass:   3, // FIXME: is it ok? and uint8? and what does the value mean?
			AuxiliaryCount: 0, // FIXME: wtf?
		}},

		StringsHeader{
			Length: uint32(binary.Size(StringsHeader{})), // empty strings table -- but we must still show size of the table's header...
		},
	}

	if len(icons) > 0 {
		//FIXME: add corresponding DataEntries
		/*
			coff.Dir.NumberOfIdEntries+=2

			coff.Dir.DirEntries = append(coff.Dir.DirEntries, DirEntry{NameOrId: RT_ICON})
			coff.Dir.Dirs = append(coff.Dir.Dirs, Dir{
				NumberOfIdEntries: len(icons),
			})
			dir := &coff.Dir.Dirs[len(coff.Dir.Dirs)-1]
			group := GRPICONDIR{
				Reserved: 0, // magic num.
				Type: 1, // magic num.
				Count: len(icons),
			}
			for i, icon := range icons {
				id := <-newid

				group.Entries = append(group.Entries, GRPICONDIRENTRY{icon, id})
				dir.DirEntries = append(dir.DirEntries, DirEntry{NameOrId: uint32(id)})
				dir.Dirs = append(dir.Dirs, Dir{
					NumberOfIdEntries: 1,
					DirEntries: LANG_ENTRY,
				})

				coff.Data = append(coff.Data, io.NewSectionReader(iconsf, icon.ImageOffset, icon.BytesInRes))
			}

			coff.Dir.DirEntries = append(coff.Dir.DirEntries, DirEntry{NameOrId: RT_GROUP_ICON})
			coff.Dir.Dirs = append(coff.Dir.Dirs, Dir{
				NumberOfIdEntries: 1,
				DirEntries: DirEntries{{NameOrId: uint32(<-newid)}},
				Dirs: Dirs{{
					NumberOfIdEntries: 1,
					DirEntries: LANG_ENTRY,
				}},
			})
			coff.Data = append(coff.Data, group)
		*/
	}

	leafwalker := make(chan *DirEntry)
	go func() {
		for _, dir1 := range coff.Dir.Dirs { // resource type
			for _, dir2 := range dir1.Dirs { // resource ID
				for i := range dir2.DirEntries { // resource lang
					leafwalker <- &dir2.DirEntries[i]
				}
			}
		}
	}()

	N := `\[(\d+)\]`
	dir_n := regexp.MustCompile("^/Dir/Dirs" + N + "$")
	dir_n_n := regexp.MustCompile("^/Dir/Dirs" + N + "/Dirs" + N + "$")
	dataentry_n := regexp.MustCompile("^/DataEntries" + N + "$")
	//dataentry_n_off := regexp.MustCompile("^/DataEntries" + N + "/OffsetToData$")
	//data_n := regexp.MustCompile("^/Data" + N + "$")

	// fill in some important offsets in resulting file
	var offset, diroff uint32
	Walk(coff, func(v reflect.Value, path string) error {
		switch path {
		case "/Dir":
			coff.SectionHeader32.PointerToRawData = offset
			diroff = offset
		//case "/Dir/Dirs[0]":
		//	coff.Dir.DirEntries[0].OffsetToData = MASK_SUBDIRECTORY | (offset - diroff)
		//case "/Dir/Dirs[0]/Dirs[0]":
		//	coff.Dir.Dirs[0].DirEntries[0].OffsetToData = MASK_SUBDIRECTORY | (offset - diroff)
		//case "/DataEntries[0]":
		//	direntry := <-leafwalker
		//	direntry.OffsetToData = offset - diroff
		case "/DataEntries[0]/OffsetToData":
			coff.Relocations[0].RVA = offset - diroff
		case "/Data[0]":
			coff.DataEntries[0].OffsetToData = offset - diroff
		case "/Relocations":
			coff.SectionHeader32.PointerToRelocations = offset
			coff.SectionHeader32.SizeOfRawData = offset - diroff
		case "/Symbols":
			coff.FileHeader.PointerToSymbolTable = offset
		}
		m := matcher{}
		switch {
		case m.Find(path, dir_n):
			coff.Dir.DirEntries[m[0]].OffsetToData = MASK_SUBDIRECTORY | (offset - diroff)
		case m.Find(path, dir_n_n):
			coff.Dir.Dirs[m[0]].DirEntries[m[1]].OffsetToData = MASK_SUBDIRECTORY | (offset - diroff)
		case m.Find(path, dataentry_n):
			direntry := <-leafwalker
			direntry.OffsetToData = offset - diroff
		}

		if Plain(v.Kind()) {
			offset += uint32(binary.Size(v.Interface())) // TODO: change to v.Type().Size() ?
			return nil
		}
		vv, ok := v.Interface().(SizedReader)
		if ok {
			offset += uint32(vv.Size())
			return WALK_SKIP
		}
		return nil
	})

	// write the resulting file to disk
	Walk(coff, func(v reflect.Value, path string) error {
		if Plain(v.Kind()) {
			w.WriteLE(v.Interface())
			return nil
		}
		vv, ok := v.Interface().(SizedReader)
		if ok {
			w.WriteFromSized(vv)
			return WALK_SKIP
		}
		return nil
	})

	if w.Err != nil {
		return fmt.Errorf("Error writing output file: %s", w.Err)
	}

	return nil
}

func Plain(kind reflect.Kind) bool {
	switch kind {
	case reflect.Bool, reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr, reflect.Float32, reflect.Float64, reflect.Complex64, reflect.Complex128:
		return true
	}
	return false
}

func MustAtoi(s string) int {
	i, err := strconv.Atoi(s)
	if err != nil {
		panic(err)
	}
	return i
}

type matcher []int

func (m *matcher) Find(s string, re *regexp.Regexp) bool {
	subs := re.FindStringSubmatch(s)
	if subs == nil {
		return false
	}

	*m = (*m)[:0]
	for i := 1; i < len(subs); i++ {
		*m = append(*m, MustAtoi(subs[i]))
	}
	return true
}
<commit_msg>enabling new approach for dataentry_n_off<commit_after>package main

// TODO:
// - to store icons, see: http://blogs.msdn.com/b/oldnewthing/archive/2012/07/20/10331787.aspx
//   - also need to first read and split input ico file

import (
	"debug/pe"
	"encoding/binary"
	"flag"
	"fmt"
	"os"
	"reflect"
	"regexp"
	"strconv"

	"github.com/akavel/rsrc/ico"
)

type Dir struct { // struct IMAGE_RESOURCE_DIRECTORY
	Characteristics      uint32
	TimeDateStamp        uint32
	MajorVersion         uint16
	MinorVersion         uint16
	NumberOfNamedEntries uint16
	NumberOfIdEntries    uint16
	DirEntries
	Dirs
}

type DirEntries []DirEntry
type Dirs []Dir

type DirEntry struct { // struct IMAGE_RESOURCE_DIRECTORY_ENTRY
	NameOrId     uint32
	OffsetToData uint32
}

type DataEntry struct { // struct IMAGE_RESOURCE_DATA_ENTRY
	OffsetToData uint32
	Size1        uint32
	CodePage     uint32
	Reserved     uint32
}

type RelocationEntry struct {
	RVA         uint32 // "offset within the Section's raw data where the address starts."
	SymbolIndex uint32 // "(zero based) index in the Symbol table to which the reference refers."
	Type        uint16
}

type Symbol struct {
	Name           [8]byte
	Value          uint32
	SectionNumber  uint16
	Type           uint16
	StorageClass   uint8
	AuxiliaryCount uint8
}

type StringsHeader struct {
	Length uint32
}

const (
	MASK_SUBDIRECTORY = 1 << 31

	RT_ICON       = 3
	RT_GROUP_ICON = 3 + 11
	RT_MANIFEST   = 24
)

var (
	STRING_RSRC = [8]byte{'.', 'r', 's', 'r', 'c', 0, 0, 0}
	LANG_ENTRY  = DirEntries{{NameOrId: 0x0409}} //FIXME: language; what value should be here?
)

type GRPICONDIR struct {
	ico.ICONDIR
	Entries []GRPICONDIRENTRY
}

type GRPICONDIRENTRY struct {
	ico.IconDirEntryCommon
	Id uint16
}

func main() {
	//TODO: allow in options advanced specification of multiple resources, as a tree (json?)
	var fnamein, fnameico, fnameout string
	flags := flag.NewFlagSet("", flag.ContinueOnError)
	flags.StringVar(&fnamein, "manifest", "", "path to a Windows manifest file to embed")
	flags.StringVar(&fnameico, "ico", "", "UNSUPPORTED: path to ICO file to embed")
	flags.StringVar(&fnameout, "o", "rsrc.syso", "name of output COFF (.res or .syso) file")
	_ = flags.Parse(os.Args[1:])
	if fnamein == "" {
		fmt.Fprintf(os.Stderr, "USAGE: %s -manifest FILE.exe.manifest [-o FILE.syso]\n"+
			"Generates a .syso file with specified resources embedded in .rsrc section,\n"+
			"aimed for consumption by Go linker when building Win32 excecutables.\n"+
			"OPTIONS:\n",
			os.Args[0])
		flags.PrintDefaults()
		os.Exit(1)
	}

	err := run(fnamein, fnameico, fnameout)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

type Coff struct {
	pe.FileHeader
	pe.SectionHeader32

	Dir
	DataEntries []DataEntry
	Data        []interface{}

	Relocations []RelocationEntry
	Symbols     []Symbol
	StringsHeader
}

func run(fnamein, fnameico, fnameout string) error {
	manifest, err := SizedOpen(fnamein)
	if err != nil {
		return fmt.Errorf("Error opening manifest file '%s': %s", fnamein, err)
	}
	defer manifest.Close()

	var icons []ico.ICONDIRENTRY
	var iconsf *os.File
	if fnameico != "" {
		iconsf, err = os.Open(fnameico)
		if err != nil {
			return err
		}
		defer iconsf.Close()
		icons, err = ico.DecodeHeaders(iconsf)
		if err != nil {
			return err
		}
	}

	newid := make(chan uint16)
	go func() {
		for i := uint16(1); ; i++ {
			newid <- i
		}
	}()

	out, err := os.Create(fnameout)
	if err != nil {
		return err
	}
	defer out.Close()
	w := Writer{W: out}

	coff := Coff{
		pe.FileHeader{
			Machine:              0x014c, //FIXME: find out how to differentiate this value, or maybe not necessary for Go
			NumberOfSections:     1,      // .rsrc
			TimeDateStamp:        0,      // was also 0 in sample data from MinGW's windres.exe
			NumberOfSymbols:      1,
			SizeOfOptionalHeader: 0,
			Characteristics:      0x0104, //FIXME: copied from windres.exe output, find out what should be here and why
		},
		pe.SectionHeader32{
			Name:                STRING_RSRC,
			NumberOfRelocations: 1,
			Characteristics:     0x40000040, // "INITIALIZED_DATA MEM_READ" ?
		},

		// now, build "directory hierarchy" of .rsrc section: first type, then id/name, then language
		Dir{
			NumberOfIdEntries: 1,
			DirEntries:        DirEntries{{NameOrId: RT_MANIFEST}},
			Dirs: Dirs{{
				NumberOfIdEntries: 1,
				DirEntries:        DirEntries{{NameOrId: uint32(<-newid)}}, // resource ID
				Dirs: Dirs{{
					NumberOfIdEntries: 1,
					DirEntries:        LANG_ENTRY,
				}},
			}},
		},
		[]DataEntry{
			DataEntry{
				Size1:    uint32(manifest.Size()),
				CodePage: 0, //FIXME: what value here? for now just tried 0
			},
		},
		[]interface{}{
			manifest,
		},

		[]RelocationEntry{RelocationEntry{
			SymbolIndex: 0, // "(zero based) index in the Symbol table to which the reference refers. Once you have loaded the COFF file into memory and know where each symbol is, you find the new updated address for the given symbol and update the reference accordingly."
			Type:        7, // according to ldpe.c, this decodes to: IMAGE_REL_I386_DIR32NB
		}},

		[]Symbol{Symbol{
			Name:           STRING_RSRC,
			Value:          0,
			SectionNumber:  1,
			Type:           0, // FIXME: wtf?
			StorageClass:   3, // FIXME: is it ok? and uint8? and what does the value mean?
			AuxiliaryCount: 0, // FIXME: wtf?
		}},

		StringsHeader{
			Length: uint32(binary.Size(StringsHeader{})), // empty strings table -- but we must still show size of the table's header...
		},
	}

	if len(icons) > 0 {
		//FIXME: add corresponding DataEntries
		//FIXME: add relocations
		/*
			coff.Dir.NumberOfIdEntries+=2

			coff.Dir.DirEntries = append(coff.Dir.DirEntries, DirEntry{NameOrId: RT_ICON})
			coff.Dir.Dirs = append(coff.Dir.Dirs, Dir{
				NumberOfIdEntries: len(icons),
			})
			dir := &coff.Dir.Dirs[len(coff.Dir.Dirs)-1]
			group := GRPICONDIR{
				Reserved: 0, // magic num.
				Type: 1, // magic num.
				Count: len(icons),
			}
			for i, icon := range icons {
				id := <-newid

				group.Entries = append(group.Entries, GRPICONDIRENTRY{icon, id})
				dir.DirEntries = append(dir.DirEntries, DirEntry{NameOrId: uint32(id)})
				dir.Dirs = append(dir.Dirs, Dir{
					NumberOfIdEntries: 1,
					DirEntries: LANG_ENTRY,
				})

				coff.Data = append(coff.Data, io.NewSectionReader(iconsf, icon.ImageOffset, icon.BytesInRes))
			}

			coff.Dir.DirEntries = append(coff.Dir.DirEntries, DirEntry{NameOrId: RT_GROUP_ICON})
			coff.Dir.Dirs = append(coff.Dir.Dirs, Dir{
				NumberOfIdEntries: 1,
				DirEntries: DirEntries{{NameOrId: uint32(<-newid)}},
				Dirs: Dirs{{
					NumberOfIdEntries: 1,
					DirEntries: LANG_ENTRY,
				}},
			})
			coff.Data = append(coff.Data, group)
		*/
	}

	leafwalker := make(chan *DirEntry)
	go func() {
		for _, dir1 := range coff.Dir.Dirs { // resource type
			for _, dir2 := range dir1.Dirs { // resource ID
				for i := range dir2.DirEntries { // resource lang
					leafwalker <- &dir2.DirEntries[i]
				}
			}
		}
	}()

	N := `\[(\d+)\]`
	dir_n := regexp.MustCompile("^/Dir/Dirs" + N + "$")
	dir_n_n := regexp.MustCompile("^/Dir/Dirs" + N + "/Dirs" + N + "$")
	dataentry_n := regexp.MustCompile("^/DataEntries" + N + "$")
	dataentry_n_off := regexp.MustCompile("^/DataEntries" + N + "/OffsetToData$")
	//data_n := regexp.MustCompile("^/Data" + N + "$")

	// fill in some important offsets in resulting file
	var offset, diroff uint32
	Walk(coff, func(v reflect.Value, path string) error {
		switch path {
		case "/Dir":
			coff.SectionHeader32.PointerToRawData = offset
			diroff = offset
		//case "/Dir/Dirs[0]":
		//	coff.Dir.DirEntries[0].OffsetToData = MASK_SUBDIRECTORY | (offset - diroff)
		//case "/Dir/Dirs[0]/Dirs[0]":
		//	coff.Dir.Dirs[0].DirEntries[0].OffsetToData = MASK_SUBDIRECTORY | (offset - diroff)
		//case "/DataEntries[0]":
		//	direntry := <-leafwalker
		//	direntry.OffsetToData = offset - diroff
		//case "/DataEntries[0]/OffsetToData":
		//	coff.Relocations[0].RVA = offset - diroff
		case "/Data[0]":
			coff.DataEntries[0].OffsetToData = offset - diroff
		case "/Relocations":
			coff.SectionHeader32.PointerToRelocations = offset
			coff.SectionHeader32.SizeOfRawData = offset - diroff
		case "/Symbols":
			coff.FileHeader.PointerToSymbolTable = offset
		}
		m := matcher{}
		switch {
		case m.Find(path, dir_n):
			coff.Dir.DirEntries[m[0]].OffsetToData = MASK_SUBDIRECTORY | (offset - diroff)
		case m.Find(path, dir_n_n):
			coff.Dir.Dirs[m[0]].DirEntries[m[1]].OffsetToData = MASK_SUBDIRECTORY | (offset - diroff)
		case m.Find(path, dataentry_n):
			direntry := <-leafwalker
			direntry.OffsetToData = offset - diroff
		case m.Find(path, dataentry_n_off):
			coff.Relocations[m[0]].RVA = offset - diroff
		}

		if Plain(v.Kind()) {
			offset += uint32(binary.Size(v.Interface())) // TODO: change to v.Type().Size() ?
			return nil
		}
		vv, ok := v.Interface().(SizedReader)
		if ok {
			offset += uint32(vv.Size())
			return WALK_SKIP
		}
		return nil
	})

	// write the resulting file to disk
	Walk(coff, func(v reflect.Value, path string) error {
		if Plain(v.Kind()) {
			w.WriteLE(v.Interface())
			return nil
		}
		vv, ok := v.Interface().(SizedReader)
		if ok {
			w.WriteFromSized(vv)
			return WALK_SKIP
		}
		return nil
	})

	if w.Err != nil {
		return fmt.Errorf("Error writing output file: %s", w.Err)
	}

	return nil
}

func Plain(kind reflect.Kind) bool {
	switch kind {
	case reflect.Bool, reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr, reflect.Float32, reflect.Float64, reflect.Complex64, reflect.Complex128:
		return true
	}
	return false
}

func MustAtoi(s string) int {
	i, err := strconv.Atoi(s)
	if err != nil {
		panic(err)
	}
	return i
}

type matcher []int

func (m *matcher) Find(s string, re *regexp.Regexp) bool {
	subs := re.FindStringSubmatch(s)
	if subs == nil {
		return false
	}

	*m = (*m)[:0]
	for i := 1; i < len(subs); i++ {
		*m = append(*m, MustAtoi(subs[i]))
	}
	return true
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/andybons/hipchat"
	"log"
	"time"
)

const (
	timeout time.Duration = 2
)

func SendMessage(msg string, config Config) {
	c := hipchat.Client{AuthToken: config.Token}
	request := hipchat.MessageRequest{
		RoomId:        config.Room,
		From:          config.From,
		Message:       msg,
		Color:         hipchat.ColorPurple,
		MessageFormat: hipchat.FormatText,
		Notify:        true,
	}

	errc := make(chan error, 1)
	okc := make(chan bool, 1)
	go func() {
		err := c.PostMessage(request)
		if err != nil {
			errc <- err
		}
		okc <- true
	}()

	select {
	case <-okc:
	case err := <-errc:
		log.Printf("Cannot send message: %s %q", msg, err)
	case <-time.After(time.Second * timeout):
		log.Printf("Timout after %d seconds", timeout)
	}
}
<commit_msg>Better error handling with else for ok channel<commit_after>package main

import (
	"log"
	"time"

	"github.com/andybons/hipchat"
)

const (
	timeout time.Duration = 2
)

func SendMessage(msg string, config Config) {
	c := hipchat.Client{AuthToken: config.Token}
	request := hipchat.MessageRequest{
		RoomId:        config.Room,
		From:          config.From,
		Message:       msg,
		Color:         hipchat.ColorPurple,
		MessageFormat: hipchat.FormatText,
		Notify:        true,
	}

	errc := make(chan error, 1)
	okc := make(chan bool, 1)
	go func() {
		err := c.PostMessage(request)
		if err != nil {
			errc <- err
		} else {
			okc <- true
		}
	}()

	select {
	case <-okc:
	case err := <-errc:
		log.Printf("Cannot send message: %s %q", msg, err)
	case <-time.After(time.Second * timeout):
		log.Printf("Timout after %d seconds", timeout)
	}
}
<|endoftext|>


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

package cmd

import (
	"fmt"
	goruntime "runtime"

	"github.com/spf13/cobra"
)

func newVersionCmd(c *Command) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "version",
		Short: "print CUE version",
		Long:  ``,
		RunE:  mkRunE(c, runVersion),
	}
	return cmd
}

// set by goreleaser or other builder using
// -ldflags='-X cuelang.org/go/cmd/cue/cmd.version=<version>'
var (
	version = "custom"
)

func runVersion(cmd *Command, args []string) error {
	w := cmd.OutOrStdout()
	fmt.Fprintf(w, "cue version %v %s/%s\n",
		version,
		goruntime.GOOS, goruntime.GOARCH,
	)
	return nil
}
<commit_msg>cmd/cue: use module information if no -X version information specified<commit_after>// Copyright 2019 CUE Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"fmt"
	goruntime "runtime"
	"runtime/debug"

	"github.com/spf13/cobra"
)

func newVersionCmd(c *Command) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "version",
		Short: "print CUE version",
		Long:  ``,
		RunE:  mkRunE(c, runVersion),
	}
	return cmd
}

const (
	defaultVersion = "devel"
)

// set by goreleaser or other builder using
// -ldflags='-X cuelang.org/go/cmd/cue/cmd.version=<version>'
var (
	version = defaultVersion
)

func runVersion(cmd *Command, args []string) error {
	w := cmd.OutOrStdout()
	if bi, ok := debug.ReadBuildInfo(); ok && version == defaultVersion {
		// No specific version provided via version
		version = bi.Main.Version
	}
	fmt.Fprintf(w, "cue version %v %s/%s\n",
		version,
		goruntime.GOOS, goruntime.GOARCH,
	)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"syscall"

	"github.com/facette/facette/pkg/config"
	"github.com/facette/facette/pkg/logger"
	"github.com/facette/facette/pkg/server"
	"github.com/facette/facette/pkg/utils"
)

const (
	cmdUsage = "Usage: %s [OPTIONS]"
)

var (
	flagConfig   string
	flagHelp     bool
	flagLog      string
	flagLogLevel string
	logLevel     int
	err          error
)

func init() {
	flag.StringVar(&flagConfig, "c", config.DefaultConfigFile, "configuration file path")
	flag.BoolVar(&flagHelp, "h", false, "display this help and exit")
	flag.StringVar(&flagLog, "l", config.DefaultLogFile, "log file path")
	flag.StringVar(&flagLogLevel, "L", config.DefaultLogLevel, "logging level")
	flag.Usage = func() { utils.PrintUsage(os.Stderr, cmdUsage) }
	flag.Parse()

	if flagHelp {
		utils.PrintUsage(os.Stdout, cmdUsage)
	} else if flagConfig == "" {
		fmt.Fprintf(os.Stderr, "Error: configuration file path is mandatory\n")
		utils.PrintUsage(os.Stderr, cmdUsage)
	}

	if logLevel, err = logger.GetLevelByName(flagLogLevel); err != nil {
		fmt.Fprintf(os.Stderr, "Error: invalid log level `%s'\n", flagLogLevel)
		os.Exit(1)
	}
}

func main() {
	// Create new server instance and load configuration
	instance := server.NewServer(flagConfig, flagLog, logLevel)

	// Reload server configuration on SIGHUP
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGHUP, syscall.SIGTERM)

	go func() {
		for sig := range sigChan {
			switch sig {
			case syscall.SIGHUP:
				instance.Reload(true)
				break

			case syscall.SIGINT, syscall.SIGTERM:
				instance.Stop()
				break
			}
		}
	}()

	// Run instance
	if err := instance.Run(); err != nil {
		fmt.Println("Error: " + err.Error())
		os.Exit(1)
	}

	os.Exit(0)
}
<commit_msg>Add `-V' flag to facette binary<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"syscall"

	"github.com/facette/facette/pkg/config"
	"github.com/facette/facette/pkg/logger"
	"github.com/facette/facette/pkg/server"
	"github.com/facette/facette/pkg/utils"
)

const (
	cmdUsage = "Usage: %s [OPTIONS]"
	version  = "0.0.0"
)

var (
	flagConfig   string
	flagHelp     bool
	flagLog      string
	flagLogLevel string
	flagVersion  bool
	logLevel     int
	err          error
)

func init() {
	flag.StringVar(&flagConfig, "c", config.DefaultConfigFile, "configuration file path")
	flag.BoolVar(&flagHelp, "h", false, "display this help and exit")
	flag.StringVar(&flagLog, "l", config.DefaultLogFile, "log file path")
	flag.StringVar(&flagLogLevel, "L", config.DefaultLogLevel, "logging level")
	flag.BoolVar(&flagVersion, "V", false, "display software version and exit")
	flag.Usage = func() { utils.PrintUsage(os.Stderr, cmdUsage) }
	flag.Parse()

	if flagHelp {
		utils.PrintUsage(os.Stdout, cmdUsage)
	} else if flagVersion {
		fmt.Printf("facette version %s\n", version)
		os.Exit(0)
	} else if flagConfig == "" {
		fmt.Fprintf(os.Stderr, "Error: configuration file path is mandatory\n")
		utils.PrintUsage(os.Stderr, cmdUsage)
	}

	if logLevel, err = logger.GetLevelByName(flagLogLevel); err != nil {
		fmt.Fprintf(os.Stderr, "Error: invalid log level `%s'\n", flagLogLevel)
		os.Exit(1)
	}
}

func main() {
	// Create new server instance and load configuration
	instance := server.NewServer(flagConfig, flagLog, logLevel)

	// Reload server configuration on SIGHUP
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGHUP, syscall.SIGTERM)

	go func() {
		for sig := range sigChan {
			switch sig {
			case syscall.SIGHUP:
				instance.Reload(true)
				break

			case syscall.SIGINT, syscall.SIGTERM:
				instance.Stop()
				break
			}
		}
	}()

	// Run instance
	if err := instance.Run(); err != nil {
		fmt.Println("Error: " + err.Error())
		os.Exit(1)
	}

	os.Exit(0)
}
<|endoftext|>


<commit_before>// Copyright (c) 2015, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"fmt"
	"log"
	"regexp"
	"strings"

	"github.com/mvdan/fdroidcl"
	"github.com/mvdan/fdroidcl/adb"
)

var cmdSearch = &Command{
	UsageLine: "search <regexp...>",
	Short:     "Search available apps",
}

var (
	quiet     = cmdSearch.Flag.Bool("q", false, "Print package names only")
	installed = cmdSearch.Flag.Bool("i", false, "Filter installed apps")
	updates   = cmdSearch.Flag.Bool("u", false, "Filter apps with updates")
)

func init() {
	cmdSearch.Run = runSearch
}

func runSearch(args []string) {
	if *installed && *updates {
		fmt.Println("-i is redundant if -u is specified")
		cmdSearch.Flag.Usage()
	}
	var device *adb.Device
	if *installed || *updates {
		device = mustOneDevice()
	}
	apps := filterAppsSearch(mustLoadIndexes(), args)
	instPkgs := mustInstalled(device)
	if *installed {
		apps = filterAppsInstalled(apps, instPkgs)
	}
	if *updates {
		apps = filterAppsUpdates(apps, instPkgs)
	}
	if *quiet {
		for _, app := range apps {
			fmt.Println(app.ID)
		}
	} else {
		printApps(apps, instPkgs)
	}
}

func filterAppsSearch(apps []fdroidcl.App, terms []string) []fdroidcl.App {
	regexes := make([]*regexp.Regexp, len(terms))
	for i, term := range terms {
		regexes[i] = regexp.MustCompile(term)
	}
	var result []fdroidcl.App
	for _, app := range apps {
		fields := []string{
			strings.ToLower(app.ID),
			strings.ToLower(app.Name),
			strings.ToLower(app.Summary),
			strings.ToLower(app.Desc),
		}
		if !appMatches(fields, regexes) {
			continue
		}
		result = append(result, app)
	}
	return result
}

func appMatches(fields []string, regexes []*regexp.Regexp) bool {
fieldLoop:
	for _, field := range fields {
		for _, regex := range regexes {
			if !regex.MatchString(field) {
				continue fieldLoop
			}
		}
		return true
	}
	return false
}

func printApps(apps []fdroidcl.App, inst map[string]adb.Package) {
	maxIDLen := 0
	for _, app := range apps {
		if len(app.ID) > maxIDLen {
			maxIDLen = len(app.ID)
		}
	}
	for _, app := range apps {
		var pkg *adb.Package
		p, e := inst[app.ID]
		if e {
			pkg = &p
		}
		printApp(app, maxIDLen, pkg)
	}
}

func descVersion(app fdroidcl.App, inst *adb.Package) string {
	cur := app.CurApk()
	if cur == nil {
		return "(no version available)"
	}
	if inst == nil {
		return fmt.Sprintf("%s (%d)", cur.VName, cur.VCode)
	}
	if !*installed && inst.VCode >= cur.VCode {
		return fmt.Sprintf("%s (%d) [installed]", cur.VName, cur.VCode)
	}
	return fmt.Sprintf("%s (%d) -> %s (%d)", inst.VName, inst.VCode,
		cur.VName, cur.VCode)
}

func printApp(app fdroidcl.App, IDLen int, inst *adb.Package) {
	fmt.Printf("%s%s %s - %s\n", app.ID, strings.Repeat(" ", IDLen-len(app.ID)),
		app.Name, descVersion(app, inst))
	fmt.Printf("    %s\n", app.Summary)
}

func mustInstalled(device *adb.Device) map[string]adb.Package {
	if device == nil {
		return nil
	}
	inst, err := device.Installed()
	if err != nil {
		log.Fatalf("Could not get installed packages: %v", err)
	}
	return inst
}

func filterAppsInstalled(apps []fdroidcl.App, inst map[string]adb.Package) []fdroidcl.App {
	var result []fdroidcl.App
	for _, app := range apps {
		if _, e := inst[app.ID]; !e {
			continue
		}
		result = append(result, app)
	}
	return result
}

func filterAppsUpdates(apps []fdroidcl.App, inst map[string]adb.Package) []fdroidcl.App {
	var result []fdroidcl.App
	for _, app := range apps {
		p, e := inst[app.ID]
		if !e {
			continue
		}
		cur := app.CurApk()
		if cur == nil {
			continue
		}
		if p.VCode >= cur.VCode {
			continue
		}
		result = append(result, app)
	}
	return result
}
<commit_msg>search: add -o for sort order<commit_after>// Copyright (c) 2015, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"fmt"
	"log"
	"os"
	"regexp"
	"sort"
	"strings"

	"github.com/mvdan/fdroidcl"
	"github.com/mvdan/fdroidcl/adb"
)

var cmdSearch = &Command{
	UsageLine: "search <regexp...>",
	Short:     "Search available apps",
}

var (
	quiet     = cmdSearch.Flag.Bool("q", false, "Print package names only")
	installed = cmdSearch.Flag.Bool("i", false, "Filter installed apps")
	updates   = cmdSearch.Flag.Bool("u", false, "Filter apps with updates")
	sortBy    = cmdSearch.Flag.String("o", "", "Sort order (added)")
)

func init() {
	cmdSearch.Run = runSearch
}

func runSearch(args []string) {
	if *installed && *updates {
		fmt.Fprintf(os.Stderr, "-i is redundant if -u is specified\n")
		cmdSearch.Flag.Usage()
	}
	sfunc, err := sortFunc(*sortBy)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		cmdSearch.Flag.Usage()
	}
	var device *adb.Device
	if *installed || *updates {
		device = mustOneDevice()
	}
	apps := filterAppsSearch(mustLoadIndexes(), args)
	instPkgs := mustInstalled(device)
	if *installed {
		apps = filterAppsInstalled(apps, instPkgs)
	}
	if *updates {
		apps = filterAppsUpdates(apps, instPkgs)
	}
	if sfunc != nil {
		apps = sortApps(apps, sfunc)
	}
	if *quiet {
		for _, app := range apps {
			fmt.Println(app.ID)
		}
	} else {
		printApps(apps, instPkgs)
	}
}

func filterAppsSearch(apps []fdroidcl.App, terms []string) []fdroidcl.App {
	regexes := make([]*regexp.Regexp, len(terms))
	for i, term := range terms {
		regexes[i] = regexp.MustCompile(term)
	}
	var result []fdroidcl.App
	for _, app := range apps {
		fields := []string{
			strings.ToLower(app.ID),
			strings.ToLower(app.Name),
			strings.ToLower(app.Summary),
			strings.ToLower(app.Desc),
		}
		if !appMatches(fields, regexes) {
			continue
		}
		result = append(result, app)
	}
	return result
}

func appMatches(fields []string, regexes []*regexp.Regexp) bool {
fieldLoop:
	for _, field := range fields {
		for _, regex := range regexes {
			if !regex.MatchString(field) {
				continue fieldLoop
			}
		}
		return true
	}
	return false
}

func printApps(apps []fdroidcl.App, inst map[string]adb.Package) {
	maxIDLen := 0
	for _, app := range apps {
		if len(app.ID) > maxIDLen {
			maxIDLen = len(app.ID)
		}
	}
	for _, app := range apps {
		var pkg *adb.Package
		p, e := inst[app.ID]
		if e {
			pkg = &p
		}
		printApp(app, maxIDLen, pkg)
	}
}

func descVersion(app fdroidcl.App, inst *adb.Package) string {
	cur := app.CurApk()
	if cur == nil {
		return "(no version available)"
	}
	if inst == nil {
		return fmt.Sprintf("%s (%d)", cur.VName, cur.VCode)
	}
	if !*installed && inst.VCode >= cur.VCode {
		return fmt.Sprintf("%s (%d) [installed]", cur.VName, cur.VCode)
	}
	return fmt.Sprintf("%s (%d) -> %s (%d)", inst.VName, inst.VCode,
		cur.VName, cur.VCode)
}

func printApp(app fdroidcl.App, IDLen int, inst *adb.Package) {
	fmt.Printf("%s%s %s - %s\n", app.ID, strings.Repeat(" ", IDLen-len(app.ID)),
		app.Name, descVersion(app, inst))
	fmt.Printf("    %s\n", app.Summary)
}

func mustInstalled(device *adb.Device) map[string]adb.Package {
	if device == nil {
		return nil
	}
	inst, err := device.Installed()
	if err != nil {
		log.Fatalf("Could not get installed packages: %v", err)
	}
	return inst
}

func filterAppsInstalled(apps []fdroidcl.App, inst map[string]adb.Package) []fdroidcl.App {
	var result []fdroidcl.App
	for _, app := range apps {
		if _, e := inst[app.ID]; !e {
			continue
		}
		result = append(result, app)
	}
	return result
}

func filterAppsUpdates(apps []fdroidcl.App, inst map[string]adb.Package) []fdroidcl.App {
	var result []fdroidcl.App
	for _, app := range apps {
		p, e := inst[app.ID]
		if !e {
			continue
		}
		cur := app.CurApk()
		if cur == nil {
			continue
		}
		if p.VCode >= cur.VCode {
			continue
		}
		result = append(result, app)
	}
	return result
}

func cmpAdded(a, b *fdroidcl.App) bool {
	return a.Added.Before(b.Added.Time)
}

func sortFunc(sortBy string) (func(a, b *fdroidcl.App) bool, error) {
	switch sortBy {
	case "added":
		return cmpAdded, nil
	case "":
		return nil, nil
	}
	return nil, fmt.Errorf("Unknown sort order: %s", sortBy)
}

type appList struct {
	l []fdroidcl.App
	f func(a, b *fdroidcl.App) bool
}

func (al appList) Len() int           { return len(al.l) }
func (al appList) Swap(i, j int)      { al.l[i], al.l[j] = al.l[j], al.l[i] }
func (al appList) Less(i, j int) bool { return al.f(&al.l[i], &al.l[j]) }

func sortApps(apps []fdroidcl.App, f func(a, b *fdroidcl.App) bool) []fdroidcl.App {
	sort.Sort(appList{l: apps, f: f})
	return apps
}
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-metrics"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/golib/signal"
)

var (
	stats *peekStats
)

type peekStats struct {
	MsgCountPerSecond metrics.Meter
	MsgBytesPerSecond metrics.Meter
}

func newPeekStats() *peekStats {
	this := &peekStats{
		MsgCountPerSecond: metrics.NewMeter(),
		MsgBytesPerSecond: metrics.NewMeter(),
	}

	metrics.Register("msg.count.per.second", this.MsgCountPerSecond)
	metrics.Register("msg.bytes.per.second", this.MsgBytesPerSecond)
	return this
}

func (this *peekStats) start() {
	metrics.Log(metrics.DefaultRegistry, time.Second*10,
		log.New(os.Stdout, "metrics: ", log.Lmicroseconds))
}

type Peek struct {
	Ui  cli.Ui
	Cmd string

	offset   int64
	lastN    int64 // peek the most recent N messages
	colorize bool
	limit    int
	quit     chan struct{}
	once     sync.Once
}

func (this *Peek) Run(args []string) (exitCode int) {
	var (
		cluster      string
		zone         string
		topicPattern string
		partitionId  int
		silence      bool
	)
	cmdFlags := flag.NewFlagSet("peek", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&cluster, "c", "", "")
	cmdFlags.StringVar(&topicPattern, "t", "", "")
	cmdFlags.IntVar(&partitionId, "p", 0, "")
	cmdFlags.BoolVar(&this.colorize, "color", true, "")
	cmdFlags.Int64Var(&this.lastN, "last", -1, "")
	cmdFlags.IntVar(&this.limit, "limit", -1, "")
	cmdFlags.Int64Var(&this.offset, "offset", sarama.OffsetNewest, "")
	cmdFlags.BoolVar(&silence, "s", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	this.quit = make(chan struct{})

	if silence {
		stats := newPeekStats()
		go stats.start()
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	msgChan := make(chan *sarama.ConsumerMessage, 20000) // msg aggerator channel
	if cluster == "" {
		zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
			this.consumeCluster(zkcluster, topicPattern, partitionId, msgChan)
		})
	} else {
		zkcluster := zkzone.NewCluster(cluster)
		this.consumeCluster(zkcluster, topicPattern, partitionId, msgChan)
	}

	signal.RegisterSignalsHandler(func(sig os.Signal) {
		log.Printf("received signal: %s", strings.ToUpper(sig.String()))
		log.Println("quiting...")

		this.once.Do(func() {
			close(this.quit)
		})
	}, syscall.SIGINT, syscall.SIGTERM)

	var (
		startAt = time.Now()
		msg     *sarama.ConsumerMessage
		total   int
	)
LOOP:
	for {
		select {
		case <-this.quit:
			this.Ui.Output(fmt.Sprintf("Total: %d msgs", total))
			this.Ui.Output(fmt.Sprintf("Throughput: %d/s", total/int(time.Since(startAt).Seconds())))

		case msg = <-msgChan:
			if silence {
				stats.MsgCountPerSecond.Mark(1)
				stats.MsgBytesPerSecond.Mark(int64(len(msg.Value)))
			} else {
				if this.colorize {
					this.Ui.Output(fmt.Sprintf("%s/%d %s k:%s, v:%s",
						color.Green(msg.Topic), msg.Partition,
						gofmt.Comma(msg.Offset), string(msg.Key), string(msg.Value)))
				} else {
					// colored UI will have invisible chars output
					fmt.Println(fmt.Sprintf("%s/%d %s k:%s, v:%s",
						msg.Topic, msg.Partition,
						gofmt.Comma(msg.Offset), string(msg.Key), string(msg.Value)))
				}
			}

			total++

			if this.limit > 0 && total >= this.limit {
				break LOOP

			}
			if this.lastN > 0 && total >= int(this.lastN) {
				break LOOP
			}

		}
	}

	return
}

func (this *Peek) consumeCluster(zkcluster *zk.ZkCluster, topicPattern string,
	partitionId int, msgChan chan *sarama.ConsumerMessage) {
	brokerList := zkcluster.BrokerList()
	if len(brokerList) == 0 {
		return
	}
	kfk, err := sarama.NewClient(brokerList, sarama.NewConfig())
	if err != nil {
		this.Ui.Output(err.Error())
		return
	}
	//defer kfk.Close() // FIXME how to close it

	topics, err := kfk.Topics()
	if err != nil {
		this.Ui.Output(err.Error())
		return
	}

	for _, t := range topics {
		if patternMatched(t, topicPattern) {
			go this.simpleConsumeTopic(kfk, t, int32(partitionId), msgChan)
		}
	}

}

func (this *Peek) simpleConsumeTopic(kfk sarama.Client, topic string, partitionId int32,
	msgCh chan *sarama.ConsumerMessage) {
	consumer, err := sarama.NewConsumerFromClient(kfk)
	if err != nil {
		panic(err)
	}
	defer consumer.Close()

	if partitionId == -1 {
		// all partitions
		partitions, err := kfk.Partitions(topic)
		if err != nil {
			panic(err)
		}

		for _, p := range partitions {
			offset := this.offset
			if this.lastN > 0 {
				latestOffset, err := kfk.GetOffset(topic, p, sarama.OffsetNewest)
				swallow(err)

				oldestOffset, err := kfk.GetOffset(topic, p, sarama.OffsetOldest)
				swallow(err)

				offset = latestOffset - this.lastN
				if offset < oldestOffset {
					offset = oldestOffset
				}

				if offset == 0 {
					// no message in store
					return
				}
			}

			go this.consumePartition(kfk, consumer, topic, p, msgCh, offset)
		}

	} else {
		offset := this.offset
		if this.lastN > 0 {
			latestOffset, err := kfk.GetOffset(topic, partitionId, sarama.OffsetNewest)
			swallow(err)
			offset = latestOffset - this.lastN
			if offset < 0 {
				offset = sarama.OffsetOldest
			}
		}
		this.consumePartition(kfk, consumer, topic, partitionId, msgCh, offset)
	}

}

func (this *Peek) consumePartition(kfk sarama.Client, consumer sarama.Consumer,
	topic string, partitionId int32, msgCh chan *sarama.ConsumerMessage, offset int64) {
	p, err := consumer.ConsumePartition(topic, partitionId, offset)
	if err != nil {
		this.Ui.Error(fmt.Sprintf("%s/%d: offset=%d %v", topic, partitionId, offset, err))
		return
	}
	defer p.Close()

	n := int64(0)
	for {
		select {
		case <-this.quit:
			return

		case msg := <-p.Messages():
			msgCh <- msg

			n++
			if this.lastN > 0 && n >= this.lastN {
				return
			}
		}
	}
}

func (*Peek) Synopsis() string {
	return "Peek kafka cluster messages ongoing from any offset"
}

func (this *Peek) Help() string {
	help := fmt.Sprintf(`
Usage: %s peek [options]

    Peek kafka cluster messages ongoing from any offset

Options:

    -z zone
      Default %s

    -c cluster

    -t topic pattern

    -p partition id
      -1 will peek all partitions of a topic

    -last n
      Peek the most recent N messages

    -offset message offset value
      -1 OffsetNewest, -2 OffsetOldest.
      You can specify your own offset.
      Default -1(OffsetNewest)

    -limit n
      Limit how many messages to consume

    -s
      Silence mode, only display statastics instead of message content

    -color
      Enable colorized output
`, this.Cmd, ctx.ZkDefaultZone())
	return strings.TrimSpace(help)
}
<commit_msg>fix the stupid bug: didn't stop<commit_after>package command

import (
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-metrics"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/golib/signal"
)

var (
	stats *peekStats
)

type peekStats struct {
	MsgCountPerSecond metrics.Meter
	MsgBytesPerSecond metrics.Meter
}

func newPeekStats() *peekStats {
	this := &peekStats{
		MsgCountPerSecond: metrics.NewMeter(),
		MsgBytesPerSecond: metrics.NewMeter(),
	}

	metrics.Register("msg.count.per.second", this.MsgCountPerSecond)
	metrics.Register("msg.bytes.per.second", this.MsgBytesPerSecond)
	return this
}

func (this *peekStats) start() {
	metrics.Log(metrics.DefaultRegistry, time.Second*10,
		log.New(os.Stdout, "metrics: ", log.Lmicroseconds))
}

type Peek struct {
	Ui  cli.Ui
	Cmd string

	offset   int64
	lastN    int64 // peek the most recent N messages
	colorize bool
	limit    int
	quit     chan struct{}
	once     sync.Once
}

func (this *Peek) Run(args []string) (exitCode int) {
	var (
		cluster      string
		zone         string
		topicPattern string
		partitionId  int
		silence      bool
	)
	cmdFlags := flag.NewFlagSet("peek", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&cluster, "c", "", "")
	cmdFlags.StringVar(&topicPattern, "t", "", "")
	cmdFlags.IntVar(&partitionId, "p", 0, "")
	cmdFlags.BoolVar(&this.colorize, "color", true, "")
	cmdFlags.Int64Var(&this.lastN, "last", -1, "")
	cmdFlags.IntVar(&this.limit, "limit", -1, "")
	cmdFlags.Int64Var(&this.offset, "offset", sarama.OffsetNewest, "")
	cmdFlags.BoolVar(&silence, "s", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	this.quit = make(chan struct{})

	if silence {
		stats := newPeekStats()
		go stats.start()
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	msgChan := make(chan *sarama.ConsumerMessage, 20000) // msg aggerator channel
	if cluster == "" {
		zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
			this.consumeCluster(zkcluster, topicPattern, partitionId, msgChan)
		})
	} else {
		zkcluster := zkzone.NewCluster(cluster)
		this.consumeCluster(zkcluster, topicPattern, partitionId, msgChan)
	}

	signal.RegisterSignalsHandler(func(sig os.Signal) {
		log.Printf("received signal: %s", strings.ToUpper(sig.String()))
		log.Println("quiting...")

		this.once.Do(func() {
			close(this.quit)
		})
	}, syscall.SIGINT, syscall.SIGTERM)

	var (
		startAt = time.Now()
		msg     *sarama.ConsumerMessage
		total   int
	)

LOOP:
	for {
		select {
		case <-this.quit:
			this.Ui.Output(fmt.Sprintf("Total: %d msgs", total))
			this.Ui.Output(fmt.Sprintf("Throughput: %d/s", total/int(time.Since(startAt).Seconds())))
			return

		case msg = <-msgChan:
			if silence {
				stats.MsgCountPerSecond.Mark(1)
				stats.MsgBytesPerSecond.Mark(int64(len(msg.Value)))
			} else {
				if this.colorize {
					this.Ui.Output(fmt.Sprintf("%s/%d %s k:%s, v:%s",
						color.Green(msg.Topic), msg.Partition,
						gofmt.Comma(msg.Offset), string(msg.Key), string(msg.Value)))
				} else {
					// colored UI will have invisible chars output
					fmt.Println(fmt.Sprintf("%s/%d %s k:%s, v:%s",
						msg.Topic, msg.Partition,
						gofmt.Comma(msg.Offset), string(msg.Key), string(msg.Value)))
				}
			}

			total++

			if this.limit > 0 && total >= this.limit {
				break LOOP

			}
			if this.lastN > 0 && total >= int(this.lastN) {
				break LOOP
			}

		}
	}

	return
}

func (this *Peek) consumeCluster(zkcluster *zk.ZkCluster, topicPattern string,
	partitionId int, msgChan chan *sarama.ConsumerMessage) {
	brokerList := zkcluster.BrokerList()
	if len(brokerList) == 0 {
		return
	}
	kfk, err := sarama.NewClient(brokerList, sarama.NewConfig())
	if err != nil {
		this.Ui.Output(err.Error())
		return
	}
	//defer kfk.Close() // FIXME how to close it

	topics, err := kfk.Topics()
	if err != nil {
		this.Ui.Output(err.Error())
		return
	}

	for _, t := range topics {
		if patternMatched(t, topicPattern) {
			go this.simpleConsumeTopic(kfk, t, int32(partitionId), msgChan)
		}
	}

}

func (this *Peek) simpleConsumeTopic(kfk sarama.Client, topic string, partitionId int32,
	msgCh chan *sarama.ConsumerMessage) {
	consumer, err := sarama.NewConsumerFromClient(kfk)
	if err != nil {
		panic(err)
	}
	defer consumer.Close()

	if partitionId == -1 {
		// all partitions
		partitions, err := kfk.Partitions(topic)
		if err != nil {
			panic(err)
		}

		for _, p := range partitions {
			offset := this.offset
			if this.lastN > 0 {
				latestOffset, err := kfk.GetOffset(topic, p, sarama.OffsetNewest)
				swallow(err)

				oldestOffset, err := kfk.GetOffset(topic, p, sarama.OffsetOldest)
				swallow(err)

				offset = latestOffset - this.lastN
				if offset < oldestOffset {
					offset = oldestOffset
				}

				if offset == 0 {
					// no message in store
					return
				}
			}

			go this.consumePartition(kfk, consumer, topic, p, msgCh, offset)
		}

	} else {
		offset := this.offset
		if this.lastN > 0 {
			latestOffset, err := kfk.GetOffset(topic, partitionId, sarama.OffsetNewest)
			swallow(err)
			offset = latestOffset - this.lastN
			if offset < 0 {
				offset = sarama.OffsetOldest
			}
		}
		this.consumePartition(kfk, consumer, topic, partitionId, msgCh, offset)
	}

}

func (this *Peek) consumePartition(kfk sarama.Client, consumer sarama.Consumer,
	topic string, partitionId int32, msgCh chan *sarama.ConsumerMessage, offset int64) {
	p, err := consumer.ConsumePartition(topic, partitionId, offset)
	if err != nil {
		this.Ui.Error(fmt.Sprintf("%s/%d: offset=%d %v", topic, partitionId, offset, err))
		return
	}
	defer p.Close()

	n := int64(0)
	for {
		select {
		case <-this.quit:
			return

		case msg := <-p.Messages():
			msgCh <- msg

			n++
			if this.lastN > 0 && n >= this.lastN {
				return
			}
		}
	}
}

func (*Peek) Synopsis() string {
	return "Peek kafka cluster messages ongoing from any offset"
}

func (this *Peek) Help() string {
	help := fmt.Sprintf(`
Usage: %s peek [options]

    Peek kafka cluster messages ongoing from any offset

Options:

    -z zone
      Default %s

    -c cluster

    -t topic pattern

    -p partition id
      -1 will peek all partitions of a topic

    -last n
      Peek the most recent N messages

    -offset message offset value
      -1 OffsetNewest, -2 OffsetOldest.
      You can specify your own offset.
      Default -1(OffsetNewest)

    -limit n
      Limit how many messages to consume

    -s
      Silence mode, only display statastics instead of message content

    -color
      Enable colorized output
`, this.Cmd, ctx.ZkDefaultZone())
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	log "github.com/funkygao/log4go"
)

type Ping struct {
	Ui  cli.Ui
	Cmd string

	zone            string
	zkzone          *zk.ZkZone
	logfile         string
	problematicMode bool
	interval        time.Duration
}

// TODO run 3 nodes in a zone to monitor as daemon
// register the 3 nodes as host service tag.
func (this *Ping) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("ping", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", "", "")
	cmdFlags.DurationVar(&this.interval, "interval", time.Minute*5, "")
	cmdFlags.StringVar(&this.logfile, "logfile", "stdout", "")
	cmdFlags.BoolVar(&this.problematicMode, "p", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		require("-z").
		invalid(args) {
		return 2
	}

	this.setupLog()
	this.zkzone = zk.NewZkZone(zk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))

	for {
		this.diagnose()
		if this.logfile == "stdout" {
			break
		}

		time.Sleep(this.interval)
	}

	return
}

func (this *Ping) setupLog() {
	if this.logfile != "stdout" {
		log.DeleteFilter("stdout")

		filer := log.NewFileLogWriter(this.logfile, true, false, 0)
		filer.SetFormat("[%d %T] [%L] (%S) %M")
		filer.SetRotateSize(0)
		filer.SetRotateLines(0)
		filer.SetRotateDaily(true)
		log.AddFilter("file", log.DEBUG, filer)
	}

}

func (this *Ping) diagnose() {
	this.zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
		registeredBrokers := zkcluster.RegisteredInfo().Roster
		for _, broker := range registeredBrokers {
			log.Debug("ping %s", broker.Addr())

			kfk, err := sarama.NewClient([]string{broker.Addr()}, sarama.NewConfig())
			if err != nil {
				log.Error("%25s %30s %s", broker.Addr(), broker.NamedAddr(), color.Red(err.Error()))

				continue
			}

			_, err = kfk.Topics() // kafka didn't provide ping, so use Topics() as ping
			if err != nil {
				log.Error("%25s %30s %s", broker.Addr(), broker.NamedAddr(), color.Red(err.Error()))
			} else {
				if !this.problematicMode {
					log.Info("%25s %30s %s", broker.Addr(), broker.NamedAddr(), color.Green("ok"))
				}
			}
			kfk.Close()
		}
	})

}

func (*Ping) Synopsis() string {
	return "Ping liveness of all registered brokers in a zone"
}

func (this *Ping) Help() string {
	help := fmt.Sprintf(`
Usage: %s ping -z zone [options]

    Ping liveness of all registered brokers in a zone

Options:

    -p
      Only show problematic brokers

    -interval duration
      Defaults 5m

    -logfile filename
      Defaults stdout in current directory

`, this.Cmd)
	return strings.TrimSpace(help)
}
<commit_msg>ping has default zone<commit_after>package command

import (
	"flag"
	"fmt"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	log "github.com/funkygao/log4go"
)

type Ping struct {
	Ui  cli.Ui
	Cmd string

	zone            string
	zkzone          *zk.ZkZone
	logfile         string
	problematicMode bool
	interval        time.Duration
}

// TODO run 3 nodes in a zone to monitor as daemon
// register the 3 nodes as host service tag.
func (this *Ping) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("ping", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.DurationVar(&this.interval, "interval", time.Minute*5, "")
	cmdFlags.StringVar(&this.logfile, "logfile", "stdout", "")
	cmdFlags.BoolVar(&this.problematicMode, "p", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		require("-z").
		invalid(args) {
		return 2
	}

	this.setupLog()
	this.zkzone = zk.NewZkZone(zk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))

	for {
		this.diagnose()
		if this.logfile == "stdout" {
			break
		}

		time.Sleep(this.interval)
	}

	return
}

func (this *Ping) setupLog() {
	if this.logfile != "stdout" {
		log.DeleteFilter("stdout")

		filer := log.NewFileLogWriter(this.logfile, true, false, 0)
		filer.SetFormat("[%d %T] [%L] (%S) %M")
		filer.SetRotateSize(0)
		filer.SetRotateLines(0)
		filer.SetRotateDaily(true)
		log.AddFilter("file", log.DEBUG, filer)
	}

}

func (this *Ping) diagnose() {
	this.zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
		registeredBrokers := zkcluster.RegisteredInfo().Roster
		for _, broker := range registeredBrokers {
			log.Debug("ping %s", broker.Addr())

			kfk, err := sarama.NewClient([]string{broker.Addr()}, sarama.NewConfig())
			if err != nil {
				log.Error("%25s %30s %s", broker.Addr(), broker.NamedAddr(), color.Red(err.Error()))

				continue
			}

			_, err = kfk.Topics() // kafka didn't provide ping, so use Topics() as ping
			if err != nil {
				log.Error("%25s %30s %s", broker.Addr(), broker.NamedAddr(), color.Red(err.Error()))
			} else {
				if !this.problematicMode {
					log.Info("%25s %30s %s", broker.Addr(), broker.NamedAddr(), color.Green("ok"))
				}
			}
			kfk.Close()
		}
	})

}

func (*Ping) Synopsis() string {
	return "Ping liveness of all registered brokers in a zone"
}

func (this *Ping) Help() string {
	help := fmt.Sprintf(`
Usage: %s ping -z zone [options]

    Ping liveness of all registered brokers in a zone

Options:

    -p
      Only show problematic brokers

    -interval duration
      Defaults 5m

    -logfile filename
      Defaults stdout in current directory

`, this.Cmd)
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/version"
	"os"
	"path/filepath"
	"time"
)

type UnitSuite struct {
	testing.JujuConnSuite
}

var _ = Suite(&UnitSuite{})

func (s *UnitSuite) TestParseSuccess(c *C) {
	create := func() (cmd.Command, *AgentConf) {
		a := &UnitAgent{}
		return a, &a.Conf
	}
	uc := CheckAgentCommand(c, create, []string{"--unit-name", "w0rd-pre55/1"})
	c.Assert(uc.(*UnitAgent).UnitName, Equals, "w0rd-pre55/1")
}

func (s *UnitSuite) TestParseMissing(c *C) {
	uc := &UnitAgent{}
	err := ParseAgentCommand(uc, []string{})
	c.Assert(err, ErrorMatches, "--unit-name option must be set")
}

func (s *UnitSuite) TestParseNonsense(c *C) {
	for _, args := range [][]string{
		[]string{"--unit-name", "wordpress"},
		[]string{"--unit-name", "wordpress/seventeen"},
		[]string{"--unit-name", "wordpress/-32"},
		[]string{"--unit-name", "wordpress/wild/9"},
		[]string{"--unit-name", "20/20"},
	} {
		err := ParseAgentCommand(&UnitAgent{}, args)
		c.Assert(err, ErrorMatches, `--unit-name option expects "<service>/<n>" argument`)
	}
}

func (s *UnitSuite) TestParseUnknown(c *C) {
	uc := &UnitAgent{}
	err := ParseAgentCommand(uc, []string{"--unit-name", "wordpress/1", "thundering typhoons"})
	c.Assert(err, ErrorMatches, `unrecognized args: \["thundering typhoons"\]`)
}

func (s *UnitSuite) TestRunStop(c *C) {
	// Set up state.
	ch := s.AddTestingCharm(c, "dummy")
	svc, err := s.Conn.AddService("dummy", ch)
	c.Assert(err, IsNil)
	unit, err := svc.AddUnit()
	c.Assert(err, IsNil)

	// Set up local environment.
	jujuDir := environs.VarDir
	vers := version.Current.String()
	toolsDir := filepath.Join(jujuDir, "tools", vers)
	err = os.MkdirAll(toolsDir, 0755)
	c.Assert(err, IsNil)
	err = ioutil.WriteFile(filepath.Join(toolsDir, "jujuc"), nil, 0644)
	c.Assert(err, IsNil)
	toolsLink := filepath.Join(jujuDir, "tools", "unit-dummy-0")
	err = os.Symlink(vers, toolsLink)
	c.Assert(err, IsNil)

	// Run a unit agent.
	a := &UnitAgent{
		Conf:     AgentConf{jujuDir, *s.StateInfo(c)},
		UnitName: unit.Name(),
	}
	done := make(chan error)
	go func() {
		done <- a.Run(nil)
	}()
	defer a.Stop()
	timeout := time.After(2000 * time.Millisecond)
	for {
		select {
		case <-timeout:
			c.Fatalf("no activity detected")
		case <-time.After(200 * time.Millisecond):
			st, info, err := unit.Status()
			c.Assert(err, IsNil)
			switch st {
			case state.UnitPending, state.UnitInstalled:
				c.Logf("waiting...")
				continue
			case state.UnitStarted:
				c.Logf("started!")
			default:
				c.Fatalf("unexpected status %s %s", st, info)
			}
		}
		break
	}
	err = a.Stop()
	c.Assert(err, IsNil)
	c.Assert(<-done, IsNil)
}
<commit_msg>address review<commit_after>package main

import (
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/version"
	"os"
	"path/filepath"
	"time"
)

type UnitSuite struct {
	testing.JujuConnSuite
}

var _ = Suite(&UnitSuite{})

func (s *UnitSuite) TestParseSuccess(c *C) {
	create := func() (cmd.Command, *AgentConf) {
		a := &UnitAgent{}
		return a, &a.Conf
	}
	uc := CheckAgentCommand(c, create, []string{"--unit-name", "w0rd-pre55/1"})
	c.Assert(uc.(*UnitAgent).UnitName, Equals, "w0rd-pre55/1")
}

func (s *UnitSuite) TestParseMissing(c *C) {
	uc := &UnitAgent{}
	err := ParseAgentCommand(uc, []string{})
	c.Assert(err, ErrorMatches, "--unit-name option must be set")
}

func (s *UnitSuite) TestParseNonsense(c *C) {
	for _, args := range [][]string{
		[]string{"--unit-name", "wordpress"},
		[]string{"--unit-name", "wordpress/seventeen"},
		[]string{"--unit-name", "wordpress/-32"},
		[]string{"--unit-name", "wordpress/wild/9"},
		[]string{"--unit-name", "20/20"},
	} {
		err := ParseAgentCommand(&UnitAgent{}, args)
		c.Assert(err, ErrorMatches, `--unit-name option expects "<service>/<n>" argument`)
	}
}

func (s *UnitSuite) TestParseUnknown(c *C) {
	uc := &UnitAgent{}
	err := ParseAgentCommand(uc, []string{"--unit-name", "wordpress/1", "thundering typhoons"})
	c.Assert(err, ErrorMatches, `unrecognized args: \["thundering typhoons"\]`)
}

func (s *UnitSuite) TestRunStop(c *C) {
	// Set up state.
	ch := s.AddTestingCharm(c, "dummy")
	svc, err := s.Conn.AddService("dummy", ch)
	c.Assert(err, IsNil)
	unit, err := svc.AddUnit()
	c.Assert(err, IsNil)

	// Set up local environment.
	jujuDir := environs.VarDir
	vers := version.Current.String()
	toolsDir := filepath.Join(jujuDir, "tools", vers)
	err = os.MkdirAll(toolsDir, 0755)
	c.Assert(err, IsNil)
	err = ioutil.WriteFile(filepath.Join(toolsDir, "jujuc"), nil, 0644)
	c.Assert(err, IsNil)
	toolsLink := filepath.Join(jujuDir, "tools", "unit-dummy-0")
	err = os.Symlink(vers, toolsLink)
	c.Assert(err, IsNil)

	// Run a unit agent.
	a := &UnitAgent{
		Conf:     AgentConf{jujuDir, *s.StateInfo(c)},
		UnitName: unit.Name(),
	}
	done := make(chan error)
	go func() {
		done <- a.Run(nil)
	}()
	defer a.Stop()
	timeout := time.After(3 * time.Second)
	for {
		select {
		case <-timeout:
			c.Fatalf("no activity detected")
		case <-time.After(50 * time.Millisecond):
			st, info, err := unit.Status()
			c.Assert(err, IsNil)
			switch st {
			case state.UnitPending, state.UnitInstalled:
				c.Logf("waiting...")
				continue
			case state.UnitStarted:
				c.Logf("started!")
			default:
				c.Fatalf("unexpected status %s %s", st, info)
			}
		}
		break
	}
	err = a.Stop()
	c.Assert(err, IsNil)
	c.Assert(<-done, IsNil)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"

	"github.com/gregjones/httpcache/diskcache"
	"github.com/miku/proxima"
)

func main() {
	addr := flag.String("addr", ":3000", "host:port to listen on")
	cacheDir := flag.String("dir", filepath.Join(os.TempDir(), "proxima.v1"), "cache directory")
	version := flag.Bool("v", false, "prints current program version")

	flag.Parse()

	if *version {
		fmt.Println(proxima.Version)
		os.Exit(0)
	}

	options := proxima.Options{Cache: diskcache.New(*cacheDir)}
	mux := http.NewServeMux()

	mux.Handle("/u", proxima.URLHandler(options))
	mux.Handle("/g/s/", proxima.SameAsHandler(options))
	mux.Handle("/g/i/", proxima.GndImageHandler(options))

	mux.HandleFunc("/g/r/", func(w http.ResponseWriter, r *http.Request) {
		gnd := r.URL.Path[len("/g/r/"):]
		url := fmt.Sprintf("http://d-nb.info/gnd/%s/about/rdf", gnd)
		http.Redirect(w, r, fmt.Sprintf("/u?url=%s", url), 303)
	})

	log.Printf("Listening on %s, %s\n", *addr, *cacheDir)
	http.ListenAndServe(*addr, mux)
}
<commit_msg>add support for base url path<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"

	"github.com/gregjones/httpcache/diskcache"
	"github.com/miku/proxima"
)

// prepend helps to run this service under a directory by appending a base
// TODO: there must be support for this in mux.Handle, or not?
func prepend(s, base string) string {
	if base == "/" {
		return s
	}
	return fmt.Sprintf("%s%s", base, s)
}

func main() {
	addr := flag.String("addr", ":3000", "host:port to listen on")
	cacheDir := flag.String("dir", filepath.Join(os.TempDir(), "proxima.v1"), "cache directory")
	version := flag.Bool("v", false, "prints current program version")
	base := flag.String("base", "/", "URL base")

	flag.Parse()

	if *version {
		fmt.Println(proxima.Version)
		os.Exit(0)
	}

	options := proxima.Options{Cache: diskcache.New(*cacheDir)}
	mux := http.NewServeMux()

	mux.Handle(prependPrefix("/u", *base), proxima.URLHandler(options))
	mux.Handle(prependPrefix("/g/s/", *base), proxima.SameAsHandler(options))
	mux.Handle(prependPrefix("/g/i/", *base), proxima.GndImageHandler(options))

	mux.HandleFunc("/g/r/", func(w http.ResponseWriter, r *http.Request) {
		gnd := r.URL.Path[len("/g/r/"):]
		url := fmt.Sprintf("http://d-nb.info/gnd/%s/about/rdf", gnd)
		http.Redirect(w, r, fmt.Sprintf("/u?url=%s", url), 303)
	})

	log.Printf("Listening on %s, %s\n", *addr, *cacheDir)
	http.ListenAndServe(*addr, mux)
}
<|endoftext|>


<commit_before>package cmdbhv

import (
	"fmt"

	"go.polydawn.net/meep"
)

const (
	EXIT_SUCCESS      = 0
	EXIT_BADARGS      = 1
	EXIT_UNKNOWNPANIC = 2  // same code as golang uses when the process dies naturally on an unhandled panic.
	EXIT_JOB          = 10 // used to indicate a job reported a nonzero exit code (from cli commands that execute a single job).
	EXIT_USER         = 3  // grab bag for general user input errors (try to make a more specific code if possible/useful)
)

type ErrExit struct {
	Message string
	Code int
}

func (e ErrExit) Error() string {
	return e.Message
}

type ErrBadArgs struct {
	meep.TraitAutodescribing
	Message string
}

func ErrMissingParameter(paramName string) error {
	return meep.Meep(&ErrBadArgs{
		Message: fmt.Sprintf("%q is a required parameter", paramName),
	})
}
<commit_msg>gofmt fix.<commit_after>package cmdbhv

import (
	"fmt"

	"go.polydawn.net/meep"
)

const (
	EXIT_SUCCESS      = 0
	EXIT_BADARGS      = 1
	EXIT_UNKNOWNPANIC = 2  // same code as golang uses when the process dies naturally on an unhandled panic.
	EXIT_JOB          = 10 // used to indicate a job reported a nonzero exit code (from cli commands that execute a single job).
	EXIT_USER         = 3  // grab bag for general user input errors (try to make a more specific code if possible/useful)
)

type ErrExit struct {
	Message string
	Code    int
}

func (e ErrExit) Error() string {
	return e.Message
}

type ErrBadArgs struct {
	meep.TraitAutodescribing
	Message string
}

func ErrMissingParameter(paramName string) error {
	return meep.Meep(&ErrBadArgs{
		Message: fmt.Sprintf("%q is a required parameter", paramName),
	})
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"net/url"
	"strings"
	"time"

	"github.com/StackExchange/scollector/collectors"
	"github.com/StackExchange/scollector/opentsdb"
	"github.com/StackExchange/scollector/queue"
	"github.com/StackExchange/slog"
)

var flagFilter = flag.String("f", "", "Filters collectors matching this term. Works with all other arguments.")
var flagTest = flag.Bool("t", false, "Test - run collectors once, print, and exit.")
var flagList = flag.Bool("l", false, "List")
var flagPrint = flag.Bool("p", false, "Print to screen instead of sending to a host")
var host = flag.String("h", "ny-devtsaf01", `OpenTSDB host. Ex: "tsdb.example.com". Can optionally specify port: "tsdb.example.com:4000", but will default to 4242 otherwise`)
var colDir = flag.String("c", "", `Passthrough collector directory. It should contain numbered directories like the OpenTSDB scollector expects. Any executable file in those directories is run every N seconds, where N is the name of the directory. Use 0 for a program that should be run continuously and simply pass data through to OpenTSDB (the program will be restarted if it exits. Data output format is: "metric timestamp value tag1=val1 tag2=val2 ...". Timestamp is in Unix format (seconds since epoch). Tags are optional. A host tag is automatically added, but overridden if specified.`)
var batchSize = flag.Int("b", 0, "OpenTSDB batch size. Used for debugging bad data.")

var mains []func()

func main() {
	flag.Parse()
	for _, m := range mains {
		m()
	}

	if *colDir != "" {
		collectors.InitPrograms(*colDir)
	}
	if *batchSize > 0 {
		queue.BatchSize = *batchSize
	}
	c := collectors.Search(*flagFilter)
	u := parseHost()
	if *flagTest {
		test(c)
		return
	} else if *flagList {
		list(c)
		return
	} else if *host != "" {
		if u == nil {
			slog.Fatal("invalid host:", *host)
		}
	}

	if *flagPrint {
		collectors.DEFAULT_FREQ = time.Second * 3
		slog.Infoln("Set default frequency to", collectors.DEFAULT_FREQ)
	}
	cdp := collectors.Run(c)
	if u != nil && !*flagPrint {
		slog.Infoln("OpenTSDB host:", u)
		queue.New(u.String(), cdp)
	} else {
		slog.Infoln("Outputting to screen")
		printPut(cdp)
	}
	select {}
}

func test(cs []collectors.Collector) {
	dpchan := make(chan *opentsdb.DataPoint)
	for _, c := range cs {
		go c.Run(dpchan)
		slog.Infoln("run", c.Name())
	}
	next := time.After(time.Second * 2)
Loop:
	for {
		select {
		case dp := <-dpchan:
			slog.Info(dp.Telnet())
		case <-next:
			break Loop
		}
	}
}

func list(cs []collectors.Collector) {
	for _, c := range cs {
		slog.Infoln(c.Name())
	}
}

func parseHost() *url.URL {
	if *host == "" {
		return nil
	}
	u := url.URL{
		Scheme: "http",
		Path:   "/api/put",
	}
	if !strings.Contains(*host, ":") {
		*host += ":4242"
	}
	u.Host = *host
	return &u
}

func printPut(c chan *opentsdb.DataPoint) {
	for dp := range c {
		slog.Info(dp.Telnet())
	}
}
<commit_msg>cmd/scollector: Bump test run time and display message<commit_after>package main

import (
	"flag"
	"net/url"
	"strings"
	"time"

	"github.com/StackExchange/scollector/collectors"
	"github.com/StackExchange/scollector/opentsdb"
	"github.com/StackExchange/scollector/queue"
	"github.com/StackExchange/slog"
)

var flagFilter = flag.String("f", "", "Filters collectors matching this term. Works with all other arguments.")
var flagTest = flag.Bool("t", false, "Test - run collectors once, print, and exit.")
var flagList = flag.Bool("l", false, "List")
var flagPrint = flag.Bool("p", false, "Print to screen instead of sending to a host")
var host = flag.String("h", "ny-devtsaf01", `OpenTSDB host. Ex: "tsdb.example.com". Can optionally specify port: "tsdb.example.com:4000", but will default to 4242 otherwise`)
var colDir = flag.String("c", "", `Passthrough collector directory. It should contain numbered directories like the OpenTSDB scollector expects. Any executable file in those directories is run every N seconds, where N is the name of the directory. Use 0 for a program that should be run continuously and simply pass data through to OpenTSDB (the program will be restarted if it exits. Data output format is: "metric timestamp value tag1=val1 tag2=val2 ...". Timestamp is in Unix format (seconds since epoch). Tags are optional. A host tag is automatically added, but overridden if specified.`)
var batchSize = flag.Int("b", 0, "OpenTSDB batch size. Used for debugging bad data.")

var mains []func()

func main() {
	flag.Parse()
	for _, m := range mains {
		m()
	}

	if *colDir != "" {
		collectors.InitPrograms(*colDir)
	}
	if *batchSize > 0 {
		queue.BatchSize = *batchSize
	}
	c := collectors.Search(*flagFilter)
	u := parseHost()
	if *flagTest {
		test(c)
		return
	} else if *flagList {
		list(c)
		return
	} else if *host != "" {
		if u == nil {
			slog.Fatal("invalid host:", *host)
		}
	}

	if *flagPrint {
		collectors.DEFAULT_FREQ = time.Second * 3
		slog.Infoln("Set default frequency to", collectors.DEFAULT_FREQ)
	}
	cdp := collectors.Run(c)
	if u != nil && !*flagPrint {
		slog.Infoln("OpenTSDB host:", u)
		queue.New(u.String(), cdp)
	} else {
		slog.Infoln("Outputting to screen")
		printPut(cdp)
	}
	select {}
}

func test(cs []collectors.Collector) {
	dpchan := make(chan *opentsdb.DataPoint)
	for _, c := range cs {
		go c.Run(dpchan)
		slog.Infoln("run", c.Name())
	}
	dur := time.Second * 10
	slog.Infoln("running for", dur)
	next := time.After(dur)
Loop:
	for {
		select {
		case dp := <-dpchan:
			slog.Info(dp.Telnet())
		case <-next:
			break Loop
		}
	}
}

func list(cs []collectors.Collector) {
	for _, c := range cs {
		slog.Infoln(c.Name())
	}
}

func parseHost() *url.URL {
	if *host == "" {
		return nil
	}
	u := url.URL{
		Scheme: "http",
		Path:   "/api/put",
	}
	if !strings.Contains(*host, ":") {
		*host += ":4242"
	}
	u.Host = *host
	return &u
}

func printPut(c chan *opentsdb.DataPoint) {
	for dp := range c {
		slog.Info(dp.Telnet())
	}
}
<|endoftext|>


<commit_before>package storage

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"

	"github.com/golang/protobuf/jsonpb"
	cmdutil "github.com/ohsu-comp-bio/funnel/cmd/util"
	"github.com/ohsu-comp-bio/funnel/config"
	"github.com/ohsu-comp-bio/funnel/storage"
	"github.com/ohsu-comp-bio/funnel/tes"
	"github.com/spf13/cobra"
)

// NewCommand returns the "storage" subcommands.
func NewCommand() *cobra.Command {

	configFile := ""
	conf := config.Config{}
	flagConf := config.Config{}

	cmd := &cobra.Command{
		Use:   "storage",
		Short: "Access storage via Funnel's client libraries.",
		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
			var err error
			conf, err = cmdutil.MergeConfigFileWithFlags(configFile, flagConf)
			if err != nil {
				return fmt.Errorf("processing config: %v", err)
			}
			return nil
		},
	}
	cmd.SetGlobalNormalizationFunc(cmdutil.NormalizeFlags)
	f := cmd.PersistentFlags()
	f.StringVarP(&configFile, "config", "c", configFile, "Config File")

	statCmd := &cobra.Command{
		Use:   "stat [url]",
		Short: "Returns information about the object at the given URL.",
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) != 1 {
				return cmd.Usage()
			}

			store, err := storage.NewMux(conf)
			if err != nil {
				return fmt.Errorf("creating storage clients: %s", err)
			}

			obj, err := store.Stat(context.Background(), args[0])
			if err != nil {
				return err
			}

			b, err := json.Marshal(obj)
			if err != nil {
				return fmt.Errorf("marshaling output: %s", err)
			}
			fmt.Println(string(b))
			return nil
		},
	}

	statTaskCmd := &cobra.Command{
		Use:   "stat-task [task file]",
		Short: "Returns information about inputs/outputs of the task.",
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) != 1 {
				return cmd.Usage()
			}

			store, err := storage.NewMux(conf)
			if err != nil {
				return fmt.Errorf("creating storage clients: %s", err)
			}

			f, err := os.Open(args[0])
			if err != nil {
				return fmt.Errorf("opening task file: %s", err)
			}
			defer f.Close()

			dec := json.NewDecoder(f)
			for {
				task := &tes.Task{}
				err := jsonpb.UnmarshalNext(dec, task)
				if err == io.EOF {
					break
				}
				if err != nil {
					fmt.Fprintf(os.Stderr, "error: %s\n", err)
					continue
				}

				for _, in := range task.Inputs {
					obj, err := store.Stat(context.Background(), in.Url)
					if err != nil {
						fmt.Fprintf(os.Stderr, "error: %s\n", err)
						continue
					}

					b, _ := json.Marshal(obj)
					fmt.Println(string(b))
				}

				for _, out := range task.Outputs {
					obj, err := store.Stat(context.Background(), out.Url)
					if err != nil {
						fmt.Fprintf(os.Stderr, "error: %s\n", err)
						continue
					}

					b, _ := json.Marshal(obj)
					fmt.Println(string(b))
				}
			}

			return nil
		},
	}

	listCmd := &cobra.Command{
		Use:   "list [url]",
		Short: "List objects at the given URL.",
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) != 1 {
				return cmd.Usage()
			}

			store, err := storage.NewMux(conf)
			if err != nil {
				return fmt.Errorf("creating storage clients: %s", err)
			}

			objects, err := store.List(context.Background(), args[0])
			if err != nil {
				return err
			}

			b, err := json.Marshal(objects)
			if err != nil {
				return fmt.Errorf("marshaling output: %s", err)
			}
			fmt.Println(string(b))
			return nil
		},
	}

	cmd.AddCommand(listCmd)
	cmd.AddCommand(statCmd)
	cmd.AddCommand(statTaskCmd)
	return cmd
}
<commit_msg>cmd/storage: add get and put<commit_after>package storage

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"

	"github.com/golang/protobuf/jsonpb"
	cmdutil "github.com/ohsu-comp-bio/funnel/cmd/util"
	"github.com/ohsu-comp-bio/funnel/config"
	"github.com/ohsu-comp-bio/funnel/storage"
	"github.com/ohsu-comp-bio/funnel/tes"
	"github.com/spf13/cobra"
)

// NewCommand returns the "storage" subcommands.
func NewCommand() *cobra.Command {

	configFile := ""
	conf := config.Config{}
	flagConf := config.Config{}

	cmd := &cobra.Command{
		Use:   "storage",
		Short: "Access storage via Funnel's client libraries.",
		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
			var err error
			conf, err = cmdutil.MergeConfigFileWithFlags(configFile, flagConf)
			if err != nil {
				return fmt.Errorf("processing config: %v", err)
			}
			return nil
		},
	}
	cmd.SetGlobalNormalizationFunc(cmdutil.NormalizeFlags)
	f := cmd.PersistentFlags()
	f.StringVarP(&configFile, "config", "c", configFile, "Config File")

	statCmd := &cobra.Command{
		Use:   "stat [url]",
		Short: "Returns information about the object at the given URL.",
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) != 1 {
				return cmd.Usage()
			}

			store, err := storage.NewMux(conf)
			if err != nil {
				return fmt.Errorf("creating storage clients: %s", err)
			}

			obj, err := store.Stat(context.Background(), args[0])
			if err != nil {
				return err
			}

			b, err := json.Marshal(obj)
			if err != nil {
				return fmt.Errorf("marshaling output: %s", err)
			}
			fmt.Println(string(b))
			return nil
		},
	}

	statTaskCmd := &cobra.Command{
		Use:   "stat-task [task file]",
		Short: "Returns information about inputs/outputs of the task.",
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) != 1 {
				return cmd.Usage()
			}

			store, err := storage.NewMux(conf)
			if err != nil {
				return fmt.Errorf("creating storage clients: %s", err)
			}

			f, err := os.Open(args[0])
			if err != nil {
				return fmt.Errorf("opening task file: %s", err)
			}
			defer f.Close()

			dec := json.NewDecoder(f)
			for {
				task := &tes.Task{}
				err := jsonpb.UnmarshalNext(dec, task)
				if err == io.EOF {
					break
				}
				if err != nil {
					fmt.Fprintf(os.Stderr, "error: %s\n", err)
					continue
				}

				for _, in := range task.Inputs {
					obj, err := store.Stat(context.Background(), in.Url)
					if err != nil {
						fmt.Fprintf(os.Stderr, "error: %s\n", err)
						continue
					}

					b, _ := json.Marshal(obj)
					fmt.Println(string(b))
				}

				for _, out := range task.Outputs {
					obj, err := store.Stat(context.Background(), out.Url)
					if err != nil {
						fmt.Fprintf(os.Stderr, "error: %s\n", err)
						continue
					}

					b, _ := json.Marshal(obj)
					fmt.Println(string(b))
				}
			}

			return nil
		},
	}

	listCmd := &cobra.Command{
		Use:   "list [url]",
		Short: "List objects at the given URL.",
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) != 1 {
				return cmd.Usage()
			}

			store, err := storage.NewMux(conf)
			if err != nil {
				return fmt.Errorf("creating storage clients: %s", err)
			}

			objects, err := store.List(context.Background(), args[0])
			if err != nil {
				return err
			}

			b, err := json.Marshal(objects)
			if err != nil {
				return fmt.Errorf("marshaling output: %s", err)
			}
			fmt.Println(string(b))
			return nil
		},
	}

	getCmd := &cobra.Command{
		Use:   "get [url] [path]",
		Short: "Get the object at the given URL.",
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) != 2 {
				return cmd.Usage()
			}

			store, err := storage.NewMux(conf)
			if err != nil {
				return fmt.Errorf("creating storage clients: %s", err)
			}

			obj, err := store.Get(context.Background(), args[0], args[1])
			if err != nil {
				return err
			}

			b, err := json.Marshal(obj)
			if err != nil {
				return fmt.Errorf("marshaling output: %s", err)
			}
			fmt.Println(string(b))
			return nil
		},
	}

	putCmd := &cobra.Command{
		Use:   "put [path] [url]",
		Short: "Put the local file to the given URL.",
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) != 2 {
				return cmd.Usage()
			}

			store, err := storage.NewMux(conf)
			if err != nil {
				return fmt.Errorf("creating storage clients: %s", err)
			}

			obj, err := store.Put(context.Background(), args[1], args[0])
			if err != nil {
				return err
			}

			b, err := json.Marshal(obj)
			if err != nil {
				return fmt.Errorf("marshaling output: %s", err)
			}
			fmt.Println(string(b))
			return nil
		},
	}

	cmd.AddCommand(getCmd)
	cmd.AddCommand(putCmd)
	cmd.AddCommand(listCmd)
	cmd.AddCommand(statCmd)
	cmd.AddCommand(statTaskCmd)
	return cmd
}
<|endoftext|>


<commit_before>package main
<commit_msg>stub wsdlgen command<commit_after>package main

func main() {
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"github.com/GeoNet/web"
	"github.com/GeoNet/web/api/apidoc"
	"html/template"
	"net/http"
	"strings"
)

const (
	siteGeoJSON = `SELECT row_to_json(fc)
                         FROM ( SELECT 'FeatureCollection' as type, array_to_json(array_agg(f)) as features
                         FROM (SELECT 'Feature' as type,
                         ST_AsGeoJSON(s.location)::json as geometry,
                         row_to_json((SELECT l FROM
                         	(
                         		SELECT
                         		siteid AS "siteID",
                                height,
                                ground_relationship AS "groundRelationship",
                                name,
                                networkID as "networkID"
                           ) as l
                         )) as properties FROM (fits.site join fits.network using (networkpk)) as s `
	fc = ` ) As f )  as fc`
)

var siteDoc = apidoc.Endpoint{Title: "Site",
	Description: `Look up site information.`,
	Queries: []*apidoc.Query{
		new(siteTypeQuery).Doc(),
		new(siteQuery).Doc(),
	},
}

var siteQueryD = &apidoc.Query{
	Accept:      web.V1GeoJSON,
	Title:       "Site",
	Description: "Find information for individual sites.",
	Example:     "/site?siteID=HOLD&networkID=CG",
	ExampleHost: exHost,
	URI:         "/site?siteID=(siteID)&networkID=(networkID)",
	Params: map[string]template.HTML{
		"siteID":    siteIDDoc,
		"networkID": networkIDDoc,
	},
	Props: siteProps,
}

type siteQuery struct {
	siteID, networkID string
}

func (q *siteQuery) Doc() *apidoc.Query {
	return siteQueryD
}

func (q *siteQuery) Validate(w http.ResponseWriter, r *http.Request) bool {
	switch {
	case len(r.URL.Query()) != 2:
		web.BadRequest(w, r, "incorrect number of query params.")
		return false
	case !web.ParamsExist(w, r, "siteID", "networkID"):
		return false
	}

	q.networkID = r.URL.Query().Get("networkID")
	q.siteID = r.URL.Query().Get("siteID")

	return validSite(w, r, q.networkID, q.siteID)
}

func (q *siteQuery) Handle(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", web.V1GeoJSON)

	var d string

	err := db.QueryRow(
		siteGeoJSON+` WHERE siteid = $1 and networkid = $2`+fc, q.siteID, q.networkID).Scan(&d)
	if err != nil {
		web.ServiceUnavailable(w, r, err)
		return
	}

	b := []byte(d)
	web.Ok(w, r, &b)
}

var siteTypeQueryD = &apidoc.Query{
	Accept:      web.V1GeoJSON,
	Title:       "Sites",
	Description: "Filter sites by observation type, method, and location.",
	Example:     "/site?typeID=e",
	ExampleHost: exHost,
	URI:         "/site?[typeID=(typeID)]&[methodID=(methodID)]&[within=POLYGON((...))]",
	Params: map[string]template.HTML{
		"typeID":   optDoc + `  ` + typeIDDoc,
		"methodID": optDoc + `  ` + methodIDDoc + `  typeID must be specified as well.`,
		"within":   optDoc + `  ` + withinDoc,
	},
	Props: siteProps,
}

type siteTypeQuery struct {
	typeID, methodID, within string
}

func (q *siteTypeQuery) Doc() *apidoc.Query {
	return siteTypeQueryD
}

func (q *siteTypeQuery) Validate(w http.ResponseWriter, r *http.Request) bool {
	rl := r.URL.Query()

	if rl.Get("methodID") != "" && rl.Get("typeID") == "" {
		web.BadRequest(w, r, "typeID must be specified when methodID is specified.")
		return false
	}

	if rl.Get("typeID") != "" {
		q.typeID = rl.Get("typeID")

		if !validType(w, r, q.typeID) {
			return false
		}

		if rl.Get("methodID") != "" {
			q.methodID = rl.Get("methodID")
			if !validTypeMethod(w, r, q.typeID, q.methodID) {
				return false
			}
		}
	}

	if rl.Get("within") != "" {
		q.within = strings.Replace(rl.Get("within"), "+", "", -1)
		if !validPoly(w, r, q.within) {
			return false
		}
	}

	// delete any query params we know how to handle and there should be nothing left.
	rl.Del("typeID")
	rl.Del("methodID")
	rl.Del("within")
	if len(rl) > 0 {
		web.BadRequest(w, r, "incorrect number of query params.")
		return false
	}

	return true
}

func (q *siteTypeQuery) Handle(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", web.V1GeoJSON)

	var d string
	var err error

	switch {
	case q.typeID == "" && q.methodID == "" && q.within == "":
		err = db.QueryRow(
			siteGeoJSON + fc).Scan(&d)
	case q.typeID == "" && q.methodID == "" && q.within != "":
		err = db.QueryRow(
			siteGeoJSON+
				`where ST_Within(location::geometry, ST_GeomFromText($1, 4326))`+
				fc, q.within).Scan(&d)
	case q.typeID != "" && q.methodID == "" && q.within == "":
		err = db.QueryRow(
			siteGeoJSON+
				` where sitepk IN
(select distinct on (sitepk) sitepk from fits.observation where observation.typepk = (select typepk from fits.type where typeid = $1))`+fc, q.typeID).Scan(&d)
	case q.typeID != "" && q.methodID == "" && q.within != "":
		err = db.QueryRow(
			siteGeoJSON+
				` where sitepk IN
(select distinct on (sitepk) sitepk from fits.observation where observation.typepk = (select typepk from fits.type where typeid = $1))
 AND ST_Within(ST_Shift_Longitude(location::geometry), ST_Shift_Longitude(ST_GeomFromText($2, 4326)))`+fc, q.typeID, q.within).Scan(&d)
	case q.typeID != "" && q.methodID != "" && q.within == "":
		err = db.QueryRow(
			siteGeoJSON+
				` where sitepk IN
(select distinct on (sitepk) sitepk from fits.observation where
	observation.typepk = (select typepk from fits.type where typeid = $1)
	AND observation.methodpk = (select methodpk from fits.method where methodid = $2))`+fc, q.typeID, q.methodID).Scan(&d)
	case q.typeID != "" && q.methodID != "" && q.within != "":
		err = db.QueryRow(
			siteGeoJSON+
				` where sitepk IN
(select distinct on (sitepk) sitepk from fits.observation where
	observation.typepk = (select typepk from fits.type where typeid = $1)
	AND observation.methodpk = (select methodpk from fits.method where methodid = $2))
		 AND ST_Within(ST_Shift_Longitude(location::geometry), ST_Shift_Longitude(ST_GeomFromText($3, 4326)))`+fc, q.typeID, q.methodID, q.within).Scan(&d)
	}

	if err != nil {
		web.ServiceUnavailable(w, r, err)
		return
	}

	b := []byte(d)
	web.Ok(w, r, &b)
}

// validSite checks that the siteID and networkID combination exists in the DB.
func validSite(w http.ResponseWriter, r *http.Request, networkID, siteID string) bool {
	var d string

	err := db.QueryRow("select siteID FROM fits.site join fits.network using (networkpk) where siteid = $2 and networkid = $1", networkID, siteID).Scan(&d)
	if err == sql.ErrNoRows {
		web.NotFound(w, r, "invalid siteID and networkID combination: "+siteID+" "+networkID)
		return false
	}
	if err != nil {
		web.ServiceUnavailable(w, r, err)
		return false
	}

	return true
}
<commit_msg>Move site geoJSON queries to own funcs.<commit_after>package main

import (
	"database/sql"
	"github.com/GeoNet/web"
	"github.com/GeoNet/web/api/apidoc"
	"html/template"
	"net/http"
	"strings"
)

const (
	siteGeoJSON = `SELECT row_to_json(fc)
                         FROM ( SELECT 'FeatureCollection' as type, array_to_json(array_agg(f)) as features
                         FROM (SELECT 'Feature' as type,
                         ST_AsGeoJSON(s.location)::json as geometry,
                         row_to_json((SELECT l FROM
                         	(
                         		SELECT
                         		siteid AS "siteID",
                                height,
                                ground_relationship AS "groundRelationship",
                                name,
                                networkID as "networkID"
                           ) as l
                         )) as properties FROM (fits.site join fits.network using (networkpk)) as s `
	fc = ` ) As f )  as fc`
)

var siteDoc = apidoc.Endpoint{Title: "Site",
	Description: `Look up site information.`,
	Queries: []*apidoc.Query{
		new(siteTypeQuery).Doc(),
		new(siteQuery).Doc(),
	},
}

var siteQueryD = &apidoc.Query{
	Accept:      web.V1GeoJSON,
	Title:       "Site",
	Description: "Find information for individual sites.",
	Example:     "/site?siteID=HOLD&networkID=CG",
	ExampleHost: exHost,
	URI:         "/site?siteID=(siteID)&networkID=(networkID)",
	Params: map[string]template.HTML{
		"siteID":    siteIDDoc,
		"networkID": networkIDDoc,
	},
	Props: siteProps,
}

type siteQuery struct {
	siteID, networkID string
}

func (q *siteQuery) Doc() *apidoc.Query {
	return siteQueryD
}

func (q *siteQuery) Validate(w http.ResponseWriter, r *http.Request) bool {
	switch {
	case len(r.URL.Query()) != 2:
		web.BadRequest(w, r, "incorrect number of query params.")
		return false
	case !web.ParamsExist(w, r, "siteID", "networkID"):
		return false
	}

	q.networkID = r.URL.Query().Get("networkID")
	q.siteID = r.URL.Query().Get("siteID")

	return validSite(w, r, q.networkID, q.siteID)
}

func (q *siteQuery) Handle(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", web.V1GeoJSON)

	b, err := geoJSONSite(q.networkID, q.siteID)
	if err != nil {
		web.ServiceUnavailable(w, r, err)
		return
	}

	web.Ok(w, r, &b)
}

var siteTypeQueryD = &apidoc.Query{
	Accept:      web.V1GeoJSON,
	Title:       "Sites",
	Description: "Filter sites by observation type, method, and location.",
	Example:     "/site?typeID=e",
	ExampleHost: exHost,
	URI:         "/site?[typeID=(typeID)]&[methodID=(methodID)]&[within=POLYGON((...))]",
	Params: map[string]template.HTML{
		"typeID":   optDoc + `  ` + typeIDDoc,
		"methodID": optDoc + `  ` + methodIDDoc + `  typeID must be specified as well.`,
		"within":   optDoc + `  ` + withinDoc,
	},
	Props: siteProps,
}

type siteTypeQuery struct {
	typeID, methodID, within string
}

func (q *siteTypeQuery) Doc() *apidoc.Query {
	return siteTypeQueryD
}

func (q *siteTypeQuery) Validate(w http.ResponseWriter, r *http.Request) bool {
	rl := r.URL.Query()

	if rl.Get("methodID") != "" && rl.Get("typeID") == "" {
		web.BadRequest(w, r, "typeID must be specified when methodID is specified.")
		return false
	}

	if rl.Get("typeID") != "" {
		q.typeID = rl.Get("typeID")

		if !validType(w, r, q.typeID) {
			return false
		}

		if rl.Get("methodID") != "" {
			q.methodID = rl.Get("methodID")
			if !validTypeMethod(w, r, q.typeID, q.methodID) {
				return false
			}
		}
	}

	if rl.Get("within") != "" {
		q.within = strings.Replace(rl.Get("within"), "+", "", -1)
		if !validPoly(w, r, q.within) {
			return false
		}
	}

	// delete any query params we know how to handle and there should be nothing left.
	rl.Del("typeID")
	rl.Del("methodID")
	rl.Del("within")
	if len(rl) > 0 {
		web.BadRequest(w, r, "incorrect number of query params.")
		return false
	}

	return true
}

func (q *siteTypeQuery) Handle(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", web.V1GeoJSON)

	b, err := q.geoJSONSites()
	if err != nil {
		web.ServiceUnavailable(w, r, err)
		return
	}
	web.Ok(w, r, &b)
}

// validSite checks that the siteID and networkID combination exists in the DB.
func validSite(w http.ResponseWriter, r *http.Request, networkID, siteID string) bool {
	var d string

	err := db.QueryRow("select siteID FROM fits.site join fits.network using (networkpk) where siteid = $2 and networkid = $1", networkID, siteID).Scan(&d)
	if err == sql.ErrNoRows {
		web.NotFound(w, r, "invalid siteID and networkID combination: "+siteID+" "+networkID)
		return false
	}
	if err != nil {
		web.ServiceUnavailable(w, r, err)
		return false
	}

	return true
}

func geoJSONSite(networkID, siteID string) ([]byte, error) {
	var d string
	err := db.QueryRow(
		siteGeoJSON+` WHERE siteid = $1 and networkid = $2`+fc, siteID, networkID).Scan(&d)

	return []byte(d), err
}

func (q *siteTypeQuery) geoJSONSites() ([]byte, error) {
	var d string
	var err error

	switch {
	case q.typeID == "" && q.methodID == "" && q.within == "":
		err = db.QueryRow(
			siteGeoJSON + fc).Scan(&d)
	case q.typeID == "" && q.methodID == "" && q.within != "":
		err = db.QueryRow(
			siteGeoJSON+
				`where ST_Within(location::geometry, ST_GeomFromText($1, 4326))`+
				fc, q.within).Scan(&d)
	case q.typeID != "" && q.methodID == "" && q.within == "":
		err = db.QueryRow(
			siteGeoJSON+
				` where sitepk IN
(select distinct on (sitepk) sitepk from fits.observation where observation.typepk = (select typepk from fits.type where typeid = $1))`+fc, q.typeID).Scan(&d)
	case q.typeID != "" && q.methodID == "" && q.within != "":
		err = db.QueryRow(
			siteGeoJSON+
				` where sitepk IN
(select distinct on (sitepk) sitepk from fits.observation where observation.typepk = (select typepk from fits.type where typeid = $1))
 AND ST_Within(ST_Shift_Longitude(location::geometry), ST_Shift_Longitude(ST_GeomFromText($2, 4326)))`+fc, q.typeID, q.within).Scan(&d)
	case q.typeID != "" && q.methodID != "" && q.within == "":
		err = db.QueryRow(
			siteGeoJSON+
				` where sitepk IN
(select distinct on (sitepk) sitepk from fits.observation where
	observation.typepk = (select typepk from fits.type where typeid = $1)
	AND observation.methodpk = (select methodpk from fits.method where methodid = $2))`+fc, q.typeID, q.methodID).Scan(&d)
	case q.typeID != "" && q.methodID != "" && q.within != "":
		err = db.QueryRow(
			siteGeoJSON+
				` where sitepk IN
(select distinct on (sitepk) sitepk from fits.observation where
	observation.typepk = (select typepk from fits.type where typeid = $1)
	AND observation.methodpk = (select methodpk from fits.method where methodid = $2))
		 AND ST_Within(ST_Shift_Longitude(location::geometry), ST_Shift_Longitude(ST_GeomFromText($3, 4326)))`+fc, q.typeID, q.methodID, q.within).Scan(&d)
	}

	return []byte(d), err
}
<|endoftext|>