<commit_before>package asm

import (
	"bytes"
	"compress/gzip"
	"crypto/sha1"
	"fmt"
	"hash/crc64"
	"io"
	"os"
	"testing"

	"github.com/vbatts/tar-split/tar/storage"
)

var entries = []struct {
	Entry storage.Entry
	Body  []byte
}{
	{
		Entry: storage.Entry{
			Type:    storage.FileType,
			Name:    "./hurr.txt",
			Payload: []byte{2, 116, 164, 177, 171, 236, 107, 78},
			Size:    20,
		},
		Body: []byte("imma hurr til I derp"),
	},
	{
		Entry: storage.Entry{
			Type:    storage.FileType,
			Name:    "./ermahgerd.txt",
			Payload: []byte{126, 72, 89, 239, 230, 252, 160, 187},
			Size:    26,
		},
		Body: []byte("café con leche, por favor"),
	},
	{
		Entry: storage.Entry{
			Type:    storage.FileType,
			NameRaw: []byte{0x66, 0x69, 0x6c, 0x65, 0x2d, 0xe4}, // this is invalid UTF-8. Just checking the round trip.
			Payload: []byte{126, 72, 89, 239, 230, 252, 160, 187},
			Size:    26,
		},
		Body: []byte("café con leche, por favor"),
	},
}
var entriesMangled = []struct {
	Entry storage.Entry
	Body  []byte
}{
	{
		Entry: storage.Entry{
			Type:    storage.FileType,
			Name:    "./hurr.txt",
			Payload: []byte{3, 116, 164, 177, 171, 236, 107, 78},
			Size:    20,
		},
		// switch
		Body: []byte("imma derp til I hurr"),
	},
	{
		Entry: storage.Entry{
			Type:    storage.FileType,
			Name:    "./ermahgerd.txt",
			Payload: []byte{127, 72, 89, 239, 230, 252, 160, 187},
			Size:    26,
		},
		// san not con
		Body: []byte("café sans leche, por favor"),
	},
	{
		Entry: storage.Entry{
			Type:    storage.FileType,
			NameRaw: []byte{0x66, 0x69, 0x6c, 0x65, 0x2d, 0xe4},
			Payload: []byte{127, 72, 89, 239, 230, 252, 160, 187},
			Size:    26,
		},
		Body: []byte("café con leche, por favor"),
	},
}

func TestTarStreamMangledGetterPutter(t *testing.T) {
	fgp := storage.NewBufferFileGetPutter()

	// first lets prep a GetPutter and Packer
	for i := range entries {
		if entries[i].Entry.Type == storage.FileType {
			j, csum, err := fgp.Put(entries[i].Entry.GetName(), bytes.NewBuffer(entries[i].Body))
			if err != nil {
				t.Error(err)
			}
			if j != entries[i].Entry.Size {
				t.Errorf("size %q: expected %d; got %d",
					entries[i].Entry.GetName(),
					entries[i].Entry.Size,
					j)
			}
			if !bytes.Equal(csum, entries[i].Entry.Payload) {
				t.Errorf("checksum %q: expected %v; got %v",
					entries[i].Entry.GetName(),
					entries[i].Entry.Payload,
					csum)
			}
		}
	}

	for _, e := range entriesMangled {
		if e.Entry.Type == storage.FileType {
			rdr, err := fgp.Get(e.Entry.GetName())
			if err != nil {
				t.Error(err)
			}
			c := crc64.New(storage.CRCTable)
			i, err := io.Copy(c, rdr)
			if err != nil {
				t.Fatal(err)
			}
			rdr.Close()

			csum := c.Sum(nil)
			if bytes.Equal(csum, e.Entry.Payload) {
				t.Errorf("wrote %d bytes. checksum for %q should not have matched! %v",
					i,
					e.Entry.GetName(),
					csum)
			}
		}
	}
}

func TestTarStream(t *testing.T) {
	testCases := []struct {
		path            string
		expectedSHA1Sum string
		expectedSize    int64
	}{
		{"./testdata/t.tar.gz", "1eb237ff69bca6e22789ecb05b45d35ca307adbd", 10240},
		{"./testdata/longlink.tar.gz", "d9f6babe107b7247953dff6b5b5ae31a3a880add", 20480},
		{"./testdata/fatlonglink.tar.gz", "8537f03f89aeef537382f8b0bb065d93e03b0be8", 26234880},
		{"./testdata/iso-8859.tar.gz", "ddafa51cb03c74ec117ab366ee2240d13bba1ec3", 10240},
	}

	for _, tc := range testCases {
		fh, err := os.Open(tc.path)
		if err != nil {
			t.Fatal(err)
		}
		defer fh.Close()
		gzRdr, err := gzip.NewReader(fh)
		if err != nil {
			t.Fatal(err)
		}
		defer gzRdr.Close()

		// Setup where we'll store the metadata
		w := bytes.NewBuffer([]byte{})
		sp := storage.NewJSONPacker(w)
		fgp := storage.NewBufferFileGetPutter()

		// wrap the disassembly stream
		tarStream, err := NewInputTarStream(gzRdr, sp, fgp)
		if err != nil {
			t.Fatal(err)
		}

		// get a sum of the stream after it has passed through to ensure it's the same.
		h0 := sha1.New()
		i, err := io.Copy(h0, tarStream)
		if err != nil {
			t.Fatal(err)
		}

		if i != tc.expectedSize {
			t.Errorf("size of tar: expected %d; got %d", tc.expectedSize, i)
		}
		if fmt.Sprintf("%x", h0.Sum(nil)) != tc.expectedSHA1Sum {
			t.Fatalf("checksum of tar: expected %s; got %x", tc.expectedSHA1Sum, h0.Sum(nil))
		}

		//t.Logf("%s", w.String()) // if we fail, then show the packed info

		// If we've made it this far, then we'll turn it around and create a tar
		// stream from the packed metadata and buffered file contents.
		r := bytes.NewBuffer(w.Bytes())
		sup := storage.NewJSONUnpacker(r)
		// and reuse the fgp that we Put the payloads to.

		rc := NewOutputTarStream(fgp, sup)
		h1 := sha1.New()
		i, err = io.Copy(h1, rc)
		if err != nil {
			t.Fatal(err)
		}

		if i != tc.expectedSize {
			t.Errorf("size of output tar: expected %d; got %d", tc.expectedSize, i)
		}
		if fmt.Sprintf("%x", h1.Sum(nil)) != tc.expectedSHA1Sum {
			t.Fatalf("checksum of output tar: expected %s; got %x", tc.expectedSHA1Sum, h1.Sum(nil))
		}
	}
}
<commit_msg>tar/asm: basic benchmark on disasm/asm of testdata<commit_after>package asm

import (
	"bytes"
	"compress/gzip"
	"crypto/sha1"
	"fmt"
	"hash/crc64"
	"io"
	"io/ioutil"
	"os"
	"testing"

	"github.com/vbatts/tar-split/tar/storage"
)

var entries = []struct {
	Entry storage.Entry
	Body  []byte
}{
	{
		Entry: storage.Entry{
			Type:    storage.FileType,
			Name:    "./hurr.txt",
			Payload: []byte{2, 116, 164, 177, 171, 236, 107, 78},
			Size:    20,
		},
		Body: []byte("imma hurr til I derp"),
	},
	{
		Entry: storage.Entry{
			Type:    storage.FileType,
			Name:    "./ermahgerd.txt",
			Payload: []byte{126, 72, 89, 239, 230, 252, 160, 187},
			Size:    26,
		},
		Body: []byte("café con leche, por favor"),
	},
	{
		Entry: storage.Entry{
			Type:    storage.FileType,
			NameRaw: []byte{0x66, 0x69, 0x6c, 0x65, 0x2d, 0xe4}, // this is invalid UTF-8. Just checking the round trip.
			Payload: []byte{126, 72, 89, 239, 230, 252, 160, 187},
			Size:    26,
		},
		Body: []byte("café con leche, por favor"),
	},
}
var entriesMangled = []struct {
	Entry storage.Entry
	Body  []byte
}{
	{
		Entry: storage.Entry{
			Type:    storage.FileType,
			Name:    "./hurr.txt",
			Payload: []byte{3, 116, 164, 177, 171, 236, 107, 78},
			Size:    20,
		},
		// switch
		Body: []byte("imma derp til I hurr"),
	},
	{
		Entry: storage.Entry{
			Type:    storage.FileType,
			Name:    "./ermahgerd.txt",
			Payload: []byte{127, 72, 89, 239, 230, 252, 160, 187},
			Size:    26,
		},
		// san not con
		Body: []byte("café sans leche, por favor"),
	},
	{
		Entry: storage.Entry{
			Type:    storage.FileType,
			NameRaw: []byte{0x66, 0x69, 0x6c, 0x65, 0x2d, 0xe4},
			Payload: []byte{127, 72, 89, 239, 230, 252, 160, 187},
			Size:    26,
		},
		Body: []byte("café con leche, por favor"),
	},
}

func TestTarStreamMangledGetterPutter(t *testing.T) {
	fgp := storage.NewBufferFileGetPutter()

	// first lets prep a GetPutter and Packer
	for i := range entries {
		if entries[i].Entry.Type == storage.FileType {
			j, csum, err := fgp.Put(entries[i].Entry.GetName(), bytes.NewBuffer(entries[i].Body))
			if err != nil {
				t.Error(err)
			}
			if j != entries[i].Entry.Size {
				t.Errorf("size %q: expected %d; got %d",
					entries[i].Entry.GetName(),
					entries[i].Entry.Size,
					j)
			}
			if !bytes.Equal(csum, entries[i].Entry.Payload) {
				t.Errorf("checksum %q: expected %v; got %v",
					entries[i].Entry.GetName(),
					entries[i].Entry.Payload,
					csum)
			}
		}
	}

	for _, e := range entriesMangled {
		if e.Entry.Type == storage.FileType {
			rdr, err := fgp.Get(e.Entry.GetName())
			if err != nil {
				t.Error(err)
			}
			c := crc64.New(storage.CRCTable)
			i, err := io.Copy(c, rdr)
			if err != nil {
				t.Fatal(err)
			}
			rdr.Close()

			csum := c.Sum(nil)
			if bytes.Equal(csum, e.Entry.Payload) {
				t.Errorf("wrote %d bytes. checksum for %q should not have matched! %v",
					i,
					e.Entry.GetName(),
					csum)
			}
		}
	}
}

var testCases = []struct {
	path            string
	expectedSHA1Sum string
	expectedSize    int64
}{
	{"./testdata/t.tar.gz", "1eb237ff69bca6e22789ecb05b45d35ca307adbd", 10240},
	{"./testdata/longlink.tar.gz", "d9f6babe107b7247953dff6b5b5ae31a3a880add", 20480},
	{"./testdata/fatlonglink.tar.gz", "8537f03f89aeef537382f8b0bb065d93e03b0be8", 26234880},
	{"./testdata/iso-8859.tar.gz", "ddafa51cb03c74ec117ab366ee2240d13bba1ec3", 10240},
}

func TestTarStream(t *testing.T) {

	for _, tc := range testCases {
		fh, err := os.Open(tc.path)
		if err != nil {
			t.Fatal(err)
		}
		defer fh.Close()
		gzRdr, err := gzip.NewReader(fh)
		if err != nil {
			t.Fatal(err)
		}
		defer gzRdr.Close()

		// Setup where we'll store the metadata
		w := bytes.NewBuffer([]byte{})
		sp := storage.NewJSONPacker(w)
		fgp := storage.NewBufferFileGetPutter()

		// wrap the disassembly stream
		tarStream, err := NewInputTarStream(gzRdr, sp, fgp)
		if err != nil {
			t.Fatal(err)
		}

		// get a sum of the stream after it has passed through to ensure it's the same.
		h0 := sha1.New()
		i, err := io.Copy(h0, tarStream)
		if err != nil {
			t.Fatal(err)
		}

		if i != tc.expectedSize {
			t.Errorf("size of tar: expected %d; got %d", tc.expectedSize, i)
		}
		if fmt.Sprintf("%x", h0.Sum(nil)) != tc.expectedSHA1Sum {
			t.Fatalf("checksum of tar: expected %s; got %x", tc.expectedSHA1Sum, h0.Sum(nil))
		}

		//t.Logf("%s", w.String()) // if we fail, then show the packed info

		// If we've made it this far, then we'll turn it around and create a tar
		// stream from the packed metadata and buffered file contents.
		r := bytes.NewBuffer(w.Bytes())
		sup := storage.NewJSONUnpacker(r)
		// and reuse the fgp that we Put the payloads to.

		rc := NewOutputTarStream(fgp, sup)
		h1 := sha1.New()
		i, err = io.Copy(h1, rc)
		if err != nil {
			t.Fatal(err)
		}

		if i != tc.expectedSize {
			t.Errorf("size of output tar: expected %d; got %d", tc.expectedSize, i)
		}
		if fmt.Sprintf("%x", h1.Sum(nil)) != tc.expectedSHA1Sum {
			t.Fatalf("checksum of output tar: expected %s; got %x", tc.expectedSHA1Sum, h1.Sum(nil))
		}
	}
}

func BenchmarkAsm(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for _, tc := range testCases {
			func() {
				fh, err := os.Open(tc.path)
				if err != nil {
					b.Fatal(err)
				}
				defer fh.Close()
				gzRdr, err := gzip.NewReader(fh)
				if err != nil {
					b.Fatal(err)
				}
				defer gzRdr.Close()

				// Setup where we'll store the metadata
				w := bytes.NewBuffer([]byte{})
				sp := storage.NewJSONPacker(w)
				fgp := storage.NewBufferFileGetPutter()

				// wrap the disassembly stream
				tarStream, err := NewInputTarStream(gzRdr, sp, fgp)
				if err != nil {
					b.Fatal(err)
				}
				// read it all to the bit bucket
				i1, err := io.Copy(ioutil.Discard, tarStream)
				if err != nil {
					b.Fatal(err)
				}

				r := bytes.NewBuffer(w.Bytes())
				sup := storage.NewJSONUnpacker(r)
				// and reuse the fgp that we Put the payloads to.

				rc := NewOutputTarStream(fgp, sup)

				i2, err := io.Copy(ioutil.Discard, rc)
				if err != nil {
					b.Fatal(err)
				}
				if i1 != i2 {
					b.Errorf("%s: input(%d) and ouput(%d) byte count didn't match", tc.path, i1, i2)
				}
			}()
		}
	}
}
<|endoftext|>


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

package gaesessions

import (
	"bytes"
	"encoding/base32"
	"encoding/gob"
	"net/http"
	"strings"
	"time"

	"appengine"
	"appengine/datastore"
	"appengine/memcache"

	"github.com/gorilla/securecookie"
	"github.com/gorilla/sessions"
)

// MemcacheDatastoreStore -----------------------------------------------------

const defaultKind = "Session"

// NewMemcacheDatastoreStore returns a new MemcacheDatastoreStore.
//
// The kind argument is the kind name used to store the session data.
// If empty it will use "Session".
//
// See NewCookieStore() for a description of the other parameters.
func NewMemcacheDatastoreStore(kind, keyPrefix string, keyPairs ...[]byte) *MemcacheDatastoreStore {
	if kind == "" {
		kind = defaultKind
	}
	if keyPrefix == "" {
		keyPrefix = "gorilla.appengine.sessions."
	}
	return &MemcacheDatastoreStore{
		Codecs: securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path:   "/",
			MaxAge: 86400 * 30,
		},
		kind:   kind,
		prefix: keyPrefix,
	}
}

type MemcacheDatastoreStore struct {
	Codecs  []securecookie.Codec
	Options *sessions.Options // default configuration
	kind    string
	prefix  string
}

// Get returns a session for the given name after adding it to the registry.
//
// See CookieStore.Get().
func (s *MemcacheDatastoreStore) Get(r *http.Request, name string) (
	*sessions.Session, error) {
	return sessions.GetRegistry(r).Get(s, name)
}

// New returns a session for the given name without adding it to the registry.
//
// See CookieStore.New().
func (s *MemcacheDatastoreStore) New(r *http.Request, name string) (*sessions.Session,
	error) {
	session := sessions.NewSession(s, name)
	session.Options = &(*s.Options)
	session.IsNew = true
	var err error
	if cookie, errCookie := r.Cookie(name); errCookie == nil {
		err = securecookie.DecodeMulti(name, cookie.Value, &session.ID,
			s.Codecs...)
		if err == nil {
			c := appengine.NewContext(r)
			err = loadFromMemcache(c, session)
			if err == memcache.ErrCacheMiss {
				err = loadFromDatastore(c, s.kind, session)
			}
			if err == nil {
				session.IsNew = false
			}
		}
	}
	return session, err
}

// Save adds a single session to the response.
func (s *MemcacheDatastoreStore) Save(r *http.Request, w http.ResponseWriter,
	session *sessions.Session) error {
	if session.ID == "" {
		session.ID = s.prefix +
			strings.TrimRight(
				base32.StdEncoding.EncodeToString(
					securecookie.GenerateRandomKey(32)), "=")
	}
	c := appengine.NewContext(r)
	if err := saveToMemcache(c, session); err != nil {
		return err
	}
	if err := saveToDatastore(c, s.kind, session); err != nil {
		return err
	}
	encoded, err := securecookie.EncodeMulti(session.Name(), session.ID,
		s.Codecs...)
	if err != nil {
		return err
	}
	http.SetCookie(w, sessions.NewCookie(session.Name(), encoded,
		session.Options))
	return nil
}

// DatastoreStore -------------------------------------------------------------

// Session is used to load and save session data in the datastore.
type Session struct {
	Date           time.Time
	ExpirationDate time.Time
	Value          []byte
}

// NewDatastoreStore returns a new DatastoreStore.
//
// The kind argument is the kind name used to store the session data.
// If empty it will use "Session".
//
// See NewCookieStore() for a description of the other parameters.
func NewDatastoreStore(kind string, keyPairs ...[]byte) *DatastoreStore {
	if kind == "" {
		kind = "Session"
	}
	return &DatastoreStore{
		Codecs: securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path:   "/",
			MaxAge: 86400 * 30,
		},
		kind: kind,
	}
}

// DatastoreStore stores sessions in the App Engine datastore.
type DatastoreStore struct {
	Codecs  []securecookie.Codec
	Options *sessions.Options // default configuration
	kind    string
}

// Get returns a session for the given name after adding it to the registry.
//
// See CookieStore.Get().
func (s *DatastoreStore) Get(r *http.Request, name string) (*sessions.Session,
	error) {
	return sessions.GetRegistry(r).Get(s, name)
}

// New returns a session for the given name without adding it to the registry.
//
// See CookieStore.New().
func (s *DatastoreStore) New(r *http.Request, name string) (*sessions.Session,
	error) {
	session := sessions.NewSession(s, name)
	session.Options = &(*s.Options)
	session.IsNew = true
	var err error
	if cookie, errCookie := r.Cookie(name); errCookie == nil {
		err = securecookie.DecodeMulti(name, cookie.Value, &session.ID,
			s.Codecs...)
		if err == nil {
			c := appengine.NewContext(r)
			err = loadFromDatastore(c, s.kind, session)
			if err == nil {
				session.IsNew = false
			}
		}
	}
	return session, err
}

// Save adds a single session to the response.
func (s *DatastoreStore) Save(r *http.Request, w http.ResponseWriter,
	session *sessions.Session) error {
	if session.ID == "" {
		session.ID =
			strings.TrimRight(
				base32.StdEncoding.EncodeToString(
					securecookie.GenerateRandomKey(32)), "=")
	}
	c := appengine.NewContext(r)
	if err := saveToDatastore(c, s.kind, session); err != nil {
		return err
	}
	encoded, err := securecookie.EncodeMulti(session.Name(), session.ID,
		s.Codecs...)
	if err != nil {
		return err
	}
	http.SetCookie(w, sessions.NewCookie(session.Name(), encoded,
		session.Options))
	return nil
}

// save writes encoded session.Values to datastore.
func saveToDatastore(c appengine.Context, kind string,
	session *sessions.Session) error {
	if len(session.Values) == 0 {
		// Don't need to write anything.
		return nil
	}
	serialized, err := serialize(session.Values)
	if err != nil {
		return err
	}
	k := datastore.NewKey(c, kind, session.ID, 0, nil)
	now := time.Now()
	var expirationDate time.Time
	if session.Options.MaxAge > 0 {
		expiration := time.Duration(session.Options.MaxAge) * time.Second
		expirationDate = now.Add(expiration)

		k, err = datastore.Put(c, k, &Session{
			Date:           now,
			ExpirationDate: expirationDate,
			Value:          serialized,
		})
		if err != nil {
			return err
		}
	} else {
		err = datastore.Delete(c, k)
		if err != nil {
			return err
		}
	}
	return nil
}

// load gets a value from datastore and decodes its content into
// session.Values.
func loadFromDatastore(c appengine.Context, kind string,
	session *sessions.Session) error {
	k := datastore.NewKey(c, kind, session.ID, 0, nil)
	entity := Session{}
	if err := datastore.Get(c, k, &entity); err != nil {
		return err
	}
	if err := deserialize(entity.Value, &session.Values); err != nil {
		return err
	}
	return nil
}

// remove expired sessions in the datastore. you can call this function
// from a cron job.
//
// sample handler config in app.yaml:
// handlers:
// - url: /tasks/removeExpiredSessions
//   script: _go_app
//   login: admin
// - url: /.*
//   script: _go_app
//
// handler registration code:
// http.HandleFunc("/tasks/removeExpiredSessions", removeExpiredSessionsHandler)
//
// sample handler:
// func removeExpiredSessionsHandler(w http.ResponseWriter, r *http.Request) {
//	c := appengine.NewContext(r)
//	gaesessions.RemoveExpiredDatastoreSessions(c, "")
// }
//
// sample cron.yaml:
// cron:
// - description: expired session removal job
//   url: /tasks/removeExpiredSessions
//   schedule: every 1 minutes
func RemoveExpiredDatastoreSessions(c appengine.Context, kind string) error {
	keys, err := findExpiredDatastoreSessionKeys(c, kind)
	if err != nil {
		return err
	}
	return datastore.DeleteMulti(c, keys)
}

func findExpiredDatastoreSessionKeys(c appengine.Context, kind string) (keys []*datastore.Key, err error) {
	if kind == "" {
		kind = defaultKind
	}
	now := time.Now()
	q := datastore.NewQuery(kind).Filter("ExpirationDate <=", now).KeysOnly()
	keys, err = q.GetAll(c, nil)
	return
}

// MemcacheStore --------------------------------------------------------------

// NewMemcacheStore returns a new MemcacheStore.
//
// The keyPrefix argument is the prefix used for memcache keys. If empty it
// will use "gorilla.appengine.sessions.".
//
// See NewCookieStore() for a description of the other parameters.
func NewMemcacheStore(keyPrefix string, keyPairs ...[]byte) *MemcacheStore {
	if keyPrefix == "" {
		keyPrefix = "gorilla.appengine.sessions."
	}
	return &MemcacheStore{
		Codecs: securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path:   "/",
			MaxAge: 86400 * 30,
		},
		prefix: keyPrefix,
	}
}

// MemcacheStore stores sessions in the App Engine memcache.
type MemcacheStore struct {
	Codecs  []securecookie.Codec
	Options *sessions.Options // default configuration
	prefix  string
}

// Get returns a session for the given name after adding it to the registry.
//
// See CookieStore.Get().
func (s *MemcacheStore) Get(r *http.Request, name string) (*sessions.Session,
	error) {
	return sessions.GetRegistry(r).Get(s, name)
}

// New returns a session for the given name without adding it to the registry.
//
// See CookieStore.New().
func (s *MemcacheStore) New(r *http.Request, name string) (*sessions.Session,
	error) {
	session := sessions.NewSession(s, name)
	session.Options = &(*s.Options)
	session.IsNew = true
	var err error
	if cookie, errCookie := r.Cookie(name); errCookie == nil {
		err = securecookie.DecodeMulti(name, cookie.Value, &session.ID,
			s.Codecs...)
		if err == nil {
			c := appengine.NewContext(r)
			err = loadFromMemcache(c, session)
			if err == nil {
				session.IsNew = false
			}
		}
	}
	return session, err
}

// Save adds a single session to the response.
func (s *MemcacheStore) Save(r *http.Request, w http.ResponseWriter,
	session *sessions.Session) error {
	if session.ID == "" {
		session.ID = s.prefix +
			strings.TrimRight(
				base32.StdEncoding.EncodeToString(
					securecookie.GenerateRandomKey(32)), "=")
	}
	c := appengine.NewContext(r)
	if err := saveToMemcache(c, session); err != nil {
		return err
	}
	encoded, err := securecookie.EncodeMulti(session.Name(), session.ID,
		s.Codecs...)
	if err != nil {
		return err
	}
	http.SetCookie(w, sessions.NewCookie(session.Name(), encoded,
		session.Options))
	return nil
}

// save writes encoded session.Values to memcache.
func saveToMemcache(c appengine.Context, session *sessions.Session) error {
	if len(session.Values) == 0 {
		// Don't need to write anything.
		return nil
	}
	serialized, err := serialize(session.Values)
	if err != nil {
		return err
	}
	var expiration time.Duration
	if session.Options.MaxAge > 0 {
		expiration = time.Duration(session.Options.MaxAge) * time.Second
		c.Debugf("MemcacheStore.save. session.ID=%s, expiration=%s",
			session.ID, expiration)
		err = memcache.Set(c, &memcache.Item{
			Key:        session.ID,
			Value:      serialized,
			Expiration: expiration,
		})
		if err != nil {
			return err
		}
	} else {
		err = memcache.Delete(c, session.ID)
		if err != nil {
			return err
		}
		c.Debugf("MemcacheStore.save. delete session.ID=%s", session.ID)
	}
	return nil
}

// load gets a value from memcache and decodes its content into session.Values.
func loadFromMemcache(c appengine.Context, session *sessions.Session) error {
	item, err := memcache.Get(c, session.ID)
	if err != nil {
		return err
	}
	if err := deserialize(item.Value, &session.Values); err != nil {
		return err
	}
	return nil
}

// Serialization --------------------------------------------------------------

// serialize encodes a value using gob.
func serialize(src interface{}) ([]byte, error) {
	buf := new(bytes.Buffer)
	enc := gob.NewEncoder(buf)
	if err := enc.Encode(src); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// deserialize decodes a value using gob.
func deserialize(src []byte, dst interface{}) error {
	dec := gob.NewDecoder(bytes.NewBuffer(src))
	if err := dec.Decode(dst); err != nil {
		return err
	}
	return nil
}
<commit_msg>Add a nonPersistentSessionDuration parameter to functions which creates session stores. It is used when session.Options.MaxAge <= 0.<commit_after>// Copyright 2012 The Gorilla Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gaesessions

import (
	"bytes"
	"encoding/base32"
	"encoding/gob"
	"net/http"
	"strings"
	"time"

	"appengine"
	"appengine/datastore"
	"appengine/memcache"

	"github.com/gorilla/securecookie"
	"github.com/gorilla/sessions"
)

// MemcacheDatastoreStore -----------------------------------------------------

const DefaultNonPersistentSessionDuration = time.Duration(24) * time.Hour
const defaultKind = "Session"

// NewMemcacheDatastoreStore returns a new MemcacheDatastoreStore.
//
// The kind argument is the kind name used to store the session data.
// If empty it will use "Session".
//
// See NewCookieStore() for a description of the other parameters.
func NewMemcacheDatastoreStore(kind, keyPrefix string, nonPersistentSessionDuration time.Duration, keyPairs ...[]byte) *MemcacheDatastoreStore {
	if kind == "" {
		kind = defaultKind
	}
	if keyPrefix == "" {
		keyPrefix = "gorilla.appengine.sessions."
	}
	return &MemcacheDatastoreStore{
		Codecs: securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path:   "/",
			MaxAge: 86400 * 30,
		},
		kind:   kind,
		prefix: keyPrefix,
		nonPersistentSessionDuration: nonPersistentSessionDuration,
	}
}

type MemcacheDatastoreStore struct {
	Codecs                       []securecookie.Codec
	Options                      *sessions.Options // default configuration
	kind                         string
	prefix                       string
	nonPersistentSessionDuration time.Duration
}

// Get returns a session for the given name after adding it to the registry.
//
// See CookieStore.Get().
func (s *MemcacheDatastoreStore) Get(r *http.Request, name string) (
	*sessions.Session, error) {
	return sessions.GetRegistry(r).Get(s, name)
}

// New returns a session for the given name without adding it to the registry.
//
// See CookieStore.New().
func (s *MemcacheDatastoreStore) New(r *http.Request, name string) (*sessions.Session,
	error) {
	session := sessions.NewSession(s, name)
	session.Options = &(*s.Options)
	session.IsNew = true
	var err error
	if cookie, errCookie := r.Cookie(name); errCookie == nil {
		err = securecookie.DecodeMulti(name, cookie.Value, &session.ID,
			s.Codecs...)
		if err == nil {
			c := appengine.NewContext(r)
			err = loadFromMemcache(c, session)
			if err == memcache.ErrCacheMiss {
				err = loadFromDatastore(c, s.kind, session)
			}
			if err == nil {
				session.IsNew = false
			}
		}
	}
	return session, err
}

// Save adds a single session to the response.
func (s *MemcacheDatastoreStore) Save(r *http.Request, w http.ResponseWriter,
	session *sessions.Session) error {
	if session.ID == "" {
		session.ID = s.prefix +
			strings.TrimRight(
				base32.StdEncoding.EncodeToString(
					securecookie.GenerateRandomKey(32)), "=")
	}
	c := appengine.NewContext(r)
	if err := saveToMemcache(c, s.nonPersistentSessionDuration, session); err != nil {
		return err
	}
	if err := saveToDatastore(c, s.kind, s.nonPersistentSessionDuration, session); err != nil {
		return err
	}
	encoded, err := securecookie.EncodeMulti(session.Name(), session.ID,
		s.Codecs...)
	if err != nil {
		return err
	}
	http.SetCookie(w, sessions.NewCookie(session.Name(), encoded,
		session.Options))
	return nil
}

// DatastoreStore -------------------------------------------------------------

// Session is used to load and save session data in the datastore.
type Session struct {
	Date           time.Time
	ExpirationDate time.Time
	Value          []byte
}

// NewDatastoreStore returns a new DatastoreStore.
//
// The kind argument is the kind name used to store the session data.
// If empty it will use "Session".
//
// See NewCookieStore() for a description of the other parameters.
func NewDatastoreStore(kind string, nonPersistentSessionDuration time.Duration, keyPairs ...[]byte) *DatastoreStore {
	if kind == "" {
		kind = "Session"
	}
	return &DatastoreStore{
		Codecs: securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path:   "/",
			MaxAge: 86400 * 30,
		},
		kind: kind,
		nonPersistentSessionDuration: nonPersistentSessionDuration,
	}
}

// DatastoreStore stores sessions in the App Engine datastore.
type DatastoreStore struct {
	Codecs                       []securecookie.Codec
	Options                      *sessions.Options // default configuration
	kind                         string
	nonPersistentSessionDuration time.Duration
}

// Get returns a session for the given name after adding it to the registry.
//
// See CookieStore.Get().
func (s *DatastoreStore) Get(r *http.Request, name string) (*sessions.Session,
	error) {
	return sessions.GetRegistry(r).Get(s, name)
}

// New returns a session for the given name without adding it to the registry.
//
// See CookieStore.New().
func (s *DatastoreStore) New(r *http.Request, name string) (*sessions.Session,
	error) {
	session := sessions.NewSession(s, name)
	session.Options = &(*s.Options)
	session.IsNew = true
	var err error
	if cookie, errCookie := r.Cookie(name); errCookie == nil {
		err = securecookie.DecodeMulti(name, cookie.Value, &session.ID,
			s.Codecs...)
		if err == nil {
			c := appengine.NewContext(r)
			err = loadFromDatastore(c, s.kind, session)
			if err == nil {
				session.IsNew = false
			}
		}
	}
	return session, err
}

// Save adds a single session to the response.
func (s *DatastoreStore) Save(r *http.Request, w http.ResponseWriter,
	session *sessions.Session) error {
	if session.ID == "" {
		session.ID =
			strings.TrimRight(
				base32.StdEncoding.EncodeToString(
					securecookie.GenerateRandomKey(32)), "=")
	}
	c := appengine.NewContext(r)
	if err := saveToDatastore(c, s.kind, s.nonPersistentSessionDuration, session); err != nil {
		return err
	}
	encoded, err := securecookie.EncodeMulti(session.Name(), session.ID,
		s.Codecs...)
	if err != nil {
		return err
	}
	http.SetCookie(w, sessions.NewCookie(session.Name(), encoded,
		session.Options))
	return nil
}

// save writes encoded session.Values to datastore.
func saveToDatastore(c appengine.Context, kind string,
	nonPersistentSessionDuration time.Duration,
	session *sessions.Session) error {
	if len(session.Values) == 0 {
		// Don't need to write anything.
		return nil
	}
	serialized, err := serialize(session.Values)
	if err != nil {
		return err
	}
	k := datastore.NewKey(c, kind, session.ID, 0, nil)
	now := time.Now()
	var expirationDate time.Time
	var expiration time.Duration
	if session.Options.MaxAge > 0 {
		expiration = time.Duration(session.Options.MaxAge) * time.Second
	} else {
		expiration = nonPersistentSessionDuration
	}
	if expiration > 0 {
		expirationDate = now.Add(expiration)
		k, err = datastore.Put(c, k, &Session{
			Date:           now,
			ExpirationDate: expirationDate,
			Value:          serialized,
		})
		if err != nil {
			return err
		}
	} else {
		err = datastore.Delete(c, k)
		if err != nil {
			return err
		}
	}
	return nil
}

// load gets a value from datastore and decodes its content into
// session.Values.
func loadFromDatastore(c appengine.Context, kind string,
	session *sessions.Session) error {
	k := datastore.NewKey(c, kind, session.ID, 0, nil)
	entity := Session{}
	if err := datastore.Get(c, k, &entity); err != nil {
		return err
	}
	if err := deserialize(entity.Value, &session.Values); err != nil {
		return err
	}
	return nil
}

// remove expired sessions in the datastore. you can call this function
// from a cron job.
//
// sample handler config in app.yaml:
// handlers:
// - url: /tasks/removeExpiredSessions
//   script: _go_app
//   login: admin
// - url: /.*
//   script: _go_app
//
// handler registration code:
// http.HandleFunc("/tasks/removeExpiredSessions", removeExpiredSessionsHandler)
//
// sample handler:
// func removeExpiredSessionsHandler(w http.ResponseWriter, r *http.Request) {
//	c := appengine.NewContext(r)
//	gaesessions.RemoveExpiredDatastoreSessions(c, "")
// }
//
// sample cron.yaml:
// cron:
// - description: expired session removal job
//   url: /tasks/removeExpiredSessions
//   schedule: every 1 minutes
func RemoveExpiredDatastoreSessions(c appengine.Context, kind string) error {
	keys, err := findExpiredDatastoreSessionKeys(c, kind)
	if err != nil {
		return err
	}
	return datastore.DeleteMulti(c, keys)
}

func findExpiredDatastoreSessionKeys(c appengine.Context, kind string) (keys []*datastore.Key, err error) {
	if kind == "" {
		kind = defaultKind
	}
	now := time.Now()
	q := datastore.NewQuery(kind).Filter("ExpirationDate <=", now).KeysOnly()
	keys, err = q.GetAll(c, nil)
	return
}

// MemcacheStore --------------------------------------------------------------

// NewMemcacheStore returns a new MemcacheStore.
//
// The keyPrefix argument is the prefix used for memcache keys. If empty it
// will use "gorilla.appengine.sessions.".
//
// See NewCookieStore() for a description of the other parameters.
func NewMemcacheStore(keyPrefix string, nonPersistentSessionDuration time.Duration, keyPairs ...[]byte) *MemcacheStore {
	if keyPrefix == "" {
		keyPrefix = "gorilla.appengine.sessions."
	}
	return &MemcacheStore{
		Codecs: securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path:   "/",
			MaxAge: 86400 * 30,
		},
		prefix: keyPrefix,
		nonPersistentSessionDuration: nonPersistentSessionDuration,
	}
}

// MemcacheStore stores sessions in the App Engine memcache.
type MemcacheStore struct {
	Codecs                       []securecookie.Codec
	Options                      *sessions.Options // default configuration
	prefix                       string
	nonPersistentSessionDuration time.Duration
}

// Get returns a session for the given name after adding it to the registry.
//
// See CookieStore.Get().
func (s *MemcacheStore) Get(r *http.Request, name string) (*sessions.Session,
	error) {
	return sessions.GetRegistry(r).Get(s, name)
}

// New returns a session for the given name without adding it to the registry.
//
// See CookieStore.New().
func (s *MemcacheStore) New(r *http.Request, name string) (*sessions.Session,
	error) {
	session := sessions.NewSession(s, name)
	session.Options = &(*s.Options)
	session.IsNew = true
	var err error
	if cookie, errCookie := r.Cookie(name); errCookie == nil {
		err = securecookie.DecodeMulti(name, cookie.Value, &session.ID,
			s.Codecs...)
		if err == nil {
			c := appengine.NewContext(r)
			err = loadFromMemcache(c, session)
			if err == nil {
				session.IsNew = false
			}
		}
	}
	return session, err
}

// Save adds a single session to the response.
func (s *MemcacheStore) Save(r *http.Request, w http.ResponseWriter,
	session *sessions.Session) error {
	if session.ID == "" {
		session.ID = s.prefix +
			strings.TrimRight(
				base32.StdEncoding.EncodeToString(
					securecookie.GenerateRandomKey(32)), "=")
	}
	c := appengine.NewContext(r)
	if err := saveToMemcache(c, s.nonPersistentSessionDuration, session); err != nil {
		return err
	}
	encoded, err := securecookie.EncodeMulti(session.Name(), session.ID,
		s.Codecs...)
	if err != nil {
		return err
	}
	http.SetCookie(w, sessions.NewCookie(session.Name(), encoded,
		session.Options))
	return nil
}

// save writes encoded session.Values to memcache.
func saveToMemcache(c appengine.Context,
	nonPersistentSessionDuration time.Duration,
	session *sessions.Session) error {
	if len(session.Values) == 0 {
		// Don't need to write anything.
		return nil
	}
	serialized, err := serialize(session.Values)
	if err != nil {
		return err
	}
	var expiration time.Duration
	if session.Options.MaxAge > 0 {
		expiration = time.Duration(session.Options.MaxAge) * time.Second
	} else {
		expiration = nonPersistentSessionDuration
	}
	if expiration > 0 {
		c.Debugf("MemcacheStore.save. session.ID=%s, expiration=%s",
			session.ID, expiration)
		err = memcache.Set(c, &memcache.Item{
			Key:        session.ID,
			Value:      serialized,
			Expiration: expiration,
		})
		if err != nil {
			return err
		}
	} else {
		err = memcache.Delete(c, session.ID)
		if err != nil {
			return err
		}
		c.Debugf("MemcacheStore.save. delete session.ID=%s", session.ID)
	}
	return nil
}

// load gets a value from memcache and decodes its content into session.Values.
func loadFromMemcache(c appengine.Context, session *sessions.Session) error {
	item, err := memcache.Get(c, session.ID)
	if err != nil {
		return err
	}
	if err := deserialize(item.Value, &session.Values); err != nil {
		return err
	}
	return nil
}

// Serialization --------------------------------------------------------------

// serialize encodes a value using gob.
func serialize(src interface{}) ([]byte, error) {
	buf := new(bytes.Buffer)
	enc := gob.NewEncoder(buf)
	if err := enc.Encode(src); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// deserialize decodes a value using gob.
func deserialize(src []byte, dst interface{}) error {
	dec := gob.NewDecoder(bytes.NewBuffer(src))
	if err := dec.Decode(dst); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 The prometheus-operator 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 main

import (
	"context"
	"fmt"
	"os"
	"regexp"
	"strings"

	"github.com/coreos/prometheus-operator/pkg/version"

	"github.com/go-kit/kit/log"
	"github.com/oklog/run"
	"github.com/thanos-io/thanos/pkg/reloader"
	"gopkg.in/alecthomas/kingpin.v2"
)

const (
	logFormatLogfmt                     = "logfmt"
	logFormatJson                       = "json"
	statefulsetOrdinalEnvvar            = "STATEFULSET_ORDINAL_NUMBER"
	statefulsetOrdinalFromEnvvarDefault = "POD_NAME"
)

var (
	availableLogFormats = []string{
		logFormatLogfmt,
		logFormatJson,
	}
)

func main() {
	app := kingpin.New("prometheus-config-reloader", "")
	cfgFile := app.Flag("config-file", "config file watched by the reloader").
		String()

	cfgSubstFile := app.Flag("config-envsubst-file", "output file for environment variable substituted config file").
		String()

	createStatefulsetOrdinalFrom := app.Flag(
		"statefulset-ordinal-from-envvar",
		fmt.Sprintf("parse this environment variable to create %s, containing the statefulset ordinal number", statefulsetOrdinalEnvvar)).
		Default(statefulsetOrdinalFromEnvvarDefault).String()

	logFormat := app.Flag(
		"log-format",
		fmt.Sprintf("Log format to use. Possible values: %s", strings.Join(availableLogFormats, ", "))).
		Default(logFormatLogfmt).String()

	reloadURL := app.Flag("reload-url", "reload URL to trigger Prometheus reload on").
		Default("http://127.0.0.1:9090/-/reload").URL()

	if _, err := app.Parse(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(2)
	}

	logger := log.NewLogfmtLogger(log.NewSyncWriter(os.Stdout))
	if *logFormat == logFormatJson {
		logger = log.NewJSONLogger(log.NewSyncWriter(os.Stdout))
	}
	logger = log.With(logger, "ts", log.DefaultTimestampUTC)
	logger = log.With(logger, "caller", log.DefaultCaller)

	if createStatefulsetOrdinalFrom != nil {
		if err := createOrdinalEnvvar(*createStatefulsetOrdinalFrom); err != nil {
			logger.Log("msg", fmt.Sprintf("Failed setting %s", statefulsetOrdinalEnvvar))
		}
	}

	logger.Log("msg", fmt.Sprintf("Starting prometheus-config-reloader version '%v'.", version.Version))

	var g run.Group
	{
		ctx, cancel := context.WithCancel(context.Background())
		rel := reloader.New(logger, *reloadURL, *cfgFile, *cfgSubstFile, []string{})

		g.Add(func() error {
			return rel.Watch(ctx)
		}, func(error) {
			cancel()
		})
	}

	if err := g.Run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func createOrdinalEnvvar(fromName string) error {
	reg := regexp.MustCompile(`\d+$`)
	val := reg.FindString(os.Getenv(fromName))
	return os.Setenv(statefulsetOrdinalEnvvar, val)
}
<commit_msg>prometheus-config-reloader: support for watching rules directory<commit_after>// Copyright 2016 The prometheus-operator 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 main

import (
	"context"
	"fmt"
	"os"
	"regexp"
	"strings"

	"github.com/coreos/prometheus-operator/pkg/version"

	"github.com/go-kit/kit/log"
	"github.com/oklog/run"
	"github.com/thanos-io/thanos/pkg/reloader"
	"gopkg.in/alecthomas/kingpin.v2"
)

const (
	logFormatLogfmt                     = "logfmt"
	logFormatJson                       = "json"
	statefulsetOrdinalEnvvar            = "STATEFULSET_ORDINAL_NUMBER"
	statefulsetOrdinalFromEnvvarDefault = "POD_NAME"
)

var (
	availableLogFormats = []string{
		logFormatLogfmt,
		logFormatJson,
	}
)

func main() {
	app := kingpin.New("prometheus-config-reloader", "")
	cfgFile := app.Flag("config-file", "config file watched by the reloader").
		String()

	cfgSubstFile := app.Flag("config-envsubst-file", "output file for environment variable substituted config file").
		String()

	rulesDir := app.Flag("rules-dir", "Rules directory to watch non-recursively").Strings()

	createStatefulsetOrdinalFrom := app.Flag(
		"statefulset-ordinal-from-envvar",
		fmt.Sprintf("parse this environment variable to create %s, containing the statefulset ordinal number", statefulsetOrdinalEnvvar)).
		Default(statefulsetOrdinalFromEnvvarDefault).String()

	logFormat := app.Flag(
		"log-format",
		fmt.Sprintf("Log format to use. Possible values: %s", strings.Join(availableLogFormats, ", "))).
		Default(logFormatLogfmt).String()

	reloadURL := app.Flag("reload-url", "reload URL to trigger Prometheus reload on").
		Default("http://127.0.0.1:9090/-/reload").URL()

	if _, err := app.Parse(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(2)
	}

	logger := log.NewLogfmtLogger(log.NewSyncWriter(os.Stdout))
	if *logFormat == logFormatJson {
		logger = log.NewJSONLogger(log.NewSyncWriter(os.Stdout))
	}
	logger = log.With(logger, "ts", log.DefaultTimestampUTC)
	logger = log.With(logger, "caller", log.DefaultCaller)

	if createStatefulsetOrdinalFrom != nil {
		if err := createOrdinalEnvvar(*createStatefulsetOrdinalFrom); err != nil {
			logger.Log("msg", fmt.Sprintf("Failed setting %s", statefulsetOrdinalEnvvar))
		}
	}

	logger.Log("msg", fmt.Sprintf("Starting prometheus-config-reloader version '%v'.", version.Version))

	var g run.Group
	{
		ctx, cancel := context.WithCancel(context.Background())
		rel := reloader.New(logger, *reloadURL, *cfgFile, *cfgSubstFile, *rulesDir)

		g.Add(func() error {
			return rel.Watch(ctx)
		}, func(error) {
			cancel()
		})
	}

	if err := g.Run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func createOrdinalEnvvar(fromName string) error {
	reg := regexp.MustCompile(`\d+$`)
	val := reg.FindString(os.Getenv(fromName))
	return os.Setenv(statefulsetOrdinalEnvvar, val)
}
<|endoftext|>


<commit_before><commit_msg>Add outdir option to pack and release<commit_after><|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/rzh/utils/go/mc/parser"
)

/*
type Testbed struct {
	Type    string                       `json:type`
	Servers map[string]map[string]string `json:"servers"`
}

type TestDriver struct {
	Version   string `json:version`
	GitSHA    string `json:git_hash`
	BuildDate string `json:build_date`
}

type DateTime struct {
	Date int64 `json:"$date"`
}

// definition of Stat to be reported to dashboard
type Stats struct {
	Harness       string                 `json:"harness"`
	Workload      string                 `json:"workload"`
	ServerVersion string                 `json:"server_version"`
	ServerGitSHA  string                 `json:"server_git_hash"`
	Attributes    map[string]interface{} `json:"attributes"`
	Testbed       Testbed                `json:"test_bed"`
	TestDriver    TestDriver             `json:"test_driver"`
	Summary       parser.StatsSummary    `json:"summary"`

	// TPS string
	// Run_Date   string
	Start_Time DateTime `json:"start_time"` //epoch time, time.Now().Unix()
	End_Time   DateTime `json:"end_time"`
	// ID         string
	// Type         string // hammertime, sysbench, mongo-sim
	History      []string
	Server_Stats map[string]parser.ServerStats `json:"server_stats"`
}
*/

func replaceDot(s string) string {
	return strings.Replace(s, ".", "_", -1)
}

// only used for clean report
type HammerTask_ struct {
	Run_id        string   `json:run_id`
	Cmd           string   `json:cmd`
	Hammer_folder string   `json:hammer_folder`
	Clients       []string `json:clients`
	Servers       []string `json:servers`

	// log files to be collected from client and server
	Client_logs []string `json:client_logs`
	Server_logs []string `json:server_logs`

	Type string `json:type`
}

func (r *TheRun) reportResults(run_id int, log_file string, run_dir string) {
	// this is the place to analyze results.
	t := strings.ToLower(r.Runs[run_id].Type)
	// r.Runs[run_id].Stats.Type = t
	r.Runs[run_id].Stats.Harness = t
	// r.Runs[run_id].Stats.ID = r.Runs[run_id].Run_id
	r.Runs[run_id].Stats.Workload = r.Runs[run_id].Run_id

	// cache run first
	//rr, _ := json.Marshal(r.Runs[run_id])
	rr := r.Runs[run_id]
	r.Runs[run_id].Stats.Attributes["run-by"] = "hammer-mc"
	r.Runs[run_id].Stats.Attributes["hammer-mc-cmd"] = HammerTask_{Run_id: rr.Run_id,
		Cmd: rr.Cmd, Clients: rr.Clients, Servers: rr.Servers,
		Client_logs: rr.Client_logs, Server_logs: rr.Server_logs,
		Type: rr.Type}

	if report_url == "" {
		report_url = "http://54.68.84.192:8080/api/v1/results"
		// report_url = "http://dyno.mongodb.parts/api/v1/results"
	}
	var err error
	switch t {
	case "sysbench":
		log.Println("Process sysbench results")
		cum, history, att := parser.ProcessSysbenchResult(log_file)

		// r.Runs[run_id].Stats.TPS = cum
		r.Runs[run_id].Stats.Summary.AllNodes.Op_throughput, err = strconv.ParseFloat(strings.Replace(cum, ",", "", -1), 64)
		if err != nil {
			log.Panicln("Error parsing op_throughput ", cum, ", error: ", err)
		}

		r.Runs[run_id].Stats.History = history

		// merge attribute into Stats
		for k, v := range att {
			r.Runs[run_id].Stats.Attributes[k] = v
		}

	case "mongo-sim":
		log.Println("Processing mongo-sim results")
		result_ := parser.ProcessMongoSIMResult(log_file)

		// need merge the two Stats together. Will copy
		r.Runs[run_id].Stats.Summary = result_.Summary
		// r.Runs[run_id].Stats.Testbed = result_.Testbed
		r.Runs[run_id].Stats.TestDriver = result_.TestDriver

		// merge attributes together
		for k, v := range result_.Attributes {
			if val, ok := r.Runs[run_id].Stats.Attributes[k]; ok {
				log.Println("Discard hammer-mc attribute[", k, "] = ", val, " with new value from mongo-sim ", v)
			}
			r.Runs[run_id].Stats.Attributes[k] = v
		}

	case "mongo-perf":
		log.Println("Processing mongo-perf results")
		result_ := parser.ProcessMongoPerfResult(log_file)

		for k, v := range result_ {
			_ = k
			_ = v
			/*
				Name    string
				Thread  int64
				Result  float64
				CV      string
				Version string
				GitSHA  string
			*/
			r.Runs[run_id].Stats.Workload = v.Name
			r.Runs[run_id].Stats.Attributes["nThread"] = v.Thread
			r.Runs[run_id].Stats.Attributes["CV"] = v.CV
			r.Runs[run_id].Stats.ServerVersion = strings.Fields(v.Version)[2]
			r.Runs[run_id].Stats.ServerGitSHA = v.GitSHA
			r.Runs[run_id].Stats.Summary.AllNodes.Op_throughput = v.Result

			// print
			s, _ := json.MarshalIndent(r.Runs[run_id].Stats, "  ", "    ")
			os.Stdout.Write(s)
			fmt.Println("\n********")

			// report to server
			if report_url != "" {
				// report to report_url if it is not empty
				r, err := http.Post(report_url, "application/json", bytes.NewBuffer(s))

				if err == nil {
					log.Println("Submit results to server succeeded with reponse:\n", r)
				} else {
					log.Panicln("Submit results failed with error: ", err)
				}
			}
		}

		// mongo-perf will not report server stats since it is not meaningful
		return

	default:
		log.Println("no type infor, ignore results analyzing")
	}

	// report pidstat here
	r.Runs[run_id].Stats.Server_Stats = make(map[string]parser.ServerStats)

	for k := 0; k < len(r.Runs[run_id].Servers); k++ {
		pidfile := run_dir + "/pidstat.log--" + r.Runs[run_id].Servers[k]
		stats := parser.ParsePIDStat(pidfile)

		// r.Runs[run_id].Stats.Server_Stats[replaceDot(r.Runs[run_id].Servers[k])] = make(ServerStats)
		// r.Runs[run_id].Stats.Server_Stats[r.Runs[run_id].Servers[k]][kk] = make([]parser.DataPoint, len(vv))
		// log.Println("++++> ", copy(r.Runs[run_id].Stats.Server_Stats[r.Runs[run_id].Servers[k]][kk], vv))
		//append(r.Runs[run_id].Stats.Server_Stats[r.Runs[run_id].Servers[k]][kk], vv)
		log.Println("\n\n++++++++-----+++++++++\n\n")
		//r.Runs[run_id].Stats.Server_Stats[replaceDot(r.Runs[run_id].Servers[k])].Cpu = make([]parser.DataPoint, len(stats["cpu"]), len(stats["cpu"]))
		// copy(r.Runs[run_id].Stats.Server_Stats[replaceDot(r.Runs[run_id].Servers[k])].Cpu, stats["cpu"])
		r.Runs[run_id].Stats.Server_Stats[replaceDot(r.Runs[run_id].Servers[k])] = stats
		log.Println(r.Runs[run_id].Stats.Server_Stats[replaceDot(r.Runs[run_id].Servers[k])])
	}

	s, _ := json.MarshalIndent(r.Runs[run_id].Stats, "  ", "    ")
	if report_url != "" {
		// report to report_url if it is not empty

		for _, rurl := range strings.Fields(report_url) {
			r, err := http.Post(rurl, "application/json", bytes.NewBuffer(s))

			if err == nil {
				log.Println("Submit results to server [", rurl, "] succeeded with reponse:\n", r)
			} else {
				log.Panicln("Submit results to server [", rurl, "] failed with error: ", err)
			}
		}
	}

	// print
	os.Stdout.Write(s)
	fmt.Println("\n********")
	//fmt.Printf("%# v", pretty.Formatter(r.Runs[run_id].Stats))
}
<commit_msg>fix url for report_url<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/rzh/utils/go/mc/parser"
)

/*
type Testbed struct {
	Type    string                       `json:type`
	Servers map[string]map[string]string `json:"servers"`
}

type TestDriver struct {
	Version   string `json:version`
	GitSHA    string `json:git_hash`
	BuildDate string `json:build_date`
}

type DateTime struct {
	Date int64 `json:"$date"`
}

// definition of Stat to be reported to dashboard
type Stats struct {
	Harness       string                 `json:"harness"`
	Workload      string                 `json:"workload"`
	ServerVersion string                 `json:"server_version"`
	ServerGitSHA  string                 `json:"server_git_hash"`
	Attributes    map[string]interface{} `json:"attributes"`
	Testbed       Testbed                `json:"test_bed"`
	TestDriver    TestDriver             `json:"test_driver"`
	Summary       parser.StatsSummary    `json:"summary"`

	// TPS string
	// Run_Date   string
	Start_Time DateTime `json:"start_time"` //epoch time, time.Now().Unix()
	End_Time   DateTime `json:"end_time"`
	// ID         string
	// Type         string // hammertime, sysbench, mongo-sim
	History      []string
	Server_Stats map[string]parser.ServerStats `json:"server_stats"`
}
*/

func replaceDot(s string) string {
	return strings.Replace(s, ".", "_", -1)
}

// only used for clean report
type HammerTask_ struct {
	Run_id        string   `json:run_id`
	Cmd           string   `json:cmd`
	Hammer_folder string   `json:hammer_folder`
	Clients       []string `json:clients`
	Servers       []string `json:servers`

	// log files to be collected from client and server
	Client_logs []string `json:client_logs`
	Server_logs []string `json:server_logs`

	Type string `json:type`
}

func (r *TheRun) reportResults(run_id int, log_file string, run_dir string) {
	// this is the place to analyze results.
	t := strings.ToLower(r.Runs[run_id].Type)
	// r.Runs[run_id].Stats.Type = t
	r.Runs[run_id].Stats.Harness = t
	// r.Runs[run_id].Stats.ID = r.Runs[run_id].Run_id
	r.Runs[run_id].Stats.Workload = r.Runs[run_id].Run_id

	// cache run first
	//rr, _ := json.Marshal(r.Runs[run_id])
	rr := r.Runs[run_id]
	r.Runs[run_id].Stats.Attributes["run-by"] = "hammer-mc"
	r.Runs[run_id].Stats.Attributes["hammer-mc-cmd"] = HammerTask_{Run_id: rr.Run_id,
		Cmd: rr.Cmd, Clients: rr.Clients, Servers: rr.Servers,
		Client_logs: rr.Client_logs, Server_logs: rr.Server_logs,
		Type: rr.Type}

	if report_url == "" {
		// report_url = "http://54.68.84.192:8080/api/v1/results"
		report_url = "http://dyno.mongodb.parts/api/v1/results"
	}
	var err error
	switch t {
	case "sysbench":
		log.Println("Process sysbench results")
		cum, history, att := parser.ProcessSysbenchResult(log_file)

		// r.Runs[run_id].Stats.TPS = cum
		r.Runs[run_id].Stats.Summary.AllNodes.Op_throughput, err = strconv.ParseFloat(strings.Replace(cum, ",", "", -1), 64)
		if err != nil {
			log.Panicln("Error parsing op_throughput ", cum, ", error: ", err)
		}

		r.Runs[run_id].Stats.History = history

		// merge attribute into Stats
		for k, v := range att {
			r.Runs[run_id].Stats.Attributes[k] = v
		}

	case "mongo-sim":
		log.Println("Processing mongo-sim results")
		result_ := parser.ProcessMongoSIMResult(log_file)

		// need merge the two Stats together. Will copy
		r.Runs[run_id].Stats.Summary = result_.Summary
		// r.Runs[run_id].Stats.Testbed = result_.Testbed
		r.Runs[run_id].Stats.TestDriver = result_.TestDriver

		// merge attributes together
		for k, v := range result_.Attributes {
			if val, ok := r.Runs[run_id].Stats.Attributes[k]; ok {
				log.Println("Discard hammer-mc attribute[", k, "] = ", val, " with new value from mongo-sim ", v)
			}
			r.Runs[run_id].Stats.Attributes[k] = v
		}

	case "mongo-perf":
		log.Println("Processing mongo-perf results")
		result_ := parser.ProcessMongoPerfResult(log_file)

		for k, v := range result_ {
			_ = k
			_ = v
			/*
				Name    string
				Thread  int64
				Result  float64
				CV      string
				Version string
				GitSHA  string
			*/
			r.Runs[run_id].Stats.Workload = v.Name
			r.Runs[run_id].Stats.Attributes["nThread"] = v.Thread
			r.Runs[run_id].Stats.Attributes["CV"] = v.CV
			r.Runs[run_id].Stats.ServerVersion = strings.Fields(v.Version)[2]
			r.Runs[run_id].Stats.ServerGitSHA = v.GitSHA
			r.Runs[run_id].Stats.Summary.AllNodes.Op_throughput = v.Result

			// print
			s, _ := json.MarshalIndent(r.Runs[run_id].Stats, "  ", "    ")
			os.Stdout.Write(s)
			fmt.Println("\n********")

			// report to server
			if report_url != "" {
				// report to report_url if it is not empty
				r, err := http.Post(report_url, "application/json", bytes.NewBuffer(s))

				if err == nil {
					log.Println("Submit results to server succeeded with reponse:\n", r)
				} else {
					log.Panicln("Submit results failed with error: ", err)
				}
			}
		}

		// mongo-perf will not report server stats since it is not meaningful
		return

	default:
		log.Println("no type infor, ignore results analyzing")
	}

	// report pidstat here
	r.Runs[run_id].Stats.Server_Stats = make(map[string]parser.ServerStats)

	for k := 0; k < len(r.Runs[run_id].Servers); k++ {
		pidfile := run_dir + "/pidstat.log--" + r.Runs[run_id].Servers[k]
		stats := parser.ParsePIDStat(pidfile)

		// r.Runs[run_id].Stats.Server_Stats[replaceDot(r.Runs[run_id].Servers[k])] = make(ServerStats)
		// r.Runs[run_id].Stats.Server_Stats[r.Runs[run_id].Servers[k]][kk] = make([]parser.DataPoint, len(vv))
		// log.Println("++++> ", copy(r.Runs[run_id].Stats.Server_Stats[r.Runs[run_id].Servers[k]][kk], vv))
		//append(r.Runs[run_id].Stats.Server_Stats[r.Runs[run_id].Servers[k]][kk], vv)
		log.Println("\n\n++++++++-----+++++++++\n\n")
		//r.Runs[run_id].Stats.Server_Stats[replaceDot(r.Runs[run_id].Servers[k])].Cpu = make([]parser.DataPoint, len(stats["cpu"]), len(stats["cpu"]))
		// copy(r.Runs[run_id].Stats.Server_Stats[replaceDot(r.Runs[run_id].Servers[k])].Cpu, stats["cpu"])
		r.Runs[run_id].Stats.Server_Stats[replaceDot(r.Runs[run_id].Servers[k])] = stats
		log.Println(r.Runs[run_id].Stats.Server_Stats[replaceDot(r.Runs[run_id].Servers[k])])
	}

	s, _ := json.MarshalIndent(r.Runs[run_id].Stats, "  ", "    ")
	if report_url != "" {
		// report to report_url if it is not empty

		for _, rurl := range strings.Fields(report_url) {
			r, err := http.Post(rurl, "application/json", bytes.NewBuffer(s))

			if err == nil {
				log.Println("Submit results to server [", rurl, "] succeeded with reponse:\n", r)
			} else {
				log.Panicln("Submit results to server [", rurl, "] failed with error: ", err)
			}
		}
	}

	// print
	os.Stdout.Write(s)
	fmt.Println("\n********")
	//fmt.Printf("%# v", pretty.Formatter(r.Runs[run_id].Stats))
}
<|endoftext|>


<commit_before>// Copyright © 2016-2021 Wei Shen <shenwei356@gmail.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"fmt"
	"strconv"
	"strings"
	"sync"

	"github.com/shenwei356/breader"
	"github.com/shenwei356/util/stringutil"
	"github.com/shenwei356/xopen"
	"github.com/spf13/cobra"
)

// flineageCmd represents the fx2tab command
var flineageCmd = &cobra.Command{
	Use:   "reformat",
	Short: "Reformat lineage in canonical ranks",
	Long: `Reformat lineage in canonical ranks

Output format can be formated by flag --format, available placeholders:

    {k}: superkingdom
    {p}: phylum
    {c}: class
    {o}: order
    {f}: family
    {g}: genus
    {s}: species
    {t}: subspecies/strain

    {S}: subspecies
    {T}: strain

Output format can contains some escape charactors like "\t".

This command appends reformated lineage to the input line.
The corresponding taxIDs of reformated lineage can be provided as another
column by flag "-t/--show-lineage-taxids".

`,
	Run: func(cmd *cobra.Command, args []string) {
		config := getConfigs(cmd)

		format := getFlagString(cmd, "format")
		delimiter := getFlagString(cmd, "delimiter")
		blank := getFlagString(cmd, "miss-rank-repl")
		prefix := getFlagString(cmd, "miss-rank-repl-prefix")
		iblank := getFlagString(cmd, "miss-taxid-repl")
		fill := getFlagBool(cmd, "fill-miss-rank")
		field := getFlagPositiveInt(cmd, "lineage-field") - 1
		printLineageInTaxid := getFlagBool(cmd, "show-lineage-taxids")

		addPrefix := getFlagBool(cmd, "add-prefix")
		prefixK := getFlagString(cmd, "prefix-k")
		prefixP := getFlagString(cmd, "prefix-p")
		prefixC := getFlagString(cmd, "prefix-c")
		prefixO := getFlagString(cmd, "prefix-o")
		prefixF := getFlagString(cmd, "prefix-f")
		prefixG := getFlagString(cmd, "prefix-g")
		prefixs := getFlagString(cmd, "prefix-s")
		prefixS := getFlagString(cmd, "prefix-S")
		prefixt := getFlagString(cmd, "prefix-t")
		prefixT := getFlagString(cmd, "prefix-T")

		trim := getFlagBool(cmd, "trim")

		prefixes := map[string]string{
			"k": prefixK,
			"p": prefixP,
			"c": prefixC,
			"o": prefixO,
			"f": prefixF,
			"g": prefixG,
			"s": prefixs,
			"S": prefixS,
			"t": prefixt,
			"T": prefixT,
		}
		// check format
		if !reRankPlaceHolder.MatchString(format) {
			checkError(fmt.Errorf("placeholder of simplified rank not found in output format: %s", format))
		}
		matches := reRankPlaceHolder.FindAllStringSubmatch(format, -1)
		for _, match := range matches {
			if _, ok := symbol2rank[match[1]]; !ok {
				checkError(fmt.Errorf("invalid placeholder: %s", match[0]))
			}
		}

		files := getFileList(args)

		if len(files) == 1 && isStdin(files[0]) && !xopen.IsStdin() {
			checkError(fmt.Errorf("stdin not detected"))
		}

		outfh, err := xopen.Wopen(config.OutFile)
		checkError(err)
		defer outfh.Close()

		taxid2taxon, name2parent2taxid, name2taxid := getName2Parent2Taxid(config)

		type line2flineage struct {
			line      string
			flineage  string
			iflineage string
		}

		unescape := stringutil.UnEscaper()

		var poolStrings = &sync.Pool{New: func() interface{} {
			return make([]string, 0, 32)
		}}

		fn := func(line string) (interface{}, bool, error) {
			if len(line) == 0 || line[0] == '#' {
				return nil, false, nil
			}
			line = strings.Trim(line, "\r\n ")
			if line == "" {
				return nil, false, nil
			}
			data := strings.Split(line, "\t")
			if len(data) < field+1 {
				return nil, false, fmt.Errorf("lineage-field (%d) out of range (%d):%s", field+1, len(data), line)
			}

			// names
			names := strings.Split(data[field], delimiter) // all names of full lineage

			// ranks := make([]string, len(names))
			ranks := poolStrings.Get().([]string)

			// sranks := make([]string, len(names))
			sranks := poolStrings.Get().([]string)

			var rank, srank string   // lower case of name : name
			var lname, plname string // lower case of name, rank and it's one-letter symbol
			var ok bool

			name2Name := make(map[string]string, len(names)) // lower case of name of parent

			srank2idx := make(map[string]int) // srank: index

			// preprare replacements.
			// find the orphan names and missing ranks
			replacements := make(map[string]string, len(matches))

			var ireplacements map[string]string
			if printLineageInTaxid {
				ireplacements = make(map[string]string, len(matches))
			}

			for _, match := range matches {
				replacements[match[1]] = blank
				if printLineageInTaxid {
					ireplacements[match[1]] = iblank
				}
			}

			var taxid uint32
			var maxRankWeight float32
			for i, name := range names {
				if name == "" {
					continue
				}

				lname = strings.ToLower(name)
				name2Name[lname] = name
				name = lname

				if _, ok = name2taxid[name]; !ok { // unofficial name
					log.Warningf(`unofficial taxon name detected: %s. Possible reasons: 1) lineages were produced with different taxonomy data files, please re-run taxonkit lineage; 2) some taxon names contain semicolon (";"), please re-run taxonkit lineage and taxonkit reformat with different flag value of -d, e.g., -d /`, name)
					return line2flineage{line, "", ""}, true, nil
				}

				if i == 0 { // root node
					taxid = name2taxid[name]
				} else {
					plname = strings.ToLower(names[i-1])
					if _, ok = name2parent2taxid[name]; !ok {
						log.Warningf(`unofficial taxon name detected: %s. Possible reasons: 1) lineages were produced with different taxonomy data files, please re-run taxonkit lineage; 2) some taxon names contain semicolon (";"), please re-run taxonkit lineage and taxonkit reformat with different flag value of -d, e.g., -d /`, name)
						return line2flineage{line, "", ""}, true, nil
					} else if taxid, ok = name2parent2taxid[name][plname]; !ok {
						log.Warningf(`unofficial taxon name detected: %s. Possible reasons: 1) lineages were produced with different taxonomy data files, please re-run taxonkit lineage; 2) some taxon names contain semicolon (";"), please re-run taxonkit lineage and taxonkit reformat with different flag value of -d, e.g., -d /`, plname)
						return line2flineage{line, "", ""}, true, nil
					}
				}
				// note that code below is computing rank of current name, not its parent.
				rank = taxid2taxon[taxid].Rank

				if rank == norank {
					continue
				}

				// ranks[i] = rank
				ranks = append(ranks, rank)
				if srank, ok = rank2symbol[rank]; ok {
					// special symbol "{t}"
					switch rank {
					case "strain":
						replacements["t"] = name2Name[name]
						if printLineageInTaxid {
							ireplacements["t"] = strconv.Itoa(int(taxid))
						}
						srank2idx["t"] = i
					case "subspecies":
						replacements["t"] = name2Name[name]
						if printLineageInTaxid {
							ireplacements["t"] = strconv.Itoa(int(taxid))
						}
						srank2idx["t"] = i
					}

					replacements[srank] = name2Name[name]
					if printLineageInTaxid {
						ireplacements[srank] = strconv.Itoa(int(taxid))
					}
					srank2idx[srank] = i
					// sranks[i] = srank
					sranks = append(sranks, srank)

					if trim && symbol2weight[srank] > maxRankWeight {
						maxRankWeight = symbol2weight[srank]
					}
				} else {
					sranks = append(sranks, "")
				}
			}

			if fill {
				var j, lastI int
				var srank2 string
				for _, srank = range srankList {
					if srank == "" {
						continue
					}

					if _, ok = srank2idx[srank]; ok {
						continue
					}

					if trim && symbol2weight[srank] > maxRankWeight {
						continue
					}

					// missing some ranks.
					// find the nearst higher formal rank
					for j, rank = range ranks {
						srank2 = sranks[j]
						if _, ok = srank2idx[srank2]; ok {
							if symbol2weight[srank2] < symbol2weight[srank] {
								lastI = j
							} else {
								break
							}
						}
					}
					replacements[srank] = prefix + names[lastI] + " " + symbol2rank[srank]
					// replacements[srank] = fmt.Sprintf("%s%s %s", prefix, names[lastI], symbol2rank[srank])
				}
			}

			flineage := format
			var iflineage string

			if printLineageInTaxid {
				iflineage = format
			}

			for srank, re := range reRankPlaceHolders {
				if addPrefix {
					flineage = re.ReplaceAllString(flineage, prefixes[srank]+replacements[srank])
				} else {
					flineage = re.ReplaceAllString(flineage, replacements[srank])
				}

				if printLineageInTaxid {
					iflineage = re.ReplaceAllString(iflineage, ireplacements[srank])
				}
			}

			// recycle
			ranks = ranks[:0]
			poolStrings.Put(ranks)
			sranks = sranks[:0]
			poolStrings.Put(sranks)

			return line2flineage{line, unescape(flineage), unescape(iflineage)}, true, nil
		}

		for _, file := range files {
			reader, err := breader.NewBufferedReader(file, config.Threads, 64, fn)
			checkError(err)

			var l2s line2flineage
			var data interface{}
			for chunk := range reader.Ch {
				checkError(chunk.Err)

				for _, data = range chunk.Data {
					l2s = data.(line2flineage)

					if printLineageInTaxid {
						outfh.WriteString(l2s.line + "\t" + l2s.flineage + "\t" + l2s.iflineage + "\n")
					} else {
						outfh.WriteString(l2s.line + "\t" + l2s.flineage + "\n")
					}
					if config.LineBuffered {
						outfh.Flush()
					}
				}
			}
		}
	},
}

func init() {
	RootCmd.AddCommand(flineageCmd)

	flineageCmd.Flags().StringP("format", "f", "{k};{p};{c};{o};{f};{g};{s}", "output format, placeholders of rank are needed")
	flineageCmd.Flags().StringP("delimiter", "d", ";", "field delimiter in input lineage")
	flineageCmd.Flags().StringP("miss-rank-repl", "r", "", `replacement string for missing rank`)
	flineageCmd.Flags().StringP("miss-rank-repl-prefix", "p", "unclassified ", `prefix for estimated taxon level`)
	flineageCmd.Flags().StringP("miss-taxid-repl", "R", "", `replacement string for missing taxid`)
	flineageCmd.Flags().BoolP("fill-miss-rank", "F", false, "fill missing rank with original lineage information (experimental)")
	flineageCmd.Flags().IntP("lineage-field", "i", 2, "field index of lineage. data should be tab-separated")
	flineageCmd.Flags().BoolP("show-lineage-taxids", "t", false, `show corresponding taxids of reformated lineage`)

	flineageCmd.Flags().BoolP("add-prefix", "P", false, `add prefixes for all ranks, single prefix for a rank is defined by flag --prefix-X`)
	flineageCmd.Flags().StringP("prefix-k", "", "k__", `prefix for superkingdom, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-p", "", "p__", `prefix for phylum, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-c", "", "c__", `prefix for class, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-o", "", "o__", `prefix for order, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-f", "", "f__", `prefix for family, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-g", "", "g__", `prefix for genus, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-s", "", "s__", `prefix for species, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-t", "", "t__", `prefix for subspecies/strain, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-S", "", "S__", `prefix for subspecies, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-T", "", "T__", `prefix for strain, used along with flag -P/--add-prefix`)

	flineageCmd.Flags().BoolP("trim", "T", false, "do not fill missing rank lower than current rank")
}
<commit_msg>reformat: fix a little bug<commit_after>// Copyright © 2016-2021 Wei Shen <shenwei356@gmail.com>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"fmt"
	"strconv"
	"strings"
	"sync"

	"github.com/shenwei356/breader"
	"github.com/shenwei356/util/stringutil"
	"github.com/shenwei356/xopen"
	"github.com/spf13/cobra"
)

// flineageCmd represents the fx2tab command
var flineageCmd = &cobra.Command{
	Use:   "reformat",
	Short: "Reformat lineage in canonical ranks",
	Long: `Reformat lineage in canonical ranks

Output format can be formated by flag --format, available placeholders:

    {k}: superkingdom
    {p}: phylum
    {c}: class
    {o}: order
    {f}: family
    {g}: genus
    {s}: species
    {t}: subspecies/strain

    {S}: subspecies
    {T}: strain

Output format can contains some escape charactors like "\t".

This command appends reformated lineage to the input line.
The corresponding taxIDs of reformated lineage can be provided as another
column by flag "-t/--show-lineage-taxids".

`,
	Run: func(cmd *cobra.Command, args []string) {
		config := getConfigs(cmd)

		format := getFlagString(cmd, "format")
		delimiter := getFlagString(cmd, "delimiter")
		blank := getFlagString(cmd, "miss-rank-repl")
		prefix := getFlagString(cmd, "miss-rank-repl-prefix")
		iblank := getFlagString(cmd, "miss-taxid-repl")
		fill := getFlagBool(cmd, "fill-miss-rank")
		field := getFlagPositiveInt(cmd, "lineage-field") - 1
		printLineageInTaxid := getFlagBool(cmd, "show-lineage-taxids")

		addPrefix := getFlagBool(cmd, "add-prefix")
		prefixK := getFlagString(cmd, "prefix-k")
		prefixP := getFlagString(cmd, "prefix-p")
		prefixC := getFlagString(cmd, "prefix-c")
		prefixO := getFlagString(cmd, "prefix-o")
		prefixF := getFlagString(cmd, "prefix-f")
		prefixG := getFlagString(cmd, "prefix-g")
		prefixs := getFlagString(cmd, "prefix-s")
		prefixS := getFlagString(cmd, "prefix-S")
		prefixt := getFlagString(cmd, "prefix-t")
		prefixT := getFlagString(cmd, "prefix-T")

		trim := getFlagBool(cmd, "trim")

		prefixes := map[string]string{
			"k": prefixK,
			"p": prefixP,
			"c": prefixC,
			"o": prefixO,
			"f": prefixF,
			"g": prefixG,
			"s": prefixs,
			"S": prefixS,
			"t": prefixt,
			"T": prefixT,
		}
		// check format
		if !reRankPlaceHolder.MatchString(format) {
			checkError(fmt.Errorf("placeholder of simplified rank not found in output format: %s", format))
		}
		matches := reRankPlaceHolder.FindAllStringSubmatch(format, -1)
		for _, match := range matches {
			if _, ok := symbol2rank[match[1]]; !ok {
				checkError(fmt.Errorf("invalid placeholder: %s", match[0]))
			}
		}

		files := getFileList(args)

		if len(files) == 1 && isStdin(files[0]) && !xopen.IsStdin() {
			checkError(fmt.Errorf("stdin not detected"))
		}

		outfh, err := xopen.Wopen(config.OutFile)
		checkError(err)
		defer outfh.Close()

		taxid2taxon, name2parent2taxid, name2taxid := getName2Parent2Taxid(config)

		type line2flineage struct {
			line      string
			flineage  string
			iflineage string
		}

		unescape := stringutil.UnEscaper()

		var poolStrings = &sync.Pool{New: func() interface{} {
			return make([]string, 0, 32)
		}}

		fn := func(line string) (interface{}, bool, error) {
			if len(line) == 0 || line[0] == '#' {
				return nil, false, nil
			}
			line = strings.Trim(line, "\r\n ")
			if line == "" {
				return nil, false, nil
			}
			data := strings.Split(line, "\t")
			if len(data) < field+1 {
				return nil, false, fmt.Errorf("lineage-field (%d) out of range (%d):%s", field+1, len(data), line)
			}

			// names
			names := strings.Split(data[field], delimiter) // all names of full lineage

			// ranks := make([]string, len(names))
			ranks := poolStrings.Get().([]string)

			// sranks := make([]string, len(names))
			sranks := poolStrings.Get().([]string)

			var rank, srank string   // lower case of name : name
			var lname, plname string // lower case of name, rank and it's one-letter symbol
			var ok bool

			name2Name := make(map[string]string, len(names)) // lower case of name of parent

			srank2idx := make(map[string]int) // srank: index

			// preprare replacements.
			// find the orphan names and missing ranks
			replacements := make(map[string]string, len(matches))

			var ireplacements map[string]string
			if printLineageInTaxid {
				ireplacements = make(map[string]string, len(matches))
			}

			for _, match := range matches {
				replacements[match[1]] = blank
				if printLineageInTaxid {
					ireplacements[match[1]] = iblank
				}
			}

			var taxid uint32
			var maxRankWeight float32
			for i, name := range names {
				if name == "" {
					continue
				}

				lname = strings.ToLower(name)
				name2Name[lname] = name
				name = lname

				if _, ok = name2taxid[name]; !ok { // unofficial name
					log.Warningf(`unofficial taxon name detected: %s. Possible reasons: 1) lineages were produced with different taxonomy data files, please re-run taxonkit lineage; 2) some taxon names contain semicolon (";"), please re-run taxonkit lineage and taxonkit reformat with different flag value of -d, e.g., -d /`, name)
					return line2flineage{line, "", ""}, true, nil
				}

				if i == 0 { // root node
					taxid = name2taxid[name]
				} else {
					plname = strings.ToLower(names[i-1])
					if _, ok = name2parent2taxid[name]; !ok {
						log.Warningf(`unofficial taxon name detected: %s. Possible reasons: 1) lineages were produced with different taxonomy data files, please re-run taxonkit lineage; 2) some taxon names contain semicolon (";"), please re-run taxonkit lineage and taxonkit reformat with different flag value of -d, e.g., -d /`, name)
						return line2flineage{line, "", ""}, true, nil
					} else if taxid, ok = name2parent2taxid[name][plname]; !ok {
						log.Warningf(`unofficial taxon name detected: %s. Possible reasons: 1) lineages were produced with different taxonomy data files, please re-run taxonkit lineage; 2) some taxon names contain semicolon (";"), please re-run taxonkit lineage and taxonkit reformat with different flag value of -d, e.g., -d /`, plname)
						return line2flineage{line, "", ""}, true, nil
					}
				}
				// note that code below is computing rank of current name, not its parent.
				rank = taxid2taxon[taxid].Rank

				if rank == norank {
					ranks = append(ranks, rank)
					sranks = append(sranks, "")
					continue
				}

				// ranks[i] = rank
				ranks = append(ranks, rank)
				if srank, ok = rank2symbol[rank]; ok {
					// special symbol "{t}"
					switch rank {
					case "strain":
						replacements["t"] = name2Name[name]
						if printLineageInTaxid {
							ireplacements["t"] = strconv.Itoa(int(taxid))
						}
						srank2idx["t"] = i
					case "subspecies":
						replacements["t"] = name2Name[name]
						if printLineageInTaxid {
							ireplacements["t"] = strconv.Itoa(int(taxid))
						}
						srank2idx["t"] = i
					}

					replacements[srank] = name2Name[name]
					if printLineageInTaxid {
						ireplacements[srank] = strconv.Itoa(int(taxid))
					}
					srank2idx[srank] = i
					// sranks[i] = srank
					sranks = append(sranks, srank)

					if trim && symbol2weight[srank] > maxRankWeight {
						maxRankWeight = symbol2weight[srank]
					}
				} else {
					sranks = append(sranks, "")
				}
			}

			if fill {
				var j, lastI int
				var srank2 string
				for _, srank = range srankList {
					if srank == "" {
						continue
					}

					if _, ok = srank2idx[srank]; ok {
						continue
					}

					if trim && symbol2weight[srank] > maxRankWeight {
						continue
					}

					// missing some ranks.
					// find the nearst higher formal rank
					for j, rank = range ranks {
						srank2 = sranks[j]
						if _, ok = srank2idx[srank2]; ok {
							if symbol2weight[srank2] < symbol2weight[srank] {
								lastI = j
							} else {
								break
							}
						}
					}
					replacements[srank] = prefix + names[lastI] + " " + symbol2rank[srank]
					// replacements[srank] = fmt.Sprintf("%s%s %s", prefix, names[lastI], symbol2rank[srank])
				}
			}

			flineage := format
			var iflineage string

			if printLineageInTaxid {
				iflineage = format
			}

			for srank, re := range reRankPlaceHolders {
				if addPrefix {
					flineage = re.ReplaceAllString(flineage, prefixes[srank]+replacements[srank])
				} else {
					flineage = re.ReplaceAllString(flineage, replacements[srank])
				}

				if printLineageInTaxid {
					iflineage = re.ReplaceAllString(iflineage, ireplacements[srank])
				}
			}

			// recycle
			ranks = ranks[:0]
			poolStrings.Put(ranks)
			sranks = sranks[:0]
			poolStrings.Put(sranks)

			return line2flineage{line, unescape(flineage), unescape(iflineage)}, true, nil
		}

		for _, file := range files {
			reader, err := breader.NewBufferedReader(file, config.Threads, 64, fn)
			checkError(err)

			var l2s line2flineage
			var data interface{}
			for chunk := range reader.Ch {
				checkError(chunk.Err)

				for _, data = range chunk.Data {
					l2s = data.(line2flineage)

					if printLineageInTaxid {
						outfh.WriteString(l2s.line + "\t" + l2s.flineage + "\t" + l2s.iflineage + "\n")
					} else {
						outfh.WriteString(l2s.line + "\t" + l2s.flineage + "\n")
					}
					if config.LineBuffered {
						outfh.Flush()
					}
				}
			}
		}
	},
}

func init() {
	RootCmd.AddCommand(flineageCmd)

	flineageCmd.Flags().StringP("format", "f", "{k};{p};{c};{o};{f};{g};{s}", "output format, placeholders of rank are needed")
	flineageCmd.Flags().StringP("delimiter", "d", ";", "field delimiter in input lineage")
	flineageCmd.Flags().StringP("miss-rank-repl", "r", "", `replacement string for missing rank`)
	flineageCmd.Flags().StringP("miss-rank-repl-prefix", "p", "unclassified ", `prefix for estimated taxon level`)
	flineageCmd.Flags().StringP("miss-taxid-repl", "R", "", `replacement string for missing taxid`)
	flineageCmd.Flags().BoolP("fill-miss-rank", "F", false, "fill missing rank with original lineage information (experimental)")
	flineageCmd.Flags().IntP("lineage-field", "i", 2, "field index of lineage. data should be tab-separated")
	flineageCmd.Flags().BoolP("show-lineage-taxids", "t", false, `show corresponding taxids of reformated lineage`)

	flineageCmd.Flags().BoolP("add-prefix", "P", false, `add prefixes for all ranks, single prefix for a rank is defined by flag --prefix-X`)
	flineageCmd.Flags().StringP("prefix-k", "", "k__", `prefix for superkingdom, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-p", "", "p__", `prefix for phylum, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-c", "", "c__", `prefix for class, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-o", "", "o__", `prefix for order, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-f", "", "f__", `prefix for family, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-g", "", "g__", `prefix for genus, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-s", "", "s__", `prefix for species, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-t", "", "t__", `prefix for subspecies/strain, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-S", "", "S__", `prefix for subspecies, used along with flag -P/--add-prefix`)
	flineageCmd.Flags().StringP("prefix-T", "", "T__", `prefix for strain, used along with flag -P/--add-prefix`)

	flineageCmd.Flags().BoolP("trim", "T", false, "do not fill missing rank lower than current rank")
}
<|endoftext|>


<commit_before>package dialects

import (
	"database/sql"
	"reflect"
	"strconv"
	"strings"

	"github.com/gernest/ngorm/model"
)

// Dialect interface contains behaviors that differ across SQL database
type Dialect interface {
	// GetName get dialect's name
	GetName() string

	// SetDB set db for dialect
	SetDB(db model.SQLCommon)

	// BindVar return the placeholder for actual values in SQL statements, in many dbs it is "?", Postgres using $1
	BindVar(i int) string
	// Quote quotes field name to avoid SQL parsing exceptions by using a reserved word as a field name
	Quote(key string) string
	// DataTypeOf return data's sql type
	DataTypeOf(field *model.StructField) (string, error)

	// HasIndex check has index or not
	HasIndex(tableName string, indexName string) bool
	// HasForeignKey check has foreign key or not
	HasForeignKey(tableName string, foreignKeyName string) bool
	// RemoveIndex remove index
	RemoveIndex(tableName string, indexName string) error
	// HasTable check has table or not
	HasTable(tableName string) bool
	// HasColumn check has column or not
	HasColumn(tableName string, columnName string) bool

	// LimitAndOffsetSQL return generated SQL with Limit and Offset, as mssql has special case
	LimitAndOffsetSQL(limit, offset interface{}) string
	// SelectFromDummyTable return select values, for most dbs, `SELECT values` just works, mysql needs `SELECT value FROM DUAL`
	SelectFromDummyTable() string
	// LastInsertIdReturningSuffix most dbs support LastInsertId, but postgres needs to use `RETURNING`
	LastInsertIDReturningSuffix(tableName, columnName string) string

	// BuildForeignKeyName returns a foreign key name for the given table, field and reference
	BuildForeignKeyName(tableName, field, dest string) string

	// CurrentDatabase return current database name
	CurrentDatabase() string
}

func ParseFieldStructForDialect(field *model.StructField) (fieldValue reflect.Value, sqlType string, size int, additionalType string) {
	// Get redirected field type
	var reflectType = field.Struct.Type
	for reflectType.Kind() == reflect.Ptr {
		reflectType = reflectType.Elem()
	}

	// Get redirected field value
	fieldValue = reflect.Indirect(reflect.New(reflectType))

	// Get scanner's real value
	var getScannerValue func(reflect.Value)
	getScannerValue = func(value reflect.Value) {
		fieldValue = value
		if _, isScanner := reflect.New(fieldValue.Type()).Interface().(sql.Scanner); isScanner && fieldValue.Kind() == reflect.Struct {
			getScannerValue(fieldValue.Field(0))
		}
	}
	getScannerValue(fieldValue)

	// Default Size
	if num, ok := field.TagSettings["SIZE"]; ok {
		size, _ = strconv.Atoi(num)
	} else {
		size = 255
	}

	// Default type from tag setting
	additionalType = field.TagSettings["NOT NULL"] + " " + field.TagSettings["UNIQUE"]
	if value, ok := field.TagSettings["DEFAULT"]; ok {
		additionalType = additionalType + " DEFAULT " + value
	}

	return fieldValue, field.TagSettings["TYPE"], size, strings.TrimSpace(additionalType)
}
<commit_msg>[dialects] Add godoc<commit_after>package dialects

import (
	"database/sql"
	"reflect"
	"strconv"
	"strings"

	"github.com/gernest/ngorm/model"
)

// Dialect interface contains behaviors that differ across SQL database
type Dialect interface {
	// GetName get dialect's name
	GetName() string

	// SetDB set db for dialect
	SetDB(db model.SQLCommon)

	// BindVar return the placeholder for actual values in SQL statements, in many dbs it is "?", Postgres using $1
	BindVar(i int) string
	// Quote quotes field name to avoid SQL parsing exceptions by using a reserved word as a field name
	Quote(key string) string
	// DataTypeOf return data's sql type
	DataTypeOf(field *model.StructField) (string, error)

	// HasIndex check has index or not
	HasIndex(tableName string, indexName string) bool
	// HasForeignKey check has foreign key or not
	HasForeignKey(tableName string, foreignKeyName string) bool
	// RemoveIndex remove index
	RemoveIndex(tableName string, indexName string) error
	// HasTable check has table or not
	HasTable(tableName string) bool
	// HasColumn check has column or not
	HasColumn(tableName string, columnName string) bool

	// LimitAndOffsetSQL return generated SQL with Limit and Offset, as mssql has special case
	LimitAndOffsetSQL(limit, offset interface{}) string
	// SelectFromDummyTable return select values, for most dbs, `SELECT values` just works, mysql needs `SELECT value FROM DUAL`
	SelectFromDummyTable() string
	// LastInsertIdReturningSuffix most dbs support LastInsertId, but postgres needs to use `RETURNING`
	LastInsertIDReturningSuffix(tableName, columnName string) string

	// BuildForeignKeyName returns a foreign key name for the given table, field and reference
	BuildForeignKeyName(tableName, field, dest string) string

	// CurrentDatabase return current database name
	CurrentDatabase() string
}

//ParseFieldStructForDialect pases metadatab enough to be used by dialects. The values
//returned are useful for implementing the DataOf method of the Dialect
//interface.
//
// The fieldValue returned is the value of the field. The sqlType value returned
// is the value specified in the tags for by TYPE key, size is the value of the
// SIZE tag key it defaults to 255 when not set.
func ParseFieldStructForDialect(field *model.StructField) (fieldValue reflect.Value, sqlType string, size int, additionalType string) {
	// Get redirected field type
	var reflectType = field.Struct.Type
	for reflectType.Kind() == reflect.Ptr {
		reflectType = reflectType.Elem()
	}

	// Get redirected field value
	fieldValue = reflect.Indirect(reflect.New(reflectType))

	// Get scanner's real value
	var getScannerValue func(reflect.Value)
	getScannerValue = func(value reflect.Value) {
		fieldValue = value
		if _, isScanner := reflect.New(fieldValue.Type()).Interface().(sql.Scanner); isScanner && fieldValue.Kind() == reflect.Struct {
			getScannerValue(fieldValue.Field(0))
		}
	}
	getScannerValue(fieldValue)

	// Default Size
	if num, ok := field.TagSettings["SIZE"]; ok {
		size, _ = strconv.Atoi(num)
	} else {
		size = 255
	}

	// Default type from tag setting
	additionalType = field.TagSettings["NOT NULL"] + " " + field.TagSettings["UNIQUE"]
	if value, ok := field.TagSettings["DEFAULT"]; ok {
		additionalType = additionalType + " DEFAULT " + value
	}

	return fieldValue, field.TagSettings["TYPE"], size, strings.TrimSpace(additionalType)
}
<|endoftext|>