<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 blob_test

import (
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/blob/mock"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestChecking(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type checkingStoreTest struct {
	wrapped mock_blob.MockStore
	store   blob.Store
}

func (t *checkingStoreTest) SetUp(i *TestInfo) {
	t.wrapped = mock_blob.NewMockStore(i.MockController, "wrapped")
	t.store = blob.NewCheckingStore(t.wrapped)
}

////////////////////////////////////////////////////////////////////////
// Store
////////////////////////////////////////////////////////////////////////

type CheckingStore_StoreTest struct {
	checkingStoreTest
}

func init() { RegisterTestSuite(&CheckingStore_StoreTest{}) }

func (t *CheckingStore_StoreTest) DoesFoo() {
	ExpectEq("TODO", "")
}

////////////////////////////////////////////////////////////////////////
// Load
////////////////////////////////////////////////////////////////////////

type CheckingStore_LoadTest struct {
	checkingStoreTest
}

func init() { RegisterTestSuite(&CheckingStore_LoadTest{}) }

func (t *CheckingStore_LoadTest) DoesFoo() {
	ExpectEq("TODO", "")
}
<commit_msg>Added test names.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 blob_test

import (
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/blob/mock"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestChecking(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type checkingStoreTest struct {
	wrapped mock_blob.MockStore
	store   blob.Store
}

func (t *checkingStoreTest) SetUp(i *TestInfo) {
	t.wrapped = mock_blob.NewMockStore(i.MockController, "wrapped")
	t.store = blob.NewCheckingStore(t.wrapped)
}

////////////////////////////////////////////////////////////////////////
// Store
////////////////////////////////////////////////////////////////////////

type CheckingStore_StoreTest struct {
	checkingStoreTest
}

func init() { RegisterTestSuite(&CheckingStore_StoreTest{}) }

func (t *CheckingStore_StoreTest) CallsWrapped() {
	ExpectEq("TODO", "")
}

func (t *CheckingStore_StoreTest) WrappedReturnsError() {
	ExpectEq("TODO", "")
}

func (t *CheckingStore_StoreTest) WrappedReturnsIncorrectScore() {
	ExpectEq("TODO", "")
}

func (t *CheckingStore_StoreTest) WrappedReturnsCorrectScore() {
	ExpectEq("TODO", "")
}

////////////////////////////////////////////////////////////////////////
// Load
////////////////////////////////////////////////////////////////////////

type CheckingStore_LoadTest struct {
	checkingStoreTest
}

func init() { RegisterTestSuite(&CheckingStore_LoadTest{}) }

func (t *CheckingStore_LoadTest) CallsWrapped() {
	ExpectEq("TODO", "")
}

func (t *CheckingStore_LoadTest) WrappedReturnsError() {
	ExpectEq("TODO", "")
}

func (t *CheckingStore_LoadTest) WrappedReturnsCorrectData() {
	ExpectEq("TODO", "")
}

func (t *CheckingStore_LoadTest) WrappedReturnsIncorrectData() {
	ExpectEq("TODO", "")
}
<|endoftext|>


<commit_before>// BaruwaAPI Golang bindings for Baruwa REST API
// Copyright (C) 2019 Andrew Colin Kissa <andrew@topdog.za.net>

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

// Package api Golang bindings for Baruwa REST API
package api

import (
	"fmt"
	"net/url"

	"github.com/google/go-querystring/query"
)

// UserDomain holds user
type UserDomain struct {
	ID   int    `json:"id"`
	Name string `json:"name"`
}

// User holds users
type User struct {
	Username    string       `json:"username"`
	Firstname   string       `json:"firstname"`
	Lastname    string       `json:"lastname"`
	Password1   string       `json:"password1,omitempty"`
	Password2   string       `json:"password2,omitempty"`
	Email       string       `json:"email"`
	Timezone    string       `json:"timezone"`
	AccountType int          `json:"account_type"`
	Domains     []UserDomain `json:"domains,omitempty"`
	Active      bool         `json:"active"`
	SendReport  bool         `json:"send_report"`
	SpamChecks  bool         `json:"spam_checks"`
	LowScore    float64      `json:"low_score"`
	HighScore   float64      `json:"high_score"`
	BlockMacros bool         `json:"block_macros"`
}

// GetUser returns a user account
func (c *Client) GetUser(id int) (user *User, err error) {
	if id <= 0 {
		err = fmt.Errorf("The id param should be > 0")
		return
	}

	err = c.get(fmt.Sprintf("users/%d", id), user)
	return
}

// CreateUser creates a user account
func (c *Client) CreateUser(user *User) (err error) {
	var v url.Values

	if user == nil {
		err = fmt.Errorf("The user param cannot be nil")
		return
	}

	if v, err = query.Values(user); err != nil {
		return
	}

	err = c.post("users", v, user)

	return
}

// UpdateUser updates a user account
func (c *Client) UpdateUser(id int, user *User) (err error) {
	var v url.Values

	if id <= 0 {
		err = fmt.Errorf("The id param should be > 0")
		return
	}

	if user == nil {
		err = fmt.Errorf("The user param cannot be nil")
		return
	}

	if v, err = query.Values(user); err != nil {
		return
	}

	err = c.put(fmt.Sprintf("users/%d", id), v, user)

	return
}

// DeleteUser deletes a user account
func (c *Client) DeleteUser(id int) (err error) {
	if id <= 0 {
		err = fmt.Errorf("The id param should be > 0")
		return
	}

	err = c.delete(fmt.Sprintf("users/%d", id))

	return
}
<commit_msg>FIX: Support user id field<commit_after>// BaruwaAPI Golang bindings for Baruwa REST API
// Copyright (C) 2019 Andrew Colin Kissa <andrew@topdog.za.net>

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

// Package api Golang bindings for Baruwa REST API
package api

import (
	"fmt"
	"net/url"

	"github.com/google/go-querystring/query"
)

// UserDomain holds user
type UserDomain struct {
	ID   int    `json:"id"`
	Name string `json:"name"`
}

// User holds users
type User struct {
	ID          int          `json:"id,omitempty"`
	Username    string       `json:"username"`
	Firstname   string       `json:"firstname"`
	Lastname    string       `json:"lastname"`
	Password1   string       `json:"password1,omitempty"`
	Password2   string       `json:"password2,omitempty"`
	Email       string       `json:"email"`
	Timezone    string       `json:"timezone"`
	AccountType int          `json:"account_type"`
	Domains     []UserDomain `json:"domains,omitempty"`
	Active      bool         `json:"active"`
	SendReport  bool         `json:"send_report"`
	SpamChecks  bool         `json:"spam_checks"`
	LowScore    float64      `json:"low_score"`
	HighScore   float64      `json:"high_score"`
	BlockMacros bool         `json:"block_macros"`
}

// GetUser returns a user account
func (c *Client) GetUser(id int) (user *User, err error) {
	if id <= 0 {
		err = fmt.Errorf("The id param should be > 0")
		return
	}

	err = c.get(fmt.Sprintf("users/%d", id), user)
	return
}

// CreateUser creates a user account
func (c *Client) CreateUser(user *User) (err error) {
	var v url.Values

	if user == nil {
		err = fmt.Errorf("The user param cannot be nil")
		return
	}

	if v, err = query.Values(user); err != nil {
		return
	}

	err = c.post("users", v, user)

	return
}

// UpdateUser updates a user account
func (c *Client) UpdateUser(user *User) (err error) {
	var v url.Values

	if user == nil {
		err = fmt.Errorf("The user param cannot be nil")
		return
	}

	if user.ID <= 0 {
		err = fmt.Errorf("The user.ID param should be > 0")
		return
	}

	if v, err = query.Values(user); err != nil {
		return
	}

	err = c.put(fmt.Sprintf("users/%d", user.ID), v, user)

	return
}

// DeleteUser deletes a user account
func (c *Client) DeleteUser(id int) (err error) {
	if id <= 0 {
		err = fmt.Errorf("The id param should be > 0")
		return
	}

	err = c.delete(fmt.Sprintf("users/%d", id))

	return
}
<|endoftext|>


<commit_before><commit_msg>Updated: make sure Close in safe.<commit_after><|endoftext|>


<commit_before>package main

import (
 "fmt"
 "os/exec"
 "log"
 "os"
 "bufio"
)

func main(){

  arg0 := "jq"
  arg1 := "'.Parameters'"
  arg2 := "file"
  arg3 := "'.Resources'"
  arg4 := "'.Mappings'"
  arg5 := "'.Metadata'"

reader := bufio.NewReader(os.Stdin)
fmt.Print("\n Enter path of your file: ")
file, _ := reader.ReadString('\n')

		fmt.Println("Checking Parameters")
        cmd1 := exec.Command(arg0, arg1, file)
        err1 := cmd1.Start()
		if err1 != nil {
            log.Fatal(err1)
			log.Printf("Parameters failed")
			os.Exit(1)
        }
		err1 = cmd1.Wait()

  fmt.Println("Parameters OK")

        fmt.Println("Checking Resources")
        cmd2 := exec.Command(arg0, arg3, file)
        err2 := cmd2.Start()
		if err2 != nil {
            log.Fatal(err2)
			log.Printf("Resources failed")
			os.Exit(1)
        }
		err1 = cmd2.Wait()

  fmt.Println("Resources OK")

  fmt.Println("Checking Mappings")
        cmd3 := exec.Command(arg0, arg4, file)
        err3 := cmd3.Start()
		if err3 != nil {
            log.Fatal(err3)
			log.Printf("Mappings failed")
			os.Exit(1)
        }
		err3 = cmd3.Wait()

  fmt.Println("Mappings OK")

fmt.Println("Checking Metadata")
        cmd4 := exec.Command(arg0, arg5, file)
        err4 := cmd4.Start()
		if err4 != nil {
            log.Fatal(err4)
			log.Printf("Metadata failed")
			os.Exit(1)
        }
		err4 = cmd4.Wait()

fmt.Println("Metadata OK")

fmt.Printf("Basic Checks completed.\nOPTIONS [.Metadata][.Resources] \nList of services: \n * .EcsCluster \n * .EcsElb \n * .ECSAutoScalingGroup \n * .ContainerInstances")
reader1 := bufio.NewReader(os.Stdin)
fmt.Print("\n Enter your service: ")
service, _ := reader1.ReadString('\n')
fmt.Println(service)

        cmdservice := exec.Command(arg0, service, arg2)
        //fmt.Println(cmdservice.Start())
        errservice := cmdservice.Start()
		if errservice != nil {
            log.Fatal(errservice)
			log.Printf("Service failed.Please check the list.")
			os.Exit(1)
        }
		errservice = cmdservice.Wait()

}
<commit_msg>adding new functions<commit_after>package main

import (
 "fmt"
 "os/exec"
 "log"
 "os"
 "bufio"
 "runtime"
)

func test(){

    fmt.Println("Checking Prerequisites ..")

    if runtime.GOOS == "windows" {
    cmdos1 := exec.Command("type", "-a ", "jq")
        erros1 := cmdos1.Start()
		if erros1 != nil {
            log.Fatal(erros1)
			log.Printf("JQ not available")
            fmt.Println("Installing JQ on windows")
            cmdinstall1 := exec.Command("sudo", "apt-get", "install", "jq")
            cmdinstall1.Start()
			os.Exit(1)
        }
		erros1 = cmdos1.Wait()

    fmt.Println("JQ installed on windows")
    }

    if runtime.GOOS == "linux" {
        cmdos2 := exec.Command("sudo", "type", "-a ", "jq")
        erros2 := cmdos2.Start()
		if erros2 != nil {
            log.Fatal(erros2)
			log.Printf("JQ not available")
            fmt.Println("Installing JQ on windows")
            cmdinstall2 := exec.Command("sudo", "apt-get", "install", "jq")
            cmdinstall2.Start()
			os.Exit(1)
        }
		erros2 = cmdos2.Wait()

    fmt.Println("JQ installed on linux")
    }

}

func main(){

  arg0 := "jq"
  arg1 := "'.Parameters'"
  //arg2 := "file"
  arg3 := "'.Resources'"
  arg4 := "'.Mappings'"
  arg5 := "'.Metadata'"

test()

reader := bufio.NewReader(os.Stdin)
fmt.Print("\n Enter path of your file: ")
file, _ := reader.ReadString('\n')

		fmt.Println("Checking Parameters")
        cmd1 := exec.Command(arg0, arg1, file)
        err1 := cmd1.Start()
		if err1 != nil {
            log.Fatal(err1)
			log.Printf("Parameters failed")
			os.Exit(1)
        }
		err1 = cmd1.Wait()

  fmt.Println("Parameters OK")

        fmt.Println("Checking Resources")
        cmd2 := exec.Command(arg0, arg3, file)
        err2 := cmd2.Start()
		if err2 != nil {
            log.Fatal(err2)
			log.Printf("Resources failed")
			os.Exit(1)
        }
		err1 = cmd2.Wait()

  fmt.Println("Resources OK")

  fmt.Println("Checking Mappings")
        cmd3 := exec.Command(arg0, arg4, file)
        err3 := cmd3.Start()
		if err3 != nil {
            log.Fatal(err3)
			log.Printf("Mappings failed")
			os.Exit(1)
        }
		err3 = cmd3.Wait()

  fmt.Println("Mappings OK")

fmt.Println("Checking Metadata")
        cmd4 := exec.Command(arg0, arg5, file)
        err4 := cmd4.Start()
		if err4 != nil {
            log.Fatal(err4)
			log.Printf("Metadata failed")
			os.Exit(1)
        }
		err4 = cmd4.Wait()

fmt.Println("Metadata OK")

readerforward := bufio.NewReader(os.Stdin)
fmt.Printf("Do you wish to check your services (y/n) ?")
service, _ := readerforward.ReadString('\n')

    if service == "y" {
        fmt.Printf("Starting service ...")
        //service();
    } else {
    os.Exit(1)
    }
}

func service(){

fmt.Printf("Basic Checks completed.\nOPTIONS [.Metadata][.Resources] \nList of services: \n * .EcsCluster \n * .EcsElb \n * .ECSAutoScalingGroup \n * .ContainerInstances")
reader1 := bufio.NewReader(os.Stdin)
fmt.Print("\n Enter your service: ")
service, _ := reader1.ReadString('\n')

        cmdservice := exec.Command("jq", service, "ecs-pipeline.json")
        fmt.Println(cmdservice.Start())
        errservice := cmdservice.Start()
		if errservice != nil {
            log.Fatal(errservice)
			log.Printf("Service failed.Please check the list.")
			os.Exit(1)
        }
		errservice = cmdservice.Wait()
}
<|endoftext|>


<commit_before><commit_msg>Simplify and upload files to the appropriate location.<commit_after><|endoftext|>


<commit_before>package interp

import (
	"common"
	"fmt"

	zygo "github.com/glycerine/zygomys/repl"
)

// EventType define la clase de eventos que se pueden emitir
type EventType int

const (
	// InterpDone señaliza que el interprete terminó de evaluar código
	InterpDone EventType = iota
	// Error reporta un error
	Error
)

// Event se utiliza para representar un evento emitido
type Event struct {
	Type EventType
	Data interface{}
}

var in chan common.Command
var out chan Event

var env *zygo.Glisp

func init() {
	env = zygo.NewGlisp()
	in = make(chan common.Command)
	out = make(chan Event)
}

// Start inicia el módulo
func Start() <-chan Event {
	go loop(in)
	return out
}

// In regresa el channel para mandar comandos al módulo
func In() chan<- common.Command {
	return in
}

func loop(input <-chan common.Command) {
	for c := range input {
		switch c.Cmd {
		case "interp":
			fmt.Println(c.Args["code"])
			if err := env.LoadString(c.Args["code"]); err != nil {
				env.Clear()
				out <- Event{
					Type: Error,
					Data: fmt.Sprintf("imposible cargar código: %s", err.Error()),
				}
			} else if expr, err := env.Run(); err != nil {
				env.Clear()
				out <- Event{
					Type: Error,
					Data: fmt.Sprintf("imposible cargar código: %s", err.Error()),
				}
			} else {
				env.Clear()
				out <- Event{
					Type: InterpDone,
					Data: map[string]string{
						"peer":   c.Args["peer"],
						"result": fmt.Sprintf(expr.SexpString()),
					},
				}
			}
		}
	}
}
<commit_msg>removed debug message<commit_after>package interp

import (
	"common"
	"fmt"

	zygo "github.com/glycerine/zygomys/repl"
)

// EventType define la clase de eventos que se pueden emitir
type EventType int

const (
	// InterpDone señaliza que el interprete terminó de evaluar código
	InterpDone EventType = iota
	// Error reporta un error
	Error
)

// Event se utiliza para representar un evento emitido
type Event struct {
	Type EventType
	Data interface{}
}

var in chan common.Command
var out chan Event

var env *zygo.Glisp

func init() {
	env = zygo.NewGlisp()
	in = make(chan common.Command)
	out = make(chan Event)
}

// Start inicia el módulo
func Start() <-chan Event {
	go loop(in)
	return out
}

// In regresa el channel para mandar comandos al módulo
func In() chan<- common.Command {
	return in
}

func loop(input <-chan common.Command) {
	for c := range input {
		switch c.Cmd {
		case "interp":
			if err := env.LoadString(c.Args["code"]); err != nil {
				env.Clear()
				out <- Event{
					Type: Error,
					Data: fmt.Sprintf("imposible cargar código: %s", err.Error()),
				}
			} else if expr, err := env.Run(); err != nil {
				env.Clear()
				out <- Event{
					Type: Error,
					Data: fmt.Sprintf("imposible cargar código: %s", err.Error()),
				}
			} else {
				env.Clear()
				out <- Event{
					Type: InterpDone,
					Data: map[string]string{
						"peer":   c.Args["peer"],
						"result": fmt.Sprintf(expr.SexpString()),
					},
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 The Serviced 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.

// +build integration

package facade

import (
	"github.com/control-center/serviced/auth"
	"github.com/control-center/serviced/datastore"
	"github.com/control-center/serviced/datastore/elastic"
	dfsmocks "github.com/control-center/serviced/dfs/mocks"
	"github.com/control-center/serviced/domain/addressassignment"
	"github.com/control-center/serviced/domain/host"
	"github.com/control-center/serviced/domain/pool"
	"github.com/control-center/serviced/domain/registry"
	"github.com/control-center/serviced/domain/service"
	"github.com/control-center/serviced/domain/serviceconfigfile"
	"github.com/control-center/serviced/domain/servicetemplate"
	"github.com/control-center/serviced/domain/user"
	zzkmocks "github.com/control-center/serviced/facade/mocks"
	"github.com/stretchr/testify/mock"
	gocheck "gopkg.in/check.v1"
)

// FacadeIntegrationTest used for running integration tests where a facade type is needed.
type FacadeIntegrationTest struct {
	elastic.ElasticTest
	CTX    datastore.Context
	Facade *Facade
	zzk    *zzkmocks.ZZK
	dfs    *dfsmocks.DFS
}

var _ = gocheck.Suite(&FacadeIntegrationTest{})

//SetUpSuite sets up test suite
func (ft *FacadeIntegrationTest) SetUpSuite(c *gocheck.C) {

	//set up index and mappings before setting up elastic
	ft.Index = "controlplane"
	if ft.Mappings == nil {
		ft.Mappings = make([]elastic.Mapping, 0)
	}
	ft.Mappings = append(ft.Mappings, host.MAPPING)
	ft.Mappings = append(ft.Mappings, pool.MAPPING)
	ft.Mappings = append(ft.Mappings, service.MAPPING)
	ft.Mappings = append(ft.Mappings, servicetemplate.MAPPING)
	ft.Mappings = append(ft.Mappings, addressassignment.MAPPING)
	ft.Mappings = append(ft.Mappings, serviceconfigfile.MAPPING)
	ft.Mappings = append(ft.Mappings, user.MAPPING)
	ft.Mappings = append(ft.Mappings, registry.MAPPING)

	ft.ElasticTest.SetUpSuite(c)
	datastore.Register(ft.Driver())
	ft.CTX = datastore.Get()

	ft.Facade = New()
	ft.dfs = &dfsmocks.DFS{}
	ft.Facade.SetDFS(ft.dfs)
	ft.setupMockDFS()

	// Create a master key pair
	pub, priv, _ := auth.GenerateRSAKeyPairPEM(nil)
	auth.LoadMasterKeysFromPEM(pub, priv)
}

func (ft *FacadeIntegrationTest) SetUpTest(c *gocheck.C) {
	ft.ElasticTest.SetUpTest(c)
	ft.zzk = &zzkmocks.ZZK{}
	ft.Facade.SetZZK(ft.zzk)
	ft.dfs = &dfsmocks.DFS{}
	ft.Facade.SetDFS(ft.dfs)
	ft.setupMockZZK()
	ft.setupMockDFS()
	LogstashContainerReloader = reloadLogstashContainerStub
}

func (ft *FacadeIntegrationTest) setupMockZZK() {
	ft.zzk.On("AddResourcePool", mock.AnythingOfType("*pool.ResourcePool")).Return(nil)
	ft.zzk.On("UpdateResourcePool", mock.AnythingOfType("*pool.ResourcePool")).Return(nil)
	ft.zzk.On("RemoveResourcePool", mock.AnythingOfType("string")).Return(nil)
	ft.zzk.On("AddVirtualIP", mock.AnythingOfType("*pool.VirtualIP")).Return(nil)
	ft.zzk.On("RemoveVirtualIP", mock.AnythingOfType("*pool.VirtualIP")).Return(nil)
	ft.zzk.On("AddHost", mock.AnythingOfType("*host.Host")).Return(nil)
	ft.zzk.On("UpdateHost", mock.AnythingOfType("*host.Host")).Return(nil)
	ft.zzk.On("RemoveHost", mock.AnythingOfType("*host.Host")).Return(nil)
	ft.zzk.On("UpdateService", mock.AnythingOfType("*datastore.Context"), mock.AnythingOfType("string"), mock.AnythingOfType("*service.Service"), mock.AnythingOfType("bool"), mock.AnythingOfType("bool")).Return(nil)
	ft.zzk.On("RemoveService", mock.AnythingOfType("string"), mock.AnythingOfType("string")).Return(nil)
	ft.zzk.On("RemoveServiceEndpoints", mock.AnythingOfType("string")).Return(nil)
	ft.zzk.On("RemoveTenantExports", mock.AnythingOfType("string")).Return(nil)
	ft.zzk.On("SetRegistryImage", mock.AnythingOfType("*registry.Image")).Return(nil)
	ft.zzk.On("DeleteRegistryImage", mock.AnythingOfType("string")).Return(nil)
	ft.zzk.On("DeleteRegistryLibrary", mock.AnythingOfType("string")).Return(nil)
	ft.zzk.On("LockServices", mock.AnythingOfType("[]service.Service")).Return(nil)
	ft.zzk.On("UnlockServices", mock.AnythingOfType("[]service.Service")).Return(nil)

}

func (ft *FacadeIntegrationTest) setupMockDFS() {
	ft.dfs.On("Destroy", mock.AnythingOfType("string")).Return(nil)
}

func (ft *FacadeIntegrationTest) TearDownTest(c *gocheck.C) {
}

func reloadLogstashContainerStub(_ datastore.Context, _ FacadeInterface) error {
	return nil
}
<commit_msg>still trying to fix broken mock<commit_after>// Copyright 2014 The Serviced 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.

// +build integration

package facade

import (
	"github.com/control-center/serviced/auth"
	"github.com/control-center/serviced/datastore"
	"github.com/control-center/serviced/datastore/elastic"
	dfsmocks "github.com/control-center/serviced/dfs/mocks"
	"github.com/control-center/serviced/domain/addressassignment"
	"github.com/control-center/serviced/domain/host"
	"github.com/control-center/serviced/domain/pool"
	"github.com/control-center/serviced/domain/registry"
	"github.com/control-center/serviced/domain/service"
	"github.com/control-center/serviced/domain/serviceconfigfile"
	"github.com/control-center/serviced/domain/servicetemplate"
	"github.com/control-center/serviced/domain/user"
	zzkmocks "github.com/control-center/serviced/facade/mocks"
	"github.com/stretchr/testify/mock"
	gocheck "gopkg.in/check.v1"
)

// FacadeIntegrationTest used for running integration tests where a facade type is needed.
type FacadeIntegrationTest struct {
	elastic.ElasticTest
	CTX    datastore.Context
	Facade *Facade
	zzk    *zzkmocks.ZZK
	dfs    *dfsmocks.DFS
}

var _ = gocheck.Suite(&FacadeIntegrationTest{})

//SetUpSuite sets up test suite
func (ft *FacadeIntegrationTest) SetUpSuite(c *gocheck.C) {

	//set up index and mappings before setting up elastic
	ft.Index = "controlplane"
	if ft.Mappings == nil {
		ft.Mappings = make([]elastic.Mapping, 0)
	}
	ft.Mappings = append(ft.Mappings, host.MAPPING)
	ft.Mappings = append(ft.Mappings, pool.MAPPING)
	ft.Mappings = append(ft.Mappings, service.MAPPING)
	ft.Mappings = append(ft.Mappings, servicetemplate.MAPPING)
	ft.Mappings = append(ft.Mappings, addressassignment.MAPPING)
	ft.Mappings = append(ft.Mappings, serviceconfigfile.MAPPING)
	ft.Mappings = append(ft.Mappings, user.MAPPING)
	ft.Mappings = append(ft.Mappings, registry.MAPPING)

	ft.ElasticTest.SetUpSuite(c)
	datastore.Register(ft.Driver())
	ft.CTX = datastore.Get()

	ft.Facade = New()
	ft.dfs = &dfsmocks.DFS{}
	ft.Facade.SetDFS(ft.dfs)
	ft.setupMockDFS()

	// Create a master key pair
	pub, priv, _ := auth.GenerateRSAKeyPairPEM(nil)
	auth.LoadMasterKeysFromPEM(pub, priv)
}

func (ft *FacadeIntegrationTest) SetUpTest(c *gocheck.C) {
	ft.ElasticTest.SetUpTest(c)
	ft.zzk = &zzkmocks.ZZK{}
	ft.Facade.SetZZK(ft.zzk)
	ft.dfs = &dfsmocks.DFS{}
	ft.Facade.SetDFS(ft.dfs)
	ft.setupMockZZK()
	ft.setupMockDFS()
	LogstashContainerReloader = reloadLogstashContainerStub
}

func (ft *FacadeIntegrationTest) setupMockZZK() {
	ft.zzk.On("AddResourcePool", mock.AnythingOfType("*pool.ResourcePool")).Return(nil)
	ft.zzk.On("UpdateResourcePool", mock.AnythingOfType("*pool.ResourcePool")).Return(nil)
	ft.zzk.On("RemoveResourcePool", mock.AnythingOfType("string")).Return(nil)
	ft.zzk.On("AddVirtualIP", mock.AnythingOfType("*pool.VirtualIP")).Return(nil)
	ft.zzk.On("RemoveVirtualIP", mock.AnythingOfType("*pool.VirtualIP")).Return(nil)
	ft.zzk.On("AddHost", mock.AnythingOfType("*host.Host")).Return(nil)
	ft.zzk.On("UpdateHost", mock.AnythingOfType("*host.Host")).Return(nil)
	ft.zzk.On("RemoveHost", mock.AnythingOfType("*host.Host")).Return(nil)
	ft.zzk.On("UpdateService", mock.AnythingOfType("*datastore.context"), mock.AnythingOfType("string"), mock.AnythingOfType("*service.Service"), mock.AnythingOfType("bool"), mock.AnythingOfType("bool")).Return(nil)
	ft.zzk.On("RemoveService", mock.AnythingOfType("string"), mock.AnythingOfType("string")).Return(nil)
	ft.zzk.On("RemoveServiceEndpoints", mock.AnythingOfType("string")).Return(nil)
	ft.zzk.On("RemoveTenantExports", mock.AnythingOfType("string")).Return(nil)
	ft.zzk.On("SetRegistryImage", mock.AnythingOfType("*registry.Image")).Return(nil)
	ft.zzk.On("DeleteRegistryImage", mock.AnythingOfType("string")).Return(nil)
	ft.zzk.On("DeleteRegistryLibrary", mock.AnythingOfType("string")).Return(nil)
	ft.zzk.On("LockServices", mock.AnythingOfType("[]service.Service")).Return(nil)
	ft.zzk.On("UnlockServices", mock.AnythingOfType("[]service.Service")).Return(nil)

}

func (ft *FacadeIntegrationTest) setupMockDFS() {
	ft.dfs.On("Destroy", mock.AnythingOfType("string")).Return(nil)
}

func (ft *FacadeIntegrationTest) TearDownTest(c *gocheck.C) {
}

func reloadLogstashContainerStub(_ datastore.Context, _ FacadeInterface) error {
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// 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 main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"

	_ "github.com/go-sql-driver/mysql"
	"google.golang.org/appengine"
	serverLog "google.golang.org/appengine/log"
)

var (
	db          *sql.DB
	queryString = "SELECT probability_of_answer, probability_of_downvote, minutes FROM questions WHERE tag = ? AND first_word = ? AND ends_question = ? AND weekday_utc = ? AND account_creation_year = ? AND question_length = ? AND hour_utc = ?"
	port        = ":8080"
)

func main() {
	var err error
	db, err = getDatabase()
	if err != nil {
		log.Printf("Failed to establish connection to Database. Server terminated")
		return
	}

	http.HandleFunc("/experiment/bqml-stackoverflow/api/", func(w http.ResponseWriter, r *http.Request) {
		ctx := appengine.NewContext(r)
		rows, err := attemptQuery(5, r)
		if err != nil {
			serverLog.Debugf(ctx, "db.Query: %v", err)

			w.Write(getEmptyAnswer())
			return
		}
		defer rows.Close()

		if rows.Next() {
			var probabilityOfAnswer float32
			var probablityOfDownvote float32
			var minutes float32

			if err := rows.Scan(&probabilityOfAnswer, &probablityOfDownvote, &minutes); err != nil {
				log.Printf("rows.Scan: %v", err)
			}

			var answer = &Answer{
				Minutes:   minutes,
				PDownvote: probablityOfDownvote,
				PAnswer:   probabilityOfAnswer,
			}

			json, _ := json.Marshal(answer)

			w.Write(json)
		} else {
			w.Write(getEmptyAnswer())
		}

	})

	fs := wrapHandler(http.FileServer(http.Dir("./dist")))
	http.HandleFunc("/", fs)

	fmt.Printf("Starting server on port %s\n", port)
	if err := http.ListenAndServe(port, nil); err != nil {
		log.Fatal(err)
	}

}

func getEmptyAnswer() []byte {
	json, _ := json.Marshal(&Answer{
		Minutes:   0,
		PAnswer:   0,
		PDownvote: 0,
	})

	return json
}

func getDatabase() (*sql.DB, error) {
	user := os.Getenv("DB_USER")
	password := os.Getenv("DB_PASSWORD")
	table := os.Getenv("DB_DATABASE")
	instance := os.Getenv("DB_INSTANCE")

	if appengine.IsDevAppServer() {
		return sql.Open("mysql", fmt.Sprintf("%s:%s@tcp([localhost]:3306)/%s", user, password, table))
	}

	return sql.Open("mysql", fmt.Sprintf("%s:%s@cloudsql(%s)/%s", user, password, instance, table))
}

func attemptQuery(attempts int, r *http.Request) (*sql.Rows, error) {
	var rows *sql.Rows
	var err error

	query := r.URL.Query()

	for i := 0; i < attempts; i++ {
		rows, err = db.Query(
			queryString,
			query.Get("tag"),
			query.Get("first_word"),
			query.Get("ends_question"),
			query.Get("weekday_utc"),
			query.Get("account_creation_year"),
			query.Get("question_length"),
			query.Get("hour_utc"),
		)

		if err == nil {
			break
		} else {
			db, err = getDatabase()
		}
	}

	return rows, err
}

// Answer respresents the api query response json.
type Answer struct {
	Minutes   float32 `json:"minutes"`
	PAnswer   float32 `json:"probability_of_answer"`
	PDownvote float32 `json:"probabiliy_of_downvote"`
}

type NotFoundRedirectRespWr struct {
	http.ResponseWriter // We embed http.ResponseWriter
	status              int
}

func (w *NotFoundRedirectRespWr) WriteHeader(status int) {
	w.status = status // Store the status for our own use
	if status != http.StatusNotFound {
		w.ResponseWriter.WriteHeader(status)
	}
}

func (w *NotFoundRedirectRespWr) Write(p []byte) (int, error) {
	if w.status != http.StatusNotFound {
		return w.ResponseWriter.Write(p)
	}
	return len(p), nil // Lie that we successfully written it
}

func wrapHandler(h http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		nfrw := &NotFoundRedirectRespWr{ResponseWriter: w}
		h.ServeHTTP(nfrw, r)
		if nfrw.status == 404 {
			http.Redirect(w, r, "/index.html", http.StatusFound)
		}
	}
}
<commit_msg>Update main.go<commit_after>// Copyright 2019 Google LLC
//
// 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 main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"

	_ "github.com/go-sql-driver/mysql"
	"google.golang.org/appengine"
	serverLog "google.golang.org/appengine/log"
)

var (
	db          *sql.DB
	queryString = "SELECT probability_of_answer, probability_of_downvote, minutes FROM questions WHERE tag = ? AND first_word = ? AND ends_question = ? AND weekday_utc = ? AND account_creation_year = ? AND question_length = ? AND hour_utc = ?"
	port        = ":8080"
)

func main() {
	var err error
	db, err = getDatabase()
	if err != nil {
		log.Printf("Failed to establish connection to Database. Server terminated")
		return
	}

	http.HandleFunc("/experiment/bqml-stackoverflow/api/", func(w http.ResponseWriter, r *http.Request) {
		ctx := appengine.NewContext(r)
		rows, err := attemptQuery(5, r)
		if err != nil {
			serverLog.Debugf(ctx, "db.Query: %v", err)

			w.Write(getEmptyAnswer())
			return
		}
		defer rows.Close()

		if rows.Next() {
			var probabilityOfAnswer float32
			var probablityOfDownvote float32
			var minutes float32

			if err := rows.Scan(&probabilityOfAnswer, &probablityOfDownvote, &minutes); err != nil {
				log.Printf("rows.Scan: %v", err)
			}

			var answer = &Answer{
				Minutes:   minutes,
				PDownvote: probablityOfDownvote,
				PAnswer:   probabilityOfAnswer,
			}

			json, _ := json.Marshal(answer)

			w.Write(json)
		} else {
			w.Write(getEmptyAnswer())
		}

	})

	fs := wrapHandler(http.FileServer(http.Dir("./dist")))
	http.HandleFunc("/experiment/bqml-stackoverflow/", fs)

	fmt.Printf("Starting server on port %s\n", port)
	if err := http.ListenAndServe(port, nil); err != nil {
		log.Fatal(err)
	}

}

func getEmptyAnswer() []byte {
	json, _ := json.Marshal(&Answer{
		Minutes:   0,
		PAnswer:   0,
		PDownvote: 0,
	})

	return json
}

func getDatabase() (*sql.DB, error) {
	user := os.Getenv("DB_USER")
	password := os.Getenv("DB_PASSWORD")
	table := os.Getenv("DB_DATABASE")
	instance := os.Getenv("DB_INSTANCE")

	if appengine.IsDevAppServer() {
		return sql.Open("mysql", fmt.Sprintf("%s:%s@tcp([localhost]:3306)/%s", user, password, table))
	}

	return sql.Open("mysql", fmt.Sprintf("%s:%s@cloudsql(%s)/%s", user, password, instance, table))
}

func attemptQuery(attempts int, r *http.Request) (*sql.Rows, error) {
	var rows *sql.Rows
	var err error

	query := r.URL.Query()

	for i := 0; i < attempts; i++ {
		rows, err = db.Query(
			queryString,
			query.Get("tag"),
			query.Get("first_word"),
			query.Get("ends_question"),
			query.Get("weekday_utc"),
			query.Get("account_creation_year"),
			query.Get("question_length"),
			query.Get("hour_utc"),
		)

		if err == nil {
			break
		} else {
			db, err = getDatabase()
		}
	}

	return rows, err
}

// Answer respresents the api query response json.
type Answer struct {
	Minutes   float32 `json:"minutes"`
	PAnswer   float32 `json:"probability_of_answer"`
	PDownvote float32 `json:"probabiliy_of_downvote"`
}

type NotFoundRedirectRespWr struct {
	http.ResponseWriter // We embed http.ResponseWriter
	status              int
}

func (w *NotFoundRedirectRespWr) WriteHeader(status int) {
	w.status = status // Store the status for our own use
	if status != http.StatusNotFound {
		w.ResponseWriter.WriteHeader(status)
	}
}

func (w *NotFoundRedirectRespWr) Write(p []byte) (int, error) {
	if w.status != http.StatusNotFound {
		return w.ResponseWriter.Write(p)
	}
	return len(p), nil // Lie that we successfully written it
}

func wrapHandler(h http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		nfrw := &NotFoundRedirectRespWr{ResponseWriter: w}
		h.ServeHTTP(nfrw, r)
		if nfrw.status == 404 {
			http.Redirect(w, r, "/index.html", http.StatusFound)
		}
	}
}
<|endoftext|>


<commit_before>package v1

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
	"time"

	redis "gopkg.in/redis.v5"
	"zxq.co/ripple/rippleapi/common"
)

type setAllowedData struct {
	UserID  int `json:"user_id"`
	Allowed int `json:"allowed"`
}

// UserManageSetAllowedPOST allows to set the allowed status of an user.
func UserManageSetAllowedPOST(md common.MethodData) common.CodeMessager {
	var data setAllowedData
	if err := md.Unmarshal(&data); err != nil {
		return ErrBadJSON
	}
	if data.Allowed < 0 || data.Allowed > 2 {
		return common.SimpleResponse(400, "Allowed status must be between 0 and 2")
	}
	var banDatetime int64
	var privsSet string
	if data.Allowed == 0 {
		banDatetime = time.Now().Unix()
		privsSet = "privileges = (privileges & ~3)"
	} else {
		banDatetime = 0
		privsSet = "privileges = (privileges | 3)"
	}
	_, err := md.DB.Exec("UPDATE users SET "+privsSet+", ban_datetime = ? WHERE id = ?", banDatetime, data.UserID)
	if err != nil {
		md.Err(err)
		return Err500
	}
	rapLog(md, fmt.Sprintf("changed UserID:%d's allowed to %d. This was done using the API's terrible ManageSetAllowed.", data.UserID, data.Allowed))
	go fixPrivileges(data.UserID, md.DB)
	query := `
SELECT users.id, users.username, register_datetime, privileges,
	latest_activity, users_stats.username_aka,
	users_stats.country
FROM users
LEFT JOIN users_stats
ON users.id=users_stats.id
WHERE users.id=?
LIMIT 1`
	return userPutsSingle(md, md.DB.QueryRowx(query, data.UserID))
}

type userEditData struct {
	ID            int          `json:"id"`
	Username      *string      `json:"username"`
	UsernameAKA   *string      `json:"username_aka"`
	Privileges    *uint64      `json:"privileges"`
	Country       *string      `json:"country"`
	SilenceInfo   *silenceInfo `json:"silence_info"`
	ResetUserpage bool         `json:"reset_userpage"`
	//ResetAvatar bool `json:"reset_avatar"`
}

var privChangeList = [...]string{
	"banned",
	"locked",
	"restricted",
	"removed all restrictions on",
}

// UserEditPOST allows to edit an user's information.
func UserEditPOST(md common.MethodData) common.CodeMessager {
	var data userEditData
	if err := md.Unmarshal(&data); err != nil {
		fmt.Println(err)
		return ErrBadJSON
	}

	if data.ID == 0 {
		return common.SimpleResponse(404, "That user could not be found")
	}

	var prevUser struct {
		Username   string
		Privileges uint64
	}
	err := md.DB.Get(&prevUser, "SELECT username, privileges FROM users WHERE id = ? LIMIT 1", data.ID)

	switch err {
	case nil: // carry on
	case sql.ErrNoRows:
		return common.SimpleResponse(404, "That user could not be found")
	default:
		md.Err(err)
		return Err500
	}

	const initQuery = "UPDATE users SET\n"
	q := initQuery
	var args []interface{}

	// totally did not realise I had to update some fields in users_stats as well
	// and just copy pasting the above code by prefixing "stats" to every
	// variable
	const statsInitQuery = "UPDATE users_stats SET\n"
	statsQ := statsInitQuery
	var statsArgs []interface{}

	if common.UserPrivileges(prevUser.Privileges)&common.AdminPrivilegeManageUsers != 0 &&
		data.ID != md.User.UserID {
		return common.SimpleResponse(403, "Can't edit that user")
	}

	var isBanned bool
	var isSilenced bool
	if data.Privileges != nil {
		// If we want to modify privileges other than Normal/Public, we need to have
		// the right privilege ourselves and AdminManageUsers won't suffice.
		if (*data.Privileges&^3) != (prevUser.Privileges&^3) &&
			md.User.UserPrivileges&common.AdminPrivilegeManagePrivilege == 0 {
			return common.SimpleResponse(403, "Can't modify user privileges without AdminManagePrivileges")
		}
		q += "privileges = ?,\n"
		args = append(args, *data.Privileges)

		// UserPublic became 0, so banned or restricted
		const uPublic = uint64(common.UserPrivilegePublic)
		if *data.Privileges&uPublic == 0 && prevUser.Privileges&uPublic != 0 {
			q += "ban_datetime = ?,\n"
			args = append(args, time.Now().Unix())
			isBanned = true
		}

		// If we modified other privileges apart from Normal and Public, we use a generic
		// "changed user's privileges". Otherwise, we are more descriptive.
		if *data.Privileges^prevUser.Privileges > 3 {
			rapLog(md, fmt.Sprintf("has changed %s's privileges", prevUser.Username))
		} else {
			rapLog(md, fmt.Sprintf("has %s %s", privChangeList[*data.Privileges&3], prevUser.Username))
		}
	}
	if data.Username != nil {
		if strings.Contains(*data.Username, " ") && strings.Contains(*data.Username, "_") {
			return common.SimpleResponse(400, "Mixed spaces and underscores")
		}
		if usernameAvailable(md, *data.Username, data.ID) {
			return common.SimpleResponse(409, "User with that username exists")
		}
		jsonData, _ := json.Marshal(struct {
			UserID      int    `json:"userID"`
			NewUsername string `json:"newUsername"`
		}{data.ID, *data.Username})
		md.R.Publish("peppy:change_username", string(jsonData))
	}
	if data.UsernameAKA != nil {
		statsQ += "username_aka = ?,\n"
		statsArgs = append(statsArgs, *data.UsernameAKA)
	}
	if data.Country != nil {
		statsQ += "country = ?,\n"
		statsArgs = append(statsArgs, *data.Country)
		rapLog(md, fmt.Sprintf("has changed %s country to %s", prevUser.Username, *data.Country))
		appendToUserNotes(md, "country changed to "+*data.Country, data.ID)
	}
	if data.SilenceInfo != nil && md.User.UserPrivileges&common.AdminPrivilegeSilenceUsers != 0 {
		q += "silence_end = ?, silence_reason = ?,\n"
		args = append(args, time.Time(data.SilenceInfo.End).Unix(), data.SilenceInfo.Reason)
		isSilenced = true
	}
	if data.ResetUserpage {
		statsQ += "userpage_content = '',\n"
	}

	if q != initQuery {
		q = q[:len(q)-2] + " WHERE id = ? LIMIT 1"
		args = append(args, data.ID)
		_, err = md.DB.Exec(q, args...)
		if err != nil {
			md.Err(err)
			return Err500
		}

	}
	if statsQ != statsInitQuery {
		statsQ = statsQ[:len(statsQ)-2] + " WHERE id = ? LIMIT 1"
		statsArgs = append(statsArgs, data.ID)
		_, err = md.DB.Exec(statsQ, statsArgs...)
		if err != nil {
			md.Err(err)
			return Err500
		}
	}

	if isBanned || isSilenced {
		if err := updateBanBancho(md.R, data.ID); err != nil {
			md.Err(err)
			return Err500
		}
	}

	rapLog(md, fmt.Sprintf("has updated user %s", prevUser.Username))

	return userPutsSingle(md, md.DB.QueryRowx(userFields+" WHERE users.id = ? LIMIT 1", data.ID))
}

func updateBanBancho(r *redis.Client, user int) error {
	return r.Publish("peppy:ban", strconv.Itoa(user)).Err()
}

type wipeUserData struct {
	ID    int   `json:"id"`
	Modes []int `json:"modes"`
}

// WipeUserPOST wipes an user's scores.
func WipeUserPOST(md common.MethodData) common.CodeMessager {
	var data wipeUserData
	if err := md.Unmarshal(&data); err != nil {
		return ErrBadJSON
	}
	if data.ID == 0 {
		return ErrMissingField("id")
	}
	if len(data.Modes) == 0 {
		return ErrMissingField("modes")
	}

	var userData struct {
		Username   string
		Privileges uint64
	}
	err := md.DB.Get(&userData, "SELECT username, privileges FROM users WHERE id = ?", data.ID)
	switch err {
	case sql.ErrNoRows:
		return common.SimpleResponse(404, "That user could not be found!")
	case nil: // carry on
	default:
		md.Err(err)
		return Err500
	}

	if common.UserPrivileges(userData.Privileges)&common.AdminPrivilegeManageUsers != 0 {
		return common.SimpleResponse(403, "Can't edit that user")
	}

	tx, err := md.DB.Beginx()
	if err != nil {
		md.Err(err)
		return Err500
	}

	for _, mode := range data.Modes {
		if mode < 0 || mode > 3 {
			continue
		}
		_, err = tx.Exec("INSERT INTO scores_removed SELECT * FROM scores WHERE userid = ? AND play_mode = ?", data.ID, mode)
		if err != nil {
			md.Err(err)
		}
		_, err = tx.Exec("DELETE FROM scores WHERE userid = ? AND play_mode = ?", data.ID, mode)
		if err != nil {
			md.Err(err)
		}
		_, err = tx.Exec(strings.Replace(
			`UPDATE users_stats SET total_score_MODE = 0, ranked_score_MODE = 0, replays_watched_MODE = 0,
			playcount_MODE = 0, avg_accuracy_MODE = 0, total_hits_MODE = 0, level_MODE = 0, pp_MODE = 0
			WHERE id = ?`, "MODE", modesToReadable[mode], -1,
		), data.ID)
		if err != nil {
			md.Err(err)
		}
	}

	if err = tx.Commit(); err != nil {
		md.Err(err)
		return Err500
	}

	rapLog(md, fmt.Sprintf("has wiped %s's account", userData.Username))

	return userPutsSingle(md, md.DB.QueryRowx(userFields+" WHERE users.id = ? LIMIT 1", data.ID))
}

func appendToUserNotes(md common.MethodData, message string, user int) {
	message = "\n[" + time.Now().Format("2006-01-02 15:04:05") + "] API: " + message
	_, err := md.DB.Exec("UPDATE users SET notes = CONCAT(COALESCE(notes, ''), ?) WHERE id = ?",
		message, user)
	if err != nil {
		md.Err(err)
	}
}

func usernameAvailable(md common.MethodData, u string, userID int) (r bool) {
	err := md.DB.QueryRow("SELECT EXISTS(SELECT 1 FROM users WHERE username_safe = ? AND id != ?)", common.SafeUsername(u), userID).Scan(&r)
	if err != nil && err != sql.ErrNoRows {
		md.Err(err)
	}
	return
}
<commit_msg>set_allowed: Allow privileges = 2 to restrict someone<commit_after>package v1

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
	"time"

	redis "gopkg.in/redis.v5"
	"zxq.co/ripple/rippleapi/common"
)

type setAllowedData struct {
	UserID  int `json:"user_id"`
	Allowed int `json:"allowed"`
}

// UserManageSetAllowedPOST allows to set the allowed status of an user.
func UserManageSetAllowedPOST(md common.MethodData) common.CodeMessager {
	var data setAllowedData
	if err := md.Unmarshal(&data); err != nil {
		return ErrBadJSON
	}
	if data.Allowed < 0 || data.Allowed > 2 {
		return common.SimpleResponse(400, "Allowed status must be between 0 and 2")
	}
	var banDatetime int64
	var privsSet string
	if data.Allowed == 0 {
		banDatetime = time.Now().Unix()
		privsSet = "privileges = (privileges & ~3)"
	} else if data.Allowed == 1 {
		banDatetime = 0
		privsSet = "privileges = (privileges | 3)"
	} else if data.Allowed == 2 {
		banDatetime = time.Now().Unix()
		privsSet = "privileges = (privileges | 2) & (privileges & ~1)"
	}
	_, err := md.DB.Exec("UPDATE users SET "+privsSet+", ban_datetime = ? WHERE id = ?", banDatetime, data.UserID)
	if err != nil {
		md.Err(err)
		return Err500
	}
	rapLog(md, fmt.Sprintf("changed UserID:%d's allowed to %d. This was done using the API's terrible ManageSetAllowed.", data.UserID, data.Allowed))
	go fixPrivileges(data.UserID, md.DB)
	query := `
SELECT users.id, users.username, register_datetime, privileges,
	latest_activity, users_stats.username_aka,
	users_stats.country
FROM users
LEFT JOIN users_stats
ON users.id=users_stats.id
WHERE users.id=?
LIMIT 1`
	return userPutsSingle(md, md.DB.QueryRowx(query, data.UserID))
}

type userEditData struct {
	ID            int          `json:"id"`
	Username      *string      `json:"username"`
	UsernameAKA   *string      `json:"username_aka"`
	Privileges    *uint64      `json:"privileges"`
	Country       *string      `json:"country"`
	SilenceInfo   *silenceInfo `json:"silence_info"`
	ResetUserpage bool         `json:"reset_userpage"`
	//ResetAvatar bool `json:"reset_avatar"`
}

var privChangeList = [...]string{
	"banned",
	"locked",
	"restricted",
	"removed all restrictions on",
}

// UserEditPOST allows to edit an user's information.
func UserEditPOST(md common.MethodData) common.CodeMessager {
	var data userEditData
	if err := md.Unmarshal(&data); err != nil {
		fmt.Println(err)
		return ErrBadJSON
	}

	if data.ID == 0 {
		return common.SimpleResponse(404, "That user could not be found")
	}

	var prevUser struct {
		Username   string
		Privileges uint64
	}
	err := md.DB.Get(&prevUser, "SELECT username, privileges FROM users WHERE id = ? LIMIT 1", data.ID)

	switch err {
	case nil: // carry on
	case sql.ErrNoRows:
		return common.SimpleResponse(404, "That user could not be found")
	default:
		md.Err(err)
		return Err500
	}

	const initQuery = "UPDATE users SET\n"
	q := initQuery
	var args []interface{}

	// totally did not realise I had to update some fields in users_stats as well
	// and just copy pasting the above code by prefixing "stats" to every
	// variable
	const statsInitQuery = "UPDATE users_stats SET\n"
	statsQ := statsInitQuery
	var statsArgs []interface{}

	if common.UserPrivileges(prevUser.Privileges)&common.AdminPrivilegeManageUsers != 0 &&
		data.ID != md.User.UserID {
		return common.SimpleResponse(403, "Can't edit that user")
	}

	var isBanned bool
	var isSilenced bool
	if data.Privileges != nil {
		// If we want to modify privileges other than Normal/Public, we need to have
		// the right privilege ourselves and AdminManageUsers won't suffice.
		if (*data.Privileges&^3) != (prevUser.Privileges&^3) &&
			md.User.UserPrivileges&common.AdminPrivilegeManagePrivilege == 0 {
			return common.SimpleResponse(403, "Can't modify user privileges without AdminManagePrivileges")
		}
		q += "privileges = ?,\n"
		args = append(args, *data.Privileges)

		// UserPublic became 0, so banned or restricted
		const uPublic = uint64(common.UserPrivilegePublic)
		if *data.Privileges&uPublic == 0 && prevUser.Privileges&uPublic != 0 {
			q += "ban_datetime = ?,\n"
			args = append(args, time.Now().Unix())
			isBanned = true
		}

		// If we modified other privileges apart from Normal and Public, we use a generic
		// "changed user's privileges". Otherwise, we are more descriptive.
		if *data.Privileges^prevUser.Privileges > 3 {
			rapLog(md, fmt.Sprintf("has changed %s's privileges", prevUser.Username))
		} else {
			rapLog(md, fmt.Sprintf("has %s %s", privChangeList[*data.Privileges&3], prevUser.Username))
		}
	}
	if data.Username != nil {
		if strings.Contains(*data.Username, " ") && strings.Contains(*data.Username, "_") {
			return common.SimpleResponse(400, "Mixed spaces and underscores")
		}
		if usernameAvailable(md, *data.Username, data.ID) {
			return common.SimpleResponse(409, "User with that username exists")
		}
		jsonData, _ := json.Marshal(struct {
			UserID      int    `json:"userID"`
			NewUsername string `json:"newUsername"`
		}{data.ID, *data.Username})
		md.R.Publish("peppy:change_username", string(jsonData))
	}
	if data.UsernameAKA != nil {
		statsQ += "username_aka = ?,\n"
		statsArgs = append(statsArgs, *data.UsernameAKA)
	}
	if data.Country != nil {
		statsQ += "country = ?,\n"
		statsArgs = append(statsArgs, *data.Country)
		rapLog(md, fmt.Sprintf("has changed %s country to %s", prevUser.Username, *data.Country))
		appendToUserNotes(md, "country changed to "+*data.Country, data.ID)
	}
	if data.SilenceInfo != nil && md.User.UserPrivileges&common.AdminPrivilegeSilenceUsers != 0 {
		q += "silence_end = ?, silence_reason = ?,\n"
		args = append(args, time.Time(data.SilenceInfo.End).Unix(), data.SilenceInfo.Reason)
		isSilenced = true
	}
	if data.ResetUserpage {
		statsQ += "userpage_content = '',\n"
	}

	if q != initQuery {
		q = q[:len(q)-2] + " WHERE id = ? LIMIT 1"
		args = append(args, data.ID)
		_, err = md.DB.Exec(q, args...)
		if err != nil {
			md.Err(err)
			return Err500
		}

	}
	if statsQ != statsInitQuery {
		statsQ = statsQ[:len(statsQ)-2] + " WHERE id = ? LIMIT 1"
		statsArgs = append(statsArgs, data.ID)
		_, err = md.DB.Exec(statsQ, statsArgs...)
		if err != nil {
			md.Err(err)
			return Err500
		}
	}

	if isBanned || isSilenced {
		if err := updateBanBancho(md.R, data.ID); err != nil {
			md.Err(err)
			return Err500
		}
	}

	rapLog(md, fmt.Sprintf("has updated user %s", prevUser.Username))

	return userPutsSingle(md, md.DB.QueryRowx(userFields+" WHERE users.id = ? LIMIT 1", data.ID))
}

func updateBanBancho(r *redis.Client, user int) error {
	return r.Publish("peppy:ban", strconv.Itoa(user)).Err()
}

type wipeUserData struct {
	ID    int   `json:"id"`
	Modes []int `json:"modes"`
}

// WipeUserPOST wipes an user's scores.
func WipeUserPOST(md common.MethodData) common.CodeMessager {
	var data wipeUserData
	if err := md.Unmarshal(&data); err != nil {
		return ErrBadJSON
	}
	if data.ID == 0 {
		return ErrMissingField("id")
	}
	if len(data.Modes) == 0 {
		return ErrMissingField("modes")
	}

	var userData struct {
		Username   string
		Privileges uint64
	}
	err := md.DB.Get(&userData, "SELECT username, privileges FROM users WHERE id = ?", data.ID)
	switch err {
	case sql.ErrNoRows:
		return common.SimpleResponse(404, "That user could not be found!")
	case nil: // carry on
	default:
		md.Err(err)
		return Err500
	}

	if common.UserPrivileges(userData.Privileges)&common.AdminPrivilegeManageUsers != 0 {
		return common.SimpleResponse(403, "Can't edit that user")
	}

	tx, err := md.DB.Beginx()
	if err != nil {
		md.Err(err)
		return Err500
	}

	for _, mode := range data.Modes {
		if mode < 0 || mode > 3 {
			continue
		}
		_, err = tx.Exec("INSERT INTO scores_removed SELECT * FROM scores WHERE userid = ? AND play_mode = ?", data.ID, mode)
		if err != nil {
			md.Err(err)
		}
		_, err = tx.Exec("DELETE FROM scores WHERE userid = ? AND play_mode = ?", data.ID, mode)
		if err != nil {
			md.Err(err)
		}
		_, err = tx.Exec(strings.Replace(
			`UPDATE users_stats SET total_score_MODE = 0, ranked_score_MODE = 0, replays_watched_MODE = 0,
			playcount_MODE = 0, avg_accuracy_MODE = 0, total_hits_MODE = 0, level_MODE = 0, pp_MODE = 0
			WHERE id = ?`, "MODE", modesToReadable[mode], -1,
		), data.ID)
		if err != nil {
			md.Err(err)
		}
	}

	if err = tx.Commit(); err != nil {
		md.Err(err)
		return Err500
	}

	rapLog(md, fmt.Sprintf("has wiped %s's account", userData.Username))

	return userPutsSingle(md, md.DB.QueryRowx(userFields+" WHERE users.id = ? LIMIT 1", data.ID))
}

func appendToUserNotes(md common.MethodData, message string, user int) {
	message = "\n[" + time.Now().Format("2006-01-02 15:04:05") + "] API: " + message
	_, err := md.DB.Exec("UPDATE users SET notes = CONCAT(COALESCE(notes, ''), ?) WHERE id = ?",
		message, user)
	if err != nil {
		md.Err(err)
	}
}

func usernameAvailable(md common.MethodData, u string, userID int) (r bool) {
	err := md.DB.QueryRow("SELECT EXISTS(SELECT 1 FROM users WHERE username_safe = ? AND id != ?)", common.SafeUsername(u), userID).Scan(&r)
	if err != nil && err != sql.ErrNoRows {
		md.Err(err)
	}
	return
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// +build darwin

package client

import (
	"github.com/keybase/cli"
	"github.com/keybase/client/go/install"
	"github.com/keybase/client/go/libcmdline"
	"github.com/keybase/client/go/libkb"
)

func NewCmdCtlStart(cl *libcmdline.CommandLine, g *libkb.GlobalContext) cli.Command {
	return cli.Command{
		Name:  "start",
		Usage: "Start the keybase service",
		Flags: []cli.Flag{
			// Using autofork bool instead of enum to be more like other autofork flags?
			// https://github.com/keybase/client/pull/2414
			cli.BoolFlag{
				Name:  "auto-fork",
				Usage: "Use auto forking",
			},
		},
		Action: func(c *cli.Context) {
			cl.ChooseCommand(NewCmdCtlStartRunner(g), "start", c)
			cl.SetForkCmd(libcmdline.NoFork)
			cl.SetLogForward(libcmdline.LogForwardNone)
		},
	}
}

type CmdCtlStart struct {
	libkb.Contextified
	autoFork bool
}

func NewCmdCtlStartRunner(g *libkb.GlobalContext) *CmdCtlStart {
	return &CmdCtlStart{Contextified: libkb.NewContextified(g)}
}

func (s *CmdCtlStart) ParseArgv(ctx *cli.Context) error {
	s.autoFork = ctx.Bool("auto-fork")
	return nil
}

func (s *CmdCtlStart) Run() error {
	if s.autoFork {
		_, err := AutoForkServer(s.G(), s.G().Env.GetCommandLine())
		return err
	}

	return StartLaunchdService(s.G(), string(install.AppServiceLabel), true)
}

func (s *CmdCtlStart) GetUsage() libkb.Usage {
	return libkb.Usage{}
}
<commit_msg>Remove auto fork flag (use env GetAutoFork)<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

// +build darwin

package client

import (
	"github.com/keybase/cli"
	"github.com/keybase/client/go/install"
	"github.com/keybase/client/go/libcmdline"
	"github.com/keybase/client/go/libkb"
)

func NewCmdCtlStart(cl *libcmdline.CommandLine, g *libkb.GlobalContext) cli.Command {
	return cli.Command{
		Name:  "start",
		Usage: "Start the keybase service",
		Action: func(c *cli.Context) {
			cl.ChooseCommand(NewCmdCtlStartRunner(g), "start", c)
			cl.SetForkCmd(libcmdline.NoFork)
			cl.SetLogForward(libcmdline.LogForwardNone)
		},
	}
}

type CmdCtlStart struct {
	libkb.Contextified
}

func NewCmdCtlStartRunner(g *libkb.GlobalContext) *CmdCtlStart {
	return &CmdCtlStart{libkb.NewContextified(g)}
}

func (s *CmdCtlStart) ParseArgv(ctx *cli.Context) error {
	return nil
}

func (s *CmdCtlStart) Run() error {
	if s.G().Env.GetAutoFork() {
		_, err := AutoForkServer(s.G(), s.G().Env.GetCommandLine())
		return err
	}

	return StartLaunchdService(s.G(), string(install.AppServiceLabel), true)
}

func (s *CmdCtlStart) GetUsage() libkb.Usage {
	return libkb.Usage{}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"log/syslog"
	"os"
	"os/exec"
	"os/user"
	"path"
	"time"

	"github.com/dotcloud/docker/pkg/mount"
)

// Must be true everywhere.
const databoxGid = 10000

var pamUser = os.Getenv("PAM_USER")
var log, _ = syslog.NewLogger(syslog.LOG_WARNING | syslog.LOG_AUTH, 0)

func Fatal(first string, args ...interface{}) {
	// TODO(pwaller): send to syslog?
	log.Fatalf("pam script: "+first, args...)
}

func isDataboxUser() bool {
	u, err := user.Lookup(pamUser)
	if err != nil {
		Fatal("Failed to obtain passwd entry for %q", pamUser)
	}
	return u.Gid == fmt.Sprint(databoxGid)
}

func initMounts() {
	home := path.Join("/var/lib/cobalt/home/", pamUser)

	mounts := []struct{ src, tgt string }{
		{"/opt/basejail", "/jail"},
		{"/dev", "/jail/dev"},
		{"/dev/pts", "/jail/dev/pts"},
		{"/proc", "/jail/proc"},
		{"/var/spool/cron/crontabs", "/jail/var/spool/cron/crontabs"},
		{"/var/lib/extrausers", "/jail/var/lib/extrausers"},
		{home, "/jail/home"},
	}

	for _, m := range mounts {
		// Note the use of recursive bind mounts.
		// We could avoid some mounts by just arranging that /opt/basejail
		// already has most of the mounts.
		err := mount.Mount(m.src, m.tgt, "", "rbind")
		if err != nil {
			log.Fatalf("pamscript: Failed to mount %s -> %s: %q", m.src, m.tgt, err)
		}
	}
}

func cgcreate() error {
	args := []string{"-t", pamUser, "-g", "memory,cpu,cpuacct:" + pamUser}
	cmd := exec.Command("cgcreate", args...)
	return cmd.Run()
}

func initCgroup() {
	_, err := exec.LookPath("cgcreate")
	if err != nil {
		// cgroups isn't installed on this machine, NOOP.
		return
	}

	if _, err := os.Stat(path.Join("/sys/fs/cgroup/cpu", pamUser)); err != nil {
		err = cgcreate()
		if err != nil {
			Fatal("Failed to create cgroup")
		}
	}

	// 512MiB
	const memoryLimit = 512 * 1024 * 1024

	f := path.Join("/sys/fs/cgroup/memory", pamUser, "/memory.limit_in_bytes")
	err = ioutil.WriteFile(f, []byte(fmt.Sprint(memoryLimit)), 0)
	if err != nil {
		Fatal("Failed to write", f, ":", err)
	}

	// echo $MemoryLimit > /sys/fs/cgroup/memory/$PAM_USER/memory.limit_in_bytes

	// TODO(pwaller): do we want this? Maybe? Is there some other way we can give
	// system things priority?

	// # CPU share is form of priority. By specifying a low number here, we
	// # ensure that important system services get a higher share of the CPU
	// # and thus remain responsive.
	// Priority=12
	// echo $Priority > /sys/fs/cgroup/cpu/$PAM_USER/cpu.shares

	// Put the owning process (usually the "su -l" or cron child process)
	// into the cgroup (and therefore all of its future children)

	files := []string{
		path.Join("/sys/fs/cgroup/cpu", pamUser, "/tasks"),
		path.Join("/sys/fs/cgroup/memory", pamUser, "/tasks"),
		path.Join("/sys/fs/cgroup/cpuacct", pamUser, "/tasks"),
	}

	parentPid := []byte(fmt.Sprint(os.Getppid()))
	for _, f := range files {
		err = ioutil.WriteFile(f, parentPid, 0)
		if err != nil {
			Fatal("Failed to write", f, ":", err)
		}
	}
}

func verifyMountNamespace() {
	// if [[ "$(readlink /proc/1/ns/mnt)" == "$(readlink /proc/self/ns/mnt)" ]]; then
	initMountNS, err := os.Readlink("/proc/1/ns/mnt")
	if err != nil {
		Fatal("Unable to readlink(/proc/1/ns/mnt). Aborting.")
	}
	myMountNS, err := os.Readlink("/proc/self/ns/mnt")
	if err != nil {
		Fatal("Unable to readlink(/proc/self/ns/mnt). Aborting.")
	}
	if initMountNS == myMountNS {
		Fatal("Not in mount namespace. Abort.")
	}
}

func main() {
	start := time.Now()
	defer func() {
		log.Println("pam_script_ses_open took", time.Since(start))
	}()

	if !isDataboxUser() {
		log.Println("Skip non-databox user")
		// skip non-databox login
		return
	}

	if pamUser == "" {
		Fatal("PAM_USER not set. Abort.")
	}

	verifyMountNamespace()

	initCgroup()
	initMounts()
}
<commit_msg>Remove clutter from pam_script_ses_open logging<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	logpkg "log"
	"log/syslog"
	"os"
	"os/exec"
	"os/user"
	"path"
	"time"

	"github.com/dotcloud/docker/pkg/mount"
)

// Must be true everywhere.
const databoxGid = 10000

var pamUser = os.Getenv("PAM_USER")

var log *logpkg.Logger

func init() {
	os.Args[0] = "pam_script_ses_open"
	log, _ = syslog.NewLogger(syslog.LOG_WARNING|syslog.LOG_AUTH, 0)
}

func Fatal(first string, args ...interface{}) {
	// TODO(pwaller): send to syslog?
	log.Fatalf("pam script: "+first, args...)
}

func isDataboxUser() bool {
	u, err := user.Lookup(pamUser)
	if err != nil {
		Fatal("Failed to obtain passwd entry for %q", pamUser)
	}
	return u.Gid == fmt.Sprint(databoxGid)
}

func initMounts() {
	home := path.Join("/var/lib/cobalt/home/", pamUser)

	mounts := []struct{ src, tgt string }{
		{"/opt/basejail", "/jail"},
		{"/dev", "/jail/dev"},
		{"/dev/pts", "/jail/dev/pts"},
		{"/proc", "/jail/proc"},
		{"/var/spool/cron/crontabs", "/jail/var/spool/cron/crontabs"},
		{"/var/lib/extrausers", "/jail/var/lib/extrausers"},
		{home, "/jail/home"},
	}

	for _, m := range mounts {
		// Note the use of recursive bind mounts.
		// We could avoid some mounts by just arranging that /opt/basejail
		// already has most of the mounts.
		err := mount.Mount(m.src, m.tgt, "", "rbind")
		if err != nil {
			log.Fatalf("pamscript: Failed to mount %s -> %s: %q", m.src, m.tgt, err)
		}
	}
}

func cgcreate() error {
	args := []string{"-t", pamUser, "-g", "memory,cpu,cpuacct:" + pamUser}
	cmd := exec.Command("cgcreate", args...)
	return cmd.Run()
}

func initCgroup() {
	_, err := exec.LookPath("cgcreate")
	if err != nil {
		// cgroups isn't installed on this machine, NOOP.
		return
	}

	if _, err := os.Stat(path.Join("/sys/fs/cgroup/cpu", pamUser)); err != nil {
		err = cgcreate()
		if err != nil {
			Fatal("Failed to create cgroup")
		}
	}

	// 512MiB
	const memoryLimit = 512 * 1024 * 1024

	f := path.Join("/sys/fs/cgroup/memory", pamUser, "/memory.limit_in_bytes")
	err = ioutil.WriteFile(f, []byte(fmt.Sprint(memoryLimit)), 0)
	if err != nil {
		Fatal("Failed to write", f, ":", err)
	}

	// echo $MemoryLimit > /sys/fs/cgroup/memory/$PAM_USER/memory.limit_in_bytes

	// TODO(pwaller): do we want this? Maybe? Is there some other way we can give
	// system things priority?

	// # CPU share is form of priority. By specifying a low number here, we
	// # ensure that important system services get a higher share of the CPU
	// # and thus remain responsive.
	// Priority=12
	// echo $Priority > /sys/fs/cgroup/cpu/$PAM_USER/cpu.shares

	// Put the owning process (usually the "su -l" or cron child process)
	// into the cgroup (and therefore all of its future children)

	files := []string{
		path.Join("/sys/fs/cgroup/cpu", pamUser, "/tasks"),
		path.Join("/sys/fs/cgroup/memory", pamUser, "/tasks"),
		path.Join("/sys/fs/cgroup/cpuacct", pamUser, "/tasks"),
	}

	parentPid := []byte(fmt.Sprint(os.Getppid()))
	for _, f := range files {
		err = ioutil.WriteFile(f, parentPid, 0)
		if err != nil {
			Fatal("Failed to write", f, ":", err)
		}
	}
}

func verifyMountNamespace() {
	// if [[ "$(readlink /proc/1/ns/mnt)" == "$(readlink /proc/self/ns/mnt)" ]]; then
	initMountNS, err := os.Readlink("/proc/1/ns/mnt")
	if err != nil {
		Fatal("Unable to readlink(/proc/1/ns/mnt). Aborting.")
	}
	myMountNS, err := os.Readlink("/proc/self/ns/mnt")
	if err != nil {
		Fatal("Unable to readlink(/proc/self/ns/mnt). Aborting.")
	}
	if initMountNS == myMountNS {
		Fatal("Not in mount namespace. Abort.")
	}
}

func main() {
	start := time.Now()
	defer func() {
		log.Println("took", time.Since(start))
	}()

	if !isDataboxUser() {
		log.Println("Skip non-databox user")
		// skip non-databox login
		return
	}

	if pamUser == "" {
		Fatal("PAM_USER not set. Abort.")
	}

	verifyMountNamespace()

	initCgroup()
	initMounts()
}
<|endoftext|>


<commit_before>package teams

import (
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/stretchr/testify/require"
	"testing"
)

func TestTeamRoleMap(t *testing.T) {
	fus, tcs, cleanup := setupNTests(t, 2)
	defer cleanup()

	t.Logf("create team")
	teamName, teamID := createTeam2(*tcs[0])
	m := make([]libkb.MetaContext, 2)
	for i, tc := range tcs {
		m[i] = libkb.NewMetaContextForTest(*tc)
	}
	t.Logf("add B to the team so they can load it")
	_, err := AddMember(m[0].Ctx(), tcs[0].G, teamName.String(), fus[1].Username, keybase1.TeamRole_ADMIN, nil)
	require.NoError(t, err)
	subteamName := "abc"
	subteamID, err := CreateSubteam(m[0].Ctx(), tcs[0].G, subteamName, teamName, keybase1.TeamRole_WRITER)
	require.NoError(t, err)

	received, err := m[1].G().GetTeamRoleMapManager().Get(m[1], false)
	require.NoError(t, err)

	expected := keybase1.TeamRoleMapAndVersion{
		Teams:   make(map[keybase1.TeamID]keybase1.TeamRolePair),
		Version: keybase1.UserTeamVersion(1),
	}
	expected.Teams[teamID] = keybase1.TeamRolePair{
		Role:         keybase1.TeamRole_ADMIN,
		ImplicitRole: keybase1.TeamRole_NONE,
	}
	expected.Teams[*subteamID] = keybase1.TeamRolePair{
		Role:         keybase1.TeamRole_NONE,
		ImplicitRole: keybase1.TeamRole_ADMIN,
	}

	require.Equal(t, expected, received)

	_, err = AddMember(m[0].Ctx(), tcs[0].G, teamName.String()+"."+subteamName, fus[1].Username, keybase1.TeamRole_READER, nil)
	require.NoError(t, err)

	expected.Teams[*subteamID] = keybase1.TeamRolePair{
		Role:         keybase1.TeamRole_READER,
		ImplicitRole: keybase1.TeamRole_ADMIN,
	}
	expected.Version++

	// In teams/*_test.go, we don't have gregor hooked up, so we have to mock out what happens
	// when a new gregpr notification comes down from the server.
	err = m[1].G().GetTeamRoleMapManager().Update(m[1], expected.Version)
	require.NoError(t, err)

	// Check that the state is eagerly refreshed.
	pollForTrue(t, m[1].G(), func(i int) bool {
		received := m[1].G().GetTeamRoleMapManager().(*TeamRoleMapManager).Query()
		if received != nil && received.Data.Version == expected.Version {
			require.Equal(t, expected.Teams, received.Data.Teams)
			return true
		}
		return false
	})

	m[1].G().GetTeamRoleMapManager().FlushCache()

	received, err = m[1].G().GetTeamRoleMapManager().Get(m[1], false)
	require.NoError(t, err)
	require.Equal(t, expected, received)
}
<commit_msg>teams: fix teamRoleMap test<commit_after>package teams

import (
	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/stretchr/testify/require"
	"testing"
)

func TestTeamRoleMap(t *testing.T) {
	fus, tcs, cleanup := setupNTests(t, 2)
	defer cleanup()

	t.Logf("create team")
	teamName, teamID := createTeam2(*tcs[0])
	m := make([]libkb.MetaContext, 2)
	for i, tc := range tcs {
		m[i] = libkb.NewMetaContextForTest(*tc)
	}
	t.Logf("add B to the team so they can load it")
	_, err := AddMember(m[0].Ctx(), tcs[0].G, teamName.String(), fus[1].Username, keybase1.TeamRole_ADMIN, nil)
	require.NoError(t, err)
	subteamName := "abc"
	subteamID, err := CreateSubteam(m[0].Ctx(), tcs[0].G, subteamName, teamName, keybase1.TeamRole_WRITER)
	require.NoError(t, err)

	received, err := m[1].G().GetTeamRoleMapManager().Get(m[1], false)
	require.NoError(t, err)

	expected := keybase1.TeamRoleMapAndVersion{
		Teams:   make(map[keybase1.TeamID]keybase1.TeamRolePair),
		Version: keybase1.UserTeamVersion(2),
	}
	expected.Teams[teamID] = keybase1.TeamRolePair{
		Role:         keybase1.TeamRole_ADMIN,
		ImplicitRole: keybase1.TeamRole_NONE,
	}
	expected.Teams[*subteamID] = keybase1.TeamRolePair{
		Role:         keybase1.TeamRole_NONE,
		ImplicitRole: keybase1.TeamRole_ADMIN,
	}

	require.Equal(t, expected, received)

	_, err = AddMember(m[0].Ctx(), tcs[0].G, teamName.String()+"."+subteamName, fus[1].Username, keybase1.TeamRole_READER, nil)
	require.NoError(t, err)

	expected.Teams[*subteamID] = keybase1.TeamRolePair{
		Role:         keybase1.TeamRole_READER,
		ImplicitRole: keybase1.TeamRole_ADMIN,
	}
	expected.Version++

	// In teams/*_test.go, we don't have gregor hooked up, so we have to mock out what happens
	// when a new gregpr notification comes down from the server.
	err = m[1].G().GetTeamRoleMapManager().Update(m[1], expected.Version)
	require.NoError(t, err)

	// Check that the state is eagerly refreshed.
	pollForTrue(t, m[1].G(), func(i int) bool {
		received := m[1].G().GetTeamRoleMapManager().(*TeamRoleMapManager).Query()
		if received != nil && received.Data.Version == expected.Version {
			require.Equal(t, expected.Teams, received.Data.Teams)
			return true
		}
		return false
	})

	m[1].G().GetTeamRoleMapManager().FlushCache()

	received, err = m[1].G().GetTeamRoleMapManager().Get(m[1], false)
	require.NoError(t, err)
	require.Equal(t, expected, received)
}
<|endoftext|>


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

package worker

import (
	"bytes"
	"fmt"
	"strings"

	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/stats"

	tabletmanagerdatapb "github.com/youtube/vitess/go/vt/proto/tabletmanagerdata"
)

// RowAggregator aggregates SQL reconciliation statements into one statement.
// Once a limit (maxRows or maxSize) is reached, the statement will be sent to
// the destination's insertChannel.
// RowAggregator is also aware of the type of statement (DiffType) and
// constructs the necessary SQL command based on that.
// Aggregating multiple statements is done to improve the overall performance.
// One RowAggregator instance is specific to one destination shard and diff
// type.
// Important: The SQL statement generation assumes that the fields of the
// provided row are in the same order as orderedColumns(td) would return on the
// provided "td" (TableDefinition). That means primary key fields come first and
// then all other fields. The order of these two sets of fields is identical
// to the order in td.PrimaryKeyColumns and td.Columns respectively.
type RowAggregator struct {
	ctx           context.Context
	maxRows       int
	maxSize       int
	insertChannel chan string
	td            *tabletmanagerdatapb.TableDefinition
	diffType      DiffType
	builder       QueryBuilder
	// statsCounters has Counters to track how many rows were changed per
	// diffTyp.
	statsCounters []*stats.Counters

	buffer       bytes.Buffer
	bufferedRows int
}

// NewRowAggregator returns a RowAggregator.
// The index of the elements in statCounters must match the elements
// in "DiffTypes" i.e. the first counter is for inserts, second for updates
// and the third for deletes.
func NewRowAggregator(ctx context.Context, maxRows, maxSize int, insertChannel chan string, dbName string, td *tabletmanagerdatapb.TableDefinition, diffType DiffType, statsCounters []*stats.Counters) *RowAggregator {
	// Construct head and tail base commands for the reconciliation statement.
	var builder QueryBuilder
	switch diffType {
	case DiffMissing:
		// Example: INSERT INTO test (id, sub_id, msg) VALUES (0, 10, 'a'), (1, 11, 'b')
		builder = NewInsertsQueryBuilder(dbName, td)
	case DiffNotEqual:
		// Example: UPDATE test SET msg='a' WHERE id=0 AND sub_id=10
		builder = NewUpdatesQueryBuilder(dbName, td)
		// UPDATE ... SET does not support multiple rows as input.
		maxRows = 1
	case DiffExtraneous:
		// Example: DELETE FROM test WHERE (id, sub_id) IN ((0, 10), (1, 11))
		builder = NewDeletesQueryBuilder(dbName, td)
	default:
		panic(fmt.Sprintf("unknown DiffType: %v", diffType))
	}

	if len(statsCounters) != len(DiffTypes) {
		panic(fmt.Sprintf("statsCounter has the wrong number of elements. got = %v, want = %v", len(statsCounters), len(DiffTypes)))
	}

	return &RowAggregator{
		ctx:           ctx,
		maxRows:       maxRows,
		maxSize:       maxSize,
		insertChannel: insertChannel,
		td:            td,
		diffType:      diffType,
		builder:       builder,
		statsCounters: statsCounters,
	}
}

// Add will add a new row which must be reconciled.
// If an error is returned, RowAggregator will be in an undefined state and must
// not be used any longer.
func (ra *RowAggregator) Add(row []sqltypes.Value) error {
	if ra.buffer.Len() == 0 {
		ra.builder.WriteHead(&ra.buffer)
	}

	if ra.bufferedRows >= 1 {
		ra.builder.WriteSeparator(&ra.buffer)
	}
	ra.builder.WriteRow(&ra.buffer, row)
	ra.bufferedRows++

	if ra.bufferedRows >= ra.maxRows || ra.buffer.Len() >= ra.maxSize {
		if err := ra.Flush(); err != nil {
			return err
		}
	}

	return nil
}

// Flush sends out the current aggregation buffer.
func (ra *RowAggregator) Flush() error {
	if ra.buffer.Len() == 0 {
		// Already flushed.
		return nil
	}

	ra.builder.WriteTail(&ra.buffer)
	// select blocks until sending the SQL succeeded or the context was canceled.
	select {
	case ra.insertChannel <- ra.buffer.String():
	case <-ra.ctx.Done():
		return fmt.Errorf("failed to flush RowAggregator and send the query to a writer thread channel: %v", ra.ctx.Err())
	}

	// Update our statistics.
	ra.statsCounters[ra.diffType].Add(ra.td.Name, int64(ra.bufferedRows))

	ra.buffer.Reset()
	ra.bufferedRows = 0
	return nil
}

// QueryBuilder defines for a given reconciliation type how we have to
// build the SQL query for one or more rows.
type QueryBuilder interface {
	// WriteHead writes the beginning of the query into the buffer.
	WriteHead(*bytes.Buffer)
	// WriteTail writes any required tailing string into the buffer.
	WriteTail(*bytes.Buffer)
	// Write the separator between two rows.
	WriteSeparator(*bytes.Buffer)
	// Write the row itself.
	WriteRow(*bytes.Buffer, []sqltypes.Value)
}

// BaseQueryBuilder partially implements the QueryBuilder interface.
// It can be used by other QueryBuilder implementations to avoid repeating
// code.
type BaseQueryBuilder struct {
	head      string
	tail      string
	separator byte
}

// WriteHead implements the QueryBuilder interface.
func (b *BaseQueryBuilder) WriteHead(buffer *bytes.Buffer) {
	buffer.WriteString(b.head)
}

// WriteTail implements the QueryBuilder interface.
func (b *BaseQueryBuilder) WriteTail(buffer *bytes.Buffer) {
	buffer.WriteString(b.tail)
}

// WriteSeparator implements the QueryBuilder interface.
func (b *BaseQueryBuilder) WriteSeparator(buffer *bytes.Buffer) {
	buffer.WriteByte(b.separator)
}

// InsertsQueryBuilder implements the QueryBuilder interface for INSERT queries.
type InsertsQueryBuilder struct {
	BaseQueryBuilder
}

// NewInsertsQueryBuilder creates a new InsertsQueryBuilder.
func NewInsertsQueryBuilder(dbName string, td *tabletmanagerdatapb.TableDefinition) *InsertsQueryBuilder {
	// Example: INSERT INTO test (id, sub_id, msg) VALUES (0, 10, 'a'), (1, 11, 'b')
	return &InsertsQueryBuilder{
		BaseQueryBuilder{
			head:      "INSERT INTO `" + dbName + "`." + td.Name + " (" + strings.Join(td.Columns, ", ") + ") VALUES ",
			separator: ',',
		},
	}
}

// WriteRow implements the QueryBuilder interface.
func (*InsertsQueryBuilder) WriteRow(buffer *bytes.Buffer, row []sqltypes.Value) {
	// Example: (0, 10, 'a'), (1, 11, 'b')
	buffer.WriteByte('(')
	for i, value := range row {
		if i > 0 {
			buffer.WriteByte(',')
		}
		value.EncodeSQL(buffer)
	}
	buffer.WriteByte(')')
}

// UpdatesQueryBuilder implements the QueryBuilder interface for UPDATE queries.
type UpdatesQueryBuilder struct {
	BaseQueryBuilder
	td *tabletmanagerdatapb.TableDefinition
	// nonPrimaryKeyColumns is td.Columns filtered by td.PrimaryKeyColumns.
	// The order of td.Columns is preserved.
	nonPrimaryKeyColumns []string
}

// NewUpdatesQueryBuilder creates a new UpdatesQueryBuilder.
func NewUpdatesQueryBuilder(dbName string, td *tabletmanagerdatapb.TableDefinition) *UpdatesQueryBuilder {
	// Example: UPDATE test SET msg='a' WHERE id=0 AND sub_id=10
	//
	// Note: We cannot use INSERT INTO ... ON DUPLICATE KEY UPDATE here because
	// it's not recommended for tables with more than one unique (i.e. the
	// primary key) index. That's because the update function would also be
	// triggered if a unique, non-primary key index matches the row. In that
	// case, we would update the wrong row (it gets selected by the unique key
	// and not the primary key).
	return &UpdatesQueryBuilder{
		BaseQueryBuilder: BaseQueryBuilder{
			head: "UPDATE `" + dbName + "`." + td.Name + " SET ",
		},
		td: td,
		// Build list of non-primary key columns (required for update statements).
		nonPrimaryKeyColumns: orderedColumnsWithoutPrimaryKeyColumns(td),
	}
}

// WriteSeparator implements the QueryBuilder interface and overrides
// the BaseQueryBuilder implementation.
func (b *UpdatesQueryBuilder) WriteSeparator(buffer *bytes.Buffer) {
	panic("UpdatesQueryBuilder does not support aggregating multiple rows in one query")
}

// WriteRow implements the QueryBuilder interface.
func (b *UpdatesQueryBuilder) WriteRow(buffer *bytes.Buffer, row []sqltypes.Value) {
	// Example: msg='a' WHERE id=0 AND sub_id=10
	nonPrimaryOffset := len(b.td.PrimaryKeyColumns)
	for i, column := range b.nonPrimaryKeyColumns {
		if i > 0 {
			buffer.WriteByte(',')
		}
		buffer.WriteString(column)
		buffer.WriteByte('=')
		row[nonPrimaryOffset+i].EncodeSQL(buffer)
	}
	buffer.WriteString(" WHERE ")
	for i, pkColumn := range b.td.PrimaryKeyColumns {
		if i > 0 {
			buffer.WriteString(" AND ")
		}
		buffer.WriteString(pkColumn)
		buffer.WriteByte('=')
		row[i].EncodeSQL(buffer)
	}
}

// DeletesQueryBuilder implements the QueryBuilder interface for DELETE queries.
type DeletesQueryBuilder struct {
	BaseQueryBuilder
	td *tabletmanagerdatapb.TableDefinition
}

// NewDeletesQueryBuilder creates a new DeletesQueryBuilder.
func NewDeletesQueryBuilder(dbName string, td *tabletmanagerdatapb.TableDefinition) *DeletesQueryBuilder {
	// Example: DELETE FROM test WHERE (id, sub_id) IN ((0, 10), (1, 11))
	return &DeletesQueryBuilder{
		BaseQueryBuilder: BaseQueryBuilder{
			head:      "DELETE FROM `" + dbName + "`." + td.Name + " WHERE (" + strings.Join(td.PrimaryKeyColumns, ", ") + ") IN (",
			tail:      ")",
			separator: ',',
		},
		td: td,
	}
}

// WriteRow implements the QueryBuilder interface.
func (b *DeletesQueryBuilder) WriteRow(buffer *bytes.Buffer, row []sqltypes.Value) {
	// Example: (0, 10), (1, 11)
	buffer.WriteByte('(')
	for i := 0; i < len(b.td.PrimaryKeyColumns); i++ {
		if i > 0 {
			buffer.WriteByte(',')
		}
		row[i].EncodeSQL(buffer)
	}
	buffer.WriteByte(')')
}
<commit_msg>worker: Fixed typo in comment.<commit_after>// Copyright 2016, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package worker

import (
	"bytes"
	"fmt"
	"strings"

	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/stats"

	tabletmanagerdatapb "github.com/youtube/vitess/go/vt/proto/tabletmanagerdata"
)

// RowAggregator aggregates SQL reconciliation statements into one statement.
// Once a limit (maxRows or maxSize) is reached, the statement will be sent to
// the destination's insertChannel.
// RowAggregator is also aware of the type of statement (DiffType) and
// constructs the necessary SQL command based on that.
// Aggregating multiple statements is done to improve the overall performance.
// One RowAggregator instance is specific to one destination shard and diff
// type.
// Important: The SQL statement generation assumes that the fields of the
// provided row are in the same order as orderedColumns(td) would return on the
// provided "td" (TableDefinition). That means primary key fields come first and
// then all other fields. The order of these two sets of fields is identical
// to the order in td.PrimaryKeyColumns and td.Columns respectively.
type RowAggregator struct {
	ctx           context.Context
	maxRows       int
	maxSize       int
	insertChannel chan string
	td            *tabletmanagerdatapb.TableDefinition
	diffType      DiffType
	builder       QueryBuilder
	// statsCounters has Counters to track how many rows were changed per
	// diffType.
	statsCounters []*stats.Counters

	buffer       bytes.Buffer
	bufferedRows int
}

// NewRowAggregator returns a RowAggregator.
// The index of the elements in statCounters must match the elements
// in "DiffTypes" i.e. the first counter is for inserts, second for updates
// and the third for deletes.
func NewRowAggregator(ctx context.Context, maxRows, maxSize int, insertChannel chan string, dbName string, td *tabletmanagerdatapb.TableDefinition, diffType DiffType, statsCounters []*stats.Counters) *RowAggregator {
	// Construct head and tail base commands for the reconciliation statement.
	var builder QueryBuilder
	switch diffType {
	case DiffMissing:
		// Example: INSERT INTO test (id, sub_id, msg) VALUES (0, 10, 'a'), (1, 11, 'b')
		builder = NewInsertsQueryBuilder(dbName, td)
	case DiffNotEqual:
		// Example: UPDATE test SET msg='a' WHERE id=0 AND sub_id=10
		builder = NewUpdatesQueryBuilder(dbName, td)
		// UPDATE ... SET does not support multiple rows as input.
		maxRows = 1
	case DiffExtraneous:
		// Example: DELETE FROM test WHERE (id, sub_id) IN ((0, 10), (1, 11))
		builder = NewDeletesQueryBuilder(dbName, td)
	default:
		panic(fmt.Sprintf("unknown DiffType: %v", diffType))
	}

	if len(statsCounters) != len(DiffTypes) {
		panic(fmt.Sprintf("statsCounter has the wrong number of elements. got = %v, want = %v", len(statsCounters), len(DiffTypes)))
	}

	return &RowAggregator{
		ctx:           ctx,
		maxRows:       maxRows,
		maxSize:       maxSize,
		insertChannel: insertChannel,
		td:            td,
		diffType:      diffType,
		builder:       builder,
		statsCounters: statsCounters,
	}
}

// Add will add a new row which must be reconciled.
// If an error is returned, RowAggregator will be in an undefined state and must
// not be used any longer.
func (ra *RowAggregator) Add(row []sqltypes.Value) error {
	if ra.buffer.Len() == 0 {
		ra.builder.WriteHead(&ra.buffer)
	}

	if ra.bufferedRows >= 1 {
		ra.builder.WriteSeparator(&ra.buffer)
	}
	ra.builder.WriteRow(&ra.buffer, row)
	ra.bufferedRows++

	if ra.bufferedRows >= ra.maxRows || ra.buffer.Len() >= ra.maxSize {
		if err := ra.Flush(); err != nil {
			return err
		}
	}

	return nil
}

// Flush sends out the current aggregation buffer.
func (ra *RowAggregator) Flush() error {
	if ra.buffer.Len() == 0 {
		// Already flushed.
		return nil
	}

	ra.builder.WriteTail(&ra.buffer)
	// select blocks until sending the SQL succeeded or the context was canceled.
	select {
	case ra.insertChannel <- ra.buffer.String():
	case <-ra.ctx.Done():
		return fmt.Errorf("failed to flush RowAggregator and send the query to a writer thread channel: %v", ra.ctx.Err())
	}

	// Update our statistics.
	ra.statsCounters[ra.diffType].Add(ra.td.Name, int64(ra.bufferedRows))

	ra.buffer.Reset()
	ra.bufferedRows = 0
	return nil
}

// QueryBuilder defines for a given reconciliation type how we have to
// build the SQL query for one or more rows.
type QueryBuilder interface {
	// WriteHead writes the beginning of the query into the buffer.
	WriteHead(*bytes.Buffer)
	// WriteTail writes any required tailing string into the buffer.
	WriteTail(*bytes.Buffer)
	// Write the separator between two rows.
	WriteSeparator(*bytes.Buffer)
	// Write the row itself.
	WriteRow(*bytes.Buffer, []sqltypes.Value)
}

// BaseQueryBuilder partially implements the QueryBuilder interface.
// It can be used by other QueryBuilder implementations to avoid repeating
// code.
type BaseQueryBuilder struct {
	head      string
	tail      string
	separator byte
}

// WriteHead implements the QueryBuilder interface.
func (b *BaseQueryBuilder) WriteHead(buffer *bytes.Buffer) {
	buffer.WriteString(b.head)
}

// WriteTail implements the QueryBuilder interface.
func (b *BaseQueryBuilder) WriteTail(buffer *bytes.Buffer) {
	buffer.WriteString(b.tail)
}

// WriteSeparator implements the QueryBuilder interface.
func (b *BaseQueryBuilder) WriteSeparator(buffer *bytes.Buffer) {
	buffer.WriteByte(b.separator)
}

// InsertsQueryBuilder implements the QueryBuilder interface for INSERT queries.
type InsertsQueryBuilder struct {
	BaseQueryBuilder
}

// NewInsertsQueryBuilder creates a new InsertsQueryBuilder.
func NewInsertsQueryBuilder(dbName string, td *tabletmanagerdatapb.TableDefinition) *InsertsQueryBuilder {
	// Example: INSERT INTO test (id, sub_id, msg) VALUES (0, 10, 'a'), (1, 11, 'b')
	return &InsertsQueryBuilder{
		BaseQueryBuilder{
			head:      "INSERT INTO `" + dbName + "`." + td.Name + " (" + strings.Join(td.Columns, ", ") + ") VALUES ",
			separator: ',',
		},
	}
}

// WriteRow implements the QueryBuilder interface.
func (*InsertsQueryBuilder) WriteRow(buffer *bytes.Buffer, row []sqltypes.Value) {
	// Example: (0, 10, 'a'), (1, 11, 'b')
	buffer.WriteByte('(')
	for i, value := range row {
		if i > 0 {
			buffer.WriteByte(',')
		}
		value.EncodeSQL(buffer)
	}
	buffer.WriteByte(')')
}

// UpdatesQueryBuilder implements the QueryBuilder interface for UPDATE queries.
type UpdatesQueryBuilder struct {
	BaseQueryBuilder
	td *tabletmanagerdatapb.TableDefinition
	// nonPrimaryKeyColumns is td.Columns filtered by td.PrimaryKeyColumns.
	// The order of td.Columns is preserved.
	nonPrimaryKeyColumns []string
}

// NewUpdatesQueryBuilder creates a new UpdatesQueryBuilder.
func NewUpdatesQueryBuilder(dbName string, td *tabletmanagerdatapb.TableDefinition) *UpdatesQueryBuilder {
	// Example: UPDATE test SET msg='a' WHERE id=0 AND sub_id=10
	//
	// Note: We cannot use INSERT INTO ... ON DUPLICATE KEY UPDATE here because
	// it's not recommended for tables with more than one unique (i.e. the
	// primary key) index. That's because the update function would also be
	// triggered if a unique, non-primary key index matches the row. In that
	// case, we would update the wrong row (it gets selected by the unique key
	// and not the primary key).
	return &UpdatesQueryBuilder{
		BaseQueryBuilder: BaseQueryBuilder{
			head: "UPDATE `" + dbName + "`." + td.Name + " SET ",
		},
		td: td,
		// Build list of non-primary key columns (required for update statements).
		nonPrimaryKeyColumns: orderedColumnsWithoutPrimaryKeyColumns(td),
	}
}

// WriteSeparator implements the QueryBuilder interface and overrides
// the BaseQueryBuilder implementation.
func (b *UpdatesQueryBuilder) WriteSeparator(buffer *bytes.Buffer) {
	panic("UpdatesQueryBuilder does not support aggregating multiple rows in one query")
}

// WriteRow implements the QueryBuilder interface.
func (b *UpdatesQueryBuilder) WriteRow(buffer *bytes.Buffer, row []sqltypes.Value) {
	// Example: msg='a' WHERE id=0 AND sub_id=10
	nonPrimaryOffset := len(b.td.PrimaryKeyColumns)
	for i, column := range b.nonPrimaryKeyColumns {
		if i > 0 {
			buffer.WriteByte(',')
		}
		buffer.WriteString(column)
		buffer.WriteByte('=')
		row[nonPrimaryOffset+i].EncodeSQL(buffer)
	}
	buffer.WriteString(" WHERE ")
	for i, pkColumn := range b.td.PrimaryKeyColumns {
		if i > 0 {
			buffer.WriteString(" AND ")
		}
		buffer.WriteString(pkColumn)
		buffer.WriteByte('=')
		row[i].EncodeSQL(buffer)
	}
}

// DeletesQueryBuilder implements the QueryBuilder interface for DELETE queries.
type DeletesQueryBuilder struct {
	BaseQueryBuilder
	td *tabletmanagerdatapb.TableDefinition
}

// NewDeletesQueryBuilder creates a new DeletesQueryBuilder.
func NewDeletesQueryBuilder(dbName string, td *tabletmanagerdatapb.TableDefinition) *DeletesQueryBuilder {
	// Example: DELETE FROM test WHERE (id, sub_id) IN ((0, 10), (1, 11))
	return &DeletesQueryBuilder{
		BaseQueryBuilder: BaseQueryBuilder{
			head:      "DELETE FROM `" + dbName + "`." + td.Name + " WHERE (" + strings.Join(td.PrimaryKeyColumns, ", ") + ") IN (",
			tail:      ")",
			separator: ',',
		},
		td: td,
	}
}

// WriteRow implements the QueryBuilder interface.
func (b *DeletesQueryBuilder) WriteRow(buffer *bytes.Buffer, row []sqltypes.Value) {
	// Example: (0, 10), (1, 11)
	buffer.WriteByte('(')
	for i := 0; i < len(b.td.PrimaryKeyColumns); i++ {
		if i > 0 {
			buffer.WriteByte(',')
		}
		row[i].EncodeSQL(buffer)
	}
	buffer.WriteByte(')')
}
<|endoftext|>