<commit_before>/*
 * Copyright (c) 2013 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package miniprofiler_gae

import (
	"appengine"
	"appengine/memcache"
	"appengine/user"
	"appengine_internal"
	"fmt"
	"github.com/mjibson/MiniProfiler/go/miniprofiler"
	"github.com/mjibson/appstats"
	"net/http"
)

func init() {
	miniprofiler.Enable = EnableIfAdminOrDev
	miniprofiler.Get = GetMemcache
	miniprofiler.Store = StoreMemcache
	miniprofiler.MachineName = Instance
}

// EnableIfAdminOrDev returns true if this is the dev server or the current
// user is an admin. This is the default for miniprofiler.Enable.
func EnableIfAdminOrDev(r *http.Request) bool {
	if appengine.IsDevAppServer() {
		return true
	}
	c := appengine.NewContext(r)
	u := user.Current(c)
	return u.Admin
}

// Instance returns the app engine instance id, or the hostname on dev.
// This is the default for miniprofiler.MachineName.
func Instance() string {
	if i := appengine.InstanceID(); i != "" {
		return i[len(i)-8:]
	}
	return miniprofiler.Hostname()
}

// StoreMemcache stores the Profile in memcache. This is the default for
// miniprofiler.Store.
func StoreMemcache(r *http.Request, p *miniprofiler.Profile) {
	item := &memcache.Item{
		Key:   mp_key(string(p.Id)),
		Value: p.Json(),
	}
	c := appengine.NewContext(r)
	memcache.Set(c, item)
}

// GetMemcache gets the Profile from memcache. This is the default for
// miniprofiler.Get.
func GetMemcache(r *http.Request, id string) *miniprofiler.Profile {
	c := appengine.NewContext(r)
	item, err := memcache.Get(c, mp_key(id))
	if err != nil {
		return nil
	}
	return miniprofiler.ProfileFromJson(item.Value)
}

type Context struct {
	appstats.Context
	P *miniprofiler.Profile
}

func (c Context) Call(service, method string, in, out appengine_internal.ProtoMessage, opts *appengine_internal.CallOptions) error {
	err := c.Context.Call(service, method, in, out, opts)
	v := c.Context.Stats.RPCStats[len(c.Context.Stats.RPCStats)-1]
	c.P.AddCustomTiming("RPC", &miniprofiler.CustomTiming{
		StartMilliseconds:    float64(v.Offset.Nanoseconds()) / 1000000,
		DurationMilliseconds: float64(v.Duration.Nanoseconds()) / 1000000,
	})
	return err
}

// NewHandler returns a profiled, appstats-aware appengine.Context.
func NewHandler(f func(Context, http.ResponseWriter, *http.Request)) appstats.Handler {
	return appstats.NewHandler(func(c appengine.Context, w http.ResponseWriter, r *http.Request) {
		pc := Context{
			Context: c.(appstats.Context),
		}

		if miniprofiler.Enable(r) {
			pc.P = miniprofiler.NewProfile(w, r, miniprofiler.FuncName(f))
			f(pc, w, r)
			pc.P.CustomLink = pc.URL()
			pc.P.CustomLinkName = "appstats"
			pc.P.Finalize()
		} else {
			f(pc, w, r)
		}
	})
}

func mp_key(id string) string {
	return fmt.Sprintf("mini-profiler-results:%s", id)
}
<commit_msg>Fix app engine enable check - allow no user<commit_after>/*
 * Copyright (c) 2013 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package miniprofiler_gae

import (
	"appengine"
	"appengine/memcache"
	"appengine/user"
	"appengine_internal"
	"fmt"
	"github.com/mjibson/MiniProfiler/go/miniprofiler"
	"github.com/mjibson/appstats"
	"net/http"
)

func init() {
	miniprofiler.Enable = EnableIfAdminOrDev
	miniprofiler.Get = GetMemcache
	miniprofiler.Store = StoreMemcache
	miniprofiler.MachineName = Instance
}

// EnableIfAdminOrDev returns true if this is the dev server or the current
// user is an admin. This is the default for miniprofiler.Enable.
func EnableIfAdminOrDev(r *http.Request) bool {
	if appengine.IsDevAppServer() {
		return true
	}
	c := appengine.NewContext(r)
	if u := user.Current(c); u != nil {
		return u.Admin
	}
	return false
}

// Instance returns the app engine instance id, or the hostname on dev.
// This is the default for miniprofiler.MachineName.
func Instance() string {
	if i := appengine.InstanceID(); i != "" {
		return i[len(i)-8:]
	}
	return miniprofiler.Hostname()
}

// StoreMemcache stores the Profile in memcache. This is the default for
// miniprofiler.Store.
func StoreMemcache(r *http.Request, p *miniprofiler.Profile) {
	item := &memcache.Item{
		Key:   mp_key(string(p.Id)),
		Value: p.Json(),
	}
	c := appengine.NewContext(r)
	memcache.Set(c, item)
}

// GetMemcache gets the Profile from memcache. This is the default for
// miniprofiler.Get.
func GetMemcache(r *http.Request, id string) *miniprofiler.Profile {
	c := appengine.NewContext(r)
	item, err := memcache.Get(c, mp_key(id))
	if err != nil {
		return nil
	}
	return miniprofiler.ProfileFromJson(item.Value)
}

type Context struct {
	appstats.Context
	P *miniprofiler.Profile
}

func (c Context) Call(service, method string, in, out appengine_internal.ProtoMessage, opts *appengine_internal.CallOptions) error {
	err := c.Context.Call(service, method, in, out, opts)
	v := c.Context.Stats.RPCStats[len(c.Context.Stats.RPCStats)-1]
	c.P.AddCustomTiming("RPC", &miniprofiler.CustomTiming{
		StartMilliseconds:    float64(v.Offset.Nanoseconds()) / 1000000,
		DurationMilliseconds: float64(v.Duration.Nanoseconds()) / 1000000,
	})
	return err
}

// NewHandler returns a profiled, appstats-aware appengine.Context.
func NewHandler(f func(Context, http.ResponseWriter, *http.Request)) appstats.Handler {
	return appstats.NewHandler(func(c appengine.Context, w http.ResponseWriter, r *http.Request) {
		pc := Context{
			Context: c.(appstats.Context),
		}

		if miniprofiler.Enable(r) {
			pc.P = miniprofiler.NewProfile(w, r, miniprofiler.FuncName(f))
			f(pc, w, r)
			pc.P.CustomLink = pc.URL()
			pc.P.CustomLinkName = "appstats"
			pc.P.Finalize()
		} else {
			f(pc, w, r)
		}
	})
}

func mp_key(id string) string {
	return fmt.Sprintf("mini-profiler-results:%s", id)
}
<|endoftext|>


<commit_before>package bongo

import (
	"encoding/json"
	"fmt"
	"strconv"
	"strings"

	"github.com/jinzhu/gorm"
)

// Fetch tries the best option available for retrieving the data first tries
// cache layer, if fails, then tries db
func (b *Bongo) Fetch(i Modellable, id int64) error {
	if id == 0 {
		return IdIsNotSet
	}

	data, ok := i.(Cacher)
	if !ok {
		return CacherIsNotImplemented
	}

	return b.GetFromBest(i, data, id)
}

// ById Fetches data from db by it's id
func (b *Bongo) ById(i Modellable, id int64) error {
	if err := b.DB.
		Table(i.BongoName()).
		Where("id = ?", id).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

func (b *Bongo) Unscoped() *gorm.DB {
	return b.DB.Unscoped()
}

func (b *Bongo) UnscopedById(i Modellable, id int64) error {
	if err := b.DB.
		Unscoped().
		Table(i.BongoName()).
		Where("id = ?", id).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// Creates a new record with the given struct and its fields
func (b *Bongo) Create(i Modellable) error {
	if err := b.DB.Save(i).Error; err != nil {
		return err
	}

	return nil
}

// Update updates all fields of a struct with assigned data
func (b *Bongo) Update(i Modellable) error {
	if i.GetId() == 0 {
		return IdIsNotSet
	}

	// Update and Create is using the Save method, so they are
	// same functions but GORM handles, AfterCreate and AfterUpdate
	// in correct manner
	return b.Create(i)
}

// Delete deletes the data by it's id, it doesnt take any other fields
// into consideration
func (b *Bongo) Delete(i Modellable) error {
	if i.GetId() == 0 {
		return IdIsNotSet
	}

	if err := b.DB.Delete(i).Error; err != nil {
		return err
	}

	return nil
}

// Purge deletes all data, and must be used with care.
func (b *Bongo) Purge(i Modellable) error {
	return b.DB.Delete(i).Error
}

// FetchByIds fetches records by their ids and returns results in the same order
// as the ids; if no records in db we don't return error
func (b *Bongo) FetchByIds(i Modellable, data interface{}, ids []int64) error {
	if len(ids) == 0 {
		return nil
	}

	orderByQuery := ""
	comma := ""
	for _, id := range ids {
		orderByQuery = orderByQuery + comma + " id = " + strconv.FormatInt(id, 10) + " desc"
		comma = ","
	}

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.BongoName())

	query = query.Order(orderByQuery)

	query = query.Where(ids)

	query = query.Find(data)

	// supress not found errors
	return CheckErr(query)

}

// This function doesnt work, after fixing we can open again
// func (b *Bongo) UpdatePartial(i Modellable, set map[string]interface{}) error {
// 	if i.GetId() == 0 {
// 		return IdIsNotSet
// 	}

// 	// init query
// 	query := b.DB

// 	query = query.Model(i)
// 	// query = query.Table(i.BongoName())

// 	query = query.Where("id = ? ", i.GetId())
// 	fmt.Println("query-->", query)
// 	fmt.Printf("set %# v", pretty.Formatter(set))
// 	if err := query.Update(set).Error; err != nil {
// 		fmt.Printf("%# v", pretty.Formatter(err))
// 		return err
// 	}

// 	if err := b.ById(i, i.GetId()); err != nil {
// 		fmt.Printf("%# v", pretty.Formatter(err))
// 		return err
// 	}

// 	fmt.Printf("i %# v", pretty.Formatter(i))
// 	// b.AfterUpdate(i)
// 	return nil
// }

// selector, set
func (b *Bongo) UpdateMulti(i Modellable, rest ...map[string]interface{}) error {
	var set, selector map[string]interface{}

	switch len(rest) {
	case 1:
		set = rest[0]
		selector = nil
	case 2:
		selector = rest[0]
		set = rest[1]
	default:
		return WrongParameter
	}

	query := b.DB.Table(i.BongoName())

	//add selector
	query = addWhere(query, selector)

	if err := query.Updates(set).Error; err != nil {
		return err
	}

	return nil
}

func (b *Bongo) Count(i Modellable, where ...interface{}) (int, error) {
	var count int

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.BongoName())

	// add query
	query = query.Where(where[0], where[1:len(where)]...)

	return count, query.Count(&count).Error
}

func (b *Bongo) CountWithQuery(i Modellable, q *Query) (int, error) {
	query := b.BuildQuery(i, q)
	var count int
	return count, query.Count(&count).Error
}

type Scope func(d *gorm.DB) *gorm.DB

type Query struct {
	Selector   map[string]interface{}
	Sort       map[string]string
	Pluck      string
	Pagination Pagination
	Scopes     []Scope
}

func (q *Query) AddScope(scope Scope) {
	if q.Scopes == nil {
		q.Scopes = make([]Scope, 0)
	}

	q.Scopes = append(q.Scopes, scope)
}

type Pagination struct {
	Limit int
	Skip  int
}

func NewPagination(limit int, skip int) *Pagination {
	return &Pagination{
		Limit: limit,
		Skip:  skip,
	}
}

func NewQS(selector map[string]interface{}) *Query {
	return &Query{
		Selector: selector,
	}
}

// selector, sort, limit, pluck,
func (b *Bongo) Some(i Modellable, data interface{}, q *Query) error {
	err := b.executeQuery(i, data, q)
	if err == gorm.RecordNotFound {
		return nil
	}
	return err
}

func (b *Bongo) One(i Modellable, data interface{}, q *Query) error {
	q.Pagination.Limit = 1
	return b.executeQuery(i, data, q)
}

func (b *Bongo) BuildQuery(i Modellable, q *Query) *gorm.DB {
	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.BongoName())

	// add sort options
	query = addSort(query, q.Sort)

	query = addSkip(query, q.Pagination.Skip)

	query = addLimit(query, q.Pagination.Limit)

	// add selector
	query = addWhere(query, q.Selector)

	// put scopes
	if q.Scopes != nil && len(q.Scopes) > 0 {
		for _, scope := range q.Scopes {
			query = query.Scopes(scope)
		}
	}

	return query
}

func (b *Bongo) executeQuery(i Modellable, data interface{}, q *Query) error {
	// init query
	query := b.BuildQuery(i, q)

	var err error
	// TODO refactor this part
	if q.Pluck != "" {
		if strings.Contains(q.Pluck, ",") {
			// add pluck data
			query = addPluck(query, q.Pluck)

			err = query.Find(data).Error
		} else {
			err = query.Pluck(q.Pluck, data).Error
		}
	} else {
		err = query.Find(data).Error
	}

	return err
}

func (b *Bongo) PublishEvent(eventName string, i Modellable) error {
	return b.Emit(i.BongoName()+"_"+eventName, i)
}

func (b *Bongo) Emit(eventName string, i interface{}) error {
	data, err := json.Marshal(i)
	if err != nil {
		b.log.Error("Error while marshalling for emitting %s", err)
		return err
	}

	err = b.Broker.Publish(eventName, data)
	if err != nil {
		b.log.Error("Error while emitting %s", err)
		return err
	}

	return nil
}

func (b *Bongo) AfterCreate(i Modellable) error {
	b.PublishEvent("created", i)
	return b.AddToCache(i)
}

func (b *Bongo) AfterUpdate(i Modellable) error {
	b.PublishEvent("updated", i)
	return b.AddToCache(i)
}

func (b *Bongo) AfterDelete(i Modellable) error {
	b.PublishEvent("deleted", i)
	return b.AddToCache(i)
}

func (b *Bongo) AddToCache(i Modellable) error {
	data, ok := i.(Cacher)
	if !ok {
		return nil
	}

	err := b.SetToCache(data)
	if err != ErrCacheIsNotEnabled {
		return err
	}

	return nil
}

// addSort injects sort parameters into query
func addSort(query *gorm.DB, options map[string]string) *gorm.DB {

	if options == nil {
		return query
	}

	if len(options) == 0 {
		return query
	}

	var opts []string
	for key, val := range options {
		opts = append(opts, fmt.Sprintf("%s %v", key, val))
	}
	return query.Order(strings.Join(opts, ","))
}

// addPluck basically adds select statement for
// only required fields
func addPluck(query *gorm.DB, plucked string) *gorm.DB {
	if plucked == "" {
		return query
	}

	return query.Select(plucked)
}

// addWhere adds where query
func addWhere(query *gorm.DB, selector map[string]interface{}) *gorm.DB {
	if selector == nil {
		return query
	}

	// instead sending one selector, do chaining here
	return query.Where(selector)
}

// addSkip adds skip parameter into sql query
func addSkip(query *gorm.DB, skip int) *gorm.DB {
	if skip > 0 {
		return query.Offset(skip)
	}

	return query
}

// addLimit adds limit into query if set
func addLimit(query *gorm.DB, limit int) *gorm.DB {
	// if limit is minus or 0 ignore
	if limit > 0 {
		return query.Limit(limit)
	}

	return query
}

// CheckErr checks error exitence and returns if found, but this function
// suppress RecordNotFound errors
func CheckErr(res *gorm.DB) error {
	if res == nil {
		return nil
	}

	if res.Error == nil {
		return nil
	}

	if res.Error == gorm.RecordNotFound {
		return nil
	}

	return res.Error
}
<commit_msg>bongo: add table name before creating data in db<commit_after>package bongo

import (
	"encoding/json"
	"fmt"
	"strconv"
	"strings"

	"github.com/jinzhu/gorm"
)

// Fetch tries the best option available for retrieving the data first tries
// cache layer, if fails, then tries db
func (b *Bongo) Fetch(i Modellable, id int64) error {
	if id == 0 {
		return IdIsNotSet
	}

	data, ok := i.(Cacher)
	if !ok {
		return CacherIsNotImplemented
	}

	return b.GetFromBest(i, data, id)
}

// ById Fetches data from db by it's id
func (b *Bongo) ById(i Modellable, id int64) error {
	if err := b.DB.
		Table(i.BongoName()).
		Where("id = ?", id).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

func (b *Bongo) Unscoped() *gorm.DB {
	return b.DB.Unscoped()
}

func (b *Bongo) UnscopedById(i Modellable, id int64) error {
	if err := b.DB.
		Unscoped().
		Table(i.BongoName()).
		Where("id = ?", id).
		Find(i).
		Error; err != nil {
		return err
	}

	return nil
}

// Creates a new record with the given struct and its fields
func (b *Bongo) Create(i Modellable) error {
	if err := b.DB.Table(i.BongoName()).Save(i).Error; err != nil {
		return err
	}

	return nil
}

// Update updates all fields of a struct with assigned data
func (b *Bongo) Update(i Modellable) error {
	if i.GetId() == 0 {
		return IdIsNotSet
	}

	// Update and Create is using the Save method, so they are
	// same functions but GORM handles, AfterCreate and AfterUpdate
	// in correct manner
	return b.Create(i)
}

// Delete deletes the data by it's id, it doesnt take any other fields
// into consideration
func (b *Bongo) Delete(i Modellable) error {
	if i.GetId() == 0 {
		return IdIsNotSet
	}

	if err := b.DB.Table(i.BongoName()).Delete(i).Error; err != nil {
		return err
	}

	return nil
}

// Purge deletes all data, and must be used with care.
func (b *Bongo) Purge(i Modellable) error {
	return b.DB.Table(i.BongoName()).Delete(i).Error
}

// FetchByIds fetches records by their ids and returns results in the same order
// as the ids; if no records in db we don't return error
func (b *Bongo) FetchByIds(i Modellable, data interface{}, ids []int64) error {
	if len(ids) == 0 {
		return nil
	}

	orderByQuery := ""
	comma := ""
	for _, id := range ids {
		orderByQuery = orderByQuery + comma + " id = " + strconv.FormatInt(id, 10) + " desc"
		comma = ","
	}

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.BongoName())

	query = query.Order(orderByQuery)

	query = query.Where(ids)

	query = query.Find(data)

	// supress not found errors
	return CheckErr(query)

}

// This function doesnt work, after fixing we can open again
// func (b *Bongo) UpdatePartial(i Modellable, set map[string]interface{}) error {
// 	if i.GetId() == 0 {
// 		return IdIsNotSet
// 	}

// 	// init query
// 	query := b.DB

// 	query = query.Model(i)
// 	// query = query.Table(i.BongoName())

// 	query = query.Where("id = ? ", i.GetId())
// 	fmt.Println("query-->", query)
// 	fmt.Printf("set %# v", pretty.Formatter(set))
// 	if err := query.Update(set).Error; err != nil {
// 		fmt.Printf("%# v", pretty.Formatter(err))
// 		return err
// 	}

// 	if err := b.ById(i, i.GetId()); err != nil {
// 		fmt.Printf("%# v", pretty.Formatter(err))
// 		return err
// 	}

// 	fmt.Printf("i %# v", pretty.Formatter(i))
// 	// b.AfterUpdate(i)
// 	return nil
// }

// selector, set
func (b *Bongo) UpdateMulti(i Modellable, rest ...map[string]interface{}) error {
	var set, selector map[string]interface{}

	switch len(rest) {
	case 1:
		set = rest[0]
		selector = nil
	case 2:
		selector = rest[0]
		set = rest[1]
	default:
		return WrongParameter
	}

	query := b.DB.Table(i.BongoName())

	//add selector
	query = addWhere(query, selector)

	if err := query.Updates(set).Error; err != nil {
		return err
	}

	return nil
}

func (b *Bongo) Count(i Modellable, where ...interface{}) (int, error) {
	var count int

	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.BongoName())

	// add query
	query = query.Where(where[0], where[1:len(where)]...)

	return count, query.Count(&count).Error
}

func (b *Bongo) CountWithQuery(i Modellable, q *Query) (int, error) {
	query := b.BuildQuery(i, q)
	var count int
	return count, query.Count(&count).Error
}

type Scope func(d *gorm.DB) *gorm.DB

type Query struct {
	Selector   map[string]interface{}
	Sort       map[string]string
	Pluck      string
	Pagination Pagination
	Scopes     []Scope
}

func (q *Query) AddScope(scope Scope) {
	if q.Scopes == nil {
		q.Scopes = make([]Scope, 0)
	}

	q.Scopes = append(q.Scopes, scope)
}

type Pagination struct {
	Limit int
	Skip  int
}

func NewPagination(limit int, skip int) *Pagination {
	return &Pagination{
		Limit: limit,
		Skip:  skip,
	}
}

func NewQS(selector map[string]interface{}) *Query {
	return &Query{
		Selector: selector,
	}
}

// selector, sort, limit, pluck,
func (b *Bongo) Some(i Modellable, data interface{}, q *Query) error {
	err := b.executeQuery(i, data, q)
	if err == gorm.RecordNotFound {
		return nil
	}
	return err
}

func (b *Bongo) One(i Modellable, data interface{}, q *Query) error {
	q.Pagination.Limit = 1
	return b.executeQuery(i, data, q)
}

func (b *Bongo) BuildQuery(i Modellable, q *Query) *gorm.DB {
	// init query
	query := b.DB.Model(i)

	// add table name
	query = query.Table(i.BongoName())

	// add sort options
	query = addSort(query, q.Sort)

	query = addSkip(query, q.Pagination.Skip)

	query = addLimit(query, q.Pagination.Limit)

	// add selector
	query = addWhere(query, q.Selector)

	// put scopes
	if q.Scopes != nil && len(q.Scopes) > 0 {
		for _, scope := range q.Scopes {
			query = query.Scopes(scope)
		}
	}

	return query
}

func (b *Bongo) executeQuery(i Modellable, data interface{}, q *Query) error {
	// init query
	query := b.BuildQuery(i, q)

	var err error
	// TODO refactor this part
	if q.Pluck != "" {
		if strings.Contains(q.Pluck, ",") {
			// add pluck data
			query = addPluck(query, q.Pluck)

			err = query.Find(data).Error
		} else {
			err = query.Pluck(q.Pluck, data).Error
		}
	} else {
		err = query.Find(data).Error
	}

	return err
}

func (b *Bongo) PublishEvent(eventName string, i Modellable) error {
	return b.Emit(i.BongoName()+"_"+eventName, i)
}

func (b *Bongo) Emit(eventName string, i interface{}) error {
	data, err := json.Marshal(i)
	if err != nil {
		b.log.Error("Error while marshalling for emitting %s", err)
		return err
	}

	err = b.Broker.Publish(eventName, data)
	if err != nil {
		b.log.Error("Error while emitting %s", err)
		return err
	}

	return nil
}

func (b *Bongo) AfterCreate(i Modellable) error {
	b.PublishEvent("created", i)
	return b.AddToCache(i)
}

func (b *Bongo) AfterUpdate(i Modellable) error {
	b.PublishEvent("updated", i)
	return b.AddToCache(i)
}

func (b *Bongo) AfterDelete(i Modellable) error {
	b.PublishEvent("deleted", i)
	return b.AddToCache(i)
}

func (b *Bongo) AddToCache(i Modellable) error {
	data, ok := i.(Cacher)
	if !ok {
		return nil
	}

	err := b.SetToCache(data)
	if err != ErrCacheIsNotEnabled {
		return err
	}

	return nil
}

// addSort injects sort parameters into query
func addSort(query *gorm.DB, options map[string]string) *gorm.DB {

	if options == nil {
		return query
	}

	if len(options) == 0 {
		return query
	}

	var opts []string
	for key, val := range options {
		opts = append(opts, fmt.Sprintf("%s %v", key, val))
	}
	return query.Order(strings.Join(opts, ","))
}

// addPluck basically adds select statement for
// only required fields
func addPluck(query *gorm.DB, plucked string) *gorm.DB {
	if plucked == "" {
		return query
	}

	return query.Select(plucked)
}

// addWhere adds where query
func addWhere(query *gorm.DB, selector map[string]interface{}) *gorm.DB {
	if selector == nil {
		return query
	}

	// instead sending one selector, do chaining here
	return query.Where(selector)
}

// addSkip adds skip parameter into sql query
func addSkip(query *gorm.DB, skip int) *gorm.DB {
	if skip > 0 {
		return query.Offset(skip)
	}

	return query
}

// addLimit adds limit into query if set
func addLimit(query *gorm.DB, limit int) *gorm.DB {
	// if limit is minus or 0 ignore
	if limit > 0 {
		return query.Limit(limit)
	}

	return query
}

// CheckErr checks error exitence and returns if found, but this function
// suppress RecordNotFound errors
func CheckErr(res *gorm.DB) error {
	if res == nil {
		return nil
	}

	if res.Error == nil {
		return nil
	}

	if res.Error == gorm.RecordNotFound {
		return nil
	}

	return res.Error
}
<|endoftext|>


<commit_before>package terminal

import (
	"bytes"
	"errors"
	"fmt"
	"koding/tools/dnode"
	"koding/tools/kite"
	"koding/tools/pty"
	"koding/tools/utils"
	"koding/virt"
	"os"
	"strconv"
	"strings"
	"syscall"
	"time"
	"unicode/utf8"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

const (
	sessionPrefix     = "koding"
	kodingScreenPath  = "/opt/koding/bin/screen"
	kodingScreenrc    = "/opt/koding/etc/screenrc"
	defaultScreenPath = "/usr/bin/screen"
)

type WebtermServer struct {
	Session          string `json:"session"`
	remote           WebtermRemote
	vm               *virt.VM
	user             *virt.User
	isForeignSession bool
	pty              *pty.PTY
	currentSecond    int64
	messageCounter   int
	byteCounter      int
	lineFeeedCounter int
	screenPath       string
}

type WebtermRemote struct {
	Output       dnode.Callback
	SessionEnded dnode.Callback
}

func webtermKillSession(args *dnode.Partial, channel *kite.Channel, vos *virt.VOS) (interface{}, error) {
	sessions := screenSessions(vos)
	if len(sessions) == 0 {
		return nil, errors.New("no sessions available")
	}

	return sessions, nil
}

func webtermGetSessions(args *dnode.Partial, channel *kite.Channel, vos *virt.VOS) (interface{}, error) {
	sessions := screenSessions(vos)
	if len(sessions) == 0 {
		return nil, errors.New("no sessions available")
	}

	return sessions, nil
}

// this method is special cased in oskite.go to allow foreign access
func webtermConnect(args *dnode.Partial, channel *kite.Channel, vos *virt.VOS) (interface{}, error) {
	var params struct {
		Remote       WebtermRemote
		Session      string
		SizeX, SizeY int
		Mode         string
		JoinUser     string
	}

	if args == nil {
		return nil, &kite.ArgumentError{Expected: "empty argument passed"}
	}

	if args.Unmarshal(&params) != nil || params.SizeX <= 0 || params.SizeY <= 0 {
		return nil, &kite.ArgumentError{Expected: "{ remote: [object], session: [string], sizeX: [integer], sizeY: [integer], noScreen: [boolean] }"}
	}

	if params.JoinUser != "" {
		if len(params.Session) != utils.RandomStringLength {
			return nil, &kite.BaseError{
				Message: "Invalid session identifier",
				CodeErr: ErrInvalidSession,
			}
		}

		user := new(virt.User)
		if err := mongodbConn.Run("jUsers", func(c *mgo.Collection) error {
			return c.Find(bson.M{"username": params.JoinUser}).One(user)
		}); err != nil {
			return nil, err
		}

		// vos.VM is replaced already in registerMethod via
		// channel.CorrelationName which is the remote VM hostnameAlias
		vos.User = user
	}

	screen, err := newScreen(vos, params.Mode, params.Session)
	if err != nil {
		return nil, err
	}

	server := &WebtermServer{
		Session:          screen.Session,
		remote:           params.Remote,
		vm:               vos.VM,
		user:             vos.User,
		isForeignSession: vos.User.Name != channel.Username,
		pty:              pty.New(vos.VM.PtsDir()),
		screenPath:       screen.ScreenPath,
	}

	server.SetSize(float64(params.SizeX), float64(params.SizeY))
	server.pty.Slave.Chown(vos.User.Uid, -1)

	cmd := vos.VM.AttachCommand(vos.User.Uid, "/dev/pts/"+strconv.Itoa(server.pty.No), screen.Command...)
	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	go func() {
		defer log.RecoverAndLog()

		cmd.Wait()
		server.pty.Slave.Close()
		server.pty.Master.Close()
		server.remote.SessionEnded()
	}()

	go func() {
		defer log.RecoverAndLog()

		buf := make([]byte, (1<<12)-utf8.UTFMax, 1<<12)
		for {
			n, err := server.pty.Master.Read(buf)
			for n < cap(buf)-1 {
				r, _ := utf8.DecodeLastRune(buf[:n])
				if r != utf8.RuneError {
					break
				}
				server.pty.Master.Read(buf[n : n+1])
				n++
			}

			s := time.Now().Unix()
			if server.currentSecond != s {
				server.currentSecond = s
				server.messageCounter = 0
				server.byteCounter = 0
				server.lineFeeedCounter = 0
			}
			server.messageCounter += 1
			server.byteCounter += n
			server.lineFeeedCounter += bytes.Count(buf[:n], []byte{'\n'})
			if server.messageCounter > 100 || server.byteCounter > 1<<18 || server.lineFeeedCounter > 300 {
				time.Sleep(time.Second / 100)
			}

			server.remote.Output(string(utils.FilterInvalidUTF8(buf[:n])))
			if err != nil {
				break
			}
		}
	}()

	channel.OnDisconnect(func() { server.Close() })

	return server, nil
}

func (server *WebtermServer) Input(data string) {
	server.pty.Master.Write([]byte(data))
}

func (server *WebtermServer) ControlSequence(data string) {
	server.pty.MasterEncoded.Write([]byte(data))
}

func (server *WebtermServer) SetSize(x, y float64) {
	server.pty.SetSize(uint16(x), uint16(y))
}

func (server *WebtermServer) Close() error {
	server.pty.Signal(syscall.SIGHUP)
	return nil
}

func (server *WebtermServer) Terminate() error {
	server.Close()
	if !server.isForeignSession {
		server.vm.AttachCommand(server.user.Uid, "", server.screenPath, "-S", sessionPrefix+"."+server.Session, "-X", "quit").Run()
	}
	return nil
}

type screen struct {
	// Binary used for starting the screen
	ScreenPath string

	// Used for remote or multiuser mode, defines the custom session name
	Session string

	// the final command to be executed
	Command []string
}

func getScreenPath(vos *virt.VOS) (string, error) {
	screenPath := kodingScreenPath

	// it can happen that the user deleted our screen binary
	// accidently, if this happens fallback to default screen binary
	_, err := vos.Stat(kodingScreenPath)
	if os.IsNotExist(err) {
		// check if the default screen binary exists too
		_, err := vos.Stat(defaultScreenPath)
		if os.IsNotExist(err) {
			return "", &kite.BaseError{
				Message: fmt.Sprintf("neither %s nor %s does exist.", kodingScreenPath, defaultScreenPath),
				CodeErr: ErrInvalidSession,
			}
		}

		screenPath = defaultScreenPath
	}

	return screenPath, nil
}

// newScreen returns a new screen instance that is used to start screen. The
// screen command line is created differently based on the incoming mode.
func newScreen(vos *virt.VOS, mode, session string) (*screen, error) {
	screenPath, err := getScreenPath(vos)
	if err != nil {
		return nil, err
	}

	cmdArgs := []string{screenPath, "-c", kodingScreenrc, "-S"}

	// check also if our custom screenrc exists before we continue
	_, err = vos.Stat(kodingScreenrc)
	if os.IsNotExist(err) {
		log.Critical("Screenrc %s does not exist. Starting screen without screenrc.", kodingScreenrc)
		cmdArgs = []string{screenPath, "-S"}
	}

	switch mode {
	case "shared", "resume":
		if session == "" {
			return nil, &kite.ArgumentError{Expected: "{ session: [string] }"}
		}

		if !sessionExists(vos, session) {
			return nil, &kite.BaseError{
				Message: fmt.Sprintf("The given session '%s' is not available.", session),
				CodeErr: ErrInvalidSession,
			}
		}

		cmdArgs = append(cmdArgs, sessionPrefix+"."+session)
		if mode == "shared" {
			cmdArgs = append(cmdArgs, "-x") // multiuser mode
		} else if mode == "resume" {
			cmdArgs = append(cmdArgs, "-raAd") // resume
		}
	case "noscreen":
		cmdArgs = nil
	case "create":
		session = utils.RandomString()
		cmdArgs = append(cmdArgs, sessionPrefix+"."+session)
	default:
		return nil, &kite.ArgumentError{Expected: "{ mode: [shared|noscreen|resume|create] }"}
	}

	return &screen{
		ScreenPath: screenPath,
		Session:    session,
		Command:    cmdArgs,
	}, nil

}

// screenSessions returns a list of sessions that belongs to the given vos
// context. The sessions are in the form of ["k7sdjv12344", "askIj12sas12", ...]
func screenSessions(vos *virt.VOS) []string {
	// Do not include dead sessions in our result
	vos.VM.AttachCommand(vos.User.Uid, "", "screen", "-wipe").Output()

	// We need to use ls here, because /var/run/screen mount is only
	// visible from inside of container. Errors are ignored.
	out, _ := vos.VM.AttachCommand(vos.User.Uid, "", "ls", "/var/run/screen/S-"+vos.User.Name).Output()
	shellOut := string(bytes.TrimSpace(out))
	if shellOut == "" {
		return []string{}
	}

	names := strings.Split(shellOut, "\n")
	sessions := make([]string, len(names))

	prefix := sessionPrefix + "."
	for i, name := range names {
		segments := strings.SplitN(name, ".", 2)
		sessions[i] = strings.TrimPrefix(segments[1], prefix)
	}

	return sessions
}

// screenExists checks whether the given session exists in the running list of
// screen sessions.
func sessionExists(vos *virt.VOS, session string) bool {
	for _, s := range screenSessions(vos) {
		if s == session {
			return true
		}
	}

	return false
}

// // killSession kills the given SessionID
// func killSession(vos *virt.VOS, sessionId string) error {
// 	vos.VM.AttachCommand(vos.User.Uid, "", "screen", "-X", "-S", sessionID, "kill").Output()
//
// }
<commit_msg>terminal: add webterm.killSession method<commit_after>package terminal

import (
	"bytes"
	"errors"
	"fmt"
	"koding/tools/dnode"
	"koding/tools/kite"
	"koding/tools/pty"
	"koding/tools/utils"
	"koding/virt"
	"os"
	"strconv"
	"strings"
	"syscall"
	"time"
	"unicode/utf8"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

const (
	sessionPrefix     = "koding"
	kodingScreenPath  = "/opt/koding/bin/screen"
	kodingScreenrc    = "/opt/koding/etc/screenrc"
	defaultScreenPath = "/usr/bin/screen"
)

type WebtermServer struct {
	Session          string `json:"session"`
	remote           WebtermRemote
	vm               *virt.VM
	user             *virt.User
	isForeignSession bool
	pty              *pty.PTY
	currentSecond    int64
	messageCounter   int
	byteCounter      int
	lineFeeedCounter int
	screenPath       string
}

type WebtermRemote struct {
	Output       dnode.Callback
	SessionEnded dnode.Callback
}

func webtermKillSession(args *dnode.Partial, channel *kite.Channel, vos *virt.VOS) (interface{}, error) {
	var params struct {
		Session string
	}

	if args == nil {
		return nil, &kite.ArgumentError{Expected: "empty argument passed"}
	}

	if args.Unmarshal(&params) != nil {
		return nil, &kite.ArgumentError{Expected: "{ session: [string]}"}
	}

	if err := killSession(vos, params.Session); err != nil {
		return nil, err
	}

	return true, nil
}

func webtermGetSessions(args *dnode.Partial, channel *kite.Channel, vos *virt.VOS) (interface{}, error) {
	sessions := screenSessions(vos)
	if len(sessions) == 0 {
		return nil, errors.New("no sessions available")
	}

	return sessions, nil
}

// this method is special cased in oskite.go to allow foreign access
func webtermConnect(args *dnode.Partial, channel *kite.Channel, vos *virt.VOS) (interface{}, error) {
	var params struct {
		Remote       WebtermRemote
		Session      string
		SizeX, SizeY int
		Mode         string
		JoinUser     string
	}

	if args == nil {
		return nil, &kite.ArgumentError{Expected: "empty argument passed"}
	}

	if args.Unmarshal(&params) != nil || params.SizeX <= 0 || params.SizeY <= 0 {
		return nil, &kite.ArgumentError{Expected: "{ remote: [object], session: [string], sizeX: [integer], sizeY: [integer], noScreen: [boolean] }"}
	}

	if params.JoinUser != "" {
		if len(params.Session) != utils.RandomStringLength {
			return nil, &kite.BaseError{
				Message: "Invalid session identifier",
				CodeErr: ErrInvalidSession,
			}
		}

		user := new(virt.User)
		if err := mongodbConn.Run("jUsers", func(c *mgo.Collection) error {
			return c.Find(bson.M{"username": params.JoinUser}).One(user)
		}); err != nil {
			return nil, err
		}

		// vos.VM is replaced already in registerMethod via
		// channel.CorrelationName which is the remote VM hostnameAlias
		vos.User = user
	}

	screen, err := newScreen(vos, params.Mode, params.Session)
	if err != nil {
		return nil, err
	}

	server := &WebtermServer{
		Session:          screen.Session,
		remote:           params.Remote,
		vm:               vos.VM,
		user:             vos.User,
		isForeignSession: vos.User.Name != channel.Username,
		pty:              pty.New(vos.VM.PtsDir()),
		screenPath:       screen.ScreenPath,
	}

	server.SetSize(float64(params.SizeX), float64(params.SizeY))
	server.pty.Slave.Chown(vos.User.Uid, -1)

	cmd := vos.VM.AttachCommand(vos.User.Uid, "/dev/pts/"+strconv.Itoa(server.pty.No), screen.Command...)
	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	go func() {
		defer log.RecoverAndLog()

		cmd.Wait()
		server.pty.Slave.Close()
		server.pty.Master.Close()
		server.remote.SessionEnded()
	}()

	go func() {
		defer log.RecoverAndLog()

		buf := make([]byte, (1<<12)-utf8.UTFMax, 1<<12)
		for {
			n, err := server.pty.Master.Read(buf)
			for n < cap(buf)-1 {
				r, _ := utf8.DecodeLastRune(buf[:n])
				if r != utf8.RuneError {
					break
				}
				server.pty.Master.Read(buf[n : n+1])
				n++
			}

			s := time.Now().Unix()
			if server.currentSecond != s {
				server.currentSecond = s
				server.messageCounter = 0
				server.byteCounter = 0
				server.lineFeeedCounter = 0
			}
			server.messageCounter += 1
			server.byteCounter += n
			server.lineFeeedCounter += bytes.Count(buf[:n], []byte{'\n'})
			if server.messageCounter > 100 || server.byteCounter > 1<<18 || server.lineFeeedCounter > 300 {
				time.Sleep(time.Second / 100)
			}

			server.remote.Output(string(utils.FilterInvalidUTF8(buf[:n])))
			if err != nil {
				break
			}
		}
	}()

	channel.OnDisconnect(func() { server.Close() })

	return server, nil
}

func (server *WebtermServer) Input(data string) {
	server.pty.Master.Write([]byte(data))
}

func (server *WebtermServer) ControlSequence(data string) {
	server.pty.MasterEncoded.Write([]byte(data))
}

func (server *WebtermServer) SetSize(x, y float64) {
	server.pty.SetSize(uint16(x), uint16(y))
}

func (server *WebtermServer) Close() error {
	server.pty.Signal(syscall.SIGHUP)
	return nil
}

func (server *WebtermServer) Terminate() error {
	server.Close()
	if !server.isForeignSession {
		server.vm.AttachCommand(server.user.Uid, "", server.screenPath, "-S", sessionPrefix+"."+server.Session, "-X", "quit").Run()
	}
	return nil
}

type screen struct {
	// Binary used for starting the screen
	ScreenPath string

	// Used for remote or multiuser mode, defines the custom session name
	Session string

	// the final command to be executed
	Command []string
}

func getScreenPath(vos *virt.VOS) (string, error) {
	screenPath := kodingScreenPath

	// it can happen that the user deleted our screen binary
	// accidently, if this happens fallback to default screen binary
	_, err := vos.Stat(kodingScreenPath)
	if os.IsNotExist(err) {
		// check if the default screen binary exists too
		_, err := vos.Stat(defaultScreenPath)
		if os.IsNotExist(err) {
			return "", &kite.BaseError{
				Message: fmt.Sprintf("neither %s nor %s does exist.", kodingScreenPath, defaultScreenPath),
				CodeErr: ErrInvalidSession,
			}
		}

		screenPath = defaultScreenPath
	}

	return screenPath, nil
}

// newScreen returns a new screen instance that is used to start screen. The
// screen command line is created differently based on the incoming mode.
func newScreen(vos *virt.VOS, mode, session string) (*screen, error) {
	screenPath, err := getScreenPath(vos)
	if err != nil {
		return nil, err
	}

	cmdArgs := []string{screenPath, "-c", kodingScreenrc, "-S"}

	// check also if our custom screenrc exists before we continue
	_, err = vos.Stat(kodingScreenrc)
	if os.IsNotExist(err) {
		log.Critical("Screenrc %s does not exist. Starting screen without screenrc.", kodingScreenrc)
		cmdArgs = []string{screenPath, "-S"}
	}

	switch mode {
	case "shared", "resume":
		if session == "" {
			return nil, &kite.ArgumentError{Expected: "{ session: [string] }"}
		}

		if !sessionExists(vos, session) {
			return nil, &kite.BaseError{
				Message: fmt.Sprintf("The given session '%s' is not available.", session),
				CodeErr: ErrInvalidSession,
			}
		}

		cmdArgs = append(cmdArgs, sessionPrefix+"."+session)
		if mode == "shared" {
			cmdArgs = append(cmdArgs, "-x") // multiuser mode
		} else if mode == "resume" {
			cmdArgs = append(cmdArgs, "-raAd") // resume
		}
	case "noscreen":
		cmdArgs = nil
	case "create":
		session = utils.RandomString()
		cmdArgs = append(cmdArgs, sessionPrefix+"."+session)
	default:
		return nil, &kite.ArgumentError{Expected: "{ mode: [shared|noscreen|resume|create] }"}
	}

	return &screen{
		ScreenPath: screenPath,
		Session:    session,
		Command:    cmdArgs,
	}, nil

}

// screenSessions returns a list of sessions that belongs to the given vos
// context. The sessions are in the form of ["k7sdjv12344", "askIj12sas12", ...]
func screenSessions(vos *virt.VOS) []string {
	// Do not include dead sessions in our result
	vos.VM.AttachCommand(vos.User.Uid, "", "screen", "-wipe").Output()

	// We need to use ls here, because /var/run/screen mount is only
	// visible from inside of container. Errors are ignored.
	out, _ := vos.VM.AttachCommand(vos.User.Uid, "", "ls", "/var/run/screen/S-"+vos.User.Name).Output()
	shellOut := string(bytes.TrimSpace(out))
	if shellOut == "" {
		return []string{}
	}

	names := strings.Split(shellOut, "\n")
	sessions := make([]string, len(names))

	prefix := sessionPrefix + "."
	for i, name := range names {
		segments := strings.SplitN(name, ".", 2)
		sessions[i] = strings.TrimPrefix(segments[1], prefix)
	}

	return sessions
}

// screenExists checks whether the given session exists in the running list of
// screen sessions.
func sessionExists(vos *virt.VOS, session string) bool {
	for _, s := range screenSessions(vos) {
		if s == session {
			return true
		}
	}

	return false
}

// killSession kills the given SessionID
func killSession(vos *virt.VOS, sessionID string) error {
	screenPath, err := getScreenPath(vos)
	if err != nil {
		return err
	}

	out, err := vos.VM.AttachCommand(vos.User.Uid, "", screenPath, "-X", "-S", sessionPrefix+"."+sessionID, "kill").Output()
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package models

import (
	"socialapi/workers/common/runner"
	"testing"

	"github.com/koding/bongo"
	. "github.com/smartystreets/goconvey/convey"
)

func TestAccountNewAccount(t *testing.T) {
	Convey("while testing new account", t, func() {
		Convey("Function call should return account", func() {
			So(NewAccount(), ShouldNotBeNil)
		})
	})
}

func TestAccountGetId(t *testing.T) {
	Convey("while testing get id", t, func() {
		Convey("Initialized struct ", func() {
			Convey("should return given id", func() {
				a := Account{Id: 42}
				So(a.GetId(), ShouldEqual, 42)
			})

			Convey("Uninitialized struct ", func() {
				Convey("should return 0", func() {
					So(NewAccount().GetId(), ShouldEqual, 0)
				})
				So(NewAccount(), ShouldNotBeNil)
			})
		})
	})
}

func TestAccountFetchOrCreate(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while fetching or creating account", t, func() {
		Convey("it should have old id", func() {
			a := NewAccount()
			err := a.FetchOrCreate()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrOldIdIsNotSet)
		})

		Convey("it should have error if nick contains guest-", func() {
			a := NewAccount()
			a.OldId = "oldestId"
			a.Nick = "guest-test"
			err := a.FetchOrCreate()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrGuestsAreNotAllowed)
		})

		Convey("it should not have error if required fields are exist", func() {
			a := NewAccount()
			a.OldId = "oldIdOfAccount"
			a.Nick = "WhatANick"
			err := a.FetchOrCreate()
			So(err, ShouldBeNil)
		})

	})
}

func TestAccountFetchAccountById(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while fetching an account by id", t, func() {
		Convey("it should not have error while fething", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			// fetch the account by id of account
			fa, err := FetchAccountById(acc.Id)
			// error should be nil
			// means that fetching is done successfully
			So(err, ShouldBeNil)
			// account in the db should equal to fetched account
			So(fa.Id, ShouldEqual, acc.Id)
			So(fa.OldId, ShouldEqual, acc.OldId)
			So(fa.Nick, ShouldEqual, acc.Nick)
		})

		Convey("it should have error if record is not found", func() {
			// init account
			a := NewAccount()
			a.Id = 12345
			a.OldId = "oldIdOfAccount"
			a.Nick = "WhatANick"

			// this account id is not exist
			_, err := FetchAccountById(a.Id)
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)

		})

	})
}

func TestAccountFetchOldIdsByAccountIds(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while fetching Old Ids by account ids", t, func() {
		Convey("it should have account id (ids of length) more than zero", func() {
			acc := []int64{}
			foi, err := FetchOldIdsByAccountIds(acc)
			So(err, ShouldBeNil)
			So(foi, ShouldBeEmpty)
		})

		Convey("it should not have error if account is exist in db", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)
			// we created slice to send to FetchOldIdsByAccountIds as argument
			idd := []int64{acc.Id}

			// FetchOldIdsByAccountIds returns as slice
			foi, err := FetchOldIdsByAccountIds(idd)
			So(err, ShouldBeNil)
			// used shouldcontain because foi is a slice
			So(foi, ShouldContain, acc.OldId)

		})

		Convey("it should append successfully", func() {
			acc1 := NewAccount()
			acc1.Id = 1
			acc1.OldId = "11"
			acc1.Nick = "acc1"
			So(acc1.Create(), ShouldBeNil)

			acc2 := NewAccount()
			acc2.Id = 2
			acc2.OldId = "22"
			acc2.Nick = "acc2"
			So(acc2.Create(), ShouldBeNil)

			idd := []int64{acc1.Id, acc2.Id}
			old := []string{acc1.OldId, acc2.OldId}

			foi, err := FetchOldIdsByAccountIds(idd)
			So(err, ShouldBeNil)
			So(foi[1], ShouldEqual, old[1])

		})

	})
}

func TestAccountCreateFollowingFeedChannel(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while creating feed channel", t, func() {
		Convey("it should have account id", func() {
			// create account
			acc := NewAccount()

			_, err := acc.CreateFollowingFeedChannel()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrAccountIdIsNotSet)
		})

		Convey("it should have creator id as account id", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			cff, err := acc.CreateFollowingFeedChannel()
			So(err, ShouldBeNil)
			So(cff.CreatorId, ShouldEqual, acc.Id)
		})

		Convey("it should have channel name as required", func() {
			// create account
			acc := NewAccount()
			acc.Id = 1
			acc.OldId = "11"
			acc.Nick = "acc1"
			So(acc.Create(), ShouldBeNil)

			cff, err := acc.CreateFollowingFeedChannel()
			So(err, ShouldBeNil)
			So(cff.Name, ShouldEqual, "1-FollowingFeedChannel")
		})

		Convey("it should have group name as Channel_KODING_NAME", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			cff, err := acc.CreateFollowingFeedChannel()
			So(err, ShouldBeNil)
			So(cff.GroupName, ShouldEqual, Channel_KODING_NAME)
		})

		Convey("it should have purpose as Following Feed for me", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			cff, err := acc.CreateFollowingFeedChannel()
			So(err, ShouldBeNil)
			So(cff.Purpose, ShouldEqual, "Following Feed for Me")
		})

		Convey("it should have type constant as Channel_TYPE_FOLLOWERS", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			cff, err := acc.CreateFollowingFeedChannel()
			So(err, ShouldBeNil)
			So(cff.TypeConstant, ShouldEqual, Channel_TYPE_FOLLOWERS)
		})

	})
}

func TestAccountUnMarkAsTroll(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while unmarking as troll", t, func() {
		Convey("it should have account id", func() {
			acc := NewAccount()

			err := acc.UnMarkAsTroll()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrAccountIdIsNotSet)
		})

		Convey("it should have account in db", func() {
			acc := NewAccount()
			acc.Id = 112233

			err := acc.UnMarkAsTroll()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)
		})

		Convey("it should have error if not troll", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			err := acc.UnMarkAsTroll()
			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldContainSubstring, "account is not a troll")
		})

		Convey("it should not have error if troll is mark as not a troll", func() {
			// create account
			acc := createAccountWithTest()
			acc.IsTroll = true
			So(acc.Create(), ShouldBeNil)

			err := acc.UnMarkAsTroll()
			So(err, ShouldBeNil)
		})
	})
}

func TestAccountMarkAsTroll(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while marking account as troll", t, func() {
		Convey("it should have account id", func() {
			acc := NewAccount()

			err := acc.MarkAsTroll()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrAccountIdIsNotSet)
		})

		Convey("it should have account in db", func() {
			acc := NewAccount()
			acc.Id = 1122312

			err := acc.MarkAsTroll()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)
		})

		Convey("it should have error if account is already troll", func() {
			// create account
			acc := createAccountWithTest()
			acc.IsTroll = true
			So(acc.Create(), ShouldBeNil)

			err := acc.MarkAsTroll()
			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldContainSubstring, "account is already a troll")
		})

		Convey("it should not have error if non-troll account is marked as troll", func() {
			// create account
			acc := createAccountWithTest()
			acc.IsTroll = false
			So(acc.Create(), ShouldBeNil)

			err := acc.MarkAsTroll()
			So(err, ShouldBeNil)
			So(acc.IsTroll, ShouldEqual, true)
		})

	})
}

func TestAccountFetchChannel(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while fetching channel", t, func() {
		Convey("it should have account id", func() {
			acc := NewAccount()

			_, err := acc.FetchChannel("ChannelType")
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrAccountIdIsNotSet)
		})

		Convey("it should have error if channel is not exist", func() {
			acc := NewAccount()
			acc.Id = 35135

			_, err := acc.FetchChannel(Channel_TYPE_GROUP)
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)
		})

		Convey("it should not have error if channel type is exist", func() {

			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			// create channel
			c := createNewChannelWithTest()
			c.CreatorId = acc.Id
			c.TypeConstant = Channel_TYPE_TOPIC
			So(c.Create(), ShouldBeNil)

			cha, err := acc.FetchChannel(Channel_TYPE_TOPIC)
			So(err, ShouldBeNil)
			So(cha.TypeConstant, ShouldEqual, c.TypeConstant)
		})

	})
}
<commit_msg>Social: add tests for ByNick function<commit_after>package models

import (
	"socialapi/workers/common/runner"
	"testing"

	"github.com/koding/bongo"
	. "github.com/smartystreets/goconvey/convey"
)

func TestAccountNewAccount(t *testing.T) {
	Convey("while testing new account", t, func() {
		Convey("Function call should return account", func() {
			So(NewAccount(), ShouldNotBeNil)
		})
	})
}

func TestAccountGetId(t *testing.T) {
	Convey("while testing get id", t, func() {
		Convey("Initialized struct ", func() {
			Convey("should return given id", func() {
				a := Account{Id: 42}
				So(a.GetId(), ShouldEqual, 42)
			})

			Convey("Uninitialized struct ", func() {
				Convey("should return 0", func() {
					So(NewAccount().GetId(), ShouldEqual, 0)
				})
				So(NewAccount(), ShouldNotBeNil)
			})
		})
	})
}

func TestAccountFetchOrCreate(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while fetching or creating account", t, func() {
		Convey("it should have old id", func() {
			a := NewAccount()
			err := a.FetchOrCreate()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrOldIdIsNotSet)
		})

		Convey("it should have error if nick contains guest-", func() {
			a := NewAccount()
			a.OldId = "oldestId"
			a.Nick = "guest-test"
			err := a.FetchOrCreate()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrGuestsAreNotAllowed)
		})

		Convey("it should not have error if required fields are exist", func() {
			a := NewAccount()
			a.OldId = "oldIdOfAccount"
			a.Nick = "WhatANick"
			err := a.FetchOrCreate()
			So(err, ShouldBeNil)
		})

	})
}

func TestAccountFetchAccountById(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while fetching an account by id", t, func() {
		Convey("it should not have error while fething", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			// fetch the account by id of account
			fa, err := FetchAccountById(acc.Id)
			// error should be nil
			// means that fetching is done successfully
			So(err, ShouldBeNil)
			// account in the db should equal to fetched account
			So(fa.Id, ShouldEqual, acc.Id)
			So(fa.OldId, ShouldEqual, acc.OldId)
			So(fa.Nick, ShouldEqual, acc.Nick)
		})

		Convey("it should have error if record is not found", func() {
			// init account
			a := NewAccount()
			a.Id = 12345
			a.OldId = "oldIdOfAccount"
			a.Nick = "WhatANick"

			// this account id is not exist
			_, err := FetchAccountById(a.Id)
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)

		})

	})
}

func TestAccountByNick(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while fetching an account by nick", t, func() {
		Convey("it should not have error while fething", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			fa := &NewAccount()

			// fetch the account by nick of account
			fa.ByNick(acc.Nick)
			// error should be nil
			// means that fetching is done successfully
			So(err, ShouldBeNil)
			// account in the db should be equal to fetched account
			So(fa.Id, ShouldEqual, acc.Id)
			So(fa.OldId, ShouldEqual, acc.OldId)
			So(fa.Nick, ShouldEqual, acc.Nick)
		})

		Convey("it should have error if record is not found", func() {

			// init account
			fa := &NewAccount()
			err := fa.ByNick("")
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)

		})

	})
}

func TestAccountFetchOldIdsByAccountIds(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while fetching Old Ids by account ids", t, func() {
		Convey("it should have account id (ids of length) more than zero", func() {
			acc := []int64{}
			foi, err := FetchOldIdsByAccountIds(acc)
			So(err, ShouldBeNil)
			So(foi, ShouldBeEmpty)
		})

		Convey("it should not have error if account is exist in db", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)
			// we created slice to send to FetchOldIdsByAccountIds as argument
			idd := []int64{acc.Id}

			// FetchOldIdsByAccountIds returns as slice
			foi, err := FetchOldIdsByAccountIds(idd)
			So(err, ShouldBeNil)
			// used shouldcontain because foi is a slice
			So(foi, ShouldContain, acc.OldId)

		})

		Convey("it should append successfully", func() {
			acc1 := NewAccount()
			acc1.Id = 1
			acc1.OldId = "11"
			acc1.Nick = "acc1"
			So(acc1.Create(), ShouldBeNil)

			acc2 := NewAccount()
			acc2.Id = 2
			acc2.OldId = "22"
			acc2.Nick = "acc2"
			So(acc2.Create(), ShouldBeNil)

			idd := []int64{acc1.Id, acc2.Id}
			old := []string{acc1.OldId, acc2.OldId}

			foi, err := FetchOldIdsByAccountIds(idd)
			So(err, ShouldBeNil)
			So(foi[1], ShouldEqual, old[1])

		})

	})
}

func TestAccountCreateFollowingFeedChannel(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while creating feed channel", t, func() {
		Convey("it should have account id", func() {
			// create account
			acc := NewAccount()

			_, err := acc.CreateFollowingFeedChannel()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrAccountIdIsNotSet)
		})

		Convey("it should have creator id as account id", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			cff, err := acc.CreateFollowingFeedChannel()
			So(err, ShouldBeNil)
			So(cff.CreatorId, ShouldEqual, acc.Id)
		})

		Convey("it should have channel name as required", func() {
			// create account
			acc := NewAccount()
			acc.Id = 1
			acc.OldId = "11"
			acc.Nick = "acc1"
			So(acc.Create(), ShouldBeNil)

			cff, err := acc.CreateFollowingFeedChannel()
			So(err, ShouldBeNil)
			So(cff.Name, ShouldEqual, "1-FollowingFeedChannel")
		})

		Convey("it should have group name as Channel_KODING_NAME", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			cff, err := acc.CreateFollowingFeedChannel()
			So(err, ShouldBeNil)
			So(cff.GroupName, ShouldEqual, Channel_KODING_NAME)
		})

		Convey("it should have purpose as Following Feed for me", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			cff, err := acc.CreateFollowingFeedChannel()
			So(err, ShouldBeNil)
			So(cff.Purpose, ShouldEqual, "Following Feed for Me")
		})

		Convey("it should have type constant as Channel_TYPE_FOLLOWERS", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			cff, err := acc.CreateFollowingFeedChannel()
			So(err, ShouldBeNil)
			So(cff.TypeConstant, ShouldEqual, Channel_TYPE_FOLLOWERS)
		})

	})
}

func TestAccountUnMarkAsTroll(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while unmarking as troll", t, func() {
		Convey("it should have account id", func() {
			acc := NewAccount()

			err := acc.UnMarkAsTroll()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrAccountIdIsNotSet)
		})

		Convey("it should have account in db", func() {
			acc := NewAccount()
			acc.Id = 112233

			err := acc.UnMarkAsTroll()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)
		})

		Convey("it should have error if not troll", func() {
			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			err := acc.UnMarkAsTroll()
			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldContainSubstring, "account is not a troll")
		})

		Convey("it should not have error if troll is mark as not a troll", func() {
			// create account
			acc := createAccountWithTest()
			acc.IsTroll = true
			So(acc.Create(), ShouldBeNil)

			err := acc.UnMarkAsTroll()
			So(err, ShouldBeNil)
		})
	})
}

func TestAccountMarkAsTroll(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while marking account as troll", t, func() {
		Convey("it should have account id", func() {
			acc := NewAccount()

			err := acc.MarkAsTroll()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrAccountIdIsNotSet)
		})

		Convey("it should have account in db", func() {
			acc := NewAccount()
			acc.Id = 1122312

			err := acc.MarkAsTroll()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)
		})

		Convey("it should have error if account is already troll", func() {
			// create account
			acc := createAccountWithTest()
			acc.IsTroll = true
			So(acc.Create(), ShouldBeNil)

			err := acc.MarkAsTroll()
			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldContainSubstring, "account is already a troll")
		})

		Convey("it should not have error if non-troll account is marked as troll", func() {
			// create account
			acc := createAccountWithTest()
			acc.IsTroll = false
			So(acc.Create(), ShouldBeNil)

			err := acc.MarkAsTroll()
			So(err, ShouldBeNil)
			So(acc.IsTroll, ShouldEqual, true)
		})

	})
}

func TestAccountFetchChannel(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while fetching channel", t, func() {
		Convey("it should have account id", func() {
			acc := NewAccount()

			_, err := acc.FetchChannel("ChannelType")
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, ErrAccountIdIsNotSet)
		})

		Convey("it should have error if channel is not exist", func() {
			acc := NewAccount()
			acc.Id = 35135

			_, err := acc.FetchChannel(Channel_TYPE_GROUP)
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)
		})

		Convey("it should not have error if channel type is exist", func() {

			// create account
			acc := createAccountWithTest()
			So(acc.Create(), ShouldBeNil)

			// create channel
			c := createNewChannelWithTest()
			c.CreatorId = acc.Id
			c.TypeConstant = Channel_TYPE_TOPIC
			So(c.Create(), ShouldBeNil)

			cha, err := acc.FetchChannel(Channel_TYPE_TOPIC)
			So(err, ShouldBeNil)
			So(cha.TypeConstant, ShouldEqual, c.TypeConstant)
		})

	})
}
<|endoftext|>


<commit_before>// Copyright 2014, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package actionnode

// This file contains utility functions to be used with actionnode /
// topology server.

import (
	"fmt"
	"sync"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/concurrency"
	"github.com/youtube/vitess/go/vt/hook"
	"github.com/youtube/vitess/go/vt/key"
	"github.com/youtube/vitess/go/vt/topo"
)

var (
	// DefaultLockTimeout is a good value to use as a default for
	// locking a shard / keyspace.
	DefaultLockTimeout = 30 * time.Second
)

// LockKeyspace will lock the keyspace in the topology server.
// UnlockKeyspace should be called if this returns no error.
func (n *ActionNode) LockKeyspace(ts topo.Server, keyspace string, lockTimeout time.Duration, interrupted chan struct{}) (lockPath string, err error) {
	log.Infof("Locking keyspace %v for action %v", keyspace, n.Action)
	return ts.LockKeyspaceForAction(keyspace, n.ToJson(), lockTimeout, interrupted)
}

// UnlockKeyspace unlocks a previously locked keyspace.
func (n *ActionNode) UnlockKeyspace(ts topo.Server, keyspace string, lockPath string, actionError error) error {
	// first update the actionNode
	if actionError != nil {
		log.Infof("Unlocking keyspace %v for action %v with error %v", keyspace, n.Action, actionError)
		n.Error = actionError.Error()
		n.State = ACTION_STATE_FAILED
	} else {
		log.Infof("Unlocking keyspace %v for successful action %v", keyspace, n.Action)
		n.Error = ""
		n.State = ACTION_STATE_DONE
	}
	err := ts.UnlockKeyspaceForAction(keyspace, lockPath, n.ToJson())
	if actionError != nil {
		if err != nil {
			// this will be masked
			log.Warningf("UnlockKeyspaceForAction failed: %v", err)
		}
		return actionError
	}
	return err
}

// LockShard will lock the shard in the topology server.
// UnlockShard should be called if this returns no error.
func (n *ActionNode) LockShard(ts topo.Server, keyspace, shard string, lockTimeout time.Duration, interrupted chan struct{}) (lockPath string, err error) {
	log.Infof("Locking shard %v/%v for action %v", keyspace, shard, n.Action)
	return ts.LockShardForAction(keyspace, shard, n.ToJson(), lockTimeout, interrupted)
}

// UnlockShard unlocks a previously locked shard.
func (n *ActionNode) UnlockShard(ts topo.Server, keyspace, shard string, lockPath string, actionError error) error {
	// first update the actionNode
	if actionError != nil {
		log.Infof("Unlocking shard %v/%v for action %v with error %v", keyspace, shard, n.Action, actionError)
		n.Error = actionError.Error()
		n.State = ACTION_STATE_FAILED
	} else {
		log.Infof("Unlocking shard %v/%v for successful action %v", keyspace, shard, n.Action)
		n.Error = ""
		n.State = ACTION_STATE_DONE
	}
	err := ts.UnlockShardForAction(keyspace, shard, lockPath, n.ToJson())
	if actionError != nil {
		if err != nil {
			// this will be masked
			log.Warningf("UnlockShardForAction failed: %v", err)
		}
		return actionError
	}
	return err
}

// ConfigureTabletHook configures the right parameters for a hook
// running locally on a tablet.
func ConfigureTabletHook(hk *hook.Hook, tabletAlias topo.TabletAlias) {
	if hk.ExtraEnv == nil {
		hk.ExtraEnv = make(map[string]string, 1)
	}
	hk.ExtraEnv["TABLET_ALIAS"] = tabletAlias.String()
}

// Scrap will update the tablet type to 'Scrap', and remove it from
// the serving graph.
//
// 'force' means we are not on the tablet being scrapped, so it is
//   probably dead.
// So if 'force' is true, we will also remove pending remote actions.
// And if 'force' is false, we also run an optional hook.
func Scrap(ts topo.Server, tabletAlias topo.TabletAlias, force bool) error {
	tablet, err := ts.GetTablet(tabletAlias)
	if err != nil {
		return err
	}

	// If you are already scrap, skip updating replication data. It won't
	// be there anyway.
	wasAssigned := tablet.IsAssigned()
	tablet.Type = topo.TYPE_SCRAP
	tablet.Parent = topo.TabletAlias{}
	// Update the tablet first, since that is canonical.
	err = topo.UpdateTablet(ts, tablet)
	if err != nil {
		return err
	}

	// Remove any pending actions. Presumably forcing a scrap
	// means you don't want the agent doing anything and the
	// machine requires manual attention.
	if force {
		err := ts.PurgeTabletActions(tabletAlias, ActionNodeCanBePurged)
		if err != nil {
			log.Warningf("purge actions failed: %v", err)
		}
	}

	if wasAssigned {
		err = topo.DeleteTabletReplicationData(ts, tablet.Tablet)
		if err != nil {
			if err == topo.ErrNoNode {
				log.V(6).Infof("no ShardReplication object for cell %v", tablet.Alias.Cell)
				err = nil
			}
			if err != nil {
				log.Warningf("remove replication data for %v failed: %v", tablet.Alias, err)
			}
		}
	}

	// run a hook for final cleanup, only in non-force mode.
	// (force mode executes on the vtctl side, not on the vttablet side)
	if !force {
		hk := hook.NewSimpleHook("postflight_scrap")
		ConfigureTabletHook(hk, tablet.Alias)
		if hookErr := hk.ExecuteOptional(); hookErr != nil {
			// we don't want to return an error, the server
			// is already in bad shape probably.
			log.Warningf("Scrap: postflight_scrap failed: %v", hookErr)
		}
	}

	return nil
}

// ChangeType changes the type of the tablet and possibly also updates
// the health informaton for it. Make this external, since these
// transitions need to be forced from time to time.
//
// - if health is nil, we don't touch the Tablet's Health record.
// - if health is an empty map, we clear the Tablet's Health record.
// - if health has values, we overwrite the Tablet's Health record.
func ChangeType(ts topo.Server, tabletAlias topo.TabletAlias, newType topo.TabletType, health map[string]string, runHooks bool) error {
	tablet, err := ts.GetTablet(tabletAlias)
	if err != nil {
		return err
	}

	if !topo.IsTrivialTypeChange(tablet.Type, newType) || !topo.IsValidTypeChange(tablet.Type, newType) {
		return fmt.Errorf("cannot change tablet type %v -> %v %v", tablet.Type, newType, tabletAlias)
	}

	if runHooks {
		// Only run the preflight_serving_type hook when
		// transitioning from non-serving to serving.
		if !topo.IsInServingGraph(tablet.Type) && topo.IsInServingGraph(newType) {
			if err := hook.NewSimpleHook("preflight_serving_type").ExecuteOptional(); err != nil {
				return err
			}
		}
	}

	tablet.Type = newType
	if newType == topo.TYPE_IDLE {
		if tablet.Parent.IsZero() {
			si, err := ts.GetShard(tablet.Keyspace, tablet.Shard)
			if err != nil {
				return err
			}
			rec := concurrency.AllErrorRecorder{}
			wg := sync.WaitGroup{}
			for _, cell := range si.Cells {
				wg.Add(1)
				go func(cell string) {
					defer wg.Done()
					sri, err := ts.GetShardReplication(cell, tablet.Keyspace, tablet.Shard)
					if err != nil {
						log.Warningf("Cannot check cell %v for extra replication paths, assuming it's good", cell)
						return
					}
					for _, rl := range sri.ReplicationLinks {
						if rl.Parent == tabletAlias {
							rec.RecordError(fmt.Errorf("Still have a ReplicationLink in cell %v", cell))
						}
					}
				}(cell)
			}
			wg.Wait()
			if rec.HasErrors() {
				return rec.Error()
			}
		}
		tablet.Parent = topo.TabletAlias{}
		tablet.Keyspace = ""
		tablet.Shard = ""
		tablet.KeyRange = key.KeyRange{}
		tablet.Health = health
	}
	if health != nil {
		if len(health) == 0 {
			tablet.Health = nil
		} else {
			tablet.Health = health
		}
	}
	return topo.UpdateTablet(ts, tablet)
}
<commit_msg>Fixing another comment.<commit_after>// Copyright 2014, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package actionnode

// This file contains utility functions to be used with actionnode /
// topology server.

import (
	"fmt"
	"sync"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/concurrency"
	"github.com/youtube/vitess/go/vt/hook"
	"github.com/youtube/vitess/go/vt/key"
	"github.com/youtube/vitess/go/vt/topo"
)

var (
	// DefaultLockTimeout is a good value to use as a default for
	// locking a shard / keyspace.
	DefaultLockTimeout = 30 * time.Second
)

// LockKeyspace will lock the keyspace in the topology server.
// UnlockKeyspace should be called if this returns no error.
func (n *ActionNode) LockKeyspace(ts topo.Server, keyspace string, lockTimeout time.Duration, interrupted chan struct{}) (lockPath string, err error) {
	log.Infof("Locking keyspace %v for action %v", keyspace, n.Action)
	return ts.LockKeyspaceForAction(keyspace, n.ToJson(), lockTimeout, interrupted)
}

// UnlockKeyspace unlocks a previously locked keyspace.
func (n *ActionNode) UnlockKeyspace(ts topo.Server, keyspace string, lockPath string, actionError error) error {
	// first update the actionNode
	if actionError != nil {
		log.Infof("Unlocking keyspace %v for action %v with error %v", keyspace, n.Action, actionError)
		n.Error = actionError.Error()
		n.State = ACTION_STATE_FAILED
	} else {
		log.Infof("Unlocking keyspace %v for successful action %v", keyspace, n.Action)
		n.Error = ""
		n.State = ACTION_STATE_DONE
	}
	err := ts.UnlockKeyspaceForAction(keyspace, lockPath, n.ToJson())
	if actionError != nil {
		if err != nil {
			// this will be masked
			log.Warningf("UnlockKeyspaceForAction failed: %v", err)
		}
		return actionError
	}
	return err
}

// LockShard will lock the shard in the topology server.
// UnlockShard should be called if this returns no error.
func (n *ActionNode) LockShard(ts topo.Server, keyspace, shard string, lockTimeout time.Duration, interrupted chan struct{}) (lockPath string, err error) {
	log.Infof("Locking shard %v/%v for action %v", keyspace, shard, n.Action)
	return ts.LockShardForAction(keyspace, shard, n.ToJson(), lockTimeout, interrupted)
}

// UnlockShard unlocks a previously locked shard.
func (n *ActionNode) UnlockShard(ts topo.Server, keyspace, shard string, lockPath string, actionError error) error {
	// first update the actionNode
	if actionError != nil {
		log.Infof("Unlocking shard %v/%v for action %v with error %v", keyspace, shard, n.Action, actionError)
		n.Error = actionError.Error()
		n.State = ACTION_STATE_FAILED
	} else {
		log.Infof("Unlocking shard %v/%v for successful action %v", keyspace, shard, n.Action)
		n.Error = ""
		n.State = ACTION_STATE_DONE
	}
	err := ts.UnlockShardForAction(keyspace, shard, lockPath, n.ToJson())
	if actionError != nil {
		if err != nil {
			// this will be masked
			log.Warningf("UnlockShardForAction failed: %v", err)
		}
		return actionError
	}
	return err
}

// ConfigureTabletHook configures the right parameters for a hook
// running locally on a tablet.
func ConfigureTabletHook(hk *hook.Hook, tabletAlias topo.TabletAlias) {
	if hk.ExtraEnv == nil {
		hk.ExtraEnv = make(map[string]string, 1)
	}
	hk.ExtraEnv["TABLET_ALIAS"] = tabletAlias.String()
}

// Scrap will update the tablet type to 'Scrap', and remove it from
// the serving graph.
//
// 'force' means we are not on the tablet being scrapped, so it is
// probably dead. So if 'force' is true, we will also remove pending
// remote actions.  And if 'force' is false, we also run an optional
// hook.
func Scrap(ts topo.Server, tabletAlias topo.TabletAlias, force bool) error {
	tablet, err := ts.GetTablet(tabletAlias)
	if err != nil {
		return err
	}

	// If you are already scrap, skip updating replication data. It won't
	// be there anyway.
	wasAssigned := tablet.IsAssigned()
	tablet.Type = topo.TYPE_SCRAP
	tablet.Parent = topo.TabletAlias{}
	// Update the tablet first, since that is canonical.
	err = topo.UpdateTablet(ts, tablet)
	if err != nil {
		return err
	}

	// Remove any pending actions. Presumably forcing a scrap
	// means you don't want the agent doing anything and the
	// machine requires manual attention.
	if force {
		err := ts.PurgeTabletActions(tabletAlias, ActionNodeCanBePurged)
		if err != nil {
			log.Warningf("purge actions failed: %v", err)
		}
	}

	if wasAssigned {
		err = topo.DeleteTabletReplicationData(ts, tablet.Tablet)
		if err != nil {
			if err == topo.ErrNoNode {
				log.V(6).Infof("no ShardReplication object for cell %v", tablet.Alias.Cell)
				err = nil
			}
			if err != nil {
				log.Warningf("remove replication data for %v failed: %v", tablet.Alias, err)
			}
		}
	}

	// run a hook for final cleanup, only in non-force mode.
	// (force mode executes on the vtctl side, not on the vttablet side)
	if !force {
		hk := hook.NewSimpleHook("postflight_scrap")
		ConfigureTabletHook(hk, tablet.Alias)
		if hookErr := hk.ExecuteOptional(); hookErr != nil {
			// we don't want to return an error, the server
			// is already in bad shape probably.
			log.Warningf("Scrap: postflight_scrap failed: %v", hookErr)
		}
	}

	return nil
}

// ChangeType changes the type of the tablet and possibly also updates
// the health informaton for it. Make this external, since these
// transitions need to be forced from time to time.
//
// - if health is nil, we don't touch the Tablet's Health record.
// - if health is an empty map, we clear the Tablet's Health record.
// - if health has values, we overwrite the Tablet's Health record.
func ChangeType(ts topo.Server, tabletAlias topo.TabletAlias, newType topo.TabletType, health map[string]string, runHooks bool) error {
	tablet, err := ts.GetTablet(tabletAlias)
	if err != nil {
		return err
	}

	if !topo.IsTrivialTypeChange(tablet.Type, newType) || !topo.IsValidTypeChange(tablet.Type, newType) {
		return fmt.Errorf("cannot change tablet type %v -> %v %v", tablet.Type, newType, tabletAlias)
	}

	if runHooks {
		// Only run the preflight_serving_type hook when
		// transitioning from non-serving to serving.
		if !topo.IsInServingGraph(tablet.Type) && topo.IsInServingGraph(newType) {
			if err := hook.NewSimpleHook("preflight_serving_type").ExecuteOptional(); err != nil {
				return err
			}
		}
	}

	tablet.Type = newType
	if newType == topo.TYPE_IDLE {
		if tablet.Parent.IsZero() {
			si, err := ts.GetShard(tablet.Keyspace, tablet.Shard)
			if err != nil {
				return err
			}
			rec := concurrency.AllErrorRecorder{}
			wg := sync.WaitGroup{}
			for _, cell := range si.Cells {
				wg.Add(1)
				go func(cell string) {
					defer wg.Done()
					sri, err := ts.GetShardReplication(cell, tablet.Keyspace, tablet.Shard)
					if err != nil {
						log.Warningf("Cannot check cell %v for extra replication paths, assuming it's good", cell)
						return
					}
					for _, rl := range sri.ReplicationLinks {
						if rl.Parent == tabletAlias {
							rec.RecordError(fmt.Errorf("Still have a ReplicationLink in cell %v", cell))
						}
					}
				}(cell)
			}
			wg.Wait()
			if rec.HasErrors() {
				return rec.Error()
			}
		}
		tablet.Parent = topo.TabletAlias{}
		tablet.Keyspace = ""
		tablet.Shard = ""
		tablet.KeyRange = key.KeyRange{}
		tablet.Health = health
	}
	if health != nil {
		if len(health) == 0 {
			tablet.Health = nil
		} else {
			tablet.Health = health
		}
	}
	return topo.UpdateTablet(ts, tablet)
}
<|endoftext|>


<commit_before>package main

var edition = "ce"
var version = "v1.3.0"
<commit_msg>release: v1.3.1<commit_after>package main

var edition = "ce"
var version = "v1.3.1"
<|endoftext|>


<commit_before>// +build doc

package cony_test

import (
	"log"
	"os"
	"time"

	"github.com/assembla/cony"
	"github.com/streadway/amqp"
)

func Example() {
	client := cony.NewClient(cony.URL(os.Getenv("AMQP_URL")), cony.Backoff(cony.DefaultBackoff))

	q := &cony.Queue{
		Name:       "", // autogenerated queue name
		AutoDelete: true,
	}

	exchange := cony.Exchange{
		Name:    "amq.topic",
		Durable: true,
	}

	b := cony.Binding{
		Queue:    q,
		Exchange: exchange,
		Key:      "something.#",
	}

	// wrap all declarations and save into slice
	declarations := []cony.Declaration{
		cony.DeclareQueue(q),
		cony.DeclareExchange(exchange),
		cony.DeclareBinding(b),
	}

	// declare consumer
	consumer := cony.NewConsumer(q,
		cony.Qos(10),
		cony.AutoTag(),
		cony.AutoAck(),
	)

	// declare publisher
	publisher := cony.NewPublisher(exchange.Name,
		"ololo.key",
		cony.PublishingTemplate(amqp.Publishing{
			ContentType: "application/json",
			AppId:       "app1",
		}), // template amqp.Publising
	)

	// let client know about declarations
	client.Declare(declarations)

	// let client know about consumers
	client.Consume(consumer)

	clientErrs := client.Errors()
	deliveries := consumer.Deliveries()
	consumerErrs := consumer.Errors()

	// connect, reconnect, or exit loop
	// run network operations such as:
	// queue, exchange, bidning, consumers declarations
	for client.Loop() {
		select {
		case msg := <-deliveries:
			log.Println(msg)
			msg.Ack(false)
			publisher.Write([]byte("ololo reply"))
		case err := <-consumerErrs:
			log.Println("CONSUMER ERROR: ", err)
		case err := <-clientErrs:
			log.Println("CLIENT ERROR: ", err)
			client.Close()
		}
	}

}

func ExampleURL() {
	cony.NewClient(cony.URL("amqp://guest:guest@localhost/"))
}

func ExampleClient_Loop() {
	client := cony.NewClient(cony.URL("amqp://guest:guest@localhost/"))

	for client.Loop() {
		select {
		case err := <-client.Errors():
			log.Println("CLIENT ERROR: ", err)
			client.Close()
		}
		time.Sleep(1 * time.Second) // naive backoff
	}
}
<commit_msg>show example code in godoc<commit_after>package cony_test

import (
	"log"
	"os"
	"time"

	"github.com/assembla/cony"
	"github.com/streadway/amqp"
)

func Example() {
	client := cony.NewClient(cony.URL(os.Getenv("AMQP_URL")), cony.Backoff(cony.DefaultBackoff))

	q := &cony.Queue{
		Name:       "", // autogenerated queue name
		AutoDelete: true,
	}

	exchange := cony.Exchange{
		Name:    "amq.topic",
		Durable: true,
	}

	b := cony.Binding{
		Queue:    q,
		Exchange: exchange,
		Key:      "something.#",
	}

	// wrap all declarations and save into slice
	declarations := []cony.Declaration{
		cony.DeclareQueue(q),
		cony.DeclareExchange(exchange),
		cony.DeclareBinding(b),
	}

	// declare consumer
	consumer := cony.NewConsumer(q,
		cony.Qos(10),
		cony.AutoTag(),
		cony.AutoAck(),
	)

	// declare publisher
	publisher := cony.NewPublisher(exchange.Name,
		"ololo.key",
		cony.PublishingTemplate(amqp.Publishing{
			ContentType: "application/json",
			AppId:       "app1",
		}), // template amqp.Publising
	)

	// let client know about declarations
	client.Declare(declarations)

	// let client know about consumers
	client.Consume(consumer)

	clientErrs := client.Errors()
	deliveries := consumer.Deliveries()
	consumerErrs := consumer.Errors()

	// connect, reconnect, or exit loop
	// run network operations such as:
	// queue, exchange, bidning, consumers declarations
	for client.Loop() {
		select {
		case msg := <-deliveries:
			log.Println(msg)
			msg.Ack(false)
			publisher.Write([]byte("ololo reply"))
		case err := <-consumerErrs:
			log.Println("CONSUMER ERROR: ", err)
		case err := <-clientErrs:
			log.Println("CLIENT ERROR: ", err)
			client.Close()
		}
	}

}

func ExampleURL() {
	cony.NewClient(cony.URL("amqp://guest:guest@localhost/"))
}

func ExampleClient_Loop() {
	client := cony.NewClient(cony.URL("amqp://guest:guest@localhost/"))

	for client.Loop() {
		select {
		case err := <-client.Errors():
			log.Println("CLIENT ERROR: ", err)
			client.Close()
		}
		time.Sleep(1 * time.Second) // naive backoff
	}
}
<|endoftext|>