<commit_before>/*
Copyright 2014 The Camlistore 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 search_test

import (
	"fmt"
	"testing"

	"camlistore.org/pkg/blob"
	"camlistore.org/pkg/index"
	"camlistore.org/pkg/index/indextest"
	"camlistore.org/pkg/search"
	"camlistore.org/pkg/test"
	"camlistore.org/pkg/types/camtypes"

	"golang.org/x/net/context"
)

func addPermanode(fi *test.FakeIndex, pnStr string, attrs ...string) {
	pn := blob.MustParse(pnStr)
	fi.AddMeta(pn, "permanode", 123)
	for len(attrs) > 0 {
		k, v := attrs[0], attrs[1]
		attrs = attrs[2:]
		fi.AddClaim(owner, pn, "add-attribute", k, v)
	}
}

func addFileWithLocation(fi *test.FakeIndex, fileStr string, lat, long float64) {
	fileRef := blob.MustParse(fileStr)
	fi.AddFileLocation(fileRef, camtypes.Location{Latitude: lat, Longitude: long})
	fi.AddMeta(fileRef, "file", 123)
}

func searchDescribeSetup(fi *test.FakeIndex) index.Interface {
	addPermanode(fi, "abc-123",
		"camliContent", "abc-123c",
		"camliImageContent", "abc-888",
	)
	addPermanode(fi, "abc-123c",
		"camliContent", "abc-123cc",
		"camliImageContent", "abc-123c1",
	)
	addPermanode(fi, "abc-123c1",
		"some", "image",
	)
	addPermanode(fi, "abc-123cc",
		"name", "leaf",
	)
	addPermanode(fi, "abc-888",
		"camliContent", "abc-8881",
	)
	addPermanode(fi, "abc-8881",
		"name", "leaf8881",
	)

	addPermanode(fi, "fourcheckin-0",
		"camliNodeType", "foursquare.com:checkin",
		"foursquareVenuePermanode", "fourvenue-123",
	)
	addPermanode(fi, "fourvenue-123",
		"camliNodeType", "foursquare.com:venue",
		"camliPath:photos", "venuepicset-123",
		"latitude", "12",
		"longitude", "34",
	)
	addPermanode(fi, "venuepicset-123",
		"camliPath:1.jpg", "venuepic-1",
	)
	addPermanode(fi, "venuepic-1",
		"camliContent", "somevenuepic-0",
	)
	addPermanode(fi, "somevenuepic-0",
		"foo", "bar",
	)
	addPermanode(fi, "venuepic-2",
		"camliContent", "somevenuepic-2",
	)
	addPermanode(fi, "somevenuepic-2",
		"foo", "baz",
	)

	addPermanode(fi, "homedir-0",
		"camliPath:subdir.1", "homedir-1",
	)
	addPermanode(fi, "homedir-1",
		"camliPath:subdir.2", "homedir-2",
	)
	addPermanode(fi, "homedir-2",
		"foo", "bar",
	)

	addPermanode(fi, "set-0",
		"camliMember", "venuepic-1",
		"camliMember", "venuepic-2",
	)

	addFileWithLocation(fi, "filewithloc-0", 45, 56)
	addPermanode(fi, "location-0",
		"camliContent", "filewithloc-0",
	)

	addPermanode(fi, "locationpriority-1",
		"latitude", "67",
		"longitude", "78",
		"camliNodeType", "foursquare.com:checkin",
		"foursquareVenuePermanode", "fourvenue-123",
		"camliContent", "filewithloc-0",
	)

	addPermanode(fi, "locationpriority-2",
		"camliNodeType", "foursquare.com:checkin",
		"foursquareVenuePermanode", "fourvenue-123",
		"camliContent", "filewithloc-0",
	)

	addPermanode(fi, "locationoverride-1",
		"latitude", "67",
		"longitude", "78",
		"camliContent", "filewithloc-0",
	)

	addPermanode(fi, "locationoverride-2",
		"latitude", "67",
		"longitude", "78",
		"camliNodeType", "foursquare.com:checkin",
		"foursquareVenuePermanode", "fourvenue-123",
	)

	return fi
}

var searchDescribeTests = []handlerTest{
	{
		name:     "null",
		postBody: marshalJSON(&search.DescribeRequest{}),
		want: jmap(&search.DescribeResponse{
			Meta: search.MetaMap{},
		}),
	},

	{
		name: "single",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRef: blob.MustParse("abc-123"),
		}),
		wantDescribed: []string{"abc-123"},
	},

	{
		name: "follow all camliContent",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRef: blob.MustParse("abc-123"),
			Rules: []*search.DescribeRule{
				{
					Attrs: []string{"camliContent"},
				},
			},
		}),
		wantDescribed: []string{"abc-123", "abc-123c", "abc-123cc"},
	},

	{
		name: "follow only root camliContent",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRef: blob.MustParse("abc-123"),
			Rules: []*search.DescribeRule{
				{
					IfResultRoot: true,
					Attrs:        []string{"camliContent"},
				},
			},
		}),
		wantDescribed: []string{"abc-123", "abc-123c"},
	},

	{
		name: "follow all root, substring",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRef: blob.MustParse("abc-123"),
			Rules: []*search.DescribeRule{
				{
					IfResultRoot: true,
					Attrs:        []string{"camli*"},
				},
			},
		}),
		wantDescribed: []string{"abc-123", "abc-123c", "abc-888"},
	},

	{
		name: "two rules, two attrs",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRef: blob.MustParse("abc-123"),
			Rules: []*search.DescribeRule{
				{
					IfResultRoot: true,
					Attrs:        []string{"camliContent", "camliImageContent"},
				},
				{
					Attrs: []string{"camliContent"},
				},
			},
		}),
		wantDescribed: []string{"abc-123", "abc-123c", "abc-123cc", "abc-888", "abc-8881"},
	},

	{
		name: "foursquare venue photos, but not recursive camliPath explosion",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRefs: []blob.Ref{
				blob.MustParse("homedir-0"),
				blob.MustParse("fourcheckin-0"),
			},
			Rules: []*search.DescribeRule{
				{
					Attrs: []string{"camliContent", "camliContentImage"},
				},
				{
					IfCamliNodeType: "foursquare.com:checkin",
					Attrs:           []string{"foursquareVenuePermanode"},
				},
				{
					IfCamliNodeType: "foursquare.com:venue",
					Attrs:           []string{"camliPath:photos"},
					Rules: []*search.DescribeRule{
						{
							Attrs: []string{"camliPath:*"},
						},
					},
				},
			},
		}),
		wantDescribed: []string{"homedir-0", "fourcheckin-0", "fourvenue-123", "venuepicset-123", "venuepic-1", "somevenuepic-0"},
	},

	{
		name: "home dirs forever",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRefs: []blob.Ref{
				blob.MustParse("homedir-0"),
			},
			Rules: []*search.DescribeRule{
				{
					Attrs: []string{"camliPath:*"},
				},
			},
		}),
		wantDescribed: []string{"homedir-0", "homedir-1", "homedir-2"},
	},

	{
		name: "find members",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRef: blob.MustParse("set-0"),
			Rules: []*search.DescribeRule{
				{
					IfResultRoot: true,
					Attrs:        []string{"camliMember"},
					Rules: []*search.DescribeRule{
						{Attrs: []string{"camliContent"}},
					},
				},
			},
		}),
		wantDescribed: []string{"set-0", "venuepic-1", "venuepic-2", "somevenuepic-0", "somevenuepic-2"},
	},
}

func init() {
	checkNoDups("searchDescribeTests", searchDescribeTests)
}

func TestSearchDescribe(t *testing.T) {
	for _, ht := range searchDescribeTests {
		if ht.setup == nil {
			ht.setup = searchDescribeSetup
		}
		if ht.query == "" {
			ht.query = "describe"
		}
		ht.test(t)
	}
}

// should be run with -race
func TestDescribeRace(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping in short mode")
	}
	idx := index.NewMemoryIndex()
	idxd := indextest.NewIndexDeps(idx)
	idxd.Fataler = t
	corpus, err := idxd.Index.KeepInMemory()
	if err != nil {
		t.Fatalf("error slurping index to memory: %v", err)
	}
	h := search.NewHandler(idx, idxd.SignerBlobRef)
	h.SetCorpus(corpus)
	donec := make(chan struct{})
	headstart := 500
	blobrefs := make([]blob.Ref, headstart)
	headstartc := make(chan struct{})
	go func() {
		for i := 0; i < headstart*2; i++ {
			nth := fmt.Sprintf("%d", i)
			// No need to lock the index here. It is already done within NewPlannedPermanode,
			// because it calls idxd.Index.ReceiveBlob.
			pn := idxd.NewPlannedPermanode(nth)
			idxd.SetAttribute(pn, "tag", nth)
			if i > headstart {
				continue
			}
			if i == headstart {
				headstartc <- struct{}{}
				continue
			}
			blobrefs[i] = pn
		}
	}()
	<-headstartc
	ctx := context.Background()
	go func() {
		for i := 0; i < headstart; i++ {
			br := blobrefs[i]
			res, err := h.Describe(ctx, &search.DescribeRequest{
				BlobRef: br,
				Depth:   1,
			})
			if err != nil {
				t.Fatal(err)
			}
			_, ok := res.Meta[br.String()]
			if !ok {
				t.Errorf("permanode %v wasn't in Describe response", br)
			}
		}
		donec <- struct{}{}
	}()
	<-donec
}

func TestDescribeLocation(t *testing.T) {
	tests := []struct {
		ref       string
		lat, long float64
		hasNoLoc  bool
	}{
		{ref: "filewithloc-0", lat: 45, long: 56},
		{ref: "location-0", lat: 45, long: 56},
		{ref: "locationpriority-1", lat: 67, long: 78},
		{ref: "locationpriority-2", lat: 12, long: 34},
		{ref: "locationoverride-1", lat: 67, long: 78},
		{ref: "locationoverride-2", lat: 67, long: 78},
		{ref: "homedir-0", hasNoLoc: true},
	}

	ix := searchDescribeSetup(test.NewFakeIndex())
	ctx := context.Background()
	h := search.NewHandler(ix, owner)

	ix.RLock()
	defer ix.RUnlock()

	for _, tt := range tests {
		var err error
		br := blob.MustParse(tt.ref)
		res, err := h.Describe(ctx, &search.DescribeRequest{
			BlobRef: br,
			Depth:   1,
		})
		if err != nil {
			t.Errorf("Describe for %v failed: %v", br, err)
			continue
		}
		db := res.Meta[br.String()]
		if db == nil {
			t.Errorf("Describe result for %v is missing", br)
			continue
		}
		loc := db.Location
		if tt.hasNoLoc {
			if loc != nil {
				t.Errorf("got location for %v, should have no location", br)
			}
		} else {
			if loc == nil {
				t.Errorf("no location in result for %v", br)
				continue
			}
			if loc.Latitude != tt.lat || loc.Longitude != tt.long {
				t.Errorf("location for %v invalid, got %f,%f want %f,%f",
					tt.ref, loc.Latitude, loc.Longitude, tt.lat, tt.long)
			}
		}
	}
}
<commit_msg>pkg/search: add test to illustrate issue 881<commit_after>/*
Copyright 2014 The Camlistore 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 search_test

import (
	"fmt"
	"testing"

	"camlistore.org/pkg/blob"
	"camlistore.org/pkg/index"
	"camlistore.org/pkg/index/indextest"
	"camlistore.org/pkg/search"
	"camlistore.org/pkg/test"
	"camlistore.org/pkg/types/camtypes"

	"golang.org/x/net/context"
)

func addPermanode(fi *test.FakeIndex, pnStr string, attrs ...string) {
	pn := blob.MustParse(pnStr)
	fi.AddMeta(pn, "permanode", 123)
	for len(attrs) > 0 {
		k, v := attrs[0], attrs[1]
		attrs = attrs[2:]
		fi.AddClaim(owner, pn, "add-attribute", k, v)
	}
}

func addFileWithLocation(fi *test.FakeIndex, fileStr string, lat, long float64) {
	fileRef := blob.MustParse(fileStr)
	fi.AddFileLocation(fileRef, camtypes.Location{Latitude: lat, Longitude: long})
	fi.AddMeta(fileRef, "file", 123)
}

func searchDescribeSetup(fi *test.FakeIndex) index.Interface {
	addPermanode(fi, "abc-123",
		"camliContent", "abc-123c",
		"camliImageContent", "abc-888",
	)
	addPermanode(fi, "abc-123c",
		"camliContent", "abc-123cc",
		"camliImageContent", "abc-123c1",
	)
	addPermanode(fi, "abc-123c1",
		"some", "image",
	)
	addPermanode(fi, "abc-123cc",
		"name", "leaf",
	)
	addPermanode(fi, "abc-888",
		"camliContent", "abc-8881",
	)
	addPermanode(fi, "abc-8881",
		"name", "leaf8881",
	)

	addPermanode(fi, "fourcheckin-0",
		"camliNodeType", "foursquare.com:checkin",
		"foursquareVenuePermanode", "fourvenue-123",
	)
	addPermanode(fi, "fourvenue-123",
		"camliNodeType", "foursquare.com:venue",
		"camliPath:photos", "venuepicset-123",
		"latitude", "12",
		"longitude", "34",
	)
	addPermanode(fi, "venuepicset-123",
		"camliPath:1.jpg", "venuepic-1",
	)
	addPermanode(fi, "venuepic-1",
		"camliContent", "somevenuepic-0",
	)
	addPermanode(fi, "somevenuepic-0",
		"foo", "bar",
	)
	addPermanode(fi, "venuepic-2",
		"camliContent", "somevenuepic-2",
	)
	addPermanode(fi, "somevenuepic-2",
		"foo", "baz",
	)

	addPermanode(fi, "homedir-0",
		"camliPath:subdir.1", "homedir-1",
	)
	addPermanode(fi, "homedir-1",
		"camliPath:subdir.2", "homedir-2",
	)
	addPermanode(fi, "homedir-2",
		"foo", "bar",
	)

	addPermanode(fi, "set-0",
		"camliMember", "venuepic-1",
		"camliMember", "venuepic-2",
	)

	addFileWithLocation(fi, "filewithloc-0", 45, 56)
	addPermanode(fi, "location-0",
		"camliContent", "filewithloc-0",
	)

	addPermanode(fi, "locationpriority-1",
		"latitude", "67",
		"longitude", "78",
		"camliNodeType", "foursquare.com:checkin",
		"foursquareVenuePermanode", "fourvenue-123",
		"camliContent", "filewithloc-0",
	)

	addPermanode(fi, "locationpriority-2",
		"camliNodeType", "foursquare.com:checkin",
		"foursquareVenuePermanode", "fourvenue-123",
		"camliContent", "filewithloc-0",
	)

	addPermanode(fi, "locationoverride-1",
		"latitude", "67",
		"longitude", "78",
		"camliContent", "filewithloc-0",
	)

	addPermanode(fi, "locationoverride-2",
		"latitude", "67",
		"longitude", "78",
		"camliNodeType", "foursquare.com:checkin",
		"foursquareVenuePermanode", "fourvenue-123",
	)

	return fi
}

var searchDescribeTests = []handlerTest{
	{
		name:     "null",
		postBody: marshalJSON(&search.DescribeRequest{}),
		want: jmap(&search.DescribeResponse{
			Meta: search.MetaMap{},
		}),
	},

	{
		name: "single",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRef: blob.MustParse("abc-123"),
		}),
		wantDescribed: []string{"abc-123"},
	},

	{
		name: "follow all camliContent",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRef: blob.MustParse("abc-123"),
			Rules: []*search.DescribeRule{
				{
					Attrs: []string{"camliContent"},
				},
			},
		}),
		wantDescribed: []string{"abc-123", "abc-123c", "abc-123cc"},
	},

	{
		name: "follow only root camliContent",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRef: blob.MustParse("abc-123"),
			Rules: []*search.DescribeRule{
				{
					IfResultRoot: true,
					Attrs:        []string{"camliContent"},
				},
			},
		}),
		wantDescribed: []string{"abc-123", "abc-123c"},
	},

	{
		name: "follow all root, substring",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRef: blob.MustParse("abc-123"),
			Rules: []*search.DescribeRule{
				{
					IfResultRoot: true,
					Attrs:        []string{"camli*"},
				},
			},
		}),
		wantDescribed: []string{"abc-123", "abc-123c", "abc-888"},
	},

	{
		name: "two rules, two attrs",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRef: blob.MustParse("abc-123"),
			Rules: []*search.DescribeRule{
				{
					IfResultRoot: true,
					Attrs:        []string{"camliContent", "camliImageContent"},
				},
				{
					Attrs: []string{"camliContent"},
				},
			},
		}),
		wantDescribed: []string{"abc-123", "abc-123c", "abc-123cc", "abc-888", "abc-8881"},
	},

	{
		name: "foursquare venue photos, but not recursive camliPath explosion",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRefs: []blob.Ref{
				blob.MustParse("homedir-0"),
				blob.MustParse("fourcheckin-0"),
			},
			Rules: []*search.DescribeRule{
				{
					Attrs: []string{"camliContent", "camliContentImage"},
				},
				{
					IfCamliNodeType: "foursquare.com:checkin",
					Attrs:           []string{"foursquareVenuePermanode"},
				},
				{
					IfCamliNodeType: "foursquare.com:venue",
					Attrs:           []string{"camliPath:photos"},
					Rules: []*search.DescribeRule{
						{
							Attrs: []string{"camliPath:*"},
						},
					},
				},
			},
		}),
		wantDescribed: []string{"homedir-0", "fourcheckin-0", "fourvenue-123", "venuepicset-123", "venuepic-1", "somevenuepic-0"},
	},

	{
		name: "home dirs forever",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRefs: []blob.Ref{
				blob.MustParse("homedir-0"),
			},
			Rules: []*search.DescribeRule{
				{
					Attrs: []string{"camliPath:*"},
				},
			},
		}),
		wantDescribed: []string{"homedir-0", "homedir-1", "homedir-2"},
	},

	{
		name: "find members",
		postBody: marshalJSON(&search.DescribeRequest{
			BlobRef: blob.MustParse("set-0"),
			Rules: []*search.DescribeRule{
				{
					IfResultRoot: true,
					Attrs:        []string{"camliMember"},
					Rules: []*search.DescribeRule{
						{Attrs: []string{"camliContent"}},
					},
				},
			},
		}),
		wantDescribed: []string{"set-0", "venuepic-1", "venuepic-2", "somevenuepic-0", "somevenuepic-2"},
	},
}

func init() {
	checkNoDups("searchDescribeTests", searchDescribeTests)
}

func TestSearchDescribe(t *testing.T) {
	for _, ht := range searchDescribeTests {
		if ht.setup == nil {
			ht.setup = searchDescribeSetup
		}
		if ht.query == "" {
			ht.query = "describe"
		}
		ht.test(t)
	}
}

// should be run with -race
func TestDescribeRace(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping in short mode")
	}
	idx := index.NewMemoryIndex()
	idxd := indextest.NewIndexDeps(idx)
	idxd.Fataler = t
	corpus, err := idxd.Index.KeepInMemory()
	if err != nil {
		t.Fatalf("error slurping index to memory: %v", err)
	}
	h := search.NewHandler(idx, idxd.SignerBlobRef)
	h.SetCorpus(corpus)
	donec := make(chan struct{})
	headstart := 500
	blobrefs := make([]blob.Ref, headstart)
	headstartc := make(chan struct{})
	go func() {
		for i := 0; i < headstart*2; i++ {
			nth := fmt.Sprintf("%d", i)
			// No need to lock the index here. It is already done within NewPlannedPermanode,
			// because it calls idxd.Index.ReceiveBlob.
			pn := idxd.NewPlannedPermanode(nth)
			idxd.SetAttribute(pn, "tag", nth)
			if i > headstart {
				continue
			}
			if i == headstart {
				headstartc <- struct{}{}
				continue
			}
			blobrefs[i] = pn
		}
	}()
	<-headstartc
	ctx := context.Background()
	go func() {
		for i := 0; i < headstart; i++ {
			br := blobrefs[i]
			res, err := h.Describe(ctx, &search.DescribeRequest{
				BlobRef: br,
				Depth:   1,
			})
			if err != nil {
				t.Fatal(err)
			}
			_, ok := res.Meta[br.String()]
			if !ok {
				t.Errorf("permanode %v wasn't in Describe response", br)
			}
		}
		donec <- struct{}{}
	}()
	<-donec
}

func TestDescribeLocation(t *testing.T) {
	tests := []struct {
		ref       string
		lat, long float64
		hasNoLoc  bool
	}{
		{ref: "filewithloc-0", lat: 45, long: 56},
		{ref: "location-0", lat: 45, long: 56},
		{ref: "locationpriority-1", lat: 67, long: 78},
		{ref: "locationpriority-2", lat: 12, long: 34},
		{ref: "locationoverride-1", lat: 67, long: 78},
		{ref: "locationoverride-2", lat: 67, long: 78},
		{ref: "homedir-0", hasNoLoc: true},
	}

	ix := searchDescribeSetup(test.NewFakeIndex())
	ctx := context.Background()
	h := search.NewHandler(ix, owner)

	ix.RLock()
	defer ix.RUnlock()

	for _, tt := range tests {
		var err error
		br := blob.MustParse(tt.ref)
		res, err := h.Describe(ctx, &search.DescribeRequest{
			BlobRef: br,
			Depth:   1,
		})
		if err != nil {
			t.Errorf("Describe for %v failed: %v", br, err)
			continue
		}
		db := res.Meta[br.String()]
		if db == nil {
			t.Errorf("Describe result for %v is missing", br)
			continue
		}
		loc := db.Location
		if tt.hasNoLoc {
			if loc != nil {
				t.Errorf("got location for %v, should have no location", br)
			}
		} else {
			if loc == nil {
				t.Errorf("no location in result for %v", br)
				continue
			}
			if loc.Latitude != tt.lat || loc.Longitude != tt.long {
				t.Errorf("location for %v invalid, got %f,%f want %f,%f",
					tt.ref, loc.Latitude, loc.Longitude, tt.lat, tt.long)
			}
		}
	}
}

// To make sure we don't regress into issue 881: i.e. a permanode with no attr
// should not lead us to call index.claimsIntfAttrValue with a nil claims argument.
func TestDescribePermNoAttr(t *testing.T) {
	ix := index.NewMemoryIndex()
	ctx := context.Background()
	h := search.NewHandler(ix, owner)
	corpus, err := ix.KeepInMemory()
	if err != nil {
		t.Fatal(err)
	}
	h.SetCorpus(corpus)
	id := indextest.NewIndexDeps(ix)
	br := id.NewPlannedPermanode("noattr-0")

	ix.RLock()
	defer ix.RUnlock()

	res, err := h.Describe(ctx, &search.DescribeRequest{
		BlobRef: br,
		Depth:   1,
	})
	if err != nil {
		t.Fatalf("Describe for %v failed: %v", br, err)
	}
	db := res.Meta[br.String()]
	if db == nil {
		t.Fatalf("Describe result for %v is missing", br)
	}
}
<|endoftext|>


<commit_before>package persist

import (
	"crypto/rand"
	"encoding/base32"
	"errors"
	"os"
	"path/filepath"
)

const (
	// persistDir defines the folder that is used for testing the persist
	// package.
	persistDir = "persist"
)

var (
	// ErrBadVersion indicates that the version number of the file is not
	// compatible with the current codebase.
	ErrBadVersion = errors.New("incompatible version")

	// ErrBadHeader indicates that the file opened is not the file that was
	// expected.
	ErrBadHeader = errors.New("wrong header")
)

// Metadata contains the header and version of the data being stored.
type Metadata struct {
	Header, Version string
}

// RandomSuffix returns a 20 character base32 suffix for a filename. There are
// 100 bits of entropy, and a very low probability of colliding with existing
// files unintentionally.
func RandomSuffix() string {
	randBytes := make([]byte, 20)
	rand.Read(randBytes)
	str := base32.StdEncoding.EncodeToString(randBytes)
	return str[:20]
}

// A safeFile is a file that is stored under a temporary filename. When Commit
// is called, the file is renamed to its "final" filename. This allows for
// atomic updating of files; otherwise, an unexpected shutdown could leave a
// valuable file in a corrupted state. Callers must still Close the file handle
// as usual.
type safeFile struct {
	*os.File
	finalName string
}

// Commit syncs the file and then renames it to the intended final filename.
// Writing to the file after calling Commit will succeed but will write to the
// final file location (sf.Name() will deceptively still point to the old file
// location).  Therefore it is recommended that the file handle be closed
// immediately after calling Commit. Note that the file must not be closed
// before calling commit as this will cause the sync to fail.
func (sf *safeFile) Commit() error {
	err := sf.Sync()
	if err != nil {
		return err
	}
	return os.Rename(sf.finalName+"_temp", sf.finalName)
}

// NewSafeFile returns a file that can atomically be written to disk,
// minimizing the risk of corruption.
func NewSafeFile(filename string) (*safeFile, error) {
	file, err := os.Create(filename + "_temp")
	if err != nil {
		return nil, err
	}

	// Get the absolute path of the filename so that calling os.Chdir in
	// between calling NewSafeFile and calling safeFile.Commit does not change
	// the final file path.
	absFilename, err := filepath.Abs(filename)
	if err != nil {
		return nil, err
	}

	return &safeFile{file, absFilename}, nil
}
<commit_msg>Close safeFile before Renaming<commit_after>package persist

import (
	"crypto/rand"
	"encoding/base32"
	"errors"
	"os"
	"path/filepath"
)

const (
	// persistDir defines the folder that is used for testing the persist
	// package.
	persistDir = "persist"
)

var (
	// ErrBadVersion indicates that the version number of the file is not
	// compatible with the current codebase.
	ErrBadVersion = errors.New("incompatible version")

	// ErrBadHeader indicates that the file opened is not the file that was
	// expected.
	ErrBadHeader = errors.New("wrong header")
)

// Metadata contains the header and version of the data being stored.
type Metadata struct {
	Header, Version string
}

// RandomSuffix returns a 20 character base32 suffix for a filename. There are
// 100 bits of entropy, and a very low probability of colliding with existing
// files unintentionally.
func RandomSuffix() string {
	randBytes := make([]byte, 20)
	rand.Read(randBytes)
	str := base32.StdEncoding.EncodeToString(randBytes)
	return str[:20]
}

// A safeFile is a file that is stored under a temporary filename. When Commit
// is called, the file is renamed to its "final" filename. This allows for
// atomic updating of files; otherwise, an unexpected shutdown could leave a
// valuable file in a corrupted state. Callers must still Close the file handle
// as usual.
type safeFile struct {
	*os.File
	finalName string
}

// Commit syncs the file and then renames it to the intended final filename.
// Writing to the file after calling Commit will succeed but will write to the
// final file location (sf.Name() will deceptively still point to the old file
// location).  Therefore it is recommended that the file handle be closed
// immediately after calling Commit. Note that the file must not be closed
// before calling commit as this will cause the sync to fail.
func (sf *safeFile) Commit() error {
	err := sf.Sync()
	if err != nil {
		return err
	}
	err = sf.Close()
	if err != nil {
		return err
	}
	return os.Rename(sf.finalName+"_temp", sf.finalName)
}

// NewSafeFile returns a file that can atomically be written to disk,
// minimizing the risk of corruption.
func NewSafeFile(filename string) (*safeFile, error) {
	file, err := os.Create(filename + "_temp")
	if err != nil {
		return nil, err
	}

	// Get the absolute path of the filename so that calling os.Chdir in
	// between calling NewSafeFile and calling safeFile.Commit does not change
	// the final file path.
	absFilename, err := filepath.Abs(filename)
	if err != nil {
		return nil, err
	}

	return &safeFile{file, absFilename}, nil
}
<|endoftext|>


<commit_before>package pagerduty

/*
 * Copyright 2016 Albert P. Tobey <atobey@netflix.com>
 *
 * 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.
 */

import (
	"bytes"
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/netflix/hal-9001/hal"
)

func Register() {
	pg := hal.Plugin{
		Name:  "page",
		Func:  page,
		Regex: "^[[:space:]]*!page",
	}
	pg.Register()

	oc := hal.Plugin{
		Name:  "oncall",
		Func:  oncall,
		Init:  cacheInit,
		Regex: "^[[:space:]]*!oncall",
	}
	oc.Register()
}

// the hal.secrets key that should contain the pagerduty auth token
const PagerdutyTokenKey = `pagerduty.token`

// the hal.secrets key that should contain the pagerduty account domain
const PagerdutyDomainKey = `pagerduty.domain`

// the key name used for caching the full escalation policy
const PolicyCacheKey = `pagerduty.policy_cache`

const PageUsage = `!page <alias> [optional message]

Send an alert via Pagerduty with an optional custom message.

Aliases that have a comma-separated list of service keys will result in one page going to each service key when the alias is paged.

!page core
!page core <message>
!pagecore HELP ME YOU ARE MY ONLY HOPE

!page add <alias> <service key>
!page add <alias> <service key>,<service_key>,<service_key>,...
!page rm <alias>
!page list
`

const OncallUsage = `!oncall <alias>

Find out who is oncall. If only one argument is provided, it must match
a known alias for a Pagerduty service. Otherwise, it is expected to be
a subcommand.

!oncall core
`

const PageDefaultMessage = `HAL: your presence is requested in the chat room.`

const cacheExpire = time.Minute * 10

// TODO: for now there is just one cache and periodic function, but once that is changed
// to allow multiple pagerduty domains/tokens, this will need to be split as well
const PeriodicFuncName = "pagerduty-cache-update-frequency"

const DefaultCacheInterval = "1h"

func page(msg hal.Evt) {
	parts := msg.BodyAsArgv()

	// detect concatenated command + team name & split them
	// e.g. "!pagecore" -> {"!page", "core"}
	if strings.HasPrefix(parts[0], "!page") && len(parts[0]) > 5 {
		team := strings.TrimPrefix(parts[0], "!page")
		parts = append([]string{"!page", team}, parts[1:]...)
	}

	// should be 2 parts now, "!page" and the target team
	if parts[0] != "!page" || len(parts) < 2 {
		msg.Reply(PageUsage)
		return
	}

	switch parts[1] {
	case "h", "help":
		msg.Reply(PageUsage)
	case "add":
		addAlias(msg, parts[2:])
	case "rm":
		rmAlias(msg, parts[2:])
	case "list":
		listAlias(msg)
	default:
		pageAlias(msg, parts[1:])
	}
}

func pageAlias(msg hal.Evt, parts []string) {
	pageMessage := PageDefaultMessage
	if len(parts) > 1 {
		pageMessage = strings.Join(parts, " ")
	}

	// map alias name to PD token via prefs
	key := aliasKey(parts[0])
	pref := msg.AsPref().FindKey(key).One()

	// make sure the query succeeded
	if !pref.Success {
		log.Println("%s", pref.String())
		msg.Replyf("Unable to access preferences: %q", pref.Error)
		return
	}

	// if qpref.Get returned the default, the alias was not found
	if pref.Value == "" {
		msg.Replyf("Alias %q not recognized. Try !page add <alias> <service key>", parts[0])
		return
	}

	// make sure the hal secrets are set up
	err := checkSecrets()
	if err != nil {
		msg.Error(err)
		return
	}

	// the value can be a list of tokens, separated by commas
	response := bytes.NewBuffer([]byte{})
	for _, svckey := range strings.Split(pref.Value, ",") {
		// get the Pagerduty auth token from the secrets API
		secrets := hal.Secrets()
		token := secrets.Get(PagerdutyTokenKey)
		if token == "" {
			msg.Replyf("Your Pagerduty auth token does not seem to be configured. Please add the %q secret.",
				PagerdutyTokenKey)
			return
		}

		// create the event and send it
		pde := NewTrigger(svckey, pageMessage) // in ./pagerduty.go
		resp, err := pde.Send(token)
		if err != nil {
			msg.Replyf("Error while communicating with Pagerduty. %d %s", resp.StatusCode, resp.Message)
			return
		}

		fmt.Fprintf(response, "%s\n", resp.Message)
	}

	// TODO: add some boilerplate around this
	msg.Reply(response.String())
}

func addAlias(msg hal.Evt, parts []string) {
	if len(parts) < 2 {
		msg.Replyf("!page add requires 2 arguments, e.g. !page add sysadmins XXXXXXX")
		return
	} else if len(parts) > 2 {
		keys := strings.Replace(strings.Join(parts[1:], ","), ",,", ",", len(parts)-2)
		parts = []string{parts[0], keys}
	}

	pref := msg.AsPref()
	pref.User = "" // filled in by AsPref and unwanted
	pref.Key = aliasKey(parts[0])
	pref.Value = parts[1]

	err := pref.Set()
	if err != nil {
		msg.Replyf("Write failed: %s", err)
	} else {
		msg.Replyf("Added alias: %q -> %q", parts[0], parts[1])
	}
}

func rmAlias(msg hal.Evt, parts []string) {
	if len(parts) != 1 {
		msg.Replyf("!page rm requires 1 argument, e.g. !page rm sysadmins")
		return
	}

	pref := msg.AsPref()
	pref.User = "" // filled in by AsPref and unwanted
	pref.Key = aliasKey(parts[0])
	pref.Delete()

	msg.Replyf("Removed alias %q", parts[0])
}

func listAlias(msg hal.Evt) {
	pref := msg.AsPref()
	pref.User = "" // filled in by AsPref and unwanted
	prefs := pref.GetPrefs()
	data := prefs.Table()
	msg.ReplyTable(data[0], data[1:])
}

func aliasKey(alias string) string {
	return fmt.Sprintf("alias.%s", alias)
}

func oncall(msg hal.Evt) {
	parts := msg.BodyAsArgv()

	if len(parts) == 1 {
		msg.Reply(OncallUsage)
		return
	} else if len(parts) != 2 {
		msg.Replyf("%s: invalid command.\n%s", parts[0], OncallUsage)
		return
	}

	// make sure the pagerduty token and domain are setup in hal.Secrets
	err := checkSecrets()
	if err != nil {
		msg.Error(err)
		return
	}

	msg.Replyf("Command: %q", parts[1])

	if parts[1] == "cache-now" {
		msg.Reply("Updating Pagerduty policy cache now.")
		cacheNow()
		msg.Reply("Pagerduty policy cache update complete.")
		return
	} else if parts[1] == "cache-status" {
		age := int(hal.Cache().Age(PolicyCacheKey).Seconds())
		next := time.Time{}
		status := "broken"
		pf := hal.GetPeriodicFunc(PeriodicFuncName)
		if pf != nil {
			next = pf.Last().Add(pf.Interval)
			status = pf.Status()
		}
		msg.Replyf("The cache is %d seconds old. Auto-update is %s and its next update is at %s.",
			age, status, next.Format(time.UnixDate))
		return
	}

	// TODO: look at the aliases set up for !page and try for an exact match
	// before doing fuzzy search -- move fuzzy search to a "search" subcommand
	// so it's clear that it is not precise
	want := strings.ToLower(parts[1])
	matches := make([]EscalationPolicy, 0)
	policies := getPolicyCache(false)

	// search over all policies looking for matching policy name, escalation
	// rule name, or service name
	for _, policy := range policies {
		// try matching the policy name
		lname := strings.ToLower(policy.Name)
		if strings.Contains(lname, want) {
			matches = append(matches, policy)
			continue
		}

		// try matching the escalation rule names
		for _, rule := range policy.EscalationRules {
			lname = strings.ToLower(rule.RuleObject.Name)
			if strings.Contains(lname, want) {
				matches = append(matches, policy)
				continue
			}
		}

		// try matching service names
		for _, svc := range policy.Services {
			lname = strings.ToLower(svc.Name)
			if strings.Contains(lname, want) {
				matches = append(matches, policy)
				continue
			}
		}
	}

	reply := formatOncallReply(want, matches)
	msg.Reply(reply)
}

func checkSecrets() error {
	secrets := hal.Secrets()
	token := secrets.Get(PagerdutyTokenKey)
	if token == "" {
		return fmt.Errorf("Your Pagerduty auth token does not seem to be configured. Please add the %q secret.", PagerdutyTokenKey)
	}

	domain := secrets.Get(PagerdutyDomainKey)
	if domain == "" {
		return fmt.Errorf("Your Pagerduty domain does not seem to be configured. Please add the %q secret.", PagerdutyDomainKey)
	}

	return nil
}

func getPolicyCache(forceUpdate bool) []EscalationPolicy {
	// see if there's a copy cached
	policies := []EscalationPolicy{}
	if hal.Cache().Exists(PolicyCacheKey) {
		ttl, _ := hal.Cache().Get(PolicyCacheKey, &policies)
		// TODO: maybe hal.Cache().Get should be careful to not modify the pointer if the ttl is expired...
		if ttl == 0 || forceUpdate {
			policies = []EscalationPolicy{}
		}
	}

	// the cache exists and is still valid, return it now
	if len(policies) > 0 {
		return policies
	}

	// TODO: consider making the token key per-room so different rooms can use different tokens
	// doing this will require a separate cache object per token...
	secrets := hal.Secrets()
	token := secrets.Get(PagerdutyTokenKey)
	domain := secrets.Get(PagerdutyDomainKey)

	// log and noop if the secrets aren't configured (yet)
	// the user-facing commands will report if they are missing
	if token == "" || domain == "" {
		log.Printf("pagerduty: Either the %s or %s is not set up in hal.Secrets. Returning empty list.",
			PagerdutyTokenKey, PagerdutyDomainKey)
		return []EscalationPolicy{}
	}

	// get all of the defined policies
	var err error
	policies, err = GetEscalationPolicies(token, domain)
	if err != nil {
		log.Printf("Returning empty list. REST call to Pagerduty failed: %s", err)
		return []EscalationPolicy{}
	}

	// TODO: make this configurable via prefs
	hal.Cache().Set(PolicyCacheKey, &policies, cacheExpire)

	return policies
}

func cacheInit(i *hal.Instance) {
	freqPref := hal.GetPref("", "", i.RoomId, "pagerduty", "cache-update-frequency", DefaultCacheInterval)

	td, err := time.ParseDuration(freqPref.Value)
	if err != nil {
		log.Panicf("BUG: could not parse freq stored in db: %q", freqPref.Value)
	}

	log.Printf("cacheInit called for pagerduty...")

	pf := hal.GetPeriodicFunc(PeriodicFuncName)
	if pf != nil {
		if pf.Status() != "running" {
			pf.Start()
		}
	} else {
		pf = &hal.PeriodicFunc{
			Name:     PeriodicFuncName,
			Interval: td,
			Function: cacheNow,
		}
		pf.Register()
		pf.Start()
	}

	// TODO: add a command to stop, etc.
}

func cacheNow() {
	getPolicyCache(true)
}

func formatOncallReply(wanted string, policies []EscalationPolicy) string {
	age := int(hal.Cache().Age(PolicyCacheKey).Seconds())

	intro := fmt.Sprintf("%d results for %q (%d seconds ago)\n", len(policies), wanted, age)
	buf := bytes.NewBufferString(intro)

	for _, policy := range policies {
		buf.WriteString(policy.Name)
		buf.WriteString("\n")

		for _, oncall := range policy.OnCall {
			times := formatTimes(oncall.Start, oncall.End)
			indent := strings.Repeat("  ", oncall.Level) // indent deeper per level
			user := fmt.Sprintf("  %s%s: %s %s\n", indent, oncall.User.Name, oncall.User.Email, times)
			buf.WriteString(user)
		}

		buf.WriteString("\n")
	}

	return buf.String()
}

func formatTimes(st, et *time.Time) string {
	var start, end string
	if st != nil {
		start = st.Local().Format("2006-01-02")
	} else {
		return "always on call"
	}

	if et != nil {
		end = et.Local().Format("2006-01-02")
	} else {
		return "always on call"
	}

	return fmt.Sprintf("%s - %s", start, end)
}
<commit_msg>remove accidental debug response<commit_after>package pagerduty

/*
 * Copyright 2016 Albert P. Tobey <atobey@netflix.com>
 *
 * 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.
 */

import (
	"bytes"
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/netflix/hal-9001/hal"
)

func Register() {
	pg := hal.Plugin{
		Name:  "page",
		Func:  page,
		Regex: "^[[:space:]]*!page",
	}
	pg.Register()

	oc := hal.Plugin{
		Name:  "oncall",
		Func:  oncall,
		Init:  cacheInit,
		Regex: "^[[:space:]]*!oncall",
	}
	oc.Register()
}

// the hal.secrets key that should contain the pagerduty auth token
const PagerdutyTokenKey = `pagerduty.token`

// the hal.secrets key that should contain the pagerduty account domain
const PagerdutyDomainKey = `pagerduty.domain`

// the key name used for caching the full escalation policy
const PolicyCacheKey = `pagerduty.policy_cache`

const PageUsage = `!page <alias> [optional message]

Send an alert via Pagerduty with an optional custom message.

Aliases that have a comma-separated list of service keys will result in one page going to each service key when the alias is paged.

!page core
!page core <message>
!pagecore HELP ME YOU ARE MY ONLY HOPE

!page add <alias> <service key>
!page add <alias> <service key>,<service_key>,<service_key>,...
!page rm <alias>
!page list
`

const OncallUsage = `!oncall <alias>

Find out who is oncall. If only one argument is provided, it must match
a known alias for a Pagerduty service. Otherwise, it is expected to be
a subcommand.

!oncall core
`

const PageDefaultMessage = `HAL: your presence is requested in the chat room.`

const cacheExpire = time.Minute * 10

// TODO: for now there is just one cache and periodic function, but once that is changed
// to allow multiple pagerduty domains/tokens, this will need to be split as well
const PeriodicFuncName = "pagerduty-cache-update-frequency"

const DefaultCacheInterval = "1h"

func page(msg hal.Evt) {
	parts := msg.BodyAsArgv()

	// detect concatenated command + team name & split them
	// e.g. "!pagecore" -> {"!page", "core"}
	if strings.HasPrefix(parts[0], "!page") && len(parts[0]) > 5 {
		team := strings.TrimPrefix(parts[0], "!page")
		parts = append([]string{"!page", team}, parts[1:]...)
	}

	// should be 2 parts now, "!page" and the target team
	if parts[0] != "!page" || len(parts) < 2 {
		msg.Reply(PageUsage)
		return
	}

	switch parts[1] {
	case "h", "help":
		msg.Reply(PageUsage)
	case "add":
		addAlias(msg, parts[2:])
	case "rm":
		rmAlias(msg, parts[2:])
	case "list":
		listAlias(msg)
	default:
		pageAlias(msg, parts[1:])
	}
}

func pageAlias(msg hal.Evt, parts []string) {
	pageMessage := PageDefaultMessage
	if len(parts) > 1 {
		pageMessage = strings.Join(parts, " ")
	}

	// map alias name to PD token via prefs
	key := aliasKey(parts[0])
	pref := msg.AsPref().FindKey(key).One()

	// make sure the query succeeded
	if !pref.Success {
		log.Println("%s", pref.String())
		msg.Replyf("Unable to access preferences: %q", pref.Error)
		return
	}

	// if qpref.Get returned the default, the alias was not found
	if pref.Value == "" {
		msg.Replyf("Alias %q not recognized. Try !page add <alias> <service key>", parts[0])
		return
	}

	// make sure the hal secrets are set up
	err := checkSecrets()
	if err != nil {
		msg.Error(err)
		return
	}

	// the value can be a list of tokens, separated by commas
	response := bytes.NewBuffer([]byte{})
	for _, svckey := range strings.Split(pref.Value, ",") {
		// get the Pagerduty auth token from the secrets API
		secrets := hal.Secrets()
		token := secrets.Get(PagerdutyTokenKey)
		if token == "" {
			msg.Replyf("Your Pagerduty auth token does not seem to be configured. Please add the %q secret.",
				PagerdutyTokenKey)
			return
		}

		// create the event and send it
		pde := NewTrigger(svckey, pageMessage) // in ./pagerduty.go
		resp, err := pde.Send(token)
		if err != nil {
			msg.Replyf("Error while communicating with Pagerduty. %d %s", resp.StatusCode, resp.Message)
			return
		}

		fmt.Fprintf(response, "%s\n", resp.Message)
	}

	// TODO: add some boilerplate around this
	msg.Reply(response.String())
}

func addAlias(msg hal.Evt, parts []string) {
	if len(parts) < 2 {
		msg.Replyf("!page add requires 2 arguments, e.g. !page add sysadmins XXXXXXX")
		return
	} else if len(parts) > 2 {
		keys := strings.Replace(strings.Join(parts[1:], ","), ",,", ",", len(parts)-2)
		parts = []string{parts[0], keys}
	}

	pref := msg.AsPref()
	pref.User = "" // filled in by AsPref and unwanted
	pref.Key = aliasKey(parts[0])
	pref.Value = parts[1]

	err := pref.Set()
	if err != nil {
		msg.Replyf("Write failed: %s", err)
	} else {
		msg.Replyf("Added alias: %q -> %q", parts[0], parts[1])
	}
}

func rmAlias(msg hal.Evt, parts []string) {
	if len(parts) != 1 {
		msg.Replyf("!page rm requires 1 argument, e.g. !page rm sysadmins")
		return
	}

	pref := msg.AsPref()
	pref.User = "" // filled in by AsPref and unwanted
	pref.Key = aliasKey(parts[0])
	pref.Delete()

	msg.Replyf("Removed alias %q", parts[0])
}

func listAlias(msg hal.Evt) {
	pref := msg.AsPref()
	pref.User = "" // filled in by AsPref and unwanted
	prefs := pref.GetPrefs()
	data := prefs.Table()
	msg.ReplyTable(data[0], data[1:])
}

func aliasKey(alias string) string {
	return fmt.Sprintf("alias.%s", alias)
}

func oncall(msg hal.Evt) {
	parts := msg.BodyAsArgv()

	if len(parts) == 1 {
		msg.Reply(OncallUsage)
		return
	} else if len(parts) != 2 {
		msg.Replyf("%s: invalid command.\n%s", parts[0], OncallUsage)
		return
	}

	// make sure the pagerduty token and domain are setup in hal.Secrets
	err := checkSecrets()
	if err != nil {
		msg.Error(err)
		return
	}

	if parts[1] == "cache-now" {
		msg.Reply("Updating Pagerduty policy cache now.")
		cacheNow()
		msg.Reply("Pagerduty policy cache update complete.")
		return
	} else if parts[1] == "cache-status" {
		age := int(hal.Cache().Age(PolicyCacheKey).Seconds())
		next := time.Time{}
		status := "broken"
		pf := hal.GetPeriodicFunc(PeriodicFuncName)
		if pf != nil {
			next = pf.Last().Add(pf.Interval)
			status = pf.Status()
		}
		msg.Replyf("The cache is %d seconds old. Auto-update is %s and its next update is at %s.",
			age, status, next.Format(time.UnixDate))
		return
	}

	// TODO: look at the aliases set up for !page and try for an exact match
	// before doing fuzzy search -- move fuzzy search to a "search" subcommand
	// so it's clear that it is not precise
	want := strings.ToLower(parts[1])
	matches := make([]EscalationPolicy, 0)
	policies := getPolicyCache(false)

	// search over all policies looking for matching policy name, escalation
	// rule name, or service name
	for _, policy := range policies {
		// try matching the policy name
		lname := strings.ToLower(policy.Name)
		if strings.Contains(lname, want) {
			matches = append(matches, policy)
			continue
		}

		// try matching the escalation rule names
		for _, rule := range policy.EscalationRules {
			lname = strings.ToLower(rule.RuleObject.Name)
			if strings.Contains(lname, want) {
				matches = append(matches, policy)
				continue
			}
		}

		// try matching service names
		for _, svc := range policy.Services {
			lname = strings.ToLower(svc.Name)
			if strings.Contains(lname, want) {
				matches = append(matches, policy)
				continue
			}
		}
	}

	reply := formatOncallReply(want, matches)
	msg.Reply(reply)
}

func checkSecrets() error {
	secrets := hal.Secrets()
	token := secrets.Get(PagerdutyTokenKey)
	if token == "" {
		return fmt.Errorf("Your Pagerduty auth token does not seem to be configured. Please add the %q secret.", PagerdutyTokenKey)
	}

	domain := secrets.Get(PagerdutyDomainKey)
	if domain == "" {
		return fmt.Errorf("Your Pagerduty domain does not seem to be configured. Please add the %q secret.", PagerdutyDomainKey)
	}

	return nil
}

func getPolicyCache(forceUpdate bool) []EscalationPolicy {
	// see if there's a copy cached
	policies := []EscalationPolicy{}
	if hal.Cache().Exists(PolicyCacheKey) {
		ttl, _ := hal.Cache().Get(PolicyCacheKey, &policies)
		// TODO: maybe hal.Cache().Get should be careful to not modify the pointer if the ttl is expired...
		if ttl == 0 || forceUpdate {
			policies = []EscalationPolicy{}
		}
	}

	// the cache exists and is still valid, return it now
	if len(policies) > 0 {
		return policies
	}

	// TODO: consider making the token key per-room so different rooms can use different tokens
	// doing this will require a separate cache object per token...
	secrets := hal.Secrets()
	token := secrets.Get(PagerdutyTokenKey)
	domain := secrets.Get(PagerdutyDomainKey)

	// log and noop if the secrets aren't configured (yet)
	// the user-facing commands will report if they are missing
	if token == "" || domain == "" {
		log.Printf("pagerduty: Either the %s or %s is not set up in hal.Secrets. Returning empty list.",
			PagerdutyTokenKey, PagerdutyDomainKey)
		return []EscalationPolicy{}
	}

	// get all of the defined policies
	var err error
	policies, err = GetEscalationPolicies(token, domain)
	if err != nil {
		log.Printf("Returning empty list. REST call to Pagerduty failed: %s", err)
		return []EscalationPolicy{}
	}

	// TODO: make this configurable via prefs
	hal.Cache().Set(PolicyCacheKey, &policies, cacheExpire)

	return policies
}

func cacheInit(i *hal.Instance) {
	freqPref := hal.GetPref("", "", i.RoomId, "pagerduty", "cache-update-frequency", DefaultCacheInterval)

	td, err := time.ParseDuration(freqPref.Value)
	if err != nil {
		log.Panicf("BUG: could not parse freq stored in db: %q", freqPref.Value)
	}

	log.Printf("cacheInit called for pagerduty...")

	pf := hal.GetPeriodicFunc(PeriodicFuncName)
	if pf != nil {
		if pf.Status() != "running" {
			pf.Start()
		}
	} else {
		pf = &hal.PeriodicFunc{
			Name:     PeriodicFuncName,
			Interval: td,
			Function: cacheNow,
		}
		pf.Register()
		pf.Start()
	}

	// TODO: add a command to stop, etc.
}

func cacheNow() {
	getPolicyCache(true)
}

func formatOncallReply(wanted string, policies []EscalationPolicy) string {
	age := int(hal.Cache().Age(PolicyCacheKey).Seconds())

	intro := fmt.Sprintf("%d results for %q (%d seconds ago)\n", len(policies), wanted, age)
	buf := bytes.NewBufferString(intro)

	for _, policy := range policies {
		buf.WriteString(policy.Name)
		buf.WriteString("\n")

		for _, oncall := range policy.OnCall {
			times := formatTimes(oncall.Start, oncall.End)
			indent := strings.Repeat("  ", oncall.Level) // indent deeper per level
			user := fmt.Sprintf("  %s%s: %s %s\n", indent, oncall.User.Name, oncall.User.Email, times)
			buf.WriteString(user)
		}

		buf.WriteString("\n")
	}

	return buf.String()
}

func formatTimes(st, et *time.Time) string {
	var start, end string
	if st != nil {
		start = st.Local().Format("2006-01-02")
	} else {
		return "always on call"
	}

	if et != nil {
		end = et.Local().Format("2006-01-02")
	} else {
		return "always on call"
	}

	return fmt.Sprintf("%s - %s", start, end)
}
<|endoftext|>


<commit_before>//
// Copyright 2016 Rackspace
//
// 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 poller

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"net"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/racker/rackspace-monitoring-poller/check"
	"github.com/racker/rackspace-monitoring-poller/config"
	"math"
)

// EleConnectionStream implements ConnectionStream
// See ConnectionStream for more information
type EleConnectionStream struct {
	LogPrefixGetter

	ctx     context.Context
	rootCAs *x509.CertPool

	stopCh chan struct{}
	config *config.Config

	connectionFactory ConnectionFactory
	connsMu           sync.Mutex
	conns             ConnectionsByHost
	wg                sync.WaitGroup

	// map is the private zone ID as a string
	schedulers map[string]Scheduler
}

// NewConnectionStream instantiates a new EleConnectionStream
// It sets up the contexts and the starts the schedulers based on configured private zones
func NewConnectionStream(config *config.Config, rootCAs *x509.CertPool) ConnectionStream {
	return NewCustomConnectionStream(config, rootCAs, nil)
}

// NewCustomConnectionStream is a variant of NewConnectionStream that allows providing a customized ConnectionFactory
func NewCustomConnectionStream(config *config.Config, rootCAs *x509.CertPool, connectionFactory ConnectionFactory) ConnectionStream {
	if connectionFactory == nil {
		connectionFactory = NewConnection
	}
	stream := &EleConnectionStream{
		config:            config,
		rootCAs:           rootCAs,
		schedulers:        make(map[string]Scheduler),
		connectionFactory: connectionFactory,
	}
	stream.ctx = context.Background()
	stream.conns = make(ConnectionsByHost)
	stream.stopCh = make(chan struct{}, 1)
	for _, pz := range config.ZoneIds {
		stream.schedulers[pz] = NewScheduler(pz, stream)
	}
	return stream
}

// GetLogPrefix returns the log prefix for this module
func (cs *EleConnectionStream) GetLogPrefix() string {
	return "stream"
}

// getRegisteredConnectionNames returns the registered connection names
func (cs *EleConnectionStream) getRegisteredConnectionNames() []string {
	names := []string{}
	for _, conn := range cs.conns {
		names = append(names, conn.GetLogPrefix())
	}
	return names
}

// RegisterConnection sets up a new connection and adds it to
// connection stream
// If no connection list has been initialized, this method will
// return an InvalidConnectionStreamError.  If that's the case,
// please instantiate a new connection stream via NewConnectionStream function
func (cs *EleConnectionStream) RegisterConnection(qry string, conn Connection) error {
	cs.connsMu.Lock()
	defer cs.connsMu.Unlock()
	if cs.conns == nil {
		return ErrInvalidConnectionStream
	}
	cs.conns[qry] = conn
	log.WithFields(log.Fields{
		"prefix":      cs.GetLogPrefix(),
		"connections": cs.getRegisteredConnectionNames(),
	}).Debug("Currently registered connections")
	return nil
}

// ReconcileChecks routes the ChecksPreparation to its schedulers.
func (cs *EleConnectionStream) ReconcileChecks(cp ChecksPrepared) {
	for _, sched := range cs.schedulers {
		sched.ReconcileChecks(cp)
	}
}

func (cs *EleConnectionStream) ValidateChecks(cp ChecksPreparing) error {
	for _, sched := range cs.schedulers {
		err := sched.ValidateChecks(cp)
		if err != nil {
			log.WithFields(log.Fields{
				"prefix":    cs.GetLogPrefix(),
				"scheduler": sched,
				"cp":        cp,
				"err":       err,
			}).Warn("Scheduler was not able to validate check preparation")
			return err
		}
	}

	return nil
}

// Stop explicitly stops all connections in the stream and notifies the channel
func (cs *EleConnectionStream) Stop() {
	if cs.conns == nil {
		return
	}
	for _, conn := range cs.conns {
		conn.Close()
	}
	cs.stopCh <- struct{}{}
}

// StopNotify returns a stop channel
func (cs *EleConnectionStream) StopNotify() chan struct{} {
	return cs.stopCh
}

// SendMetrics sends a CheckResultSet via the first connection it can
// retrieve in the connection list
func (cs *EleConnectionStream) SendMetrics(crs *check.ResultSet) error {
	if cs.conns == nil || len(cs.conns) == 0 {
		return ErrNoConnections
	}

	if conn := cs.conns.ChooseBest(); conn != nil {
		conn.GetSession().Send(check.NewMetricsPostRequest(crs, conn.GetClockOffset()))
	}

	return nil

}

// Connect connects to configured endpoints.
// There are 2 ways to connect:
// 1. You can utilize SRV records defined in the configuration
// to dynamically find endpoints
// 2. You can explicitly specify endpoint addresses and connect
// to them directly
// DEFAULT: Using SRV records
func (cs *EleConnectionStream) Connect() {
	if cs.config.UseSrv {
		for _, qry := range cs.config.SrvQueries {
			cs.wg.Add(1)
			go cs.connectBySrv(qry)
		}
	} else {
		for _, addr := range cs.config.Addresses {
			cs.wg.Add(1)
			go cs.connectByHost(addr)
		}
	}
}

// WaitCh provides a channel for waiting on connection establishment
func (cs *EleConnectionStream) WaitCh() <-chan struct{} {
	c := make(chan struct{}, 1)
	go func() {
		cs.wg.Wait()
		c <- struct{}{}
	}()
	return c
}

func (cs *EleConnectionStream) connectBySrv(qry string) {
	_, addrs, err := net.LookupSRV("", "", qry)
	if err != nil {
		log.WithFields(log.Fields{
			"prefix": cs.GetLogPrefix(),
		}).Errorf("SRV Lookup Failure : %v", err)
		return
	}
	if len(addrs) == 0 {
		log.WithFields(log.Fields{
			"prefix": cs.GetLogPrefix(),
		}).Error("no addresses returned")
		return
	}
	addr := net.JoinHostPort(addrs[0].Target, fmt.Sprintf("%v", addrs[0].Port))
	log.WithFields(log.Fields{
		"prefix": cs.GetLogPrefix(),
		"query":  qry,
		"addr":   addr,
	}).Debug("Connecting")
	cs.connectByHost(addr)
}

func (cs *EleConnectionStream) connectByHost(addr string) {
	defer cs.wg.Done()
	for {
		conn := cs.connectionFactory(addr, cs.config.Guid, cs)
		err := conn.Connect(cs.ctx, cs.config, cs.buildTLSConfig(addr))
		if err != nil {
			goto conn_error
		}
		err = cs.RegisterConnection(addr, conn)
		if err != nil {
			goto conn_error
		}
		conn.Wait()
		goto new_connection
	conn_error:
		log.WithFields(log.Fields{
			"prefix":  cs.GetLogPrefix(),
			"address": addr,
		}).Errorf("Error: %v", err)
	new_connection:
		log.WithFields(log.Fields{
			"prefix":  cs.GetLogPrefix(),
			"address": addr,
			"timeout": ReconnectTimeout,
		}).Debug("Connection sleeping")
		for {
			select {
			case <-cs.ctx.Done():
				log.WithFields(log.Fields{
					"prefix":  cs.GetLogPrefix(),
					"address": addr,
				}).Debug("Connection cancelled")
				return
			case <-time.After(ReconnectTimeout):
				log.WithField("prefix", cs.GetLogPrefix()).Debug("Reconnecting")
				log.WithFields(log.Fields{
					"prefix":  cs.GetLogPrefix(),
					"address": addr,
				}).Debug("Connection cancelled")
				continue
			}
		}
	}
}

func (cs *EleConnectionStream) buildTLSConfig(addr string) *tls.Config {
	host, _, _ := net.SplitHostPort(addr)
	conf := &tls.Config{
		InsecureSkipVerify: cs.rootCAs == nil,
		ServerName:         host,
		RootCAs:            cs.rootCAs,
	}
	return conf
}

// ChooseBest selects the best of its connections for posting metrics, etc.
// Returns nil if no connections were present.
func (conns ConnectionsByHost) ChooseBest() Connection {
	var minLatency int64 = math.MaxInt64
	var best Connection

	for _, conn := range conns {
		latency := conn.GetLatency()
		if latency < minLatency {
			minLatency = latency
			best = conn
		}
	}

	return best
}
<commit_msg>jump to correct reconnect label (#106)<commit_after>//
// Copyright 2016 Rackspace
//
// 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 poller

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"net"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/racker/rackspace-monitoring-poller/check"
	"github.com/racker/rackspace-monitoring-poller/config"
	"math"
)

// EleConnectionStream implements ConnectionStream
// See ConnectionStream for more information
type EleConnectionStream struct {
	LogPrefixGetter

	ctx     context.Context
	rootCAs *x509.CertPool

	stopCh chan struct{}
	config *config.Config

	connectionFactory ConnectionFactory
	connsMu           sync.Mutex
	conns             ConnectionsByHost
	wg                sync.WaitGroup

	// map is the private zone ID as a string
	schedulers map[string]Scheduler
}

// NewConnectionStream instantiates a new EleConnectionStream
// It sets up the contexts and the starts the schedulers based on configured private zones
func NewConnectionStream(config *config.Config, rootCAs *x509.CertPool) ConnectionStream {
	return NewCustomConnectionStream(config, rootCAs, nil)
}

// NewCustomConnectionStream is a variant of NewConnectionStream that allows providing a customized ConnectionFactory
func NewCustomConnectionStream(config *config.Config, rootCAs *x509.CertPool, connectionFactory ConnectionFactory) ConnectionStream {
	if connectionFactory == nil {
		connectionFactory = NewConnection
	}
	stream := &EleConnectionStream{
		config:            config,
		rootCAs:           rootCAs,
		schedulers:        make(map[string]Scheduler),
		connectionFactory: connectionFactory,
	}
	stream.ctx = context.Background()
	stream.conns = make(ConnectionsByHost)
	stream.stopCh = make(chan struct{}, 1)
	for _, pz := range config.ZoneIds {
		stream.schedulers[pz] = NewScheduler(pz, stream)
	}
	return stream
}

// GetLogPrefix returns the log prefix for this module
func (cs *EleConnectionStream) GetLogPrefix() string {
	return "stream"
}

// getRegisteredConnectionNames returns the registered connection names
func (cs *EleConnectionStream) getRegisteredConnectionNames() []string {
	names := []string{}
	for _, conn := range cs.conns {
		names = append(names, conn.GetLogPrefix())
	}
	return names
}

// RegisterConnection sets up a new connection and adds it to
// connection stream
// If no connection list has been initialized, this method will
// return an InvalidConnectionStreamError.  If that's the case,
// please instantiate a new connection stream via NewConnectionStream function
func (cs *EleConnectionStream) RegisterConnection(qry string, conn Connection) error {
	cs.connsMu.Lock()
	defer cs.connsMu.Unlock()
	if cs.conns == nil {
		return ErrInvalidConnectionStream
	}
	cs.conns[qry] = conn
	log.WithFields(log.Fields{
		"prefix":      cs.GetLogPrefix(),
		"connections": cs.getRegisteredConnectionNames(),
	}).Debug("Currently registered connections")
	return nil
}

// ReconcileChecks routes the ChecksPreparation to its schedulers.
func (cs *EleConnectionStream) ReconcileChecks(cp ChecksPrepared) {
	for _, sched := range cs.schedulers {
		sched.ReconcileChecks(cp)
	}
}

func (cs *EleConnectionStream) ValidateChecks(cp ChecksPreparing) error {
	for _, sched := range cs.schedulers {
		err := sched.ValidateChecks(cp)
		if err != nil {
			log.WithFields(log.Fields{
				"prefix":    cs.GetLogPrefix(),
				"scheduler": sched,
				"cp":        cp,
				"err":       err,
			}).Warn("Scheduler was not able to validate check preparation")
			return err
		}
	}

	return nil
}

// Stop explicitly stops all connections in the stream and notifies the channel
func (cs *EleConnectionStream) Stop() {
	if cs.conns == nil {
		return
	}
	for _, conn := range cs.conns {
		conn.Close()
	}
	cs.stopCh <- struct{}{}
}

// StopNotify returns a stop channel
func (cs *EleConnectionStream) StopNotify() chan struct{} {
	return cs.stopCh
}

// SendMetrics sends a CheckResultSet via the first connection it can
// retrieve in the connection list
func (cs *EleConnectionStream) SendMetrics(crs *check.ResultSet) error {
	if cs.conns == nil || len(cs.conns) == 0 {
		return ErrNoConnections
	}

	if conn := cs.conns.ChooseBest(); conn != nil {
		conn.GetSession().Send(check.NewMetricsPostRequest(crs, conn.GetClockOffset()))
	}

	return nil

}

// Connect connects to configured endpoints.
// There are 2 ways to connect:
// 1. You can utilize SRV records defined in the configuration
// to dynamically find endpoints
// 2. You can explicitly specify endpoint addresses and connect
// to them directly
// DEFAULT: Using SRV records
func (cs *EleConnectionStream) Connect() {
	if cs.config.UseSrv {
		for _, qry := range cs.config.SrvQueries {
			cs.wg.Add(1)
			go cs.connectBySrv(qry)
		}
	} else {
		for _, addr := range cs.config.Addresses {
			cs.wg.Add(1)
			go cs.connectByHost(addr)
		}
	}
}

// WaitCh provides a channel for waiting on connection establishment
func (cs *EleConnectionStream) WaitCh() <-chan struct{} {
	c := make(chan struct{}, 1)
	go func() {
		cs.wg.Wait()
		c <- struct{}{}
	}()
	return c
}

func (cs *EleConnectionStream) connectBySrv(qry string) {
	_, addrs, err := net.LookupSRV("", "", qry)
	if err != nil {
		log.WithFields(log.Fields{
			"prefix": cs.GetLogPrefix(),
		}).Errorf("SRV Lookup Failure : %v", err)
		return
	}
	if len(addrs) == 0 {
		log.WithFields(log.Fields{
			"prefix": cs.GetLogPrefix(),
		}).Error("no addresses returned")
		return
	}
	addr := net.JoinHostPort(addrs[0].Target, fmt.Sprintf("%v", addrs[0].Port))
	log.WithFields(log.Fields{
		"prefix": cs.GetLogPrefix(),
		"query":  qry,
		"addr":   addr,
	}).Debug("Connecting")
	cs.connectByHost(addr)
}

func (cs *EleConnectionStream) connectByHost(addr string) {
	defer cs.wg.Done()
reconnect:
	for {
		conn := cs.connectionFactory(addr, cs.config.Guid, cs)
		err := conn.Connect(cs.ctx, cs.config, cs.buildTLSConfig(addr))
		if err != nil {
			goto conn_error
		}
		err = cs.RegisterConnection(addr, conn)
		if err != nil {
			goto conn_error
		}
		conn.Wait()
		goto new_connection
	conn_error:
		log.WithFields(log.Fields{
			"prefix":  cs.GetLogPrefix(),
			"address": addr,
		}).Errorf("Error: %v", err)
	new_connection:
		log.WithFields(log.Fields{
			"prefix":  cs.GetLogPrefix(),
			"address": addr,
			"timeout": ReconnectTimeout,
		}).Debug("Connection sleeping")
		for {
			select {
			case <-cs.ctx.Done():
				log.WithFields(log.Fields{
					"prefix":  cs.GetLogPrefix(),
					"address": addr,
				}).Debug("Connection cancelled")
				return
			case <-time.After(ReconnectTimeout):
				log.WithField("prefix", cs.GetLogPrefix()).Debug("Reconnecting")
				log.WithFields(log.Fields{
					"prefix":  cs.GetLogPrefix(),
					"address": addr,
				}).Debug("Connection cancelled")
				continue reconnect
			}
		}
	}
}

func (cs *EleConnectionStream) buildTLSConfig(addr string) *tls.Config {
	host, _, _ := net.SplitHostPort(addr)
	conf := &tls.Config{
		InsecureSkipVerify: cs.rootCAs == nil,
		ServerName:         host,
		RootCAs:            cs.rootCAs,
	}
	return conf
}

// ChooseBest selects the best of its connections for posting metrics, etc.
// Returns nil if no connections were present.
func (conns ConnectionsByHost) ChooseBest() Connection {
	var minLatency int64 = math.MaxInt64
	var best Connection

	for _, conn := range conns {
		latency := conn.GetLatency()
		if latency < minLatency {
			minLatency = latency
			best = conn
		}
	}

	return best
}
<|endoftext|>


<commit_before>/*
Copyright The Helm 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 action

import (
	"github.com/pkg/errors"

	"k8s.io/helm/pkg/release"
	reltesting "k8s.io/helm/pkg/releasetesting"
)

// Test is the action for testing a given release.
//
// It provides the implementation of 'helm test'.
type Test struct {
	cfg *Configuration

	Timeout int64
	Cleanup bool
}

// NewTest creates a new Test object with the given configuration.
func NewTest(cfg *Configuration) *Test {
	return &Test{
		cfg: cfg,
	}
}

// Run executes 'helm test' against the given release.
func (t *Test) Run(name string) (<-chan *release.TestReleaseResponse, <-chan error) {
	errc := make(chan error, 1)
	if err := validateReleaseName(name); err != nil {
		errc <- errors.Errorf("releaseTest: Release name is invalid: %s", name)
		return nil, errc
	}

	// finds the non-deleted release with the given name
	rel, err := t.cfg.Releases.Last(name)
	if err != nil {
		errc <- err
		return nil, errc
	}

	ch := make(chan *release.TestReleaseResponse, 1)
	testEnv := &reltesting.Environment{
		Namespace:  rel.Namespace,
		KubeClient: t.cfg.KubeClient,
		Timeout:    t.Timeout,
		Messages:   ch,
	}
	t.cfg.Log("running tests for release %s", rel.Name)
	tSuite := reltesting.NewTestSuite(rel)

	go func() {
		defer close(errc)
		defer close(ch)

		if err := tSuite.Run(testEnv); err != nil {
			errc <- errors.Wrapf(err, "error running test suite for %s", rel.Name)
			return
		}

		rel.Info.LastTestSuiteRun = &release.TestSuite{
			StartedAt:   tSuite.StartedAt,
			CompletedAt: tSuite.CompletedAt,
			Results:     tSuite.Results,
		}

		if t.Cleanup {
			testEnv.DeleteTestPods(tSuite.TestManifests)
		}

		if err := t.cfg.Releases.Update(rel); err != nil {
			t.cfg.Log("test: Failed to store updated release: %s", err)
		}
	}()
	return ch, errc
}
<commit_msg>fix(action): remove test.go<commit_after><|endoftext|>


<commit_before>package api

import (
	"encoding/base64"
	log "github.com/sirupsen/logrus"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/mssola/user_agent"
	"github.com/usefathom/fathom/pkg/aggregator"
	"github.com/usefathom/fathom/pkg/datastore"
	"github.com/usefathom/fathom/pkg/models"
)

func shouldCollect(r *http.Request) bool {
	// abort if DNT header is set to "1" (these should have been filtered client-side already)
	if r.Header.Get("DNT") == "1" {
		return false
	}

	// don't track prerendered pages, see https://github.com/usefathom/fathom/issues/13
	if r.Header.Get("X-Moz") == "prefetch" || r.Header.Get("X-Purpose") == "preview" {
		return false
	}

	// abort if this is a bot.
	ua := user_agent.New(r.UserAgent())
	if ua.Bot() {
		return false
	}

	return true
}

func parsePathname(p string) string {
	return "/" + strings.TrimLeft(p, "/")
}

// TODO: Move this to aggregator, as we need this endpoint to be as fast as possible
func parseReferrer(r string) string {
	u, err := url.Parse(r)
	if err != nil {
		return ""
	}

	// remove AMP & UTM vars
	q := u.Query()
	keys := []string{"amp", "utm_campaign", "utm_medium", "utm_source"}
	for _, k := range keys {
		q.Del(k)
	}
	u.RawQuery = q.Encode()

	// remove /amp/
	if strings.HasSuffix(u.Path, "/amp/") {
		u.Path = u.Path[0:(len(u.Path) - 5)]
	}

	return u.String()
}

func parseHostname(r string) string {
	u, err := url.Parse(r)
	if err != nil {
		return ""
	}

	return u.Scheme + "://" + u.Host
}

func (api *API) NewCollectHandler() http.Handler {
	pageviews := make(chan *models.Pageview, 10)
	go aggregate(api.database)
	go collect(api.database, pageviews)

	return HandlerFunc(func(w http.ResponseWriter, r *http.Request) error {
		if !shouldCollect(r) {
			return nil
		}

		q := r.URL.Query()
		now := time.Now()

		// get pageview details
		pageview := &models.Pageview{
			ID:           q.Get("id"),
			Hostname:     parseHostname(q.Get("h")),
			Pathname:     parsePathname(q.Get("p")),
			IsNewVisitor: q.Get("nv") == "1",
			IsNewSession: q.Get("ns") == "1",
			IsUnique:     q.Get("u") == "1",
			IsBounce:     q.Get("b") != "0",
			Referrer:     parseReferrer(q.Get("r")),
			Duration:     0,
			Timestamp:    now,
		}

		// find previous pageview by same visitor
		previousPageviewID := q.Get("pid")
		if !pageview.IsNewSession && previousPageviewID != "" {
			previousPageview, err := api.database.GetPageview(previousPageviewID)
			if err != nil && err != datastore.ErrNoResults {
				return err
			}

			// if we have a recent pageview that is less than 30 minutes old
			if previousPageview != nil && previousPageview.Timestamp.After(now.Add(-30*time.Minute)) {
				previousPageview.Duration = (now.Unix() - previousPageview.Timestamp.Unix())
				previousPageview.IsBounce = false

				// push onto channel to be updated (in batch) later
				pageviews <- previousPageview
			}
		}

		// push pageview onto channel to be inserted (in batch) later
		pageviews <- pageview

		// indicate that we're not tracking user data, see https://github.com/usefathom/fathom/issues/65
		w.Header().Set("Tk", "N")

		// headers to prevent caching
		w.Header().Set("Content-Type", "image/gif")
		w.Header().Set("Expires", "Mon, 01 Jan 1990 00:00:00 GMT")
		w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
		w.Header().Set("Pragma", "no-cache")

		// response
		w.WriteHeader(http.StatusOK)

		// 1x1 px transparent GIF
		b, _ := base64.StdEncoding.DecodeString("R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7")
		w.Write(b)
		return nil
	})
}

// runs the aggregate func every minute
func aggregate(db datastore.Datastore) {
	agg := aggregator.New(db)
	agg.Run()

	timeout := 1 * time.Minute

	for {
		select {
		case <-time.After(timeout):
			agg.Run()
		}
	}
}

func collect(db datastore.Datastore, pageviews chan *models.Pageview) {
	var buffer []*models.Pageview
	var size = 250
	var timeout = 500 * time.Millisecond

	for {
		select {
		case pageview := <-pageviews:
			buffer = append(buffer, pageview)
			if len(buffer) >= size {
				persist(db, buffer)
				buffer = buffer[:0]
			}
		case <-time.After(timeout):
			if len(buffer) > 0 {
				persist(db, buffer)
				buffer = buffer[:0]
			}
		}
	}
}

func persist(db datastore.Datastore, pageviews []*models.Pageview) {
	n := len(pageviews)
	updates := make([]*models.Pageview, 0, n)
	inserts := make([]*models.Pageview, 0, n)

	for _, p := range pageviews {
		if !p.IsBounce {
			updates = append(updates, p)
		} else {
			inserts = append(inserts, p)
		}
	}

	log.Debugf("persisting %d pageviews (%d inserts, %d updates)", len(pageviews), len(inserts), len(updates))

	var err error
	err = db.InsertPageviews(inserts)
	if err != nil {
		log.Error(err)
	}

	err = db.UpdatePageviews(updates)
	if err != nil {
		log.Error(err)
	}
}
<commit_msg>increase buffer size & buffer timeout for collecting pageviews<commit_after>package api

import (
	"encoding/base64"
	log "github.com/sirupsen/logrus"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/mssola/user_agent"
	"github.com/usefathom/fathom/pkg/aggregator"
	"github.com/usefathom/fathom/pkg/datastore"
	"github.com/usefathom/fathom/pkg/models"
)

func shouldCollect(r *http.Request) bool {
	// abort if DNT header is set to "1" (these should have been filtered client-side already)
	if r.Header.Get("DNT") == "1" {
		return false
	}

	// don't track prerendered pages, see https://github.com/usefathom/fathom/issues/13
	if r.Header.Get("X-Moz") == "prefetch" || r.Header.Get("X-Purpose") == "preview" {
		return false
	}

	// abort if this is a bot.
	ua := user_agent.New(r.UserAgent())
	if ua.Bot() {
		return false
	}

	return true
}

func parsePathname(p string) string {
	return "/" + strings.TrimLeft(p, "/")
}

// TODO: Move this to aggregator, as we need this endpoint to be as fast as possible
func parseReferrer(r string) string {
	u, err := url.Parse(r)
	if err != nil {
		return ""
	}

	// remove AMP & UTM vars
	q := u.Query()
	keys := []string{"amp", "utm_campaign", "utm_medium", "utm_source"}
	for _, k := range keys {
		q.Del(k)
	}
	u.RawQuery = q.Encode()

	// remove /amp/
	if strings.HasSuffix(u.Path, "/amp/") {
		u.Path = u.Path[0:(len(u.Path) - 5)]
	}

	return u.String()
}

func parseHostname(r string) string {
	u, err := url.Parse(r)
	if err != nil {
		return ""
	}

	return u.Scheme + "://" + u.Host
}

func (api *API) NewCollectHandler() http.Handler {
	pageviews := make(chan *models.Pageview, 10)
	go aggregate(api.database)
	go collect(api.database, pageviews)

	return HandlerFunc(func(w http.ResponseWriter, r *http.Request) error {
		if !shouldCollect(r) {
			return nil
		}

		q := r.URL.Query()
		now := time.Now()

		// get pageview details
		pageview := &models.Pageview{
			ID:           q.Get("id"),
			Hostname:     parseHostname(q.Get("h")),
			Pathname:     parsePathname(q.Get("p")),
			IsNewVisitor: q.Get("nv") == "1",
			IsNewSession: q.Get("ns") == "1",
			IsUnique:     q.Get("u") == "1",
			IsBounce:     q.Get("b") != "0",
			Referrer:     parseReferrer(q.Get("r")),
			Duration:     0,
			Timestamp:    now,
		}

		// find previous pageview by same visitor
		previousPageviewID := q.Get("pid")
		if !pageview.IsNewSession && previousPageviewID != "" {
			previousPageview, err := api.database.GetPageview(previousPageviewID)
			if err != nil && err != datastore.ErrNoResults {
				return err
			}

			// if we have a recent pageview that is less than 30 minutes old
			if previousPageview != nil && previousPageview.Timestamp.After(now.Add(-30*time.Minute)) {
				previousPageview.Duration = (now.Unix() - previousPageview.Timestamp.Unix())
				previousPageview.IsBounce = false

				// push onto channel to be updated (in batch) later
				pageviews <- previousPageview
			}
		}

		// push pageview onto channel to be inserted (in batch) later
		pageviews <- pageview

		// indicate that we're not tracking user data, see https://github.com/usefathom/fathom/issues/65
		w.Header().Set("Tk", "N")

		// headers to prevent caching
		w.Header().Set("Content-Type", "image/gif")
		w.Header().Set("Expires", "Mon, 01 Jan 1990 00:00:00 GMT")
		w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
		w.Header().Set("Pragma", "no-cache")

		// response
		w.WriteHeader(http.StatusOK)

		// 1x1 px transparent GIF
		b, _ := base64.StdEncoding.DecodeString("R0lGODlhAQABAIAAAAAAAP///yH5BAEAAAAALAAAAAABAAEAAAIBRAA7")
		w.Write(b)
		return nil
	})
}

// runs the aggregate func every minute
func aggregate(db datastore.Datastore) {
	agg := aggregator.New(db)
	agg.Run()

	timeout := 1 * time.Minute

	for {
		select {
		case <-time.After(timeout):
			agg.Run()
		}
	}
}

func collect(db datastore.Datastore, pageviews chan *models.Pageview) {
	var buffer []*models.Pageview
	var size = 800
	var timeout = 600 * time.Millisecond

	for {
		select {
		case pageview := <-pageviews:
			buffer = append(buffer, pageview)
			if len(buffer) >= size {
				persist(db, buffer)
				buffer = buffer[:0]
			}
		case <-time.After(timeout):
			if len(buffer) > 0 {
				persist(db, buffer)
				buffer = buffer[:0]
			}
		}
	}
}

func persist(db datastore.Datastore, pageviews []*models.Pageview) {
	n := len(pageviews)
	updates := make([]*models.Pageview, 0, n)
	inserts := make([]*models.Pageview, 0, n)

	for _, p := range pageviews {
		if !p.IsBounce {
			updates = append(updates, p)
		} else {
			inserts = append(inserts, p)
		}
	}

	log.Debugf("persisting %d pageviews (%d inserts, %d updates)", len(pageviews), len(inserts), len(updates))

	var err error
	err = db.InsertPageviews(inserts)
	if err != nil {
		log.Error(err)
	}

	err = db.UpdatePageviews(updates)
	if err != nil {
		log.Error(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2017-2020 The Usacloud 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 cli

import (
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	"github.com/sacloud/libsacloud/v2"

	"github.com/fatih/color"
	"github.com/sacloud/libsacloud/v2/helper/api"
	"github.com/sacloud/libsacloud/v2/sacloud"
	"github.com/sacloud/libsacloud/v2/sacloud/types"
	"github.com/sacloud/usacloud/pkg/config"
	"github.com/sacloud/usacloud/pkg/output"
	"github.com/sacloud/usacloud/pkg/version"
	"github.com/spf13/pflag"
)

type Context interface {
	Option() *config.Config
	Output() output.Output
	Client() sacloud.APICaller
	Zone() string
	IO() IO
	context.Context

	Args() []string

	ResourceName() string
	CommandName() string

	ID() types.ID
	SetID(id types.ID)
	WithID(id types.ID) Context

	ExecWithProgress(func() error) error

	// TODO v0との互換性維持用、あとで消す
	PrintWarning(warn string)
}

type cliContext struct {
	parentCtx     context.Context
	option        *config.Config
	output        output.Output
	cliIO         IO
	args          []string
	changeHandler changeHandler

	resourceName string
	commandName  string
	id           types.ID
}

// changeHandler usacloud v0の互換性維持のための実装
type changeHandler interface {
	Changed(string) bool
}

func NewCLIContext(resourceName, commandName string, globalFlags *pflag.FlagSet, args []string, parameter interface{}) (Context, error) {
	// TODO あとでグローバルなタイムアウトなどを実装する
	ctx := context.TODO()

	io := newIO()

	option, err := config.LoadConfigValue(globalFlags, io.Err())
	if err != nil {
		return nil, err
	}

	return &cliContext{
		parentCtx:    ctx,
		option:       option,
		output:       getOutputWriter(io, parameter),
		resourceName: resourceName,
		commandName:  commandName,
		cliIO:        io,
		args:         args,
	}, nil
}

func (c *cliContext) IO() IO {
	return c.cliIO
}

func (c *cliContext) Option() *config.Config {
	return c.option
}

func (c *cliContext) Output() output.Output {
	return c.output
}

func (c *cliContext) ResourceName() string {
	return c.resourceName
}

func (c *cliContext) CommandName() string {
	return c.commandName
}

func (c *cliContext) ID() types.ID {
	return c.id
}

func (c *cliContext) SetID(id types.ID) {
	c.id = id
}

func (c *cliContext) WithID(id types.ID) Context {
	return &cliContext{
		parentCtx:    c,
		option:       c.option,
		output:       c.output,
		cliIO:        c.cliIO,
		args:         c.args,
		resourceName: c.resourceName,
		commandName:  c.commandName,
		id:           id,
	}
}

func (c *cliContext) ExecWithProgress(f func() error) error {
	return NewProgress(c).Exec(f)
}

func (c *cliContext) Client() sacloud.APICaller {
	o := c.Option()
	return api.NewCaller(&api.CallerOptions{
		AccessToken:          o.AccessToken,
		AccessTokenSecret:    o.AccessTokenSecret,
		APIRootURL:           o.APIRootURL,
		AcceptLanguage:       o.AcceptLanguage,
		HTTPClient:           http.DefaultClient,
		HTTPRequestTimeout:   o.HTTPRequestTimeout,
		HTTPRequestRateLimit: o.HTTPRequestRateLimit,
		RetryMax:             o.RetryMax,
		RetryWaitMax:         o.RetryWaitMax,
		RetryWaitMin:         o.RetryWaitMin,
		UserAgent:            fmt.Sprintf("Usacloud/v%s (+https://github.com/sacloud/usacloud) libsacloud/%s", version.Version, libsacloud.Version),
		TraceAPI:             o.EnableAPITrace(),
		TraceHTTP:            o.EnableHTTPTrace(),
		FakeMode:             o.FakeMode,
		FakeStorePath:        o.FakeStorePath,
	})
}

func (c *cliContext) Zone() string {
	return c.Option().Zone
}

func (c *cliContext) Deadline() (deadline time.Time, ok bool) {
	return c.parentCtx.Deadline()
}

func (c *cliContext) Done() <-chan struct{} {
	return c.parentCtx.Done()
}

func (c *cliContext) Err() error {
	return c.parentCtx.Err()
}

func (c *cliContext) Value(key interface{}) interface{} {
	return c.parentCtx.Value(key)
}

func (c *cliContext) Args() []string {
	return c.args
}

func getOutputWriter(io IO, rawFormatter interface{}) output.Output {
	if rawFormatter == nil {
		return nil
	}
	formatter, ok := rawFormatter.(output.Formatter)
	if !ok {
		return nil
	}

	out := io.Out()
	err := io.Err()

	if formatter.GetQuiet() {
		return output.NewIDOutput(out, err)
	}
	if formatter.GetFormat() != "" || formatter.GetFormatFile() != "" {
		return output.NewFreeOutput(out, err, formatter)
	}
	switch formatter.GetOutputType() {
	case "csv":
		return output.NewRowOutput(out, err, ',', formatter)
	case "tsv":
		return output.NewRowOutput(out, err, '\t', formatter)
	case "json":
		query := formatter.GetQuery()
		if query == "" {
			bQuery, _ := ioutil.ReadFile(formatter.GetQueryFile()) // nolint: err was already checked
			query = string(bQuery)
		}
		return output.NewJSONOutput(out, err, query)
	case "yaml":
		return output.NewYAMLOutput(out, err)
	default:
		return output.NewTableOutput(out, err, formatter)
	}
}

// TODO v0との互換性維持用、実装する場所を再考
func (c *cliContext) PrintWarning(warn string) {
	if warn == "" {
		return
	}
	if c.option.NoColor {
		fmt.Fprintf(c.IO().Err(), "[WARN] %s\n", warn)
	} else {
		out := color.New(color.FgYellow)
		out.Fprintf(c.IO().Err(), "[WARN] %s\n", warn)
	}
}
<commit_msg>lint: structcheck<commit_after>// Copyright 2017-2020 The Usacloud 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 cli

import (
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	"github.com/sacloud/libsacloud/v2"

	"github.com/fatih/color"
	"github.com/sacloud/libsacloud/v2/helper/api"
	"github.com/sacloud/libsacloud/v2/sacloud"
	"github.com/sacloud/libsacloud/v2/sacloud/types"
	"github.com/sacloud/usacloud/pkg/config"
	"github.com/sacloud/usacloud/pkg/output"
	"github.com/sacloud/usacloud/pkg/version"
	"github.com/spf13/pflag"
)

type Context interface {
	Option() *config.Config
	Output() output.Output
	Client() sacloud.APICaller
	Zone() string
	IO() IO
	context.Context

	Args() []string

	ResourceName() string
	CommandName() string

	ID() types.ID
	SetID(id types.ID)
	WithID(id types.ID) Context

	ExecWithProgress(func() error) error

	// TODO v0との互換性維持用、あとで消す
	PrintWarning(warn string)
}

type cliContext struct {
	parentCtx context.Context
	option    *config.Config
	output    output.Output
	cliIO     IO
	args      []string

	resourceName string
	commandName  string
	id           types.ID
}

func NewCLIContext(resourceName, commandName string, globalFlags *pflag.FlagSet, args []string, parameter interface{}) (Context, error) {
	// TODO あとでグローバルなタイムアウトなどを実装する
	ctx := context.TODO()

	io := newIO()

	option, err := config.LoadConfigValue(globalFlags, io.Err())
	if err != nil {
		return nil, err
	}

	return &cliContext{
		parentCtx:    ctx,
		option:       option,
		output:       getOutputWriter(io, parameter),
		resourceName: resourceName,
		commandName:  commandName,
		cliIO:        io,
		args:         args,
	}, nil
}

func (c *cliContext) IO() IO {
	return c.cliIO
}

func (c *cliContext) Option() *config.Config {
	return c.option
}

func (c *cliContext) Output() output.Output {
	return c.output
}

func (c *cliContext) ResourceName() string {
	return c.resourceName
}

func (c *cliContext) CommandName() string {
	return c.commandName
}

func (c *cliContext) ID() types.ID {
	return c.id
}

func (c *cliContext) SetID(id types.ID) {
	c.id = id
}

func (c *cliContext) WithID(id types.ID) Context {
	return &cliContext{
		parentCtx:    c,
		option:       c.option,
		output:       c.output,
		cliIO:        c.cliIO,
		args:         c.args,
		resourceName: c.resourceName,
		commandName:  c.commandName,
		id:           id,
	}
}

func (c *cliContext) ExecWithProgress(f func() error) error {
	return NewProgress(c).Exec(f)
}

func (c *cliContext) Client() sacloud.APICaller {
	o := c.Option()
	return api.NewCaller(&api.CallerOptions{
		AccessToken:          o.AccessToken,
		AccessTokenSecret:    o.AccessTokenSecret,
		APIRootURL:           o.APIRootURL,
		AcceptLanguage:       o.AcceptLanguage,
		HTTPClient:           http.DefaultClient,
		HTTPRequestTimeout:   o.HTTPRequestTimeout,
		HTTPRequestRateLimit: o.HTTPRequestRateLimit,
		RetryMax:             o.RetryMax,
		RetryWaitMax:         o.RetryWaitMax,
		RetryWaitMin:         o.RetryWaitMin,
		UserAgent:            fmt.Sprintf("Usacloud/v%s (+https://github.com/sacloud/usacloud) libsacloud/%s", version.Version, libsacloud.Version),
		TraceAPI:             o.EnableAPITrace(),
		TraceHTTP:            o.EnableHTTPTrace(),
		FakeMode:             o.FakeMode,
		FakeStorePath:        o.FakeStorePath,
	})
}

func (c *cliContext) Zone() string {
	return c.Option().Zone
}

func (c *cliContext) Deadline() (deadline time.Time, ok bool) {
	return c.parentCtx.Deadline()
}

func (c *cliContext) Done() <-chan struct{} {
	return c.parentCtx.Done()
}

func (c *cliContext) Err() error {
	return c.parentCtx.Err()
}

func (c *cliContext) Value(key interface{}) interface{} {
	return c.parentCtx.Value(key)
}

func (c *cliContext) Args() []string {
	return c.args
}

func getOutputWriter(io IO, rawFormatter interface{}) output.Output {
	if rawFormatter == nil {
		return nil
	}
	formatter, ok := rawFormatter.(output.Formatter)
	if !ok {
		return nil
	}

	out := io.Out()
	err := io.Err()

	if formatter.GetQuiet() {
		return output.NewIDOutput(out, err)
	}
	if formatter.GetFormat() != "" || formatter.GetFormatFile() != "" {
		return output.NewFreeOutput(out, err, formatter)
	}
	switch formatter.GetOutputType() {
	case "csv":
		return output.NewRowOutput(out, err, ',', formatter)
	case "tsv":
		return output.NewRowOutput(out, err, '\t', formatter)
	case "json":
		query := formatter.GetQuery()
		if query == "" {
			bQuery, _ := ioutil.ReadFile(formatter.GetQueryFile()) // nolint: err was already checked
			query = string(bQuery)
		}
		return output.NewJSONOutput(out, err, query)
	case "yaml":
		return output.NewYAMLOutput(out, err)
	default:
		return output.NewTableOutput(out, err, formatter)
	}
}

// TODO v0との互換性維持用、実装する場所を再考
func (c *cliContext) PrintWarning(warn string) {
	if warn == "" {
		return
	}
	if c.option.NoColor {
		fmt.Fprintf(c.IO().Err(), "[WARN] %s\n", warn)
	} else {
		out := color.New(color.FgYellow)
		out.Fprintf(c.IO().Err(), "[WARN] %s\n", warn)
	}
}
<|endoftext|>


<commit_before>package db

import (
	r "github.com/dancannon/gorethink"
	"github.com/materials-commons/config"
)

// RSession creates a new RethinkDB session.
func RSession() (*r.Session, error) {
	return r.Connect(
		r.ConnectOpts{
			Address:  config.GetString("MCDB_CONNECTION"),
			Database: config.GetString("MCDB_NAME"),
		})
}

// RSessionMust creates a new RethinkDB session and panics if it cannot
// allocate it.
func RSessionMust() *r.Session {
	session, err := RSession()
	if err != nil {
		panic("Couldn't get new rethinkdb session")
	}
	return session
}

// RSessionUsing createa new RethinkDB session using the passed in parameters
func RSessionUsing(address, db string) (*r.Session, error) {
	return r.Connect(
		r.ConnectOpts{
			Address:  address,
			Database: db,
		})
}

// RSessionUsingMust creates a new RethinkDB session and panics if it cannot
// allocate it.
func RSessionUsingMust(address, db string) *r.Session {
	session, err := RSessionUsing(address, db)
	if err != nil {
		panic("Couldn't get new rethinkdb session")
	}
	return session
}
<commit_msg>Fix documentation comment.<commit_after>package db

import (
	r "github.com/dancannon/gorethink"
	"github.com/materials-commons/config"
)

// RSession creates a new RethinkDB session.
func RSession() (*r.Session, error) {
	return r.Connect(
		r.ConnectOpts{
			Address:  config.GetString("MCDB_CONNECTION"),
			Database: config.GetString("MCDB_NAME"),
		})
}

// RSessionMust creates a new RethinkDB session and panics if it cannot
// allocate it.
func RSessionMust() *r.Session {
	session, err := RSession()
	if err != nil {
		panic("Couldn't get new rethinkdb session")
	}
	return session
}

// RSessionUsing create a new RethinkDB session using the passed in parameters
func RSessionUsing(address, db string) (*r.Session, error) {
	return r.Connect(
		r.ConnectOpts{
			Address:  address,
			Database: db,
		})
}

// RSessionUsingMust creates a new RethinkDB session and panics if it cannot
// allocate it.
func RSessionUsingMust(address, db string) *r.Session {
	session, err := RSessionUsing(address, db)
	if err != nil {
		panic("Couldn't get new rethinkdb session")
	}
	return session
}
<|endoftext|>