<commit_before>// SPDX-License-Identifier: MIT

package server

import (
	"encoding/json"
	"encoding/xml"
	"sort"
	"testing"

	"github.com/issue9/assert/v2"
	"github.com/issue9/localeutil"
	"golang.org/x/text/language"
	"google.golang.org/protobuf/proto"
	"gopkg.in/yaml.v2"

	"github.com/issue9/web/serialization/form"
	"github.com/issue9/web/server/testdata"
)

var (
	_ BuildResultFunc  = DefaultResultBuilder
	_ form.Marshaler   = &defaultResult{}
	_ form.Unmarshaler = &defaultResult{}
	_ proto.Message    = &defaultResult{}
)

var (
	mimetypeResult = &defaultResult{
		Code:    "400",
		Message: "400",
		Fields: []*fieldDetail{
			{
				Name:    "field1",
				Message: []string{"message1", "message2"},
			},
			{
				Name:    "field2",
				Message: []string{"message2"},
			},
		},
	}

	simpleMimetypeResult = &defaultResult{
		Code:    "400",
		Message: "400",
	}
)

func TestDefaultResult(t *testing.T) {
	a := assert.New(t, false)

	rslt := DefaultResultBuilder(500, "50001", "error message")
	a.False(rslt.HasFields()).
		Equal(rslt.Status(), 500)

	rslt.Add("f1", "f1 msg1")
	rslt.Add("f1", "f1 msg2")
	a.True(rslt.HasFields())
	r, ok := rslt.(*defaultResult)
	a.True(ok).Equal(2, len(r.Fields[0].Message))

	rslt.Set("f1", "f1 msg")
	a.True(rslt.HasFields())
	r, ok = rslt.(*defaultResult)
	a.True(ok).Equal(1, len(r.Fields[0].Message))

	rslt = DefaultResultBuilder(400, "40001", "400")
	rslt.Set("f1", "f1 msg1")
	a.True(rslt.HasFields())
	r, ok = rslt.(*defaultResult)
	a.True(ok).
		Equal(1, len(r.Fields[0].Message)).
		Equal("f1 msg1", r.Fields[0].Message[0])

	rslt.Set("f1", "f1 msg2")
	a.True(rslt.HasFields())
	r, ok = rslt.(*defaultResult)
	a.True(ok).
		Equal(1, len(r.Fields[0].Message)).
		Equal("f1 msg2", r.Fields[0].Message[0])
}

func TestDefaultResultJSON(t *testing.T) {
	a := assert.New(t, false)

	// marshal mimetypeResult
	bs, err := json.Marshal(mimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `{"message":"400","code":"400","fields":[{"name":"field1","message":["message1","message2"]},{"name":"field2","message":["message2"]}]}`)

	// unmarshal mimetypeResult
	obj := &defaultResult{}
	a.NotError(json.Unmarshal(bs, obj))
	a.Equal(obj, mimetypeResult)

	// marshal simpleMimetypesResult
	bs, err = json.Marshal(simpleMimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `{"message":"400","code":"400"}`)

	// unmarshal simpleMimetypesResult
	obj = &defaultResult{}
	a.NotError(json.Unmarshal(bs, obj))
	a.Equal(obj, simpleMimetypeResult)
}

func TestDefaultResultXML(t *testing.T) {
	a := assert.New(t, false)

	// marshal mimetypeResult
	bs, err := xml.Marshal(mimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `<result code="400"><message>400</message><field name="field1"><message>message1</message><message>message2</message></field><field name="field2"><message>message2</message></field></result>`)

	// unmarshal mimetypeResult
	obj := &defaultResult{}
	a.NotError(xml.Unmarshal(bs, obj))
	a.Equal(obj, mimetypeResult)

	// marshal simpleMimetypesResult
	bs, err = xml.Marshal(simpleMimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `<result code="400"><message>400</message></result>`)

	// unmarshal simpleMimetypesResult
	obj = &defaultResult{}
	a.NotError(xml.Unmarshal(bs, obj))
	a.Equal(obj, simpleMimetypeResult)
}

func TestDefaultResultProtobuf(t *testing.T) {
	a := assert.New(t, false)

	// marshal mimetypeResult
	bs, err := proto.Marshal(mimetypeResult)
	a.NotError(err).NotNil(bs)

	// unmarshal mimetypeResult
	obj := &testdata.Result{}
	a.NotError(proto.Unmarshal(bs, obj))
	a.Equal(obj.Message, mimetypeResult.Message).
		Equal(obj.Code, mimetypeResult.Code).
		Equal(2, len(obj.Fields)).
		Equal(obj.Fields[0].Name, "field1").
		Equal(obj.Fields[0].Message, []string{"message1", "message2"})

	// marshal simpleMimetypesResult
	bs, err = proto.Marshal(simpleMimetypeResult)
	a.NotError(err).NotNil(bs)

	// unmarshal simpleMimetypesResult
	obj = &testdata.Result{}
	a.NotError(proto.Unmarshal(bs, obj))
	a.Equal(obj.Message, simpleMimetypeResult.Message).
		Equal(obj.Code, simpleMimetypeResult.Code)
}

func TestDefaultResultYAML(t *testing.T) {
	a := assert.New(t, false)

	// marshal mimetypeResult
	bs, err := yaml.Marshal(mimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `message: "400"
code: "400"
fields:
- name: field1
  message:
  - message1
  - message2
- name: field2
  message:
  - message2
`)

	// unmarshal mimetypeResult
	obj := &defaultResult{}
	a.NotError(yaml.Unmarshal(bs, obj))
	a.Equal(obj, mimetypeResult)

	// marshal simpleMimetypesResult
	bs, err = yaml.Marshal(simpleMimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `message: "400"
code: "400"
`)

	// unmarshal simpleMimetypesResult
	obj = &defaultResult{}
	a.NotError(yaml.Unmarshal(bs, obj))
	a.Equal(obj, simpleMimetypeResult)
}

func TestDefaultResultForm(t *testing.T) {
	a := assert.New(t, false)

	// marshal mimetypeResult
	bs, err := form.Marshal(mimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `code=400&fields.field1=message1&fields.field1=message2&fields.field2=message2&message=400`)

	// unmarshal mimetypeResult
	obj := &defaultResult{}
	a.NotError(form.Unmarshal(bs, obj))
	sort.SliceStable(obj.Fields, func(i, j int) bool { return obj.Fields[i].Name < obj.Fields[j].Name }) // 顺序一致才能相等
	a.Equal(obj, mimetypeResult)

	// marshal simpleMimetypesResult
	bs, err = form.Marshal(simpleMimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `code=400&message=400`)

	// unmarshal simpleMimetypesResult
	obj = &defaultResult{}
	a.NotError(form.Unmarshal(bs, obj))
	a.Equal(obj, simpleMimetypeResult)
}

func TestServer_Result(t *testing.T) {
	a := assert.New(t, false)
	srv := newServer(a, nil)

	srv.AddResult(400, "40000", localeutil.Phrase("lang")) // lang 有翻译

	// 能正常翻译错误信息
	rslt, ok := srv.Result(srv.Locale().Printer(language.SimplifiedChinese), "40000", nil).(*defaultResult)
	a.True(ok).NotNil(rslt)
	a.Equal(rslt.Message, "hans")

	// 采用 und
	rslt, ok = srv.Result(srv.Locale().Printer(language.Und), "40000", nil).(*defaultResult)
	a.True(ok).NotNil(rslt)
	a.Equal(rslt.Message, "und")

	// 不存在的本地化信息，采用默认的 und
	rslt, ok = srv.Result(srv.Locale().Printer(language.Afrikaans), "40000", nil).(*defaultResult)
	a.True(ok).NotNil(rslt)
	a.Equal(rslt.Message, "und")

	// 不存在
	a.Panic(func() { srv.Result(srv.Locale().Printer(language.Afrikaans), "400", nil) })
	a.Panic(func() { srv.Result(srv.Locale().Printer(language.Afrikaans), "50000", nil) })

	// with fields

	fields := map[string][]string{"f1": {"v1", "v2"}}

	// 能正常翻译错误信息
	rslt, ok = srv.Result(srv.Locale().Printer(language.SimplifiedChinese), "40000", fields).(*defaultResult)
	a.True(ok).NotNil(rslt)
	a.Equal(rslt.Message, "hans").
		Equal(rslt.Fields, []*fieldDetail{{Name: "f1", Message: []string{"v1", "v2"}}})

	// 采用 und
	rslt, ok = srv.Result(srv.Locale().Printer(language.Und), "40000", fields).(*defaultResult)
	a.True(ok).NotNil(rslt)
	a.Equal(rslt.Message, "und").
		Equal(rslt.Fields, []*fieldDetail{{Name: "f1", Message: []string{"v1", "v2"}}})
}

func TestServer_AddResult(t *testing.T) {
	a := assert.New(t, false)
	srv := newServer(a, &Options{Tag: language.SimplifiedChinese})

	a.NotPanic(func() {
		srv.AddResult(400, "1", localeutil.Phrase("1"))
		srv.AddResult(400, "100", localeutil.Phrase("100"))
	})

	msg, found := srv.resultMessages["1"]
	a.True(found).
		Equal(msg.status, 400)

	msg, found = srv.resultMessages["401"]
	a.False(found).Nil(msg)

	// 重复的 ID
	a.Panic(func() {
		srv.AddResult(400, "1", localeutil.Phrase("40010"))
	})
}

func TestServer_Results(t *testing.T) {
	a := assert.New(t, false)
	c := newServer(a, &Options{Tag: language.SimplifiedChinese})

	a.NotPanic(func() {
		c.AddResults(400, map[string]localeutil.LocaleStringer{"40010": localeutil.Phrase("lang")})
	})

	msg := c.Results(c.Locale().Printer(language.Und))
	a.Equal(msg["40010"], "und")

	msg = c.Results(c.Locale().Printer(language.SimplifiedChinese))
	a.Equal(msg["40010"], "hans")

	msg = c.Results(c.Locale().Printer(language.TraditionalChinese))
	a.Equal(msg["40010"], "hant")

	msg = c.Results(c.Locale().Printer(language.English))
	a.Equal(msg["40010"], "und")

	a.Panic(func() {
		c.Results(nil)
	})
}
<commit_msg>test(server): 采用 serialization/protobuf 用于测试<commit_after>// SPDX-License-Identifier: MIT

package server

import (
	"encoding/json"
	"encoding/xml"
	"sort"
	"testing"

	"github.com/issue9/assert/v2"
	"github.com/issue9/localeutil"
	"golang.org/x/text/language"
	"google.golang.org/protobuf/proto"
	"gopkg.in/yaml.v2"

	"github.com/issue9/web/serialization/form"
	"github.com/issue9/web/serialization/protobuf"
	"github.com/issue9/web/server/testdata"
)

var (
	_ BuildResultFunc  = DefaultResultBuilder
	_ form.Marshaler   = &defaultResult{}
	_ form.Unmarshaler = &defaultResult{}
	_ proto.Message    = &defaultResult{}
)

var (
	mimetypeResult = &defaultResult{
		Code:    "400",
		Message: "400",
		Fields: []*fieldDetail{
			{
				Name:    "field1",
				Message: []string{"message1", "message2"},
			},
			{
				Name:    "field2",
				Message: []string{"message2"},
			},
		},
	}

	simpleMimetypeResult = &defaultResult{
		Code:    "400",
		Message: "400",
	}
)

func TestDefaultResult(t *testing.T) {
	a := assert.New(t, false)

	rslt := DefaultResultBuilder(500, "50001", "error message")
	a.False(rslt.HasFields()).
		Equal(rslt.Status(), 500)

	rslt.Add("f1", "f1 msg1")
	rslt.Add("f1", "f1 msg2")
	a.True(rslt.HasFields())
	r, ok := rslt.(*defaultResult)
	a.True(ok).Equal(2, len(r.Fields[0].Message))

	rslt.Set("f1", "f1 msg")
	a.True(rslt.HasFields())
	r, ok = rslt.(*defaultResult)
	a.True(ok).Equal(1, len(r.Fields[0].Message))

	rslt = DefaultResultBuilder(400, "40001", "400")
	rslt.Set("f1", "f1 msg1")
	a.True(rslt.HasFields())
	r, ok = rslt.(*defaultResult)
	a.True(ok).
		Equal(1, len(r.Fields[0].Message)).
		Equal("f1 msg1", r.Fields[0].Message[0])

	rslt.Set("f1", "f1 msg2")
	a.True(rslt.HasFields())
	r, ok = rslt.(*defaultResult)
	a.True(ok).
		Equal(1, len(r.Fields[0].Message)).
		Equal("f1 msg2", r.Fields[0].Message[0])
}

func TestDefaultResultJSON(t *testing.T) {
	a := assert.New(t, false)

	// marshal mimetypeResult
	bs, err := json.Marshal(mimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `{"message":"400","code":"400","fields":[{"name":"field1","message":["message1","message2"]},{"name":"field2","message":["message2"]}]}`)

	// unmarshal mimetypeResult
	obj := &defaultResult{}
	a.NotError(json.Unmarshal(bs, obj))
	a.Equal(obj, mimetypeResult)

	// marshal simpleMimetypesResult
	bs, err = json.Marshal(simpleMimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `{"message":"400","code":"400"}`)

	// unmarshal simpleMimetypesResult
	obj = &defaultResult{}
	a.NotError(json.Unmarshal(bs, obj))
	a.Equal(obj, simpleMimetypeResult)
}

func TestDefaultResultXML(t *testing.T) {
	a := assert.New(t, false)

	// marshal mimetypeResult
	bs, err := xml.Marshal(mimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `<result code="400"><message>400</message><field name="field1"><message>message1</message><message>message2</message></field><field name="field2"><message>message2</message></field></result>`)

	// unmarshal mimetypeResult
	obj := &defaultResult{}
	a.NotError(xml.Unmarshal(bs, obj))
	a.Equal(obj, mimetypeResult)

	// marshal simpleMimetypesResult
	bs, err = xml.Marshal(simpleMimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `<result code="400"><message>400</message></result>`)

	// unmarshal simpleMimetypesResult
	obj = &defaultResult{}
	a.NotError(xml.Unmarshal(bs, obj))
	a.Equal(obj, simpleMimetypeResult)
}

func TestDefaultResultProtobuf(t *testing.T) {
	a := assert.New(t, false)

	// marshal mimetypeResult
	bs, err := protobuf.Marshal(mimetypeResult)
	a.NotError(err).NotNil(bs)

	// unmarshal mimetypeResult
	obj := &testdata.Result{}
	a.NotError(protobuf.Unmarshal(bs, obj))
	a.Equal(obj.Message, mimetypeResult.Message).
		Equal(obj.Code, mimetypeResult.Code).
		Equal(2, len(obj.Fields)).
		Equal(obj.Fields[0].Name, "field1").
		Equal(obj.Fields[0].Message, []string{"message1", "message2"})

	// marshal simpleMimetypesResult
	bs, err = protobuf.Marshal(simpleMimetypeResult)
	a.NotError(err).NotNil(bs)

	// unmarshal simpleMimetypesResult
	obj = &testdata.Result{}
	a.NotError(protobuf.Unmarshal(bs, obj))
	a.Equal(obj.Message, simpleMimetypeResult.Message).
		Equal(obj.Code, simpleMimetypeResult.Code)
}

func TestDefaultResultYAML(t *testing.T) {
	a := assert.New(t, false)

	// marshal mimetypeResult
	bs, err := yaml.Marshal(mimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `message: "400"
code: "400"
fields:
- name: field1
  message:
  - message1
  - message2
- name: field2
  message:
  - message2
`)

	// unmarshal mimetypeResult
	obj := &defaultResult{}
	a.NotError(yaml.Unmarshal(bs, obj))
	a.Equal(obj, mimetypeResult)

	// marshal simpleMimetypesResult
	bs, err = yaml.Marshal(simpleMimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `message: "400"
code: "400"
`)

	// unmarshal simpleMimetypesResult
	obj = &defaultResult{}
	a.NotError(yaml.Unmarshal(bs, obj))
	a.Equal(obj, simpleMimetypeResult)
}

func TestDefaultResultForm(t *testing.T) {
	a := assert.New(t, false)

	// marshal mimetypeResult
	bs, err := form.Marshal(mimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `code=400&fields.field1=message1&fields.field1=message2&fields.field2=message2&message=400`)

	// unmarshal mimetypeResult
	obj := &defaultResult{}
	a.NotError(form.Unmarshal(bs, obj))
	sort.SliceStable(obj.Fields, func(i, j int) bool { return obj.Fields[i].Name < obj.Fields[j].Name }) // 顺序一致才能相等
	a.Equal(obj, mimetypeResult)

	// marshal simpleMimetypesResult
	bs, err = form.Marshal(simpleMimetypeResult)
	a.NotError(err).NotNil(bs)
	a.Equal(string(bs), `code=400&message=400`)

	// unmarshal simpleMimetypesResult
	obj = &defaultResult{}
	a.NotError(form.Unmarshal(bs, obj))
	a.Equal(obj, simpleMimetypeResult)
}

func TestServer_Result(t *testing.T) {
	a := assert.New(t, false)
	srv := newServer(a, nil)

	srv.AddResult(400, "40000", localeutil.Phrase("lang")) // lang 有翻译

	// 能正常翻译错误信息
	rslt, ok := srv.Result(srv.Locale().Printer(language.SimplifiedChinese), "40000", nil).(*defaultResult)
	a.True(ok).NotNil(rslt)
	a.Equal(rslt.Message, "hans")

	// 采用 und
	rslt, ok = srv.Result(srv.Locale().Printer(language.Und), "40000", nil).(*defaultResult)
	a.True(ok).NotNil(rslt)
	a.Equal(rslt.Message, "und")

	// 不存在的本地化信息，采用默认的 und
	rslt, ok = srv.Result(srv.Locale().Printer(language.Afrikaans), "40000", nil).(*defaultResult)
	a.True(ok).NotNil(rslt)
	a.Equal(rslt.Message, "und")

	// 不存在
	a.Panic(func() { srv.Result(srv.Locale().Printer(language.Afrikaans), "400", nil) })
	a.Panic(func() { srv.Result(srv.Locale().Printer(language.Afrikaans), "50000", nil) })

	// with fields

	fields := map[string][]string{"f1": {"v1", "v2"}}

	// 能正常翻译错误信息
	rslt, ok = srv.Result(srv.Locale().Printer(language.SimplifiedChinese), "40000", fields).(*defaultResult)
	a.True(ok).NotNil(rslt)
	a.Equal(rslt.Message, "hans").
		Equal(rslt.Fields, []*fieldDetail{{Name: "f1", Message: []string{"v1", "v2"}}})

	// 采用 und
	rslt, ok = srv.Result(srv.Locale().Printer(language.Und), "40000", fields).(*defaultResult)
	a.True(ok).NotNil(rslt)
	a.Equal(rslt.Message, "und").
		Equal(rslt.Fields, []*fieldDetail{{Name: "f1", Message: []string{"v1", "v2"}}})
}

func TestServer_AddResult(t *testing.T) {
	a := assert.New(t, false)
	srv := newServer(a, &Options{Tag: language.SimplifiedChinese})

	a.NotPanic(func() {
		srv.AddResult(400, "1", localeutil.Phrase("1"))
		srv.AddResult(400, "100", localeutil.Phrase("100"))
	})

	msg, found := srv.resultMessages["1"]
	a.True(found).
		Equal(msg.status, 400)

	msg, found = srv.resultMessages["401"]
	a.False(found).Nil(msg)

	// 重复的 ID
	a.Panic(func() {
		srv.AddResult(400, "1", localeutil.Phrase("40010"))
	})
}

func TestServer_Results(t *testing.T) {
	a := assert.New(t, false)
	c := newServer(a, &Options{Tag: language.SimplifiedChinese})

	a.NotPanic(func() {
		c.AddResults(400, map[string]localeutil.LocaleStringer{"40010": localeutil.Phrase("lang")})
	})

	msg := c.Results(c.Locale().Printer(language.Und))
	a.Equal(msg["40010"], "und")

	msg = c.Results(c.Locale().Printer(language.SimplifiedChinese))
	a.Equal(msg["40010"], "hans")

	msg = c.Results(c.Locale().Printer(language.TraditionalChinese))
	a.Equal(msg["40010"], "hant")

	msg = c.Results(c.Locale().Printer(language.English))
	a.Equal(msg["40010"], "und")

	a.Panic(func() {
		c.Results(nil)
	})
}
<|endoftext|>


<commit_before>package server

import (
	"log"
	"testing"

	"os"

	"crypto/rand"

	"io"

	"net"

	"time"

	"github.com/pinfake/pes6go/client"
	"github.com/pinfake/pes6go/data/block"
	"github.com/pinfake/pes6go/storage"
)

const port = 19780

var s *Server

type emptyServer struct {
}

func (s emptyServer) Config() ServerConfig {
	return nil
}

func (s emptyServer) Storage() storage.Storage {
	return nil
}

func (s emptyServer) Handlers() map[uint16]Handler {
	return map[uint16]Handler{}
}

func NewEmptyServer() *Server {
	return NewServer(
		log.New(os.Stdout, "test: ", log.LstdFlags),
		emptyServer{},
	)
}

func init() {
	s = NewEmptyServer()
	go s.Serve(port)
}

func getRandom(size int) []byte {
	data := make([]byte, size)
	rand.Read(data)

	return data
}

func craftBlock(query uint16, size uint16, data []byte) *block.Block {
	b := block.Block{
		Header: block.Header{
			Query:    query,
			Size:     size,
			Sequence: 0,
			Hash:     [16]byte{},
		},
		Body: block.GenericBody{
			Data: data,
		},
	}

	return &b
}

func assertDisconnected(c *client.Client, t *testing.T) {
	defer c.Close()
	_, err := c.Read()
	if err == nil {
		t.Error("still connected: no error reading")
	} else {
		if err == io.EOF {
			return
		}
		if !err.(*net.OpError).Timeout() {
			return
		}
		t.Error("still connected")
	}
}

func connect(c *client.Client, t *testing.T) {
	err := c.Connect("localhost", port)
	if err != nil {
		t.Error("Error connecting: %s", err.Error())
	}
}

func TestConnect(t *testing.T) {
	t.Run("Should connect", func(t *testing.T) {
		c := client.NewClient()
		connect(c, t)
		c.Close()
	})
}

func TestSendInvalidData(t *testing.T) {
	c := client.NewClient()
	connect(c, t)
	c.Write([]byte{0x01, 0x02, 0x03})
	assertDisconnected(c, t)
	t.Run("Should be kicked out", func(t *testing.T) {
		assertDisconnected(c, t)
	})
}

func TestSendProperHeadLongerBody(t *testing.T) {
	t.Run("Shouldn't crash", func(t *testing.T) {
		b := craftBlock(0x3001, 10, getRandom(100))
		c := client.NewClient()
		connect(c, t)
		c.WriteBlock(b)
		c.Close()
	})
}

func TestSendProperHeadShorterBody(t *testing.T) {
	b := craftBlock(0x3001, 100, getRandom(10))
	c := client.NewClient()
	connect(c, t)
	c.WriteBlock(b)
	t.Run("Should be kicked out", func(t *testing.T) {
		assertDisconnected(c, t)
	})
}

func TestSendMoreThanReadBuffer(t *testing.T) {
	c := client.NewClient()
	connect(c, t)
	c.Write(getRandom(10000))
	t.Run("Should be kicked out", func(t *testing.T) {
		assertDisconnected(c, t)
	})
}

func TestSend1Megabyte(t *testing.T) {
	c := client.NewClient()
	connect(c, t)
	c.Write(getRandom(1 * 1024 * 1024))
	t.Run("Should be kicked out", func(t *testing.T) {
		assertDisconnected(c, t)
	})
}

func TestSendUnknownQuery(t *testing.T) {
	b := craftBlock(0x1234, 100, getRandom(100))
	c := client.NewClient()
	connect(c, t)
	c.WriteBlock(b)
	t.Run("Should be kicked out", func(t *testing.T) {
		assertDisconnected(c, t)
	})
}

func Test1KConnections(t *testing.T) {
	for i := 0; i < 1000; i++ {
		go func() {
			c := client.NewClient()
			connect(c, t)
			select {}
		}()
	}
	time.Sleep(1 * time.Second)
}
<commit_msg>Fixing tests.<commit_after>package server

import (
	"log"
	"testing"

	"os"

	"crypto/rand"

	"io"

	"net"

	"time"

	"github.com/pinfake/pes6go/client"
	"github.com/pinfake/pes6go/data/block"
	"github.com/pinfake/pes6go/storage"
)

const port = 19780

var s *Server

type emptyServer struct {
}

func (s emptyServer) Config() ServerConfig {
	return nil
}

func (s emptyServer) Storage() storage.Storage {
	return nil
}

func (s emptyServer) Handlers() map[uint16]Handler {
	return map[uint16]Handler{}
}

func (s emptyServer) Data() interface{} {
	return nil
}

func NewEmptyServer() *Server {
	return NewServer(
		log.New(os.Stdout, "test: ", log.LstdFlags),
		emptyServer{},
	)
}

func init() {
	s = NewEmptyServer()
	go s.Serve(port)
}

func getRandom(size int) []byte {
	data := make([]byte, size)
	rand.Read(data)

	return data
}

func craftBlock(query uint16, size uint16, data []byte) *block.Block {
	b := block.Block{
		Header: block.Header{
			Query:    query,
			Size:     size,
			Sequence: 0,
			Hash:     [16]byte{},
		},
		Body: block.GenericBody{
			Data: data,
		},
	}

	return &b
}

func assertDisconnected(c *client.Client, t *testing.T) {
	defer c.Close()
	_, err := c.Read()
	if err == nil {
		t.Error("still connected: no error reading")
	} else {
		if err == io.EOF {
			return
		}
		if !err.(*net.OpError).Timeout() {
			return
		}
		t.Error("still connected")
	}
}

func connect(c *client.Client, t *testing.T) {
	err := c.Connect("localhost", port)
	if err != nil {
		t.Error("Error connecting: %s", err.Error())
	}
}

func TestConnect(t *testing.T) {
	t.Run("Should connect", func(t *testing.T) {
		c := client.NewClient()
		connect(c, t)
		c.Close()
	})
}

func TestSendInvalidData(t *testing.T) {
	c := client.NewClient()
	connect(c, t)
	c.Write([]byte{0x01, 0x02, 0x03})
	assertDisconnected(c, t)
	t.Run("Should be kicked out", func(t *testing.T) {
		assertDisconnected(c, t)
	})
}

func TestSendProperHeadLongerBody(t *testing.T) {
	t.Run("Shouldn't crash", func(t *testing.T) {
		b := craftBlock(0x3001, 10, getRandom(100))
		c := client.NewClient()
		connect(c, t)
		c.WriteBlock(b)
		c.Close()
	})
}

func TestSendProperHeadShorterBody(t *testing.T) {
	b := craftBlock(0x3001, 100, getRandom(10))
	c := client.NewClient()
	connect(c, t)
	c.WriteBlock(b)
	t.Run("Should be kicked out", func(t *testing.T) {
		assertDisconnected(c, t)
	})
}

func TestSendMoreThanReadBuffer(t *testing.T) {
	c := client.NewClient()
	connect(c, t)
	c.Write(getRandom(10000))
	t.Run("Should be kicked out", func(t *testing.T) {
		assertDisconnected(c, t)
	})
}

func TestSend1Megabyte(t *testing.T) {
	c := client.NewClient()
	connect(c, t)
	c.Write(getRandom(1 * 1024 * 1024))
	t.Run("Should be kicked out", func(t *testing.T) {
		assertDisconnected(c, t)
	})
}

func TestSendUnknownQuery(t *testing.T) {
	b := craftBlock(0x1234, 100, getRandom(100))
	c := client.NewClient()
	connect(c, t)
	c.WriteBlock(b)
	t.Run("Should be kicked out", func(t *testing.T) {
		assertDisconnected(c, t)
	})
}

func Test1KConnections(t *testing.T) {
	for i := 0; i < 1000; i++ {
		go func() {
			c := client.NewClient()
			connect(c, t)
			select {}
		}()
	}
	time.Sleep(1 * time.Second)
}
<|endoftext|>


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

package server

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"go.chromium.org/luci/common/clock/testclock"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/logging/gkelogger"
	"go.chromium.org/luci/server/router"
	"go.chromium.org/luci/server/secrets"

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

func TestServer(t *testing.T) {
	t.Parallel()

	Convey("Works", t, func() {
		ctx, tc := testclock.UseTime(context.Background(), testclock.TestRecentTimeUTC)

		tmpSecret, err := tempSecret()
		So(err, ShouldBeNil)
		defer os.Remove(tmpSecret.Name())

		stdoutLogs := logsRecorder{}
		stderrLogs := logsRecorder{}

		srv := New(Options{
			Prod:           true,
			HTTPAddr:       "main_addr",
			AdminAddr:      "admin_addr",
			RootSecretPath: tmpSecret.Name(),

			testCtx:    ctx,
			testSeed:   1,
			testStdout: &stdoutLogs,
			testStderr: &stderrLogs,

			// Bind to auto-assigned ports.
			testListeners: map[string]net.Listener{
				"main_addr":  setupListener(),
				"admin_addr": setupListener(),
			},
		})

		mainPort := srv.opts.testListeners["main_addr"].Addr().(*net.TCPAddr).Port
		mainAddr := fmt.Sprintf("http://127.0.0.1:%d", mainPort)

		// Run the serving loop in parallel.
		serveErr := errorEvent{signal: make(chan struct{})}
		go func() { serveErr.Set(srv.ListenAndServe()) }()

		Reset(func() {
			srv.Shutdown()
			So(serveErr.Get(), ShouldBeNil)
		})

		Convey("Logging", func() {
			srv.Routes.GET("/test", router.MiddlewareChain{}, func(c *router.Context) {
				logging.Infof(c.Context, "Info log")
				tc.Add(time.Second)
				logging.Warningf(c.Context, "Warn log")
				c.Writer.WriteHeader(201)
				c.Writer.Write([]byte("Hello, world"))
			})

			resp, err := httpGet(mainAddr+"/test", &serveErr, map[string]string{
				"User-Agent":            "Test-user-agent",
				"X-Cloud-Trace-Context": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa/00001;trace=TRUE",
				"X-Forwarded-For":       "1.1.1.1,2.2.2.2,3.3.3.3",
			})
			So(err, ShouldBeNil)
			So(resp, ShouldEqual, "Hello, world")

			// Stderr log captures details about the request.
			So(stderrLogs.Last(1), ShouldResemble, []gkelogger.LogEntry{
				{
					Severity: "warning",
					Time:     "1454472307.7",
					TraceID:  "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
					RequestInfo: &gkelogger.RequestInfo{
						Method:       "GET",
						URL:          mainAddr + "/test",
						Status:       201,
						RequestSize:  "0",
						ResponseSize: "12", // len("Hello, world")
						UserAgent:    "Test-user-agent",
						RemoteIP:     "2.2.2.2",
						Latency:      "1.000000s",
					},
				},
			})
			// Stdout log captures individual log lines.
			So(stdoutLogs.Last(2), ShouldResemble, []gkelogger.LogEntry{
				{
					Severity: "info",
					Message:  "Info log",
					Time:     "1454472306.7",
					TraceID:  "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
					Operation: &gkelogger.Operation{
						ID: "9566c74d10037c4d7bbb0407d1e2c649",
					},
				},
				{
					Severity: "warning",
					Message:  "Warn log",
					Time:     "1454472307.7",
					TraceID:  "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
					Operation: &gkelogger.Operation{
						ID: "9566c74d10037c4d7bbb0407d1e2c649",
					},
				},
			})
		})

		Convey("Secrets", func() {
			srv.Routes.GET("/secret", router.MiddlewareChain{}, func(c *router.Context) {
				s, err := secrets.GetSecret(c.Context, "secret_name")
				if err != nil {
					c.Writer.WriteHeader(500)
				} else {
					c.Writer.Write([]byte(s.Current))
				}
			})
			resp, err := httpGet(mainAddr+"/secret", &serveErr, nil)
			So(err, ShouldBeNil)
			So(resp, ShouldNotBeEmpty)
		})
	})
}

func tempSecret() (out *os.File, err error) {
	var f *os.File
	defer func() {
		if f != nil && err != nil {
			os.Remove(f.Name())
		}
	}()
	f, err = ioutil.TempFile("", "luci-server-test")
	if err != nil {
		return nil, err
	}
	secret := secrets.Secret{Current: []byte("test secret")}
	if err := json.NewEncoder(f).Encode(&secret); err != nil {
		return nil, err
	}
	return f, f.Close()
}

func setupListener() net.Listener {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		panic(err)
	}
	return l
}

type errorEvent struct {
	err    atomic.Value
	signal chan struct{} // closed after 'err' is populated
}

func (e *errorEvent) Set(err error) {
	if err != nil {
		e.err.Store(err)
	}
	close(e.signal)
}

func (e *errorEvent) Get() error {
	<-e.signal
	err, _ := e.err.Load().(error)
	return err
}

// httpGet makes a blocking request, aborting it if 'abort' is signaled.
func httpGet(addr string, abort *errorEvent, headers map[string]string) (resp string, err error) {
	done := make(chan struct{})
	go func() {
		defer close(done)
		var req *http.Request
		if req, err = http.NewRequest("GET", addr, nil); err != nil {
			return
		}
		for k, v := range headers {
			req.Header.Set(k, v)
		}
		var res *http.Response
		if res, err = http.DefaultClient.Do(req); err != nil {
			return
		}
		defer res.Body.Close()
		var blob []byte
		if blob, err = ioutil.ReadAll(res.Body); err != nil {
			return
		}
		if res.StatusCode >= 400 {
			err = fmt.Errorf("unexpected status %d", res.StatusCode)
		}
		resp = string(blob)
	}()

	select {
	case <-abort.signal:
		err = abort.Get()
	case <-done:
	}
	return
}

type logsRecorder struct {
	m    sync.Mutex
	logs []gkelogger.LogEntry
}

func (r *logsRecorder) Write(e *gkelogger.LogEntry) {
	r.m.Lock()
	r.logs = append(r.logs, *e)
	r.m.Unlock()
}

func (r *logsRecorder) Last(n int) []gkelogger.LogEntry {
	entries := make([]gkelogger.LogEntry, n)
	r.m.Lock()
	copy(entries, r.logs[len(r.logs)-n:])
	r.m.Unlock()
	return entries
}
<commit_msg>[server] Add a benchmark for a simple handler, refactor tests.<commit_after>// Copyright 2019 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"go.chromium.org/luci/common/clock/testclock"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/logging/gkelogger"
	"go.chromium.org/luci/server/router"
	"go.chromium.org/luci/server/secrets"

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

func TestServer(t *testing.T) {
	t.Parallel()

	Convey("Works", t, func() {
		ctx, tc := testclock.UseTime(context.Background(), testclock.TestRecentTimeUTC)

		srv, err := newTestServer(ctx)
		So(err, ShouldBeNil)
		defer srv.cleanup()
		srv.ServeInBackground()
		Reset(func() { So(srv.StopBackgroundServing(), ShouldBeNil) })

		Convey("Logging", func() {
			srv.Routes.GET("/test", router.MiddlewareChain{}, func(c *router.Context) {
				logging.Infof(c.Context, "Info log")
				tc.Add(time.Second)
				logging.Warningf(c.Context, "Warn log")
				c.Writer.WriteHeader(201)
				c.Writer.Write([]byte("Hello, world"))
			})

			resp, err := srv.Get("/test", map[string]string{
				"User-Agent":            "Test-user-agent",
				"X-Cloud-Trace-Context": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa/00001;trace=TRUE",
				"X-Forwarded-For":       "1.1.1.1,2.2.2.2,3.3.3.3",
			})
			So(err, ShouldBeNil)
			So(resp, ShouldEqual, "Hello, world")

			// Stderr log captures details about the request.
			So(srv.stderr.Last(1), ShouldResemble, []gkelogger.LogEntry{
				{
					Severity: "warning",
					Time:     "1454472307.7",
					TraceID:  "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
					RequestInfo: &gkelogger.RequestInfo{
						Method:       "GET",
						URL:          srv.mainAddr + "/test",
						Status:       201,
						RequestSize:  "0",
						ResponseSize: "12", // len("Hello, world")
						UserAgent:    "Test-user-agent",
						RemoteIP:     "2.2.2.2",
						Latency:      "1.000000s",
					},
				},
			})
			// Stdout log captures individual log lines.
			So(srv.stdout.Last(2), ShouldResemble, []gkelogger.LogEntry{
				{
					Severity: "info",
					Message:  "Info log",
					Time:     "1454472306.7",
					TraceID:  "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
					Operation: &gkelogger.Operation{
						ID: "9566c74d10037c4d7bbb0407d1e2c649",
					},
				},
				{
					Severity: "warning",
					Message:  "Warn log",
					Time:     "1454472307.7",
					TraceID:  "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
					Operation: &gkelogger.Operation{
						ID: "9566c74d10037c4d7bbb0407d1e2c649",
					},
				},
			})
		})

		Convey("Secrets", func() {
			srv.Routes.GET("/secret", router.MiddlewareChain{}, func(c *router.Context) {
				s, err := secrets.GetSecret(c.Context, "secret_name")
				if err != nil {
					c.Writer.WriteHeader(500)
				} else {
					c.Writer.Write([]byte(s.Current))
				}
			})
			resp, err := srv.Get("/secret", nil)
			So(err, ShouldBeNil)
			So(resp, ShouldNotBeEmpty)
		})
	})
}

func BenchmarkServer(b *testing.B) {
	srv, err := newTestServer(context.Background())
	if err != nil {
		b.Fatal(err)
	}
	defer srv.cleanup()

	// The route we are going to hit from the benchmark.
	srv.Routes.GET("/test", router.MiddlewareChain{}, func(c *router.Context) {
		logging.Infof(c.Context, "Hello, world")
		secrets.GetSecret(c.Context, "key-name") // e.g. checking XSRF token
		c.Writer.Write([]byte("Hello, world"))
	})

	// Don't actually store logs from all many-many iterations of the loop below.
	srv.stdout.discard = true
	srv.stderr.discard = true

	// Launch the server and wait for it to start serving to make sure all guts
	// are initialized.
	srv.ServeInBackground()
	defer srv.StopBackgroundServing()
	if _, err = srv.Get("/health", nil); err != nil {
		b.Fatal(err)
	}

	// Actual benchmark loop. Note that we bypass network layer here completely
	// (by not using http.DefaultClient).
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		req, err := http.NewRequest("GET", "/test", nil)
		if err != nil {
			b.Fatal(err)
		}
		req.Header.Set("X-Cloud-Trace-Context", "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa/00001;trace=TRUE")
		req.Header.Set("X-Forwarded-For", "1.1.1.1,2.2.2.2,3.3.3.3")
		rr := httptest.NewRecorder()
		srv.Routes.ServeHTTP(rr, req)
		if rr.Code != http.StatusOK {
			b.Fatalf("unexpected status %d", rr.Code)
		}
	}
}

////////////////////////////////////////////////////////////////////////////////

type testServer struct {
	*Server

	stdout logsRecorder
	stderr logsRecorder

	mainAddr string
	cleanup  func()
	serveErr errorEvent
}

func newTestServer(ctx context.Context) (*testServer, error) {
	srv := &testServer{
		serveErr: errorEvent{signal: make(chan struct{})},
	}

	tmpSecret, err := tempSecret()
	if err != nil {
		return nil, err
	}
	srv.cleanup = func() { os.Remove(tmpSecret.Name()) }

	srv.Server = New(Options{
		Prod:           true,
		HTTPAddr:       "main_addr",
		AdminAddr:      "admin_addr",
		RootSecretPath: tmpSecret.Name(),

		testCtx:    ctx,
		testSeed:   1,
		testStdout: &srv.stdout,
		testStderr: &srv.stderr,

		// Bind to auto-assigned ports.
		testListeners: map[string]net.Listener{
			"main_addr":  setupListener(),
			"admin_addr": setupListener(),
		},
	})

	mainPort := srv.opts.testListeners["main_addr"].Addr().(*net.TCPAddr).Port
	srv.mainAddr = fmt.Sprintf("http://127.0.0.1:%d", mainPort)

	return srv, nil
}

func (s *testServer) ServeInBackground() {
	go func() { s.serveErr.Set(s.ListenAndServe()) }()
}

func (s *testServer) StopBackgroundServing() error {
	s.Shutdown()
	return s.serveErr.Get()
}

// Get makes a blocking request, aborting it if the server dies.
func (s *testServer) Get(uri string, headers map[string]string) (resp string, err error) {
	done := make(chan struct{})
	go func() {
		defer close(done)
		var req *http.Request
		if req, err = http.NewRequest("GET", s.mainAddr+uri, nil); err != nil {
			return
		}
		for k, v := range headers {
			req.Header.Set(k, v)
		}
		var res *http.Response
		if res, err = http.DefaultClient.Do(req); err != nil {
			return
		}
		defer res.Body.Close()
		var blob []byte
		if blob, err = ioutil.ReadAll(res.Body); err != nil {
			return
		}
		if res.StatusCode >= 400 {
			err = fmt.Errorf("unexpected status %d", res.StatusCode)
		}
		resp = string(blob)
	}()

	select {
	case <-s.serveErr.signal:
		err = s.serveErr.Get()
	case <-done:
	}
	return
}

////////////////////////////////////////////////////////////////////////////////

func tempSecret() (out *os.File, err error) {
	var f *os.File
	defer func() {
		if f != nil && err != nil {
			os.Remove(f.Name())
		}
	}()
	f, err = ioutil.TempFile("", "luci-server-test")
	if err != nil {
		return nil, err
	}
	secret := secrets.Secret{Current: []byte("test secret")}
	if err := json.NewEncoder(f).Encode(&secret); err != nil {
		return nil, err
	}
	return f, f.Close()
}

func setupListener() net.Listener {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		panic(err)
	}
	return l
}

////////////////////////////////////////////////////////////////////////////////

type errorEvent struct {
	err    atomic.Value
	signal chan struct{} // closed after 'err' is populated
}

func (e *errorEvent) Set(err error) {
	if err != nil {
		e.err.Store(err)
	}
	close(e.signal)
}

func (e *errorEvent) Get() error {
	<-e.signal
	err, _ := e.err.Load().(error)
	return err
}

////////////////////////////////////////////////////////////////////////////////

type logsRecorder struct {
	discard bool
	m       sync.Mutex
	logs    []gkelogger.LogEntry
}

func (r *logsRecorder) Write(e *gkelogger.LogEntry) {
	if r.discard {
		return
	}
	r.m.Lock()
	r.logs = append(r.logs, *e)
	r.m.Unlock()
}

func (r *logsRecorder) Last(n int) []gkelogger.LogEntry {
	entries := make([]gkelogger.LogEntry, n)
	r.m.Lock()
	copy(entries, r.logs[len(r.logs)-n:])
	r.m.Unlock()
	return entries
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/kidstuff/auth-mongo-mngr"
	"github.com/kidstuff/auth/model"
	"labix.org/v2/mgo"
	"log"
	"os"
)

func main() {
	MONGODB_URL := os.Getenv("MONGODB_URL")
	SERVER_URL := os.Getenv("SERVER_URL")
	DB_NAME := os.Getenv("DB_NAME")

	if len(MONGODB_URL) == 0 {
		MONGODB_URL = "localhost"
	}

	if len(SERVER_URL) == 0 {
		SERVER_URL = ":8080"
	}

	if len(DB_NAME) == 0 {
		DB_NAME = "kidstuff_auth"
	}

	session, err := mgo.Dial(MONGODB_URL)
	if err != nil {
		panic(err)
	}
	defer session.Close()
	session.SetMode(mgo.Monotonic, true)
	db := session.DB(DB_NAME)

	err = mgoauth.EnsureIndex(db)
	if err != nil {
		log.Println(err)
	}

	conf := mgoauth.NewMgoConfigMngr(db)
	settings := map[string]string{
		"auth_full_path":              "http://localhost:8080/auth",
		"auth_activate_page":          "http://localhost:8082/#!/user/%s/active?code=%s",
		"auth_approve_new_user":       "false",
		"auth_email_from":             "nvcnvn1@gmail.com",
		"auth_send_activate_email":    "true",
		"auth_activate_email_subject": "Active your account",
		"auth_activate_email_message": "Hi!\nPlease active your account by cliking here:\n%s",
		"auth_send_welcome_email":     "true",
		"auth_welcome_email_subject":  "Welcome!",
		"auth_welcome_email_message":  "Hi!\nWelcome you to join our community :)",
	}
	err = conf.SetMulti(settings)
	if err != nil {
		log.Println(err)
	}

	groupMngr := mgoauth.NewMgoGroupManager(db)
	groupName := "admin"
	g, err := groupMngr.FindByName(groupName)
	if err != nil {
		g = &model.Group{}
		g.Name = &groupName
		g.Privilege = []string{"manage_user"}
		g, err = groupMngr.AddDetail(g)
		if err != nil {
			log.Println(err)
		}
	}

	userMngr := mgoauth.NewMgoUserManager(db, groupMngr)
	email := "nvcnvn1@gmail.com"
	u, err := userMngr.FindByEmail(email)
	if err != nil {
		u = &model.User{}
		u.Email = &email
		u.ChangePassword("zaq123edc")
		g2 := model.Group{}
		g2.Id = g.Id
		g2.Name = g.Name
		u.Groups = []model.Group{g2}
		t := true
		u.Approved = &t
		u.Privilege = g.Privilege
		u, err = userMngr.AddDetail(u)
		if err != nil {
			log.Println(err)
		}
	}
}
<commit_msg>use new AddDetail interface<commit_after>package main

import (
	"github.com/kidstuff/auth-mongo-mngr"
	"github.com/kidstuff/auth/model"
	"labix.org/v2/mgo"
	"log"
	"os"
)

func main() {
	MONGODB_URL := os.Getenv("MONGODB_URL")
	SERVER_URL := os.Getenv("SERVER_URL")
	DB_NAME := os.Getenv("DB_NAME")

	if len(MONGODB_URL) == 0 {
		MONGODB_URL = "localhost"
	}

	if len(SERVER_URL) == 0 {
		SERVER_URL = ":8080"
	}

	if len(DB_NAME) == 0 {
		DB_NAME = "kidstuff_auth"
	}

	session, err := mgo.Dial(MONGODB_URL)
	if err != nil {
		panic(err)
	}
	defer session.Close()
	session.SetMode(mgo.Monotonic, true)
	db := session.DB(DB_NAME)

	err = mgoauth.EnsureIndex(db)
	if err != nil {
		log.Println(err)
	}

	conf := mgoauth.NewMgoConfigMngr(db)
	settings := map[string]string{
		"auth_full_path":              "http://localhost:8080/auth",
		"auth_activate_page":          "http://localhost:8082/#!/user/%s/active?code=%s",
		"auth_approve_new_user":       "false",
		"auth_email_from":             "nvcnvn1@gmail.com",
		"auth_send_activate_email":    "true",
		"auth_activate_email_subject": "Active your account",
		"auth_activate_email_message": "Hi!\nPlease active your account by cliking here:\n%s",
		"auth_send_welcome_email":     "true",
		"auth_welcome_email_subject":  "Welcome!",
		"auth_welcome_email_message":  "Hi!\nWelcome you to join our community :)",
	}
	err = conf.SetMulti(settings)
	if err != nil {
		log.Println(err)
	}

	groupMngr := mgoauth.NewMgoGroupManager(db)
	groupName := "admin"
	g, err := groupMngr.FindByName(groupName)
	if err != nil {
		g = &model.Group{}
		g.Name = &groupName
		g.Privilege = []string{"manage_user"}
		g, err = groupMngr.AddDetail(g)
		if err != nil {
			log.Println(err)
		}
	}

	userMngr := mgoauth.NewMgoUserManager(db, groupMngr)
	email := "nvcnvn1@gmail.com"
	_, err = userMngr.FindByEmail(email)
	if err != nil {
		g2 := model.Group{}
		g2.Id = g.Id
		g2.Name = g.Name
		_, err = userMngr.AddDetail(email, "zaq123edc", true, g.Privilege, nil,
			nil, []model.Group{g2})
		if err != nil {
			log.Println(err)
		}
	}
}
<|endoftext|>


<commit_before>package persist

import (
	"fmt"
	"io"
	"regexp"
	"time"

	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/uuid"
	"github.com/pachyderm/pachyderm/src/server/pfs/drive"

	"github.com/dancannon/gorethink"
	"go.pedge.io/pb/go/google/protobuf"
	"go.pedge.io/proto/time"
	"google.golang.org/grpc"
)

// A Table is a rethinkdb table name.
type Table string

// A PrimaryKey is a rethinkdb primary key identifier.
type PrimaryKey string

// An Index is a rethinkdb index.
type Index string

const (
	repoTable   Table = "Repos"
	branchTable Table = "Branches"

	commitTable Table = "Commits"
	// commitBranchIndex maps commits to branches
	commitBranchIndex Index = "CommitBranchIndex"

	diffTable Table = "Diffs"

	connectTimeoutSeconds = 5
)

var (
	tables = []Table{
		repoTable,
		branchTable,
		commitTable,
		diffTable,
	}

	tableToTableCreateOpts = map[Table][]gorethink.TableCreateOpts{
		repoTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "Name",
			},
		},
		branchTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "ID",
			},
		},
		commitTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "ID",
			},
		},
		diffTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "ID",
			},
		},
	}
)

type driver struct {
	blockAddress string
	blockClient  pfs.BlockAPIClient

	dbAddress string
	dbName    string
	dbClient  *gorethink.Session
}

func NewDriver(blockAddress string, dbAddress string, dbName string) (drive.Driver, error) {
	clientConn, err := grpc.Dial(blockAddress, grpc.WithInsecure())
	if err != nil {
		return nil, err
	}

	dbClient, err := dbConnect(dbAddress)
	if err != nil {
		return nil, err
	}

	return &driver{
		blockAddress: blockAddress,
		blockClient:  pfs.NewBlockAPIClient(clientConn),
		dbAddress:    dbAddress,
		dbName:       dbName,
		dbClient:     dbClient,
	}, nil
}

func InitDB(address string, databaseName string) error {
	session, err := dbConnect(address)
	if err != nil {
		return err
	}
	defer session.Close()

	// Create the database
	if _, err := gorethink.DBCreate(databaseName).RunWrite(session); err != nil {
		return err
	}

	// Create tables
	for _, table := range tables {
		tableCreateOpts := tableToTableCreateOpts[table]
		if _, err := gorethink.DB(databaseName).TableCreate(table, tableCreateOpts...).RunWrite(session); err != nil {
			return err
		}
	}

	// Create indexes
	if _, err := gorethink.DB(databaseName).Table(commitTable).IndexCreateFunc(commitBranchIndex, func(row gorethink.Term) interface{} {
		return row.Field("BranchClocks").Map(func(branchClock gorethink.Term) interface{} {
			lastClock := branchClock.Field("Clocks").Nth(-1)
			return []interface{}{
				lastClock.Field("Branch"),
				lastClock.Field("Clock"),
			}
		})
	}, gorethink.IndexCreateOpts{
		Multi: true,
	}).RunWrite(session); err != nil {
		return err
	}

	return nil
}

func RemoveDB(address string, databaseName string) error {
	session, err := dbConnect(address)
	if err != nil {
		return err
	}
	defer session.Close()

	// Create the database
	if _, err := gorethink.DBDrop(databaseName).RunWrite(session); err != nil {
		return err
	}

	return nil
}

func dbConnect(address string) (*gorethink.Session, error) {
	return gorethink.Connect(gorethink.ConnectOpts{
		Address: address,
		Timeout: connectTimeoutSeconds * time.Second,
	})
}

func validateRepoName(name string) error {
	match, _ := regexp.MatchString("^[a-zA-Z0-9_]+$", name)

	if !match {
		return fmt.Errorf("repo name (%v) invalid: only alphanumeric and underscore characters allowed", name)
	}

	return nil
}

func (d *driver) getTerm(table Table) gorethink.Term {
	return gorethink.DB(d.dbName).Table(table)
}

func (d *driver) CreateRepo(repo *pfs.Repo, created *google_protobuf.Timestamp,
	provenance []*pfs.Repo, shards map[uint64]bool) error {

	err := validateRepoName(repo.Name)
	if err != nil {
		return err
	}

	_, err = d.getTerm(repoTable).Insert(&Repo{
		Name:    repo.Name,
		Created: created,
	}).RunWrite(d.dbClient)
	return err
}

func (d *driver) InspectRepo(repo *pfs.Repo, shards map[uint64]bool) (repoInfo *pfs.RepoInfo, retErr error) {
	cursor, err := d.getTerm(repoTable).Get(repo.Name).Default(gorethink.Error("value not found")).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	rawRepo := &Repo{}
	cursor.Next(rawRepo)
	if err := cursor.Err(); err != nil {
		return nil, err
	}
	repoInfo = &pfs.RepoInfo{
		Repo:    &pfs.Repo{rawRepo.Name},
		Created: rawRepo.Created,
	}
	return repoInfo, nil
}

func (d *driver) ListRepo(provenance []*pfs.Repo, shards map[uint64]bool) (repoInfos []*pfs.RepoInfo, retErr error) {
	cursor, err := d.getTerm(repoTable).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	for {
		repo := &Repo{}
		if !cursor.Next(repo) {
			break
		}
		repoInfos = append(repoInfos, &pfs.RepoInfo{
			Repo:    &pfs.Repo{repo.Name},
			Created: repo.Created,
		})
	}
	if err := cursor.Err(); err != nil {
		return nil, err
	}
	return repoInfos, nil
}

func (d *driver) DeleteRepo(repo *pfs.Repo, shards map[uint64]bool, force bool) error {
	_, err := d.getTerm(repoTable).Get(repo.Name).Delete().RunWrite(d.dbClient)
	return err
}

func (d *driver) StartCommit(repo *pfs.Repo, commitID string, parentID string, branch string, started *google_protobuf.Timestamp, provenance []*pfs.Commit, shards map[uint64]bool) (retErr error) {
	if commitID == "" {
		commitID = uuid.NewWithoutDashes()
	}

	var _provenance []string
	for _, c := range provenance {
		_provenance = append(_provenance, c.ID)
	}
	commit := &Commit{
		ID:         commitID,
		Repo:       repo.Name,
		Started:    prototime.TimeToTimestamp(time.Now()),
		Provenance: _provenance,
	}

	if parentID == "" {
		if branch == "" {
			branch = uuid.NewWithoutDashes()
		}

		_, err := d.getTerm(branchTable).Insert(&Branch{
			ID: branchID(repo.Name, branch),
		}, gorethink.InsertOpts{
			// Set conflict to "replace" because we don't want to get an error
			// when the branch already exists.
			Conflict: "replace",
		}).RunWrite(d.dbClient)
		if err != nil {
			return err
		}

		cursor, err := d.getTerm(commitTable).OrderBy(gorethink.OrderByOpts{
			Index: gorethink.Desc(commitBranchIndex),
		}).Between(
			[]interface{}{branch, 0},
			[]interface{}{branch, gorethink.MaxVal},
		).Run(d.dbClient)
		if err != nil {
			return err
		}
		defer func() {
			if err := cursor.Close(); err != nil && retErr == nil {
				retErr = err
			}
		}()

		// The last commit on this branch will be our parent commit
		parentCommit := &Commit{}
		found := cursor.Next(parentCommit)
		if err := cursor.Err(); err != nil {
			return err
		}
		if !found {
			// we don't have a parent :(
			// so we create a new BranchClock
			commit.BranchClocks = []*BranchClock{{
				Clocks: []*Clock{{
					Branch: branch,
					Clock:  0,
				}},
			}}
		} else {
			// we do have a parent :D
			// so we inherit our parent's branch clock for this particular branch,
			// and increment the last component by 1
			var set bool
			for _, branchClock := range parentCommit.BranchClocks {
				if branchClock.Clocks[len(branchClock.Clocks)-1].Branch == branch {
					branchClock.Clocks[len(branchClock.Clocks)-1].Clock += 1
					commit.BranchClocks = []*BranchClock{branchClock}
					set = true
				}
			}
			if !set {
				return fmt.Errorf("commitBranchIndex returned a parent commit, but the parent commit is not on the branch that we are operating on; this is a bug")
			}
		}
	}

	_, err := d.getTerm(commitTable).Insert(commit).RunWrite(d.dbClient)

	return err
}

func branchID(repo string, name string) string {
	return fmt.Sprintf("%s/%s", repo, name)
}

// FinishCommit blocks until its parent has been finished/cancelled
func (d *driver) FinishCommit(commit *pfs.Commit, finished *google_protobuf.Timestamp, cancel bool, shards map[uint64]bool) error {
	_, err := d.getTerm(commitTable).Get(commit.ID).Update(
		map[string]interface{}{
			"Finished": finished,
		},
	).RunWrite(d.dbClient)

	return err
}

func (d *driver) InspectCommit(commit *pfs.Commit, shards map[uint64]bool) (commitInfo *pfs.CommitInfo, retErr error) {
	cursor, err := d.getTerm(commitTable).Get(commit.ID).Default(gorethink.Error("value not found")).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()

	rawCommit := &Commit{}
	cursor.Next(rawCommit)
	if err := cursor.Err(); err != nil {
		return nil, err
	}

	return rawCommitToCommitInfo(rawCommit), nil
}

func rawCommitToCommitInfo(rawCommit *Commit) *pfs.CommitInfo {
	commitType := pfs.CommitType_COMMIT_TYPE_READ
	if rawCommit.Finished == nil {
		commitType = pfs.CommitType_COMMIT_TYPE_WRITE
	}
	return &pfs.CommitInfo{
		Commit: &pfs.Commit{
			Repo: &pfs.Repo{rawCommit.Repo},
			ID:   rawCommit.ID,
		},
		Started:    rawCommit.Started,
		Finished:   rawCommit.Finished,
		CommitType: commitType,
	}
}

func (d *driver) ListCommit(repos []*pfs.Repo, commitType pfs.CommitType, fromCommit []*pfs.Commit,
	provenance []*pfs.Commit, all bool, shards map[uint64]bool) (commitInfos []*pfs.CommitInfo, retErr error) {
	cursor, err := d.getTerm(commitTable).Filter(func(commit gorethink.Term) gorethink.Term {
		var predicates []interface{}
		for _, repo := range repos {
			predicates = append(predicates, commit.Field("Repo").Eq(repo.Name))
		}
		return gorethink.Or(predicates...)
	}).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	for {
		rawCommit := &Commit{}
		if !cursor.Next(rawCommit) {
			break
		}
		commitInfos = append(commitInfos, rawCommitToCommitInfo(rawCommit))
	}
	if err := cursor.Err(); err != nil {
		return nil, err
	}
	return commitInfos, nil
}

func (d *driver) ListBranch(repo *pfs.Repo, shards map[uint64]bool) ([]*pfs.CommitInfo, error) {
	return nil, nil
}

func (d *driver) DeleteCommit(commit *pfs.Commit, shards map[uint64]bool) error {
	return nil
}

func (d *driver) PutFile(file *pfs.File, handle string,
	delimiter pfs.Delimiter, shard uint64, reader io.Reader) (retErr error) {
	return nil
}

func (d *driver) MakeDirectory(file *pfs.File, shard uint64) (retErr error) {
	return nil
}

func (d *driver) GetFile(file *pfs.File, filterShard *pfs.Shard, offset int64,
	size int64, from *pfs.Commit, shard uint64, unsafe bool, handle string) (io.ReadCloser, error) {
	return nil, nil
}

func (d *driver) InspectFile(file *pfs.File, filterShard *pfs.Shard, from *pfs.Commit, shard uint64, unsafe bool, handle string) (*pfs.FileInfo, error) {
	return nil, nil
}

func (d *driver) ListFile(file *pfs.File, filterShard *pfs.Shard, from *pfs.Commit, shard uint64, recurse bool, unsafe bool, handle string) ([]*pfs.FileInfo, error) {
	return nil, nil
}

func (d *driver) DeleteFile(file *pfs.File, shard uint64, unsafe bool, handle string) error {
	return nil
}

func (d *driver) DeleteAll(shards map[uint64]bool) error {
	return nil
}

func (d *driver) AddShard(shard uint64) error {
	return nil
}

func (d *driver) DeleteShard(shard uint64) error {
	return nil
}

func (d *driver) Dump() {
}
<commit_msg>Fix to make TestStartCommit pass<commit_after>package persist

import (
	"fmt"
	"io"
	"regexp"
	"time"

	"github.com/pachyderm/pachyderm/src/client/pfs"
	"github.com/pachyderm/pachyderm/src/client/pkg/uuid"
	"github.com/pachyderm/pachyderm/src/server/pfs/drive"

	"github.com/dancannon/gorethink"
	"go.pedge.io/pb/go/google/protobuf"
	"go.pedge.io/proto/time"
	"google.golang.org/grpc"
)

// A Table is a rethinkdb table name.
type Table string

// A PrimaryKey is a rethinkdb primary key identifier.
type PrimaryKey string

// An Index is a rethinkdb index.
type Index string

const (
	repoTable   Table = "Repos"
	branchTable Table = "Branches"

	commitTable Table = "Commits"
	// commitBranchIndex maps commits to branches
	commitBranchIndex Index = "CommitBranchIndex"

	diffTable Table = "Diffs"

	connectTimeoutSeconds = 5
)

var (
	tables = []Table{
		repoTable,
		branchTable,
		commitTable,
		diffTable,
	}

	tableToTableCreateOpts = map[Table][]gorethink.TableCreateOpts{
		repoTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "Name",
			},
		},
		branchTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "ID",
			},
		},
		commitTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "ID",
			},
		},
		diffTable: []gorethink.TableCreateOpts{
			gorethink.TableCreateOpts{
				PrimaryKey: "ID",
			},
		},
	}
)

type driver struct {
	blockAddress string
	blockClient  pfs.BlockAPIClient

	dbAddress string
	dbName    string
	dbClient  *gorethink.Session
}

func NewDriver(blockAddress string, dbAddress string, dbName string) (drive.Driver, error) {
	clientConn, err := grpc.Dial(blockAddress, grpc.WithInsecure())
	if err != nil {
		return nil, err
	}

	dbClient, err := dbConnect(dbAddress)
	if err != nil {
		return nil, err
	}

	return &driver{
		blockAddress: blockAddress,
		blockClient:  pfs.NewBlockAPIClient(clientConn),
		dbAddress:    dbAddress,
		dbName:       dbName,
		dbClient:     dbClient,
	}, nil
}

func InitDB(address string, databaseName string) error {
	session, err := dbConnect(address)
	if err != nil {
		return err
	}
	defer session.Close()

	// Create the database
	if _, err := gorethink.DBCreate(databaseName).RunWrite(session); err != nil {
		return err
	}

	// Create tables
	for _, table := range tables {
		tableCreateOpts := tableToTableCreateOpts[table]
		if _, err := gorethink.DB(databaseName).TableCreate(table, tableCreateOpts...).RunWrite(session); err != nil {
			return err
		}
	}

	// Create indexes
	if _, err := gorethink.DB(databaseName).Table(commitTable).IndexCreateFunc(commitBranchIndex, func(row gorethink.Term) interface{} {
		return row.Field("BranchClocks").Map(func(branchClock gorethink.Term) interface{} {
			lastClock := branchClock.Field("Clocks").Nth(-1)
			return []interface{}{
				lastClock.Field("Branch"),
				lastClock.Field("Clock"),
			}
		})
	}, gorethink.IndexCreateOpts{
		Multi: true,
	}).RunWrite(session); err != nil {
		return err
	}

	// Wait for indexes to be ready
	if _, err := gorethink.DB(databaseName).Table(commitTable).IndexWait(commitBranchIndex).RunWrite(session); err != nil {
		return err
	}

	return nil
}

func RemoveDB(address string, databaseName string) error {
	session, err := dbConnect(address)
	if err != nil {
		return err
	}
	defer session.Close()

	// Create the database
	if _, err := gorethink.DBDrop(databaseName).RunWrite(session); err != nil {
		return err
	}

	return nil
}

func dbConnect(address string) (*gorethink.Session, error) {
	return gorethink.Connect(gorethink.ConnectOpts{
		Address: address,
		Timeout: connectTimeoutSeconds * time.Second,
	})
}

func validateRepoName(name string) error {
	match, _ := regexp.MatchString("^[a-zA-Z0-9_]+$", name)

	if !match {
		return fmt.Errorf("repo name (%v) invalid: only alphanumeric and underscore characters allowed", name)
	}

	return nil
}

func (d *driver) getTerm(table Table) gorethink.Term {
	return gorethink.DB(d.dbName).Table(table)
}

func (d *driver) CreateRepo(repo *pfs.Repo, created *google_protobuf.Timestamp,
	provenance []*pfs.Repo, shards map[uint64]bool) error {

	err := validateRepoName(repo.Name)
	if err != nil {
		return err
	}

	_, err = d.getTerm(repoTable).Insert(&Repo{
		Name:    repo.Name,
		Created: created,
	}).RunWrite(d.dbClient)
	return err
}

func (d *driver) InspectRepo(repo *pfs.Repo, shards map[uint64]bool) (repoInfo *pfs.RepoInfo, retErr error) {
	cursor, err := d.getTerm(repoTable).Get(repo.Name).Default(gorethink.Error("value not found")).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	rawRepo := &Repo{}
	cursor.Next(rawRepo)
	if err := cursor.Err(); err != nil {
		return nil, err
	}
	repoInfo = &pfs.RepoInfo{
		Repo:    &pfs.Repo{rawRepo.Name},
		Created: rawRepo.Created,
	}
	return repoInfo, nil
}

func (d *driver) ListRepo(provenance []*pfs.Repo, shards map[uint64]bool) (repoInfos []*pfs.RepoInfo, retErr error) {
	cursor, err := d.getTerm(repoTable).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	for {
		repo := &Repo{}
		if !cursor.Next(repo) {
			break
		}
		repoInfos = append(repoInfos, &pfs.RepoInfo{
			Repo:    &pfs.Repo{repo.Name},
			Created: repo.Created,
		})
	}
	if err := cursor.Err(); err != nil {
		return nil, err
	}
	return repoInfos, nil
}

func (d *driver) DeleteRepo(repo *pfs.Repo, shards map[uint64]bool, force bool) error {
	_, err := d.getTerm(repoTable).Get(repo.Name).Delete().RunWrite(d.dbClient)
	return err
}

func (d *driver) StartCommit(repo *pfs.Repo, commitID string, parentID string, branch string, started *google_protobuf.Timestamp, provenance []*pfs.Commit, shards map[uint64]bool) (retErr error) {
	if commitID == "" {
		commitID = uuid.NewWithoutDashes()
	}

	var _provenance []string
	for _, c := range provenance {
		_provenance = append(_provenance, c.ID)
	}
	commit := &Commit{
		ID:         commitID,
		Repo:       repo.Name,
		Started:    prototime.TimeToTimestamp(time.Now()),
		Provenance: _provenance,
	}

	if parentID == "" {
		if branch == "" {
			branch = uuid.NewWithoutDashes()
		}

		_, err := d.getTerm(branchTable).Insert(&Branch{
			ID: branchID(repo.Name, branch),
		}, gorethink.InsertOpts{
			// Set conflict to "replace" because we don't want to get an error
			// when the branch already exists.
			Conflict: "replace",
		}).RunWrite(d.dbClient)
		if err != nil {
			return err
		}

		cursor, err := d.getTerm(commitTable).OrderBy(gorethink.OrderByOpts{
			Index: gorethink.Desc(commitBranchIndex),
		}).Between(
			[]interface{}{branch, 0},
			[]interface{}{branch, gorethink.MaxVal},
		).Run(d.dbClient)
		if err != nil {
			return err
		}
		defer func() {
			if err := cursor.Close(); err != nil && retErr == nil {
				retErr = err
			}
		}()

		// The last commit on this branch will be our parent commit
		parentCommit := &Commit{}
		found := cursor.Next(parentCommit)
		if err := cursor.Err(); err != nil {
			return err
		}
		if !found {
			// we don't have a parent :(
			// so we create a new BranchClock
			commit.BranchClocks = []*BranchClock{{
				Clocks: []*Clock{{
					Branch: branch,
					Clock:  0,
				}},
			}}
		} else {
			// we do have a parent :D
			// so we inherit our parent's branch clock for this particular branch,
			// and increment the last component by 1
			var set bool
			for _, branchClock := range parentCommit.BranchClocks {
				if branchClock.Clocks[len(branchClock.Clocks)-1].Branch == branch {
					branchClock.Clocks[len(branchClock.Clocks)-1].Clock += 1
					commit.BranchClocks = []*BranchClock{branchClock}
					set = true
				}
			}
			if !set {
				return fmt.Errorf("commitBranchIndex returned a parent commit, but the parent commit is not on the branch that we are operating on; this is a bug")
			}
		}
	}

	_, err := d.getTerm(commitTable).Insert(commit).RunWrite(d.dbClient)

	return err
}

func branchID(repo string, name string) string {
	return fmt.Sprintf("%s/%s", repo, name)
}

// FinishCommit blocks until its parent has been finished/cancelled
func (d *driver) FinishCommit(commit *pfs.Commit, finished *google_protobuf.Timestamp, cancel bool, shards map[uint64]bool) error {
	_, err := d.getTerm(commitTable).Get(commit.ID).Update(
		map[string]interface{}{
			"Finished": finished,
		},
	).RunWrite(d.dbClient)

	return err
}

func (d *driver) InspectCommit(commit *pfs.Commit, shards map[uint64]bool) (commitInfo *pfs.CommitInfo, retErr error) {
	cursor, err := d.getTerm(commitTable).Get(commit.ID).Default(gorethink.Error("value not found")).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()

	rawCommit := &Commit{}
	cursor.Next(rawCommit)
	if err := cursor.Err(); err != nil {
		return nil, err
	}

	return rawCommitToCommitInfo(rawCommit), nil
}

func rawCommitToCommitInfo(rawCommit *Commit) *pfs.CommitInfo {
	commitType := pfs.CommitType_COMMIT_TYPE_READ
	if rawCommit.Finished == nil {
		commitType = pfs.CommitType_COMMIT_TYPE_WRITE
	}
	return &pfs.CommitInfo{
		Commit: &pfs.Commit{
			Repo: &pfs.Repo{rawCommit.Repo},
			ID:   rawCommit.ID,
		},
		Started:    rawCommit.Started,
		Finished:   rawCommit.Finished,
		CommitType: commitType,
	}
}

func (d *driver) ListCommit(repos []*pfs.Repo, commitType pfs.CommitType, fromCommit []*pfs.Commit,
	provenance []*pfs.Commit, all bool, shards map[uint64]bool) (commitInfos []*pfs.CommitInfo, retErr error) {
	cursor, err := d.getTerm(commitTable).Filter(func(commit gorethink.Term) gorethink.Term {
		var predicates []interface{}
		for _, repo := range repos {
			predicates = append(predicates, commit.Field("Repo").Eq(repo.Name))
		}
		return gorethink.Or(predicates...)
	}).Run(d.dbClient)
	if err != nil {
		return nil, err
	}
	defer func() {
		if err := cursor.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	for {
		rawCommit := &Commit{}
		if !cursor.Next(rawCommit) {
			break
		}
		commitInfos = append(commitInfos, rawCommitToCommitInfo(rawCommit))
	}
	if err := cursor.Err(); err != nil {
		return nil, err
	}
	return commitInfos, nil
}

func (d *driver) ListBranch(repo *pfs.Repo, shards map[uint64]bool) ([]*pfs.CommitInfo, error) {
	return nil, nil
}

func (d *driver) DeleteCommit(commit *pfs.Commit, shards map[uint64]bool) error {
	return nil
}

func (d *driver) PutFile(file *pfs.File, handle string,
	delimiter pfs.Delimiter, shard uint64, reader io.Reader) (retErr error) {
	return nil
}

func (d *driver) MakeDirectory(file *pfs.File, shard uint64) (retErr error) {
	return nil
}

func (d *driver) GetFile(file *pfs.File, filterShard *pfs.Shard, offset int64,
	size int64, from *pfs.Commit, shard uint64, unsafe bool, handle string) (io.ReadCloser, error) {
	return nil, nil
}

func (d *driver) InspectFile(file *pfs.File, filterShard *pfs.Shard, from *pfs.Commit, shard uint64, unsafe bool, handle string) (*pfs.FileInfo, error) {
	return nil, nil
}

func (d *driver) ListFile(file *pfs.File, filterShard *pfs.Shard, from *pfs.Commit, shard uint64, recurse bool, unsafe bool, handle string) ([]*pfs.FileInfo, error) {
	return nil, nil
}

func (d *driver) DeleteFile(file *pfs.File, shard uint64, unsafe bool, handle string) error {
	return nil
}

func (d *driver) DeleteAll(shards map[uint64]bool) error {
	return nil
}

func (d *driver) AddShard(shard uint64) error {
	return nil
}

func (d *driver) DeleteShard(shard uint64) error {
	return nil
}

func (d *driver) Dump() {
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"io"
)

var closeHeadTag = []byte(`</head>`) // TODO also look for </HEAD>

// TagInjector wraps a writer and adds a script tag to its content.
// It depends on the fact that dynamic page rendering makes a single Write call,
// so that it's guaranteed to find the marker within a single invocation argument.
// It doesn't parse HTML, so it could be spoofed but probably only intentionally.
type TagInjector struct {
	w         io.Writer
	insertion []byte
}

// Write injects a livereload script tag at the end of the HTML head, if present,
// else at the beginning of the document.
func (i TagInjector) Write(b []byte) (n int, err error) {
	if !bytes.Contains(b, i.insertion) && bytes.Contains(b, closeHeadTag) {
		r := append(i.insertion, closeHeadTag...)
		b = bytes.Replace(b, closeHeadTag, r, 1)
	}
	if !bytes.Contains(b, i.insertion) {
		b = append(i.insertion, b...)
	}
	return i.w.Write(b)
}
<commit_msg>fix short write<commit_after>package server

import (
	"bytes"
	"io"
)

var closeHeadTag = []byte(`</head>`) // TODO also look for </HEAD>

// TagInjector wraps a writer and adds a script tag to its content.
// It depends on the fact that dynamic page rendering makes a single Write call,
// so that it's guaranteed to find the marker within a single invocation argument.
// It doesn't parse HTML, so it could be spoofed but probably only intentionally.
type TagInjector struct {
	w         io.Writer
	insertion []byte
}

// Write injects a livereload script tag at the end of the HTML head, if present,
// else at the beginning of the document.
func (i TagInjector) Write(b []byte) (n int, err error) {
	n = len(b)
	if !bytes.Contains(b, i.insertion) && bytes.Contains(b, closeHeadTag) {
		r := append(i.insertion, closeHeadTag...)
		b = bytes.Replace(b, closeHeadTag, r, 1)
	}
	if !bytes.Contains(b, i.insertion) {
		b = append(i.insertion, b...)
	}
	_, err = i.w.Write(b)
	return
}
<|endoftext|>


<commit_before>package service

import (
	"context"
	"fmt"
	"os/exec"
	"testing"
	"time"

	"github.com/docker/docker/client"
	"github.com/stretchr/testify/suite"
)

type AlertTestSuite struct {
	suite.Suite
	url          string
	alertService AlertServicer
	client       *client.Client
}

func TestAlertUnitTestSuite(t *testing.T) {
	suite.Run(t, new(AlertTestSuite))
}

func (s *AlertTestSuite) SetupSuite() {
	client, _ := NewDockerClientFromEnv()
	_, err := client.Info(context.Background())
	if err != nil {
		s.T().Skipf("Unable to connect to Docker Client")
	}
	s.url = "http://localhost:9093"
	s.alertService = NewAlertService(s.url, time.Second*15)
	s.client = client
}

func (s *AlertTestSuite) TearDownSuite() {
	s.client.Close()
}

func (s *AlertTestSuite) Test_SendAlert() {

	defer func() {
		cmd := "docker container rm -f am9093"
		exec.Command("/bin/sh", "-c", cmd).Output()
	}()
	cmd := `docker run --name am9093 -p 9093:9093 \
			-d prom/alertmanager:v0.13.0`
	_, err := exec.Command("/bin/sh", "-c", cmd).Output()
	if err != nil {
		s.T().Skipf(fmt.Sprintf("Unable to create alertmanager: %s", err.Error()))
		return
	}

	// Wait for am to come online
	for i := 1; i <= 60; i++ {
		info, _ := s.client.ContainerInspect(context.Background(), "am9093")
		if info.State.Running {
			time.Sleep(1 * time.Second)
			return
		}
		time.Sleep(1 * time.Second)
	}
	s.T().Skipf("Unable to create alertmanager")

	require := s.Require()
	serviceName := "web"
	alertname := "service_scaler"
	status := "success"
	summary := "Scaled web from 3 to 4 replicas"
	request := "Scale web with delta=1"

	s.alertService.Send(alertname, serviceName, request, status, summary)
	time.Sleep(1 * time.Second)

	alerts, err := FetchAlerts(s.url, alertname, status, serviceName)
	require.NoError(err)
	require.Equal(1, len(alerts))

	alert := alerts[0]
	s.Equal(alertname, string(alert.Labels["alertname"]))
	s.Equal(serviceName, string(alert.Labels["service"]))
	s.Equal(status, string(alert.Labels["status"]))
	s.Equal(summary, string(alert.Annotations["summary"]))
	s.Equal(request, string(alert.Annotations["request"]))
	s.Equal("", alert.GeneratorURL)
}

func (s *AlertTestSuite) Test_generateAlert() {

	serviceName := "web"
	alertname := "service_scaler"
	status := "success"
	summary := "Scaled web from 3 to 4 replicas"
	request := "Scale web with delta=1"
	startsAt := time.Now().UTC()
	timeout := time.Second
	endsAt := startsAt.Add(timeout)

	alert := generateAlert(alertname, serviceName, request, status, summary, startsAt, timeout)
	s.Require().NotNil(alert)
	s.Equal(alertname, string(alert.Labels["alertname"]))
	s.Equal(serviceName, string(alert.Labels["service"]))
	s.Equal(status, string(alert.Labels["status"]))
	s.Equal(summary, string(alert.Annotations["summary"]))
	s.Equal(request, string(alert.Annotations["request"]))
	s.Equal(startsAt, alert.StartsAt)
	s.Equal(endsAt, alert.EndsAt)
	s.Equal("", alert.GeneratorURL)
}
<commit_msg>TST: Fix<commit_after>package service

import (
	"context"
	"fmt"
	"os/exec"
	"testing"
	"time"

	"github.com/docker/docker/client"
	"github.com/stretchr/testify/suite"
)

type AlertTestSuite struct {
	suite.Suite
	url          string
	alertService AlertServicer
	client       *client.Client
}

func TestAlertUnitTestSuite(t *testing.T) {
	suite.Run(t, new(AlertTestSuite))
}

func (s *AlertTestSuite) SetupSuite() {
	client, _ := NewDockerClientFromEnv()
	_, err := client.Info(context.Background())
	if err != nil {
		s.T().Skipf("Unable to connect to Docker Client")
	}
	s.url = "http://localhost:9093"
	s.alertService = NewAlertService(s.url, time.Second*15)
	s.client = client
}

func (s *AlertTestSuite) TearDownSuite() {
	s.client.Close()
}

func (s *AlertTestSuite) Test_SendAlert() {

	defer func() {
		cmd := "docker container rm -f am9093"
		exec.Command("/bin/sh", "-c", cmd).Output()
	}()
	cmd := `docker run --name am9093 -p 9093:9093 \
			-d prom/alertmanager:v0.13.0`
	_, err := exec.Command("/bin/sh", "-c", cmd).Output()
	if err != nil {
		s.T().Skipf(fmt.Sprintf("Unable to create alertmanager: %s", err.Error()))
		return
	}

	running := false
	// Wait for am to come online
	for i := 1; i <= 60; i++ {
		info, _ := s.client.ContainerInspect(context.Background(), "am9093")
		if info.State.Running {
			running = true
			break
		}
		time.Sleep(1 * time.Second)
	}
	if !running {
		s.T().Skipf(fmt.Sprintf("Alertmanager not created"))
		return
	}

	require := s.Require()
	serviceName := "web"
	alertname := "service_scaler"
	status := "success"
	summary := "Scaled web from 3 to 4 replicas"
	request := "Scale web with delta=1"

	err = s.alertService.Send(alertname, serviceName, request, status, summary)
	require.NoError(err)
	time.Sleep(1 * time.Second)

	alerts, err := FetchAlerts(s.url, alertname, status, serviceName)
	require.NoError(err)
	require.Equal(1, len(alerts))

	alert := alerts[0]
	s.Equal(alertname, string(alert.Labels["alertname"]))
	s.Equal(serviceName, string(alert.Labels["service"]))
	s.Equal(status, string(alert.Labels["status"]))
	s.Equal(summary, string(alert.Annotations["summary"]))
	s.Equal(request, string(alert.Annotations["request"]))
	s.Equal("", alert.GeneratorURL)
}

func (s *AlertTestSuite) Test_generateAlert() {

	serviceName := "web"
	alertname := "service_scaler"
	status := "success"
	summary := "Scaled web from 3 to 4 replicas"
	request := "Scale web with delta=1"
	startsAt := time.Now().UTC()
	timeout := time.Second
	endsAt := startsAt.Add(timeout)

	alert := generateAlert(alertname, serviceName, request, status, summary, startsAt, timeout)
	s.Require().NotNil(alert)
	s.Equal(alertname, string(alert.Labels["alertname"]))
	s.Equal(serviceName, string(alert.Labels["service"]))
	s.Equal(status, string(alert.Labels["status"]))
	s.Equal(summary, string(alert.Annotations["summary"]))
	s.Equal(request, string(alert.Annotations["request"]))
	s.Equal(startsAt, alert.StartsAt)
	s.Equal(endsAt, alert.EndsAt)
	s.Equal("", alert.GeneratorURL)
}
<|endoftext|>


<commit_before>package rel

import (
	"log"
)

type SelectManager struct {
	Engine Engine
	Ast    *SelectStatementNode
	Ctx    *SelectCoreNode
	BaseVisitable
}

func Select(visitables ...Visitable) *SelectManager {
	return NewSelectManager(RelEngine, &Table{}).Select(visitables...)
}

func NewSelectManager(engine Engine, table *Table) *SelectManager {
	if engine == nil {
		engine = RelEngine
	}
	stmt := NewSelectStatementNode()
	manager := SelectManager{
		Engine: engine,
		Ast:    stmt,
		Ctx:    stmt.Cores[len(stmt.Cores)-1],
	}
	// setup initial join source
	manager.From(table)
	return &manager
}

func (mgr *SelectManager) ToSql() string {
	return mgr.Engine.Visitor().Accept(mgr.Ast)
}

func (mgr *SelectManager) Project(visitables ...Visitable) *SelectManager {
	return mgr.Select(visitables...)
}

func (mgr *SelectManager) Select(visitables ...Visitable) *SelectManager {
	for _, selection := range visitables {
		if mgr.Ctx.Selections == nil {
			mgr.Ctx.Selections = &[]Visitable{}
		}

		*mgr.Ctx.Selections = append(*mgr.Ctx.Selections, selection)
	}
	return mgr
}

func (mgr *SelectManager) From(table interface{}) *SelectManager {
	var v Visitable
	switch t := table.(type) {
	case *Table:
		v = t
	case Table:
		v = &t
	case string:
		v = NewTable(t)
	}
	mgr.Ctx.Source.Left = v
	return mgr
}

func (mgr *SelectManager) As(name string) *TableAliasNode {
	return &TableAliasNode{
		Relation: &GroupingNode{Expr: []Visitable{mgr.Ast}},
		Name:     name,
	}
}

func (mgr *SelectManager) On(visitables ...Visitable) *SelectManager {
	right := mgr.Ctx.Source.Right

	if len(right) > 0 {
		last := right[len(right)-1]
		switch val := last.(type) {
		case *InnerJoinNode:
			val.Right = mgr.NewOnNode(mgr.collapse(visitables...))
		case *OuterJoinNode:
			val.Right = mgr.NewOnNode(mgr.collapse(visitables...))
		default:
			log.Fatalf("Unable to call On with input type %T", val)
		}
	}

	return mgr
}

func (mgr *SelectManager) Using(str string) *SelectManager {
	right := mgr.Ctx.Source.Right

	if len(right) > 0 {
		last := right[len(right)-1]
		switch val := last.(type) {
		case *InnerJoinNode:
			val.Right = &UsingNode{Expr: &QuotedNode{Raw: str}}
		case *OuterJoinNode:
			val.Right = &UsingNode{Expr: &QuotedNode{Raw: str}}
		default:
			log.Fatalf("Unable to call On with input type %T", val)
		}
	}

	return mgr
}

func (mgr *SelectManager) Join(visitable Visitable) *SelectManager {
	return mgr.InnerJoin(visitable)
}

func (mgr *SelectManager) InnerJoin(visitable Visitable) *SelectManager {
	mgr.Ctx.Source.Right = append(mgr.Ctx.Source.Right, &InnerJoinNode{Left: visitable})
	return mgr
}

func (mgr *SelectManager) OuterJoin(visitable Visitable) *SelectManager {
	mgr.Ctx.Source.Right = append(mgr.Ctx.Source.Right, &OuterJoinNode{Left: visitable})
	return mgr
}

func (mgr *SelectManager) Lock(node SqlLiteralNode) *SelectManager {
	mgr.Ast.Lock = NewLockNode(node)
	return mgr
}

func (mgr *SelectManager) LockForUpdate() *SelectManager {
	mgr.Ast.Lock = NewLockNode(Sql("FOR UPDATE"))
	return mgr
}

func (mgr *SelectManager) Take(i int) *SelectManager {
	mgr.Ast.Limit = NewLimitNode(Sql(i))
	return mgr
}

func (mgr *SelectManager) Exists() *ExistsNode {
	return NewExistsNode(mgr.Ast)
}

func (mgr *SelectManager) Order(visitables ...Visitable) *SelectManager {
	if len(visitables) > 0 {
		if mgr.Ast.Orders == nil {
			mgr.Ast.Orders = &[]Visitable{}
		}
		for _, v := range visitables {
			*mgr.Ast.Orders = append(*mgr.Ast.Orders, v)
		}
	}
	return mgr
}

func (mgr *SelectManager) Where(visitable Visitable) *SelectManager {
	if mgr.Ctx.Wheres == nil {
		mgr.Ctx.Wheres = &[]Visitable{}
	}

	if expr, ok := visitable.(SelectManager); ok {
		*mgr.Ctx.Wheres = append(*mgr.Ctx.Wheres, expr.Ast)
	} else {
		*mgr.Ctx.Wheres = append(*mgr.Ctx.Wheres, visitable)
	}

	return mgr
}

func (mgr *SelectManager) GroupBy(visitables ...Visitable) *SelectManager {
	return mgr.Group(visitables...)
}

func (mgr *SelectManager) Group(visitables ...Visitable) *SelectManager {
	if len(visitables) > 0 {
		if mgr.Ctx.Groups == nil {
			mgr.Ctx.Groups = &[]Visitable{}
		}
		for _, v := range visitables {
			*mgr.Ctx.Groups = append(*mgr.Ctx.Groups, NewGroupNode(v))
		}
	}
	return mgr
}

func (mgr *SelectManager) Intersect(stmt1 Visitable, stmt2 Visitable) *MultiStatementManager {
	return NewMultiStatementManager(mgr.Engine).Intersect(stmt1, stmt2)
}

func (mgr *SelectManager) Union(stmt1 Visitable, stmt2 Visitable) *MultiStatementManager {
	return NewMultiStatementManager(mgr.Engine).Union(stmt1, stmt2)
}

func (mgr *SelectManager) UnionAll(stmt1 Visitable, stmt2 Visitable) *MultiStatementManager {
	return NewMultiStatementManager(mgr.Engine).UnionAll(stmt1, stmt2)
}

func (mgr *SelectManager) Except(stmt1 Visitable, stmt2 Visitable) *MultiStatementManager {
	return NewMultiStatementManager(mgr.Engine).Except(stmt1, stmt2)
}

func (mgr *SelectManager) Skip(i int) *SelectManager {
	mgr.Ast.Offset = NewOffsetNode(Sql(i))
	return mgr
}

func (mgr *SelectManager) Offset(i int) *SelectManager {
	return mgr.Skip(i)
}

func (mgr *SelectManager) Having(visitables ...Visitable) *SelectManager {
	mgr.Ctx.Having = NewHavingNode(mgr.collapse(visitables...))
	return mgr
}

func (mgr *SelectManager) Distinct() *SelectManager {
	mgr.Ctx.SetQuantifier = &DistinctNode{}
	return mgr
}

func (mgr *SelectManager) NotDistinct() *SelectManager {
	mgr.Ctx.SetQuantifier = nil
	return mgr
}

func (mgr *SelectManager) With(visitable Visitable) *SelectManager {
	mgr.Ast.With = &WithNode{Expr: visitable}
	return mgr
}

func (mgr *SelectManager) WithRecursive(visitable Visitable) *SelectManager {
	mgr.Ast.With = &WithRecursiveNode{Expr: visitable}
	return mgr
}

func (mgr *SelectManager) Window(node SqlLiteralNode) *NamedWindowNode {
	if mgr.Ctx.Windows == nil {
		mgr.Ctx.Windows = &[]Visitable{}
	}
	window := &NamedWindowNode{Name: node}
	*mgr.Ctx.Windows = append(*mgr.Ctx.Windows, window)
	return window
}

func (mgr *SelectManager) collapse(visitables ...Visitable) Visitable {
	var v Visitable

	// use the first Node if there is only one
	// else create and And node
	if len(visitables) == 1 {
		v = visitables[0]
	} else {
		v = mgr.NewAndNode(visitables...)
	}
	return v
}
<commit_msg>SelectManager should have a Limit method<commit_after>package rel

import (
	"log"
)

type SelectManager struct {
	Engine Engine
	Ast    *SelectStatementNode
	Ctx    *SelectCoreNode
	BaseVisitable
}

func Select(visitables ...Visitable) *SelectManager {
	return NewSelectManager(RelEngine, &Table{}).Select(visitables...)
}

func NewSelectManager(engine Engine, table *Table) *SelectManager {
	if engine == nil {
		engine = RelEngine
	}
	stmt := NewSelectStatementNode()
	manager := SelectManager{
		Engine: engine,
		Ast:    stmt,
		Ctx:    stmt.Cores[len(stmt.Cores)-1],
	}
	// setup initial join source
	manager.From(table)
	return &manager
}

func (mgr *SelectManager) ToSql() string {
	return mgr.Engine.Visitor().Accept(mgr.Ast)
}

func (mgr *SelectManager) Project(visitables ...Visitable) *SelectManager {
	return mgr.Select(visitables...)
}

func (mgr *SelectManager) Select(visitables ...Visitable) *SelectManager {
	for _, selection := range visitables {
		if mgr.Ctx.Selections == nil {
			mgr.Ctx.Selections = &[]Visitable{}
		}

		*mgr.Ctx.Selections = append(*mgr.Ctx.Selections, selection)
	}
	return mgr
}

func (mgr *SelectManager) From(table interface{}) *SelectManager {
	var v Visitable
	switch t := table.(type) {
	case *Table:
		v = t
	case Table:
		v = &t
	case string:
		v = NewTable(t)
	}
	mgr.Ctx.Source.Left = v
	return mgr
}

func (mgr *SelectManager) As(name string) *TableAliasNode {
	return &TableAliasNode{
		Relation: &GroupingNode{Expr: []Visitable{mgr.Ast}},
		Name:     name,
	}
}

func (mgr *SelectManager) On(visitables ...Visitable) *SelectManager {
	right := mgr.Ctx.Source.Right

	if len(right) > 0 {
		last := right[len(right)-1]
		switch val := last.(type) {
		case *InnerJoinNode:
			val.Right = mgr.NewOnNode(mgr.collapse(visitables...))
		case *OuterJoinNode:
			val.Right = mgr.NewOnNode(mgr.collapse(visitables...))
		default:
			log.Fatalf("Unable to call On with input type %T", val)
		}
	}

	return mgr
}

func (mgr *SelectManager) Using(str string) *SelectManager {
	right := mgr.Ctx.Source.Right

	if len(right) > 0 {
		last := right[len(right)-1]
		switch val := last.(type) {
		case *InnerJoinNode:
			val.Right = &UsingNode{Expr: &QuotedNode{Raw: str}}
		case *OuterJoinNode:
			val.Right = &UsingNode{Expr: &QuotedNode{Raw: str}}
		default:
			log.Fatalf("Unable to call On with input type %T", val)
		}
	}

	return mgr
}

func (mgr *SelectManager) Join(visitable Visitable) *SelectManager {
	return mgr.InnerJoin(visitable)
}

func (mgr *SelectManager) InnerJoin(visitable Visitable) *SelectManager {
	mgr.Ctx.Source.Right = append(mgr.Ctx.Source.Right, &InnerJoinNode{Left: visitable})
	return mgr
}

func (mgr *SelectManager) OuterJoin(visitable Visitable) *SelectManager {
	mgr.Ctx.Source.Right = append(mgr.Ctx.Source.Right, &OuterJoinNode{Left: visitable})
	return mgr
}

func (mgr *SelectManager) Lock(node SqlLiteralNode) *SelectManager {
	mgr.Ast.Lock = NewLockNode(node)
	return mgr
}

func (mgr *SelectManager) LockForUpdate() *SelectManager {
	mgr.Ast.Lock = NewLockNode(Sql("FOR UPDATE"))
	return mgr
}

func (mgr *SelectManager) Take(i int) *SelectManager {
	return mgr.Limit(i)
}

func (mgr *SelectManager) Limit(i int) *SelectManager {
	mgr.Ast.Limit = NewLimitNode(Sql(i))
	return mgr
}

func (mgr *SelectManager) Exists() *ExistsNode {
	return NewExistsNode(mgr.Ast)
}

func (mgr *SelectManager) Order(visitables ...Visitable) *SelectManager {
	if len(visitables) > 0 {
		if mgr.Ast.Orders == nil {
			mgr.Ast.Orders = &[]Visitable{}
		}
		for _, v := range visitables {
			*mgr.Ast.Orders = append(*mgr.Ast.Orders, v)
		}
	}
	return mgr
}

func (mgr *SelectManager) Where(visitable Visitable) *SelectManager {
	if mgr.Ctx.Wheres == nil {
		mgr.Ctx.Wheres = &[]Visitable{}
	}

	if expr, ok := visitable.(SelectManager); ok {
		*mgr.Ctx.Wheres = append(*mgr.Ctx.Wheres, expr.Ast)
	} else {
		*mgr.Ctx.Wheres = append(*mgr.Ctx.Wheres, visitable)
	}

	return mgr
}

func (mgr *SelectManager) GroupBy(visitables ...Visitable) *SelectManager {
	return mgr.Group(visitables...)
}

func (mgr *SelectManager) Group(visitables ...Visitable) *SelectManager {
	if len(visitables) > 0 {
		if mgr.Ctx.Groups == nil {
			mgr.Ctx.Groups = &[]Visitable{}
		}
		for _, v := range visitables {
			*mgr.Ctx.Groups = append(*mgr.Ctx.Groups, NewGroupNode(v))
		}
	}
	return mgr
}

func (mgr *SelectManager) Intersect(stmt1 Visitable, stmt2 Visitable) *MultiStatementManager {
	return NewMultiStatementManager(mgr.Engine).Intersect(stmt1, stmt2)
}

func (mgr *SelectManager) Union(stmt1 Visitable, stmt2 Visitable) *MultiStatementManager {
	return NewMultiStatementManager(mgr.Engine).Union(stmt1, stmt2)
}

func (mgr *SelectManager) UnionAll(stmt1 Visitable, stmt2 Visitable) *MultiStatementManager {
	return NewMultiStatementManager(mgr.Engine).UnionAll(stmt1, stmt2)
}

func (mgr *SelectManager) Except(stmt1 Visitable, stmt2 Visitable) *MultiStatementManager {
	return NewMultiStatementManager(mgr.Engine).Except(stmt1, stmt2)
}

func (mgr *SelectManager) Skip(i int) *SelectManager {
	mgr.Ast.Offset = NewOffsetNode(Sql(i))
	return mgr
}

func (mgr *SelectManager) Offset(i int) *SelectManager {
	return mgr.Skip(i)
}

func (mgr *SelectManager) Having(visitables ...Visitable) *SelectManager {
	mgr.Ctx.Having = NewHavingNode(mgr.collapse(visitables...))
	return mgr
}

func (mgr *SelectManager) Distinct() *SelectManager {
	mgr.Ctx.SetQuantifier = &DistinctNode{}
	return mgr
}

func (mgr *SelectManager) NotDistinct() *SelectManager {
	mgr.Ctx.SetQuantifier = nil
	return mgr
}

func (mgr *SelectManager) With(visitable Visitable) *SelectManager {
	mgr.Ast.With = &WithNode{Expr: visitable}
	return mgr
}

func (mgr *SelectManager) WithRecursive(visitable Visitable) *SelectManager {
	mgr.Ast.With = &WithRecursiveNode{Expr: visitable}
	return mgr
}

func (mgr *SelectManager) Window(node SqlLiteralNode) *NamedWindowNode {
	if mgr.Ctx.Windows == nil {
		mgr.Ctx.Windows = &[]Visitable{}
	}
	window := &NamedWindowNode{Name: node}
	*mgr.Ctx.Windows = append(*mgr.Ctx.Windows, window)
	return window
}

func (mgr *SelectManager) collapse(visitables ...Visitable) Visitable {
	var v Visitable

	// use the first Node if there is only one
	// else create and And node
	if len(visitables) == 1 {
		v = visitables[0]
	} else {
		v = mgr.NewAndNode(visitables...)
	}
	return v
}
<|endoftext|>