<commit_before>package filesys

import (
	"context"
	"io"
	"os"
	"sort"
	"time"

	"github.com/seaweedfs/fuse"
	"github.com/seaweedfs/fuse/fs"

	"github.com/chrislusf/seaweedfs/weed/filer2"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
)

const blockSize = 512

var _ = fs.Node(&File{})
var _ = fs.NodeOpener(&File{})
var _ = fs.NodeFsyncer(&File{})
var _ = fs.NodeSetattrer(&File{})
var _ = fs.NodeGetxattrer(&File{})
var _ = fs.NodeSetxattrer(&File{})
var _ = fs.NodeRemovexattrer(&File{})
var _ = fs.NodeListxattrer(&File{})
var _ = fs.NodeForgetter(&File{})

type File struct {
	Name           string
	dir            *Dir
	wfs            *WFS
	entry          *filer_pb.Entry
	entryViewCache []filer2.VisibleInterval
	isOpen         int
	reader         io.ReaderAt
	dirtyMetadata  bool
}

func (file *File) fullpath() util.FullPath {
	return util.NewFullPath(file.dir.FullPath(), file.Name)
}

func (file *File) Attr(ctx context.Context, attr *fuse.Attr) error {

	glog.V(5).Infof("file Attr %s, open:%v, existing attr: %+v", file.fullpath(), file.isOpen, attr)

	if file.isOpen <= 0 {
		if err := file.maybeLoadEntry(ctx); err != nil {
			return err
		}
	}

	attr.Inode = file.fullpath().AsInode()
	attr.Valid = time.Second
	attr.Mode = os.FileMode(file.entry.Attributes.FileMode)
	attr.Size = filer2.FileSize(file.entry)
	if file.isOpen > 0 {
		attr.Size = file.entry.Attributes.FileSize
		glog.V(4).Infof("file Attr %s, open:%v, size: %d", file.fullpath(), file.isOpen, attr.Size)
	}
	attr.Crtime = time.Unix(file.entry.Attributes.Crtime, 0)
	attr.Mtime = time.Unix(file.entry.Attributes.Mtime, 0)
	attr.Gid = file.entry.Attributes.Gid
	attr.Uid = file.entry.Attributes.Uid
	attr.Blocks = attr.Size/blockSize + 1
	attr.BlockSize = uint32(file.wfs.option.ChunkSizeLimit)

	return nil

}

func (file *File) Getxattr(ctx context.Context, req *fuse.GetxattrRequest, resp *fuse.GetxattrResponse) error {

	glog.V(4).Infof("file Getxattr %s", file.fullpath())

	if err := file.maybeLoadEntry(ctx); err != nil {
		return err
	}

	return getxattr(file.entry, req, resp)
}

func (file *File) Open(ctx context.Context, req *fuse.OpenRequest, resp *fuse.OpenResponse) (fs.Handle, error) {

	glog.V(4).Infof("file %v open %+v", file.fullpath(), req)

	file.isOpen++

	handle := file.wfs.AcquireHandle(file, req.Uid, req.Gid)

	resp.Handle = fuse.HandleID(handle.handle)

	glog.V(4).Infof("%v file open handle id = %d", file.fullpath(), handle.handle)

	return handle, nil

}

func (file *File) Setattr(ctx context.Context, req *fuse.SetattrRequest, resp *fuse.SetattrResponse) error {

	glog.V(5).Infof("%v file setattr %+v", file.fullpath(), req)

	if err := file.maybeLoadEntry(ctx); err != nil {
		return err
	}

	if req.Valid.Size() {

		glog.V(4).Infof("%v file setattr set size=%v chunks=%d", file.fullpath(), req.Size, len(file.entry.Chunks))
		if req.Size < filer2.TotalSize(file.entry.Chunks) {
			// fmt.Printf("truncate %v \n", fullPath)
			var chunks []*filer_pb.FileChunk
			var truncatedChunks []*filer_pb.FileChunk
			for _, chunk := range file.entry.Chunks {
				int64Size := int64(chunk.Size)
				if chunk.Offset+int64Size > int64(req.Size) {
					// this chunk is truncated
					int64Size = int64(req.Size) - chunk.Offset
					if int64Size > 0 {
						chunks = append(chunks, chunk)
						glog.V(4).Infof("truncated chunk %+v from %d to %d\n", chunk.GetFileIdString(), chunk.Size, int64Size)
						chunk.Size = uint64(int64Size)
					} else {
						glog.V(4).Infof("truncated whole chunk %+v\n", chunk.GetFileIdString())
						truncatedChunks = append(truncatedChunks, chunk)
					}
				}
			}
			file.wfs.deleteFileChunks(truncatedChunks)
			file.entry.Chunks = chunks
			file.entryViewCache = nil
			file.reader = nil
		}
		file.entry.Attributes.FileSize = req.Size
		file.dirtyMetadata = true
	}

	if req.Valid.Mode() {
		file.entry.Attributes.FileMode = uint32(req.Mode)
		file.dirtyMetadata = true
	}

	if req.Valid.Uid() {
		file.entry.Attributes.Uid = req.Uid
		file.dirtyMetadata = true
	}

	if req.Valid.Gid() {
		file.entry.Attributes.Gid = req.Gid
		file.dirtyMetadata = true
	}

	if req.Valid.Crtime() {
		file.entry.Attributes.Crtime = req.Crtime.Unix()
		file.dirtyMetadata = true
	}

	if req.Valid.Mtime() {
		file.entry.Attributes.Mtime = req.Mtime.Unix()
		file.dirtyMetadata = true
	}

	if req.Valid.Handle() {
		// fmt.Printf("file handle => %d\n", req.Handle)
	}

	if file.isOpen > 0 {
		return nil
	}

	if !file.dirtyMetadata {
		return nil
	}

	return file.saveEntry()

}

func (file *File) Setxattr(ctx context.Context, req *fuse.SetxattrRequest) error {

	glog.V(4).Infof("file Setxattr %s: %s", file.fullpath(), req.Name)

	if err := file.maybeLoadEntry(ctx); err != nil {
		return err
	}

	if err := setxattr(file.entry, req); err != nil {
		return err
	}

	return file.saveEntry()

}

func (file *File) Removexattr(ctx context.Context, req *fuse.RemovexattrRequest) error {

	glog.V(4).Infof("file Removexattr %s: %s", file.fullpath(), req.Name)

	if err := file.maybeLoadEntry(ctx); err != nil {
		return err
	}

	if err := removexattr(file.entry, req); err != nil {
		return err
	}

	return file.saveEntry()

}

func (file *File) Listxattr(ctx context.Context, req *fuse.ListxattrRequest, resp *fuse.ListxattrResponse) error {

	glog.V(4).Infof("file Listxattr %s", file.fullpath())

	if err := file.maybeLoadEntry(ctx); err != nil {
		return err
	}

	if err := listxattr(file.entry, req, resp); err != nil {
		return err
	}

	return nil

}

func (file *File) Fsync(ctx context.Context, req *fuse.FsyncRequest) error {
	// fsync works at OS level
	// write the file chunks to the filerGrpcAddress
	glog.V(4).Infof("%s/%s fsync file %+v", file.dir.FullPath(), file.Name, req)

	return nil
}

func (file *File) Forget() {
	t := util.NewFullPath(file.dir.FullPath(), file.Name)
	glog.V(5).Infof("Forget file %s", t)
	file.wfs.fsNodeCache.DeleteFsNode(t)
}

func (file *File) maybeLoadEntry(ctx context.Context) error {
	if file.entry == nil || file.isOpen <= 0 {
		entry, err := file.wfs.maybeLoadEntry(file.dir.FullPath(), file.Name)
		if err != nil {
			glog.V(3).Infof("maybeLoadEntry file %s/%s: %v", file.dir.FullPath(), file.Name, err)
			return err
		}
		if entry != nil {
			file.setEntry(entry)
		}
	}
	return nil
}

func (file *File) addChunks(chunks []*filer_pb.FileChunk) {

	sort.Slice(chunks, func(i, j int) bool {
		return chunks[i].Fid.FileKey < chunks[j].Fid.FileKey
	})

	var newVisibles []filer2.VisibleInterval
	for _, chunk := range chunks {
		newVisibles = filer2.MergeIntoVisibles(file.entryViewCache, newVisibles, chunk)
		t := file.entryViewCache[:0]
		file.entryViewCache = newVisibles
		newVisibles = t
	}

	file.reader = nil

	glog.V(4).Infof("%s existing %d chunks adds %d more", file.fullpath(), len(file.entry.Chunks), len(chunks))

	file.entry.Chunks = append(file.entry.Chunks, chunks...)
}

func (file *File) setEntry(entry *filer_pb.Entry) {
	file.entry = entry
	file.entryViewCache, _ = filer2.NonOverlappingVisibleIntervals(filer2.LookupFn(file.wfs), file.entry.Chunks)
	file.reader = nil
}

func (file *File) saveEntry() error {
	return file.wfs.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		request := &filer_pb.UpdateEntryRequest{
			Directory: file.dir.FullPath(),
			Entry:     file.entry,
		}

		glog.V(4).Infof("save file entry: %v", request)
		_, err := client.UpdateEntry(context.Background(), request)
		if err != nil {
			glog.Errorf("UpdateEntry file %s/%s: %v", file.dir.FullPath(), file.Name, err)
			return fuse.EIO
		}

		file.wfs.metaCache.UpdateEntry(context.Background(), filer2.FromPbEntry(request.Directory, request.Entry))

		return nil
	})
}
<commit_msg>add file handle locking when changing file entry<commit_after>package filesys

import (
	"context"
	"io"
	"os"
	"sort"
	"time"

	"github.com/seaweedfs/fuse"
	"github.com/seaweedfs/fuse/fs"

	"github.com/chrislusf/seaweedfs/weed/filer2"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
)

const blockSize = 512

var _ = fs.Node(&File{})
var _ = fs.NodeOpener(&File{})
var _ = fs.NodeFsyncer(&File{})
var _ = fs.NodeSetattrer(&File{})
var _ = fs.NodeGetxattrer(&File{})
var _ = fs.NodeSetxattrer(&File{})
var _ = fs.NodeRemovexattrer(&File{})
var _ = fs.NodeListxattrer(&File{})
var _ = fs.NodeForgetter(&File{})

type File struct {
	Name           string
	dir            *Dir
	wfs            *WFS
	entry          *filer_pb.Entry
	entryViewCache []filer2.VisibleInterval
	isOpen         int
	reader         io.ReaderAt
	dirtyMetadata  bool
}

func (file *File) fullpath() util.FullPath {
	return util.NewFullPath(file.dir.FullPath(), file.Name)
}

func (file *File) Attr(ctx context.Context, attr *fuse.Attr) error {

	glog.V(5).Infof("file Attr %s, open:%v, existing attr: %+v", file.fullpath(), file.isOpen, attr)

	if file.isOpen <= 0 {
		if err := file.maybeLoadEntry(ctx); err != nil {
			return err
		}
	}

	attr.Inode = file.fullpath().AsInode()
	attr.Valid = time.Second
	attr.Mode = os.FileMode(file.entry.Attributes.FileMode)
	attr.Size = filer2.FileSize(file.entry)
	if file.isOpen > 0 {
		attr.Size = file.entry.Attributes.FileSize
		glog.V(4).Infof("file Attr %s, open:%v, size: %d", file.fullpath(), file.isOpen, attr.Size)
	}
	attr.Crtime = time.Unix(file.entry.Attributes.Crtime, 0)
	attr.Mtime = time.Unix(file.entry.Attributes.Mtime, 0)
	attr.Gid = file.entry.Attributes.Gid
	attr.Uid = file.entry.Attributes.Uid
	attr.Blocks = attr.Size/blockSize + 1
	attr.BlockSize = uint32(file.wfs.option.ChunkSizeLimit)

	return nil

}

func (file *File) Getxattr(ctx context.Context, req *fuse.GetxattrRequest, resp *fuse.GetxattrResponse) error {

	glog.V(4).Infof("file Getxattr %s", file.fullpath())

	if err := file.maybeLoadEntry(ctx); err != nil {
		return err
	}

	return getxattr(file.entry, req, resp)
}

func (file *File) Open(ctx context.Context, req *fuse.OpenRequest, resp *fuse.OpenResponse) (fs.Handle, error) {

	glog.V(4).Infof("file %v open %+v", file.fullpath(), req)

	file.isOpen++

	handle := file.wfs.AcquireHandle(file, req.Uid, req.Gid)

	resp.Handle = fuse.HandleID(handle.handle)

	glog.V(4).Infof("%v file open handle id = %d", file.fullpath(), handle.handle)

	return handle, nil

}

func (file *File) Setattr(ctx context.Context, req *fuse.SetattrRequest, resp *fuse.SetattrResponse) error {

	glog.V(5).Infof("%v file setattr %+v", file.fullpath(), req)

	if err := file.maybeLoadEntry(ctx); err != nil {
		return err
	}
	if file.isOpen > 0 {
		file.wfs.handlesLock.Lock()
		fileHandle := file.wfs.handles[file.fullpath().AsInode()]
		file.wfs.handlesLock.Unlock()

		if fileHandle != nil {
			fileHandle.Lock()
			defer fileHandle.Unlock()
		}
	}

	if req.Valid.Size() {

		glog.V(4).Infof("%v file setattr set size=%v chunks=%d", file.fullpath(), req.Size, len(file.entry.Chunks))
		if req.Size < filer2.TotalSize(file.entry.Chunks) {
			// fmt.Printf("truncate %v \n", fullPath)
			var chunks []*filer_pb.FileChunk
			var truncatedChunks []*filer_pb.FileChunk
			for _, chunk := range file.entry.Chunks {
				int64Size := int64(chunk.Size)
				if chunk.Offset+int64Size > int64(req.Size) {
					// this chunk is truncated
					int64Size = int64(req.Size) - chunk.Offset
					if int64Size > 0 {
						chunks = append(chunks, chunk)
						glog.V(4).Infof("truncated chunk %+v from %d to %d\n", chunk.GetFileIdString(), chunk.Size, int64Size)
						chunk.Size = uint64(int64Size)
					} else {
						glog.V(4).Infof("truncated whole chunk %+v\n", chunk.GetFileIdString())
						truncatedChunks = append(truncatedChunks, chunk)
					}
				}
			}
			file.wfs.deleteFileChunks(truncatedChunks)
			file.entry.Chunks = chunks
			file.entryViewCache = nil
			file.reader = nil
		}
		file.entry.Attributes.FileSize = req.Size
		file.dirtyMetadata = true
	}

	if req.Valid.Mode() {
		file.entry.Attributes.FileMode = uint32(req.Mode)
		file.dirtyMetadata = true
	}

	if req.Valid.Uid() {
		file.entry.Attributes.Uid = req.Uid
		file.dirtyMetadata = true
	}

	if req.Valid.Gid() {
		file.entry.Attributes.Gid = req.Gid
		file.dirtyMetadata = true
	}

	if req.Valid.Crtime() {
		file.entry.Attributes.Crtime = req.Crtime.Unix()
		file.dirtyMetadata = true
	}

	if req.Valid.Mtime() {
		file.entry.Attributes.Mtime = req.Mtime.Unix()
		file.dirtyMetadata = true
	}

	if req.Valid.Handle() {
		// fmt.Printf("file handle => %d\n", req.Handle)
	}

	if file.isOpen > 0 {
		return nil
	}

	if !file.dirtyMetadata {
		return nil
	}

	return file.saveEntry()

}

func (file *File) Setxattr(ctx context.Context, req *fuse.SetxattrRequest) error {

	glog.V(4).Infof("file Setxattr %s: %s", file.fullpath(), req.Name)

	if err := file.maybeLoadEntry(ctx); err != nil {
		return err
	}

	if err := setxattr(file.entry, req); err != nil {
		return err
	}

	return file.saveEntry()

}

func (file *File) Removexattr(ctx context.Context, req *fuse.RemovexattrRequest) error {

	glog.V(4).Infof("file Removexattr %s: %s", file.fullpath(), req.Name)

	if err := file.maybeLoadEntry(ctx); err != nil {
		return err
	}

	if err := removexattr(file.entry, req); err != nil {
		return err
	}

	return file.saveEntry()

}

func (file *File) Listxattr(ctx context.Context, req *fuse.ListxattrRequest, resp *fuse.ListxattrResponse) error {

	glog.V(4).Infof("file Listxattr %s", file.fullpath())

	if err := file.maybeLoadEntry(ctx); err != nil {
		return err
	}

	if err := listxattr(file.entry, req, resp); err != nil {
		return err
	}

	return nil

}

func (file *File) Fsync(ctx context.Context, req *fuse.FsyncRequest) error {
	// fsync works at OS level
	// write the file chunks to the filerGrpcAddress
	glog.V(4).Infof("%s/%s fsync file %+v", file.dir.FullPath(), file.Name, req)

	return nil
}

func (file *File) Forget() {
	t := util.NewFullPath(file.dir.FullPath(), file.Name)
	glog.V(5).Infof("Forget file %s", t)
	file.wfs.fsNodeCache.DeleteFsNode(t)
}

func (file *File) maybeLoadEntry(ctx context.Context) error {
	if file.entry == nil || file.isOpen <= 0 {
		entry, err := file.wfs.maybeLoadEntry(file.dir.FullPath(), file.Name)
		if err != nil {
			glog.V(3).Infof("maybeLoadEntry file %s/%s: %v", file.dir.FullPath(), file.Name, err)
			return err
		}
		if entry != nil {
			file.setEntry(entry)
		}
	}
	return nil
}

func (file *File) addChunks(chunks []*filer_pb.FileChunk) {

	sort.Slice(chunks, func(i, j int) bool {
		return chunks[i].Fid.FileKey < chunks[j].Fid.FileKey
	})

	var newVisibles []filer2.VisibleInterval
	for _, chunk := range chunks {
		newVisibles = filer2.MergeIntoVisibles(file.entryViewCache, newVisibles, chunk)
		t := file.entryViewCache[:0]
		file.entryViewCache = newVisibles
		newVisibles = t
	}

	file.reader = nil

	glog.V(4).Infof("%s existing %d chunks adds %d more", file.fullpath(), len(file.entry.Chunks), len(chunks))

	file.entry.Chunks = append(file.entry.Chunks, chunks...)
}

func (file *File) setEntry(entry *filer_pb.Entry) {
	file.entry = entry
	file.entryViewCache, _ = filer2.NonOverlappingVisibleIntervals(filer2.LookupFn(file.wfs), file.entry.Chunks)
	file.reader = nil
}

func (file *File) saveEntry() error {
	return file.wfs.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		request := &filer_pb.UpdateEntryRequest{
			Directory: file.dir.FullPath(),
			Entry:     file.entry,
		}

		glog.V(4).Infof("save file entry: %v", request)
		_, err := client.UpdateEntry(context.Background(), request)
		if err != nil {
			glog.Errorf("UpdateEntry file %s/%s: %v", file.dir.FullPath(), file.Name, err)
			return fuse.EIO
		}

		file.wfs.metaCache.UpdateEntry(context.Background(), filer2.FromPbEntry(request.Directory, request.Entry))

		return nil
	})
}
<|endoftext|>


<commit_before>package welove

import (
	"testing"
	"io/ioutil"
	"github.com/bitly/go-simplejson"
)

func TestRandomHouse(t *testing.T) {
	_, ok := RandomHouse("562949961343086-2ca7e299a09974dd0")
	if !ok {
		t.Error()
	}
}

func TestVisit(t *testing.T) {
	id, ok := RandomHouse("562949961343086-2ca7e299a09974dd0")
	if !ok {
		t.Error()
	}
	res, err := Visit("562949961343086-2ca7e299a09974dd0", id)
	if err != nil {
		t.Error(err)
	}
	bytes, _ := ioutil.ReadAll(res.Body)
	js, err := simplejson.NewJson(bytes)
	if err != nil {
		t.Error(err)
	}
	result, _ := js.Get("result").Int()
	if result != 1 {
		t.Error("响应值result错误.")
	}
}<commit_msg>完善单元测试<commit_after>package welove

import (
	"testing"
	"io/ioutil"
	"github.com/bitly/go-simplejson"
)

func TestRandomHouse(t *testing.T) {
	_, ok := RandomHouse("562949961343086-2ca7e299a09974dd0")
	if !ok {
		t.Error()
	}
}

func TestVisit(t *testing.T) {
	id, ok := RandomHouse("562949961343086-2ca7e299a09974dd0")
	if !ok {
		t.Error()
	}
	res, err := Visit("562949961343086-2ca7e299a09974dd0", id)
	if err != nil {
		t.Error(err)
	}
	bytes, _ := ioutil.ReadAll(res.Body)
	js, err := simplejson.NewJson(bytes)
	if err != nil {
		t.Error(err)
	}
	result, _ := js.Get("result").Int()
	if result != 1 && result != 1201 {
		t.Error("响应值result错误.")
	}
}<|endoftext|>


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

package router

import (
	"bytes"
	"testing"
	"time"

	"github.com/tsuru/planb/log"
	"github.com/tsuru/planb/reverseproxy"
	"gopkg.in/check.v1"
	"gopkg.in/redis.v3"
)

type S struct {
	redis *redis.Client
}

var _ = check.Suite(&S{})

func Test(t *testing.T) {
	check.TestingT(t)
}

func clearKeys(r *redis.Client) error {
	val := r.Keys("frontend:*").Val()
	val = append(val, r.Keys("dead:*").Val()...)
	if len(val) > 0 {
		return r.Del(val...).Err()
	}
	return nil
}

func redisConn() (*redis.Client, error) {
	return redis.NewClient(&redis.Options{Addr: "127.0.0.1:6379", DB: 0}), nil
}

func (s *S) SetUpTest(c *check.C) {
	var err error
	s.redis, err = redisConn()
	c.Assert(err, check.IsNil)
	err = clearKeys(s.redis)
	c.Assert(err, check.IsNil)
}

func (s *S) TearDownTest(c *check.C) {
	s.redis.Close()
}

func (s *S) TestInit(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	c.Assert(router.roundRobin, check.DeepEquals, map[string]*int32{})
	c.Assert(router.logger, check.IsNil)
	c.Assert(router.cache, check.IsNil)
	c.Assert(router.Backend, check.NotNil)
}

func (s *S) TestInitCacheEnabled(c *check.C) {
	router := Router{CacheEnabled: true}
	err := router.Init()
	c.Assert(err, check.IsNil)
	c.Assert(router.roundRobin, check.DeepEquals, map[string]*int32{})
	c.Assert(router.logger, check.IsNil)
	c.Assert(router.cache, check.NotNil)
	c.Assert(router.Backend, check.NotNil)
}

func (s *S) TestChooseBackend(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	err = s.redis.RPush("frontend:myfrontend.com", "myfrontend", "http://url1:123").Err()
	c.Assert(err, check.IsNil)
	reqData, err := router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.StartTime.IsZero(), check.Equals, false)
	reqData.StartTime = time.Time{}
	c.Assert(reqData, check.DeepEquals, &reverseproxy.RequestData{
		Backend:    "http://url1:123",
		BackendIdx: 0,
		BackendKey: "myfrontend.com",
		BackendLen: 1,
		Host:       "myfrontend.com",
	})
}

func (s *S) TestChooseBackendNotFound(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	reqData, err := router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.Equals, reverseproxy.ErrNoRegisteredBackends)
	c.Assert(reqData.StartTime.IsZero(), check.Equals, false)
	reqData.StartTime = time.Time{}
	c.Assert(reqData, check.DeepEquals, &reverseproxy.RequestData{
		Backend:    "",
		BackendIdx: 0,
		BackendLen: 0,
		BackendKey: "",
		Host:       "myfrontend.com",
	})
}

func (s *S) TestChooseBackendNoBackends(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	err = s.redis.RPush("frontend:myfrontend.com", "myfrontend").Err()
	c.Assert(err, check.IsNil)
	reqData, err := router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.Equals, reverseproxy.ErrNoRegisteredBackends)
	c.Assert(reqData.StartTime.IsZero(), check.Equals, false)
	reqData.StartTime = time.Time{}
	c.Assert(reqData, check.DeepEquals, &reverseproxy.RequestData{
		Backend:    "",
		BackendIdx: 0,
		BackendLen: 0,
		BackendKey: "",
		Host:       "myfrontend.com",
	})
}

func (s *S) TestChooseBackendAllDead(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	err = s.redis.RPush("frontend:myfrontend.com", "myfrontend", "http://url1:123").Err()
	c.Assert(err, check.IsNil)
	err = s.redis.SAdd("dead:myfrontend.com", "0").Err()
	c.Assert(err, check.IsNil)
	reqData, err := router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.Equals, reverseproxy.ErrAllBackendsDead)
	c.Assert(reqData.StartTime.IsZero(), check.Equals, false)
	reqData.StartTime = time.Time{}
	c.Assert(reqData, check.DeepEquals, &reverseproxy.RequestData{
		Backend:    "",
		BackendIdx: 0,
		BackendLen: 1,
		BackendKey: "myfrontend.com",
		Host:       "myfrontend.com",
	})
}

func (s *S) TestChooseBackendRoundRobin(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	err = s.redis.RPush("frontend:myfrontend.com", "myfrontend", "http://url1:123", "http://url2:123", "http://url3:123").Err()
	c.Assert(err, check.IsNil)
	reqData, err := router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.StartTime.IsZero(), check.Equals, false)
	reqData.StartTime = time.Time{}
	c.Assert(reqData, check.DeepEquals, &reverseproxy.RequestData{
		Backend:    "http://url1:123",
		BackendIdx: 0,
		BackendKey: "myfrontend.com",
		BackendLen: 3,
		Host:       "myfrontend.com",
	})
	err = s.redis.RPush("frontend:myfrontend.com", "http://url4:123").Err()
	c.Assert(err, check.IsNil)
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url2:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url3:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url4:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url1:123")
}

func (s *S) TestChooseBackendRoundRobinWithCache(c *check.C) {
	router := Router{CacheEnabled: true}
	err := router.Init()
	c.Assert(err, check.IsNil)
	err = s.redis.RPush("frontend:myfrontend.com", "myfrontend", "http://url1:123", "http://url2:123", "http://url3:123").Err()
	c.Assert(err, check.IsNil)
	reqData, err := router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url1:123")
	err = s.redis.RPush("frontend:myfrontend.com", "http://url4:123").Err()
	c.Assert(err, check.IsNil)
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url2:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url3:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url1:123")
	time.Sleep(cacheTTLExpires)
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url1:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url2:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url3:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url4:123")
}

type bufferCloser struct {
	bytes.Buffer
}

func (b *bufferCloser) Close() error {
	return nil
}

func (s *S) TestEndRequest(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	buf := bufferCloser{}
	router.logger = log.NewWriterLogger(&buf)
	data := &reverseproxy.RequestData{
		Host: "myfe.com",
	}
	err = router.EndRequest(data, false, nil)
	c.Assert(err, check.IsNil)
	members := s.redis.SMembers("dead:myfe.com").Val()
	c.Assert(members, check.DeepEquals, []string{})
	router.Stop()
	c.Assert(buf.String(), check.Equals, "")
}

func (s *S) TestEndRequestWithLogFunc(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	buf := bufferCloser{}
	router.logger = log.NewWriterLogger(&buf)
	data := &reverseproxy.RequestData{
		Host: "myfe.com",
	}
	err = router.EndRequest(data, false, func() *log.LogEntry { return &log.LogEntry{} })
	c.Assert(err, check.IsNil)
	members := s.redis.SMembers("dead:myfe.com").Val()
	c.Assert(members, check.DeepEquals, []string{})
	router.Stop()
	c.Assert(buf.String(), check.Equals, "::ffff: - - [01/Jan/0001:00:00:00 +0000] \"  \" 0 0 \"\" \"\" \":\" \"\" 0.000 0.000\n")
}

func (s *S) TestEndRequestWithError(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	data := &reverseproxy.RequestData{
		Host: "myfe.com",
	}
	err = router.EndRequest(data, true, nil)
	c.Assert(err, check.IsNil)
	members := s.redis.SMembers("dead:myfe.com").Val()
	c.Assert(members, check.DeepEquals, []string{"0"})
}

func BenchmarkChooseBackend(b *testing.B) {
	r, err := redisConn()
	if err != nil {
		b.Fatal(err)
	}
	defer clearKeys(r)
	err = r.RPush("frontend:myfrontend.com", "myfrontend", "http://url1:123", "http://url2:123", "http://url3:123").Err()
	if err != nil {
		b.Fatal(err)
	}
	router := Router{}
	err = router.Init()
	if err != nil {
		b.Fatal(err)
	}
	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			router.ChooseBackend("myfrontend.com")
		}
	})
	b.StopTimer()
}

func BenchmarkChooseBackendNoCache(b *testing.B) {
	r, err := redisConn()
	if err != nil {
		b.Fatal(err)
	}
	defer clearKeys(r)
	err = r.RPush("frontend:myfrontend.com", "myfrontend", "http://url1:123", "http://url2:123", "http://url3:123").Err()
	if err != nil {
		b.Fatal(err)
	}
	router := Router{}
	err = router.Init()
	if err != nil {
		b.Fatal(err)
	}
	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			router.ChooseBackend("myfrontend.com")
			router.cache.Purge()
		}
	})
	b.StopTimer()
}

func BenchmarkChooseBackendManyNoCache(b *testing.B) {
	r, err := redisConn()
	if err != nil {
		b.Fatal(err)
	}
	defer clearKeys(r)
	backends := make([]string, 100)
	for i := range backends {
		backends[i] = "http://urlx:123"
	}
	backends = append([]string{"benchfrontend"}, backends...)
	err = r.RPush("frontend:myfrontend.com", backends...).Err()
	if err != nil {
		b.Fatal(err)
	}
	router := Router{}
	err = router.Init()
	if err != nil {
		b.Fatal(err)
	}
	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			router.ChooseBackend("myfrontend.com")
			router.cache.Purge()
		}
	})
	b.StopTimer()
}
<commit_msg>router: fix panic in tests without cache<commit_after>// Copyright 2016 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package router

import (
	"bytes"
	"testing"
	"time"

	"github.com/tsuru/planb/log"
	"github.com/tsuru/planb/reverseproxy"
	"gopkg.in/check.v1"
	"gopkg.in/redis.v3"
)

type S struct {
	redis *redis.Client
}

var _ = check.Suite(&S{})

func Test(t *testing.T) {
	check.TestingT(t)
}

func clearKeys(r *redis.Client) error {
	val := r.Keys("frontend:*").Val()
	val = append(val, r.Keys("dead:*").Val()...)
	if len(val) > 0 {
		return r.Del(val...).Err()
	}
	return nil
}

func redisConn() (*redis.Client, error) {
	return redis.NewClient(&redis.Options{Addr: "127.0.0.1:6379", DB: 0}), nil
}

func (s *S) SetUpTest(c *check.C) {
	var err error
	s.redis, err = redisConn()
	c.Assert(err, check.IsNil)
	err = clearKeys(s.redis)
	c.Assert(err, check.IsNil)
}

func (s *S) TearDownTest(c *check.C) {
	s.redis.Close()
}

func (s *S) TestInit(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	c.Assert(router.roundRobin, check.DeepEquals, map[string]*int32{})
	c.Assert(router.logger, check.IsNil)
	c.Assert(router.cache, check.IsNil)
	c.Assert(router.Backend, check.NotNil)
}

func (s *S) TestInitCacheEnabled(c *check.C) {
	router := Router{CacheEnabled: true}
	err := router.Init()
	c.Assert(err, check.IsNil)
	c.Assert(router.roundRobin, check.DeepEquals, map[string]*int32{})
	c.Assert(router.logger, check.IsNil)
	c.Assert(router.cache, check.NotNil)
	c.Assert(router.Backend, check.NotNil)
}

func (s *S) TestChooseBackend(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	err = s.redis.RPush("frontend:myfrontend.com", "myfrontend", "http://url1:123").Err()
	c.Assert(err, check.IsNil)
	reqData, err := router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.StartTime.IsZero(), check.Equals, false)
	reqData.StartTime = time.Time{}
	c.Assert(reqData, check.DeepEquals, &reverseproxy.RequestData{
		Backend:    "http://url1:123",
		BackendIdx: 0,
		BackendKey: "myfrontend.com",
		BackendLen: 1,
		Host:       "myfrontend.com",
	})
}

func (s *S) TestChooseBackendNotFound(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	reqData, err := router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.Equals, reverseproxy.ErrNoRegisteredBackends)
	c.Assert(reqData.StartTime.IsZero(), check.Equals, false)
	reqData.StartTime = time.Time{}
	c.Assert(reqData, check.DeepEquals, &reverseproxy.RequestData{
		Backend:    "",
		BackendIdx: 0,
		BackendLen: 0,
		BackendKey: "",
		Host:       "myfrontend.com",
	})
}

func (s *S) TestChooseBackendNoBackends(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	err = s.redis.RPush("frontend:myfrontend.com", "myfrontend").Err()
	c.Assert(err, check.IsNil)
	reqData, err := router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.Equals, reverseproxy.ErrNoRegisteredBackends)
	c.Assert(reqData.StartTime.IsZero(), check.Equals, false)
	reqData.StartTime = time.Time{}
	c.Assert(reqData, check.DeepEquals, &reverseproxy.RequestData{
		Backend:    "",
		BackendIdx: 0,
		BackendLen: 0,
		BackendKey: "",
		Host:       "myfrontend.com",
	})
}

func (s *S) TestChooseBackendAllDead(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	err = s.redis.RPush("frontend:myfrontend.com", "myfrontend", "http://url1:123").Err()
	c.Assert(err, check.IsNil)
	err = s.redis.SAdd("dead:myfrontend.com", "0").Err()
	c.Assert(err, check.IsNil)
	reqData, err := router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.Equals, reverseproxy.ErrAllBackendsDead)
	c.Assert(reqData.StartTime.IsZero(), check.Equals, false)
	reqData.StartTime = time.Time{}
	c.Assert(reqData, check.DeepEquals, &reverseproxy.RequestData{
		Backend:    "",
		BackendIdx: 0,
		BackendLen: 1,
		BackendKey: "myfrontend.com",
		Host:       "myfrontend.com",
	})
}

func (s *S) TestChooseBackendRoundRobin(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	err = s.redis.RPush("frontend:myfrontend.com", "myfrontend", "http://url1:123", "http://url2:123", "http://url3:123").Err()
	c.Assert(err, check.IsNil)
	reqData, err := router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.StartTime.IsZero(), check.Equals, false)
	reqData.StartTime = time.Time{}
	c.Assert(reqData, check.DeepEquals, &reverseproxy.RequestData{
		Backend:    "http://url1:123",
		BackendIdx: 0,
		BackendKey: "myfrontend.com",
		BackendLen: 3,
		Host:       "myfrontend.com",
	})
	err = s.redis.RPush("frontend:myfrontend.com", "http://url4:123").Err()
	c.Assert(err, check.IsNil)
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url2:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url3:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url4:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url1:123")
}

func (s *S) TestChooseBackendRoundRobinWithCache(c *check.C) {
	router := Router{CacheEnabled: true}
	err := router.Init()
	c.Assert(err, check.IsNil)
	err = s.redis.RPush("frontend:myfrontend.com", "myfrontend", "http://url1:123", "http://url2:123", "http://url3:123").Err()
	c.Assert(err, check.IsNil)
	reqData, err := router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url1:123")
	err = s.redis.RPush("frontend:myfrontend.com", "http://url4:123").Err()
	c.Assert(err, check.IsNil)
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url2:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url3:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url1:123")
	time.Sleep(cacheTTLExpires)
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url1:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url2:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url3:123")
	reqData, err = router.ChooseBackend("myfrontend.com")
	c.Assert(err, check.IsNil)
	c.Assert(reqData.Backend, check.Equals, "http://url4:123")
}

type bufferCloser struct {
	bytes.Buffer
}

func (b *bufferCloser) Close() error {
	return nil
}

func (s *S) TestEndRequest(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	buf := bufferCloser{}
	router.logger = log.NewWriterLogger(&buf)
	data := &reverseproxy.RequestData{
		Host: "myfe.com",
	}
	err = router.EndRequest(data, false, nil)
	c.Assert(err, check.IsNil)
	members := s.redis.SMembers("dead:myfe.com").Val()
	c.Assert(members, check.DeepEquals, []string{})
	router.Stop()
	c.Assert(buf.String(), check.Equals, "")
}

func (s *S) TestEndRequestWithLogFunc(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	buf := bufferCloser{}
	router.logger = log.NewWriterLogger(&buf)
	data := &reverseproxy.RequestData{
		Host: "myfe.com",
	}
	err = router.EndRequest(data, false, func() *log.LogEntry { return &log.LogEntry{} })
	c.Assert(err, check.IsNil)
	members := s.redis.SMembers("dead:myfe.com").Val()
	c.Assert(members, check.DeepEquals, []string{})
	router.Stop()
	c.Assert(buf.String(), check.Equals, "::ffff: - - [01/Jan/0001:00:00:00 +0000] \"  \" 0 0 \"\" \"\" \":\" \"\" 0.000 0.000\n")
}

func (s *S) TestEndRequestWithError(c *check.C) {
	router := Router{}
	err := router.Init()
	c.Assert(err, check.IsNil)
	data := &reverseproxy.RequestData{
		Host: "myfe.com",
	}
	err = router.EndRequest(data, true, nil)
	c.Assert(err, check.IsNil)
	members := s.redis.SMembers("dead:myfe.com").Val()
	c.Assert(members, check.DeepEquals, []string{"0"})
}

func BenchmarkChooseBackend(b *testing.B) {
	r, err := redisConn()
	if err != nil {
		b.Fatal(err)
	}
	defer clearKeys(r)
	err = r.RPush("frontend:myfrontend.com", "myfrontend", "http://url1:123", "http://url2:123", "http://url3:123").Err()
	if err != nil {
		b.Fatal(err)
	}
	router := Router{
		CacheEnabled: true,
	}
	err = router.Init()
	if err != nil {
		b.Fatal(err)
	}
	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			router.ChooseBackend("myfrontend.com")
		}
	})
	b.StopTimer()
}

func BenchmarkChooseBackendNoCache(b *testing.B) {
	r, err := redisConn()
	if err != nil {
		b.Fatal(err)
	}
	defer clearKeys(r)
	err = r.RPush("frontend:myfrontend.com", "myfrontend", "http://url1:123", "http://url2:123", "http://url3:123").Err()
	if err != nil {
		b.Fatal(err)
	}
	router := Router{}
	err = router.Init()
	if err != nil {
		b.Fatal(err)
	}
	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			router.ChooseBackend("myfrontend.com")
		}
	})
	b.StopTimer()
}

func BenchmarkChooseBackendManyNoCache(b *testing.B) {
	r, err := redisConn()
	if err != nil {
		b.Fatal(err)
	}
	defer clearKeys(r)
	backends := make([]string, 100)
	for i := range backends {
		backends[i] = "http://urlx:123"
	}
	backends = append([]string{"benchfrontend"}, backends...)
	err = r.RPush("frontend:myfrontend.com", backends...).Err()
	if err != nil {
		b.Fatal(err)
	}
	router := Router{}
	err = router.Init()
	if err != nil {
		b.Fatal(err)
	}
	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			router.ChooseBackend("myfrontend.com")
		}
	})
	b.StopTimer()
}
<|endoftext|>


<commit_before>package parser

import (
  "bytes"
  "rubex"
  //"strconv"
)

// Type tags so we know what kind of token we have
type lexeme int
const (
  LPAREN = iota
  RPAREN
  LBRACE
  RBRACE
  COMMA
  DOT
  EQUAL
  STRING
  REGEXP
  POS
  GVAR
  LVAR
  KWD
  ID
  FUNC
  TYPE
  IMPORT
  READ
  EOF
  ERROR
)

// A token has a type (aka lexeme), a value, and a line number
type token struct {
  Lexeme lexeme
  Value string
  LineNum int
}

/*
  Represent a tokenizer with a struct containing the remaining source text and
  the line number. Easier than using a stateless tokenizing function that
  returns them as extra values and requires the parser to keep track of them.
*/
type Tokenizer struct {
  Source []byte
  LineNum int
  Lookahead *token
}

func (t *Tokenizer) hasPrefix(s string) bool {
  return bytes.HasPrefix(t.Source, []byte(s))
}

// Discard leading spaces (excluding newlines) in the source text.
func (t *Tokenizer) discardSpaces() {
  t.Source = bytes.TrimLeft(t.Source, " \t")
}

// Discard leading text until a newline (or EOF) is found.
func(t *Tokenizer) discardLine() {
  if i := bytes.IndexByte(t.Source, '\n'); i >= 0 {
    t.Source = t.Source[i:]
  } else {
    t.Source = t.Source[len(t.Source):]
  }
}

// Discard the leading comment in the source text.
func (t *Tokenizer) discardComment() {
  if t.hasPrefix("#") || t.hasPrefix("//") {
    t.discardLine()
  } else if t.hasPrefix("/*") {
    t.discardBlockComment()
  }
}

// Helper for discarding block comments.
// TO DO: ERROR HANDLING FOR UNTERMINATED COMMENTS
func (t *Tokenizer) discardBlockComment() {
  depth, i, length := 1, 2, len(t.Source)
  for depth > 0 {
    if i >= length {
      // ERROR
    }
    switch t.Source[i] {
    case '\n':
      t.LineNum++
    case '/':
      i++
      if i >= length {
        // ERROR
      }
      if t.Source[i] == '*' {
        depth++
      }
    case '*':
      i++
      if i >= length {
        // ERROR
      }
      if t.Source[i] == '/' {
        depth--
      }
    }
    i++
  }
  t.Source = t.Source[i:]
}

// Discard all leading whitespace and comments from the source text. Need to
// tally up the newlines to keep LineNum up to date.
func (t *Tokenizer) discardWhitespaceAndComments() {
  for len(t.Source) > 0 {
    switch {
    case t.hasPrefix("\n"):
      t.LineNum++
      t.Source = t.Source[1:]
    case t.hasPrefix(" ") || t.hasPrefix("\t"):
      t.discardSpaces()
    case t.hasPrefix("#") || t.hasPrefix("//") || t.hasPrefix("/*"):
      t.discardComment()
    default:
      return
    }
  }
}

// The heart of the tokenizer. This function tries to munch off a token from
// the head of the source text.
// func (t *Tokenizer) munch() {
//   tSrc = t.Source
//   switch {
//   case len(tSrc) == 0:

var lexemeName [20]string
var matcher [20]*rubex.Regexp

func init() {
  // Is there a more elegant way to do this?
  lexemeName[LPAREN] = "LPAREN"
  lexemeName[RPAREN] = "RPAREN"
  lexemeName[LBRACE] = "LBRACE"
  lexemeName[RBRACE] = "RBRACE"
  lexemeName[COMMA]  = "COMMA"
  lexemeName[DOT]    = "DOT"
  lexemeName[EQUAL]  = "EQUAL"
  lexemeName[STRING] = "STRING"
  lexemeName[REGEXP] = "REGEXP"
  lexemeName[POS]    = "POS"
  lexemeName[GVAR]   = "GVAR"
  lexemeName[LVAR]   = "LVAR"
  lexemeName[KWD]    = "KWD"
  lexemeName[ID]     = "ID"
  lexemeName[FUNC]   = "FUNC"
  lexemeName[TYPE]   = "TYPE"
  lexemeName[IMPORT] = "IMPORT"
  lexemeName[READ]   = "READ"
  lexemeName[EOF]    = "EOF"
  lexemeName[ERROR]  = "ERROR"

  // Inline comments below indicate which captures to use
  matcher[STRING], _ = rubex.Compile(`^"(\\.|[^"\\])*"|^'(\\.|[^'\\])*'`)  // 0
  matcher[REGEXP], _ = rubex.Compile(`^\/((\\.|[^\/\\])*)\/([imxouesn]*)`) // 1,3
  matcher[POS], _    = rubex.Compile("^(top|bottom|before|after)")         // 0
  matcher[GVAR], _   = rubex.Compile(`^\$(\w+)`)                           // 1
  matcher[LVAR], _   = rubex.Compile(`^%(\w+)`)                            // 1
  matcher[KWD], _    = rubex.Compile(`^([a-zA-Z_:][-\w:.]*):`)             // 1
  matcher[ID], _     = rubex.Compile(`^\$|[_a-z](\w|\$)*`)                 // 0
  matcher[TYPE], _   = rubex.Compile(`^[A-Z](\w*)`)                        // 0
}
<commit_msg>Workin on the meat of the tokenizer.<commit_after>package parser

import (
  "bytes"
  "rubex"
  "strconv"
)

// Type tags so we know what kind of token we have
type Lexeme int
const (
  LPAREN = iota
  RPAREN
  LBRACE
  RBRACE
  COMMA
  DOT
  EQUAL
  STRING
  REGEXP
  POS
  GVAR
  LVAR
  KWD
  ID
  FUNC
  TYPE
  IMPORT
  READ
  EOF
  ERROR
)

var lexemeName [20]string
var matcher [20]*rubex.Regexp
var symbolLexeme map[string]Lexeme
var symbolPattern *rubex.Regexp
var numberPattern *rubex.Regexp

func init() {
  // Is there a more elegant way to do this?
  lexemeName[LPAREN] = "LPAREN"
  lexemeName[RPAREN] = "RPAREN"
  lexemeName[LBRACE] = "LBRACE"
  lexemeName[RBRACE] = "RBRACE"
  lexemeName[COMMA]  = "COMMA"
  lexemeName[DOT]    = "DOT"
  lexemeName[EQUAL]  = "EQUAL"
  lexemeName[STRING] = "STRING"
  lexemeName[REGEXP] = "REGEXP"
  lexemeName[POS]    = "POS"
  lexemeName[GVAR]   = "GVAR"
  lexemeName[LVAR]   = "LVAR"
  lexemeName[KWD]    = "KWD"
  lexemeName[ID]     = "ID"
  lexemeName[FUNC]   = "FUNC"
  lexemeName[TYPE]   = "TYPE"
  lexemeName[IMPORT] = "IMPORT"
  lexemeName[READ]   = "READ"
  lexemeName[EOF]    = "EOF"
  lexemeName[ERROR]  = "ERROR"

  // Inline comments below indicate which captures to use. Should probably
  // check those error codes sometime....
  matcher[STRING], _ = rubex.Compile(`^"(\\.|[^"\\])*"|^'(\\.|[^'\\])*'`)  // 0
  matcher[REGEXP], _ = rubex.Compile(`^\/((\\.|[^\/\\])*)\/([imxouesn]*)`) // 1,3
  matcher[POS], _    = rubex.Compile("^(top|bottom|before|after)")         // 0
  matcher[GVAR], _   = rubex.Compile(`^\$(\w+)`)                           // 1
  matcher[LVAR], _   = rubex.Compile(`^%(\w+)`)                            // 1
  matcher[KWD], _    = rubex.Compile(`^([a-zA-Z_:][-\w:.]*):`)             // 1
  matcher[ID], _     = rubex.Compile(`^\$|[_a-z](\w|\$)*`)                 // 0
  matcher[TYPE], _   = rubex.Compile(`^[A-Z](\w*)`)                        // 0

  // Map parens, braces, etc to their lexemes
  symbolLexeme = make(map[string]Lexeme, 7)
  symbolLexeme["("] = LPAREN
  symbolLexeme[")"] = RPAREN
  symbolLexeme["{"] = LBRACE
  symbolLexeme["}"] = RBRACE
  symbolLexeme[","] = COMMA
  symbolLexeme["."] = DOT
  symbolLexeme["="] = EQUAL
  symbolPattern, _ = rubex.Compile(`^[(){},\.=]`)

  numberPattern, _ = rubex.Compile(`\d+`)
}

// A token has a type (aka lexeme), a value, and a line number
type Token struct {
  Lexeme
  Value string
  ExtraValue string
  LineNum int
}

/*
  Represent a tokenizer with a struct containing the remaining source text and
  the line number. Easier than using a stateless tokenizing function that
  returns them as extra values and requires the parser to keep track of them.
*/
type Tokenizer struct {
  Source []byte
  LineNum int
  Lookahead *Token
}

func (t *Tokenizer) hasPrefix(s string) bool {
  return bytes.HasPrefix(t.Source, []byte(s))
}

// Discard leading spaces (excluding newlines) in the source text.
func (t *Tokenizer) discardSpaces() {
  t.Source = bytes.TrimLeft(t.Source, " \t")
}

// Discard leading text until a newline (or EOF) is found.
func(t *Tokenizer) discardLine() {
  if i := bytes.IndexByte(t.Source, '\n'); i >= 0 {
    t.Source = t.Source[i:]
  } else {
    t.Source = t.Source[len(t.Source):]
  }
}

// Discard the leading comment in the source text.
func (t *Tokenizer) discardComment() {
  if t.hasPrefix("#") || t.hasPrefix("//") {
    t.discardLine()
  } else if t.hasPrefix("/*") {
    t.discardBlockComment()
  }
}

// Helper for discarding block comments.
// TO DO: ERROR HANDLING FOR UNTERMINATED COMMENTS
func (t *Tokenizer) discardBlockComment() {
  depth, i, length := 1, 2, len(t.Source)
  for depth > 0 {
    if i >= length {
      // ERROR
    }
    switch t.Source[i] {
    case '\n':
      t.LineNum++
    case '/':
      i++
      if i >= length {
        // ERROR
      }
      if t.Source[i] == '*' {
        depth++
      }
    case '*':
      i++
      if i >= length {
        // ERROR
      }
      if t.Source[i] == '/' {
        depth--
      }
    }
    i++
  }
  t.Source = t.Source[i:]
}

// Discard all leading whitespace and comments from the source text. Need to
// tally up the newlines to keep LineNum up to date.
func (t *Tokenizer) discardWhitespaceAndComments() {
  for len(t.Source) > 0 {
    switch {
    case t.hasPrefix("\n"):
      t.LineNum++
      t.Source = t.Source[1:]
    case t.hasPrefix(" ") || t.hasPrefix("\t"):
      t.discardSpaces()
    case t.hasPrefix("#") || t.hasPrefix("//") || t.hasPrefix("/*"):
      t.discardComment()
    default:
      return
    }
  }
}

func (t *Tokenizer) popToken(lexeme Lexeme, value string, length int) *Token {
  val := &Token { Lexeme: lexeme, Value: value, ExtraValue: "", LineNum: t.LineNum }
  t.Source = t.Source[length:]
  return val
}

func (t *Tokenizer) popError(message string) *Token {
  val := &Token { Lexeme: ERROR, Value: message, ExtraValue: "", LineNum: t.LineNum }
  t.discardLine()
  return val
}

// The heart of the tokenizer. This function tries to munch off a token from
// the head of the source text.
func (t *Tokenizer) munch() *Token {
  src := t.Source
  if len(src) == 0 {
    return t.popToken(EOF, "", 0)
  } else if t.hasPrefix("*/") {
    return t.popError("unmatched comment terminator")
  } else if c := string(symbolPattern.Find(src)); len(c) > 0 {
    return t.popToken(symbolLexeme[c], c, 1)
  } else if c := string(numberPattern.Find(src)); len(c) > 0 {
    return t.popToken(STRING, c, len(c))
  } else if t.hasPrefix("'") || t.hasPrefix("\"") {
    if c := string(matcher[STRING].Find(src)); len(c) > 0 {
      unquoted, _ := strconv.Unquote(c)
      return t.popToken(STRING, unquoted, len(c))
    } else {
      return t.popError("unterminated string literal")
    }
  } else if t.hasPrefix("/") {
    if cs := matcher[REGEXP].FindSubmatch(src); len(cs) > 0 {
      pattern := cs[1][1:len(cs[1])-1]
      options := cs[3]
      val := t.popToken(REGEXP, string(pattern), len(cs[0]))
      val.ExtraValue = string(options)
      return val
    } else {
      return t.popError("unterminated regular expression literal")
    }
  } else if c := matcher[KWD].Find(src); len(c) > 0 {
    return t.popToken(KWD, string(c[:len(c)-1]), len(c))
  } else if c := string(matcher[ID].Find(src)); len(c) > 0 {
    if matcher[POS].MatchString(c) {
      return t.popToken(POS, c, len(c))
    } else if c == "read" {
      return t.popToken(READ, "", len(c))
    } else {
      return t.popToken(ID, c, len(c))
    }
  } else if c := matcher[GVAR].Find(src); len(c) > 0 {
    return t.popToken(GVAR, string(c[1:]), len(c))
  } else if c := matcher[LVAR].Find(src); len(c) > 0 {
    return t.popToken(LVAR, string(c[1:]), len(c))
  } else if c := string(matcher[TYPE].Find(src)); len(c) > 0 {
    return t.popToken(TYPE, c, len(c))
  }

  return t.popToken(ERROR, "blah", 0)
}

<|endoftext|>


<commit_before>package cloudformation

import (
	"encoding/json"
	"io/ioutil"
	"reflect"
	"testing"
)

var testFilePaths = []string{
	"examples/minimal.template",
	"examples/minimal_outputs.template",
	//"examples/app.template",
	"examples/app-minimal.template",
	"examples/LAMP_Multi_AZ.template",
	"examples/LAMP_Single_Instance.template",
	"examples/Rails_Multi_AZ.template",
	"examples/Rails_Single_Instance.template",
	"examples/Windows_Roles_And_Features.template",
	"examples/Windows_Single_Server_Active_Directory.template",
	"examples/Windows_Single_Server_SharePoint_Foundation.template",
	"examples/WordPress_Chef.template",
	"examples/WordPress_Multi_AZ.template",
	"examples/WordPress_Single_Instance.template",
}

func TestRoundtrips(t *testing.T) {
	for _, testFilePath := range testFilePaths {
		buf, err := ioutil.ReadFile(testFilePath)
		if err != nil {
			panic(err)
		}
		testMarshalResourcesPieceByPiece(t, testFilePath, buf)
		//testOne(t, buf)
	}
}

func testMarshalResourcesPieceByPiece(t *testing.T, path string, input []byte) {
	v := map[string]interface{}{}
	err := json.Unmarshal(input, &v)
	resources := v["Resources"].(map[string]interface{})
	for name, resource := range resources {
		buf, _ := json.Marshal(resource)
		r := Resource{}
		err = json.Unmarshal(buf, &r)
		if err != nil {
			t.Errorf("marshal: %s %s: %s", path, name, err)
			return
		}
	}
}

func testOne(t *testing.T, input []byte) {
	templ := Template{}
	err := json.Unmarshal(input, &templ)
	if err != nil {
		t.Errorf("decode: %s", err)
		return
	}

	output, err := json.Marshal(templ)
	if err != nil {
		t.Errorf("marshal: %s", err)
		return
	}

	parsedInput := map[string]interface{}{}
	json.Unmarshal(input, &parsedInput)

	parsedOutput := map[string]interface{}{}
	json.Unmarshal(output, &parsedOutput)

	if !reflect.DeepEqual(parsedInput, parsedOutput) {
		t.Errorf("expected %#v, got %#v", parsedInput, parsedOutput)
	}
}
<commit_msg>Enable app.template test<commit_after>package cloudformation

import (
	"encoding/json"
	"io/ioutil"
	"reflect"
	"testing"
)

var testFilePaths = []string{
	"examples/minimal.template",
	"examples/minimal_outputs.template",
	"examples/app.template",
	"examples/app-minimal.template",
	"examples/LAMP_Multi_AZ.template",
	"examples/LAMP_Single_Instance.template",
	"examples/Rails_Multi_AZ.template",
	"examples/Rails_Single_Instance.template",
	"examples/Windows_Roles_And_Features.template",
	"examples/Windows_Single_Server_Active_Directory.template",
	"examples/Windows_Single_Server_SharePoint_Foundation.template",
	"examples/WordPress_Chef.template",
	"examples/WordPress_Multi_AZ.template",
	"examples/WordPress_Single_Instance.template",
}

func TestRoundtrips(t *testing.T) {
	for _, testFilePath := range testFilePaths {
		buf, err := ioutil.ReadFile(testFilePath)
		if err != nil {
			panic(err)
		}
		testMarshalResourcesPieceByPiece(t, testFilePath, buf)
		//testOne(t, buf)
	}
}

func testMarshalResourcesPieceByPiece(t *testing.T, path string, input []byte) {
	v := map[string]interface{}{}
	err := json.Unmarshal(input, &v)
	resources := v["Resources"].(map[string]interface{})
	for name, resource := range resources {
		buf, _ := json.Marshal(resource)
		r := Resource{}
		err = json.Unmarshal(buf, &r)
		if err != nil {
			t.Errorf("marshal: %s %s: %s", path, name, err)
			return
		}
	}
}

func testOne(t *testing.T, input []byte) {
	templ := Template{}
	err := json.Unmarshal(input, &templ)
	if err != nil {
		t.Errorf("decode: %s", err)
		return
	}

	output, err := json.Marshal(templ)
	if err != nil {
		t.Errorf("marshal: %s", err)
		return
	}

	parsedInput := map[string]interface{}{}
	json.Unmarshal(input, &parsedInput)

	parsedOutput := map[string]interface{}{}
	json.Unmarshal(output, &parsedOutput)

	if !reflect.DeepEqual(parsedInput, parsedOutput) {
		t.Errorf("expected %#v, got %#v", parsedInput, parsedOutput)
	}
}
<|endoftext|>


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

package list

import (
	"github.com/andreaskoch/dotman/projects"
	"github.com/andreaskoch/dotman/ui"
)

const (
	ActionName        = "list"
	ActionDescription = "Get a list of all projects in the current dotfile collection."
)

type List struct {
	projectCollectionProvider func() *projects.Collection
}

func New(projectCollectionProvider func() *projects.Collection) *List {
	return &List{
		projectCollectionProvider: projectCollectionProvider,
	}
}

func (list *List) Name() string {
	return ActionName
}

func (list *List) Description() string {
	return ActionDescription
}

func (list *List) Execute(arguments []string) {
	list.execute(arguments)
}

func (list *List) DryRun(arguments []string) {
	list.execute(arguments)
}

func (list *List) execute(arguments []string) {
	projects := list.projectCollectionProvider()

	for _, project := range projects.Collection {
		ui.Message("%s", project)
	}
}
<commit_msg>Refactored the list action to use the Base action.<commit_after>// Copyright 2013 Andreas Koch. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package list

import (
	"github.com/andreaskoch/dotman/actions/base"
	"github.com/andreaskoch/dotman/projects"
	"github.com/andreaskoch/dotman/ui"
)

const (
	ActionName        = "list"
	ActionDescription = "Get a list of all projects in the current dotfile collection."
)

type List struct {
	*base.Action
}

func New(projectCollectionProvider func() *projects.Collection) *List {
	return &List{
		base.New(ActionName, ActionDescription, projectCollectionProvider, func(project *projects.Project, executeADryRunOnly bool) {
			ui.Message("%s", project)
		}),
	}
}
<|endoftext|>


<commit_before>package mongo

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path"
	"path/filepath"

	"github.com/juju/loggo"
	"labix.org/v2/mgo"

	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/replicaset"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/upstart"
	"launchpad.net/juju-core/utils"
	"launchpad.net/juju-core/version"
)

const (
	maxFiles = 65000
	maxProcs = 20000

	serviceName = "juju-db"

	// SharedSecretFile is the name of the Mongo shared secret file
	// located within the Juju data directory.
	SharedSecretFile = "shared-secret"

	ReplicaSetName = "juju"
)

var (
	logger          = loggo.GetLogger("juju.agent.mongo")
	mongoConfigPath = "/etc/default/mongodb"

	// JujuMongodPath holds the default path to the juju-specific mongod.
	JujuMongodPath = "/usr/lib/juju/bin/mongod"

	upstartConfInstall          = (*upstart.Conf).Install
	upstartServiceStopAndRemove = (*upstart.Service).StopAndRemove
)

// WithAddresses represents an entity that has a set of
// addresses. e.g. a state Machine object
type WithAddresses interface {
	Addresses() []instance.Address
}

// IsMaster returns a boolean that represents whether the given
// machine's peer address is the primary mongo host for the replicaset
func IsMaster(session *mgo.Session, obj WithAddresses) (bool, error) {
	addrs := obj.Addresses()

	masterHostPort, err := replicaset.MasterHostPort(session)

	// If the replica set has not been configured, then we
	// can have only one master and the caller must
	// be that master.
	if err == replicaset.ErrMasterNotConfigured {
		return true, nil
	}

	if err != nil {
		return false, err
	}

	masterAddr, _, err := net.SplitHostPort(masterHostPort)
	if err != nil {
		return false, err
	}

	machinePeerAddr := SelectPeerAddress(addrs)
	return machinePeerAddr == masterAddr, nil
}

// SelectPeerAddress returns the address to use as the
// mongo replica set peer address by selecting it from the given addresses.
func SelectPeerAddress(addrs []instance.Address) string {
	return instance.SelectInternalAddress(addrs, false)
}

// SelectPeerHostPort returns the HostPort to use as the
// mongo replica set peer by selecting it from the given hostPorts.
func SelectPeerHostPort(hostPorts []instance.HostPort) string {
	return instance.SelectInternalHostPort(hostPorts, false)
}

// GenerateSharedSecret generates a pseudo-random shared secret (keyfile)
// for use with Mongo replica sets.
func GenerateSharedSecret() (string, error) {
	// "A key’s length must be between 6 and 1024 characters and may
	// only contain characters in the base64 set."
	//   -- http://docs.mongodb.org/manual/tutorial/generate-key-file/
	buf := make([]byte, base64.StdEncoding.DecodedLen(1024))
	if _, err := rand.Read(buf); err != nil {
		return "", fmt.Errorf("cannot read random secret: %v", err)
	}
	return base64.StdEncoding.EncodeToString(buf), nil
}

// MongoPath returns the executable path to be used to run mongod on this
// machine. If the juju-bundled version of mongo exists, it will return that
// path, otherwise it will return the command to run mongod from the path.
func MongodPath() (string, error) {
	if _, err := os.Stat(JujuMongodPath); err == nil {
		return JujuMongodPath, nil
	}

	path, err := exec.LookPath("mongod")
	if err != nil {
		return "", err
	}
	return path, nil
}

// RemoveService removes the mongoDB upstart service from this machine.
func RemoveService(namespace string) error {
	svc := upstart.NewService(ServiceName(namespace))
	return upstartServiceStopAndRemove(svc)
}

// EnsureMongoServer ensures that the correct mongo upstart script is installed
// and running.
//
// This method will remove old versions of the mongo upstart script as necessary
// before installing the new version.
//
// The namespace is a unique identifier to prevent multiple instances of mongo
// on this machine from colliding. This should be empty unless using
// the local provider.
func EnsureMongoServer(dataDir string, namespace string, info params.StateServingInfo) error {
	logger.Infof("Ensuring mongo server is running; dataDir %s; port %d", dataDir, info.StatePort)
	dbDir := filepath.Join(dataDir, "db")

	if err := os.MkdirAll(dbDir, 0700); err != nil {
		return fmt.Errorf("cannot create mongo dbdir: %v", err)
	}

	certKey := info.Cert + "\n" + info.PrivateKey
	err := utils.AtomicWriteFile(sslKeyPath(dataDir), []byte(certKey), 0600)
	if err != nil {
		return fmt.Errorf("cannot write SSL key: %v", err)
	}

	err = utils.AtomicWriteFile(sharedSecretPath(dataDir), []byte(info.SharedSecret), 0600)
	if err != nil {
		return fmt.Errorf("cannot write mongod shared secret: %v", err)
	}

	// Disable the default mongodb installed by the mongodb-server package.
	// Only do this if the file doesn't exist already, so users can run
	// their own mongodb server if they wish to.
	if _, err := os.Stat(mongoConfigPath); os.IsNotExist(err) {
		err = ioutil.WriteFile(
			mongoConfigPath,
			[]byte("ENABLE_MONGODB=no"),
			0644,
		)
		if err != nil {
			return err
		}
	}

	if err := aptGetInstallMongod(); err != nil {
		return fmt.Errorf("cannot install mongod: %v", err)
	}

	upstartConf, err := mongoUpstartService(namespace, dataDir, dbDir, info.StatePort)
	if err != nil {
		return err
	}

	// TODO(natefinch) 2014-04-12 https://launchpad.net/bugs/1306902
	// remove this once we support upgrading to HA
	if service.Installed() {
		return nil
	}

	if err := makeJournalDirs(dbDir); err != nil {
		return fmt.Errorf("Error creating journal directories: %v", err)
	}
	return upstartConfInstall(upstartConf)
}

// ServiceName returns the name of the upstart service config for mongo using
// the given namespace.
func ServiceName(namespace string) string {
	if namespace != "" {
		return fmt.Sprintf("%s-%s", serviceName, namespace)
	}
	return serviceName
}

func makeJournalDirs(dataDir string) error {
	journalDir := path.Join(dataDir, "journal")

	if err := os.MkdirAll(journalDir, 0700); err != nil {
		logger.Errorf("failed to make mongo journal dir %s: %v", journalDir, err)
		return err
	}

	// manually create the prealloc files, since otherwise they get created as 100M files.
	zeroes := make([]byte, 64*1024) // should be enough for anyone
	for x := 0; x < 3; x++ {
		name := fmt.Sprintf("prealloc.%d", x)
		filename := filepath.Join(journalDir, name)
		f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0700)
		// TODO(jam) 2014-04-12 https://launchpad.net/bugs/1306902
		// When we support upgrading Mongo into Replica mode, we should
		// start rewriting the upstart config
		if os.IsExist(err) {
			// already exists, don't overwrite
			continue
		}
		if err != nil {
			return fmt.Errorf("failed to open mongo prealloc file %q: %v", filename, err)
		}
		defer f.Close()
		for total := 0; total < 1024*1024; {
			n, err := f.Write(zeroes)
			if err != nil {
				return fmt.Errorf("failed to write to mongo prealloc file %q: %v", filename, err)
			}
			total += n
		}
	}
	return nil
}

func sslKeyPath(dataDir string) string {
	return filepath.Join(dataDir, "server.pem")
}

func sharedSecretPath(dataDir string) string {
	return filepath.Join(dataDir, SharedSecretFile)
}

// mongoUpstartService returns the upstart config for the mongo state service.
//
func mongoUpstartService(namespace, dataDir, dbDir string, port int) (*upstart.Conf, error) {
	svc := upstart.NewService(ServiceName(namespace))

	mongoPath, err := MongodPath()
	if err != nil {
		return nil, err
	}

	conf := &upstart.Conf{
		Service: *svc,
		Desc:    "juju state database",
		Limit: map[string]string{
			"nofile": fmt.Sprintf("%d %d", maxFiles, maxFiles),
			"nproc":  fmt.Sprintf("%d %d", maxProcs, maxProcs),
		},
		Cmd: mongoPath + " --auth" +
			" --dbpath=" + dbDir +
			" --sslOnNormalPorts" +
			" --sslPEMKeyFile " + utils.ShQuote(sslKeyPath(dataDir)) +
			" --sslPEMKeyPassword ignored" +
			" --bind_ip 0.0.0.0" +
			" --port " + fmt.Sprint(port) +
			" --noprealloc" +
			" --syslog" +
			" --smallfiles" +
			" --replSet " + ReplicaSetName +
			" --keyFile " + utils.ShQuote(sharedSecretPath(dataDir)),
	}
	return conf, nil
}

func aptGetInstallMongod() error {
	// Only Quantal requires the PPA.
	if version.Current.Series == "quantal" {
		if err := addAptRepository("ppa:juju/stable"); err != nil {
			return err
		}
	}
	cmds := utils.AptGetPreparePackages([]string{"mongodb-server"}, version.Current.Series)
	logger.Infof("installing mongodb-server")
	for _, cmd := range cmds {
		if err := utils.AptGetInstall(cmd...); err != nil {
			return err
		}
	}
	return nil
}

func addAptRepository(name string) error {
	// add-apt-repository requires python-software-properties
	cmds := utils.AptGetPreparePackages(
		[]string{"python-software-properties"},
		version.Current.Series,
	)
	logger.Infof("installing python-software-properties")
	for _, cmd := range cmds {
		if err := utils.AptGetInstall(cmd...); err != nil {
			return err
		}
	}

	logger.Infof("adding apt repository %q", name)
	cmd := exec.Command("add-apt-repository", "-y", name)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("cannot add apt repository: %v (output %s)", err, bytes.TrimSpace(out))
	}
	return nil
}
<commit_msg>agent/mongo: minor fix<commit_after>package mongo

import (
	"bytes"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path"
	"path/filepath"

	"github.com/juju/loggo"
	"labix.org/v2/mgo"

	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/replicaset"
	"launchpad.net/juju-core/state/api/params"
	"launchpad.net/juju-core/upstart"
	"launchpad.net/juju-core/utils"
	"launchpad.net/juju-core/version"
)

const (
	maxFiles = 65000
	maxProcs = 20000

	serviceName = "juju-db"

	// SharedSecretFile is the name of the Mongo shared secret file
	// located within the Juju data directory.
	SharedSecretFile = "shared-secret"

	ReplicaSetName = "juju"
)

var (
	logger          = loggo.GetLogger("juju.agent.mongo")
	mongoConfigPath = "/etc/default/mongodb"

	// JujuMongodPath holds the default path to the juju-specific mongod.
	JujuMongodPath = "/usr/lib/juju/bin/mongod"

	upstartConfInstall          = (*upstart.Conf).Install
	upstartServiceStopAndRemove = (*upstart.Service).StopAndRemove
)

// WithAddresses represents an entity that has a set of
// addresses. e.g. a state Machine object
type WithAddresses interface {
	Addresses() []instance.Address
}

// IsMaster returns a boolean that represents whether the given
// machine's peer address is the primary mongo host for the replicaset
func IsMaster(session *mgo.Session, obj WithAddresses) (bool, error) {
	addrs := obj.Addresses()

	masterHostPort, err := replicaset.MasterHostPort(session)

	// If the replica set has not been configured, then we
	// can have only one master and the caller must
	// be that master.
	if err == replicaset.ErrMasterNotConfigured {
		return true, nil
	}

	if err != nil {
		return false, err
	}

	masterAddr, _, err := net.SplitHostPort(masterHostPort)
	if err != nil {
		return false, err
	}

	machinePeerAddr := SelectPeerAddress(addrs)
	return machinePeerAddr == masterAddr, nil
}

// SelectPeerAddress returns the address to use as the
// mongo replica set peer address by selecting it from the given addresses.
func SelectPeerAddress(addrs []instance.Address) string {
	return instance.SelectInternalAddress(addrs, false)
}

// SelectPeerHostPort returns the HostPort to use as the
// mongo replica set peer by selecting it from the given hostPorts.
func SelectPeerHostPort(hostPorts []instance.HostPort) string {
	return instance.SelectInternalHostPort(hostPorts, false)
}

// GenerateSharedSecret generates a pseudo-random shared secret (keyfile)
// for use with Mongo replica sets.
func GenerateSharedSecret() (string, error) {
	// "A key’s length must be between 6 and 1024 characters and may
	// only contain characters in the base64 set."
	//   -- http://docs.mongodb.org/manual/tutorial/generate-key-file/
	buf := make([]byte, base64.StdEncoding.DecodedLen(1024))
	if _, err := rand.Read(buf); err != nil {
		return "", fmt.Errorf("cannot read random secret: %v", err)
	}
	return base64.StdEncoding.EncodeToString(buf), nil
}

// MongoPath returns the executable path to be used to run mongod on this
// machine. If the juju-bundled version of mongo exists, it will return that
// path, otherwise it will return the command to run mongod from the path.
func MongodPath() (string, error) {
	if _, err := os.Stat(JujuMongodPath); err == nil {
		return JujuMongodPath, nil
	}

	path, err := exec.LookPath("mongod")
	if err != nil {
		return "", err
	}
	return path, nil
}

// RemoveService removes the mongoDB upstart service from this machine.
func RemoveService(namespace string) error {
	svc := upstart.NewService(ServiceName(namespace))
	return upstartServiceStopAndRemove(svc)
}

// EnsureMongoServer ensures that the correct mongo upstart script is installed
// and running.
//
// This method will remove old versions of the mongo upstart script as necessary
// before installing the new version.
//
// The namespace is a unique identifier to prevent multiple instances of mongo
// on this machine from colliding. This should be empty unless using
// the local provider.
func EnsureMongoServer(dataDir string, namespace string, info params.StateServingInfo) error {
	logger.Infof("Ensuring mongo server is running; dataDir %s; port %d", dataDir, info.StatePort)
	dbDir := filepath.Join(dataDir, "db")

	if err := os.MkdirAll(dbDir, 0700); err != nil {
		return fmt.Errorf("cannot create mongo dbdir: %v", err)
	}

	certKey := info.Cert + "\n" + info.PrivateKey
	err := utils.AtomicWriteFile(sslKeyPath(dataDir), []byte(certKey), 0600)
	if err != nil {
		return fmt.Errorf("cannot write SSL key: %v", err)
	}

	err = utils.AtomicWriteFile(sharedSecretPath(dataDir), []byte(info.SharedSecret), 0600)
	if err != nil {
		return fmt.Errorf("cannot write mongod shared secret: %v", err)
	}

	// Disable the default mongodb installed by the mongodb-server package.
	// Only do this if the file doesn't exist already, so users can run
	// their own mongodb server if they wish to.
	if _, err := os.Stat(mongoConfigPath); os.IsNotExist(err) {
		err = ioutil.WriteFile(
			mongoConfigPath,
			[]byte("ENABLE_MONGODB=no"),
			0644,
		)
		if err != nil {
			return err
		}
	}

	if err := aptGetInstallMongod(); err != nil {
		return fmt.Errorf("cannot install mongod: %v", err)
	}

	upstartConf, err := mongoUpstartService(namespace, dataDir, dbDir, info.StatePort)
	if err != nil {
		return err
	}

	// TODO(natefinch) 2014-04-12 https://launchpad.net/bugs/1306902
	// remove this once we support upgrading to HA
	if upstartConf.Installed() {
		return nil
	}

	if err := makeJournalDirs(dbDir); err != nil {
		return fmt.Errorf("Error creating journal directories: %v", err)
	}
	return upstartConfInstall(upstartConf)
}

// ServiceName returns the name of the upstart service config for mongo using
// the given namespace.
func ServiceName(namespace string) string {
	if namespace != "" {
		return fmt.Sprintf("%s-%s", serviceName, namespace)
	}
	return serviceName
}

func makeJournalDirs(dataDir string) error {
	journalDir := path.Join(dataDir, "journal")

	if err := os.MkdirAll(journalDir, 0700); err != nil {
		logger.Errorf("failed to make mongo journal dir %s: %v", journalDir, err)
		return err
	}

	// manually create the prealloc files, since otherwise they get created as 100M files.
	zeroes := make([]byte, 64*1024) // should be enough for anyone
	for x := 0; x < 3; x++ {
		name := fmt.Sprintf("prealloc.%d", x)
		filename := filepath.Join(journalDir, name)
		f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0700)
		// TODO(jam) 2014-04-12 https://launchpad.net/bugs/1306902
		// When we support upgrading Mongo into Replica mode, we should
		// start rewriting the upstart config
		if os.IsExist(err) {
			// already exists, don't overwrite
			continue
		}
		if err != nil {
			return fmt.Errorf("failed to open mongo prealloc file %q: %v", filename, err)
		}
		defer f.Close()
		for total := 0; total < 1024*1024; {
			n, err := f.Write(zeroes)
			if err != nil {
				return fmt.Errorf("failed to write to mongo prealloc file %q: %v", filename, err)
			}
			total += n
		}
	}
	return nil
}

func sslKeyPath(dataDir string) string {
	return filepath.Join(dataDir, "server.pem")
}

func sharedSecretPath(dataDir string) string {
	return filepath.Join(dataDir, SharedSecretFile)
}

// mongoUpstartService returns the upstart config for the mongo state service.
//
func mongoUpstartService(namespace, dataDir, dbDir string, port int) (*upstart.Conf, error) {
	svc := upstart.NewService(ServiceName(namespace))

	mongoPath, err := MongodPath()
	if err != nil {
		return nil, err
	}

	conf := &upstart.Conf{
		Service: *svc,
		Desc:    "juju state database",
		Limit: map[string]string{
			"nofile": fmt.Sprintf("%d %d", maxFiles, maxFiles),
			"nproc":  fmt.Sprintf("%d %d", maxProcs, maxProcs),
		},
		Cmd: mongoPath + " --auth" +
			" --dbpath=" + dbDir +
			" --sslOnNormalPorts" +
			" --sslPEMKeyFile " + utils.ShQuote(sslKeyPath(dataDir)) +
			" --sslPEMKeyPassword ignored" +
			" --bind_ip 0.0.0.0" +
			" --port " + fmt.Sprint(port) +
			" --noprealloc" +
			" --syslog" +
			" --smallfiles" +
			" --replSet " + ReplicaSetName +
			" --keyFile " + utils.ShQuote(sharedSecretPath(dataDir)),
	}
	return conf, nil
}

func aptGetInstallMongod() error {
	// Only Quantal requires the PPA.
	if version.Current.Series == "quantal" {
		if err := addAptRepository("ppa:juju/stable"); err != nil {
			return err
		}
	}
	cmds := utils.AptGetPreparePackages([]string{"mongodb-server"}, version.Current.Series)
	logger.Infof("installing mongodb-server")
	for _, cmd := range cmds {
		if err := utils.AptGetInstall(cmd...); err != nil {
			return err
		}
	}
	return nil
}

func addAptRepository(name string) error {
	// add-apt-repository requires python-software-properties
	cmds := utils.AptGetPreparePackages(
		[]string{"python-software-properties"},
		version.Current.Series,
	)
	logger.Infof("installing python-software-properties")
	for _, cmd := range cmds {
		if err := utils.AptGetInstall(cmd...); err != nil {
			return err
		}
	}

	logger.Infof("adding apt repository %q", name)
	cmd := exec.Command("add-apt-repository", "-y", name)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("cannot add apt repository: %v (output %s)", err, bytes.TrimSpace(out))
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) Mainflux
// SPDX-License-Identifier: Apache-2.0

package postgres

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"

	"github.com/gofrs/uuid"
	"github.com/lib/pq"
	"github.com/mainflux/mainflux/pkg/errors"
	"github.com/mainflux/mainflux/users"
)

var (
	errDeleteGroupDB = errors.New("delete group failed")
	errSelectDb      = errors.New("select group from db error")

	errFK         = "foreign_key_violation"
	errInvalid    = "invalid_text_representation"
	errTruncation = "string_data_right_truncation"
)

var _ users.GroupRepository = (*groupRepository)(nil)

type groupRepository struct {
	db Database
}

// NewGroupRepo instantiates a PostgreSQL implementation of group
// repository.
func NewGroupRepo(db Database) users.GroupRepository {
	return &groupRepository{
		db: db,
	}
}

func (gr groupRepository) Save(ctx context.Context, group users.Group) (users.Group, error) {
	var id string
	q := `INSERT INTO groups (name, description, id, owner_id, parent_id, metadata) VALUES (:name, :description, :id, :owner_id, :parent_id, :metadata) RETURNING id`
	if group.ParentID == "" {
		q = `INSERT INTO groups (name, description, id, owner_id, metadata) VALUES (:name, :description, :id, :owner_id, :metadata) RETURNING id`
	}

	dbu, err := toDBGroup(group)
	if err != nil {
		return users.Group{}, err
	}

	row, err := gr.db.NamedQueryContext(ctx, q, dbu)
	if err != nil {
		pqErr, ok := err.(*pq.Error)
		if ok {
			switch pqErr.Code.Name() {
			case errInvalid, errTruncation:
				return users.Group{}, errors.Wrap(users.ErrMalformedEntity, err)
			case errDuplicate:
				return users.Group{}, errors.Wrap(users.ErrGroupConflict, err)
			}
		}

		return users.Group{}, errors.Wrap(users.ErrCreateGroup, err)
	}

	defer row.Close()
	row.Next()
	if err := row.Scan(&id); err != nil {
		return users.Group{}, err
	}
	group.ID = id
	return group, nil
}

func (gr groupRepository) Update(ctx context.Context, group users.Group) error {
	q := `UPDATE groups SET(name, description, metadata) VALUES (:name, :description, :metadata) WHERE id = :id`

	dbu, err := toDBGroup(group)
	if err != nil {
		return errors.Wrap(errUpdateDB, err)
	}

	if _, err := gr.db.NamedExecContext(ctx, q, dbu); err != nil {
		return errors.Wrap(errUpdateDB, err)
	}

	return nil
}

func (gr groupRepository) Delete(ctx context.Context, groupID string) error {
	qd := `DELETE FROM groups WHERE id = :id`
	dbg, err := toDBGroup(users.Group{ID: groupID})
	if err != nil {
		return errors.Wrap(errUpdateDB, err)
	}

	res, err := gr.db.NamedExecContext(ctx, qd, dbg)
	if err != nil {
		return errors.Wrap(errDeleteGroupDB, err)
	}

	cnt, err := res.RowsAffected()
	if err != nil {
		return errors.Wrap(errDeleteGroupDB, err)
	}

	if cnt != 1 {
		return errors.Wrap(users.ErrDeleteGroupMissing, err)
	}
	return nil
}

func (gr groupRepository) RetrieveByID(ctx context.Context, id string) (users.Group, error) {
	q := `SELECT id, name, owner_id, parent_id, description, metadata FROM groups WHERE id = $1`
	dbu := dbGroup{
		ID: id,
	}

	if err := gr.db.QueryRowxContext(ctx, q, id).StructScan(&dbu); err != nil {
		if err == sql.ErrNoRows {
			return users.Group{}, errors.Wrap(users.ErrNotFound, err)

		}
		return users.Group{}, errors.Wrap(errRetrieveDB, err)
	}

	return toGroup(dbu), nil
}

func (gr groupRepository) RetrieveByName(ctx context.Context, name string) (users.Group, error) {
	q := `SELECT id, name, description, metadata FROM groups WHERE name = $1`

	dbu := dbGroup{
		Name: name,
	}

	if err := gr.db.QueryRowxContext(ctx, q, name).StructScan(&dbu); err != nil {
		if err == sql.ErrNoRows {
			return users.Group{}, errors.Wrap(users.ErrNotFound, err)

		}
		return users.Group{}, errors.Wrap(errRetrieveDB, err)
	}

	group := toGroup(dbu)
	return group, nil
}

func (gr groupRepository) RetrieveAllWithAncestors(ctx context.Context, groupID string, offset, limit uint64, gm users.Metadata) (users.GroupPage, error) {
	_, mq, err := getGroupsMetadataQuery(gm)
	if err != nil {
		return users.GroupPage{}, errors.Wrap(errRetrieveDB, err)
	}

	q := fmt.Sprintf(`WITH RECURSIVE subordinates AS (
						SELECT id, owner_id, parent_id, name, description, metadata
						FROM groups
						WHERE id = :id
						UNION
							SELECT groups.id, groups.owner_id, groups.parent_id, groups.name, groups.description, groups.metadata
							FROM groups
							INNER JOIN subordinates s ON s.id = groups.parent_id %s
					) SELECT * FROM subordinates ORDER BY id LIMIT :limit OFFSET :offset`, mq)

	dbPage, err := toDBGroupPage("", groupID, offset, limit, gm)
	if err != nil {
		return users.GroupPage{}, errors.Wrap(errSelectDb, err)
	}

	rows, err := gr.db.NamedQueryContext(ctx, q, dbPage)
	if err != nil {
		return users.GroupPage{}, errors.Wrap(errSelectDb, err)
	}

	defer rows.Close()

	var items []users.Group
	for rows.Next() {
		dbgr := dbGroup{}
		if err := rows.StructScan(&dbgr); err != nil {
			return users.GroupPage{}, errors.Wrap(errSelectDb, err)
		}
		gr := toGroup(dbgr)
		if err != nil {
			return users.GroupPage{}, err
		}
		items = append(items, gr)
	}

	cq := fmt.Sprintf(`WITH RECURSIVE subordinates AS (
						SELECT id, owner_id, parent_id, name, description, metadata
						FROM groups
						WHERE id = :id
						UNION
						SELECT groups.id, groups.owner_id, groups.parent_id, groups.name, groups.description, groups.metadata
						FROM groups
						INNER JOIN subordinates s ON s.id = groups.parent_id %s
					) SELECT COUNT(*) FROM subordinates`, mq)

	total, err := total(ctx, gr.db, cq, dbPage)
	if err != nil {
		return users.GroupPage{}, errors.Wrap(errSelectDb, err)
	}

	page := users.GroupPage{
		Groups: items,
		PageMetadata: users.PageMetadata{
			Total:  total,
			Offset: offset,
			Limit:  limit,
		},
	}

	return page, nil
}

func (gr groupRepository) Memberships(ctx context.Context, userID string, offset, limit uint64, gm users.Metadata) (users.GroupPage, error) {
	m, mq, err := getGroupsMetadataQuery(gm)
	if err != nil {
		return users.GroupPage{}, errors.Wrap(errRetrieveDB, err)
	}

	q := fmt.Sprintf(`SELECT g.id, g.owner_id, g.parent_id, g.name, g.description, g.metadata
					  FROM group_relations gr, groups g
					  WHERE gr.group_id = g.id and gr.user_id = :userID
		  			  %s ORDER BY id LIMIT :limit OFFSET :offset;`, mq)

	params := map[string]interface{}{
		"userID":   userID,
		"limit":    limit,
		"offset":   offset,
		"metadata": m,
	}

	rows, err := gr.db.NamedQueryContext(ctx, q, params)
	if err != nil {
		return users.GroupPage{}, errors.Wrap(errSelectDb, err)
	}
	defer rows.Close()

	var items []users.Group
	for rows.Next() {
		dbgr := dbGroup{}
		if err := rows.StructScan(&dbgr); err != nil {
			return users.GroupPage{}, errors.Wrap(errSelectDb, err)
		}
		gr := toGroup(dbgr)
		if err != nil {
			return users.GroupPage{}, err
		}
		items = append(items, gr)
	}

	cq := fmt.Sprintf(`SELECT COUNT(*)
					   FROM group_relations gr, groups g
					   WHERE gr.group_id = g.id and gr.user_id = :userID %s;`, mq)

	total, err := total(ctx, gr.db, cq, params)
	if err != nil {
		return users.GroupPage{}, errors.Wrap(errSelectDb, err)
	}

	page := users.GroupPage{
		Groups: items,
		PageMetadata: users.PageMetadata{
			Total:  total,
			Offset: offset,
			Limit:  limit,
		},
	}

	return page, nil
}

func (gr groupRepository) Assign(ctx context.Context, userID, groupID string) error {
	dbr, err := toDBGroupRelation(userID, groupID)
	if err != nil {
		return errors.Wrap(users.ErrAssignUserToGroup, err)
	}

	qIns := `INSERT INTO group_relations (group_id, user_id) VALUES (:group_id, :user_id)`
	_, err = gr.db.NamedQueryContext(ctx, qIns, dbr)
	if err != nil {
		pqErr, ok := err.(*pq.Error)
		if ok {
			switch pqErr.Code.Name() {
			case errInvalid, errTruncation:
				return errors.Wrap(users.ErrMalformedEntity, err)
			case errDuplicate:
				return errors.Wrap(users.ErrGroupConflict, err)
			case errFK:
				return errors.Wrap(users.ErrNotFound, err)
			}
		}
		return errors.Wrap(users.ErrAssignUserToGroup, err)
	}

	return nil
}

func (gr groupRepository) Unassign(ctx context.Context, userID, groupID string) error {
	q := `DELETE FROM group_relations WHERE user_id = :user_id AND group_id = :group_id`
	dbr, err := toDBGroupRelation(userID, groupID)
	if err != nil {
		return errors.Wrap(users.ErrNotFound, err)
	}
	if _, err := gr.db.NamedExecContext(ctx, q, dbr); err != nil {
		return errors.Wrap(users.ErrConflict, err)
	}
	return nil
}

type dbGroup struct {
	ID          string        `db:"id"`
	Name        string        `db:"name"`
	OwnerID     uuid.NullUUID `db:"owner_id"`
	ParentID    uuid.NullUUID `db:"parent_id"`
	Description string        `db:"description"`
	Metadata    dbMetadata    `db:"metadata"`
}

type dbGroupPage struct {
	ID       uuid.NullUUID `db:"id"`
	OwnerID  uuid.NullUUID `db:"owner_id"`
	ParentID uuid.NullUUID `db:"parent_id"`
	Metadata dbMetadata    `db:"metadata"`
	Limit    uint64
	Offset   uint64
	Size     uint64
}

func toUUID(id string) (uuid.NullUUID, error) {
	var parentID uuid.NullUUID
	if err := parentID.Scan(id); err != nil {
		if id != "" {
			return parentID, err
		}
		if err := parentID.Scan(nil); err != nil {
			return parentID, err
		}
	}
	return parentID, nil
}

func toDBGroup(g users.Group) (dbGroup, error) {
	parentID := ""
	if g.ParentID != "" {
		parentID = g.ParentID
	}
	parent, err := toUUID(parentID)
	if err != nil {
		return dbGroup{}, err
	}
	owner, err := toUUID(g.OwnerID)
	if err != nil {
		return dbGroup{}, err
	}

	return dbGroup{
		ID:          g.ID,
		Name:        g.Name,
		ParentID:    parent,
		OwnerID:     owner,
		Description: g.Description,
		Metadata:    g.Metadata,
	}, nil
}

func toDBGroupPage(ownerID, groupID string, offset, limit uint64, metadata users.Metadata) (dbGroupPage, error) {
	owner, err := toUUID(ownerID)
	if err != nil {
		return dbGroupPage{}, err
	}
	group, err := toUUID(groupID)
	if err != nil {
		return dbGroupPage{}, err
	}
	if err != nil {
		return dbGroupPage{}, err
	}
	return dbGroupPage{
		ID:       group,
		Metadata: dbMetadata(metadata),
		OwnerID:  owner,
		Offset:   offset,
		Limit:    limit,
	}, nil
}

func toGroup(dbu dbGroup) users.Group {
	return users.Group{
		ID:          dbu.ID,
		Name:        dbu.Name,
		ParentID:    dbu.ParentID.UUID.String(),
		OwnerID:     dbu.OwnerID.UUID.String(),
		Description: dbu.Description,
		Metadata:    dbu.Metadata,
	}
}

type dbGroupRelation struct {
	Group uuid.UUID `db:"group_id"`
	User  uuid.UUID `db:"user_id"`
}

func toDBGroupRelation(userID, groupID string) (dbGroupRelation, error) {
	group, err := uuid.FromString(groupID)
	if err != nil {
		return dbGroupRelation{}, err
	}
	user, err := uuid.FromString(userID)
	if err != nil {
		return dbGroupRelation{}, err
	}
	return dbGroupRelation{
		Group: group,
		User:  user,
	}, nil
}

func getGroupsMetadataQuery(m users.Metadata) ([]byte, string, error) {
	mq := ""
	mb := []byte("{}")
	if len(m) > 0 {
		mq = ` AND groups.metadata @> :metadata`

		b, err := json.Marshal(m)
		if err != nil {
			return nil, "", err
		}
		mb = b
	}
	return mb, mq, nil
}

func total(ctx context.Context, db Database, query string, params interface{}) (uint64, error) {
	rows, err := db.NamedQueryContext(ctx, query, params)
	if err != nil {
		return 0, err
	}
	defer rows.Close()
	total := uint64(0)
	if rows.Next() {
		if err := rows.Scan(&total); err != nil {
			return 0, err
		}
	}
	return total, nil
}
<commit_msg>NOISSUE - Fix group retrieval when parent id is not specified (#1247)<commit_after>// Copyright (c) Mainflux
// SPDX-License-Identifier: Apache-2.0

package postgres

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"

	"github.com/gofrs/uuid"
	"github.com/lib/pq"
	"github.com/mainflux/mainflux/pkg/errors"
	"github.com/mainflux/mainflux/users"
)

var (
	errDeleteGroupDB = errors.New("delete group failed")
	errSelectDb      = errors.New("select group from db error")

	errFK         = "foreign_key_violation"
	errInvalid    = "invalid_text_representation"
	errTruncation = "string_data_right_truncation"
)

var _ users.GroupRepository = (*groupRepository)(nil)

type groupRepository struct {
	db Database
}

// NewGroupRepo instantiates a PostgreSQL implementation of group
// repository.
func NewGroupRepo(db Database) users.GroupRepository {
	return &groupRepository{
		db: db,
	}
}

func (gr groupRepository) Save(ctx context.Context, group users.Group) (users.Group, error) {
	var id string
	q := `INSERT INTO groups (name, description, id, owner_id, parent_id, metadata) VALUES (:name, :description, :id, :owner_id, :parent_id, :metadata) RETURNING id`
	if group.ParentID == "" {
		q = `INSERT INTO groups (name, description, id, owner_id, metadata) VALUES (:name, :description, :id, :owner_id, :metadata) RETURNING id`
	}

	dbu, err := toDBGroup(group)
	if err != nil {
		return users.Group{}, err
	}

	row, err := gr.db.NamedQueryContext(ctx, q, dbu)
	if err != nil {
		pqErr, ok := err.(*pq.Error)
		if ok {
			switch pqErr.Code.Name() {
			case errInvalid, errTruncation:
				return users.Group{}, errors.Wrap(users.ErrMalformedEntity, err)
			case errDuplicate:
				return users.Group{}, errors.Wrap(users.ErrGroupConflict, err)
			}
		}

		return users.Group{}, errors.Wrap(users.ErrCreateGroup, err)
	}

	defer row.Close()
	row.Next()
	if err := row.Scan(&id); err != nil {
		return users.Group{}, err
	}
	group.ID = id
	return group, nil
}

func (gr groupRepository) Update(ctx context.Context, group users.Group) error {
	q := `UPDATE groups SET(name, description, metadata) VALUES (:name, :description, :metadata) WHERE id = :id`

	dbu, err := toDBGroup(group)
	if err != nil {
		return errors.Wrap(errUpdateDB, err)
	}

	if _, err := gr.db.NamedExecContext(ctx, q, dbu); err != nil {
		return errors.Wrap(errUpdateDB, err)
	}

	return nil
}

func (gr groupRepository) Delete(ctx context.Context, groupID string) error {
	qd := `DELETE FROM groups WHERE id = :id`
	dbg, err := toDBGroup(users.Group{ID: groupID})
	if err != nil {
		return errors.Wrap(errUpdateDB, err)
	}

	res, err := gr.db.NamedExecContext(ctx, qd, dbg)
	if err != nil {
		return errors.Wrap(errDeleteGroupDB, err)
	}

	cnt, err := res.RowsAffected()
	if err != nil {
		return errors.Wrap(errDeleteGroupDB, err)
	}

	if cnt != 1 {
		return errors.Wrap(users.ErrDeleteGroupMissing, err)
	}
	return nil
}

func (gr groupRepository) RetrieveByID(ctx context.Context, id string) (users.Group, error) {
	q := `SELECT id, name, owner_id, parent_id, description, metadata FROM groups WHERE id = $1`
	dbu := dbGroup{
		ID: id,
	}

	if err := gr.db.QueryRowxContext(ctx, q, id).StructScan(&dbu); err != nil {
		if err == sql.ErrNoRows {
			return users.Group{}, errors.Wrap(users.ErrNotFound, err)

		}
		return users.Group{}, errors.Wrap(errRetrieveDB, err)
	}

	return toGroup(dbu), nil
}

func (gr groupRepository) RetrieveByName(ctx context.Context, name string) (users.Group, error) {
	q := `SELECT id, name, description, metadata FROM groups WHERE name = $1`

	dbu := dbGroup{
		Name: name,
	}

	if err := gr.db.QueryRowxContext(ctx, q, name).StructScan(&dbu); err != nil {
		if err == sql.ErrNoRows {
			return users.Group{}, errors.Wrap(users.ErrNotFound, err)

		}
		return users.Group{}, errors.Wrap(errRetrieveDB, err)
	}

	group := toGroup(dbu)
	return group, nil
}

func (gr groupRepository) RetrieveAllWithAncestors(ctx context.Context, groupID string, offset, limit uint64, gm users.Metadata) (users.GroupPage, error) {
	_, mq, err := getGroupsMetadataQuery(gm)
	if err != nil {
		return users.GroupPage{}, errors.Wrap(errRetrieveDB, err)
	}
	if mq != "" {
		mq = fmt.Sprintf("WHERE %s", mq)
	}

	cq := fmt.Sprintf("SELECT COUNT(*) FROM groups %s", mq)
	sq := fmt.Sprintf("SELECT id, owner_id, parent_id, name, description, metadata FROM groups %s", mq)
	q := fmt.Sprintf("%s ORDER BY id LIMIT :limit OFFSET :offset", sq)

	if groupID != "" {
		sq = fmt.Sprintf(
			`WITH RECURSIVE subordinates AS (
				SELECT id, owner_id, parent_id, name, description, metadata
				FROM groups
				WHERE id = :id
				UNION
					SELECT groups.id, groups.owner_id, groups.parent_id, groups.name, groups.description, groups.metadata
					FROM groups
					INNER JOIN subordinates s ON s.id = groups.parent_id %s
			)`, mq)
		q = fmt.Sprintf("%s SELECT * FROM subordinates ORDER BY id LIMIT :limit OFFSET :offset", sq)
		cq = fmt.Sprintf("%s SELECT COUNT(*) FROM subordinates", sq)
	}

	dbPage, err := toDBGroupPage("", groupID, offset, limit, gm)
	if err != nil {
		return users.GroupPage{}, errors.Wrap(errSelectDb, err)
	}

	rows, err := gr.db.NamedQueryContext(ctx, q, dbPage)
	if err != nil {
		return users.GroupPage{}, errors.Wrap(errSelectDb, err)
	}
	defer rows.Close()

	var items []users.Group
	for rows.Next() {
		dbgr := dbGroup{}
		if err := rows.StructScan(&dbgr); err != nil {
			return users.GroupPage{}, errors.Wrap(errSelectDb, err)
		}
		gr := toGroup(dbgr)
		if err != nil {
			return users.GroupPage{}, err
		}
		items = append(items, gr)
	}

	total, err := total(ctx, gr.db, cq, dbPage)
	if err != nil {
		return users.GroupPage{}, errors.Wrap(errSelectDb, err)
	}

	page := users.GroupPage{
		Groups: items,
		PageMetadata: users.PageMetadata{
			Total:  total,
			Offset: offset,
			Limit:  limit,
		},
	}

	return page, nil
}

func (gr groupRepository) Memberships(ctx context.Context, userID string, offset, limit uint64, gm users.Metadata) (users.GroupPage, error) {
	m, mq, err := getGroupsMetadataQuery(gm)
	if err != nil {
		return users.GroupPage{}, errors.Wrap(errRetrieveDB, err)
	}

	if mq != "" {
		mq = fmt.Sprintf("AND %s", mq)
	}
	q := fmt.Sprintf(`SELECT g.id, g.owner_id, g.parent_id, g.name, g.description, g.metadata
					  FROM group_relations gr, groups g
					  WHERE gr.group_id = g.id and gr.user_id = :userID
		  			  %s ORDER BY id LIMIT :limit OFFSET :offset;`, mq)

	params := map[string]interface{}{
		"userID":   userID,
		"limit":    limit,
		"offset":   offset,
		"metadata": m,
	}

	rows, err := gr.db.NamedQueryContext(ctx, q, params)
	if err != nil {
		return users.GroupPage{}, errors.Wrap(errSelectDb, err)
	}
	defer rows.Close()

	var items []users.Group
	for rows.Next() {
		dbgr := dbGroup{}
		if err := rows.StructScan(&dbgr); err != nil {
			return users.GroupPage{}, errors.Wrap(errSelectDb, err)
		}
		gr := toGroup(dbgr)
		if err != nil {
			return users.GroupPage{}, err
		}
		items = append(items, gr)
	}

	cq := fmt.Sprintf(`SELECT COUNT(*)
					   FROM group_relations gr, groups g
					   WHERE gr.group_id = g.id and gr.user_id = :userID %s;`, mq)

	total, err := total(ctx, gr.db, cq, params)
	if err != nil {
		return users.GroupPage{}, errors.Wrap(errSelectDb, err)
	}

	page := users.GroupPage{
		Groups: items,
		PageMetadata: users.PageMetadata{
			Total:  total,
			Offset: offset,
			Limit:  limit,
		},
	}

	return page, nil
}

func (gr groupRepository) Assign(ctx context.Context, userID, groupID string) error {
	dbr, err := toDBGroupRelation(userID, groupID)
	if err != nil {
		return errors.Wrap(users.ErrAssignUserToGroup, err)
	}

	qIns := `INSERT INTO group_relations (group_id, user_id) VALUES (:group_id, :user_id)`
	_, err = gr.db.NamedQueryContext(ctx, qIns, dbr)
	if err != nil {
		pqErr, ok := err.(*pq.Error)
		if ok {
			switch pqErr.Code.Name() {
			case errInvalid, errTruncation:
				return errors.Wrap(users.ErrMalformedEntity, err)
			case errDuplicate:
				return errors.Wrap(users.ErrGroupConflict, err)
			case errFK:
				return errors.Wrap(users.ErrNotFound, err)
			}
		}
		return errors.Wrap(users.ErrAssignUserToGroup, err)
	}

	return nil
}

func (gr groupRepository) Unassign(ctx context.Context, userID, groupID string) error {
	q := `DELETE FROM group_relations WHERE user_id = :user_id AND group_id = :group_id`
	dbr, err := toDBGroupRelation(userID, groupID)
	if err != nil {
		return errors.Wrap(users.ErrNotFound, err)
	}
	if _, err := gr.db.NamedExecContext(ctx, q, dbr); err != nil {
		return errors.Wrap(users.ErrConflict, err)
	}
	return nil
}

type dbGroup struct {
	ID          string        `db:"id"`
	Name        string        `db:"name"`
	OwnerID     uuid.NullUUID `db:"owner_id"`
	ParentID    uuid.NullUUID `db:"parent_id"`
	Description string        `db:"description"`
	Metadata    dbMetadata    `db:"metadata"`
}

type dbGroupPage struct {
	ID       uuid.NullUUID `db:"id"`
	OwnerID  uuid.NullUUID `db:"owner_id"`
	ParentID uuid.NullUUID `db:"parent_id"`
	Metadata dbMetadata    `db:"metadata"`
	Limit    uint64
	Offset   uint64
	Size     uint64
}

func toUUID(id string) (uuid.NullUUID, error) {
	var parentID uuid.NullUUID
	if err := parentID.Scan(id); err != nil {
		if id != "" {
			return parentID, err
		}
		if err := parentID.Scan(nil); err != nil {
			return parentID, err
		}
	}
	return parentID, nil
}

func toDBGroup(g users.Group) (dbGroup, error) {
	parentID := ""
	if g.ParentID != "" {
		parentID = g.ParentID
	}
	parent, err := toUUID(parentID)
	if err != nil {
		return dbGroup{}, err
	}
	owner, err := toUUID(g.OwnerID)
	if err != nil {
		return dbGroup{}, err
	}

	return dbGroup{
		ID:          g.ID,
		Name:        g.Name,
		ParentID:    parent,
		OwnerID:     owner,
		Description: g.Description,
		Metadata:    g.Metadata,
	}, nil
}

func toDBGroupPage(ownerID, groupID string, offset, limit uint64, metadata users.Metadata) (dbGroupPage, error) {
	owner, err := toUUID(ownerID)
	if err != nil {
		return dbGroupPage{}, err
	}
	group, err := toUUID(groupID)
	if err != nil {
		return dbGroupPage{}, err
	}
	if err != nil {
		return dbGroupPage{}, err
	}
	return dbGroupPage{
		ID:       group,
		Metadata: dbMetadata(metadata),
		OwnerID:  owner,
		Offset:   offset,
		Limit:    limit,
	}, nil
}

func toGroup(dbu dbGroup) users.Group {
	return users.Group{
		ID:          dbu.ID,
		Name:        dbu.Name,
		ParentID:    dbu.ParentID.UUID.String(),
		OwnerID:     dbu.OwnerID.UUID.String(),
		Description: dbu.Description,
		Metadata:    dbu.Metadata,
	}
}

type dbGroupRelation struct {
	Group uuid.UUID `db:"group_id"`
	User  uuid.UUID `db:"user_id"`
}

func toDBGroupRelation(userID, groupID string) (dbGroupRelation, error) {
	group, err := uuid.FromString(groupID)
	if err != nil {
		return dbGroupRelation{}, err
	}
	user, err := uuid.FromString(userID)
	if err != nil {
		return dbGroupRelation{}, err
	}
	return dbGroupRelation{
		Group: group,
		User:  user,
	}, nil
}

func getGroupsMetadataQuery(m users.Metadata) ([]byte, string, error) {
	mq := ""
	mb := []byte("{}")
	if len(m) > 0 {
		mq = `groups.metadata @> :metadata`

		b, err := json.Marshal(m)
		if err != nil {
			return nil, "", err
		}
		mb = b
	}
	return mb, mq, nil
}

func total(ctx context.Context, db Database, query string, params interface{}) (uint64, error) {
	rows, err := db.NamedQueryContext(ctx, query, params)
	if err != nil {
		return 0, err
	}
	defer rows.Close()
	total := uint64(0)
	if rows.Next() {
		if err := rows.Scan(&total); err != nil {
			return 0, err
		}
	}
	return total, nil
}
<|endoftext|>