<commit_before>/*-
 * Copyright (c) 2016, 1&1 Internet SE
 * Copyright (c) 2016, Jörg Pernfuß
 *
 * Use of this source code is governed by a 2-clause BSD license
 * that can be found in the LICENSE file.
 */

package adm

import (
	"fmt"
	"net/mail"
	"regexp"
	"strconv"
	"unicode/utf8"

	"github.com/1and1/soma/lib/proto"
	resty "gopkg.in/resty.v0"
)

// ValidateRuneCount tests if a string's number of unicode runes is
// below an upper limit
func ValidateRuneCount(s string, l int) error {
	if cnt := utf8.RuneCountInString(s); cnt > l {
		return fmt.Errorf("Validation error, string '%s' above character limit %d.",
			s, l)
	}
	return nil
}

// ValidateRuneCountRange tests if a string's number of unicode runes is
// between an upper and lower bound
func ValidateRuneCountRange(s string, lower, higher int) error {
	if utf8.RuneCountInString(s) < lower || utf8.RuneCountInString(s) > higher {
		return fmt.Errorf("Validation error, string '%s' outside permitted length."+
			"Required: %d < len(%s) < %d.", s, lower, s, higher)
	}
	return nil
}

// IsUUID validates if a string is one very narrow formatting of a UUID,
// namely the one used by the server. Other valid formats with braces etc
// are not accepted
func IsUUID(s string) bool {
	const reUUID string = `^[[:xdigit:]]{8}-[[:xdigit:]]{4}-[1-5][[:xdigit:]]{3}-[[:xdigit:]]{4}-[[:xdigit:]]{12}$`
	const reUNIL string = `^0{8}-0{4}-0{4}-0{4}-0{12}$`
	re := regexp.MustCompile(fmt.Sprintf("%s|%s", reUUID, reUNIL))

	return re.MatchString(s)
}

// isUint64 validates if a string is a uint64 number
func isUint64(s string) (bool, uint64) {
	if u, err := strconv.ParseUint(s, 10, 64); err != nil {
		return false, 0
	} else {
		return true, u
	}
}

// ValidateOncallNumber tests if a string is a 4 digit number
func ValidateOncallNumber(n string) error {
	num, err := strconv.Atoi(n)
	if err != nil {
		return fmt.Errorf("Syntax error, argument is not a number: %s", err.Error())
	}

	if num <= 0 || num > 9999 {
		return fmt.Errorf("Phone number must be 4-digit extension")
	}
	return nil
}

// ValidateEmployeeNumber tests if a string is a non-negative number
func ValidateEmployeeNumber(s string) error {
	employeeNumber, err := strconv.Atoi(s)
	if err != nil {
		return fmt.Errorf("Error: employeenr argument not a number")
	}
	if employeeNumber < 0 {
		return fmt.Errorf("Negative employee number is not allowed")
	}
	return nil
}

// ValidateMailAddress tests if a string is a mail address
func ValidateMailAddress(s string) error {
	if _, err := mail.ParseAddress(s); err != nil {
		return fmt.Errorf("Error: mailaddr does not parse as RFC 5322 address")
	}
	return nil
}

// ValidateBool tests that string s represents and bool and sets
// its value in b.
// Returns error if s could not be parsed as bool.
func ValidateBool(s string, b *bool) error {
	var err error
	*b, err = strconv.ParseBool(s)
	return err
}

// ValidateLBoundUint64 tests that string s is an unsigned number
// with a specific minimum value and sets it in i.
// Returns error if the number conversion failed or the minimum
// value was not reached.
func ValidateLBoundUint64(s string, i *uint64, min uint64) error {
	var err error
	if *i, err = strconv.ParseUint(s, 10, 64); err != nil {
		return err
	}
	if *i < min {
		return fmt.Errorf("Error, value %d is less than the minimun %d", i, min)
	}
	return nil
}

// ValidateUnit tests against the server if string s is a valid
// unit.
func ValidateUnit(s string) error {
	res, err := fetchObjList(`/units/`)
	if err != nil {
		return err
	}

	if res.Units != nil {
		for _, unit := range *res.Units {
			if unit.Unit == s {
				return nil
			}
		}
	}
	return fmt.Errorf("Value %s is not a valid unit", s)
}

// ValidateProvider tests aginst the server if string s is a valid
// provider.
func ValidateProvider(s string) error {
	res, err := fetchObjList(`/providers/`)
	if err != nil {
		return err
	}

	if res.Providers != nil {
		for _, prov := range *res.Providers {
			if prov.Name == s {
				return nil
			}
		}
	}
	return fmt.Errorf("Value %s is not a valid provider", s)
}

// ValidateSystemProperty tests against the server if string s is a
// valid system property.
func ValidateSystemProperty(s string) error {
	res, err := fetchObjList(`/property/system/`)
	if err != nil {
		return err
	}

	if res.Properties != nil {
		for _, prop := range *res.Properties {
			if prop.System.Name == s {
				return nil
			}
		}
	}
	return fmt.Errorf("Invalid system property requested: %s", s)
}

// ValidateEnvironment tests against the server if string s is a
// valid system property
func ValidateEnvironment(s string) error {
	res, err := fetchObjList(`/environments/`)
	if err != nil {
		return err
	}

	if res.Environments != nil {
		for _, env := range *res.Environments {
			if env.Name == s {
				return nil
			}
		}
	}
	return fmt.Errorf("Invalid environment requested: %s", s)
}

// ValidateStatus tests against the server if string s is a
// valid check instance status
func ValidateStatus(s string) error {
	res, err := fetchObjList(`/status/`)
	if err != nil {
		return err
	}

	if res.Status != nil {
		for _, st := range *res.Status {
			if st.Name == s {
				return nil
			}
		}
	}
	return fmt.Errorf("Invalid instance status requested: %s", s)
}

// ValidatePredicate tests against the server if s is a valid
// predicate
func ValidatePredicate(s string) error {
	res, err := fetchObjList(fmt.Sprintf("/predicates/%s", s))
	if err != nil {
		return err
	}

	if res.Predicates != nil || len(*res.Predicates) == 0 {
		return fmt.Errorf(`no object returned`)
	}

	if s == (*res.Predicates)[0].Symbol {
		return nil
	}
	return fmt.Errorf("Invalid predicate requested: %s", s)
}

// ValidateCheckConstraints tests that all specified check constraints
// resolve to a valid property or attribute.
func ValidateCheckConstraints(repoId, teamId string,
	constraints []proto.CheckConfigConstraint) (
	[]proto.CheckConfigConstraint, error) {
	valid := []proto.CheckConfigConstraint{}

	for _, prop := range constraints {
		switch prop.ConstraintType {
		case `native`:
			if _, err := fetchObjList(
				fmt.Sprintf("/property/native/%s", prop.Native.Name),
			); err != nil {
				return nil, err
			}
			valid = append(valid, prop)

		case `system`:
			if _, err := fetchObjList(
				fmt.Sprintf("/property/system/%s", prop.System.Name),
			); err != nil {
				return nil, err
			}
			valid = append(valid, prop)

		case `attribute`:
			if _, err := fetchObjList(
				fmt.Sprintf("/attributes/%s", prop.Attribute.Name),
			); err != nil {
				return nil, err
			}
			valid = append(valid, prop)

		case `oncall`:
			oncall := proto.PropertyOncall{}
			var err error
			if prop.Oncall.Name != `` {
				if oncall.Id, err = LookupOncallId(
					prop.Oncall.Name); err != nil {
					return nil, err
				}
			} else if prop.Oncall.Id != `` {
				if oncall.Id, err = LookupOncallId(
					prop.Oncall.Id); err != nil {
					return nil, err
				}
			} else {
				return nil, fmt.Errorf(
					`Invalid oncall constraint spec`)
			}
			valid = append(valid, proto.CheckConfigConstraint{
				ConstraintType: prop.ConstraintType,
				Oncall:         &oncall,
			})

		case `service`:
			service := proto.PropertyService{}
			var err error
			if service.Name, err = LookupServicePropertyId(
				prop.Service.Name, teamId); err != nil {
				return nil, err
			}
			service.TeamId = teamId
			valid = append(valid, proto.CheckConfigConstraint{
				ConstraintType: prop.ConstraintType,
				Service:        &service,
			})

		case `custom`:
			custom := proto.PropertyCustom{}
			var err error
			if custom.Id, err = LookupCustomPropertyId(
				prop.Custom.Name, repoId); err != nil {
				return nil, err
			}
			custom.RepositoryId = repoId
			custom.Value = prop.Custom.Value
			valid = append(valid, proto.CheckConfigConstraint{
				ConstraintType: prop.ConstraintType,
				Custom:         &custom,
			})
		}
	}
	return valid, nil
}

// ValidateThresholds tests that all thresholds use the same predicate
// and all referenced levels exist. It normalizes the possible mixed
// use of long and short level names.
func ValidateThresholds(thresholds []proto.CheckConfigThreshold) (
	[]proto.CheckConfigThreshold, error) {
	valid := []proto.CheckConfigThreshold{}
	pred := ``

	for _, thr := range thresholds {
		var err error
		if pred == `` {
			pred = thr.Predicate.Symbol
		} else {
			if pred != thr.Predicate.Symbol {
				return nil, fmt.Errorf(
					"Error, threshold specification is using"+
						" multiple predicates: %s, %s",
					pred,
					thr.Predicate.Symbol,
				)
			}
		}
		t := proto.CheckConfigThreshold{
			Value:     thr.Value,
			Predicate: proto.Predicate{},
			Level:     proto.Level{},
		}
		if t.Level.Name, err = LookupLevelName(
			thr.Level.Name); err != nil {
			return nil, err
		}
		if err = ValidatePredicate(thr.Predicate.Symbol); err != nil {
			return nil, err
		}
		t.Predicate.Symbol = thr.Predicate.Symbol
		valid = append(valid, t)
	}
	return valid, nil
}

// fetchObjList is a helper for ValidateUnit and ValidateProvider
func fetchObjList(path string) (*proto.Result, error) {
	var (
		err  error
		resp *resty.Response
		res  *proto.Result
	)
	if resp, err = GetReq(path); err != nil {
		return nil, err
	}
	if res, err = decodeResponse(resp); err != nil {
		return nil, err
	}
	if err = checkApplicationError(res); err != nil {
		return nil, err
	}
	return res, nil
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>internal/adm: add ValidateInstance()<commit_after>/*-
 * Copyright (c) 2016, 1&1 Internet SE
 * Copyright (c) 2016, Jörg Pernfuß
 *
 * Use of this source code is governed by a 2-clause BSD license
 * that can be found in the LICENSE file.
 */

package adm

import (
	"fmt"
	"net/mail"
	"regexp"
	"strconv"
	"unicode/utf8"

	"github.com/1and1/soma/lib/proto"
	resty "gopkg.in/resty.v0"
)

// ValidateRuneCount tests if a string's number of unicode runes is
// below an upper limit
func ValidateRuneCount(s string, l int) error {
	if cnt := utf8.RuneCountInString(s); cnt > l {
		return fmt.Errorf("Validation error, string '%s' above character limit %d.",
			s, l)
	}
	return nil
}

// ValidateRuneCountRange tests if a string's number of unicode runes is
// between an upper and lower bound
func ValidateRuneCountRange(s string, lower, higher int) error {
	if utf8.RuneCountInString(s) < lower || utf8.RuneCountInString(s) > higher {
		return fmt.Errorf("Validation error, string '%s' outside permitted length."+
			"Required: %d < len(%s) < %d.", s, lower, s, higher)
	}
	return nil
}

// IsUUID validates if a string is one very narrow formatting of a UUID,
// namely the one used by the server. Other valid formats with braces etc
// are not accepted
func IsUUID(s string) bool {
	const reUUID string = `^[[:xdigit:]]{8}-[[:xdigit:]]{4}-[1-5][[:xdigit:]]{3}-[[:xdigit:]]{4}-[[:xdigit:]]{12}$`
	const reUNIL string = `^0{8}-0{4}-0{4}-0{4}-0{12}$`
	re := regexp.MustCompile(fmt.Sprintf("%s|%s", reUUID, reUNIL))

	return re.MatchString(s)
}

// isUint64 validates if a string is a uint64 number
func isUint64(s string) (bool, uint64) {
	if u, err := strconv.ParseUint(s, 10, 64); err != nil {
		return false, 0
	} else {
		return true, u
	}
}

// ValidateOncallNumber tests if a string is a 4 digit number
func ValidateOncallNumber(n string) error {
	num, err := strconv.Atoi(n)
	if err != nil {
		return fmt.Errorf("Syntax error, argument is not a number: %s", err.Error())
	}

	if num <= 0 || num > 9999 {
		return fmt.Errorf("Phone number must be 4-digit extension")
	}
	return nil
}

// ValidateEmployeeNumber tests if a string is a non-negative number
func ValidateEmployeeNumber(s string) error {
	employeeNumber, err := strconv.Atoi(s)
	if err != nil {
		return fmt.Errorf("Error: employeenr argument not a number")
	}
	if employeeNumber < 0 {
		return fmt.Errorf("Negative employee number is not allowed")
	}
	return nil
}

// ValidateMailAddress tests if a string is a mail address
func ValidateMailAddress(s string) error {
	if _, err := mail.ParseAddress(s); err != nil {
		return fmt.Errorf("Error: mailaddr does not parse as RFC 5322 address")
	}
	return nil
}

// ValidateBool tests that string s represents and bool and sets
// its value in b.
// Returns error if s could not be parsed as bool.
func ValidateBool(s string, b *bool) error {
	var err error
	*b, err = strconv.ParseBool(s)
	return err
}

// ValidateLBoundUint64 tests that string s is an unsigned number
// with a specific minimum value and sets it in i.
// Returns error if the number conversion failed or the minimum
// value was not reached.
func ValidateLBoundUint64(s string, i *uint64, min uint64) error {
	var err error
	if *i, err = strconv.ParseUint(s, 10, 64); err != nil {
		return err
	}
	if *i < min {
		return fmt.Errorf("Error, value %d is less than the minimun %d", i, min)
	}
	return nil
}

// ValidateUnit tests against the server if string s is a valid
// unit.
func ValidateUnit(s string) error {
	res, err := fetchObjList(`/units/`)
	if err != nil {
		return err
	}

	if res.Units != nil {
		for _, unit := range *res.Units {
			if unit.Unit == s {
				return nil
			}
		}
	}
	return fmt.Errorf("Value %s is not a valid unit", s)
}

// ValidateProvider tests aginst the server if string s is a valid
// provider.
func ValidateProvider(s string) error {
	res, err := fetchObjList(`/providers/`)
	if err != nil {
		return err
	}

	if res.Providers != nil {
		for _, prov := range *res.Providers {
			if prov.Name == s {
				return nil
			}
		}
	}
	return fmt.Errorf("Value %s is not a valid provider", s)
}

// ValidateSystemProperty tests against the server if string s is a
// valid system property.
func ValidateSystemProperty(s string) error {
	res, err := fetchObjList(`/property/system/`)
	if err != nil {
		return err
	}

	if res.Properties != nil {
		for _, prop := range *res.Properties {
			if prop.System.Name == s {
				return nil
			}
		}
	}
	return fmt.Errorf("Invalid system property requested: %s", s)
}

// ValidateEnvironment tests against the server if string s is a
// valid system property
func ValidateEnvironment(s string) error {
	res, err := fetchObjList(`/environments/`)
	if err != nil {
		return err
	}

	if res.Environments != nil {
		for _, env := range *res.Environments {
			if env.Name == s {
				return nil
			}
		}
	}
	return fmt.Errorf("Invalid environment requested: %s", s)
}

// ValidateStatus tests against the server if string s is a
// valid check instance status
func ValidateStatus(s string) error {
	res, err := fetchObjList(`/status/`)
	if err != nil {
		return err
	}

	if res.Status != nil {
		for _, st := range *res.Status {
			if st.Name == s {
				return nil
			}
		}
	}
	return fmt.Errorf("Invalid instance status requested: %s", s)
}

// ValidateInstance tests against the server if string s is a
// valid check instance id
func ValidateInstance(s string) error {
	if !IsUUID(s) {
		return fmt.Errorf("Argument is not a UUID: %s", s)
	}
	path := fmt.Sprintf("/instances/%s", s)

	if _, err := fetchObjList(path); err != nil {
		return err
	}
	// instanceId is valid if there was no 404/NotFound
	return nil
}

// ValidatePredicate tests against the server if s is a valid
// predicate
func ValidatePredicate(s string) error {
	res, err := fetchObjList(fmt.Sprintf("/predicates/%s", s))
	if err != nil {
		return err
	}

	if res.Predicates != nil || len(*res.Predicates) == 0 {
		return fmt.Errorf(`no object returned`)
	}

	if s == (*res.Predicates)[0].Symbol {
		return nil
	}
	return fmt.Errorf("Invalid predicate requested: %s", s)
}

// ValidateCheckConstraints tests that all specified check constraints
// resolve to a valid property or attribute.
func ValidateCheckConstraints(repoId, teamId string,
	constraints []proto.CheckConfigConstraint) (
	[]proto.CheckConfigConstraint, error) {
	valid := []proto.CheckConfigConstraint{}

	for _, prop := range constraints {
		switch prop.ConstraintType {
		case `native`:
			if _, err := fetchObjList(
				fmt.Sprintf("/property/native/%s", prop.Native.Name),
			); err != nil {
				return nil, err
			}
			valid = append(valid, prop)

		case `system`:
			if _, err := fetchObjList(
				fmt.Sprintf("/property/system/%s", prop.System.Name),
			); err != nil {
				return nil, err
			}
			valid = append(valid, prop)

		case `attribute`:
			if _, err := fetchObjList(
				fmt.Sprintf("/attributes/%s", prop.Attribute.Name),
			); err != nil {
				return nil, err
			}
			valid = append(valid, prop)

		case `oncall`:
			oncall := proto.PropertyOncall{}
			var err error
			if prop.Oncall.Name != `` {
				if oncall.Id, err = LookupOncallId(
					prop.Oncall.Name); err != nil {
					return nil, err
				}
			} else if prop.Oncall.Id != `` {
				if oncall.Id, err = LookupOncallId(
					prop.Oncall.Id); err != nil {
					return nil, err
				}
			} else {
				return nil, fmt.Errorf(
					`Invalid oncall constraint spec`)
			}
			valid = append(valid, proto.CheckConfigConstraint{
				ConstraintType: prop.ConstraintType,
				Oncall:         &oncall,
			})

		case `service`:
			service := proto.PropertyService{}
			var err error
			if service.Name, err = LookupServicePropertyId(
				prop.Service.Name, teamId); err != nil {
				return nil, err
			}
			service.TeamId = teamId
			valid = append(valid, proto.CheckConfigConstraint{
				ConstraintType: prop.ConstraintType,
				Service:        &service,
			})

		case `custom`:
			custom := proto.PropertyCustom{}
			var err error
			if custom.Id, err = LookupCustomPropertyId(
				prop.Custom.Name, repoId); err != nil {
				return nil, err
			}
			custom.RepositoryId = repoId
			custom.Value = prop.Custom.Value
			valid = append(valid, proto.CheckConfigConstraint{
				ConstraintType: prop.ConstraintType,
				Custom:         &custom,
			})
		}
	}
	return valid, nil
}

// ValidateThresholds tests that all thresholds use the same predicate
// and all referenced levels exist. It normalizes the possible mixed
// use of long and short level names.
func ValidateThresholds(thresholds []proto.CheckConfigThreshold) (
	[]proto.CheckConfigThreshold, error) {
	valid := []proto.CheckConfigThreshold{}
	pred := ``

	for _, thr := range thresholds {
		var err error
		if pred == `` {
			pred = thr.Predicate.Symbol
		} else {
			if pred != thr.Predicate.Symbol {
				return nil, fmt.Errorf(
					"Error, threshold specification is using"+
						" multiple predicates: %s, %s",
					pred,
					thr.Predicate.Symbol,
				)
			}
		}
		t := proto.CheckConfigThreshold{
			Value:     thr.Value,
			Predicate: proto.Predicate{},
			Level:     proto.Level{},
		}
		if t.Level.Name, err = LookupLevelName(
			thr.Level.Name); err != nil {
			return nil, err
		}
		if err = ValidatePredicate(thr.Predicate.Symbol); err != nil {
			return nil, err
		}
		t.Predicate.Symbol = thr.Predicate.Symbol
		valid = append(valid, t)
	}
	return valid, nil
}

// fetchObjList is a helper for ValidateUnit and ValidateProvider
func fetchObjList(path string) (*proto.Result, error) {
	var (
		err  error
		resp *resty.Response
		res  *proto.Result
	)
	if resp, err = GetReq(path); err != nil {
		return nil, err
	}
	if res, err = decodeResponse(resp); err != nil {
		return nil, err
	}
	if err = checkApplicationError(res); err != nil {
		return nil, err
	}
	return res, nil
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>package conn

import (
	"github.com/garyburd/redigo/redis"
	"log"
)

type Memo struct {
	Time    string `redis:"time"`
	Content string `redis:"content"`
}

type Task struct {
	Id     int
	ChatId int    `redis:"chatId"`
	Owner  string `redis:"owner"`
	Desc   string `redis:"content"`
	When   string `redis:"time"`
}

//All redis actions

func SetMasterId(id int) {
	c := Pool.Get()
	defer c.Close()
	c.Do("SET", "evolsnowChatId", id)
}

func GetMasterId() int {
	c := Pool.Get()
	defer c.Close()
	id, _ := redis.Int(c.Do("GET", "evolsnowChatId"))
	return id
}

func SetUserChatId(user string, id int) {
	c := Pool.Get()
	defer c.Close()
	key := user + "ChatId"
	c.Do("SET", key, id)
}

func GetUserChatId(user string) int {
	c := Pool.Get()
	defer c.Close()
	key := user + "ChatId"
	id, _ := redis.Int(c.Do("GET", key))
	return id
}

func HSetMemo(user, time, memo string) {
	c := Pool.Get()
	defer c.Close()
	var setMemoLua = `
	local id = redis.call("INCR", "memoIncrId")
	redis.call("RPUSH", KEYS[1]..":memos", id)
	redis.call("HMSET", "memo:"..id, "time", KEYS[2], "content", KEYS[3])
	`
	script := redis.NewScript(3, setMemoLua)
	script.Do(c, user, time, memo)
}

func GetTaskId() int {
	c := Pool.Get()
	defer c.Close()
	id, _ := redis.Int(c.Do("INCR", "taskIncrId"))
	return id
}

func HSetTask(ts Task) {
	c := Pool.Get()
	defer c.Close()
	log.Println("save task")
	var setTaskLua = `
	redis.call("RPUSH", KEYS[2]..":tasks", id)
	redis.call("HMSET", "task:"..KEYS[1], "owner", KEYS[2], "time", KEYS[3], "content", KEYS[4], "chatID", KEYS[5])
	`
	script := redis.NewScript(5, setTaskLua)
	script.Do(c, ts.Id, ts.Owner, ts.When, ts.Desc, ts.ChatId)
}

func RemoveTask(ts Task) {
	c := Pool.Get()
	defer c.Close()
	log.Println("remove", ts.Id)
	var removeTaskLua = `
	redis.call("LREM", KEYS[1]..":tasks", 1, KEYS[2])
	redis.call("DEL", "task:"..KEYS[2])
	`
	script := redis.NewScript(2, removeTaskLua)
	script.Do(c, ts.Owner, ts.Id)
}

func HGetUserTasks(user string) []Task {
	c := Pool.Get()
	defer c.Close()
	var multiGetTaskLua = `
	local data = redis.call("LRANGE", KEYS[1]..":tasks", "0", "-1")
	local ret = {}
  	for idx=1, #data do
  		ret[idx] = redis.call("HGETALL", "task:"..data[idx])
  		print(ret[idx])
  	end
  	return ret
   `
	var tasks []Task
	script := redis.NewScript(1, multiGetTaskLua)
	values, err := redis.Values(script.Do(c, user))
	if err != nil {
		log.Println(err)
	}
	for i := range values {
		t := new(Task)
		redis.ScanStruct(values[i].([]interface{}), t)
		log.Println("task:", *t)
		tasks = append(tasks, *t)
	}
	return tasks
}

func HGetAllMemos(user string) []Memo {
	c := Pool.Get()
	defer c.Close()
	var multiGetMemoLua = `
	local data = redis.call("LRANGE", KEYS[1]..":memos", "0", "-1")
	local ret = {}
  	for idx=1, #data do
  		ret[idx] = redis.call("HGETALL", "memo:"..data[idx])
  	end
  	return ret
   `
	var memos []Memo
	script := redis.NewScript(1, multiGetMemoLua)
	values, err := redis.Values(script.Do(c, user))
	if err != nil {
		log.Println(err)
	}
	for i := range values {
		m := new(Memo)
		redis.ScanStruct(values[i].([]interface{}), m)
		memos = append(memos, *m)
	}
	return memos
}

//
//var multiGetScript = redis.NewScript(0, multiGetMemoLua)
<commit_msg>get tasks<commit_after>package conn

import (
	"github.com/garyburd/redigo/redis"
	"log"
)

type Memo struct {
	Time    string `redis:"time"`
	Content string `redis:"content"`
}

type Task struct {
	Id     int
	ChatId int    `redis:"chatId"`
	Owner  string `redis:"owner"`
	Desc   string `redis:"content"`
	When   string `redis:"time"`
}

//All redis actions

func SetMasterId(id int) {
	c := Pool.Get()
	defer c.Close()
	c.Do("SET", "evolsnowChatId", id)
}

func GetMasterId() int {
	c := Pool.Get()
	defer c.Close()
	id, _ := redis.Int(c.Do("GET", "evolsnowChatId"))
	return id
}

func SetUserChatId(user string, id int) {
	c := Pool.Get()
	defer c.Close()
	key := user + "ChatId"
	c.Do("SET", key, id)
}

func GetUserChatId(user string) int {
	c := Pool.Get()
	defer c.Close()
	key := user + "ChatId"
	id, _ := redis.Int(c.Do("GET", key))
	return id
}

func HSetMemo(user, time, memo string) {
	c := Pool.Get()
	defer c.Close()
	var setMemoLua = `
	local id = redis.call("INCR", "memoIncrId")
	redis.call("RPUSH", KEYS[1]..":memos", id)
	redis.call("HMSET", "memo:"..id, "time", KEYS[2], "content", KEYS[3])
	`
	script := redis.NewScript(3, setMemoLua)
	script.Do(c, user, time, memo)
}

func GetTaskId() int {
	c := Pool.Get()
	defer c.Close()
	id, _ := redis.Int(c.Do("INCR", "taskIncrId"))
	return id
}

func HSetTask(ts Task) {
	c := Pool.Get()
	defer c.Close()
	log.Println("save task")
	var setTaskLua = `
	redis.call("RPUSH", KEYS[2]..":tasks", id)
	redis.call("HMSET", "task:"..KEYS[1], "owner", KEYS[2], "time", KEYS[3], "content", KEYS[4], "chatID", KEYS[5])
	`
	script := redis.NewScript(5, setTaskLua)
	if _, err := script.Do(c, ts.Id, ts.Owner, ts.When, ts.Desc, ts.ChatId); err != nil {
		log.Println(err)
	}
}

func RemoveTask(ts Task) {
	c := Pool.Get()
	defer c.Close()
	log.Println("remove", ts.Id)
	var removeTaskLua = `
	redis.call("LREM", KEYS[1]..":tasks", 1, KEYS[2])
	redis.call("DEL", "task:"..KEYS[2])
	`
	script := redis.NewScript(2, removeTaskLua)
	script.Do(c, ts.Owner, ts.Id)
}

func HGetUserTasks(user string) []Task {
	c := Pool.Get()
	defer c.Close()
	var multiGetTaskLua = `
	local data = redis.call("LRANGE", KEYS[1]..":tasks", "0", "-1")
	local ret = {}
  	for idx=1, #data do
  		ret[idx] = redis.call("HGETALL", "task:"..data[idx])
  		print(ret[idx])
  	end
  	return ret
   `
	var tasks []Task
	script := redis.NewScript(1, multiGetTaskLua)
	values, err := redis.Values(script.Do(c, user))
	if err != nil {
		log.Println(err)
	}
	for i := range values {
		t := new(Task)
		redis.ScanStruct(values[i].([]interface{}), t)
		log.Println("task:", *t)
		tasks = append(tasks, *t)
	}
	return tasks
}

func HGetAllMemos(user string) []Memo {
	c := Pool.Get()
	defer c.Close()
	var multiGetMemoLua = `
	local data = redis.call("LRANGE", KEYS[1]..":memos", "0", "-1")
	local ret = {}
  	for idx=1, #data do
  		ret[idx] = redis.call("HGETALL", "memo:"..data[idx])
  	end
  	return ret
   `
	var memos []Memo
	script := redis.NewScript(1, multiGetMemoLua)
	values, err := redis.Values(script.Do(c, user))
	if err != nil {
		log.Println(err)
	}
	for i := range values {
		m := new(Memo)
		redis.ScanStruct(values[i].([]interface{}), m)
		memos = append(memos, *m)
	}
	return memos
}

//
//var multiGetScript = redis.NewScript(0, multiGetMemoLua)
<|endoftext|>


<commit_before>//
// Copyright 2018, Sander van Harmelen
//
// 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 gitlab

import (
	"fmt"
	"time"
)

// DeploymentsService handles communication with the deployment related methods
// of the GitLab API.
//
// GitLab API docs: https://docs.gitlab.com/ce/api/deployments.html
type DeploymentsService struct {
	client *Client
}

// Deployment represents the Gitlab deployment
type Deployment struct {
	ID          int          `json:"id"`
	IID         int          `json:"iid"`
	Ref         string       `json:"ref"`
	SHA         string       `json:"sha"`
	CreatedAt   *time.Time   `json:"created_at"`
	UpdatedAt   *time.Time   `json:"updated_at"`
	User        *ProjectUser `json:"user"`
	Environment *Environment `json:"environment"`
	Deployable  struct {
		ID         int        `json:"id"`
		Status     string     `json:"status"`
		Stage      string     `json:"stage"`
		Name       string     `json:"name"`
		Ref        string     `json:"ref"`
		Tag        bool       `json:"tag"`
		Coverage   float64    `json:"coverage"`
		CreatedAt  *time.Time `json:"created_at"`
		StartedAt  *time.Time `json:"started_at"`
		FinishedAt *time.Time `json:"finished_at"`
		Duration   float64    `json:"duration"`
		User       *User      `json:"user"`
		Commit     *Commit    `json:"commit"`
		Pipeline   struct {
			ID        int    `json:"id"`
			SHA       string `json:"sha"`
			Ref       string `json:"ref"`
			Status    string `json:"status"`
			CreatedAt string `json:"created_at"`
			UpdatedAt string `json:"updated_at"`
		} `json:"pipeline"`
		Runner *Runner `json:"runner"`
	} `json:"deployable"`
}

// ListProjectDeploymentsOptions represents the available ListProjectDeployments() options.
//
// GitLab API docs:
// https://docs.gitlab.com/ce/api/deployments.html#list-project-deployments
type ListProjectDeploymentsOptions struct {
	ListOptions
	OrderBy       *string    `url:"order_by,omitempty" json:"order_by,omitempty"`
	Sort          *string    `url:"sort,omitempty" json:"sort,omitempty"`
	UpdatedAfter  *time.Time `url:"updated_after,omitempty" json:"updated_after,omitempty"`
	UpdatedBefore *time.Time `url:"update_before,omitempty" json:"updated_before,omitempty"`
	Environment   *string    `url:"environment,omitempty" json:"environment,omitempty"`
	Status        *string    `url:"status,omitempty" json:"status,omitempty"`
}

// ListProjectDeployments gets a list of deployments in a project.
//
// GitLab API docs: https://docs.gitlab.com/ce/api/deployments.html#list-project-deployments
func (s *DeploymentsService) ListProjectDeployments(pid interface{}, opts *ListProjectDeploymentsOptions, options ...RequestOptionFunc) ([]*Deployment, *Response, error) {
	project, err := parseID(pid)
	if err != nil {
		return nil, nil, err
	}
	u := fmt.Sprintf("projects/%s/deployments", pathEscape(project))

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

	var ds []*Deployment
	resp, err := s.client.Do(req, &ds)
	if err != nil {
		return nil, resp, err
	}

	return ds, resp, err
}

// GetProjectDeployment get a deployment for a project.
//
// GitLab API docs: https://docs.gitlab.com/ce/api/deployments.html#get-a-specific-deployment
func (s *DeploymentsService) GetProjectDeployment(pid interface{}, deployment int, options ...RequestOptionFunc) (*Deployment, *Response, error) {
	project, err := parseID(pid)
	if err != nil {
		return nil, nil, err
	}
	u := fmt.Sprintf("projects/%s/deployments/%d", pathEscape(project), deployment)

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

	d := new(Deployment)
	resp, err := s.client.Do(req, d)
	if err != nil {
		return nil, resp, err
	}

	return d, resp, err
}

// CreateProjectDeploymentOptions represents the available
// CreateProjectDeployment() options.
//
// GitLab API docs: https://docs.gitlab.com/ee/api/deployments.html#create-a-deployment
type CreateProjectDeploymentOptions struct {
	Environment *string                `url:"environment,omitempty" json:"environment,omitempty"`
	Ref         *string                `url:"ref,omitempty" json:"ref,omitempty"`
	SHA         *string                `url:"sha,omitempty" json:"sha,omitempty"`
	Tag         *bool                  `url:"tag,omitempty" json:"tag,omitempty"`
	Status      *DeploymentStatusValue `url:"status,omitempty" json:"status,omitempty"`
}

// CreateProjectDeployment creates a project deployment.
//
// GitLab API docs: https://docs.gitlab.com/ee/api/deployments.html#create-a-deployment
func (s *DeploymentsService) CreateProjectDeployment(pid interface{}, opt *CreateProjectDeploymentOptions, options ...RequestOptionFunc) (*Deployment, *Response, error) {
	project, err := parseID(pid)
	if err != nil {
		return nil, nil, err
	}
	u := fmt.Sprintf("projects/%s/deployments", pathEscape(project))

	req, err := s.client.NewRequest("POST", u, opt, options)
	if err != nil {
		return nil, nil, err
	}

	d := new(Deployment)
	resp, err := s.client.Do(req, &d)
	if err != nil {
		return nil, resp, err
	}

	return d, resp, err
}

// UpdateProjectDeploymentOptions represents the available
// UpdateProjectDeployment() options.
//
// GitLab API docs: https://docs.gitlab.com/ee/api/deployments.html#updating-a-deployment
type UpdateProjectDeploymentOptions struct {
	Status *DeploymentStatusValue `url:"status,omitempty" json:"status,omitempty"`
}

// UpdateProjectDeployment updates a project deployment.
//
// GitLab API docs: https://docs.gitlab.com/ee/api/deployments.html#updating-a-deployment
func (s *DeploymentsService) UpdateProjectDeployment(pid interface{}, deployment int, opt *UpdateProjectDeploymentOptions, options ...RequestOptionFunc) (*Deployment, *Response, error) {
	project, err := parseID(pid)
	if err != nil {
		return nil, nil, err
	}
	u := fmt.Sprintf("projects/%s/deployments/%d", pathEscape(project), deployment)

	req, err := s.client.NewRequest("PUT", u, opt, options)
	if err != nil {
		return nil, nil, err
	}

	d := new(Deployment)
	resp, err := s.client.Do(req, &d)
	if err != nil {
		return nil, resp, err
	}

	return d, resp, err
}
<commit_msg>update data type for CreatedAt and UpdatedAt fields in Pipeline in Deployment<commit_after>//
// Copyright 2018, Sander van Harmelen
//
// 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 gitlab

import (
	"fmt"
	"time"
)

// DeploymentsService handles communication with the deployment related methods
// of the GitLab API.
//
// GitLab API docs: https://docs.gitlab.com/ce/api/deployments.html
type DeploymentsService struct {
	client *Client
}

// Deployment represents the Gitlab deployment
type Deployment struct {
	ID          int          `json:"id"`
	IID         int          `json:"iid"`
	Ref         string       `json:"ref"`
	SHA         string       `json:"sha"`
	CreatedAt   *time.Time   `json:"created_at"`
	UpdatedAt   *time.Time   `json:"updated_at"`
	User        *ProjectUser `json:"user"`
	Environment *Environment `json:"environment"`
	Deployable  struct {
		ID         int        `json:"id"`
		Status     string     `json:"status"`
		Stage      string     `json:"stage"`
		Name       string     `json:"name"`
		Ref        string     `json:"ref"`
		Tag        bool       `json:"tag"`
		Coverage   float64    `json:"coverage"`
		CreatedAt  *time.Time `json:"created_at"`
		StartedAt  *time.Time `json:"started_at"`
		FinishedAt *time.Time `json:"finished_at"`
		Duration   float64    `json:"duration"`
		User       *User      `json:"user"`
		Commit     *Commit    `json:"commit"`
		Pipeline   struct {
			ID        int        `json:"id"`
			SHA       string     `json:"sha"`
			Ref       string     `json:"ref"`
			Status    string     `json:"status"`
			CreatedAt *time.Time `json:"created_at"`
			UpdatedAt *time.Time `json:"updated_at"`
		} `json:"pipeline"`
		Runner *Runner `json:"runner"`
	} `json:"deployable"`
}

// ListProjectDeploymentsOptions represents the available ListProjectDeployments() options.
//
// GitLab API docs:
// https://docs.gitlab.com/ce/api/deployments.html#list-project-deployments
type ListProjectDeploymentsOptions struct {
	ListOptions
	OrderBy       *string    `url:"order_by,omitempty" json:"order_by,omitempty"`
	Sort          *string    `url:"sort,omitempty" json:"sort,omitempty"`
	UpdatedAfter  *time.Time `url:"updated_after,omitempty" json:"updated_after,omitempty"`
	UpdatedBefore *time.Time `url:"update_before,omitempty" json:"updated_before,omitempty"`
	Environment   *string    `url:"environment,omitempty" json:"environment,omitempty"`
	Status        *string    `url:"status,omitempty" json:"status,omitempty"`
}

// ListProjectDeployments gets a list of deployments in a project.
//
// GitLab API docs: https://docs.gitlab.com/ce/api/deployments.html#list-project-deployments
func (s *DeploymentsService) ListProjectDeployments(pid interface{}, opts *ListProjectDeploymentsOptions, options ...RequestOptionFunc) ([]*Deployment, *Response, error) {
	project, err := parseID(pid)
	if err != nil {
		return nil, nil, err
	}
	u := fmt.Sprintf("projects/%s/deployments", pathEscape(project))

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

	var ds []*Deployment
	resp, err := s.client.Do(req, &ds)
	if err != nil {
		return nil, resp, err
	}

	return ds, resp, err
}

// GetProjectDeployment get a deployment for a project.
//
// GitLab API docs: https://docs.gitlab.com/ce/api/deployments.html#get-a-specific-deployment
func (s *DeploymentsService) GetProjectDeployment(pid interface{}, deployment int, options ...RequestOptionFunc) (*Deployment, *Response, error) {
	project, err := parseID(pid)
	if err != nil {
		return nil, nil, err
	}
	u := fmt.Sprintf("projects/%s/deployments/%d", pathEscape(project), deployment)

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

	d := new(Deployment)
	resp, err := s.client.Do(req, d)
	if err != nil {
		return nil, resp, err
	}

	return d, resp, err
}

// CreateProjectDeploymentOptions represents the available
// CreateProjectDeployment() options.
//
// GitLab API docs: https://docs.gitlab.com/ee/api/deployments.html#create-a-deployment
type CreateProjectDeploymentOptions struct {
	Environment *string                `url:"environment,omitempty" json:"environment,omitempty"`
	Ref         *string                `url:"ref,omitempty" json:"ref,omitempty"`
	SHA         *string                `url:"sha,omitempty" json:"sha,omitempty"`
	Tag         *bool                  `url:"tag,omitempty" json:"tag,omitempty"`
	Status      *DeploymentStatusValue `url:"status,omitempty" json:"status,omitempty"`
}

// CreateProjectDeployment creates a project deployment.
//
// GitLab API docs: https://docs.gitlab.com/ee/api/deployments.html#create-a-deployment
func (s *DeploymentsService) CreateProjectDeployment(pid interface{}, opt *CreateProjectDeploymentOptions, options ...RequestOptionFunc) (*Deployment, *Response, error) {
	project, err := parseID(pid)
	if err != nil {
		return nil, nil, err
	}
	u := fmt.Sprintf("projects/%s/deployments", pathEscape(project))

	req, err := s.client.NewRequest("POST", u, opt, options)
	if err != nil {
		return nil, nil, err
	}

	d := new(Deployment)
	resp, err := s.client.Do(req, &d)
	if err != nil {
		return nil, resp, err
	}

	return d, resp, err
}

// UpdateProjectDeploymentOptions represents the available
// UpdateProjectDeployment() options.
//
// GitLab API docs: https://docs.gitlab.com/ee/api/deployments.html#updating-a-deployment
type UpdateProjectDeploymentOptions struct {
	Status *DeploymentStatusValue `url:"status,omitempty" json:"status,omitempty"`
}

// UpdateProjectDeployment updates a project deployment.
//
// GitLab API docs: https://docs.gitlab.com/ee/api/deployments.html#updating-a-deployment
func (s *DeploymentsService) UpdateProjectDeployment(pid interface{}, deployment int, opt *UpdateProjectDeploymentOptions, options ...RequestOptionFunc) (*Deployment, *Response, error) {
	project, err := parseID(pid)
	if err != nil {
		return nil, nil, err
	}
	u := fmt.Sprintf("projects/%s/deployments/%d", pathEscape(project), deployment)

	req, err := s.client.NewRequest("PUT", u, opt, options)
	if err != nil {
		return nil, nil, err
	}

	d := new(Deployment)
	resp, err := s.client.Do(req, &d)
	if err != nil {
		return nil, resp, err
	}

	return d, resp, err
}
<|endoftext|>


<commit_before>// Package depth parallelizes calls to samtools depths and outputs:
// 1) $prefix.callable.bed that contains collapsed per-base regions of NO/LOW/or CALLABLE coverage.
// where low is < MinCov.
// 2) $prefix.depth.bed that contains the average depth for each window interval specified by WindowSize.
// TODO: output gc-content in depth windows.
package depth

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"

	arg "github.com/alexflint/go-arg"
	"github.com/brentp/faidx"
	"github.com/brentp/gargs/process"
	"github.com/brentp/xopen"
	"github.com/fatih/color"
)

type dargs struct {
	WindowSize   int    `arg:"-w,help:window size in which to calculate high-depth regions"`
	MaxMeanDepth int    `arg:"-m,help:windows with depth > than this are high-depth. The default reports the depth of all regions."`
	Q            int    `arg:"-Q,help:mapping quality cutoff"`
	Chrom        string `arg:"-c,help:optional chromosome to limit analysis"`
	MinCov       int    `arg:"help:minimum depth considered callable"`
	Stats        bool   `arg:"-s,help:report sequence stats [GC CpG masked] for each window"`
	Reference    string `arg:"-r,required,help:path to reference fasta"`
	Processes    int    `arg:"-p,help:number of processors to parallelize."`
	Bed          string `arg:"-b,help:optional file of positions or regions to restrict depth calculations."`
	Prefix       string `arg:"positional,required,help:prefix for output files [\-depth.bed, \-callable.bed]"`
	Bam          string `arg:"positional,required,help:bam for which to calculate depth"`
}

// we echo the region first so the callback knows the full extents even if there is NOTE
// coverage for part of it.
const command = "echo %s; samtools depth -Q %d -d %d -r %s --reference %s %s"

// this is the size in basepairs of the genomic chunks for parallelization.
const step = 5000000

var exitCode = 0

func pcheck(e error) {
	if e != nil {
		log.Fatal(e)
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// match chrom:start-end and chrom\tstart\tend
var re = regexp.MustCompile("(.+?)[:\t](\\d+)([\\-\t])(\\d+).*?")

func chromStartEndFromLine(line []byte) (string, int, int) {
	ret := re.FindSubmatch(line)
	if len(ret) != 5 {
		log.Fatal("couldn't get region from line", line)
	}
	chrom, start, isep, end := ret[1], ret[2], ret[3], ret[4]
	// convert from bed to chrom:start-end region so add 1 to start
	istart, err := strconv.Atoi(string(start))
	if err != nil {
		log.Fatal(err)
	}
	if bytes.Equal(isep, []byte{'-'}) {
		istart--
	}
	iend, err := strconv.Atoi(string(end))
	if err != nil {
		log.Fatal(err)
	}
	return string(chrom), istart, iend
}

func regionFromLine(line []byte) string {
	chrom, start, end := chromStartEndFromLine(line)
	// convert from bed to chrom:start-end region so add 1 to start
	return fmt.Sprintf("%s:%d-%d", chrom, start+1, end)
}

// when the user specified a Bed file of regions for coverage, this is used.
func genFromBed(ch chan string, args dargs) {
	rdr, err := xopen.Ropen(args.Bed)
	pcheck(err)
	for {
		line, err := rdr.ReadBytes('\n')
		if err == io.EOF {
			break
		}
		pcheck(err)
		if len(line) == 0 {
			continue
		}
		region := regionFromLine(line)
		ch <- fmt.Sprintf(command, region, args.Q, max(args.MaxMeanDepth+1000, 8000),
			region, args.Reference, args.Bam)
	}
	close(ch)
}

func genCommands(args dargs) chan string {
	ch := make(chan string)
	if args.Bed != "" {
		go genFromBed(ch, args)
		return ch
	}
	go func() {

		rdr, err := xopen.Ropen(args.Reference + ".fai")
		pcheck(err)
		for {
			line, err := rdr.ReadString('\n')
			if err == io.EOF {
				break
			}
			pcheck(err)
			toks := strings.Split(line, "\t")

			chrom := toks[0]
			if args.Chrom != "" && chrom != args.Chrom {
				continue
			}
			length, err := strconv.Atoi(toks[1])
			pcheck(err)
			for i := 0; i < length; i += step {
				region := fmt.Sprintf("%s:%d-%d", chrom, i, min(i+step, length))
				ch <- fmt.Sprintf(command, region, args.Q, max(args.MaxMeanDepth+1000, 8000),
					region, args.Reference, args.Bam)
			}
		}
		close(ch)
	}()
	return ch
}

// Main is run from the dispatcher
func Main() {

	args := dargs{WindowSize: 250,
		MaxMeanDepth: 0,
		MinCov:       4,
		Q:            1}
	p := arg.MustParse(&args)
	if args.Prefix == "" {
		p.Fail("you must specify an output prefix")
	}
	runtime.GOMAXPROCS(args.Processes)
	run(args)
	os.Exit(exitCode)
}

type ipos struct {
	chrom string
	pos   int
	depth int
}

func mean(sl []int) float64 {
	if len(sl) == 0 {
		return 0
	}
	sum := 0
	for _, v := range sl {
		sum += v
	}
	return float64(sum) / float64(len(sl))
}

func getStats(fa *faidx.Faidx, chrom string, start, end int) string {
	if fa == nil {
		return ""
	}
	st, err := fa.Stats(chrom, start, end)
	if err != nil {
		log.Println(err)
	}
	return fmt.Sprintf("\t%.2f\t%.2f\t%.2f", st.GC, st.CpG, st.Masked)
}

func run(args dargs) {

	callback := func(r io.Reader, w io.WriteCloser) error {
		rdr := bufio.NewReader(r)
		wtr := bufio.NewWriter(w)

		var fa *faidx.Faidx
		var err error
		if args.Stats {
			fa, err = faidx.New(args.Reference)
			defer fa.Close()
		}
		if err != nil {
			return err
		}

		defer w.Close()
		defer wtr.Flush()
		cache := make([]ipos, 2)
		depthCache := make([]int, 0, args.WindowSize)
		var depth, pos int
		var lastCovClass, lastChrom string
		lastWindow := -1

		region, err := rdr.ReadBytes('\n')
		if err != nil {
			return err
		}
		chrom, start, regionEnd := chromStartEndFromLine(region)
		lastWindow = start/args.WindowSize - 1

		hdPath := filepath.Join(os.TempDir(), fmt.Sprintf("%s-%s-%s.depth.bed", chrom, start, regionEnd))
		fhHD, ferr := xopen.Wopen(hdPath)
		if ferr != nil {
			return ferr
		}
		caPath := filepath.Join(os.TempDir(), fmt.Sprintf("%s-%s-%s.callable.bed", chrom, start, regionEnd))
		fhCA, ferr := xopen.Wopen(caPath)
		if ferr != nil {
			return ferr
		}
		defer fhCA.Close()
		defer fhHD.Close()

		line, err := rdr.ReadString('\n')
		for err == nil {

			chrom = line[:strings.Index(line, "\t")]

			// toks starts after chrom. so [0] is pos and [1] is depth.
			toks := strings.SplitN(line[len(chrom)+1:], "\t", 2)
			pos, err = strconv.Atoi(toks[0])
			if err != nil {
				break
			}

			if len(toks[1]) > 1 && toks[1][len(toks[1])-1] == '\n' {
				toks[1] = toks[1][:len(toks[1])-1]
			}

			depth, err = strconv.Atoi(toks[1])
			if err != nil {
				break
			}
			// if we have a full window...
			if pos/args.WindowSize != lastWindow {
				thisWindow := pos / args.WindowSize
				// fill in empty windows for which we had no coverage.
				for iwindow := lastWindow + 1; iwindow < thisWindow; iwindow++ {
					s := iwindow * args.WindowSize
					stats := getStats(fa, chrom, s, s+args.WindowSize)
					fhHD.WriteString(fmt.Sprintf("%s\t%d\t%d\t%d%s\n", chrom, s, s+args.WindowSize, 0, stats))
				}
				// fill in this window:
				s := thisWindow * args.WindowSize
				// we dont just break because we still need to print the stuff to the callable file below.
				if s < regionEnd {
					// NOTE: we include the full bin, even for the last postion in the region.w
					// may want to truncate to min(s, s+args.WindowSize)
					stats := getStats(fa, chrom, thisWindow, thisWindow+args.WindowSize)
					fhHD.WriteString(fmt.Sprintf("%s\t%d\t%d\t%.2f%s\n", chrom, s, s+args.WindowSize, mean(depthCache), stats))
					depthCache = depthCache[:0]
					lastWindow = thisWindow
				}
			}
			depthCache = append(depthCache, depth)

			covClass := "CALLABLE"
			if depth == 0 {
				covClass = "NO_COVERAGE"
			} else if depth < args.MinCov {
				covClass = "LOW_COVERAGE"
			} else if args.MaxMeanDepth > 0 && depth >= args.MaxMeanDepth {
				covClass = "EXCESSIVE_COVERAGE"
			}
			// check for a gap or a change in the coverage class.
			if covClass != lastCovClass || pos != cache[1].pos+1 {
				if lastChrom != "" && cache[1].pos != 0 {
					fhCA.WriteString(fmt.Sprintf("%s\t%d\t%d\t%s\n", lastChrom, cache[0].pos-1, cache[1].pos, lastCovClass))
					// also fill in block without any coverage.
					if pos != cache[1].pos+1 {
						fhCA.WriteString(fmt.Sprintf("%s\t%d\t%d\t%s\n", lastChrom, cache[1].pos-1, pos, "NO_COVERAGE"))
					}
				}
				lastCovClass = covClass
				lastChrom = chrom
				cache = cache[:0]
				// append twice so we can safely use cache[1] above.
				cache = append(cache, ipos{chrom, pos, depth})
				cache = append(cache, ipos{chrom, pos, depth})
			} else {
				cache[1].pos = pos
			}
			line, err = rdr.ReadString('\n')
		}
		if len(cache) > 1 {
			fhCA.WriteString(fmt.Sprintf("%s\t%d\t%d\t%s\n", lastChrom, cache[0].pos-1, cache[1].pos, lastCovClass))
		}
		// we didn't get data for the full region, so it must end in no-coverage.
		if cache[1].pos < regionEnd {
			fhCA.WriteString(fmt.Sprintf("%s\t%d\t%d\tNO_COVERAGE\n", lastChrom, cache[1].pos, regionEnd))
			for ds := cache[1].pos / args.WindowSize * args.WindowSize; ds < regionEnd; ds += args.WindowSize {
				thisWindow := ds / args.WindowSize
				stats := getStats(fa, chrom, thisWindow, thisWindow+args.WindowSize)
				log.Println(stats)
				fhHD.WriteString(fmt.Sprintf("%s\t%d\t%d\t%.2f%s\n", chrom, ds, ds+args.WindowSize, mean(depthCache), stats))
				depthCache = depthCache[:0]
			}
		}
		wtr.WriteString(caPath + "\n")
		wtr.WriteString(hdPath + "\n")
		wtr.Flush()
		return w.Close()
	}

	cancel := make(chan bool)
	defer close(cancel)
	stdout := bufio.NewWriter(os.Stdout)
	defer stdout.Flush()

	chrom := ""
	if args.Chrom != "" {
		chrom = "." + args.Chrom
	}
	fhca, err := xopen.Wopen(fmt.Sprintf("%s%s.callable.bed", args.Prefix, chrom))
	pcheck(err)
	fhhd, err := xopen.Wopen(fmt.Sprintf("%s%s.depth.bed", args.Prefix, chrom))
	pcheck(err)
	defer fhca.Flush()
	defer fhhd.Flush()

	opts := process.Options{Retries: 1, CallBack: callback, Ordered: false}

	for cmd := range process.Runner(genCommands(args), cancel, &opts) {
		if ex := cmd.ExitCode(); ex != 0 && cmd.Err != io.EOF {
			c := color.New(color.BgRed).Add(color.Bold)
			fmt.Fprintf(os.Stderr, "%s\n", c.SprintFunc()(fmt.Sprintf("ERROR with command: %s", cmd)))
			exitCode = max(exitCode, ex)
		}
		if cmd.Err == io.EOF {
			continue
		}
		caPath, err := cmd.ReadString('\n')
		if err != nil {
			log.Println(cmd.CmdStr, err, cmd.Err)
		}
		caSrc, err := xopen.Ropen(strings.TrimSpace(caPath))
		pcheck(err)
		io.Copy(fhca, caSrc)
		os.Remove(strings.TrimSpace(caPath))

		hdPath, err := cmd.ReadString('\n')
		if err != nil {
			log.Println(err)
		}
		hdSrc, err := xopen.Ropen(strings.TrimSpace(hdPath))
		pcheck(err)
		io.Copy(fhhd, hdSrc)
		os.Remove(strings.TrimSpace(hdPath))
	}
	fhca.Flush()
	fhca.Close()
	fhhd.Flush()
	fhhd.Close()

}
<commit_msg>add -o option to force ordered output<commit_after>// Package depth parallelizes calls to samtools depths and outputs:
// 1) $prefix.callable.bed that contains collapsed per-base regions of NO/LOW/or CALLABLE coverage.
// where low is < MinCov.
// 2) $prefix.depth.bed that contains the average depth for each window interval specified by WindowSize.
// TODO: output gc-content in depth windows.
package depth

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"

	arg "github.com/alexflint/go-arg"
	"github.com/brentp/faidx"
	"github.com/brentp/gargs/process"
	"github.com/brentp/xopen"
	"github.com/fatih/color"
)

type dargs struct {
	WindowSize   int    `arg:"-w,help:window size in which to calculate high-depth regions"`
	MaxMeanDepth int    `arg:"-m,help:windows with depth > than this are high-depth. The default reports the depth of all regions."`
	Ordered      bool   `args:-o,help:force output to be in same order as input even with -p."`
	Q            int    `arg:"-Q,help:mapping quality cutoff"`
	Chrom        string `arg:"-c,help:optional chromosome to limit analysis"`
	MinCov       int    `arg:"help:minimum depth considered callable"`
	Stats        bool   `arg:"-s,help:report sequence stats [GC CpG masked] for each window"`
	Reference    string `arg:"-r,required,help:path to reference fasta"`
	Processes    int    `arg:"-p,help:number of processors to parallelize."`
	Bed          string `arg:"-b,help:optional file of positions or regions to restrict depth calculations."`
	Prefix       string `arg:"positional,required,help:prefix for output files [\-depth.bed, \-callable.bed]"`
	Bam          string `arg:"positional,required,help:bam for which to calculate depth"`
}

// we echo the region first so the callback knows the full extents even if there is NOTE
// coverage for part of it.
const command = "echo %s; samtools depth -Q %d -d %d -r %s --reference %s %s"

// this is the size in basepairs of the genomic chunks for parallelization.
const step = 5000000

var exitCode = 0

func pcheck(e error) {
	if e != nil {
		log.Fatal(e)
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// match chrom:start-end and chrom\tstart\tend
var re = regexp.MustCompile("(.+?)[:\t](\\d+)([\\-\t])(\\d+).*?")

func chromStartEndFromLine(line []byte) (string, int, int) {
	ret := re.FindSubmatch(line)
	if len(ret) != 5 {
		log.Fatal("couldn't get region from line", line)
	}
	chrom, start, isep, end := ret[1], ret[2], ret[3], ret[4]
	// convert from bed to chrom:start-end region so add 1 to start
	istart, err := strconv.Atoi(string(start))
	if err != nil {
		log.Fatal(err)
	}
	if bytes.Equal(isep, []byte{'-'}) {
		istart--
	}
	iend, err := strconv.Atoi(string(end))
	if err != nil {
		log.Fatal(err)
	}
	return string(chrom), istart, iend
}

func regionFromLine(line []byte) string {
	chrom, start, end := chromStartEndFromLine(line)
	// convert from bed to chrom:start-end region so add 1 to start
	return fmt.Sprintf("%s:%d-%d", chrom, start+1, end)
}

// when the user specified a Bed file of regions for coverage, this is used.
func genFromBed(ch chan string, args dargs) {
	rdr, err := xopen.Ropen(args.Bed)
	pcheck(err)
	for {
		line, err := rdr.ReadBytes('\n')
		if err == io.EOF {
			break
		}
		pcheck(err)
		if len(line) == 0 {
			continue
		}
		region := regionFromLine(line)
		ch <- fmt.Sprintf(command, region, args.Q, args.MaxMeanDepth+1000,
			region, args.Reference, args.Bam)
	}
	close(ch)
}

func genCommands(args dargs) chan string {
	ch := make(chan string)
	if args.Bed != "" {
		go genFromBed(ch, args)
		return ch
	}
	go func() {

		rdr, err := xopen.Ropen(args.Reference + ".fai")
		pcheck(err)
		for {
			line, err := rdr.ReadString('\n')
			if err == io.EOF {
				break
			}
			pcheck(err)
			toks := strings.Split(line, "\t")

			chrom := toks[0]
			if args.Chrom != "" && chrom != args.Chrom {
				continue
			}
			length, err := strconv.Atoi(toks[1])
			pcheck(err)
			for i := 0; i < length; i += step {
				region := fmt.Sprintf("%s:%d-%d", chrom, i, min(i+step, length))
				ch <- fmt.Sprintf(command, region, args.Q, args.MaxMeanDepth+1000,
					region, args.Reference, args.Bam)
			}
		}
		close(ch)
	}()
	return ch
}

// Main is run from the dispatcher
func Main() {

	args := dargs{WindowSize: 250,
		MaxMeanDepth: 0,
		MinCov:       4,
		Q:            1}
	p := arg.MustParse(&args)
	if args.Prefix == "" {
		p.Fail("you must specify an output prefix")
	}
	runtime.GOMAXPROCS(args.Processes)
	run(args)
	os.Exit(exitCode)
}

type ipos struct {
	chrom string
	pos   int
	depth int
}

func mean(sl []int) float64 {
	if len(sl) == 0 {
		return 0
	}
	sum := 0
	for _, v := range sl {
		sum += v
	}
	return float64(sum) / float64(len(sl))
}

func getStats(fa *faidx.Faidx, chrom string, start, end int) string {
	if fa == nil {
		return ""
	}
	st, err := fa.Stats(chrom, start, end)
	if err != nil {
		log.Println(err)
	}
	return fmt.Sprintf("\t%.2f\t%.2f\t%.2f", st.GC, st.CpG, st.Masked)
}

func run(args dargs) {

	callback := func(r io.Reader, w io.WriteCloser) error {
		rdr := bufio.NewReader(r)
		wtr := bufio.NewWriter(w)

		var fa *faidx.Faidx
		var err error
		if args.Stats {
			fa, err = faidx.New(args.Reference)
			defer fa.Close()
		}
		if err != nil {
			return err
		}

		defer w.Close()
		defer wtr.Flush()
		cache := make([]ipos, 2)
		depthCache := make([]int, 0, args.WindowSize)
		var depth, pos int
		var lastCovClass, lastChrom string
		lastWindow := -1

		region, err := rdr.ReadBytes('\n')
		if err != nil {
			return err
		}
		chrom, start, regionEnd := chromStartEndFromLine(region)
		lastWindow = start/args.WindowSize - 1

		hdPath := filepath.Join(os.TempDir(), fmt.Sprintf("%s-%s-%s.depth.bed", chrom, start, regionEnd))
		fhHD, ferr := xopen.Wopen(hdPath)
		if ferr != nil {
			return ferr
		}
		caPath := filepath.Join(os.TempDir(), fmt.Sprintf("%s-%s-%s.callable.bed", chrom, start, regionEnd))
		fhCA, ferr := xopen.Wopen(caPath)
		if ferr != nil {
			return ferr
		}
		defer fhCA.Close()
		defer fhHD.Close()

		line, err := rdr.ReadString('\n')
		for err == nil {

			chrom = line[:strings.Index(line, "\t")]

			// toks starts after chrom. so [0] is pos and [1] is depth.
			toks := strings.SplitN(line[len(chrom)+1:], "\t", 2)
			pos, err = strconv.Atoi(toks[0])
			if err != nil {
				break
			}

			if len(toks[1]) > 1 && toks[1][len(toks[1])-1] == '\n' {
				toks[1] = toks[1][:len(toks[1])-1]
			}

			depth, err = strconv.Atoi(toks[1])
			if err != nil {
				break
			}
			// if we have a full window...
			if pos/args.WindowSize != lastWindow {
				thisWindow := pos / args.WindowSize
				// fill in empty windows for which we had no coverage.
				for iwindow := lastWindow + 1; iwindow < thisWindow; iwindow++ {
					s := iwindow * args.WindowSize
					stats := getStats(fa, chrom, s, s+args.WindowSize)
					fhHD.WriteString(fmt.Sprintf("%s\t%d\t%d\t%d%s\n", chrom, s, s+args.WindowSize, 0, stats))
				}
				// fill in this window:
				s := thisWindow * args.WindowSize
				// we dont just break because we still need to print the stuff to the callable file below.
				if s < regionEnd {
					// NOTE: we include the full bin, even for the last postion in the region.w
					// may want to truncate to min(s, s+args.WindowSize)
					stats := getStats(fa, chrom, thisWindow, thisWindow+args.WindowSize)
					fhHD.WriteString(fmt.Sprintf("%s\t%d\t%d\t%.2f%s\n", chrom, s, s+args.WindowSize, mean(depthCache), stats))
					depthCache = depthCache[:0]
					lastWindow = thisWindow
				}
			}
			depthCache = append(depthCache, depth)

			covClass := "CALLABLE"
			if depth == 0 {
				covClass = "NO_COVERAGE"
			} else if depth < args.MinCov {
				covClass = "LOW_COVERAGE"
			} else if args.MaxMeanDepth > 0 && depth >= args.MaxMeanDepth {
				covClass = "EXCESSIVE_COVERAGE"
			}
			// check for a gap or a change in the coverage class.
			if covClass != lastCovClass || pos != cache[1].pos+1 {
				if lastChrom != "" && cache[1].pos != 0 {
					fhCA.WriteString(fmt.Sprintf("%s\t%d\t%d\t%s\n", lastChrom, cache[0].pos-1, cache[1].pos, lastCovClass))
					// also fill in block without any coverage.
					if pos != cache[1].pos+1 {
						fhCA.WriteString(fmt.Sprintf("%s\t%d\t%d\t%s\n", lastChrom, cache[1].pos-1, pos, "NO_COVERAGE"))
					}
				}
				lastCovClass = covClass
				lastChrom = chrom
				cache = cache[:0]
				// append twice so we can safely use cache[1] above.
				cache = append(cache, ipos{chrom, pos, depth})
				cache = append(cache, ipos{chrom, pos, depth})
			} else {
				cache[1].pos = pos
			}
			line, err = rdr.ReadString('\n')
		}
		if len(cache) > 1 {
			fhCA.WriteString(fmt.Sprintf("%s\t%d\t%d\t%s\n", lastChrom, cache[0].pos-1, cache[1].pos, lastCovClass))
		}
		// we didn't get data for the full region, so it must end in no-coverage.
		if cache[1].pos < regionEnd {
			fhCA.WriteString(fmt.Sprintf("%s\t%d\t%d\tNO_COVERAGE\n", lastChrom, cache[1].pos, regionEnd))
			for ds := cache[1].pos / args.WindowSize * args.WindowSize; ds < regionEnd; ds += args.WindowSize {
				thisWindow := ds / args.WindowSize
				stats := getStats(fa, chrom, thisWindow, thisWindow+args.WindowSize)
				fhHD.WriteString(fmt.Sprintf("%s\t%d\t%d\t%.2f%s\n", chrom, ds, ds+args.WindowSize, mean(depthCache), stats))
				depthCache = depthCache[:0]
			}
		}
		wtr.WriteString(caPath + "\n")
		wtr.WriteString(hdPath + "\n")
		wtr.Flush()
		return w.Close()
	}

	cancel := make(chan bool)
	defer close(cancel)
	stdout := bufio.NewWriter(os.Stdout)
	defer stdout.Flush()

	chrom := ""
	if args.Chrom != "" {
		chrom = "." + args.Chrom
	}
	fhca, err := xopen.Wopen(fmt.Sprintf("%s%s.callable.bed", args.Prefix, chrom))
	pcheck(err)
	fhhd, err := xopen.Wopen(fmt.Sprintf("%s%s.depth.bed", args.Prefix, chrom))
	pcheck(err)
	defer fhca.Flush()
	defer fhhd.Flush()

	opts := process.Options{Retries: 1, CallBack: callback, Ordered: args.Ordered}

	for cmd := range process.Runner(genCommands(args), cancel, &opts) {
		if ex := cmd.ExitCode(); ex != 0 && cmd.Err != io.EOF {
			c := color.New(color.BgRed).Add(color.Bold)
			fmt.Fprintf(os.Stderr, "%s\n", c.SprintFunc()(fmt.Sprintf("ERROR with command: %s", cmd)))
			exitCode = max(exitCode, ex)
		}
		if cmd.Err == io.EOF {
			continue
		}
		caPath, err := cmd.ReadString('\n')
		if err != nil {
			log.Println(cmd.CmdStr, err, cmd.Err)
		}
		caSrc, err := xopen.Ropen(strings.TrimSpace(caPath))
		pcheck(err)
		io.Copy(fhca, caSrc)
		os.Remove(strings.TrimSpace(caPath))

		hdPath, err := cmd.ReadString('\n')
		if err != nil {
			log.Println(err)
		}
		hdSrc, err := xopen.Ropen(strings.TrimSpace(hdPath))
		pcheck(err)
		io.Copy(fhhd, hdSrc)
		os.Remove(strings.TrimSpace(hdPath))
	}
	fhca.Flush()
	fhca.Close()
	fhhd.Flush()
	fhhd.Close()

}
<|endoftext|>


<commit_before>package wemdigo

import (
	"log"
	"time"

	"github.com/davecgh/go-spew/spew"
	"github.com/gorilla/websocket"
)

type connection struct {
	ws   *websocket.Conn
	h    MessageHandler
	name string
	send chan *Message // Channel to of messages to write.
	proc chan *Message // Channel to process messages with the handler.
	conf *Config       // Middleware conf so we may access keep alive times.
}

func (c *connection) write(msg *Message) error {
	c.ws.SetWriteDeadline(time.Now().Add(c.conf.WriteWait))
	return c.ws.WriteMessage(msg.Type, msg.Data)
}

// processMessages ingests messages in the proc channel, processes them,
// and sends good messages to the Middle (via the out channel) for
// re-direction.
func (c *connection) processMessages(out chan<- *Message) {
	defer c.ws.Close()
	for msg := range c.proc {
		log.Println("received message on proc chan")
		spew.Dump(msg)
		pmsg, pass, err := c.h(msg)

		// Close the socket in case of an error.
		if err != nil {
			control := &Message{websocket.CloseMessage, nil, c.name}
			c.send <- control
			log.Println("Handler error:", err)
			return
		}

		// Only pass messages the handler deems worthy.
		if !pass {
			continue
		}

		// Ensure the message origin is preserved so that writers of
		// message handlers can ignore this parameter if they choose.
		pmsg.Origin = c.name
		out <- pmsg
	}
}

// Read incoming messages from the websocket connection and
// push them onto an internal channel for further processing.
// This ensures that processing messages does not block reading.
func (c *connection) readMessages() {
	ws := c.ws
	pongHandler := func(string) error {
		// log.Println("Connection", c.name, "received a pong message")
		// Reset deadline.
		ws.SetReadDeadline(time.Now().Add(c.conf.PongWait))
		return nil
	}

	defer func() {
		ws.Close()
		close(c.proc)
	}()

	ws.SetReadDeadline(time.Now().Add(c.conf.PongWait))
	ws.SetPongHandler(pongHandler)
	for {
		mt, raw, err := ws.ReadMessage()
		if err != nil {
			control := &Message{websocket.CloseMessage, nil, c.name}
			c.send <- control
			break
		}

		log.Println("sending message to proc channel")
		c.proc <- &Message{mt, raw, c.name}
	}
}

// writeMessages pumps messages from the Middle hub to the  websocket.
// It also keeps the underlying websocket connection alive by sending pings.
func (c *connection) writeMessages() {
	ws := c.ws
	ticker := time.NewTicker(c.conf.PingPeriod)

	defer func() {
		ticker.Stop()
		ws.Close()
	}()

	for {
		select {
		case msg, ok := <-c.send:
			if !ok {
				ws.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			if err := c.write(msg); err != nil {
				return
			}

		case <-ticker.C:
			// FIXME
			log.Println(c.name, "sending a ping.")
			control := &Message{websocket.PingMessage, nil, ""}
			if err := c.write(control); err != nil {
				return
			}
		}
	}
}

func newConnection(ws *websocket.Conn, h MessageHandler, name string, conf *Config) *connection {
	return &connection{
		ws:   ws,
		h:    h,
		name: name,
		send: make(chan *Message),
		proc: make(chan *Message),
		conf: conf,
	}
}
<commit_msg>Remove some comments.<commit_after>package wemdigo

import (
	"log"
	"time"

	"github.com/davecgh/go-spew/spew"
	"github.com/gorilla/websocket"
)

type connection struct {
	ws   *websocket.Conn
	h    MessageHandler
	name string
	send chan *Message // Channel to of messages to write.
	proc chan *Message // Channel to process messages with the handler.
	conf *Config       // Middleware conf so we may access keep alive times.
}

func (c *connection) write(msg *Message) error {
	c.ws.SetWriteDeadline(time.Now().Add(c.conf.WriteWait))
	return c.ws.WriteMessage(msg.Type, msg.Data)
}

// processMessages ingests messages in the proc channel, processes them,
// and sends good messages to the Middle (via the out channel) for
// re-direction.
func (c *connection) processMessages(out chan<- *Message) {
	defer c.ws.Close()
	for msg := range c.proc {
		log.Println("received message on proc chan")
		spew.Dump(msg)
		pmsg, pass, err := c.h(msg)

		// Close the socket in case of an error.
		if err != nil {
			control := &Message{websocket.CloseMessage, nil, c.name}
			c.send <- control
			log.Println("Handler error:", err)
			return
		}

		// Only pass messages the handler deems worthy.
		if !pass {
			continue
		}

		// Ensure the message origin is preserved so that writers of
		// message handlers can ignore this parameter if they choose.
		pmsg.Origin = c.name
		out <- pmsg
	}
}

// Read incoming messages from the websocket connection and
// push them onto an internal channel for further processing.
// This ensures that processing messages does not block reading.
func (c *connection) readMessages() {
	ws := c.ws
	pongHandler := func(string) error {
		// log.Println("Connection", c.name, "received a pong message")
		// Reset deadline.
		ws.SetReadDeadline(time.Now().Add(c.conf.PongWait))
		return nil
	}

	defer func() {
		ws.Close()
		close(c.proc)
	}()

	ws.SetReadDeadline(time.Now().Add(c.conf.PongWait))
	ws.SetPongHandler(pongHandler)
	for {
		mt, raw, err := ws.ReadMessage()
		if err != nil {
			control := &Message{websocket.CloseMessage, nil, c.name}
			c.send <- control
			break
		}

		c.proc <- &Message{mt, raw, c.name}
	}
}

// writeMessages pumps messages from the Middle hub to the  websocket.
// It also keeps the underlying websocket connection alive by sending pings.
func (c *connection) writeMessages() {
	ws := c.ws
	ticker := time.NewTicker(c.conf.PingPeriod)

	defer func() {
		ticker.Stop()
		ws.Close()
	}()

	for {
		select {
		case msg, ok := <-c.send:
			if !ok {
				ws.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			if err := c.write(msg); err != nil {
				return
			}
		case <-ticker.C:
			// log.Println(c.name, "sending a ping.")
			control := &Message{websocket.PingMessage, nil, ""}
			if err := c.write(control); err != nil {
				return
			}
		}
	}
}

func newConnection(ws *websocket.Conn, h MessageHandler, name string, conf *Config) *connection {
	return &connection{
		ws:   ws,
		h:    h,
		name: name,
		send: make(chan *Message),
		proc: make(chan *Message),
		conf: conf,
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2017-2019 Snowflake Computing Inc. All right reserved.

package gosnowflake

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"github.com/google/uuid"
	"net/url"
	"strconv"
	"strings"
	"sync/atomic"
)

const (
	statementTypeIDDml              = int64(0x3000)
	statementTypeIDInsert           = statementTypeIDDml + int64(0x100)
	statementTypeIDUpdate           = statementTypeIDDml + int64(0x200)
	statementTypeIDDelete           = statementTypeIDDml + int64(0x300)
	statementTypeIDMerge            = statementTypeIDDml + int64(0x400)
	statementTypeIDMultiTableInsert = statementTypeIDDml + int64(0x500)
)

const (
	sessionClientSessionKeepAlive          = "client_session_keep_alive"
	sessionClientValidateDefaultParameters = "CLIENT_VALIDATE_DEFAULT_PARAMETERS"
	serviceName                            = "service_name"
)

type snowflakeConn struct {
	cfg             *Config
	rest            *snowflakeRestful
	SequenceCounter uint64
	QueryID         string
	SQLState        string
}

// isDml returns true if the statement type code is in the range of DML.
func (sc *snowflakeConn) isDml(v int64) bool {
	switch v {
	case statementTypeIDDml, statementTypeIDInsert,
		statementTypeIDUpdate, statementTypeIDDelete,
		statementTypeIDMerge, statementTypeIDMultiTableInsert:
		return true
	}
	return false
}

func (sc *snowflakeConn) exec(
	ctx context.Context,
	query string, noResult bool, isInternal bool, parameters []driver.NamedValue) (*execResponse, error) {
	var err error
	counter := atomic.AddUint64(&sc.SequenceCounter, 1) // query sequence counter

	req := execRequest{
		SQLText:    query,
		AsyncExec:  noResult,
		SequenceID: counter,
	}
	req.IsInternal = isInternal
	tsmode := "TIMESTAMP_NTZ"
	idx := 1
	if len(parameters) > 0 {
		req.Bindings = make(map[string]execBindParameter, len(parameters))
		for i, n := 0, len(parameters); i < n; i++ {
			t := goTypeToSnowflake(parameters[i].Value, tsmode)
			glog.V(2).Infof("tmode: %v\n", t)
			if t == "CHANGE_TYPE" {
				tsmode, err = dataTypeMode(parameters[i].Value)
				if err != nil {
					return nil, err
				}
			} else {
				v1, err := valueToString(parameters[i].Value, tsmode)
				if err != nil {
					return nil, err
				}
				req.Bindings[strconv.Itoa(idx)] = execBindParameter{
					Type:  t,
					Value: v1,
				}
				idx++
			}
		}
	}
	glog.V(2).Infof("bindings: %v", req.Bindings)

	headers := make(map[string]string)
	headers["Content-Type"] = headerContentTypeApplicationJSON
	headers["accept"] = headerAcceptTypeApplicationSnowflake // TODO v1.1: change to JSON in case of PUT/GET
	headers["User-Agent"] = userAgent
	if serviceName, ok := sc.cfg.Params[serviceName]; ok {
		headers["X-Snowflake-Service"] = *serviceName
	}

	jsonBody, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	var data *execResponse

	requestID := uuid.New()
	data, err = sc.rest.FuncPostQuery(ctx, sc.rest, &url.Values{}, headers, jsonBody, sc.rest.RequestTimeout, &requestID)
	if err != nil {
		return nil, err
	}
	var code int
	if data.Code != "" {
		code, err = strconv.Atoi(data.Code)
		if err != nil {
			code = -1
			return nil, err
		}
	} else {
		code = -1
	}
	glog.V(2).Infof("Success: %v, Code: %v", data.Success, code)
	if !data.Success {
		return nil, &SnowflakeError{
			Number:   code,
			SQLState: data.Data.SQLState,
			Message:  data.Message,
			QueryID:  data.Data.QueryID,
		}
	}
	glog.V(2).Info("Exec/Query SUCCESS")
	sc.cfg.Database = data.Data.FinalDatabaseName
	sc.cfg.Schema = data.Data.FinalSchemaName
	sc.cfg.Role = data.Data.FinalRoleName
	sc.cfg.Warehouse = data.Data.FinalWarehouseName
	sc.QueryID = data.Data.QueryID
	sc.SQLState = data.Data.SQLState
	sc.populateSessionParameters(data.Data.Parameters)
	return data, err
}

func (sc *snowflakeConn) Begin() (driver.Tx, error) {
	return sc.BeginTx(context.TODO(), driver.TxOptions{})
}

func (sc *snowflakeConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
	glog.V(2).Info("BeginTx")
	if opts.ReadOnly {
		return nil, &SnowflakeError{
			Number:   ErrNoReadOnlyTransaction,
			SQLState: SQLStateFeatureNotSupported,
			Message:  errMsgNoReadOnlyTransaction,
		}
	}
	if int(opts.Isolation) != int(sql.LevelDefault) {
		return nil, &SnowflakeError{
			Number:   ErrNoDefaultTransactionIsolationLevel,
			SQLState: SQLStateFeatureNotSupported,
			Message:  errMsgNoDefaultTransactionIsolationLevel,
		}
	}
	if sc.rest == nil {
		return nil, driver.ErrBadConn
	}
	_, err := sc.exec(ctx, "BEGIN", false, false, nil)
	if err != nil {
		return nil, err
	}
	return &snowflakeTx{sc}, err
}

func (sc *snowflakeConn) cleanup() {
	glog.Flush() // must flush log buffer while the process is running.
	sc.rest = nil
	sc.cfg = nil
}

func (sc *snowflakeConn) Close() (err error) {
	glog.V(2).Infoln("Close")
	sc.stopHeartBeat()

	err = sc.rest.FuncCloseSession(sc.rest)
	if err != nil {
		glog.V(2).Info(err)
	}
	sc.cleanup()
	return nil
}
func (sc *snowflakeConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	glog.V(2).Infoln("Prepare")
	if sc.rest == nil {
		return nil, driver.ErrBadConn
	}
	stmt := &snowflakeStmt{
		sc:    sc,
		query: query,
	}
	return stmt, nil
}

func (sc *snowflakeConn) Prepare(query string) (driver.Stmt, error) {
	return sc.PrepareContext(context.TODO(), query)
}

func (sc *snowflakeConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
	glog.V(2).Infof("Exec: %#v, %v", query, args)
	if sc.rest == nil {
		return nil, driver.ErrBadConn
	}
	// TODO: handle noResult and isInternal
	data, err := sc.exec(ctx, query, false, false, args)
	if err != nil {
		return nil, err
	}
	var updatedRows int64
	if sc.isDml(data.Data.StatementTypeID) {
		// collects all values from the returned row sets
		updatedRows = 0
		for i, n := 0, len(data.Data.RowType); i < n; i++ {
			v, err := strconv.ParseInt(*data.Data.RowSet[0][i], 10, 64)
			if err != nil {
				return nil, err
			}
			updatedRows += v
		}
		glog.V(2).Infof("number of updated rows: %#v", updatedRows)
		return &snowflakeResult{
			affectedRows: updatedRows,
			insertID:     -1}, nil // last insert id is not supported by Snowflake
	}
	glog.V(2).Info("DDL")
	return driver.ResultNoRows, nil
}

func (sc *snowflakeConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	glog.V(2).Infof("Query: %#v, %v", query, args)
	if sc.rest == nil {
		return nil, driver.ErrBadConn
	}
	// TODO: handle noResult and isInternal
	data, err := sc.exec(ctx, query, false, false, args)
	if err != nil {
		glog.V(2).Infof("error: %v", err)
		return nil, err
	}

	rows := new(snowflakeRows)
	rows.sc = sc
	rows.RowType = data.Data.RowType
	rows.ChunkDownloader = &snowflakeChunkDownloader{
		sc:                 sc,
		ctx:                ctx,
		CurrentChunk:       make([]chunkRowType, len(data.Data.RowSet)),
		ChunkMetas:         data.Data.Chunks,
		Total:              data.Data.Total,
		TotalRowIndex:      int64(-1),
		CellCount:          len(data.Data.RowType),
		Qrmk:               data.Data.Qrmk,
		QueryResultFormat:  data.Data.QueryResultFormat,
		ChunkHeader:        data.Data.ChunkHeaders,
		FuncDownload:       downloadChunk,
		FuncDownloadHelper: downloadChunkHelper,
		FuncGet:            getChunk,
		RowSet: rowSetType{RowType: data.Data.RowType,
			JSON:         data.Data.RowSet,
			RowSetBase64: data.Data.RowSetBase64,
		},
	}
	//populateJSONRowSet(rows.ChunkDownloader.CurrentChunk, data.Data.RowSet)
	rows.ChunkDownloader.start()
	return rows, err
}

func (sc *snowflakeConn) Exec(
	query string,
	args []driver.Value) (
	driver.Result, error) {
	return sc.ExecContext(context.TODO(), query, toNamedValues(args))
}

func (sc *snowflakeConn) Query(
	query string,
	args []driver.Value) (
	driver.Rows, error) {
	return sc.QueryContext(context.TODO(), query, toNamedValues(args))
}

func (sc *snowflakeConn) Ping(ctx context.Context) error {
	glog.V(2).Infoln("Ping")
	if sc.rest == nil {
		return driver.ErrBadConn
	}
	// TODO: handle noResult and isInternal
	_, err := sc.exec(ctx, "SELECT 1", false, false, []driver.NamedValue{})
	return err
}

func (sc *snowflakeConn) populateSessionParameters(parameters []nameValueParameter) {
	// other session parameters (not all)
	glog.V(2).Infof("params: %#v", parameters)
	for _, param := range parameters {
		v := ""
		switch param.Value.(type) {
		case int64:
			if vv, ok := param.Value.(int64); ok {
				v = strconv.FormatInt(vv, 10)
			}
		case float64:
			if vv, ok := param.Value.(float64); ok {
				v = strconv.FormatFloat(vv, 'g', -1, 64)
			}
		case bool:
			if vv, ok := param.Value.(bool); ok {
				v = strconv.FormatBool(vv)
			}
		default:
			if vv, ok := param.Value.(string); ok {
				v = vv
			}
		}
		glog.V(3).Infof("parameter. name: %v, value: %v", param.Name, v)
		sc.cfg.Params[strings.ToLower(param.Name)] = &v
	}
}

func (sc *snowflakeConn) isClientSessionKeepAliveEnabled() bool {
	v, ok := sc.cfg.Params[sessionClientSessionKeepAlive]
	if !ok {
		return false
	}
	return strings.Compare(*v, "true") == 0
}

func (sc *snowflakeConn) startHeartBeat() {
	if !sc.isClientSessionKeepAliveEnabled() {
		return
	}
	sc.rest.HeartBeat = &heartbeat{
		restful: sc.rest,
	}
	sc.rest.HeartBeat.start()
}

func (sc *snowflakeConn) stopHeartBeat() {
	if !sc.isClientSessionKeepAliveEnabled() {
		return
	}
	sc.rest.HeartBeat.stop()
}
<commit_msg>Remove redundant chunkDownloader rowset population<commit_after>// Copyright (c) 2017-2019 Snowflake Computing Inc. All right reserved.

package gosnowflake

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"github.com/google/uuid"
	"net/url"
	"strconv"
	"strings"
	"sync/atomic"
)

const (
	statementTypeIDDml              = int64(0x3000)
	statementTypeIDInsert           = statementTypeIDDml + int64(0x100)
	statementTypeIDUpdate           = statementTypeIDDml + int64(0x200)
	statementTypeIDDelete           = statementTypeIDDml + int64(0x300)
	statementTypeIDMerge            = statementTypeIDDml + int64(0x400)
	statementTypeIDMultiTableInsert = statementTypeIDDml + int64(0x500)
)

const (
	sessionClientSessionKeepAlive          = "client_session_keep_alive"
	sessionClientValidateDefaultParameters = "CLIENT_VALIDATE_DEFAULT_PARAMETERS"
	serviceName                            = "service_name"
)

type snowflakeConn struct {
	cfg             *Config
	rest            *snowflakeRestful
	SequenceCounter uint64
	QueryID         string
	SQLState        string
}

// isDml returns true if the statement type code is in the range of DML.
func (sc *snowflakeConn) isDml(v int64) bool {
	switch v {
	case statementTypeIDDml, statementTypeIDInsert,
		statementTypeIDUpdate, statementTypeIDDelete,
		statementTypeIDMerge, statementTypeIDMultiTableInsert:
		return true
	}
	return false
}

func (sc *snowflakeConn) exec(
	ctx context.Context,
	query string, noResult bool, isInternal bool, parameters []driver.NamedValue) (*execResponse, error) {
	var err error
	counter := atomic.AddUint64(&sc.SequenceCounter, 1) // query sequence counter

	req := execRequest{
		SQLText:    query,
		AsyncExec:  noResult,
		SequenceID: counter,
	}
	req.IsInternal = isInternal
	tsmode := "TIMESTAMP_NTZ"
	idx := 1
	if len(parameters) > 0 {
		req.Bindings = make(map[string]execBindParameter, len(parameters))
		for i, n := 0, len(parameters); i < n; i++ {
			t := goTypeToSnowflake(parameters[i].Value, tsmode)
			glog.V(2).Infof("tmode: %v\n", t)
			if t == "CHANGE_TYPE" {
				tsmode, err = dataTypeMode(parameters[i].Value)
				if err != nil {
					return nil, err
				}
			} else {
				v1, err := valueToString(parameters[i].Value, tsmode)
				if err != nil {
					return nil, err
				}
				req.Bindings[strconv.Itoa(idx)] = execBindParameter{
					Type:  t,
					Value: v1,
				}
				idx++
			}
		}
	}
	glog.V(2).Infof("bindings: %v", req.Bindings)

	headers := make(map[string]string)
	headers["Content-Type"] = headerContentTypeApplicationJSON
	headers["accept"] = headerAcceptTypeApplicationSnowflake // TODO v1.1: change to JSON in case of PUT/GET
	headers["User-Agent"] = userAgent
	if serviceName, ok := sc.cfg.Params[serviceName]; ok {
		headers["X-Snowflake-Service"] = *serviceName
	}

	jsonBody, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	var data *execResponse

	requestID := uuid.New()
	data, err = sc.rest.FuncPostQuery(ctx, sc.rest, &url.Values{}, headers, jsonBody, sc.rest.RequestTimeout, &requestID)
	if err != nil {
		return nil, err
	}
	var code int
	if data.Code != "" {
		code, err = strconv.Atoi(data.Code)
		if err != nil {
			code = -1
			return nil, err
		}
	} else {
		code = -1
	}
	glog.V(2).Infof("Success: %v, Code: %v", data.Success, code)
	if !data.Success {
		return nil, &SnowflakeError{
			Number:   code,
			SQLState: data.Data.SQLState,
			Message:  data.Message,
			QueryID:  data.Data.QueryID,
		}
	}
	glog.V(2).Info("Exec/Query SUCCESS")
	sc.cfg.Database = data.Data.FinalDatabaseName
	sc.cfg.Schema = data.Data.FinalSchemaName
	sc.cfg.Role = data.Data.FinalRoleName
	sc.cfg.Warehouse = data.Data.FinalWarehouseName
	sc.QueryID = data.Data.QueryID
	sc.SQLState = data.Data.SQLState
	sc.populateSessionParameters(data.Data.Parameters)
	return data, err
}

func (sc *snowflakeConn) Begin() (driver.Tx, error) {
	return sc.BeginTx(context.TODO(), driver.TxOptions{})
}

func (sc *snowflakeConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
	glog.V(2).Info("BeginTx")
	if opts.ReadOnly {
		return nil, &SnowflakeError{
			Number:   ErrNoReadOnlyTransaction,
			SQLState: SQLStateFeatureNotSupported,
			Message:  errMsgNoReadOnlyTransaction,
		}
	}
	if int(opts.Isolation) != int(sql.LevelDefault) {
		return nil, &SnowflakeError{
			Number:   ErrNoDefaultTransactionIsolationLevel,
			SQLState: SQLStateFeatureNotSupported,
			Message:  errMsgNoDefaultTransactionIsolationLevel,
		}
	}
	if sc.rest == nil {
		return nil, driver.ErrBadConn
	}
	_, err := sc.exec(ctx, "BEGIN", false, false, nil)
	if err != nil {
		return nil, err
	}
	return &snowflakeTx{sc}, err
}

func (sc *snowflakeConn) cleanup() {
	glog.Flush() // must flush log buffer while the process is running.
	sc.rest = nil
	sc.cfg = nil
}

func (sc *snowflakeConn) Close() (err error) {
	glog.V(2).Infoln("Close")
	sc.stopHeartBeat()

	err = sc.rest.FuncCloseSession(sc.rest)
	if err != nil {
		glog.V(2).Info(err)
	}
	sc.cleanup()
	return nil
}
func (sc *snowflakeConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	glog.V(2).Infoln("Prepare")
	if sc.rest == nil {
		return nil, driver.ErrBadConn
	}
	stmt := &snowflakeStmt{
		sc:    sc,
		query: query,
	}
	return stmt, nil
}

func (sc *snowflakeConn) Prepare(query string) (driver.Stmt, error) {
	return sc.PrepareContext(context.TODO(), query)
}

func (sc *snowflakeConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
	glog.V(2).Infof("Exec: %#v, %v", query, args)
	if sc.rest == nil {
		return nil, driver.ErrBadConn
	}
	// TODO: handle noResult and isInternal
	data, err := sc.exec(ctx, query, false, false, args)
	if err != nil {
		return nil, err
	}
	var updatedRows int64
	if sc.isDml(data.Data.StatementTypeID) {
		// collects all values from the returned row sets
		updatedRows = 0
		for i, n := 0, len(data.Data.RowType); i < n; i++ {
			v, err := strconv.ParseInt(*data.Data.RowSet[0][i], 10, 64)
			if err != nil {
				return nil, err
			}
			updatedRows += v
		}
		glog.V(2).Infof("number of updated rows: %#v", updatedRows)
		return &snowflakeResult{
			affectedRows: updatedRows,
			insertID:     -1}, nil // last insert id is not supported by Snowflake
	}
	glog.V(2).Info("DDL")
	return driver.ResultNoRows, nil
}

func (sc *snowflakeConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	glog.V(2).Infof("Query: %#v, %v", query, args)
	if sc.rest == nil {
		return nil, driver.ErrBadConn
	}
	// TODO: handle noResult and isInternal
	data, err := sc.exec(ctx, query, false, false, args)
	if err != nil {
		glog.V(2).Infof("error: %v", err)
		return nil, err
	}

	rows := new(snowflakeRows)
	rows.sc = sc
	rows.RowType = data.Data.RowType
	rows.ChunkDownloader = &snowflakeChunkDownloader{
		sc:                 sc,
		ctx:                ctx,
		CurrentChunk:       make([]chunkRowType, len(data.Data.RowSet)),
		ChunkMetas:         data.Data.Chunks,
		Total:              data.Data.Total,
		TotalRowIndex:      int64(-1),
		CellCount:          len(data.Data.RowType),
		Qrmk:               data.Data.Qrmk,
		QueryResultFormat:  data.Data.QueryResultFormat,
		ChunkHeader:        data.Data.ChunkHeaders,
		FuncDownload:       downloadChunk,
		FuncDownloadHelper: downloadChunkHelper,
		FuncGet:            getChunk,
		RowSet: rowSetType{RowType: data.Data.RowType,
			JSON:         data.Data.RowSet,
			RowSetBase64: data.Data.RowSetBase64,
		},
	}
	rows.ChunkDownloader.start()
	return rows, err
}

func (sc *snowflakeConn) Exec(
	query string,
	args []driver.Value) (
	driver.Result, error) {
	return sc.ExecContext(context.TODO(), query, toNamedValues(args))
}

func (sc *snowflakeConn) Query(
	query string,
	args []driver.Value) (
	driver.Rows, error) {
	return sc.QueryContext(context.TODO(), query, toNamedValues(args))
}

func (sc *snowflakeConn) Ping(ctx context.Context) error {
	glog.V(2).Infoln("Ping")
	if sc.rest == nil {
		return driver.ErrBadConn
	}
	// TODO: handle noResult and isInternal
	_, err := sc.exec(ctx, "SELECT 1", false, false, []driver.NamedValue{})
	return err
}

func (sc *snowflakeConn) populateSessionParameters(parameters []nameValueParameter) {
	// other session parameters (not all)
	glog.V(2).Infof("params: %#v", parameters)
	for _, param := range parameters {
		v := ""
		switch param.Value.(type) {
		case int64:
			if vv, ok := param.Value.(int64); ok {
				v = strconv.FormatInt(vv, 10)
			}
		case float64:
			if vv, ok := param.Value.(float64); ok {
				v = strconv.FormatFloat(vv, 'g', -1, 64)
			}
		case bool:
			if vv, ok := param.Value.(bool); ok {
				v = strconv.FormatBool(vv)
			}
		default:
			if vv, ok := param.Value.(string); ok {
				v = vv
			}
		}
		glog.V(3).Infof("parameter. name: %v, value: %v", param.Name, v)
		sc.cfg.Params[strings.ToLower(param.Name)] = &v
	}
}

func (sc *snowflakeConn) isClientSessionKeepAliveEnabled() bool {
	v, ok := sc.cfg.Params[sessionClientSessionKeepAlive]
	if !ok {
		return false
	}
	return strings.Compare(*v, "true") == 0
}

func (sc *snowflakeConn) startHeartBeat() {
	if !sc.isClientSessionKeepAliveEnabled() {
		return
	}
	sc.rest.HeartBeat = &heartbeat{
		restful: sc.rest,
	}
	sc.rest.HeartBeat.start()
}

func (sc *snowflakeConn) stopHeartBeat() {
	if !sc.isClientSessionKeepAliveEnabled() {
		return
	}
	sc.rest.HeartBeat.stop()
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 fuse

import (
	"fmt"
	"log"
	"path"
	"runtime"
	"sync"

	"golang.org/x/net/context"

	"github.com/jacobsa/bazilfuse"
	"github.com/jacobsa/fuse/fuseops"
)

// A connection to the fuse kernel process.
type Connection struct {
	logger      *log.Logger
	wrapped     *bazilfuse.Conn
	opsInFlight sync.WaitGroup

	// The context from which all op contexts inherit.
	parentCtx context.Context

	// For logging purposes only.
	nextOpID uint32
}

// Responsibility for closing the wrapped connection is transferred to the
// result. You must call c.close() eventually.
func newConnection(
	parentCtx context.Context,
	logger *log.Logger,
	wrapped *bazilfuse.Conn) (c *Connection, err error) {
	c = &Connection{
		logger:    logger,
		wrapped:   wrapped,
		parentCtx: parentCtx,
	}

	return
}

// Log information for an operation with the given ID. calldepth is the depth
// to use when recovering file:line information with runtime.Caller.
func (c *Connection) log(
	opID uint32,
	calldepth int,
	format string,
	v ...interface{}) {
	// Get file:line info.
	var file string
	var line int
	var ok bool

	_, file, line, ok = runtime.Caller(calldepth)
	if !ok {
		file = "???"
	}

	fileLine := fmt.Sprintf("%v:%v", path.Base(file), line)

	// Format the actual message to be printed.
	msg := fmt.Sprintf(
		"Op 0x%08x %24s] %v",
		opID,
		fileLine,
		fmt.Sprintf(format, v...))

	// Print it.
	c.logger.Println(msg)
}

// Set up state for an op that is about to be returned to the user, given its
// bazilfuse request ID.
//
// Return a context that should be used for the op.
func (c *Connection) beginOp(reqID bazilfuse.RequestID) (ctx context.Context) {
	c.opsInFlight.Add(1)

	// TODO(jacobsa): Use WithCancel and stash a cancellation function.
	ctx = c.parentCtx

	return
}

// Clean up all state associated with an op to which the user has responded.
func (c *Connection) finishOp(reqID bazilfuse.RequestID) {
	c.opsInFlight.Done()
}

// Read the next op from the kernel process. Return io.EOF if the kernel has
// closed the connection.
//
// This function delivers ops in exactly the order they are received from
// /dev/fuse. It must not be called multiple times concurrently.
func (c *Connection) ReadOp() (op fuseops.Op, err error) {
	var bfReq bazilfuse.Request

	// Keep going until we find a request we know how to convert.
	for {
		// Read a bazilfuse request.
		bfReq, err = c.wrapped.ReadRequest()
		if err != nil {
			return
		}

		// Choose an ID for this operation.
		opID := c.nextOpID
		c.nextOpID++

		// Log the receipt of the operation.
		c.log(opID, 1, "<- %v", bfReq)

		// Special case: responding to this is required to make mounting work on OS
		// X. We don't currently expose the capability for the file system to
		// intercept this.
		if statfsReq, ok := bfReq.(*bazilfuse.StatfsRequest); ok {
			c.log(opID, 1, "-> (Statfs) OK")
			statfsReq.Respond(&bazilfuse.StatfsResponse{})
			continue
		}

		// Convert it.
		logForOp := func(calldepth int, format string, v ...interface{}) {
			c.log(opID, calldepth+1, format, v...)
		}

		finished := func(err error) { c.finishOp() }

		op = fuseops.Convert(c.parentCtx, bfReq, logForOp, finished)
		c.beginOp()
		return
	}
}

func (c *Connection) waitForReady() (err error) {
	<-c.wrapped.Ready
	err = c.wrapped.MountError
	return
}

// Close the connection and wait for in-flight ops.
func (c *Connection) close() (err error) {
	err = c.wrapped.Close()
	c.opsInFlight.Wait()
	return
}
<commit_msg>Use new beginOp/finishOp.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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 fuse

import (
	"fmt"
	"log"
	"path"
	"runtime"
	"sync"

	"golang.org/x/net/context"

	"github.com/jacobsa/bazilfuse"
	"github.com/jacobsa/fuse/fuseops"
)

// A connection to the fuse kernel process.
type Connection struct {
	logger      *log.Logger
	wrapped     *bazilfuse.Conn
	opsInFlight sync.WaitGroup

	// The context from which all op contexts inherit.
	parentCtx context.Context

	// For logging purposes only.
	nextOpID uint32
}

// Responsibility for closing the wrapped connection is transferred to the
// result. You must call c.close() eventually.
func newConnection(
	parentCtx context.Context,
	logger *log.Logger,
	wrapped *bazilfuse.Conn) (c *Connection, err error) {
	c = &Connection{
		logger:    logger,
		wrapped:   wrapped,
		parentCtx: parentCtx,
	}

	return
}

// Log information for an operation with the given ID. calldepth is the depth
// to use when recovering file:line information with runtime.Caller.
func (c *Connection) log(
	opID uint32,
	calldepth int,
	format string,
	v ...interface{}) {
	// Get file:line info.
	var file string
	var line int
	var ok bool

	_, file, line, ok = runtime.Caller(calldepth)
	if !ok {
		file = "???"
	}

	fileLine := fmt.Sprintf("%v:%v", path.Base(file), line)

	// Format the actual message to be printed.
	msg := fmt.Sprintf(
		"Op 0x%08x %24s] %v",
		opID,
		fileLine,
		fmt.Sprintf(format, v...))

	// Print it.
	c.logger.Println(msg)
}

// Set up state for an op that is about to be returned to the user, given its
// bazilfuse request ID.
//
// Return a context that should be used for the op.
func (c *Connection) beginOp(reqID bazilfuse.RequestID) (ctx context.Context) {
	c.opsInFlight.Add(1)

	// TODO(jacobsa): Use WithCancel and stash a cancellation function.
	ctx = c.parentCtx

	return
}

// Clean up all state associated with an op to which the user has responded.
func (c *Connection) finishOp(reqID bazilfuse.RequestID) {
	c.opsInFlight.Done()
}

// Read the next op from the kernel process. Return io.EOF if the kernel has
// closed the connection.
//
// This function delivers ops in exactly the order they are received from
// /dev/fuse. It must not be called multiple times concurrently.
func (c *Connection) ReadOp() (op fuseops.Op, err error) {
	// Keep going until we find a request we know how to convert.
	for {
		// Read a bazilfuse request.
		var bfReq bazilfuse.Request
		bfReq, err = c.wrapped.ReadRequest()

		if err != nil {
			return
		}

		// Choose an ID for this operation.
		opID := c.nextOpID
		c.nextOpID++

		// Log the receipt of the operation.
		c.log(opID, 1, "<- %v", bfReq)

		// Special case: responding to this is required to make mounting work on OS
		// X. We don't currently expose the capability for the file system to
		// intercept this.
		if statfsReq, ok := bfReq.(*bazilfuse.StatfsRequest); ok {
			c.log(opID, 1, "-> (Statfs) OK")
			statfsReq.Respond(&bazilfuse.StatfsResponse{})
			continue
		}

		// Set up op dependencies.
		var reqID bazilfuse.RequestID = bfReq.Hdr().ID
		opCtx := c.beginOp(reqID)

		logForOp := func(calldepth int, format string, v ...interface{}) {
			c.log(opID, calldepth+1, format, v...)
		}

		finished := func(err error) { c.finishOp(reqID) }

		op = fuseops.Convert(opCtx, bfReq, logForOp, finished)
		return
	}
}

func (c *Connection) waitForReady() (err error) {
	<-c.wrapped.Ready
	err = c.wrapped.MountError
	return
}

// Close the connection and wait for in-flight ops.
func (c *Connection) close() (err error) {
	err = c.wrapped.Close()
	c.opsInFlight.Wait()
	return
}
<|endoftext|>