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

package command

import (
	"fmt"
	"os"
	"strings"
	"testing"
	"time"
)

type MockCommand struct {
	callCount      int
	returnedOutput string
	returnedError  error
}

type MockRunner struct {
	commands          map[string]*MockCommand
	statFiles         map[string]os.FileInfo
	expectedMkdirAlls map[string]bool
	t                 *testing.T
}

type minimalFileInfo struct {
	name     string
	size     int64
	fileMode os.FileMode
	modTime  time.Time
}

func (f minimalFileInfo) Name() string {
	return f.name
}

func (f minimalFileInfo) Size() int64 {
	return f.size
}

func (f minimalFileInfo) Mode() os.FileMode {
	return f.fileMode
}

func (f minimalFileInfo) ModTime() time.Time {
	return f.modTime
}

func (f minimalFileInfo) IsDir() bool {
	return f.fileMode.IsDir()
}

func (f minimalFileInfo) Sys() interface{} {
	return nil
}

func NewMockRunner(t *testing.T) *MockRunner {
	return &MockRunner{commands: map[string]*MockCommand{}, statFiles: map[string]os.FileInfo{}, expectedMkdirAlls: map[string]bool{}, t: t}
}

func (m *MockRunner) Run(commandAndArgs ...string) (string, error) {
	commandString := strings.Join(commandAndArgs, " ")
	if _, found := m.commands[commandString]; !found {
		m.t.Fatal(fmt.Sprintf("Unexpected os command called: %s", commandString))
	}
	m.commands[commandString].callCount += 1
	return m.commands[commandString].returnedOutput, m.commands[commandString].returnedError
}

func (m *MockRunner) MkdirAll(path string, perm os.FileMode) error {
	if _, found := m.expectedMkdirAlls[path]; found {
		m.expectedMkdirAlls[path] = true
		return nil
	} else {
		return fmt.Errorf("MkdirAll() called on unexpected path: %s", path)
	}
}

func (m *MockRunner) Stat(path string) (os.FileInfo, error) {
	if fileInfo, found := m.statFiles[path]; found {
		return fileInfo, nil
	} else {
		return minimalFileInfo{}, fmt.Errorf("MockRunner.Stat(): No such file or directory: %s", path)
	}
}

func (m *MockRunner) OutputOnCall(commandAndArgs string, output string) {
	m.commands[commandAndArgs] = &MockCommand{callCount: 0, returnedOutput: output, returnedError: nil}
}

func (m *MockRunner) ErrorOnCall(commandAndArgs string, err error) {
	m.commands[commandAndArgs] = &MockCommand{callCount: 0, returnedOutput: "", returnedError: err}
}

func (m *MockRunner) RegisterMkdirAll(path string) {
	m.expectedMkdirAlls[path] = false
}

func (m *MockRunner) RegisterDeviceForStat(path string) {
	m.statFiles[path] = minimalFileInfo{name: path, fileMode: os.ModeDevice}
}

func (m *MockRunner) RegisterDirectoryForStat(path string) {
	m.statFiles[path] = minimalFileInfo{name: path, fileMode: os.ModeDir}
}

func (m *MockRunner) AssertCalled(commandAndArgs string) {
	command, found := m.commands[commandAndArgs]
	if !found || command.callCount == 0 {
		m.t.Fatal(fmt.Sprintf("Expected os command not called: %s", commandAndArgs))
	}
}

func (m *MockRunner) AssertAllCalled() {
	for commandAndArgs, command := range m.commands {
		if command.callCount == 0 {
			m.t.Fatal(fmt.Sprintf("Expected os command not called: %s", commandAndArgs))
		}
	}
	for mkdirAllPath, called := range m.expectedMkdirAlls {
		if !called {
			m.t.Fatal(fmt.Sprintf("Expected os.MkdirAll() not called: %s", mkdirAllPath))
		}
	}
}
<commit_msg>Add a flag to command.MockRunner<commit_after>// Copyright 2018 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package command

import (
	"fmt"
	"os"
	"strings"
	"testing"
	"time"
)

type MockCommand struct {
	callCount      int
	returnedOutput string
	returnedError  error
}

type MockRunner struct {
	FailOnUnexpectedCalls bool
	commands              map[string]*MockCommand
	statFiles             map[string]os.FileInfo
	expectedMkdirAlls     map[string]bool
	t                     *testing.T
}

type minimalFileInfo struct {
	name     string
	size     int64
	fileMode os.FileMode
	modTime  time.Time
}

func (f minimalFileInfo) Name() string {
	return f.name
}

func (f minimalFileInfo) Size() int64 {
	return f.size
}

func (f minimalFileInfo) Mode() os.FileMode {
	return f.fileMode
}

func (f minimalFileInfo) ModTime() time.Time {
	return f.modTime
}

func (f minimalFileInfo) IsDir() bool {
	return f.fileMode.IsDir()
}

func (f minimalFileInfo) Sys() interface{} {
	return nil
}

func NewMockRunner(t *testing.T) *MockRunner {
	return &MockRunner{FailOnUnexpectedCalls: true, commands: map[string]*MockCommand{}, statFiles: map[string]os.FileInfo{}, expectedMkdirAlls: map[string]bool{}, t: t}
}

func (m *MockRunner) Run(commandAndArgs ...string) (string, error) {
	commandString := strings.Join(commandAndArgs, " ")
	if _, found := m.commands[commandString]; !found && m.FailOnUnexpectedCalls {
		m.t.Fatal(fmt.Sprintf("Unexpected os command called: %s", commandString))
	}
	m.commands[commandString].callCount += 1
	return m.commands[commandString].returnedOutput, m.commands[commandString].returnedError
}

func (m *MockRunner) MkdirAll(path string, perm os.FileMode) error {
	if _, found := m.expectedMkdirAlls[path]; !found && m.FailOnUnexpectedCalls {
		return fmt.Errorf("MkdirAll() called on unexpected path: %s", path)
	}
	m.expectedMkdirAlls[path] = true
	return nil
}

func (m *MockRunner) Stat(path string) (os.FileInfo, error) {
	fileInfo, found := m.statFiles[path]
	if !found && m.FailOnUnexpectedCalls {
		return minimalFileInfo{}, fmt.Errorf("MockRunner.Stat(): No such file or directory: %s", path)
	}
	return fileInfo, nil
}

func (m *MockRunner) OutputOnCall(commandAndArgs string, output string) {
	m.commands[commandAndArgs] = &MockCommand{callCount: 0, returnedOutput: output, returnedError: nil}
}

func (m *MockRunner) ErrorOnCall(commandAndArgs string, err error) {
	m.commands[commandAndArgs] = &MockCommand{callCount: 0, returnedOutput: "", returnedError: err}
}

func (m *MockRunner) RegisterMkdirAll(path string) {
	m.expectedMkdirAlls[path] = false
}

func (m *MockRunner) RegisterDeviceForStat(path string) {
	m.statFiles[path] = minimalFileInfo{name: path, fileMode: os.ModeDevice}
}

func (m *MockRunner) RegisterDirectoryForStat(path string) {
	m.statFiles[path] = minimalFileInfo{name: path, fileMode: os.ModeDir}
}

func (m *MockRunner) AssertCalled(commandAndArgs string) {
	command, found := m.commands[commandAndArgs]
	if !found || command.callCount == 0 {
		m.t.Fatal(fmt.Sprintf("Expected os command not called: %s", commandAndArgs))
	}
}

func (m *MockRunner) AssertAllCalled() {
	for commandAndArgs, command := range m.commands {
		if command.callCount == 0 {
			m.t.Fatal(fmt.Sprintf("Expected os command not called: %s", commandAndArgs))
		}
	}
	for mkdirAllPath, called := range m.expectedMkdirAlls {
		if !called {
			m.t.Fatal(fmt.Sprintf("Expected os.MkdirAll() not called: %s", mkdirAllPath))
		}
	}
}
<|endoftext|>


<commit_before>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package common

import (
	stdErr "errors"

	"github.com/webx-top/captcha"
	"github.com/webx-top/echo"
	"github.com/webx-top/echo/middleware/tplfunc"
	"github.com/webx-top/echo/subdomains"
)

// Ok 操作成功
func Ok(v string) Successor {
	return NewOk(v)
}

// Err 获取错误信息
func Err(ctx echo.Context, err error) (ret interface{}) {
	if err == nil {
		flash := ctx.Flash()
		if flash != nil {
			if errMsg, ok := flash.(string); ok {
				ret = stdErr.New(errMsg)
			} else {
				ret = flash
			}
		}
	} else {
		ret = err
	}
	return
}

// SendOk 记录成功信息
func SendOk(ctx echo.Context, msg string) {
	if ctx.IsAjax() || ctx.Format() != `html` {
		ctx.Data().SetInfo(msg, 1)
		return
	}
	ctx.Session().AddFlash(Ok(msg))
}

// SendFail 记录失败信息
func SendFail(ctx echo.Context, msg string) {
	if ctx.IsAjax() || ctx.Format() != `html` {
		ctx.Data().SetInfo(msg, 0)
		return
	}
	ctx.Session().AddFlash(msg)
}

// SendErr 记录错误信息 (SendFail的别名)
func SendErr(ctx echo.Context, err error) {
	SendFail(ctx, err.Error())
}

// VerifyCaptcha 验证码验证
func VerifyCaptcha(ctx echo.Context, hostAlias string, captchaName string, args ...string) echo.Data {
	data := ctx.Data()
	idGet := ctx.Form
	if len(args) > 0 {
		idGet = func(_ string, defaults ...string) string {
			return ctx.Form(args[0], defaults...)
		}
	}
	code := ctx.Form(captchaName)
	var err error
	if len(code) == 0 {
		err = ctx.E(`请输入验证码`)
	} else if !tplfunc.CaptchaVerify(code, idGet) {
		err = ctx.E(`验证码不正确`)
	}
	if err != nil {
		data.SetZone(captchaName)
		data.SetData(CaptchaInfo(hostAlias, captchaName, args...))
		data.SetError(err, StatusCaptchaError)
	}
	return data
}

// VerifyAndSetCaptcha 验证码验证并设置新验证码信息
func VerifyAndSetCaptcha(ctx echo.Context, hostAlias string, captchaName string, args ...string) echo.Data {
	data := VerifyCaptcha(ctx, hostAlias, captchaName, args...)
	if data.GetCode() != StatusCaptchaError {
		data.SetData(CaptchaInfo(hostAlias, captchaName, args...))
	}
	return data
}

// CaptchaInfo 新验证码信息
func CaptchaInfo(hostAlias string, captchaName string, args ...string) echo.H {
	captchaID := captcha.New()
	captchaIdent := `captchaId`
	if len(args) > 0 {
		captchaIdent = args[0]
	}
	return echo.H{
		`captchaName`:  captchaName,
		`captchaIdent`: captchaIdent,
		`captchaID`:    captchaID,
		`captchaURL`:   subdomains.Default.URL(`/captcha/`+captchaID+`.png`, hostAlias),
	}
}
<commit_msg>update<commit_after>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package common

import (
	stdErr "errors"

	"github.com/webx-top/captcha"
	"github.com/webx-top/echo"
	"github.com/webx-top/echo/middleware/tplfunc"
	"github.com/webx-top/echo/subdomains"
)

// Ok 操作成功
func Ok(v string) Successor {
	return NewOk(v)
}

// Err 获取错误信息
func Err(ctx echo.Context, err error) (ret interface{}) {
	if err == nil {
		flash := ctx.Flash()
		if flash != nil {
			if errMsg, ok := flash.(string); ok {
				ret = stdErr.New(errMsg)
			} else {
				ret = flash
			}
		}
	} else {
		ret = err
	}
	return
}

// SendOk 记录成功信息
func SendOk(ctx echo.Context, msg string) {
	if ctx.IsAjax() || ctx.Format() != `html` {
		ctx.Data().SetInfo(msg, 1)
		return
	}
	ctx.Session().AddFlash(Ok(msg))
}

// SendFail 记录失败信息
func SendFail(ctx echo.Context, msg string) {
	if ctx.IsAjax() || ctx.Format() != `html` {
		ctx.Data().SetInfo(msg, 0)
		return
	}
	ctx.Session().AddFlash(msg)
}

// SendErr 记录错误信息 (SendFail的别名)
func SendErr(ctx echo.Context, err error) {
	SendFail(ctx, err.Error())
}

// VerifyCaptcha 验证码验证
func VerifyCaptcha(ctx echo.Context, hostAlias string, captchaName string, args ...string) echo.Data {
	data := ctx.Data()
	idGet := ctx.Form
	if len(args) > 0 {
		idGet = func(_ string, defaults ...string) string {
			return ctx.Form(args[0], defaults...)
		}
	}
	code := ctx.Form(captchaName)
	var err error
	if len(code) == 0 {
		err = ctx.E(`请输入验证码`)
	} else if !tplfunc.CaptchaVerify(code, idGet) {
		err = ctx.E(`验证码不正确`)
	}
	if err != nil {
		data.SetZone(captchaName)
		data.SetData(CaptchaInfo(hostAlias, captchaName, args...))
		data.SetError(err, StatusCaptchaError)
	}
	return data
}

// VerifyAndSetCaptcha 验证码验证并设置新验证码信息
func VerifyAndSetCaptcha(ctx echo.Context, hostAlias string, captchaName string, args ...string) echo.Data {
	data := VerifyCaptcha(ctx, hostAlias, captchaName, args...)
	if data.GetCode() != StatusCaptchaError {
		data.SetData(CaptchaInfo(hostAlias, captchaName, args...))
	}
	return data
}

// CaptchaInfo 新验证码信息
func CaptchaInfo(hostAlias string, captchaName string, args ...string) echo.H {
	captchaID := captcha.New()
	captchaIdent := `captchaId`
	if len(args) > 0 {
		captchaIdent = args[0]
	}
	return echo.H{
		`captchaName`:  captchaName,
		`captchaIdent`: captchaIdent,
		`captchaID`:    captchaID,
		`captchaURL`:   subdomains.Default.URL(`/captcha/`+captchaID+`.png`, hostAlias),
	}
}

type ConfigFromDB interface {
	ConfigFromDB() echo.H
}
<|endoftext|>


<commit_before>/*
  A simple kafka consumer-group client

  Copyright 2016 MistSys
*/

package stable_test

import (
	"fmt"
	"testing"

	"github.com/Shopify/sarama"
	consumer "github.com/mistsys/sarama-consumer"
	"github.com/mistsys/sarama-consumer/stable"
)

func TestGreenfield(t *testing.T) {
	var rr consumer.Partitioner = stable.Stable

	all_topics := []string{"topic1", "topic2", "topic3", "topic4"}

	var mock_client = mockClient{
		config: sarama.NewConfig(),
		partitions: map[string][]int32{
			"topic1": []int32{0, 1 /*,no 2*/, 3, 4, 5, 6, 7}, // note we pretend partition 2 is offline
			"topic2": []int32{0, 1},
			"topic3": []int32{0, 1, 2, 3, 4, 5}, // exactly 2 partitions per member
			"topic4": []int32{0, 1, 2},          // exactly 1 partitions per member
		},
	}

	// pretend to have 3 members, with no current assignments, all asking for the four topics
	var jreqs [3]sarama.JoinGroupRequest
	for i := range jreqs {
		jreqs[i].GroupId = "group"
		jreqs[i].MemberId = fmt.Sprintf("member%d", i)
		jreqs[i].ProtocolType = "consumer"
		rr.PrepareJoin(&jreqs[i], all_topics, nil)

		t.Logf("JoinGroupRequests[%d] = %v\n", i, jreqs[i])
	}

	var jresp = sarama.JoinGroupResponse{
		GenerationId:  1,
		GroupProtocol: string(stable.Stable),
		Members:       make(map[string][]byte),
	}
	for i := range jreqs {
		jresp.Members[jreqs[i].MemberId] = jreqs[i].GroupProtocols[string(stable.Stable)]
	}
	t.Logf("JoinGroupResponse = %v\n", jresp)

	var sreq = sarama.SyncGroupRequest{
		GroupId:      "group",
		GenerationId: 1,
		MemberId:     "member0",
	}
	err := rr.Partition(&sreq, &jresp, &mock_client)
	t.Logf("SyncGroupRequest = %v\n", sreq)
	if err != nil {
		t.Fatal(err)
	}

	assignments := make(map[string]map[string][]int32) // map of member to topic to the assigned list of partitions
	for i := range jreqs {
		id := jreqs[i].MemberId
		var sresp = sarama.SyncGroupResponse{
			MemberAssignment: sreq.GroupAssignments[id],
		}

		act, err := rr.ParseSync(&sresp)
		if err != nil {
			t.Fatal(err)
		}

		t.Logf("%s assignment %v\n", id, act)

		assignments[id] = act
	}

	for _, topic := range all_topics {
		low := 1 << 30
		high := 0
		used := make(map[int32]string)
		for id, act := range assignments {
			a := act[topic]

			// keep track of the high and low # of partitions that have been assigned
			n := len(a)
			if low > n {
				low = n
			}
			if high < n {
				high = n
			}

			// and make sure no partition is assigned twice
			for id2, p := range a {
				if _, ok := used[p]; ok {
					t.Errorf("Partition %d of topic %q is assigned to multiple consumers, including %v and %v", p, topic, id, id2)
				}
				used[p] = id
			}
		}
		t.Logf("topic %q assigned partitions as %v", topic, used)
		if low != high && low != high-1 {
			t.Errorf("Partition assignment of topic %q is uneven. Some have %d and some have %d", topic, low, high)
		}
	}
}

// mock sarama.Client which implements the metadata API sufficiently for our unit test purposes
type mockClient struct {
	config     *sarama.Config
	partitions map[string][]int32
}

func (mc *mockClient) Config() *sarama.Config {
	return mc.config
}

func (mc *mockClient) Topics() ([]string, error) {
	var topics = make([]string, 0, len(mc.partitions))
	for t := range mc.partitions {
		topics = append(topics, t)
	}
	return topics, nil
}

func (mc *mockClient) Partitions(topic string) ([]int32, error) {
	if p, ok := mc.partitions[topic]; ok {
		return p, nil
	}
	return nil, sarama.ErrUnknownTopicOrPartition
}

func (mc *mockClient) WritablePartitions(topic string) ([]int32, error) {
	return mc.Partitions(topic)
}

func (mc *mockClient) Leader(topic string, part int32) (*sarama.Broker, error) {
	return nil, sarama.ErrBrokerNotAvailable
}

func (mc *mockClient) Replicas(topic string, part int32) ([]int32, error) {
	return nil, sarama.ErrNotEnoughReplicas
}

func (mc *mockClient) RefreshMetadata(topics ...string) error {
	return nil
}

func (mc *mockClient) GetOffset(topic string, part int32, time int64) (int64, error) {
	return 0, nil
}

func (mc *mockClient) Coordinator(group string) (*sarama.Broker, error) {
	return nil, sarama.ErrBrokerNotAvailable
}

func (mc *mockClient) RefreshCoordinator(group string) error {
	return nil
}

func (mc *mockClient) Close() error {
	return nil
}

func (mc *mockClient) Closed() bool {
	return false
}
<commit_msg>move sanity checks into their own function<commit_after>/*
  A simple kafka consumer-group client

  Copyright 2016 MistSys
*/

package stable_test

import (
	"fmt"
	"testing"

	"github.com/Shopify/sarama"
	consumer "github.com/mistsys/sarama-consumer"
	"github.com/mistsys/sarama-consumer/stable"
)

func TestGreenfield(t *testing.T) {
	var partitioner consumer.Partitioner = stable.Stable

	all_topics := []string{"topic1", "topic2", "topic3", "topic4"}

	var mock_client = mockClient{
		config: sarama.NewConfig(),
		partitions: map[string][]int32{
			"topic1": []int32{0, 1 /*,no 2*/, 3, 4, 5, 6, 7}, // note we pretend partition 2 is offline
			"topic2": []int32{0, 1},
			"topic3": []int32{0, 1, 2, 3, 4, 5}, // exactly 2 partitions per member
			"topic4": []int32{0, 1, 2},          // exactly 1 partitions per member
		},
	}

	// pretend to have 3 members, with no current assignments, all asking for the four topics
	var jreqs [3]sarama.JoinGroupRequest
	for i := range jreqs {
		jreqs[i].GroupId = "group"
		jreqs[i].MemberId = fmt.Sprintf("member%d", i)
		jreqs[i].ProtocolType = "consumer"
		partitioner.PrepareJoin(&jreqs[i], all_topics, nil)

		t.Logf("JoinGroupRequests[%d] = %v\n", i, jreqs[i])
	}

	var jresp = sarama.JoinGroupResponse{
		GenerationId:  1,
		GroupProtocol: string(stable.Stable),
		Members:       make(map[string][]byte),
	}
	for i := range jreqs {
		jresp.Members[jreqs[i].MemberId] = jreqs[i].GroupProtocols[string(stable.Stable)]
	}
	t.Logf("JoinGroupResponse = %v\n", jresp)

	var sreq = sarama.SyncGroupRequest{
		GroupId:      "group",
		GenerationId: 1,
		MemberId:     "member0",
	}
	err := partitioner.Partition(&sreq, &jresp, &mock_client)
	t.Logf("SyncGroupRequest = %v\n", sreq)
	if err != nil {
		t.Fatal(err)
	}

	sanity_check(&sreq, partitioner, jreqs[:], all_topics, t)
}

// sanity check a set of assignments for basic correctness
func sanity_check(sreq *sarama.SyncGroupRequest, partitioner consumer.Partitioner, jreqs []sarama.JoinGroupRequest, topics []string, t *testing.T) {

	assignments := make(map[string]map[string][]int32) // map of member to topic to the assigned list of partitions
	for i := range jreqs {
		id := jreqs[i].MemberId
		var sresp = sarama.SyncGroupResponse{
			MemberAssignment: sreq.GroupAssignments[id],
		}

		act, err := partitioner.ParseSync(&sresp)
		if err != nil {
			t.Errorf("failed to parse SyncGroupRespons data: %v", err)
		}

		t.Logf("%s assignment %v\n", id, act)

		assignments[id] = act
	}

	for _, topic := range topics {
		low := 1 << 30
		high := 0
		used := make(map[int32]string)
		for id, act := range assignments {
			a := act[topic]

			// keep track of the high and low # of partitions that have been assigned
			n := len(a)
			if low > n {
				low = n
			}
			if high < n {
				high = n
			}

			// and make sure no partition is assigned twice
			for id2, p := range a {
				if _, ok := used[p]; ok {
					t.Errorf("Partition %d of topic %q is assigned to multiple consumers, including %v and %v", p, topic, id, id2)
				}
				used[p] = id
			}
		}
		t.Logf("topic %q assigned partitions as %v", topic, used)
		if low != high && low != high-1 {
			t.Errorf("Partition assignment of topic %q is uneven. Some have %d and some have %d", topic, low, high)
		}
	}
}

// mock sarama.Client which implements the metadata API sufficiently for our unit test purposes
type mockClient struct {
	config     *sarama.Config
	partitions map[string][]int32
}

func (mc *mockClient) Config() *sarama.Config {
	return mc.config
}

func (mc *mockClient) Topics() ([]string, error) {
	var topics = make([]string, 0, len(mc.partitions))
	for t := range mc.partitions {
		topics = append(topics, t)
	}
	return topics, nil
}

func (mc *mockClient) Partitions(topic string) ([]int32, error) {
	if p, ok := mc.partitions[topic]; ok {
		return p, nil
	}
	return nil, sarama.ErrUnknownTopicOrPartition
}

func (mc *mockClient) WritablePartitions(topic string) ([]int32, error) {
	return mc.Partitions(topic)
}

func (mc *mockClient) Leader(topic string, part int32) (*sarama.Broker, error) {
	return nil, sarama.ErrBrokerNotAvailable
}

func (mc *mockClient) Replicas(topic string, part int32) ([]int32, error) {
	return nil, sarama.ErrNotEnoughReplicas
}

func (mc *mockClient) RefreshMetadata(topics ...string) error {
	return nil
}

func (mc *mockClient) GetOffset(topic string, part int32, time int64) (int64, error) {
	return 0, nil
}

func (mc *mockClient) Coordinator(group string) (*sarama.Broker, error) {
	return nil, sarama.ErrBrokerNotAvailable
}

func (mc *mockClient) RefreshCoordinator(group string) error {
	return nil
}

func (mc *mockClient) Close() error {
	return nil
}

func (mc *mockClient) Closed() bool {
	return false
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/containers/standard/schema"
	"github.com/containers/standard/schema/types"
	"github.com/coreos-inc/rkt/rkt"
	"github.com/coreos/go-systemd/unit"
)

var (
	startedOnToType = map[string]string{
		"fork": "simple",
		"exit": "oneshot",
	}
)

// Container encapsulates a ContainerRuntimeManifest and AppManifests
type Container struct {
	Root     string // root directory where the container will be located
	Manifest *schema.ContainerRuntimeManifest
	Apps     map[string]*schema.AppManifest
}

// LoadContainer loads a Container Runtime Manifest (as prepared by stage0) and
// its associated Application Manifests, under $root/stage1/opt/stage1/$apphash
func LoadContainer(root string) (*Container, error) {
	c := &Container{
		Root: root,
		Apps: make(map[string]*schema.AppManifest),
	}

	buf, err := ioutil.ReadFile(rkt.ContainerManifestPath(c.Root))
	if err != nil {
		return nil, fmt.Errorf("failed reading container runtime manifest: %v", err)
	}

	cm := &schema.ContainerRuntimeManifest{}
	if err := json.Unmarshal(buf, cm); err != nil {
		return nil, fmt.Errorf("failed unmarshalling container runtime manifest: %v", err)
	}
	c.Manifest = cm

	for _, app := range c.Manifest.Apps {
		ampath := rkt.AppManifestPath(c.Root, app.ImageID.String())
		buf, err := ioutil.ReadFile(ampath)
		if err != nil {
			return nil, fmt.Errorf("failed reading app manifest %q: %v", ampath, err)
		}

		am := &schema.AppManifest{}
		if err = json.Unmarshal(buf, am); err != nil {
			return nil, fmt.Errorf("failed unmarshalling app manifest %q: %v", ampath, err)
		}
		name := am.Name.String()
		if _, ok := c.Apps[name]; ok {
			return nil, fmt.Errorf("got multiple definitions for app: %s", name)
		}
		c.Apps[name] = am
	}

	return c, nil
}

// appToSystemd transforms the provided app manifest into a systemd service unit
func (c *Container) appToSystemd(am *schema.AppManifest, id types.Hash) error {
	name := am.Name.String()
	typ, ok := startedOnToType[am.StartedOn.String()]
	if !ok {
		return fmt.Errorf("unrecognized StartedOn: %s", am.StartedOn)
	}
	execStart := fmt.Sprintf(`%s`, strings.Join(am.Exec, `" "`))
	opts := []*unit.UnitOption{
		&unit.UnitOption{"Unit", "Description", name},
		&unit.UnitOption{"Unit", "DefaultDependencies", "false"},
		&unit.UnitOption{"Service", "Type", typ},
		&unit.UnitOption{"Service", "Restart", "no"},
		&unit.UnitOption{"Service", "RootDirectory", rkt.RelAppRootfsPath(id.String())},
		&unit.UnitOption{"Service", "ExecStart", execStart},
		&unit.UnitOption{"Service", "User", am.User},
		&unit.UnitOption{"Service", "Group", am.Group},
	}

	for ek, ev := range am.Environment {
		ee := fmt.Sprintf(`"%s=%s"`, ek, ev)
		opts = append(opts, &unit.UnitOption{"Service", "Environment", ee})
	}

	file, err := os.OpenFile(rkt.ServiceFilePath(c.Root, name), os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return fmt.Errorf("failed to create service file: %v", err)
	}
	defer file.Close()

	_, err = io.Copy(file, unit.Serialize(opts))
	if err != nil {
		return fmt.Errorf("failed to write service file: %v", err)
	}

	if err = os.Symlink(path.Join("..", rkt.ServiceName(name)), rkt.WantLinkPath(c.Root, name)); err != nil {
		return fmt.Errorf("failed to link service want: %v", err)
	}

	return nil
}

// ContainerToSystemd creates the appropriate systemd service unit files for
// all the constituent apps of the Container
func (c *Container) ContainerToSystemd() error {
	if err := os.MkdirAll(rkt.ServicesPath(c.Root), 0640); err != nil {
		return fmt.Errorf("failed to create services directory: %v", err)
	}
	if err := os.MkdirAll(rkt.WantsPath(c.Root), 0640); err != nil {
		return fmt.Errorf("failed to create wants directory: %v", err)
	}
	for _, am := range c.Apps {
		id := c.Manifest.Apps[am.Name].ImageID
		if err := c.appToSystemd(am, id); err != nil {
			return fmt.Errorf("failed to transform app %q into systemd service: %v", am.Name, err)
		}
	}

	return nil
}

// appToNspawnArgs transforms the given app manifest, with the given associated
// app image id, into a subset of applicable systemd-nspawn argument
func (c *Container) appToNspawnArgs(am *schema.AppManifest, id types.Hash) ([]string, error) {
	args := []string{}
	name := am.Name.String()

	vols := make(map[types.ACLabel]types.Volume)
	for _, v := range c.Manifest.Volumes {
		for _, f := range v.Fulfills {
			vols[f] = v
		}
	}

	for key, mp := range am.MountPoints {
		vol, ok := vols[key]
		if !ok {
			return nil, fmt.Errorf("no volume for mountpoint %q in app %q", key, name)
		}
		opt := make([]string, 4)

		if mp.ReadOnly {
			opt[0] = "--bind-ro="
		} else {
			opt[0] = "--bind="
		}

		opt[1] = vol.Source
		opt[2] = ":"
		opt[3] = filepath.Join(rkt.RelAppRootfsPath(id.String()), mp.Path)

		args = append(args, strings.Join(opt, ""))
	}

	return args, nil
}

// ContainerToNspawnArgs renders a prepared Container as a systemd-nspawn
// argument list ready to be executed
func (c *Container) ContainerToNspawnArgs() ([]string, error) {
	args := []string{
		"--boot",
		"--uuid=" + c.Manifest.UUID.String(),
		"--directory=" + rkt.Stage1RootfsPath(c.Root),
		"--",
		"--default-standard-output=tty",
	}

	for _, am := range c.Apps {
		id := c.Manifest.Apps[am.Name].ImageID
		aa, err := c.appToNspawnArgs(am, id)
		if err != nil {
			return nil, fmt.Errorf("failed to construct args for app %q: %v", am.Name, err)
		}
		args = append(args, aa...)
	}

	return args, nil
}
<commit_msg>stage1: add AC_APP_NAME to the container<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/containers/standard/schema"
	"github.com/containers/standard/schema/types"
	"github.com/coreos-inc/rkt/rkt"
	"github.com/coreos/go-systemd/unit"
)

var (
	startedOnToType = map[string]string{
		"fork": "simple",
		"exit": "oneshot",
	}
)

// Container encapsulates a ContainerRuntimeManifest and AppManifests
type Container struct {
	Root     string // root directory where the container will be located
	Manifest *schema.ContainerRuntimeManifest
	Apps     map[string]*schema.AppManifest
}

// LoadContainer loads a Container Runtime Manifest (as prepared by stage0) and
// its associated Application Manifests, under $root/stage1/opt/stage1/$apphash
func LoadContainer(root string) (*Container, error) {
	c := &Container{
		Root: root,
		Apps: make(map[string]*schema.AppManifest),
	}

	buf, err := ioutil.ReadFile(rkt.ContainerManifestPath(c.Root))
	if err != nil {
		return nil, fmt.Errorf("failed reading container runtime manifest: %v", err)
	}

	cm := &schema.ContainerRuntimeManifest{}
	if err := json.Unmarshal(buf, cm); err != nil {
		return nil, fmt.Errorf("failed unmarshalling container runtime manifest: %v", err)
	}
	c.Manifest = cm

	for _, app := range c.Manifest.Apps {
		ampath := rkt.AppManifestPath(c.Root, app.ImageID.String())
		buf, err := ioutil.ReadFile(ampath)
		if err != nil {
			return nil, fmt.Errorf("failed reading app manifest %q: %v", ampath, err)
		}

		am := &schema.AppManifest{}
		if err = json.Unmarshal(buf, am); err != nil {
			return nil, fmt.Errorf("failed unmarshalling app manifest %q: %v", ampath, err)
		}
		name := am.Name.String()
		if _, ok := c.Apps[name]; ok {
			return nil, fmt.Errorf("got multiple definitions for app: %s", name)
		}
		c.Apps[name] = am
	}

	return c, nil
}

// appToSystemd transforms the provided app manifest into a systemd service unit
func (c *Container) appToSystemd(am *schema.AppManifest, id types.Hash) error {
	name := am.Name.String()
	typ, ok := startedOnToType[am.StartedOn.String()]
	if !ok {
		return fmt.Errorf("unrecognized StartedOn: %s", am.StartedOn)
	}
	execStart := fmt.Sprintf(`%s`, strings.Join(am.Exec, `" "`))
	opts := []*unit.UnitOption{
		&unit.UnitOption{"Unit", "Description", name},
		&unit.UnitOption{"Unit", "DefaultDependencies", "false"},
		&unit.UnitOption{"Service", "Type", typ},
		&unit.UnitOption{"Service", "Restart", "no"},
		&unit.UnitOption{"Service", "RootDirectory", rkt.RelAppRootfsPath(id.String())},
		&unit.UnitOption{"Service", "ExecStart", execStart},
		&unit.UnitOption{"Service", "User", am.User},
		&unit.UnitOption{"Service", "Group", am.Group},
	}

	env := am.Environment
	env["AC_APP_NAME"] = string(am.Name)
	for ek, ev := range env {
		ee := fmt.Sprintf(`"%s=%s"`, ek, ev)
		opts = append(opts, &unit.UnitOption{"Service", "Environment", ee})
	}

	file, err := os.OpenFile(rkt.ServiceFilePath(c.Root, name), os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return fmt.Errorf("failed to create service file: %v", err)
	}
	defer file.Close()

	_, err = io.Copy(file, unit.Serialize(opts))
	if err != nil {
		return fmt.Errorf("failed to write service file: %v", err)
	}

	if err = os.Symlink(path.Join("..", rkt.ServiceName(name)), rkt.WantLinkPath(c.Root, name)); err != nil {
		return fmt.Errorf("failed to link service want: %v", err)
	}

	return nil
}

// ContainerToSystemd creates the appropriate systemd service unit files for
// all the constituent apps of the Container
func (c *Container) ContainerToSystemd() error {
	if err := os.MkdirAll(rkt.ServicesPath(c.Root), 0640); err != nil {
		return fmt.Errorf("failed to create services directory: %v", err)
	}
	if err := os.MkdirAll(rkt.WantsPath(c.Root), 0640); err != nil {
		return fmt.Errorf("failed to create wants directory: %v", err)
	}
	for _, am := range c.Apps {
		id := c.Manifest.Apps[am.Name].ImageID
		if err := c.appToSystemd(am, id); err != nil {
			return fmt.Errorf("failed to transform app %q into systemd service: %v", am.Name, err)
		}
	}

	return nil
}

// appToNspawnArgs transforms the given app manifest, with the given associated
// app image id, into a subset of applicable systemd-nspawn argument
func (c *Container) appToNspawnArgs(am *schema.AppManifest, id types.Hash) ([]string, error) {
	args := []string{}
	name := am.Name.String()

	vols := make(map[types.ACLabel]types.Volume)
	for _, v := range c.Manifest.Volumes {
		for _, f := range v.Fulfills {
			vols[f] = v
		}
	}

	for key, mp := range am.MountPoints {
		vol, ok := vols[key]
		if !ok {
			return nil, fmt.Errorf("no volume for mountpoint %q in app %q", key, name)
		}
		opt := make([]string, 4)

		if mp.ReadOnly {
			opt[0] = "--bind-ro="
		} else {
			opt[0] = "--bind="
		}

		opt[1] = vol.Source
		opt[2] = ":"
		opt[3] = filepath.Join(rkt.RelAppRootfsPath(id.String()), mp.Path)

		args = append(args, strings.Join(opt, ""))
	}

	return args, nil
}

// ContainerToNspawnArgs renders a prepared Container as a systemd-nspawn
// argument list ready to be executed
func (c *Container) ContainerToNspawnArgs() ([]string, error) {
	args := []string{
		"--boot",
		"--uuid=" + c.Manifest.UUID.String(),
		"--directory=" + rkt.Stage1RootfsPath(c.Root),
		"--",
		"--default-standard-output=tty",
	}

	for _, am := range c.Apps {
		id := c.Manifest.Apps[am.Name].ImageID
		aa, err := c.appToNspawnArgs(am, id)
		if err != nil {
			return nil, fmt.Errorf("failed to construct args for app %q: %v", am.Name, err)
		}
		args = append(args, aa...)
	}

	return args, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017, 2018 Ankyra

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 state

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

	"github.com/ankyra/escape-core"
	"github.com/ankyra/escape-core/state/validate"
)

type DeploymentState struct {
	Name        string                 `json:"name"`
	Release     string                 `json:"release,omitempty"`
	Stages      map[string]*StageState `json:"stages,omitempty"`
	Inputs      map[string]interface{} `json:"inputs,omitempty"`
	environment *EnvironmentState      `json:"-"`
	parent      *DeploymentState       `json:"-"`
	parentStage *StageState            `json:"-"`
}

func NewDeploymentState(env *EnvironmentState, name, release string) (*DeploymentState, error) {
	d := &DeploymentState{
		Name:        name,
		Release:     release,
		Stages:      map[string]*StageState{},
		Inputs:      map[string]interface{}{},
		environment: env,
	}
	return d, d.validateAndFix(name, env)
}

func (d *DeploymentState) GetName() string {
	return d.Name
}

func (d *DeploymentState) Summarize() *DeploymentState {
	result, _ := NewDeploymentState(d.environment, d.Name, d.Release)
	for name, stage := range d.Stages {
		result.Stages[name] = stage.Summarize()
	}
	return result
}

func (d *DeploymentState) GetRootDeploymentName() string {
	prev := d
	p := prev
	for p != nil {
		prev = p
		p = p.parent
	}
	return prev.Name
}

func (d *DeploymentState) GetRootDeploymentStage() string {
	stage := ""
	prev := d
	p := prev
	for prev.parentStage != nil {
		stage = prev.parentStage.Name
		prev = p
		p = p.parent
	}
	return stage
}

func (d *DeploymentState) GetDeploymentPath() string {
	result := []string{}
	p := d
	for p != nil {
		result = append(result, p.Name)
		p = p.parent
	}
	for i := len(result)/2 - 1; i >= 0; i-- {
		opp := len(result) - 1 - i
		result[i], result[opp] = result[opp], result[i]
	}
	return strings.Join(result, ":")
}

func (d *DeploymentState) GetReleaseId(stage string) string {
	return d.Release + "-v" + d.GetVersion(stage)
}

func (d *DeploymentState) GetVersion(stage string) string {
	return d.GetStageOrCreateNew(stage).Version
}

func (d *DeploymentState) GetEnvironmentState() *EnvironmentState {
	return d.environment
}

func (d *DeploymentState) GetDeployment(stage, deploymentName string) (*DeploymentState, error) {
	st := d.GetStageOrCreateNew(stage)
	depl, ok := st.Deployments[deploymentName]
	if !ok {
		return nil, DeploymentDoesNotExistError(deploymentName)
	}
	depl.parentStage = st
	st.Deployments[deploymentName] = depl
	return depl, nil
}

func (d *DeploymentState) GetDeploymentOrMakeNew(stage, deploymentName string) (*DeploymentState, error) {
	st := d.GetStageOrCreateNew(stage)
	depl, ok := st.Deployments[deploymentName]
	if !ok {
		newDepl, err := NewDeploymentState(d.environment, deploymentName, deploymentName)
		if err != nil {
			return nil, err
		}
		depl = newDepl
		depl.parent = d
	}
	depl.parentStage = st
	st.Deployments[deploymentName] = depl
	return depl, nil
}

func (d *DeploymentState) GetUserInputs(stage string) map[string]interface{} {
	return d.GetStageOrCreateNew(stage).UserInputs
}

func (d *DeploymentState) GetCalculatedInputs(stage string) map[string]interface{} {
	return d.GetStageOrCreateNew(stage).Inputs
}

func (d *DeploymentState) GetCalculatedOutputs(stage string) map[string]interface{} {
	return d.GetStageOrCreateNew(stage).Outputs
}

func (d *DeploymentState) UpdateInputs(stage string, inputs map[string]interface{}) error {
	d.GetStageOrCreateNew(stage).SetInputs(inputs)
	return d.Save()
}

func (d *DeploymentState) UpdateUserInputs(stage string, inputs map[string]interface{}) error {
	d.GetStageOrCreateNew(stage).SetUserInputs(inputs)
	return d.Save()
}

func (d *DeploymentState) UpdateOutputs(stage string, outputs map[string]interface{}) error {
	d.GetStageOrCreateNew(stage).SetOutputs(outputs)
	return d.Save()
}

func (d *DeploymentState) CommitVersion(stage string, metadata *core.ReleaseMetadata) error {
	d.GetStageOrCreateNew(stage).SetVersion(metadata.Version)
	d.GetStageOrCreateNew(stage).Provides = metadata.GetProvides()
	return nil
}

func (d *DeploymentState) UpdateStatus(stage string, status *Status) error {
	d.GetStageOrCreateNew(stage).Status = status
	return d.Save()
}
func (d *DeploymentState) GetStatus(stage string) *Status {
	return d.GetStageOrCreateNew(stage).Status
}

func (d *DeploymentState) IsDeployed(stage string, metadata *core.ReleaseMetadata) bool {
	return d.GetStageOrCreateNew(stage).Version == metadata.Version
}

func (d *DeploymentState) Save() error {
	return d.environment.Save(d)
}

func (p *DeploymentState) ToJson() string {
	str, err := json.MarshalIndent(p, "", "   ")
	if err != nil {
		panic(err)
	}
	return string(str)
}

func (d *DeploymentState) SetProvider(stage, name, deplName string) {
	d.GetStageOrCreateNew(stage).Providers[name] = deplName
}

func (d *DeploymentState) GetProviders(stage string) map[string]string {
	result := map[string]string{}
	d.walkStatesAndStages(stage, func(p *DeploymentState, st *StageState) {
		for key, val := range st.Providers {
			result[key] = val
		}
	})
	return result
}

func (d *DeploymentState) ConfigureProviders(metadata *core.ReleaseMetadata, stage string, extraProviders map[string]string) error {
	configuredProviders := d.GetProviders(stage)
	availableProviders := d.environment.GetProviders()
	for _, consumerCfg := range metadata.GetConsumerConfig(stage) {
		c := consumerCfg.Name
		variable := consumerCfg.VariableName
		provider, override := extraProviders[variable]
		if override {
			d.SetProvider(stage, variable, provider)
			continue
		}
		_, configured := configuredProviders[variable]
		if configured {
			continue
		}
		implementations := availableProviders[c]
		if len(implementations) == 1 {
			d.SetProvider(stage, variable, implementations[0])
		} else {
			if variable != c {
				return fmt.Errorf("Missing provider '%s' of type '%s'. This can be configured using the -p / --extra-provider flag.", variable, c)
			}
			return fmt.Errorf("Missing provider of type '%s'. This can be configured using the -p / --extra-provider flag.", c)
		}
	}
	return nil
}

func (d *DeploymentState) GetPreStepInputs(stage string) map[string]interface{} {
	result := map[string]interface{}{}
	for key, val := range d.environment.Inputs {
		result[key] = val
	}
	d.walkStatesAndStages(stage, func(p *DeploymentState, st *StageState) {
		if p.Inputs != nil {
			for key, val := range p.Inputs {
				result[key] = val
			}
		}
		if st.UserInputs != nil {
			for key, val := range st.UserInputs {
				result[key] = val
			}
		}
	})
	return result
}

func (d *DeploymentState) walkStatesAndStages(startStage string, cb func(*DeploymentState, *StageState)) {
	deps := d.getDependencyStates()
	stages := d.getDependencyStages(startStage)
	for i := len(deps) - 1; i >= 0; i-- {
		p := deps[i]
		stage := stages[i]
		cb(p, stage)
	}
}

func (d *DeploymentState) getDependencyStates() []*DeploymentState {
	deps := []*DeploymentState{}
	p := d
	for p != nil {
		deps = append(deps, p)
		p = p.parent
	}
	return deps
}

func (d *DeploymentState) getDependencyStages(startStage string) []*StageState {
	stages := []*StageState{d.GetStageOrCreateNew(startStage)}
	p := d
	for p != nil {
		stages = append(stages, p.parentStage)
		p = p.parent
	}
	return stages
}

func (d *DeploymentState) ValidateNames() error {
	if !validate.IsValidDeploymentName(d.Name) {
		return validate.InvalidDeploymentNameError(d.Name)
	}
	for _, st := range d.Stages {
		if err := st.ValidateNames(); err != nil {
			return err
		}
	}
	return nil
}

func (d *DeploymentState) validateAndFix(name string, env *EnvironmentState) error {
	if !validate.IsValidDeploymentName(name) {
		return validate.InvalidDeploymentNameError(name)
	}
	d.Name = name
	d.environment = env
	if d.Release == "" {
		d.Release = name
	}
	if d.Inputs == nil {
		d.Inputs = map[string]interface{}{}
	}
	if d.Stages == nil {
		d.Stages = map[string]*StageState{}
	}
	for name, st := range d.Stages {
		st.validateAndFix(name, env, d)
	}
	return nil
}

func (d *DeploymentState) validateAndFixSubDeployment(stage *StageState, env *EnvironmentState, parent *DeploymentState) error {
	d.parent = parent
	d.parentStage = stage
	return d.validateAndFix(d.Name, env)
}

func (d *DeploymentState) GetStageOrCreateNew(stage string) *StageState {
	st, ok := d.Stages[stage]
	if !ok || st == nil {
		st = newStage()
		d.Stages[stage] = st
	}
	st.validateAndFix(stage, d.environment, d)
	return st
}
<commit_msg>Set stage name before validating it.<commit_after>/*
Copyright 2017, 2018 Ankyra

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 state

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

	"github.com/ankyra/escape-core"
	"github.com/ankyra/escape-core/state/validate"
)

type DeploymentState struct {
	Name        string                 `json:"name"`
	Release     string                 `json:"release,omitempty"`
	Stages      map[string]*StageState `json:"stages,omitempty"`
	Inputs      map[string]interface{} `json:"inputs,omitempty"`
	environment *EnvironmentState      `json:"-"`
	parent      *DeploymentState       `json:"-"`
	parentStage *StageState            `json:"-"`
}

func NewDeploymentState(env *EnvironmentState, name, release string) (*DeploymentState, error) {
	d := &DeploymentState{
		Name:        name,
		Release:     release,
		Stages:      map[string]*StageState{},
		Inputs:      map[string]interface{}{},
		environment: env,
	}
	return d, d.validateAndFix(name, env)
}

func (d *DeploymentState) GetName() string {
	return d.Name
}

func (d *DeploymentState) Summarize() *DeploymentState {
	result, _ := NewDeploymentState(d.environment, d.Name, d.Release)
	for name, stage := range d.Stages {
		result.Stages[name] = stage.Summarize()
	}
	return result
}

func (d *DeploymentState) GetRootDeploymentName() string {
	prev := d
	p := prev
	for p != nil {
		prev = p
		p = p.parent
	}
	return prev.Name
}

func (d *DeploymentState) GetRootDeploymentStage() string {
	stage := ""
	prev := d
	p := prev
	for prev.parentStage != nil {
		stage = prev.parentStage.Name
		prev = p
		p = p.parent
	}
	return stage
}

func (d *DeploymentState) GetDeploymentPath() string {
	result := []string{}
	p := d
	for p != nil {
		result = append(result, p.Name)
		p = p.parent
	}
	for i := len(result)/2 - 1; i >= 0; i-- {
		opp := len(result) - 1 - i
		result[i], result[opp] = result[opp], result[i]
	}
	return strings.Join(result, ":")
}

func (d *DeploymentState) GetReleaseId(stage string) string {
	return d.Release + "-v" + d.GetVersion(stage)
}

func (d *DeploymentState) GetVersion(stage string) string {
	return d.GetStageOrCreateNew(stage).Version
}

func (d *DeploymentState) GetEnvironmentState() *EnvironmentState {
	return d.environment
}

func (d *DeploymentState) GetDeployment(stage, deploymentName string) (*DeploymentState, error) {
	st := d.GetStageOrCreateNew(stage)
	depl, ok := st.Deployments[deploymentName]
	if !ok {
		return nil, DeploymentDoesNotExistError(deploymentName)
	}
	depl.parentStage = st
	st.Deployments[deploymentName] = depl
	return depl, nil
}

func (d *DeploymentState) GetDeploymentOrMakeNew(stage, deploymentName string) (*DeploymentState, error) {
	st := d.GetStageOrCreateNew(stage)
	depl, ok := st.Deployments[deploymentName]
	if !ok {
		newDepl, err := NewDeploymentState(d.environment, deploymentName, deploymentName)
		if err != nil {
			return nil, err
		}
		depl = newDepl
		depl.parent = d
	}
	depl.parentStage = st
	st.Deployments[deploymentName] = depl
	return depl, nil
}

func (d *DeploymentState) GetUserInputs(stage string) map[string]interface{} {
	return d.GetStageOrCreateNew(stage).UserInputs
}

func (d *DeploymentState) GetCalculatedInputs(stage string) map[string]interface{} {
	return d.GetStageOrCreateNew(stage).Inputs
}

func (d *DeploymentState) GetCalculatedOutputs(stage string) map[string]interface{} {
	return d.GetStageOrCreateNew(stage).Outputs
}

func (d *DeploymentState) UpdateInputs(stage string, inputs map[string]interface{}) error {
	d.GetStageOrCreateNew(stage).SetInputs(inputs)
	return d.Save()
}

func (d *DeploymentState) UpdateUserInputs(stage string, inputs map[string]interface{}) error {
	d.GetStageOrCreateNew(stage).SetUserInputs(inputs)
	return d.Save()
}

func (d *DeploymentState) UpdateOutputs(stage string, outputs map[string]interface{}) error {
	d.GetStageOrCreateNew(stage).SetOutputs(outputs)
	return d.Save()
}

func (d *DeploymentState) CommitVersion(stage string, metadata *core.ReleaseMetadata) error {
	d.GetStageOrCreateNew(stage).SetVersion(metadata.Version)
	d.GetStageOrCreateNew(stage).Provides = metadata.GetProvides()
	return nil
}

func (d *DeploymentState) UpdateStatus(stage string, status *Status) error {
	d.GetStageOrCreateNew(stage).Status = status
	return d.Save()
}
func (d *DeploymentState) GetStatus(stage string) *Status {
	return d.GetStageOrCreateNew(stage).Status
}

func (d *DeploymentState) IsDeployed(stage string, metadata *core.ReleaseMetadata) bool {
	return d.GetStageOrCreateNew(stage).Version == metadata.Version
}

func (d *DeploymentState) Save() error {
	return d.environment.Save(d)
}

func (p *DeploymentState) ToJson() string {
	str, err := json.MarshalIndent(p, "", "   ")
	if err != nil {
		panic(err)
	}
	return string(str)
}

func (d *DeploymentState) SetProvider(stage, name, deplName string) {
	d.GetStageOrCreateNew(stage).Providers[name] = deplName
}

func (d *DeploymentState) GetProviders(stage string) map[string]string {
	result := map[string]string{}
	d.walkStatesAndStages(stage, func(p *DeploymentState, st *StageState) {
		for key, val := range st.Providers {
			result[key] = val
		}
	})
	return result
}

func (d *DeploymentState) ConfigureProviders(metadata *core.ReleaseMetadata, stage string, extraProviders map[string]string) error {
	configuredProviders := d.GetProviders(stage)
	availableProviders := d.environment.GetProviders()
	for _, consumerCfg := range metadata.GetConsumerConfig(stage) {
		c := consumerCfg.Name
		variable := consumerCfg.VariableName
		provider, override := extraProviders[variable]
		if override {
			d.SetProvider(stage, variable, provider)
			continue
		}
		_, configured := configuredProviders[variable]
		if configured {
			continue
		}
		implementations := availableProviders[c]
		if len(implementations) == 1 {
			d.SetProvider(stage, variable, implementations[0])
		} else {
			if variable != c {
				return fmt.Errorf("Missing provider '%s' of type '%s'. This can be configured using the -p / --extra-provider flag.", variable, c)
			}
			return fmt.Errorf("Missing provider of type '%s'. This can be configured using the -p / --extra-provider flag.", c)
		}
	}
	return nil
}

func (d *DeploymentState) GetPreStepInputs(stage string) map[string]interface{} {
	result := map[string]interface{}{}
	for key, val := range d.environment.Inputs {
		result[key] = val
	}
	d.walkStatesAndStages(stage, func(p *DeploymentState, st *StageState) {
		if p.Inputs != nil {
			for key, val := range p.Inputs {
				result[key] = val
			}
		}
		if st.UserInputs != nil {
			for key, val := range st.UserInputs {
				result[key] = val
			}
		}
	})
	return result
}

func (d *DeploymentState) walkStatesAndStages(startStage string, cb func(*DeploymentState, *StageState)) {
	deps := d.getDependencyStates()
	stages := d.getDependencyStages(startStage)
	for i := len(deps) - 1; i >= 0; i-- {
		p := deps[i]
		stage := stages[i]
		cb(p, stage)
	}
}

func (d *DeploymentState) getDependencyStates() []*DeploymentState {
	deps := []*DeploymentState{}
	p := d
	for p != nil {
		deps = append(deps, p)
		p = p.parent
	}
	return deps
}

func (d *DeploymentState) getDependencyStages(startStage string) []*StageState {
	stages := []*StageState{d.GetStageOrCreateNew(startStage)}
	p := d
	for p != nil {
		stages = append(stages, p.parentStage)
		p = p.parent
	}
	return stages
}

func (d *DeploymentState) ValidateNames() error {
	if !validate.IsValidDeploymentName(d.Name) {
		return validate.InvalidDeploymentNameError(d.Name)
	}
	for name, st := range d.Stages {
		st.Name = name
		if err := st.ValidateNames(); err != nil {
			return err
		}
	}
	return nil
}

func (d *DeploymentState) validateAndFix(name string, env *EnvironmentState) error {
	if !validate.IsValidDeploymentName(name) {
		return validate.InvalidDeploymentNameError(name)
	}
	d.Name = name
	d.environment = env
	if d.Release == "" {
		d.Release = name
	}
	if d.Inputs == nil {
		d.Inputs = map[string]interface{}{}
	}
	if d.Stages == nil {
		d.Stages = map[string]*StageState{}
	}
	for name, st := range d.Stages {
		st.validateAndFix(name, env, d)
	}
	return nil
}

func (d *DeploymentState) validateAndFixSubDeployment(stage *StageState, env *EnvironmentState, parent *DeploymentState) error {
	d.parent = parent
	d.parentStage = stage
	return d.validateAndFix(d.Name, env)
}

func (d *DeploymentState) GetStageOrCreateNew(stage string) *StageState {
	st, ok := d.Stages[stage]
	if !ok || st == nil {
		st = newStage()
		d.Stages[stage] = st
	}
	st.validateAndFix(stage, d.environment, d)
	return st
}
<|endoftext|>


<commit_before>package kloud

import (
	"testing"

	"koding/kites/kloud/protocol"
	"koding/kites/kloud/provider/amazon"
	"koding/kites/kloud/provider/digitalocean"
	"koding/kites/kloud/provider/openstack"
)

func TestProviders(t *testing.T) {
	var _ protocol.Builder = &digitalocean.Provider{}
	var _ protocol.Builder = &amazon.Provider{}
	var _ protocol.Builder = &openstack.Provider{}
}
<commit_msg>kloud: remove unused file<commit_after><|endoftext|>


<commit_before>package models

import (
	"errors"
	"fmt"
	"socialapi/db"
	"time"
)

type ChannelParticipant struct {
	// unique identifier of the channel
	Id int64

	// Id of the channel
	ChannelId int64

	// Id of the account
	AccountId int64

	// Status of the participant in the channel
	Status int

	// date of the user's last access to regarding channel
	LastSeenAt time.Time

	// Creation date of the channel channel participant
	CreatedAt time.Time

	// Modification date of the channel participant's status
	ModifiedAt time.Time
}

const (
	ACTIVE int = iota
	LEFT
	REQUEST_PENDING
)

func NewChannelParticipant() *ChannelParticipant {
	now := time.Now().UTC()
	return &ChannelParticipant{
		ChannelId:  1,
		AccountId:  1,
		Status:     ACTIVE,
		LastSeenAt: now,
		CreatedAt:  now,
		ModifiedAt: now,
	}
}

func (c *ChannelParticipant) Create() error {
	now := time.Now().UTC()
	// created at shouldnt be updated
	c.LastSeenAt = now
	c.CreatedAt = now
	c.ModifiedAt = now
	return c.Save()
}

func (c *ChannelParticipant) Update() error {
	if c.Id == 0 {
		return errors.New("ChannelParticipant id is not set")
	}

	now := time.Now().UTC()
	c.LastSeenAt = now
	c.ModifiedAt = now

	return c.Save()
}

func (c *ChannelParticipant) Save() error {
	if err := db.DB.Save(c).Error; err != nil {
		return err
	}
	return nil
}

func (c *ChannelParticipant) Delete() error {
	c.Status = LEFT
	return c.Update()
}

func (c *ChannelParticipant) List() ([]ChannelParticipant, error) {
	var participants []ChannelParticipant

	if c.ChannelId == 0 {
		return participants, errors.New("ChannelId is not set")
	}

	// change this usage to a better one
	// we shouldnt use table name directly
	err := db.DB.Order("created_at desc").Table("channel_participant").Where(
		map[string]interface{}{
			"channel_id": c.ChannelId,
			"status":     ACTIVE,
		},
	).Find(&participants).Error

	fmt.Println(participants)
	return participants, err
}
<commit_msg>Social: use before update and save magic functions<commit_after>package models

import (
	"errors"
	"fmt"
	"socialapi/db"
	"time"
)

type ChannelParticipant struct {
	// unique identifier of the channel
	Id int64

	// Id of the channel
	ChannelId int64

	// Id of the account
	AccountId int64

	// Status of the participant in the channel
	Status int

	// date of the user's last access to regarding channel
	LastSeenAt time.Time

	// Creation date of the channel channel participant
	CreatedAt time.Time

	// Modification date of the channel participant's status
	UpdatedAt time.Time
}

const (
	ACTIVE int = iota
	LEFT
	REQUEST_PENDING
)

func NewChannelParticipant() *ChannelParticipant {
	return &ChannelParticipant{}
}

func (c *ChannelParticipant) BeforeSave() {
	c.LastSeenAt = time.Now().UTC()
}

func (c *ChannelParticipant) BeforeUpdate() {
	c.LastSeenAt = time.Now().UTC()
}

func (c *ChannelParticipant) Create() error {
	return c.Save()
}

func (c *ChannelParticipant) Update() error {
	if c.Id == 0 {
		return errors.New("ChannelParticipant id is not set")
	}

	return c.Save()
}

func (c *ChannelParticipant) Save() error {
	if err := Save(c); err != nil {
		return err
	}

	return nil
}

func (c *ChannelParticipant) Fetch() error {
	if c.AccountId == 0 {
		return errors.New("AccountId is not set")
	}

	if c.ChannelId == 0 {
		return errors.New("ChannelId is not set")
	}

	cp := NewChannelParticipant()
	err := db.DB.
		Where("account_id = ? and channel_id = ?", c.AccountId, c.ChannelId).
		Find(&cp).Error

	if err != nil {
		return err
	}

	// override channel participant
	c = cp
	return nil
}

func (c *ChannelParticipant) Delete() error {
	if err := c.Fetch(); err != nil {
		return err
	}
	c.Status = LEFT
	return c.Update()
}

func (c *ChannelParticipant) List() ([]ChannelParticipant, error) {
	var participants []ChannelParticipant

	if c.ChannelId == 0 {
		return participants, errors.New("ChannelId is not set")
	}

	// change this usage to a better one
	// we shouldnt use table name directly
	err := db.DB.Order("created_at desc").Table("channel_participant").Where(
		map[string]interface{}{
			"channel_id": c.ChannelId,
			"status":     ACTIVE,
		},
	).Find(&participants).Error

	fmt.Println(participants)
	return participants, err
}
<|endoftext|>


<commit_before>package models

import (
	"fmt"
	"socialapi/request"
	"time"

	"github.com/koding/bongo"
)

// todo Scope function for this struct
// in order not to fetch passive accounts
type ChannelParticipant struct {
	// unique identifier of the channel
	Id int64 `json:"id,string"`

	// Id of the channel
	ChannelId int64 `json:"channelId,string"       sql:"NOT NULL"`

	// Id of the account
	AccountId int64 `json:"accountId,string"       sql:"NOT NULL"`

	// Status of the participant in the channel
	StatusConstant string `json:"statusConstant"   sql:"NOT NULL;TYPE:VARCHAR(100);"`

	// holds troll, unsafe, etc
	MetaBits MetaBits `json:"metaBits"`

	// date of the user's last access to regarding channel
	LastSeenAt time.Time `json:"lastSeenAt"        sql:"NOT NULL"`

	// Creation date of the channel channel participant
	CreatedAt time.Time `json:"createdAt"          sql:"NOT NULL"`

	// Modification date of the channel participant's status
	UpdatedAt time.Time `json:"updatedAt"          sql:"NOT NULL"`
}

// here is why i did this not-so-good constants
// https://code.google.com/p/go/issues/detail?id=359
const (
	ChannelParticipant_STATUS_ACTIVE              = "active"
	ChannelParticipant_STATUS_LEFT                = "left"
	ChannelParticipant_STATUS_REQUEST_PENDING     = "requestpending"
	ChannelParticipant_Added_To_Channel_Event     = "added_to_channel"
	ChannelParticipant_Removed_From_Channel_Event = "removed_from_channel"
)

func NewChannelParticipant() *ChannelParticipant {
	return &ChannelParticipant{
		StatusConstant: ChannelParticipant_STATUS_ACTIVE,
		LastSeenAt:     time.Now().UTC(),
		CreatedAt:      time.Now().UTC(),
		UpdatedAt:      time.Now().UTC(),
	}
}

// Create creates a participant in the db as active
// multiple call of this function will result
func (c *ChannelParticipant) Create() error {
	err := c.FetchParticipant()

	// if err is nil
	// it means we already have that user in the channel
	if err == nil {
		// if the participant is already in the channel, and active do nothing
		if c.StatusConstant == ChannelParticipant_STATUS_ACTIVE {
			return nil
		}

		c.StatusConstant = ChannelParticipant_STATUS_ACTIVE
		if err := c.Update(); err != nil {
			return err
		}

		if err := bongo.B.PublishEvent(
			ChannelParticipant_Added_To_Channel_Event, c,
		); err != nil {
			// log here
		}

		return nil
	}

	if err != bongo.RecordNotFound {
		return err
	}

	return bongo.B.Create(c)
}

func (c *ChannelParticipant) CreateRaw() error {
	insertSql := "INSERT INTO " +
		c.TableName() +
		` ("channel_id","account_id", "status_constant", "last_seen_at","created_at", "updated_at") ` +
		"VALUES ($1,$2,$3,$4,$5,$6) " +
		"RETURNING ID"

	return bongo.B.DB.CommonDB().
		QueryRow(insertSql, c.ChannelId, c.AccountId, c.StatusConstant, c.LastSeenAt, c.CreatedAt, c.UpdatedAt).
		Scan(&c.Id)
}

// Tests are done.
func (c *ChannelParticipant) FetchParticipant() error {
	if c.ChannelId == 0 {
		return ErrChannelIdIsNotSet
	}

	if c.AccountId == 0 {
		return ErrAccountIdIsNotSet
	}

	selector := map[string]interface{}{
		"channel_id": c.ChannelId,
		"account_id": c.AccountId,
	}

	return c.fetchParticipant(selector)
}

// Tests are done.
func (c *ChannelParticipant) FetchActiveParticipant() error {
	selector := map[string]interface{}{
		"channel_id":      c.ChannelId,
		"account_id":      c.AccountId,
		"status_constant": ChannelParticipant_STATUS_ACTIVE,
	}

	return c.fetchParticipant(selector)
}

func (c *ChannelParticipant) fetchParticipant(selector map[string]interface{}) error {
	if c.ChannelId == 0 {
		return ErrChannelIdIsNotSet
	}

	if c.AccountId == 0 {
		return ErrAccountIdIsNotSet
	}

	// TODO do we need to add isExempt scope here?
	err := c.One(bongo.NewQS(selector))
	if err != nil {
		return err
	}

	return nil
}

// Tests are done in channelmessagelist.
func (c *ChannelParticipant) FetchUnreadCount() (int, error) {
	cml := NewChannelMessageList()
	return cml.UnreadCount(c)
}

func (c *ChannelParticipant) Delete() error {
	if err := c.FetchParticipant(); err != nil {
		return err
	}

	c.StatusConstant = ChannelParticipant_STATUS_LEFT
	if err := c.Update(); err != nil {
		return err
	}

	if err := bongo.B.PublishEvent(
		ChannelParticipant_Removed_From_Channel_Event, c,
	); err != nil {
		// log here
	}

	return nil

}

func (c *ChannelParticipant) List(q *request.Query) ([]ChannelParticipant, error) {
	var participants []ChannelParticipant

	if c.ChannelId == 0 {
		return participants, ErrChannelIdIsNotSet
	}

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id":      c.ChannelId,
			"status_constant": ChannelParticipant_STATUS_ACTIVE,
		},
	}

	// add filter for troll content
	query.AddScope(RemoveTrollContent(c, q.ShowExempt))

	err := bongo.B.Some(c, &participants, query)
	if err != nil {
		return nil, err
	}

	return participants, nil
}

func (c *ChannelParticipant) ListAccountIds(limit int) ([]int64, error) {
	var participants []int64

	if c.ChannelId == 0 {
		return participants, ErrChannelIdIsNotSet
	}

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id":      c.ChannelId,
			"status_constant": ChannelParticipant_STATUS_ACTIVE,
		},
		Pluck: "account_id",
	}

	if limit != 0 {
		query.Pagination = *bongo.NewPagination(limit, 0)
	}

	// do not include troll content
	query.AddScope(RemoveTrollContent(c, false))

	err := bongo.B.Some(c, &participants, query)
	if err != nil {
		return nil, err
	}

	return participants, nil
}

func (c *ChannelParticipant) FetchParticipatedChannelIds(a *Account, q *request.Query) ([]int64, error) {
	if a.Id == 0 {
		return nil, ErrAccountIdIsNotSet
	}

	channelIds := make([]int64, 0)

	// var results []ChannelParticipant
	query := bongo.B.DB.
		Model(c).
		Table(c.TableName()).
		Select("api.channel_participant.channel_id").
		Joins(
		`left join api.channel on
		 api.channel_participant.channel_id = api.channel.id`).
		Where(
		`api.channel_participant.account_id = ? and
		 api.channel.group_name = ? and
		 api.channel.type_constant = ? and
		 api.channel_participant.status_constant = ?`,
		a.Id,
		q.GroupName,
		q.Type,
		ChannelParticipant_STATUS_ACTIVE,
	)

	// add exempt clause if needed
	if !q.ShowExempt {
		query = query.Where("api.channel.meta_bits = ?", Safe)
	}

	rows, err := query.
		Limit(q.Limit).
		Offset(q.Skip).
		Rows()

	defer rows.Close()
	if err != nil {
		return channelIds, err
	}

	if rows == nil {
		return nil, nil
	}

	var channelId int64
	for rows.Next() {
		rows.Scan(&channelId)
		channelIds = append(channelIds, channelId)
	}

	// if this is the first query for listing the channels
	// add default channels into the result set
	if q.Skip == 0 {
		defaultChannels, err := c.fetchDefaultChannels(q)
		if err != nil {
			fmt.Println(err.Error())
		} else {
			for _, item := range channelIds {
				defaultChannels = append(defaultChannels, item)
			}
			return defaultChannels, nil
		}
	}

	return channelIds, nil
}

func (c *ChannelParticipant) fetchDefaultChannels(q *request.Query) ([]int64, error) {
	var channelIds []int64
	channel := NewChannel()
	bongoQuery := &bongo.Query{
		Selector: map[string]interface{}{
			"group_name":    q.GroupName,
			"type_constant": Channel_TYPE_GROUP,
		},
		Pluck:      "id",
		Pagination: *bongo.NewPagination(1, 0),
	}

	err := channel.Some(&channelIds, bongoQuery)
	if err != nil {
		return nil, err
	}

	return channelIds, nil
}

// FetchParticipantCount fetchs the participant count in the channel
// if there is no participant in the channel, then returns zero value
//
// Tests are done.
func (c *ChannelParticipant) FetchParticipantCount() (int, error) {
	if c.ChannelId == 0 {
		return 0, ErrChannelIdIsNotSet
	}

	return c.Count("channel_id = ? and status_constant = ?", c.ChannelId, ChannelParticipant_STATUS_ACTIVE)
}

// Tests are done.
func (c *ChannelParticipant) IsParticipant(accountId int64) (bool, error) {
	if c.ChannelId == 0 {
		return false, ErrChannelIdIsNotSet
	}

	selector := map[string]interface{}{
		"channel_id":      c.ChannelId,
		"account_id":      accountId,
		"status_constant": ChannelParticipant_STATUS_ACTIVE,
	}

	err := c.One(bongo.NewQS(selector))
	if err == nil {
		return true, nil
	}

	if err == bongo.RecordNotFound {
		return false, nil
	}

	return false, err
}

// Put them all behind an interface
// channels, messages, lists, participants, etc
//
// Tests are done.
func (c *ChannelParticipant) MarkIfExempt() error {
	isExempt, err := c.isExempt()
	if err != nil {
		return err
	}

	if isExempt {
		c.MetaBits.Mark(Troll)
	}

	return nil
}

// Tests are done.
func (c *ChannelParticipant) isExempt() (bool, error) {
	// return early if channel is already exempt
	if c.MetaBits.Is(Troll) {
		return true, nil
	}

	accountId, err := c.getAccountId()
	if err != nil {
		return false, err
	}

	account, err := ResetAccountCache(accountId)
	if err != nil {
		return false, err
	}

	if account == nil {
		return false, fmt.Errorf("account is nil, accountId:%d", accountId)
	}

	if account.IsTroll {
		return true, nil
	}

	return false, nil
}

// Tests are done.
func (c *ChannelParticipant) getAccountId() (int64, error) {
	if c.AccountId != 0 {
		return c.AccountId, nil
	}

	if c.Id == 0 {
		return 0, fmt.Errorf("couldnt find accountId from content %+v", c)
	}

	cp := NewChannelParticipant()
	if err := cp.ById(c.Id); err != nil {
		return 0, err
	}

	return cp.AccountId, nil
}

func (c *ChannelParticipant) RawUpdateLastSeenAt(t time.Time) error {
	if c.Id == 0 {
		return ErrIdIsNotSet
	}

	query := fmt.Sprintf("UPDATE %s SET last_seen_at = ? WHERE id = ?", c.TableName())
	return bongo.B.DB.Exec(query, t, c.Id).Error
}
<commit_msg>Social: add announcement channel into default channels and make sure the requester will be a participant of the channel.<commit_after>package models

import (
	"fmt"
	"socialapi/request"
	"time"

	"github.com/koding/bongo"
)

// todo Scope function for this struct
// in order not to fetch passive accounts
type ChannelParticipant struct {
	// unique identifier of the channel
	Id int64 `json:"id,string"`

	// Id of the channel
	ChannelId int64 `json:"channelId,string"       sql:"NOT NULL"`

	// Id of the account
	AccountId int64 `json:"accountId,string"       sql:"NOT NULL"`

	// Status of the participant in the channel
	StatusConstant string `json:"statusConstant"   sql:"NOT NULL;TYPE:VARCHAR(100);"`

	// holds troll, unsafe, etc
	MetaBits MetaBits `json:"metaBits"`

	// date of the user's last access to regarding channel
	LastSeenAt time.Time `json:"lastSeenAt"        sql:"NOT NULL"`

	// Creation date of the channel channel participant
	CreatedAt time.Time `json:"createdAt"          sql:"NOT NULL"`

	// Modification date of the channel participant's status
	UpdatedAt time.Time `json:"updatedAt"          sql:"NOT NULL"`
}

// here is why i did this not-so-good constants
// https://code.google.com/p/go/issues/detail?id=359
const (
	ChannelParticipant_STATUS_ACTIVE              = "active"
	ChannelParticipant_STATUS_LEFT                = "left"
	ChannelParticipant_STATUS_REQUEST_PENDING     = "requestpending"
	ChannelParticipant_Added_To_Channel_Event     = "added_to_channel"
	ChannelParticipant_Removed_From_Channel_Event = "removed_from_channel"
)

func NewChannelParticipant() *ChannelParticipant {
	return &ChannelParticipant{
		StatusConstant: ChannelParticipant_STATUS_ACTIVE,
		LastSeenAt:     time.Now().UTC(),
		CreatedAt:      time.Now().UTC(),
		UpdatedAt:      time.Now().UTC(),
	}
}

// Create creates a participant in the db as active
// multiple call of this function will result
func (c *ChannelParticipant) Create() error {
	err := c.FetchParticipant()

	// if err is nil
	// it means we already have that user in the channel
	if err == nil {
		// if the participant is already in the channel, and active do nothing
		if c.StatusConstant == ChannelParticipant_STATUS_ACTIVE {
			return nil
		}

		c.StatusConstant = ChannelParticipant_STATUS_ACTIVE
		if err := c.Update(); err != nil {
			return err
		}

		if err := bongo.B.PublishEvent(
			ChannelParticipant_Added_To_Channel_Event, c,
		); err != nil {
			// log here
		}

		return nil
	}

	if err != bongo.RecordNotFound {
		return err
	}

	return bongo.B.Create(c)
}

func (c *ChannelParticipant) CreateRaw() error {
	insertSql := "INSERT INTO " +
		c.TableName() +
		` ("channel_id","account_id", "status_constant", "last_seen_at","created_at", "updated_at") ` +
		"VALUES ($1,$2,$3,$4,$5,$6) " +
		"RETURNING ID"

	return bongo.B.DB.CommonDB().
		QueryRow(insertSql, c.ChannelId, c.AccountId, c.StatusConstant, c.LastSeenAt, c.CreatedAt, c.UpdatedAt).
		Scan(&c.Id)
}

// Tests are done.
func (c *ChannelParticipant) FetchParticipant() error {
	if c.ChannelId == 0 {
		return ErrChannelIdIsNotSet
	}

	if c.AccountId == 0 {
		return ErrAccountIdIsNotSet
	}

	selector := map[string]interface{}{
		"channel_id": c.ChannelId,
		"account_id": c.AccountId,
	}

	return c.fetchParticipant(selector)
}

// Tests are done.
func (c *ChannelParticipant) FetchActiveParticipant() error {
	selector := map[string]interface{}{
		"channel_id":      c.ChannelId,
		"account_id":      c.AccountId,
		"status_constant": ChannelParticipant_STATUS_ACTIVE,
	}

	return c.fetchParticipant(selector)
}

func (c *ChannelParticipant) fetchParticipant(selector map[string]interface{}) error {
	if c.ChannelId == 0 {
		return ErrChannelIdIsNotSet
	}

	if c.AccountId == 0 {
		return ErrAccountIdIsNotSet
	}

	// TODO do we need to add isExempt scope here?
	err := c.One(bongo.NewQS(selector))
	if err != nil {
		return err
	}

	return nil
}

// Tests are done in channelmessagelist.
func (c *ChannelParticipant) FetchUnreadCount() (int, error) {
	cml := NewChannelMessageList()
	return cml.UnreadCount(c)
}

func (c *ChannelParticipant) Delete() error {
	if err := c.FetchParticipant(); err != nil {
		return err
	}

	c.StatusConstant = ChannelParticipant_STATUS_LEFT
	if err := c.Update(); err != nil {
		return err
	}

	if err := bongo.B.PublishEvent(
		ChannelParticipant_Removed_From_Channel_Event, c,
	); err != nil {
		// log here
	}

	return nil

}

func (c *ChannelParticipant) List(q *request.Query) ([]ChannelParticipant, error) {
	var participants []ChannelParticipant

	if c.ChannelId == 0 {
		return participants, ErrChannelIdIsNotSet
	}

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id":      c.ChannelId,
			"status_constant": ChannelParticipant_STATUS_ACTIVE,
		},
	}

	// add filter for troll content
	query.AddScope(RemoveTrollContent(c, q.ShowExempt))

	err := bongo.B.Some(c, &participants, query)
	if err != nil {
		return nil, err
	}

	return participants, nil
}

func (c *ChannelParticipant) ListAccountIds(limit int) ([]int64, error) {
	var participants []int64

	if c.ChannelId == 0 {
		return participants, ErrChannelIdIsNotSet
	}

	query := &bongo.Query{
		Selector: map[string]interface{}{
			"channel_id":      c.ChannelId,
			"status_constant": ChannelParticipant_STATUS_ACTIVE,
		},
		Pluck: "account_id",
	}

	if limit != 0 {
		query.Pagination = *bongo.NewPagination(limit, 0)
	}

	// do not include troll content
	query.AddScope(RemoveTrollContent(c, false))

	err := bongo.B.Some(c, &participants, query)
	if err != nil {
		return nil, err
	}

	return participants, nil
}

func (c *ChannelParticipant) FetchParticipatedChannelIds(a *Account, q *request.Query) ([]int64, error) {
	if a.Id == 0 {
		return nil, ErrAccountIdIsNotSet
	}

	channelIds := make([]int64, 0)

	// var results []ChannelParticipant
	query := bongo.B.DB.
		Model(c).
		Table(c.TableName()).
		Select("api.channel_participant.channel_id").
		Joins(
		`left join api.channel on
		 api.channel_participant.channel_id = api.channel.id`).
		Where(
		`api.channel_participant.account_id = ? and
		 api.channel.group_name = ? and
		 api.channel.type_constant = ? and
		 api.channel_participant.status_constant = ?`,
		a.Id,
		q.GroupName,
		q.Type,
		ChannelParticipant_STATUS_ACTIVE,
	)

	// add exempt clause if needed
	if !q.ShowExempt {
		query = query.Where("api.channel.meta_bits = ?", Safe)
	}

	rows, err := query.
		Limit(q.Limit).
		Offset(q.Skip).
		Rows()

	defer rows.Close()
	if err != nil {
		return channelIds, err
	}

	if rows == nil {
		return nil, nil
	}

	var channelId int64
	for rows.Next() {
		rows.Scan(&channelId)
		channelIds = append(channelIds, channelId)
	}

	// if this is the first query for listing the channels
	// add default channels into the result set
	if q.Skip == 0 {
		defaultChannels, err := c.fetchDefaultChannels(q)
		if err != nil {
			fmt.Println(err.Error())
		} else {
			for _, item := range channelIds {
				defaultChannels = append(defaultChannels, item)
			}
			return defaultChannels, nil
		}
	}

	return channelIds, nil
}

// fetchDefaultChannels fetchs the default channels of the system, currently we
// have two different default channels, group channel and announcement channel
// that everyone in the system should be a member of them, they cannot opt-out,
// they will be able to see the contents of it, they will get the notifications,
// they will see the unread count
func (c *ChannelParticipant) fetchDefaultChannels(q *request.Query) ([]int64, error) {
	var channelIds []int64
	channel := NewChannel()
	res := bongo.B.DB.
		Model(channel).
		Table(channel.TableName()).
		Where(
		"group_name = ? AND type_constant IN (?)",
		q.GroupName,
		[]string{Channel_TYPE_GROUP, Channel_TYPE_ANNOUNCEMENT},
	).
		// no need to traverse all database, limit with a known count
		Limit(2).
		// only select ids
		Pluck("id", &channelIds)

	if err := bongo.CheckErr(res); err != nil {
		return nil, err
	}

	// be sure that this account is a participant of default channels
	if err := c.ensureParticipation(q.AccountId, channelIds); err != nil {
		return nil, err
	}

	return channelIds, nil
}

func (c *ChannelParticipant) ensureParticipation(accountId int64, channelIds []int64) error {
	for _, channelId := range channelIds {
		cp := NewChannelParticipant()
		cp.ChannelId = channelId
		cp.AccountId = accountId
		// create is idempotent, multiple calls wont cause any problem, if the
		// user is already a participant, will return as if a succesful request
		if err := cp.Create(); err != nil {
			return err
		}
	}

	return nil
}

// FetchParticipantCount fetchs the participant count in the channel
// if there is no participant in the channel, then returns zero value
//
// Tests are done.
func (c *ChannelParticipant) FetchParticipantCount() (int, error) {
	if c.ChannelId == 0 {
		return 0, ErrChannelIdIsNotSet
	}

	return c.Count("channel_id = ? and status_constant = ?", c.ChannelId, ChannelParticipant_STATUS_ACTIVE)
}

// Tests are done.
func (c *ChannelParticipant) IsParticipant(accountId int64) (bool, error) {
	if c.ChannelId == 0 {
		return false, ErrChannelIdIsNotSet
	}

	selector := map[string]interface{}{
		"channel_id":      c.ChannelId,
		"account_id":      accountId,
		"status_constant": ChannelParticipant_STATUS_ACTIVE,
	}

	err := c.One(bongo.NewQS(selector))
	if err == nil {
		return true, nil
	}

	if err == bongo.RecordNotFound {
		return false, nil
	}

	return false, err
}

// Put them all behind an interface
// channels, messages, lists, participants, etc
//
// Tests are done.
func (c *ChannelParticipant) MarkIfExempt() error {
	isExempt, err := c.isExempt()
	if err != nil {
		return err
	}

	if isExempt {
		c.MetaBits.Mark(Troll)
	}

	return nil
}

// Tests are done.
func (c *ChannelParticipant) isExempt() (bool, error) {
	// return early if channel is already exempt
	if c.MetaBits.Is(Troll) {
		return true, nil
	}

	accountId, err := c.getAccountId()
	if err != nil {
		return false, err
	}

	account, err := ResetAccountCache(accountId)
	if err != nil {
		return false, err
	}

	if account == nil {
		return false, fmt.Errorf("account is nil, accountId:%d", accountId)
	}

	if account.IsTroll {
		return true, nil
	}

	return false, nil
}

// Tests are done.
func (c *ChannelParticipant) getAccountId() (int64, error) {
	if c.AccountId != 0 {
		return c.AccountId, nil
	}

	if c.Id == 0 {
		return 0, fmt.Errorf("couldnt find accountId from content %+v", c)
	}

	cp := NewChannelParticipant()
	if err := cp.ById(c.Id); err != nil {
		return 0, err
	}

	return cp.AccountId, nil
}

func (c *ChannelParticipant) RawUpdateLastSeenAt(t time.Time) error {
	if c.Id == 0 {
		return ErrIdIsNotSet
	}

	query := fmt.Sprintf("UPDATE %s SET last_seen_at = ? WHERE id = ?", c.TableName())
	return bongo.B.DB.Exec(query, t, c.Id).Error
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Vitess 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 vtgate

import (
	"context"
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"

	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/test/endtoend/cluster"
)

// TestCheckConstraint test check constraints on CREATE TABLE
// This feature is supported from MySQL 8.0.16 and MariaDB 10.2.1.
func TestCheckConstraint(t *testing.T) {
	// Skipping as tests are run against MySQL 5.7
	t.Skip()

	conn, err := mysql.Connect(context.Background(), &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	query := `CREATE TABLE t7 (CHECK (c1 <> c2), c1 INT CHECK (c1 > 10), c2 INT CONSTRAINT c2_positive CHECK (c2 > 0), c3 INT CHECK (c3 < 100), CONSTRAINT c1_nonzero CHECK (c1 <> 0), CHECK (c1 > c3));`
	exec(t, conn, query)

	checkQuery := `SELECT CONSTRAINT_NAME FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS WHERE TABLE_NAME = 't7';`
	expected := `[[VARCHAR("t7_chk_1")] [VARCHAR("t7_chk_2")] [VARCHAR("c2_positive")] [VARCHAR("t7_chk_3")] [VARCHAR("c1_nonzero")] [VARCHAR("t7_chk_4")]]`

	assertMatches(t, conn, checkQuery, expected)

	cleanup := `DROP TABLE t7`
	exec(t, conn, cleanup)
}

func TestDbNameOverride(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.Nil(t, err)
	defer conn.Close()
	qr, err := conn.ExecuteFetch("SELECT distinct database() FROM information_schema.tables WHERE table_schema = database()", 1000, true)

	require.Nil(t, err)
	assert.Equal(t, 1, len(qr.Rows), "did not get enough rows back")
	assert.Equal(t, "vt_ks", qr.Rows[0][0].ToString())
}

func TestInformationSchemaQuery(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	assertSingleRowIsReturned(t, conn, "table_schema = 'ks'", "vt_ks")
	assertSingleRowIsReturned(t, conn, "table_schema = 'vt_ks'", "vt_ks")
	assertResultIsEmpty(t, conn, "table_schema = 'NONE'")
	assertSingleRowIsReturned(t, conn, "table_schema = 'performance_schema'", "performance_schema")
	assertResultIsEmpty(t, conn, "table_schema = 'PERFORMANCE_SCHEMA'")
}

func assertResultIsEmpty(t *testing.T, conn *mysql.Conn, pre string) {
	t.Run(pre, func(t *testing.T) {
		qr, err := conn.ExecuteFetch("SELECT distinct table_schema FROM information_schema.tables WHERE "+pre, 1000, true)
		require.NoError(t, err)
		assert.Empty(t, qr.Rows)
	})
}

func assertSingleRowIsReturned(t *testing.T, conn *mysql.Conn, predicate string, expectedKs string) {
	t.Run(predicate, func(t *testing.T) {
		qr, err := conn.ExecuteFetch("SELECT distinct table_schema FROM information_schema.tables WHERE "+predicate, 1000, true)
		require.NoError(t, err)
		assert.Equal(t, 1, len(qr.Rows), "did not get enough rows back")
		assert.Equal(t, expectedKs, qr.Rows[0][0].ToString())
	})
}

func TestInformationSchemaWithSubquery(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	result := exec(t, conn, "SELECT column_name FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = (SELECT SCHEMA()) AND TABLE_NAME = 'not_exists'")
	assert.Empty(t, result.Rows)
}

func TestInformationSchemaQueryGetsRoutedToTheRightTableAndKeyspace(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	_ = exec(t, conn, "SELECT * FROM t1000") // test that the routed table is available to us
	result := exec(t, conn, "SELECT * FROM information_schema.tables WHERE table_schema = database() and table_name='t1000'")
	assert.NotEmpty(t, result.Rows)
}

func TestFKConstraintUsingInformationSchema(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	query := "select fk.referenced_table_name as to_table, fk.referenced_column_name as primary_key, fk.column_name as `column`, fk.constraint_name as name, rc.update_rule as on_update, rc.delete_rule as on_delete from information_schema.referential_constraints as rc join information_schema.key_column_usage as fk using (constraint_schema, constraint_name) where fk.referenced_column_name is not null and fk.table_schema = database() and fk.table_name = 't7_fk' and rc.constraint_schema = database() and rc.table_name = 't7_fk'"
	assertMatches(t, conn, query, `[[VARCHAR("t7_xxhash") VARCHAR("uid") VARCHAR("t7_uid") VARCHAR("t7_fk_ibfk_1") VARCHAR("CASCADE") VARCHAR("SET NULL")]]`)
}

func TestConnectWithSystemSchema(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	for _, dbname := range []string{"information_schema", "mysql", "performance_schema", "sys"} {
		connParams := vtParams
		connParams.DbName = dbname
		conn, err := mysql.Connect(ctx, &connParams)
		require.NoError(t, err)
		conn.Close()
	}
}

func TestUseSystemSchema(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	for _, dbname := range []string{"information_schema", "mysql", "performance_schema", "sys"} {
		conn, err := mysql.Connect(ctx, &vtParams)
		require.NoError(t, err)

		exec(t, conn, fmt.Sprintf("use %s", dbname))
		conn.Close()
	}
}

func TestSystemSchemaQueryWithoutQualifier(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	queryWithQualifier := fmt.Sprintf("select t.table_schema,t.table_name,c.column_name,c.column_type "+
		"from information_schema.tables t "+
		"join information_schema.columns c "+
		"on c.table_schema = t.table_schema and c.table_name = t.table_name "+
		"where t.table_schema = '%s' and c.table_schema = '%s'", KeyspaceName, KeyspaceName)
	qr1 := exec(t, conn, queryWithQualifier)

	queryWithoutQualifier := fmt.Sprintf("select t.table_schema,t.table_name,c.column_name,c.column_type "+
		"from tables t "+
		"join columns c "+
		"on c.table_schema = t.table_schema and c.table_name = t.table_name "+
		"where t.table_schema = '%s' and c.table_schema = '%s'", KeyspaceName, KeyspaceName)
	exec(t, conn, "use information_schema")
	qr2 := exec(t, conn, queryWithoutQualifier)
	require.Equal(t, qr1, qr2)

	connParams := vtParams
	connParams.DbName = "information_schema"
	conn2, err := mysql.Connect(ctx, &connParams)
	require.NoError(t, err)
	defer conn2.Close()

	qr3 := exec(t, conn2, queryWithoutQualifier)
	require.Equal(t, qr2, qr3)
}
<commit_msg>added e2e test having table_schema and table_name in where clause<commit_after>/*
Copyright 2020 The Vitess 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 vtgate

import (
	"context"
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"

	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/test/endtoend/cluster"
)

// TestCheckConstraint test check constraints on CREATE TABLE
// This feature is supported from MySQL 8.0.16 and MariaDB 10.2.1.
func TestCheckConstraint(t *testing.T) {
	// Skipping as tests are run against MySQL 5.7
	t.Skip()

	conn, err := mysql.Connect(context.Background(), &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	query := `CREATE TABLE t7 (CHECK (c1 <> c2), c1 INT CHECK (c1 > 10), c2 INT CONSTRAINT c2_positive CHECK (c2 > 0), c3 INT CHECK (c3 < 100), CONSTRAINT c1_nonzero CHECK (c1 <> 0), CHECK (c1 > c3));`
	exec(t, conn, query)

	checkQuery := `SELECT CONSTRAINT_NAME FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS WHERE TABLE_NAME = 't7';`
	expected := `[[VARCHAR("t7_chk_1")] [VARCHAR("t7_chk_2")] [VARCHAR("c2_positive")] [VARCHAR("t7_chk_3")] [VARCHAR("c1_nonzero")] [VARCHAR("t7_chk_4")]]`

	assertMatches(t, conn, checkQuery, expected)

	cleanup := `DROP TABLE t7`
	exec(t, conn, cleanup)
}

func TestDbNameOverride(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.Nil(t, err)
	defer conn.Close()
	qr, err := conn.ExecuteFetch("SELECT distinct database() FROM information_schema.tables WHERE table_schema = database()", 1000, true)

	require.Nil(t, err)
	assert.Equal(t, 1, len(qr.Rows), "did not get enough rows back")
	assert.Equal(t, "vt_ks", qr.Rows[0][0].ToString())
}

func TestInformationSchemaQuery(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	assertSingleRowIsReturned(t, conn, "table_schema = 'ks'", "vt_ks")
	assertSingleRowIsReturned(t, conn, "table_schema = 'vt_ks'", "vt_ks")
	assertResultIsEmpty(t, conn, "table_schema = 'NONE'")
	assertSingleRowIsReturned(t, conn, "table_schema = 'performance_schema'", "performance_schema")
	assertResultIsEmpty(t, conn, "table_schema = 'PERFORMANCE_SCHEMA'")
	assertSingleRowIsReturned(t, conn, "table_schema = 'performance_schema' and table_name = 'users'", "performance_schema")
	assertResultIsEmpty(t, conn, "table_schema = 'performance_schema' and table_name = 'foo'")
}

func assertResultIsEmpty(t *testing.T, conn *mysql.Conn, pre string) {
	t.Run(pre, func(t *testing.T) {
		qr, err := conn.ExecuteFetch("SELECT distinct table_schema FROM information_schema.tables WHERE "+pre, 1000, true)
		require.NoError(t, err)
		assert.Empty(t, qr.Rows)
	})
}

func assertSingleRowIsReturned(t *testing.T, conn *mysql.Conn, predicate string, expectedKs string) {
	t.Run(predicate, func(t *testing.T) {
		qr, err := conn.ExecuteFetch("SELECT distinct table_schema FROM information_schema.tables WHERE "+predicate, 1000, true)
		require.NoError(t, err)
		assert.Equal(t, 1, len(qr.Rows), "did not get enough rows back")
		assert.Equal(t, expectedKs, qr.Rows[0][0].ToString())
	})
}

func TestInformationSchemaWithSubquery(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	result := exec(t, conn, "SELECT column_name FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = (SELECT SCHEMA()) AND TABLE_NAME = 'not_exists'")
	assert.Empty(t, result.Rows)
}

func TestInformationSchemaQueryGetsRoutedToTheRightTableAndKeyspace(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	_ = exec(t, conn, "SELECT * FROM t1000") // test that the routed table is available to us
	result := exec(t, conn, "SELECT * FROM information_schema.tables WHERE table_schema = database() and table_name='t1000'")
	assert.NotEmpty(t, result.Rows)
}

func TestFKConstraintUsingInformationSchema(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	query := "select fk.referenced_table_name as to_table, fk.referenced_column_name as primary_key, fk.column_name as `column`, fk.constraint_name as name, rc.update_rule as on_update, rc.delete_rule as on_delete from information_schema.referential_constraints as rc join information_schema.key_column_usage as fk using (constraint_schema, constraint_name) where fk.referenced_column_name is not null and fk.table_schema = database() and fk.table_name = 't7_fk' and rc.constraint_schema = database() and rc.table_name = 't7_fk'"
	assertMatches(t, conn, query, `[[VARCHAR("t7_xxhash") VARCHAR("uid") VARCHAR("t7_uid") VARCHAR("t7_fk_ibfk_1") VARCHAR("CASCADE") VARCHAR("SET NULL")]]`)
}

func TestConnectWithSystemSchema(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	for _, dbname := range []string{"information_schema", "mysql", "performance_schema", "sys"} {
		connParams := vtParams
		connParams.DbName = dbname
		conn, err := mysql.Connect(ctx, &connParams)
		require.NoError(t, err)
		conn.Close()
	}
}

func TestUseSystemSchema(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	for _, dbname := range []string{"information_schema", "mysql", "performance_schema", "sys"} {
		conn, err := mysql.Connect(ctx, &vtParams)
		require.NoError(t, err)

		exec(t, conn, fmt.Sprintf("use %s", dbname))
		conn.Close()
	}
}

func TestSystemSchemaQueryWithoutQualifier(t *testing.T) {
	defer cluster.PanicHandler(t)
	ctx := context.Background()
	conn, err := mysql.Connect(ctx, &vtParams)
	require.NoError(t, err)
	defer conn.Close()

	queryWithQualifier := fmt.Sprintf("select t.table_schema,t.table_name,c.column_name,c.column_type "+
		"from information_schema.tables t "+
		"join information_schema.columns c "+
		"on c.table_schema = t.table_schema and c.table_name = t.table_name "+
		"where t.table_schema = '%s' and c.table_schema = '%s'", KeyspaceName, KeyspaceName)
	qr1 := exec(t, conn, queryWithQualifier)

	queryWithoutQualifier := fmt.Sprintf("select t.table_schema,t.table_name,c.column_name,c.column_type "+
		"from tables t "+
		"join columns c "+
		"on c.table_schema = t.table_schema and c.table_name = t.table_name "+
		"where t.table_schema = '%s' and c.table_schema = '%s'", KeyspaceName, KeyspaceName)
	exec(t, conn, "use information_schema")
	qr2 := exec(t, conn, queryWithoutQualifier)
	require.Equal(t, qr1, qr2)

	connParams := vtParams
	connParams.DbName = "information_schema"
	conn2, err := mysql.Connect(ctx, &connParams)
	require.NoError(t, err)
	defer conn2.Close()

	qr3 := exec(t, conn2, queryWithoutQualifier)
	require.Equal(t, qr2, qr3)
}
<|endoftext|>